Метод

Общий фреймворк выглядит следующим образом:
▶️ Функция преобразования данных 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%)

Выводы

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

Transformer Layers as Painters
Qi Sun, Marc Pickett, Aakash Kumar Nain, Llion Jones
Статья: https://arxiv.org/abs/2407.09298
Twitter: https://x.com/A_K_Nain/status/1812684597248831912
Код: https://github.com/floatingbigcat/transformer-as-painter (пока нет)

Вторая работа про творческий подход к вычислению трансформерных слоёв. Первая была про LayerShuffle (https://www.tgoop.com/gonzo_ML/2845), в текущей работе от коллег GDE подход более радикальный, чем просто перемешивание слоёв.

Авторы предлагают интересную метафору для средних слоёв сети — конвейер художников. Входное полотно проходит через цепочку художников. Некоторые из них специализируются на рисовании птиц, другие на рисовании колёс. Получая полотно от предыдущего художника, каждый решает дорисовать ли чего где или отправить дальше без изменений. У всех художников общий словарь для понимания рисунков, так что не страшно, если художник получит на вход полотно от более раннего чем обычно художника. Их также можно поменять местами, большой катастрофы не произойдёт (даже если части фона будут дорисованы поверх уже нарисованных объектов). Художники даже могут рисовать одновременно впараллель.

Аналогия не претендует на научную точность, но помогает поднять кучу интересных вопросов: Используют ли слои общее пространство репрезентаций? Все ли слои необходимы? Вычисляют ли все средние слои одну и ту же функцию? Важен ли порядок слоёв? Можно ли вычислять слои впараллель? Более ли важен порядок вычисления для одних задач чем для других? Помогают ли циклы параллельно вычисляемым слоям? Какие варианты наименее вредят качеству?

Хорошие вопросы. Предыдущие работы типа LayerDrop и LayerShuffle покрывают лишь небольшую часть из этого.

LayerDrop проверяли на ViT, текущая работа сделана на предобученных LLM (декодерная Llama2 7B с 32 слоями и энкодерный BERT-Large 340M с 24 слоями), никакого файнтюнинга (кроме оценки на GLUE, где подразумевается файнтюнинг для берта). Для ламы проверяли на ARC (который от AI2, https://arxiv.org/abs/1803.05457, а не от Франсуа Шолле, https://github.com/fchollet/ARC-AGI), HellaSwag, GSM8K, WinoGrande, LAMBADA. Для берта вышеупомянутый GLUE.

Если посмотреть на результирующее качество при выкидывании слоя или при перестановке слоя с соседним, то оказывается, что начальный и конечный слои важны, а между ними всё более-менее устойчиво к таким манипуляциям. Выглядит так, что срединные слои используют общее пространство репрезентаций. Если дополнительно посмотреть на косинусную близость активаций слоёв внутри модели, то групп слоёв может даже больше: входной слой 0, слои 1-3, средние слои, последний или пара последних. Возможно у модели три разных пространства репрезентаций: для начальных, средних и конечных слоёв.

Если пропустить M слоёв (выкидываем слои из середины от N+1 до T-N, где T — это общее число слоёв в модели), то качество постепенно деградирует по мере увеличения M. Получается, что не все промежуточные слои необходимы, по крайней мере несколько средних слоёв могут быть выкинуты без катастрофической деградации.

Если в предыдущем эксперименте вместо выкидывания средних слоёв заменить их на копию центрального слоя, то результат получается намного хуже, чем если бы выкинуть. Это наиболее сильная деградация среди всех рассмотренных в работе. Отсюда вывод, что средние слои реализуют разные функции. В приложении дополнительно смотрят на косинусные близости и статистики активаций и делают вывод, что повторение среднего слоя выталкивает вход из общего пространства репрезентаций. Если художник рисует колёса, то рисование большего числа колёс на полотне приводит к тому, что художники после него будут вынуждены работать с тем, на чём не обучались.

Выходит, что средние слои разделяют общее пространство репрезентаций, но реализуют разные операции в этом пространстве. Отсюда возникает вопрос, важен ли порядок этих операций? Для проверки средние слои исполняют в 1) обратном или 2) рандомном порядке (с усреднением по 10 сидам). В обоих случаях деградация постепенная и оба варианта лучше, чем просто выкидывание слоёв. Значит слои всё ещё вносят свой вклад, хоть и вызываются в другом порядке. Случайный порядок при этом лучше обратного. Получается, что до некоторой степени порядок слоёв важен, но деградация от случайного и обратного порядка постепенная.

Интересный вопрос про порядок, а можно ли слои вычислить впараллель, а затем смёржить (через усреднение)? Снова постепенная деградация кроме GSM8K, там быстрая. Это работает лучше, чем пропуск слоёв, но хуже, чем обратный порядок. В итоге вычислять слои параллельно можно, кроме случаев нагруженных математикой бенчмарков.

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

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

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

Хорошая, легко читаемая работа, которую в принципе мог сделать любой в своём гараже, не требуется быть для этого гуглом. Это круто, есть ещё место для независимых исследователей. Авторы хотят копнуть глубже и понять, почему трансформеры устойчивы к таким вариациям. Одна возможная гипотеза в том, что residual connections в обучении необходимы для того, чтобы слои шарили общие репрезентации. Интересно будет проверить на моделях без таких соединений. Также планируют разморозить модели и посмотреть как быстро модели восстанавливаются во время файнтюнинга. Варианты с параллельным вычислением и пропуском слоёв хороши с точки зрения латенси, здесь может быть практическая польза. Ещё потенциально интересная тема — роутинг для выбора слоёв по аналогии со Switch Transformers (https://www.tgoop.com/gonzo_ML/472).

Just in case, вдруг кто не видел

https://www.youtube.com/watch?v=_cZa_7KaQ3c

Я сначала не узнал дока...

Детали соревнования тут:
https://ai.google.dev/competition

Официально анонсировали Llama 3.1, версии 8B, 70B и 405B.

Тулы, мультиязычность, контекст 128k

Сайт: https://llama.meta.com
Пост: https://ai.meta.com/blog/meta-llama-3-1/
Статья: https://ai.meta.com/research/publications/the-llama-3-herd-of-models/