Language models use fewer context tokens when focusing attention
Retrieval Capacity of Self-Attention Under Competition
Computation and LanguageMachine Learning
Summary
Language models look at many words when trying to understand text, but they don’t always use all of them equally. The authors find that by only keeping the tokens (words or pieces of words) that get the most attention, the model performs almost as well as when using all tokens. The number of important tokens changes depending on the situation, and the way the model combines them matters too. They also created a theory to explain why models need more tokens in some cases even when no new information appears. This helps us understand how language models decide what context to use when predicting text.
What this means in practice
- •For natural language engineers: Optimize transformer models by focusing computation on fewer high-attention tokens, improving efficiency without retraining.
- •For ai system designers: Design language model architectures with better context handling by understanding how token competition affects attention allocation.
Authors
Timur Mudarisov, Mikhail Burtsev, Radu State
Abstract
How many tokens from its context does a language model actually use, and what determines that number? We study this question through self-attention. Without retraining, we retain only the tokens with the highest attention weights at each head, layer, and query, keeping their original weights unchanged. By varying the selected set size and measuring the increase in negative log-likelihood (NLL), we estimate the effective attention set size needed to stay within a chosen loss tolerance. Relatively small selected sets can keep NLL close to the full-attention baseline, although the required size varies across models. Attention-based selection substantially outperforms random selection. Selected sets exhibit geometric structure, although geometric separation alone does not establish that model loss is preserved. Extending context while evaluating the same prediction targets increases the required set size, while its fraction of context decreases over the tested range. Experiments with a fixed supporting fact show that additional background pushes its tokens down the attention ranking and reduces their attention mass. Renormalizing the retained weights can substantially reduce the required set size, showing that it also depends on how selected representations are combined. Conditional theoretical models explain how competition and attention-mass retention can produce growing set sizes without more distinct information to retrieve. These results provide a way to measure effective attention set size in language models and investigate its dependence on context, competition, and aggregation.