1.02K subscribers
62 photos
3 files
148 links
Квантование & Прунинг & Дистилляция

Блог про сжатие сетей и не только.
От древнейших времен по настоящее время.
Download Telegram
Бог создал слабые и сильные методы сжатия, но файтьюн уравнял их.
Прунинг LLM в глубину

Данный пост был в замыслах уже пару месяцев, но все никак не доходили руки.

Ниже приведен разбор нескольких статей, посвященных одной и той же теме - depth pruning. Так как задача и мотивация у всех статей общая, то и введение дедуплицированное, общее для всех работ.

Введение

При структурированном прунинге веса отбрасываются не поодиночке, а группами - поканально, поголовно 🗿 (трансформерно поголовно), иногда даже весь слой целиком отправляется в царство Аида ☠️. Структурированный прунинг хорош тем, что явно уменьшает размерности матриц и тензоров в операциях, потому дает неплохое ускорение на разнообразном железе и движках. Увы, серьезного сжатия без серьезной просадки в качестве обычно сложно достичь без серьезных вложений в дообучение.

Трансформеры содержат skip-connection как в Attention, так и в MLP. Residual type архитектуры, как было замечено еще давно , являются в некотором смысле ансамблями, где конечный выход можно воспринимать как агрегацию знания от нескольких неглубоких экспертов. Одна голова хорошо, две лучше, еще больше - еще лучше, но и иногда и пары специалистов достаточно. Потому есть основания полагать, что прунинг в глубину - прореживание целых блоков - имеет шанс завестись и сжатая модель будет выдавать адекватную точность.
Shortened LLaMA: A Simple Depth Pruning for Large Language Models
[Статья][Код]

Первая статья в хронологическом порядке и с открытым исходным кодом.

Метод

Хотим мы, значится, выбрасывать блоки целиком - но какие брать?
Авторы рассматривают 3 критерия
1️⃣ Среднюю величину весов (Magntude pruning)
2️⃣ Taylor (не Свифт, а норму градиента на вес, вдв кароч)
3️⃣ Изменение перплексии при выкидывании каждого блока по отдельности
В первых двух случаях временами выпиливаются первые и последние блоки, которые оказываются позарез важными, поэтому предлагается убрать их рассмотрения.

Чтобы восстановить качество модели после сжатия, навешивают LoRA на то, что выжило и дообучают на небольшом количестве данных.

Эксперименты

Берут LLama-1 и Vicuna-7B, 13B. Для калибровки (оценки важности блоков) берут 10 последовательностей длины 128 из BookCorpus, и потом дообучают на Alpaca.

Предложенный метод не хуже, а то и лучше LLM-Pruner и структурированный Wanda (где критерием важности параметра выступает норма веса на норму активации).

Просадка при 20% прореживания заметная, но все же довольно умеренная. На малых степенях сжатия лучше всего работает критерий по градиенту умножить на вес, при запрете на отбрасывание первых четырех и последних двух блоков, при больших - оценка важности по перплексии.

Файнтьюн с LoRA неплохо восстанавливает метрики после сжатия, особенно с увеличением степени сжатия. Тем не менее разрыв между исходной моделью и сжатой все еще существенен (просадки больше, чем при SOTA 2-битной квантизации).

Данных требуется совсем немного для калибровки, качество существенно не меняется если взять больше, чем 10 примеров. Запрещать сжимать для Taylor и Magnitude первые и последние блоки важно - иначе перплексия взлетает до нескольких тысяч.

Комментарий

Метод прост и логичен. Однако просадка, на самом деле несколько больше, чем декларируется. На собственных замерах с lm_eval=0.4.0 поверх дообученных чекпоинтов (26 блоков против исходных 32) замерил просадку в 3% по сравнению с исходной моделью.
ShortGPT: Layers in Large Language Models are More Redundant Than You Expect
[Статья][Нет кода]

Метод

Трансформерные блоки, как известно обладают аддитивными skip connections следующего вида (для простоты изложения забьем на нормализацию):

x = x + attn(x)
x = x + mlp(x)


Выход mlp, attn нередко инициализируется нулем и остается небольшим по норме в процессе обучения. Следовательно, каждый блок меняет понемножку внутренне представление.

В данной работе предлагают выбрасывать те блоки, которые изменяют меньше всего внутреннее представление на некоторой выборке (примерах из PG19). В качестве расстояния используют среднее косинусное расстояние между токенами перед и после данного блока. В итоге убираются блоки с наименьшим косинусным расстоянием.

Результаты

Так как работа от китайских коллег, то замеряют качество как на Llama-2, так и на Baichuan. Для оценки качества используют стандарнтые 0-шоты, MMLU и CMMLU.

Предложенная метрика выбирает блоки ближе к концу сети (но не самые последние).

