Multi-Token Attention | Summary

This article explains the key points of Multi-Token Attention.

• 2025-04-01 (arXiv), COLM 2025
• Golovneva, Olga, Wang, Tianlu, Weston, Jason, Sukhbaatar, Sainbayar.
• FAIR at Meta
• Paper

Read this article in Korean

Summary

• Language models can struggle to locate a passage that must match several pieces of information simultaneously. With 880M-parameter models trained on the same 105B tokens from SlimPajama, a mixed-source text dataset, Multi-Token Attention (MTA) reduces average validation perplexity, a measure of next-token prediction uncertainty, from 11.25 for the standard Transformer to 10.91.
• MTA applies learned convolutions, local weighted combinations, across query positions, key positions, and attention heads: queries represent retrieval requests, keys represent candidate context tokens, and heads are parallel attention computations within a layer. It raises the standard benchmark average from 43.7 to 44.9 and shows larger advantages when retrieving multiple targets or handling distracting context, although it does not lead every task.
• The implementation has substantial runtime and memory overhead despite nearly unchanged parameter counts. During training on 32 NVIDIA H200 GPUs, the reported MTA implementation processes 5.7 × 10³ tokens per second and uses 73.8 GB, compared with 54.3 × 10³ tokens per second and 17.5 GB for the standard Transformer; the authors attribute much of the gap to MTA’s inability to use optimized attention kernels, the GPU routines that execute attention.

1 Introduction

Standard attention computes each unnormalized relevance score from one query vector and one key vector. The authors argue that identifying a sentence containing both “Alice” and “rabbit” requires a single vector to encode both conditions; retrieving each word independently does not directly identify their co-occurrence.

• MTA lets attention maps, the arrays of query-to-key relevance weights, interact across neighboring keys, previous queries, and heads, so separately detected evidence can contribute to a shared score.
• The paper tests this motivation through a controlled letter-block task, matched language-model pretraining, long-range dependency evaluations, and model-size comparisons.
• Code: https://github.com/facebookresearch/RAM/tree/main/projects/mta

2 Background on multi-head attention

For hidden states H, the token representations within a layer, each attention head constructs keys K = HW_k, values V = HW_v, and queries Q = HW_q; values carry the information aggregated into the output. With hidden dimension D and M heads, each head has dimension d = D/M, and its attention logits, the scores before probability normalization, are  = QKᵀ/√d.

• Softmax converts logits into weights summing to one over keys, while a causal mask assigns −∞ to future positions so predictions cannot use subsequent tokens.
• Each head produces AV, and the outputs are concatenated and multiplied by an output projection.

3 Multi-Token Attention

MTA adds three components to multi-head attention: key-query convolution, head mixing convolution, and group normalization with scalar gating. These combine nearby token matches, exchange attention information between heads, and rescale head outputs while learning how strongly each contributes.

• The “Alice” and “rabbit” example illustrates combining attention peaks at nearby context locations instead of requiring one query to encode all retrieval conditions.
• Figures 1 and 2 show where the operations enter the attention pipeline and which dimensions they mix.

MTA mixes attention scores before softmax and attention weights afterward, before value aggregation. Gated normalization then rescales the head outputs before concatenation.

3.1 Key-query convolution

Key-query convolution applies a separate learned two-dimensional kernel, an array of filter weights, to each head’s attention map across query and key positions. Before softmax, it combines logits and therefore produces multiplicative interactions after exponentiation; after softmax, it combines attention weights additively.

• The practical pre-softmax implementation is \(A=\operatorname{Softmax}\!\left(\operatorname{Mask}_{-\infty}\!\left(\operatorname{Conv2d}_{\theta}\!\left(\operatorname{Mask}_{0}(\hat A)\right)\right)\right)\). Zero masking before convolution prevents future scores from contributing, and the second mask excludes future keys before normalization.
• Only current and past query positions are used. The simpler masking implementation removes slightly more information than necessary but avoids modifying convolution kernels.
• Kernel sizes determine the range of token interactions: the paper’s word-tokenized example uses c_q = 4 to cover the relevant query words and c_k = 5 to combine the relevant context words.

3.2 Head mixing convolution

