Многотокеновое внимание

Multi-Token Attention

April 1, 2025

Авторы: Olga Golovneva, Tianlu Wang, Jason Weston, Sainbayar Sukhbaatar

cs.AI

Аннотация

Мягкое внимание (soft attention) является ключевым механизмом, позволяющим большим языковым моделям (LLMs) находить релевантные части в заданном контексте. Однако отдельные веса внимания определяются сходством только одного вектора запроса и ключевого токена. Такое "внимание на один токен" ограничивает объем информации, используемой для выделения релевантной части из остального контекста. Чтобы решить эту проблему, мы предлагаем новый метод внимания — Мультитокенное внимание (Multi-Token Attention, MTA), который позволяет LLMs учитывать веса внимания на основе нескольких векторов запросов и ключей одновременно. Это достигается за счет применения операций свертки к запросам, ключам и головам внимания, что позволяет соседним запросам и ключам влиять на веса внимания друг друга для более точного фокусирования. В результате наш метод может находить релевантный контекст, используя более богатую и детализированную информацию, которая выходит за пределы возможностей одного вектора. В ходе обширных экспериментов мы демонстрируем, что MTA достигает улучшенной производительности на ряде популярных бенчмарков. В частности, он превосходит базовые модели Transformer на стандартных задачах языкового моделирования, а также на задачах, требующих поиска информации в длинных контекстах, где способность нашего метода использовать более богатую информацию оказывается особенно полезной.

English

Soft attention is a critical mechanism powering LLMs to locate relevant parts within a given context. However, individual attention weights are determined by the similarity of only a single query and key token vector. This "single token attention" bottlenecks the amount of information used in distinguishing a relevant part from the rest of the context. To address this issue, we propose a new attention method, Multi-Token Attention (MTA), which allows LLMs to condition their attention weights on multiple query and key vectors simultaneously. This is achieved by applying convolution operations over queries, keys and heads, allowing nearby queries and keys to affect each other's attention weights for more precise attention. As a result, our method can locate relevant context using richer, more nuanced information that can exceed a single vector's capacity. Through extensive evaluations, we demonstrate that MTA achieves enhanced performance on a range of popular benchmarks. Notably, it outperforms Transformer baseline models on standard language modeling tasks, and on tasks that require searching for information within long contexts, where our method's ability to leverage richer information proves particularly beneficial.