Механизм внимания (attention) давно стал сердцем и душой всех современных больших языковых моделей (LLM). С момента публикации знаменитой статьи «Attention is All You Need» в 2017 году подход к вычислению внимания практически не менялся: берём один токен-запрос, сопоставляем его с ключами и определяем «ценность» (value). Но недавно исследователи из FAIR (Meta AI) предложили новый взгляд на эту проблему, представив механизм Multi-Token Attention (MTA).
📌 Что нового предложили исследователи?
Стандартный подход к вниманию опирается на сравнение одного векторного представления токена-запроса (query) с одним вектором ключа (key). Это означает, что каждая пара запрос-ключ анализируется отдельно. А что делать, если для определения релевантного контекста необходимо одновременно учитывать несколько токенов?
Простой пример, предложенный авторами статьи, показывает уязвимость такого подхода:
Если вопрос звучит как: «Где Алиса увидела кролика?», то модель должна одновременно учитывать два объекта: «Алиса» и «кролик». Обычная система внимания может эффективно найти отдельно упоминания Алисы и отдельно — кролика, но ей сложно понять, когда оба объекта встречаются одновременно в одном контексте.
Именно тут и проявляется ограниченность традиционного внимания. Multi-Token Attention решает эту проблему, позволяя учитывать несколько токенов одновременно при расчёте весов внимания.
🧩 Как это реализовано технически?
Интересно рассмотреть детали реализации этого подхода:
🔹 Свёрточные операции (convolutions)
MTA использует двумерные свёртки (Conv2d), применяемые к матрице весов внимания. Представьте себе, что вместо того, чтобы рассматривать каждую пару токенов отдельно, вы объединяете группы соседних запросов и ключей. Получается, что информация от близлежащих токенов теперь влияет на то, как будет вычислен вес внимания.
🔹 Обмен информацией между головами (head mixing)
Стандартный Transformer разделяет внимание на отдельные головы (multi-head attention). В MTA предусмотрен обмен информацией и между этими головами при помощи дополнительных свёрток — это позволяет головам «советоваться» друг с другом, усиливая полезную информацию и убирая шум.
🔹 Групповая нормализация с масштабированием
Дополнительно вводится так называемая групповая нормализация, позволяющая улучшить распространение градиентов в модели, что делает обучение более стабильным и эффективным.
🎯 Что это даёт на практике?
Авторы провели обширные эксперименты, показавшие впечатляющие результаты:
- 🧪 Тест на игрушечной задаче:
При поиске блоков текста с определёнными буквами классический Transformer ошибался примерно в половине случаев, тогда как MTA практически не допускал ошибок. - 📚 Большие эксперименты с языковыми моделями (880 млн параметров):
Обучая модели на 105 млрд токенов, исследователи получили заметное улучшение перплексии (perplexity) и производительности на популярных бенчмарках (BoolQ, PIQA, WinoGrande и др.). - 🕵️ Тесты с длинным контекстом:
На таких задачах, как поиск информации в больших текстах (Needle-In-A-Haystack), MTA показывал существенное превосходство, особенно в случаях, когда требуется точное нахождение информации среди множества отвлекающих данных.
📈 В чём главные преимущества MTA?
✅ Лучшее понимание контекста:
MTA гораздо эффективнее распознаёт сложные ситуации, где требуется одновременное внимание к нескольким токенам.
✅ Гибкость и мощность:
Использование свёрток позволяет модели учитывать взаимодействия токенов и усиливать внимание там, где действительно нужно.
✅ Незначительное усложнение:
Несмотря на всю свою мощь, MTA практически не увеличивает число параметров модели — рост составляет лишь 0,001% относительно базовой архитектуры Transformer.
🚨 Есть ли недостатки у MTA?
Конечно, любые инновации требуют дальнейшего изучения и оптимизации. Главным недостатком авторы называют производительность: текущая реализация Multi-Token Attention не оптимизирована и требует большего количества вычислительных ресурсов и памяти. Однако это скорее техническая задача, которую можно решить со временем путём улучшения CUDA-ядра и программной оптимизации.
💬 Личное мнение автора
Как мне кажется, подход, предложенный исследователями FAIR, имеет потенциал стать новым стандартом в ближайшие годы. Возможность учитывать не один, а сразу несколько токенов при расчёте внимания выглядит крайне перспективно и логично. Уже сегодня видно, что MTA показывает существенный прирост в качестве при решении задач с длинным контекстом и сложными запросами.
Кроме того, идея с использованием свёрточных слоёв и взаимодействием голов внимания кажется вполне логичной и закономерной эволюцией механизма внимания, и это будет интересно не только с научной, но и с инженерной точки зрения.
✨ Что нас ждёт в будущем?
Вполне возможно, что следующие поколения популярных языковых моделей — от GPT до открытых аналогов — начнут внедрять подходы наподобие MTA. Это позволит получить модели, гораздо лучше понимающие сложные контексты, которые ближе к человеческому восприятию текста и речи.
Нам остаётся с нетерпением наблюдать за тем, как MTA и подобные механизмы начнут постепенно менять мир больших языковых моделей и искусственного интеллекта в целом.
📌 Полезные ссылки и источники: