Multi-Token Aandacht

Multi-Token Attention

April 1, 2025

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

cs.AI

Samenvatting

Zachte aandacht is een cruciaal mechanisme dat LLM's in staat stelt om relevante delen binnen een gegeven context te lokaliseren. Echter, individuele aandachtswaarden worden bepaald door de gelijkenis van slechts één query- en sleuteltokenvector. Deze "enkele token aandacht" beperkt de hoeveelheid informatie die wordt gebruikt om een relevant deel van de rest van de context te onderscheiden. Om dit probleem aan te pakken, stellen we een nieuwe aandachtmethode voor, Multi-Token Attention (MTA), waarmee LLM's hun aandachtswaarden kunnen conditioneren op meerdere query- en sleutelvectoren tegelijkertijd. Dit wordt bereikt door convolutiebewerkingen toe te passen op queries, sleutels en heads, waardoor nabijgelegen queries en sleutels elkaars aandachtswaarden kunnen beïnvloeden voor preciezere aandacht. Als gevolg hiervan kan onze methode relevante context lokaliseren met behulp van rijkere, genuanceerdere informatie die de capaciteit van een enkele vector kan overstijgen. Door uitgebreide evaluaties tonen we aan dat MTA een verbeterde prestaties behaalt op een reeks populaire benchmarks. Met name overtreft het Transformer-basismodellen op standaard taalmodelleertaken, en op taken die vereisen dat er informatie wordt gezocht binnen lange contexten, waar het vermogen van onze methode om rijkere informatie te benutten bijzonder nuttig blijkt.

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.