Результаты бенчмарков вызывают вопросы, утверждается, что качество на MMLU почти не просаживается вплоть до 28% sparsity, при этом перплексия на wikitext2 возрастает довольно заметно. При таких значениях обычно сеть с трудом два слова связать может, не то что решать задачи на логику). lm-eval-harness нигде не цитируется. По всей видимости используют какой-то иной или самописный фреймворк.

Безлайны - методы структурированного прунинга в ширину, естественно, бьют. С ShortenedLLama, вышедшей чуть раньше, не сравниваются.
The Unreasonable Ineffectiveness of the Deeper Layers
[Статья][Нет кода]

Метод

Идея по существу та же, что и выше, но выбрасывают блоки не по отдельности, а пачкой - т.е находим группу последовательных трансформерных блоков, которые минимально меняют представление. Критерий тот же самый - косинусная близость.

Как и в ShortenedLlama для восстановления качества сжатой модели дообучают с LoRA. И дообучение называется лечением (healing 🚑).

Результаты

Замеряют качество языкового моделирование (вместо перплексии смотрят на изменение кросс-энтропии) и MMLU/BoolQ. Рассматривают Llama-2, Mistral, Phi-2.

Кросс энтропия просаживается, но неплохо лечится после файнтьюна.

Результаты на MMLU вызывают большие сомнения 🤔 - спад качества имеет скачкообразный характер. По личным наблюдениям спад MMLU довольно монотонный, постепенный и коррелировал с ростом перплексии. Утверждается, что некоторые модели можно сжать до 40% без просадки на этом бенчмарке.

Далее смотрят на корреляцию признаков в разных блоках. Наибольшая корреляция (наименьшее расстояние) в блоках ближе к концу, согласуясь с наблюдениями из двух прошлых работ. Отсюда предлагают простое правило - пруньте блоки с конца, исключая последний, и типа хорошо работает.
Your Transformer is Secretly Linear
[Статья][Пост в телеге][Пост на Хабре][Код]

Интересное релевантное исследование от коллег из AIRI/Сколтеха. Кратко, ибо все хорошо и доступно изложено самим автором в постах в телеге и на Хабре.

Метод

Оказывается, что активации трансформера в соседних блоках связаны почти что линейным преобразованием. В качестве меры линейности используется linearity score, который по существу является r2 (коэффициентом детерминации) . Сам по себе факт не столь удивителен, ибо норма skip-connection обыкновенно значительно больше преобразования в attn/mlp. Но даже если вычесть skip connection выделяется явная линейная и нелинейная компонента.

Проанализировав ряд промежуточных чекпоинтов открытых моделей, авторы замечают что норма линейной компоненты уменьшается в процессе предобучения, но возрастает в ходе файнтьюна . И предложенная регуляризация на нелинейность улучшает качество.

Исходя из наблюдений, предлагается прунить самые “линейные” слои, а для восстановления качества заместо запруненного блока учится линейный слой. И такой адаптер неплохо восстанавливает качество.

Эксперименты

Рассматривают большое число открытых моделей - от GPT-2 (когда-то OpenAI был действительно открытым), Bloom до свежих Phi-3.

Рост линейности имеет место для всех моделей и задач. Прирост качества от регуляризации хоть и невелик, но все же заметен. Интересно, насколько он будет иметь место на большом масштабе?

Несколько нетипичный сетап замера перплексии - короткие последовательности вместо обычно используемого 2к/4к контекста (как на обучении), но общая суть от сего не меняется.
Вывод

Трансформеры, вероятно, являясь в каком-то смысле ансамблем, устойчивы, как печень, к ампутации крупных кусков. Однако просадки (там где результатам можно доверять - первая и последняя работа), все же значительны при степени сжатия более 20%. Как ни крути, scaling laws существуют не просто так, и представления, выученные разными блоками, хоть имеют и разную полезность, но все же содержат в себе релевантное знание, хоть для какого-то входа. Отсюда и успех контекстуальных методов сжатия, вроде Deja Vu и PowerInfer.
QTIP: Quantization with Trellises and Incoherence Processing
[Статья][Кода нет, но обещают, что будетт]

Было очевидно, что после выхода обновленного AQLM и PV-tuning, что ответочка от конкурентов из Cornell University - это лишь вопрос времени. Тем более, что в issue в PEFT первый автор QuIP# дал ясно понять, что новый, более совершенный, метод квантизации скоро выйдет на свет.

И тихо, незаметно, так что сам я обнаружил лишь по чистой случайности, на архиве появилась статья и классическая репа с ридмишкой.

Метод

