Atención Multi-Token

Multi-Token Attention

April 1, 2025

Autores: Olga Golovneva, Tianlu Wang, Jason Weston, Sainbayar Sukhbaatar

cs.AI

Resumen

La atención suave es un mecanismo crucial que permite a los LLMs localizar partes relevantes dentro de un contexto dado. Sin embargo, los pesos de atención individuales se determinan únicamente por la similitud entre un solo vector de consulta y un vector de clave. Esta "atención de token único" limita la cantidad de información utilizada para distinguir una parte relevante del resto del contexto. Para abordar este problema, proponemos un nuevo método de atención, la Atención Multi-Token (MTA), que permite a los LLMs condicionar sus pesos de atención en múltiples vectores de consulta y clave simultáneamente. Esto se logra aplicando operaciones de convolución sobre las consultas, claves y cabezales, permitiendo que consultas y claves cercanas influyan en los pesos de atención de las demás para lograr una atención más precisa. Como resultado, nuestro método puede localizar contextos relevantes utilizando información más rica y matizada que puede superar la capacidad de un solo vector. A través de evaluaciones exhaustivas, demostramos que MTA logra un rendimiento mejorado en una variedad de benchmarks populares. En particular, supera a los modelos base de Transformer en tareas estándar de modelado de lenguaje y en tareas que requieren buscar información dentro de contextos largos, donde la capacidad de nuestro método para aprovechar información más rica resulta especialmente beneficiosa.

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.