Attenzione Multi-Token

Multi-Token Attention

April 1, 2025

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

cs.AI

Abstract

L'attenzione soft è un meccanismo cruciale che consente ai LLM di individuare le parti rilevanti all'interno di un determinato contesto. Tuttavia, i pesi di attenzione individuali sono determinati dalla similarità di un solo vettore di token query e key. Questa "attenzione a singolo token" limita la quantità di informazioni utilizzate per distinguere una parte rilevante dal resto del contesto. Per affrontare questo problema, proponiamo un nuovo metodo di attenzione, Multi-Token Attention (MTA), che consente ai LLM di condizionare i propri pesi di attenzione su più vettori query e key simultaneamente. Questo è ottenuto applicando operazioni di convoluzione su query, key e testine, permettendo alle query e key vicine di influenzare reciprocamente i pesi di attenzione per un'attenzione più precisa. Di conseguenza, il nostro metodo può individuare il contesto rilevante utilizzando informazioni più ricche e sfumate che possono superare la capacità di un singolo vettore. Attraverso valutazioni estensive, dimostriamo che MTA raggiunge prestazioni migliorate su una gamma di benchmark popolari. In particolare, supera i modelli baseline Transformer nei compiti standard di modellazione del linguaggio e nei compiti che richiedono la ricerca di informazioni all'interno di contesti lunghi, dove la capacità del nostro metodo di sfruttare informazioni più ricche si rivela particolarmente vantaggiosa.

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.