Векторная квантизация при той же битности точнее (если реализована не криворуко), чем скалярная. Раньше я думал, что данное явление обусловлено некоторыми корреляциями между измерениями, но на самом деле векторная квантизация работает точнее даже для последовательностей i.i.d (хорошее обьяснение здесь) . Чем больше группа - тем более высокой точности при заданном числе бит можно добиться потенциально. Однако, большие кодбуки становятся довольно обьемными и вносят нетривиальный вклад в общую битность (в AQLM 1x16 кодбуки накидывают ~0.3 дополнительных бит на параметр) и не влезают в кэш GPU.

В идеале бы хотелось иметь компактный кодбук и большие группы. И авторы QTIP (названного в честь рэпера?) прибегают к красивой идее из теории кодирования - trellis coding.

Обычно при кодировке последовательностей в L бит, каждый элемент может принимать любое из 2^L значений, и последовательность длины кодируется TL битами. Если представить последовательность в виде блуждания по графу, то граф полносвязный - из любой вершины можно прийти в любую другую.

В trellis coding (названном в честь архитектурной решетки), каждый узел графа (соответствующий одному из 2^L) значений соединен только с 2^k (где k<L) ребрами с другими вершинами. Последовательность кодируется как индекс начального узла и индексы ребер при блуждании по графу. Итого расход памяти L + (T-1)k, что может быть заметно меньше TL.

Однако, наивная реализация будет чрезмерно тяжеловесной для произвольного trellis при достаточно большом L и k, так как где-то надо хранить граф связей. Кроме того, декодирование для произвольного trellis - много последовательных операций, что непозволительно медленно на практике, и QTIP предлагает несколько хитрых идей, чтобы это работало эффективно.

Как и в QuIP/QuIP# квантуется не исходный вес, а обработанный случайным ортогональным преобразованием (Адамаровыми матрицами), так что элементы матриц близки к i.i.d нормальному распределению.

Далее, чтобы не хранить полный произвольный граф, авторы рассматривают специальный анзатц - bit-shift trellis, где связаны только вершины с индексами, отличающимися на битовый сдвиг. Однако, такой выбор слишком ограничителен, и авторы предлагают “перемешивать” связи некоторой хитрой случайной перестановкой и предлагают три метода:
1️⃣ 1MAD
2️⃣ 3INST
3️⃣ HYB
В первых двух случаях в качестве индексов в случайной перестановке выступает некоторое фиксированное преобразование использующее линейный конгруэнтный генератор случайных чисел и побитовые операции. А в третьем случае - некоторое обучаемое преобразование.

Для определения оптимальной конфигурации минимизируется ошибка на выходе слоя с помощью алгоритма Витерби (динамическое программирование).

Полученный метод выдает MSE ошибку близкую к минимально теоретически возможной.
Эксперименты

Метод валидируют на 🦙-2, 3, и в качестве бейзлайнов выступает QuIP# и AQLM первой версии. С PV-tuning не сравниваются, шельмецы 👺, хоть и лежит он на архиве 3 недели с их публикации) Квантуют группами 16x16 (число весов по входной и выходной размерности) с L=16, и Q=9 (гиперпараметр в обучаемом гибридном коде). Разные битности отличаются числом ребер k из каждой вершины.

Метод демонстрирует выдающееся качество даже без какого-то дообучения, в 4 бита почти без просадки, а в 2 бита почти на уровне AQLM и QuIP# c end-to-end finetuning (жулики в FT репортят цифры AQLM только с блочным файтьюном 😈 ).

С полноценным файнтьюном (дообучение знаков в Адамаровых матрицах и параметров HYB преобразования) метод вырывается далеко вперед при низких битностях от QuIP#/AQLM. Неупомянутый PV-tuning работает несколько лучше на меньших моделях - 7(8B), 13B, и сравнивается на 70B (используя, правда, больше данных и вычислений).

На 🦙-3 просадки больше, данное семейство моделей, согласно наблюдениям извне и личному опыту, тяжелее квантизуются.

Скорость инференса (замеряют на RTX4090) на уровне QuIP#. Сравнение с AQLM полная шляпа 🎩, чуваки, небось коряво поставили либу для инференса, либо что еще наворотили, ибо tokens/s в нашей статье (на более медленной RTX3090) куда больше.

Выводы

Несмотря на корявое сравнение с AQLM, в общем и целом, конкуренты проделали отличную работу. Здорово, когда красивые идеи из математики находят применение на практике. Кажется, что стоит реально засесть за ботанье 👨‍🏫 матана и чтение статей из 80-х)
Accelerating Neural Network Training with Semi-Structured (2:4) Sparsity
[Блог]

Торч недавно выкатил блог, где показывает, что 2:4 sparsity паттерн может быть полезен не только для инференса, но и на стадии обучения модели.

Напомню, что 2:4 sparsity - это когда из 4 подряд идущих весов 2-нулевые. Начиная с Ampere нвидиевские GPU имеют поддержку такого формата. Веса хранятся в виде сжатого вдвое тензора и еще есть тензор индексов элементов.