Head mixing convolution divides M heads into M/c_h non-overlapping groups, each containing c_h heads, and learns linear combinations within each group. For two heads, the new maps are A¹_new = w₁₁A¹ + w₁₂A² and A²_new = w₂₁A¹ + w₂₂A².

• The operation can mix either logits before softmax or attention weights after softmax.
• Because the kernel spans the entire head group, it can also be viewed as a fully connected transformation within that group.

3.3 Putting everything together

When key-query mixing and head mixing occur on the same side of softmax, they can be implemented together as a three-dimensional convolution across queries, keys, and head groups. The paper also considers applying key-query convolution before softmax and head mixing afterward.

• Group normalization rescales attention outputs, while a learned sigmoid scalar gate, a bounded multiplier, controls their contribution before concatenation.
• The authors replace the layer-dependent scaling used in Differential Transformer, an attention variant that subtracts two softmax attention maps, with scalar gating and compare normalization and scaling choices experimentally.

4 Experiments

The experiments cover synthetic retrieval, language-model pretraining, context-length extension, long-range dependency tasks, learned-kernel inspection, and ablations, controlled comparisons that change individual design choices. Unless otherwise specified, they apply MTA convolution both before and after softmax.

4.1 Motivating toy task

The toy task asks a model to find the unique block of N random letters containing L = 2 question letters and output either all letters, the first letter, or the last letter. A decoder-only Transformer, which predicts tokens from preceding tokens, has 4 layers, 2 heads, and 256 hidden dimensions; the MTA variant uses only pre-softmax key-query convolution, with c_q = 2 and c_k = 2N − 1.

• For block sizes N = 5 and N = 8, MTA achieves 0.0–0.1% mean error across all output variants over three seeds.
• The standard Transformer’s mean errors range from 1.3% to 78.2%, with large seed variation in several conditions. This establishes an advantage in the tested setting, rather than a universal inability of standard attention to solve conjunction tasks.

Across both block sizes and all three output variants, MTA reduces mean error to at most 0.1%. Several standard Transformer conditions have large standard deviations, showing substantial variation across the three training seeds.

4.2 Large language modeling

The main comparison trains 880M-size models on 105B SlimPajama tokens using Lingua, the training framework. Baselines include the standard Transformer, Differential Transformer, and Talking heads attention, which learns projections across heads before and after softmax.

• MTA uses c_q = 6, c_k = 11, and c_h = 16. Key-query convolution appears every fourth layer, while head mixing appears in every layer.
• Across two training runs, average validation perplexity is 11.25 for Transformer, 11.15 for Differential Transformer, 11.04 for Talking heads, and 10.91 for MTA; MTA has the lowest perplexity in every reported validation domain.
• The average standard benchmark score is 44.9 for MTA versus 43.7, 43.9, and 44.4 for the three baselines. MTA does not lead on SIQA, a social commonsense benchmark; ARC challenge, a science question-answering benchmark; or OpenBookQA, a question-answering benchmark using elementary facts. Appendix D specifies that MMLU, a multitask language-understanding benchmark, uses five examples, qualifying the main text’s general zero-shot description, meaning evaluation without task examples.

4.3 Long context finetuning

Long context finetuning continues training on the same data mixture for another 10.5B tokens while extending context from 2048 to 4096 tokens. The authors increase the theta parameter of Rotary Position Embeddings (RoPE), the model’s positional representation, to 500000, set weight decay to 0, and reduce warmup to 50 steps.

• MTA reaches average validation perplexity 10.65, compared with 11.02 for Transformer, 10.89 for Differential Transformer, and 10.88 for Talking heads.
• MTA’s language-modeling advantage persists after this context extension.

MTA has the lowest validation perplexity in every listed domain after both pretraining and context extension. Its average advantage persists after the additional 10.5B-token training stage.

MTA achieves the highest average score, 44.9, but gains vary by task. Standard attention leads on ARC challenge and OpenBookQA, while Talking heads leads on SIQA.

4.4 Long-range dependency tasks

