И для разнообразия не про LLM.

Если это не первоапрельская шутка, то Kawasaki показал концепт нового мотоцикла транспортного средства, Corleo:

https://youtu.be/vQDhzbTz-9k?si=oC8mOuc-KfpIqkNa

А кому этого мало, есть Docker контейнерные перевозки для людей:

https://youtu.be/RFZ1aGqzIN4?si=lofqVvmzlJDHjhyg

Блин, круто. Особенно Corleo.

Multi-Token Attention
Olga Golovneva, Tianlu Wang, Jason Weston, Sainbayar Sukhbaatar
Статья: https://arxiv.org/abs/2504.00927

Продолжаем разборы архитектур.

Как известно, веса внимания в классическом механизме внимания определяются одним вектором значений query и одним вектором значений key. Этот “single token attention” является своеобразным боттлнеком для отделения важных частей от всего остального. Новый подход Multi-Token Attention (MTA) позволяет устранить боттлнек и определять веса внимания на основе нескольких векторов query и keys одновременно

Напомним, что в стандартном внимании веса внимания определяются как softmax(QK/sqrt(d)). Для каждого токена есть вектор эмбеддинга, этот вектор проецируется в три отдельных вектора Q, K и V, и по скалярному произведению векторов Q и K различных токенов определяется их “похожесть” или “важность”. После нормализации на корень от размерности эмбеддинга и взятию софтмакса от результата получаются веса внимания A. Далее с этими весами взвешиваются и суммируются вектора V и генерятся новые эмбеддинги для каждого токена. На наличие множества голов, маски декодера и прочего мы в этом объяснении забиваем, если хотите лучше понять/вспомнить этот процесс, отсылаю к классике (https://jalammar.github.io/illustrated-transformer/).

Внутри и снаружи этого базового механизма внимания можно много чего модифицировать -- мы писали про температуру в софтмаксе (https://www.group-telegram.com/gonzo_ML.com/3013), про отказ от нормализации до или после слоёв внимания (https://www.group-telegram.com/gonzo_ML.com/3478), 100500 вариантов разреженного и прочего модифицированного внимания, которые даже перечислять долго (просто как пример -- Reformer, https://www.group-telegram.com/gonzo_ML.com/176, далее воспользуйтесь поиском по каналу). Текущая работа тоже где-то в этом пуле.

Допустим, мы хотим найти предложение, содержащее несколько элементов. Пусть для примера это будет предложение “Where did Alice see the rabbit?” и мы хотим найти одновременное упоминание Алисы и кролика, им соответствуют query вектора q_a и q_r. Стандартный механизм считает веса внимания описанным выше способом, мы можем “найти” места в контексте, содержащие эти слова, и нам надо бы проверить, что они находятся где-то в соседних позициях. Но стандартный механизм внимания не даёт этого сделать в пределах одного слоя (через увеличение глубины можно, но хотелось бы и без), поскольку никаких взаимодействий между отдельными attention maps в нём нет, и даже если мы используем отдельные головы внимания для обнаружения Алисы и кролика, то нет механизма для комбинирования этих весов внимания. Модификация внимания в MTA позволяет добавить это взаимодействие между attention maps для соседних позиций Q и K или между отдельными головами.

На уровне конкретных модификаций внутри стандартного механизма внимания появляются три новых блока:
1) key-query convolution: комбинирует несколько key и query внутри головы
2) head mixing convolution: шарит информацию между головами и усиливает важную
3) group normalization with depth scaling: улучшает поток градиентов

Key-query convolution перемешивает веса внимания от разных временных шагов и работает так: к логитам внимания перед софтсаксом (QK/sqrt(d)) применяется двумерная обучаемая свёртка по измерениям q и k, измерения для батча и голов внимания не трогаются. Каждая голова внимания учит свою свёртку. Внутри свёртки используется маска с занулением элементов, чтобы не залезать в будущее. Это был pre-softmax convolution, он будет использоваться по дефолту. Можно также сделать post-softmax convolution, тогда свёртка считается не поверх логитов, а уже после софтмакса. Это делает взаимодействия между весами внимания аддитивными, а не мультипликативными. Я кстати не до конца понял, почему они до софтмакса прям мультипликативные...

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

Итого, возможны четыре варианта блока MTA с разными комбинациями pre/post свёрток. Тут есть простор для оптимизации, так если оба варианта pre или post, то можно объединить это в одну трёхмерную свёртку.