Ранее было показано, что такой паттерн выдает ускорение 10-30% по сравнению с fp16 на инференсе, но чтобы добиться эффективного обучения необходимо учесть еще ряд нюансов.

Обучение с фиксированной маской работает не очень, и нужно ее обновлять по ходу дела чтобы не просесть по качеству. Наивный kernel слишком медленный и по порядку величины близок к времени forward пасса. И заслуга авторов в разработке эффективной процедуры прореживания.

1️⃣ На обратном проходе при подсчете градиента участвует не W, W.T, которая вообще говоря не 2:4. Авторы ограничивают паттерн прореживания так, что 2:4 паттерн имеет место как для W, так и для W.T (прореживая блоки 4x4). Еще есть нюанс в том, что на уровне железа реализовано умножение sparse матриц на dense, но не наоборот. Потому предлагается считать транспонированное умножение, а затем его транспонировать. В cuSPARSELt есть опция выбирать индексирование в памяти как по колонкам, так и по столбцам, что позволяет иметь непрерывное расположение элементов в обоих случаях.

2️⃣ Потоки CUDA считывают за раз по 128 байт, и чтобы отпимизировать операции чтения/записи каждый поток берет на себя по 4 блока 4x4 (4x4x4x2 (fp16) = 128 байт).

3️⃣ if/else операции, когда разные потоки в группе (warp) попадают в разные условные ветви - работают сильно неэффективно на GPU. Используется sorting network для определения важных/неважных элементов без условных операторов.

4️⃣ Сжатые матрицы и метаданные (индексы ненулевых элементов) хранятся в column-major для оптимизации операций записи.

Эксперименты

Обучают DINOv2 (ViT-L) c разными пропорциями sparse/dense обучения. Сначала идет sparse обучение, а затем dense. Для оценки качества смотрят на точность лог-регрессии на ImageNet-1k поверх признаков обученной модели.

Обучение в sparse режиме на протяжении 40-70% обучения не просаживается по качеству по сути по сравнению с dense обучением.
Всегда sparse - теряет полпроцента качества.

Ускорение времени обучения при 70% sparse обучения порядка 6% (на А100).

Вывод

2:4 паттерн не дает какого-то впечатляющего ускорения, но может слегка повысить эффективность. Интересно 🤔, насколько полученные результаты масштабируются на LLMки с миллиардами параметров. Будет ли выигрыш от 2:4 компенсироваться более медленной сходимостью?
Real-Time Video Generation with Pyramid Attention Broadcast
[Cтатьи Нет][Код есть]

В ряде отечественных тг каналов (пост на эйай ньюз, пост на Machine Learning) упомянули проект Real-Time Video Generation with Pyramid Attention Broadcast.

Самой папиры еще нет и потому какие-то нюансы могут быть неизвестны.
От себя добавлю несколько деталей.

Суть подхода заключается в следующем. В видео диффузии есть 3 вида attention операций:

1️⃣ Пространственное 🌌
2️⃣ Временное
3️⃣ Перекрестное внимание на condition 🤞

Ранее в Cache Me If You Can было замечено, что карты attention между соседними шагами диффузии мало меняются между соседними шагами на большей части процесса (за исключением начала и конца) при картиночной генерации. В данном проекте авторы замечают, что для скорости изменения attention карт справедливо следующее неравенство:

v_cross < v_time < v_spatial

И соотвественно, чем медленее меняется attention, тем чаще он переиспользуется.

За счет переиспользования карт attention можно параллелить эффективно между разными GPU разные шаги генерации по времени (уменьшении оверхеда на 50% без переиспользования карт).

Переиспользование карт дает ускорение ~30%. А распаралелливание на 8 GPU дает почти линейное ускорение (8.4x-10.6x) по сравнению с генерацией на одной GPU наивным способом. То есть ускорение достигается в первую очередь за счет эффективного параллелизма.

Тем не менее, достойный инженерный результат.
Sparse-Marlin
[Код]

Некоторое время назад я затрагивал Marlin - быстрый кернел для батчового инференса INT4xFP16.

Пару месяцев назад коллеги из IST выпустили SparseMarlin - где дополнительно к квантизации весов добавляется 2:4 sparsity, за счет чего достигается еще большее ускорение по сравнению с fp16.

Как я понял, основные инженерные наработки следующие:
1️⃣ Эффективная обработка 2:4 метаданных - позиций нулевых и ненулевых весов
2️⃣ Использование Sparse Tensor Cores (SPTCs)
3️⃣ Умное расположение квантованных весов, метаданных для sparsity, и статистик квантования