The long-range evaluations use LAMBADA, a next-word prediction dataset requiring broader discourse context; Needle-In-A-Haystack, retrieval of inserted target facts called needles; and BabiLong, reasoning over facts embedded in distracting documents. MTA lowers LAMBADA standard/OpenAI perplexity to 13.2/8.4, compared with 17.6/9.5 for the standard Transformer.

• Needle retrieval uses 2K pretrained and 4K finetuned contexts, shuffled target ordering, 500 samples per setup, and accuracy averaged over insertion depths. Figure 3 reports 2, 4, 6, 8, and 10 needles, although the prose initially lists only 2, 4, and 6.
• Figure 3 shows clear MTA advantages when several needles must be retrieved; at two needles, other baselines perform better.
• On BabiLong QA1-5, five tasks involving supporting facts or argument relations, MTA’s aggregate advantage grows as distraction increases to 4K tokens. Appendix E provides task-specific results that reveal remaining reasoning weaknesses.

MTA achieves the lowest perplexity on both LAMBADA variants, outperforming the standard Transformer and both attention baselines. This extends the evidence to next-word prediction requiring discourse context.

MTA’s advantage is pronounced when retrieving four or more targets in both context settings. Other methods perform better with two needles, qualifying the text’s claim of improvement across all counts.

4.5 Kernel patterns

Learned key-query kernels include near-identity patterns, which largely preserve the original score, and structured patterns that combine neighboring matches. Figure 4 shows a diagonal kernel and its corresponding attention map concentrating on “The magic number of San Francisco is 8”, supporting the interpretation that this kernel reinforces matching token sequences.

• The authors interpret other examples as priming previously attended keys or detecting boundaries of contiguous high-attention regions; these are qualitative interpretations.
• Small head-mixing kernels often contrast attention maps through subtraction. Without group normalization, kernel scales increase with depth; the suggestion that they compensate for the residual stream, the representations carried through skip connections, remains a hypothesis.

The diagonal kernel combines aligned neighboring query-key matches, and the associated attention map emphasizes the target city-number sentence. This illustrates the authors’ interpretation of sequence matching through score convolution.

4.6 Ablations

Ablations show that adding key-query convolution to only two layers improves validation perplexity over the plotted Transformer and Differential Transformer baselines, with six layers selected as a performance-complexity compromise. Larger head groups improve perplexity, and the authors report better convergence and final performance with identity initialization than with the tested zero and constant 0.3 initializations.

• Table 5 reports perplexity 10.90 when both convolution types operate before and after softmax with scalar gating, compared with 10.92 using depth scaling.
• Normalization and scaling affect results: one tested layer-norm scaling configuration reaches 11.41, while changing convolution placement produces smaller differences in the compared configurations.
• Figure 6 shows a relative perplexity advantage for MTA across the plotted model sizes. The text names 300M, 550M, and 1B additional configurations, while the largest plotted point is near 1.4B and Appendix I describes a 1.4B model; the paper does not reconcile these labels.

The left panel shows MTA retaining more aggregate accuracy as distracting text grows. The right panel shows diminishing perplexity improvements as more layers receive key-query convolution, supporting the choice to use it in only some layers.

Increasing head-group size lowers perplexity with diminishing returns. MTA also achieves a larger relative perplexity reduction than the compared attention variants at each plotted parameter scale.

The best listed configuration combines both convolution types before and after softmax with scalar gating. Results vary with normalization and scaling, so the gains depend on the surrounding attention-output design as well as convolution.

5 Related work

The paper relates MTA to sharper attention normalization, removal of irrelevant memories, contextual position encodings, head mixing, and noise-canceling attention. Talking heads attention is closely related to its head-mixing component, while Differential Transformer uses subtraction to suppress irrelevant attention.

• Earlier convolution-attention combinations compress keys and values, interleave convolution with Transformer layers, or convolve attention along the key dimension.
• The authors distinguish MTA through convolution of attention scores across multiple query positions, key positions, and heads together.

6 Conclusion

The authors conclude that combining attention evidence across multiple queries, keys, and heads improves the localization of relevant context. Their experiments show near-zero error on the synthetic conjunction task, lower language-model perplexity, and stronger retrieval with multiple targets and distracting text.

Appendix

