Atenção Multi-Token

Multi-Token Attention

April 1, 2025

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

cs.AI

Resumo

A atenção suave é um mecanismo crucial que capacita os LLMs a localizar partes relevantes dentro de um determinado contexto. No entanto, os pesos individuais de atenção são determinados pela similaridade de apenas um único vetor de token de consulta e chave. Essa "atenção de token único" limita a quantidade de informação utilizada para distinguir uma parte relevante do restante do contexto. Para resolver esse problema, propomos um novo método de atenção, a Atenção Multi-Token (MTA), que permite que os LLMs condicionem seus pesos de atenção em múltiplos vetores de consulta e chave simultaneamente. Isso é alcançado aplicando operações de convolução sobre consultas, chaves e cabeças, permitindo que consultas e chaves próximas afetem os pesos de atenção umas das outras para uma atenção mais precisa. Como resultado, nosso método pode localizar contextos relevantes utilizando informações mais ricas e matizadas que podem exceder a capacidade de um único vetor. Através de avaliações extensas, demonstramos que a MTA alcança um desempenho aprimorado em uma variedade de benchmarks populares. Notavelmente, ela supera os modelos de linha de base Transformer em tarefas padrão de modelagem de linguagem e em tarefas que exigem a busca de informações em contextos longos, onde a capacidade do nosso método de aproveitar informações mais ricas se mostra particularmente benéfica.

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.