Итоговое ускорение до 5.3x 🏃‍♂️ на RTX3090 (у Marlin ~4x) при перемножении больших матриц. Интересно🥴 , какое end-2-end ускорение может быть достигнуто на эффективном движке для инференса типа vLLM?

На Hopper и более новые архитектуры пока не завезли

Таким образом, 2:4 может давать определенный выигрыш по скорости. Основная проблема в том, что на текущий момент 2:4 прунинг (post-training) сильно просаживает качество LLM ☹️.
Sparse maximal update parameterization: A holistic approach to sparse training dynamics
[Папира][Кода нет]

Введение

Подбор шага обучения - занятие довольно потное 🥲 и утомительное 😞. Learning rate, хороший для меньших моделей, зачастую приводит к расходимости для больших. Потому большие модели обычно обучают с меньшим learning rate. Но подбор по сетке (или с помощью байесовской оптимизации) может оказаться слишком накладным на практике, а иметь хорошие значения гиперпараметров все-таки хочется.

В Tensor Programs V предложили параметризацию весов, градиентов и learning rate (μP), такую, что оптимальные гиперпараметры обучения переносятся с маленьких моделей на большие.

Динамика обучения в разреженных сетях существенно отличается от плотных. Потому естественно предположить, что обобщаемая параметризация на разные степени прореживания потребует дополнительных телодвижений. И в данной статье выводят такую параметризацию.

Метод

Параметризация, при которой learning rate и прочие гиперпараметры обучения будут
переноситься на разные размеры моделей и степени прореживания, должна удовлетворять следующим свойствам:

1️⃣ Нормы весов не зависят от ширины и sparsity
2️⃣ Нормы градиентов по весам не зависят от ширины и sparsity
3️⃣ Изменения весов (для произвольного алгоритма оптимизации) не зависят от ширины и sparsity

Из нехитрой математики следует, что variance весов при иниациализации и learning rate следует масштабировать как 1 / (ширину_слоя * (1 - sparsity)). Чем шире сеть - тем меньше разброс параметров и learning rate, чем более разреженная сеть - тем, наоборот, больше разброс параметров и learning rate. При sparsity = 0, предложенная SμPar вырождается в μP.

Эксперименты

Метод валидируют на LMках c SwiGLU и Alibi, обучаемых на токенах из Slim Pajama. В качестве маленькой прокси модели для тюнинга гиперпараметрво берут LMку с 40M параметрами, а основную серию экспериментов проводят на 610M модели. Замеряют loss на обучении (датасет большой, потому переобучения нет).

При стандартной параметризации и μP оптимальный learning rate приходится подбирать для каждой степени сжатия свой, в то время, как для SμPar оптимальные learning rate зафиксирован.

Перенесенные с меньшей модели параметры на большую для разных степеней прореживания позволяют достичь лучшего качества по сравнению с попытками затюнить стандартную параметризацию и μP.

Авторы - ребята из Cerebras, потому умеют вполне эффективно эксплуатировать разреженные операции 😉.

Вывод

Полезная идея при обучении разреженных сетей. Ранее в Sparsity Scaling Laws было показано, что при очень большом числе данных sparse модели более compute оптимальны, чем плотные. Если в будущем будут активно учить LLMки с использованием специализированного железа SμPar будет весьма кстати для подбора гиперпараметров.
Fast as CHITA: Neural Network Pruning with Combinatorial Optimization
[Статья][Без Кода]

CHITA - это гепард, а не административный центр Забайкальского края.
Статья 2023 года от Google Research про прунинг.

Введение

Optimal 🧠 Surgeon, использующий приближение 2️⃣ порядка для определения оптимальной сжатой конфигурации, лежит в основе многих методов прунинга и квантизации. Алгоритм Iterative Hard Thresholding (IHT), в котором на каждом шаге делается шаг оптимизации и прунинг самых маленьких весов, лежит в основе некоторых методов sparse training - в частности, AC/DC.

В данной статье решили поженить OBS и IHT и назвали полученную сущность CHITA (Combinatorial Hessian-free Iterative Thresholding Algorithm).

Метод

Основная проблема с Гессианом нейронной сети, что его хрен посчитаешь честно. Существуют различные приближения; в данной статье опираются на Фишеровское (оценка Гессиана сумой внешних произведений градиентов). Обычно градиентов (n) много меньше, чем параметров (d), потому полученная матрица низкоранговая. Можно ее даже не материализовывать, так как в итоговом алгоритме потребуются только матрично-векторные произведения, а хранить лишь n градиентов в матрице A \in R^{n x d}. Фишеровская матрица выражается как A^T A.

CHITA на каждом шаге делает шаг градиентного спуска для квадратичного разложения в окрестности оптимума (т.е L(w) = L_0 + g^ T + 1/2 w^T H w) с последующим прореживанием как в IHT. Метод требует времени и памяти линейной по числу параметров и сохраненных градиентов.