• A Limitations: The authors state that MTA is not currently compatible with popular optimized attention kernels. They leave runtime optimization to future work; the reported quality gains therefore come from an implementation with substantial training overhead.
• B MTA Architecture: Mixing two heads after softmax can be represented using value/output projections with twice the rank, the dimensional capacity of the corresponding matrix transformation. The appendix gives an analogous derivation for query/key projections before softmax, while cautioning that shared parameters prevent equivalence to independently parameterized doubled-rank heads.
• C Toy task details: Sequences contain up to 50 period-separated blocks with N ∈ {5, 8} lowercase letters, followed by “#” and L = 2 question letters, with exactly one target block. Training uses 1M sequences, a held-out set of 1K sequences, batch size 64, 100K steps, and three random seeds.
• D Language model training details: Training follows the LLaMa architecture with root mean square normalization (RMSNorm) before sublayers, SwiGLU gated activations, and Rotary Embeddings. The 880M configuration has dimension 1536, 24 layers, 16 heads, 262144 tokens per batch, context length 2048, learning rate 1.5e−4, weight decay 0.05, 375 warmup steps, and RoPE theta 100000; Differential Transformer uses half as many heads to accommodate doubled head dimension, and MMLU evaluation is 5-shot.
• E BabiLong details: Table 6 illustrates QA1-5 tasks involving one, two, or three supporting facts and two- or three-argument relations. Figure 7 shows that aggregate improvements leave reasoning weaknesses: QA3 accuracy remains near its 16.67% random baseline, and Differential Transformer outperforms MTA on QA4.
• F Complexity evaluation: Table 8 lists 876553728 parameters for Transformer and 876583320 for MTA, an increase of 29592 parameters, approximately 0.0034%; this differs from the introduction’s stated 0.001%. Table 9 reports 73.8 GB and 5.7 × 10³ tokens per second for MTA versus 17.5 GB and 54.3 × 10³ for Transformer during training on 32 NVIDIA H200 GPUs, with the baseline benefiting from optimized attention kernels unavailable to this MTA implementation.
• G Multi-needle evaluation results: Figure 8 breaks down 2K-context retrieval by insertion depths of 0%, 25%, 50%, 75%, and 100% for 2, 4, 6, and 8 needles. MTA maintains high accuracy across depths in the larger multi-needle conditions, while other methods outperform it in some two- and four-needle conditions.
• H MTA kernel patterns.: Figures 9–14 display key-query kernels across heads at layers 3, 7, 11, 15, 19, and 23; Figure 15 displays head kernels with c_h = 2 in the first three layers. Figures 16 and 17 compare diagonal kernel magnitudes and non-diagonal/diagonal ratios with and without group normalization, illustrating depth-dependent scaling and contrasting head interactions.
• I Finetuning with MTA: Identity-initialized convolutions can be inserted into pretrained Transformer layers and updated through continued training. With context length 2048, MTA adaptation yields average perplexity 10.61 versus 10.69 for continued standard-Transformer training after 10.5B tokens on the authors’ 1.4B model; after 5.3B tokens on Llama 3.2 1B, Llama 3.2 3B, and Llama 3.1 8B, the corresponding comparisons are 10.46 versus 10.53, 8.89 versus 8.93, and 8.48 versus 8.53.

Brief Thoughts

The contribution is the explicit interaction between attention maps, supported by a targeted synthetic task and improvements over standard attention and head-mixing baselines. The small parameter increase does not imply low computational cost: Table 9 shows substantial memory and throughput penalties, with different implementations’ access to optimized kernels affecting the comparison.

Memory and throughput measurements reveal an implementation cost that parameter counts conceal. The standard Transformer uses optimized attention kernels while MTA does not, so these results compare the available implementations rather than equally optimized mechanisms.

The evidence supports improved retrieval and language modeling under the tested conditions, with narrower support for general reasoning or much longer contexts. Main retrieval evaluations use 2K and 4K contexts, BabiLong retains task-specific failures, and Table 10 establishes modest perplexity gains from adaptation without testing corresponding retrieval gains.

Adding MTA during continued training lowers average perplexity for each tested pretrained model. The authors’ model pretrained with MTA remains better than the adapted standard Transformer; retrieval after adaptation is not evaluated.