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/

Теперь у нас есть две хорошие модели с мультиязычностью и function calling (писал про эту комбинацию тут https://www.tgoop.com/gonzo_ML/2821), Mistral NeMo и Llama 3.1. Ожидаю интересных развитий и решений!