Оптимальный шаг оптимизации (learning rate в IHT) находят с помощью умного 🧠 алгоритма поиска на прямой.

Эксперименты

Метод валидируют на небольших CNNках на CIFAR-10 и ImageNet. Метод быстрее M-FAC и при этом достигает лучшего качества. Однако честность сравнения в one-shot pruning, учитывая, что метод итеративный, в отличие от M-FAC, вызывает вопросы.

Метод можно применять итеративно, постепенно повышая степень сжатия и дообучая сеть с фиксированной маской. Лучше повышать линейно уровень прореживания линейно, чем экспоненциально или за раз пытаться сжать все.

При итеративном сжатии удается добиться 1% просадки на MobileNet-v1 при 75% cжатии и 5% просадки при 89%. В этом режиме почему-то забывают в сравнениях про M-FAC 😂, который достигает близких результатов по качеству. Метод довольно неплох и на ResNet50 - с умеренными просадками при 90-95% sparsity.

Выводы

Метод, неплох и не слишком сложен в реализации. Основная проблема - необходимость хранить много градиентов, что не позволяет масштабировать метод на LLMки и интересные в 2024 году модели. У вашего покорного слуги есть одна идейка, как это обойти. Но это пока секрет)
Flash Attention 3 вышел!
[Статья][Блог]

Не прошло и года с выхода Flash Attention 2️⃣, как вышло продолжение, доведя серию до трилогии.

Метод

Основная идея первого Flash Attention в уменьшении передачи памяти 📝 между HBM и кэшом GPU, а второго - в отпимизации не matmul операций. Flash Attention 2 позволяет довольно эффективно использовать вычислительные возможности A100, но для более современной H100 все еще имеет место сильная недоутилизация - всего 35%.

В Flash Attention 3 отпимизируют Attention с учетом новых архитектурных особенностей Hopper+ - новых тензорных ядер с большим throughput, тензорным ускорителем памяти (Tensor Memory Accelerator) и поддержкой fp8.

Основных источника ускорения два:

1️⃣ Паралеллизация вычислений между GEMM (General Matrix Multiply) и иными операциями (softmax, RMSNorm).

Хоть по flops softmax и нормализации кажутся мизерными, они выполняются в 256 раз медленее на H100, потому могут занимать до половины времени матричного умножения. За их выполнение отвечают другие компоненты GPU, чем тензорные ядра для перемножения матриц, потому их можно обрабатывать параллельно.

2️⃣ Использование fp8.

fp8 позволяет почти удвоить скорость. Однако есть нюанс - выбросы, приводящие к большим ошибкам квантования. Чтобы избавиться от них используют incoherence processing из QuIP. Благодаря этому ошибка квантизации уменьшается до 2.6 раз.

Результаты

С fp16 удается достичь ускорения до 1.6-1.8 раз по сравнению с Flash Attention 2.

И под 1.2 PFLOPs c fp8 на больших последовательностях и с большими головами трансформера.

Обещают в будущем накатить в торч.

Вывод

Сильно! Ждем через год Flash Attention 4️⃣!
Vision language models are blind
[Статья][Код]

Наряду с чисто языковыми моделями в последнее время интенсивно развиваются и мультимодальные модели, умеющие работать с текстовой 📝 и визуальной 👀 информацией одновременно. Флагманские проприетарные модели (GPT-4(V/O), Claude, Gemini 1.5-Pro) неплохо умеют решать нетривиальные задачи, связанные с обработкой визуальных данных - понимание содержимого сцены, нахождение нужного объекта, извлечение информации из графиков и табличек, и самое главное - понимание мемов 😄.

Однако данные задачи еще более менее имеют формулировку с точки зрения естественного языка. Кроме того, имеет место утечка 💧 в процессе обучения. И авторы задаются вопросом - как хорошо современные модели справляются с задачи, требующие умение опираться сугубо на визуальную информацию, причем те, что под силу даже детсадовцу 👶?

И оказывается, что несмотря на простоту рассматриваемых задач, нынешняя SOTA не то чтобы блестяще справляется с ними. И GPT-4o не всегда лучшая)

Метод

Предложенный бенчмарк состоит из 7 задач:
1️⃣ Подсчет числа пересечений двух графиков из 3 точек (от 0 до 3)
2️⃣ Определение того, пересекаются ли 2 круга или нет
3️⃣ Определение буквы, на которую наложили круг 🔴
4️⃣ Подсчет числа пересекающихся геометрических примитивов
5️⃣ Подсчет количества вложенных квадратов
6️⃣ Определение числа строк и столбцов в таблице
7️⃣ Умение следовать за разноцветными ломаными (сколько путей следует из одной вершины в другую). Путей от 0 до 3.