Group normalization with depth scaling использует GroupNorm и независимый скейлинг для каждой головы по рецепту от Differential Transformer (https://arxiv.org/abs/2410.05258, может кстати тоже его разобрать?).

Эксперименты начинают с игрушечной задачи: модели дают последовательность блоков, каждый из N случайных букв. Далее следует L<N букв вопроса. Задача -- найти блок, содержащий все буквы из вопроса в любом порядке. Модель должна вывести все буквы целевого блока, или только его первый либо последний токен (три разные варианта задачи). Для стандартного трансформера задача сложная, так как требует L кусочков информации для определения целевого блока, и их надо закодировать в один query вектор. С MTA должно быть проще, так как он может найти позицию каждой буквы, а потом свёрткой увеличить вес внимания, если все L букв найдены вместе.

Проверили на N=5 и 8, L=2. Пример задачи (надо найти блок с pb):

hjnvt.qfjgt.whftb.bjtpq. ...(many blocks)... .pxjvf.ulhik.qoiax#pb

Обучали на 1M таких блоков, тестировали на отложенных 1K. Трансформер 4 слоя, 2 головы, размерность 256.

У MTA ошибка почти везде ноль или рядом, у обычного трансформера почти везде двузначные числа процентов. Размеры свёрток были c_q=2 (как L), c_k=2N-1, чтобы можно было покрыть весь блок. Свёртка для голов не использовалась.

Следующий эксперимент с LLM. Предобучили 880M модели с архитектурой LLaMa и сравнили обычный трансформер, Differential Transformer и MTA. Обучали на SlimPajama на 105B токенов. В MTA key-query convolution использовали в каждом четвёртом слое, а head convolution в каждом. Свёртки c_q=6, c_k=11, размер группы 2.

По перплексии MTA лучше (GroupNorm при этом важен). На наборе бенчмарков тоже обычно бьёт остальных, но не везде и разница часто в последней цифре (и непонятно какой доверительный интервал -- обучали дважды). В среднем лучше.

Проверили на отдельном пуле long-range dependency задач: LAMBADA, NeedleIn-A-Haystack и BabiLong. На ламбаде однозначно бьёт, на multi-needle (2,4,6) retrieval точность MTA обычно выше, причём без GroupNorm часто лучше. На BabiLong и QA1-5 у MTA тоже всё хорошо.

Приложили сколько-то визуализаций свёрточных ядер, заметное число близко к identity, но есть и более хитрые. Например, один с диагональной структурой, удобен чтобы находить точное совпадение с паттерном. Или есть аналог edge detection, усиливающий первый или последний из последовательных ключей с высоким вниманием. В свёртках по головам частый паттерн это контраст, вычитание одной головы из другой.

Абляции показали, что даже пары MTA слоёв достаточно для превосходства над бейзлайнами. Все предложенные компоненты что-то улучшают по перплексии.

В целом забавно. Кажется, свёртки по q/k это ещё не предел. Для каких-то задач и языков не удивлюсь, если более забористые и менее локальные интеракции рулят. Главное чтоб параметров много не добавляли. Здесь в примере с LLM разница была на уровне 0.001% (+10K параметров на фоне 880M).

По памяти и FLOPS текущая неоптимизированная имплементация сильно проигрывает у использующих обычное scaled dot product attention: памяти раза в три больше надо, флопсов меньше раз в пять. Но это скорее проблема отсутствия оптимизированного ядра для CUDA. Интересно, компиляция через XLA что бы дала.

Brand new 2025 AI Index Report is released!

https://hai.stanford.edu/ai-index/2025-ai-index-report

Вчера в разборе Multi-Token Attention упоминалась статья про Differential Transformer. До её разбора я всё-таки не доберусь, так что продолжаю эксперименты с автоматизированным разбором.

Результат тут.

Пообщались с Иваром Максутовым и Постнаукой о хренах и пряниках разном

Что ждёт профессии, связанные с переводами, в ближайшие годы? Каким станет мышление и коммуникация в эпоху AI? Как LLM справляются с пониманием языка?

https://youtu.be/jWVbaCiN0Tc

Об этом — в подкасте с Григорием Сапуновым, соучредителем и техническим директором компании Intento, специалистом в области ИИ и анализа данных.