Multi-Token-Aufmerksamkeit

Multi-Token Attention

April 1, 2025

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

cs.AI

Zusammenfassung

Soft Attention ist ein entscheidender Mechanismus, der es großen Sprachmodellen (LLMs) ermöglicht, relevante Teile innerhalb eines gegebenen Kontexts zu lokalisieren. Allerdings werden die individuellen Aufmerksamkeitsgewichte lediglich durch die Ähnlichkeit eines einzelnen Query- und Key-Token-Vektors bestimmt. Diese „Single-Token-Attention“ begrenzt die Menge an Informationen, die zur Unterscheidung eines relevanten Teils vom restlichen Kontext verwendet wird. Um dieses Problem zu lösen, schlagen wir eine neue Aufmerksamkeitsmethode vor, die Multi-Token-Attention (MTA), die es LLMs ermöglicht, ihre Aufmerksamkeitsgewichte gleichzeitig auf mehrere Query- und Key-Vektoren zu konditionieren. Dies wird durch die Anwendung von Faltungsoperationen auf Queries, Keys und Heads erreicht, wodurch benachbarte Queries und Keys die Aufmerksamkeitsgewichte gegenseitig beeinflussen können, um eine präzisere Aufmerksamkeit zu erzielen. Dadurch kann unsere Methode relevante Kontexte mithilfe von reichhaltigeren und nuancenreicheren Informationen lokalisieren, die die Kapazität eines einzelnen Vektors übersteigen können. Durch umfangreiche Evaluierungen zeigen wir, dass MTA eine verbesserte Leistung auf einer Reihe von populären Benchmarks erzielt. Insbesondere übertrifft es Transformer-Baselinemodelle bei Standardaufgaben zur Sprachmodellierung sowie bei Aufgaben, die die Suche nach Informationen in langen Kontexten erfordern, wo sich die Fähigkeit unserer Methode, reichhaltigere Informationen zu nutzen, als besonders vorteilhaft erweist.

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.