Рассматривают 4 прориетарные VLM (Vision Language Model) модели:
1️⃣ GPT-4o
2️⃣ Gemini-1.5 Pro
3️⃣ Claude Sonnet 3
4️⃣ Claude Sonnet 3,5

Эксперименты

Несмотря на простоту перечисленных выше задач, модели справляются с ними далеко не идеально.

В частности, в задаче определения числа пересечений GPT-4o достигает качества 48%, а лучшая из моделей - Claude-3.5 Sonnet – 77%, что далеко от 100%. Чем ближе графики и чем толще линии - тем больше ошибка.

Модели хорошо понимают, что такое круг, и умеют читать текст без ошибок, но, тем не менее определение того, какую же все-таки букву закрыл кружочек, оказывается не такой уж простой задачей. И здесь лидирует Gemini-1.5 Pro (со средней точностью около 93%).

VLM-ки умеют хорошо считать разнесенные в пространстве объекты. Но как только они пересекаются, или оказываются вложенными - качество заметно проседает. И здесь снова побеждает в Claude-3.5 Sonnet.

Такая элементарная задача, как подсчет числа столбцов и строк, тоже дается нелегко, при том, что перечисленные сетки умеют решать куда, казалось бы, более сложные задачи про обработку табличных данных. Умница Claude-3.5 Sonnet выдает точность в среднем 74.26%, в то время как остальные модели между 35% и 40%. Если в ячейках есть текст, точность немного выше.

Определение числа путей по картинке, тоже дается нелегко. Claude-3.5 Sonnet (не подкупили Antropic авторов?) снова побеждает с точностью в среднем 50%. Если путь один, то Claude-3.5 с большим отрывом точнее всех, на большом числе путей - 3, GPT-4o лидирует.

Вывод

Любопытное исследование, хоть результаты и во многом ожидаемы. Текущие модели основаны преимущественно на достижениях в NLP, способности LLM улавливать сложные закономерности в последовательностях. Но задачи, опирающиеся сугубо на зрительную информацию, по всей видимости, требуют от модели обучения с упором на геометрию, выявления пространственных закономерностей. Вопрос 🤔 в том, насколько это важно в прикладных задачах. Тем не менее разработка такой процедуры обучения может стать следующим шагом в развитии VLM.
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
[Статья][Код]

Современные нейронные сети, обрабатывающие пространственно-временные данные различной природы будь то текст 📝, изображения 📷, аудио 🎵 и видео 📹 так или иначе обладают механизмом перемешивания каналов (channel mixing), обрабатывающим независимо признаки для каждого элемента последовательности, и механизмом обработки последовательности (sequence mixing), использования взаимосвязей между элементами.

В сегодняшнем рассказе речь пойдет про sequence mixing.

Существуют разнообразные опции sequence mixing. Операция смешивания может не зависеть от входа, как например свертка или обучаемая матрица L x L (L - длина последовательности) в MLP-Mixer, S4 и H3 state-space модели, или зависеть - attention механизм в трансформерах или Mamba (Selective State Spaces).

Кроме того, разные механизмы обладают разной сложностью от длины последовательности. Sequence mixing в Attention или MLP-Mixer требует квадратичного по длине последовательности числа элементарных операций с плавающей точкой (FLOPs), так как используют матричную операцию довольно общего вида. Sequence mixers, обладающие некоторой структурой (низкоранговые, Toeplitz матрицы, DFT, бабочки) позволяют добиваться субквадратичной сложности (обычно с некоторой просадкой в качестве).

И sequence mixing может быть как причинным (causal attention, большинство SSM, в частности, модная нынче Mamba 🐍), где текущий элемент последовательности может смотреть только в прошлое, и двунаправленным (как в masked language modelling, и большинстве задач с ViTами), где элементы последовательности могут изменять свое состояние, как глядя как на прошлые, так и на будущие токены.

И задача, которую, перед собой ставят авторы в данной работе - получение эффективного механизма двунаправленного sequence mixing, такого, чтобы он был с одной стороны субквадратичным (в идеале линейным по длине последовательности) и в то же время выразительным.
Метод

Общий фреймворк выглядит следующим образом:
▶️ Функция преобразования данных f_X (X - входные данные)
▶️ Функция конструкции смешивающей матрицы f_M L x L, которая может быть постоянной или зависеть от входов
▶️ Результат sequence mixing имеет вид f_M(f_X (X))

Далее авторы вводят термин Sequence Aligned Matrices (SAM, еще один… 🥱) означающий, что матрица смешивания зависит от входных данных. Такие sequece миксеры хороши с одной стороны тем, что более адаптивно подстраиваются под входы, и, кроме того, работают с последовательностями разной длины.

