Многоцелевое Внимание (Multi-Token Attention)

Представляю вам новую эру в архитектурах внимания, о которой идет речь в статье "Multi-Token Attention" авторов Ольги Головневой, Тянлу Вана, Джейсона Уэстона и Сайнбаяра Сухбаатара. Вы можете ознакомиться с полным текстом работы по ссылке. Это исследование вносит значительные изменения в понимание и использование механизма внимания, и я уверен, что его стоит подробно разобрать.

Переосмысленный Механизм Внимания

Как известное в сообществе, в классическом механизме внимания распределение весов происходит на основании одного вектора запросов (query) и одного вектора ключей (key). Такой подход, известный как "одиночное внимание" (single token attention), создает определённые ограничения в выделении значимых элементов на фоне остальных. Но вот появился новый метод — Многоцелевое Внимание (Multi-Token Attention, MTA) — который позволяет избавиться от этого узкого места, задавая веса внимания на основе нескольких векторов запросов и ключей одновременно.

Для лучшего понимания, давайте вспомним, как веса определяются в традиционном механизме. Важно отметить, что они вычисляются по формуле softmax(QK/sqrt(d)). Здесь каждый токен имеет свой эмбеддинг, который проецируется в три вектора: Q, K и V. Скалярное произведение векторов Q и K разных токенов показывает их "сходство" или "важность". После нормализации по корню от размерности эмбеддинга и применения софтмакса, мы получаем веса внимания A. Эти веса затем используются для взвешивания и суммирования векторов V, что позволяет генерировать новые эмбеддинги.

Разумеется, в данном механизме можно вносить множество модификаций. Мы уже рассматривали разные аспекты, такие как температура в софтмаксе, примеры разреженного внимания, а также другие специальные подходы, как например Reformer. Эта новая работа также входит в число таких модификаций.

Применение Многоцелевого Внимания

Теперь давайте представим, что мы хотим выделить определённые элементы в предложении, например, "Где Алиса увидела кролика?". При этом нам нужны векторы запросов q_a и q_r для Алисы и кролика соответственно. Традиционный механизм будет определять веса внимания, описанным ранее способом. Однако задача заключается в том, чтобы убедиться, что эти сущности находятся рядом Ближе всего в контексте. Стандартный механизм внимания, к сожалению, не способен это сделать на одном слое, так как отсутствует взаимодействие между различными картами внимания. Хотя использование нескольких слоев может помочь, хотелось бы избежать подобной нагрузки.

В MTA это взаимодействие становится возможным благодаря новым блокам, которые включают:

1. Свертка key-query: объединяет несколько ключей и запросов внутри одной головы.
2. Свертка смешивания голов: обменивает важную информацию между разными головами внимания.
3. Групповая нормализация с масштабированием по глубине: улучшает поток градиентов, что существенно повышает эффективность модели.

Первый блок: Свертка key-query

Свертка key-query является важным новшеством, которое помогает "перемешивать" веса внимания от различных временных шагов. Она применяется к логитам внимания до софтмакса (то есть к QK/sqrt(d)), используя двумерную обучаемую свертку по измерениям запросов и ключей, оставаясь при этом неизменной для батча и голов. Каждая голова внимания учит собственную свертку, и в ней используется маска, чтобы избежать перетасовки с будущими значениями. Это можно считать предварительной сверткой (pre-softmax convolution). Так же возможно применение свертки после софтмакса, что делает взаимодействия между весами аддитивными, а не мультипликативными. Само по себе это изменение вызывает у меня неудовлетворенность в объяснении: остается незавершённым вопрос, почему же до софтмакса оно понимается именно как мультипликативное.

Итоги

Multi-Token Attention является значительным шагом вперёд в области машинного обучения и обработки естественного языка.