Авторы рассматривают разные механизмы из литературы:
⭐️ MLP-Mixer, S4, H3, Monarch, Vandermonde и Cauchy миксеры - не SAM
⭐️ Attention, Linear Attention, S6, SSD - SAM

Потому хороший двунаправленный sequence mixer должен быть SAM, и представлять собой некоторую структурированную матрицу. В частности, предлагаются способы сделать Vandermonde и Cauchy миксеры зависящими от входов, но основной упор делается на прокачку SSD (не твердотельного жесткого диска, а механизма в Mamba-2!) под двунаправленность.

Напомним, что SSD во второй мамбе является полуразложимыми (semi-separable) матрицей - каждый блок является низкоранговой матрицей. Для двунаправленности можно было бы чередовать слои SSD (один бегущий слева направо, другой справа налево), но здесь предлагают использование одной матрицы смешивания, такой что любой ее блок в верхне-треугольной и нижне-треугольной части является низкоранговой матрицей. Иначе говоря, получается нечто типа суммы исходной SSD из Mamba-2 (нижнетреугольной матрицы) и транспонированной по длине последовательности (верхнетреугольной матрицы) и диагональной части. Такие матрицы называют квазиразложимыми (QS). Данная модификация требует всего пары дополнительных строчек в реализации по сравнению с исходным SSD слоем (shift - сдвиг на один элемент, flip - разворот последовательности задом наперед, DX - диагональная добавка).

QS(X) = shift(SS(X)) + flip(shift(SS(flip(X)))) + DX

Называют гидрой, потому что много голов, как в SSD, и звучит красиво 😹.

Эксперименты

Метод валидируют на задаче Masked Language Modelling, где в качестве бейзлайнов берется BERT, обученный по рецепту от MosaicML, и иные варианты sequence mixerов из литературы. Для оценки качества моделей смотрят на валидационную кросс-энтропию на C4 (на train set которого обучают) и точность на бенчах из GLUE. Все модели имеют размер порядка 70M параметров (несколько меньше, чем BERT-Base), так что хрен вам SOTA на LMSYS. Hydra модели глубже трансформеров примерно в 2️⃣ раза при примерно том же числе параметров.

SAM модели стабильно опережают свои не SAM версии (Toeplitz, Cauchy, Vandermonde с параметрами, зависящими от входа, заметно точнее версии с обучаемыми, не зависящими). Hydra (естесна), лучше всех, и на втором месте любимый нами трансформер. Однако, памятуя о недавнем результате MobileLLM, где более глубокие и тонкие трансформеры, оказываются лучше по качеству более коротких и жирных при том же числе параметров, задаешься вопросом - можно ли устранить разрыв в качестве за счет изменения конфигурации трансформера 🤔?

Исходная Mamba-2 не очень сильна в MLM, так как умеет обрабатывать информацию только в одном направлении, и предложенный способ (Hydra) лучше вариантов с суммой, конкатенацией, и перемножением результатов двух мамб-2.

Далее метод проверяют на задаче классификации ImageNet-1k, где обучают модели размера порядка ViT-Base (87-91M параметров) и Hydra опережает ViT-B, Hyena и S4. Однако, ViT бейзлайн вызывает вопросы, ибо согласно их результатам ViT-B имеет top-1 точность 78.8%, а его EMA 80.6%, в то время, как c рецептом обучения из Swin, на который они ссылаются (унаследованный в свою очередь из DeiT) выдает 81.8% (их лучший результат 81.6%)
Выводы

Данная статья полезна с формальной точки зрения, ибо дает некоторый обобщенный взгляд на современные варианты нейросетевых архитектур обрабатывающие последовательности, и немножечко вкусненькой 😋 математики. Результаты экспериментов вызывают некоторые вопросы, да и валидация на небольшом, ограниченном сетапе. Как мне кажется, пока преждевременно говорить об “убийце 🔫 трансформеров”. Тем не менее, интересно и занимательно.
Вот мой скромный канал преодолел рубеж в 1️⃣0️⃣0️⃣0️⃣ подписчиков) 🍾

Когда я создавал его, я и не мечтал о таком количестве. Думал, что это будет глубоко андерграундная гаражная телеграм-канава на пару десятков заинтересованных, в силу специфики тематики и моей лени в отношении заморочек с PR. Но несколько упоминаний на знаменитых и влиятельных отечественных телегах про DL и связанные темы, за что огромное спасибо @che_shr_cat, @alexwortega, @nadlsk, @milana_shhanukova, открыли мои безделицы широкой аудитории.

Всем спасибо за то, что вы здесь) Мне очень приятно видеть ваши реакции, комментарии к постам. Желание опубликовать разбор очередной статьи, наверное, лучший стимул поглубже в ней разобраться.

Будем дальше стараться. Быть добру!