Article Overview: Multi-Token Attention

Link to article: https://arxiv.org/abs/2504.00927

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

Abstract

Multi-Token Attention (MTA) extends the standard scaled dot-product attention by conditioning each attention weight on multiple queries, keys, and heads via convolutional operations. This alleviates the “single-token” bottleneck, enabling richer context retrieval and improved performance on both toy tasks and large-scale language modeling benchmarks.

1. Motivation and Background

• Standard multi-head attention computes each weight from a single query–key vector pair, limiting its ability to identify context defined by multiple elements.
• Examples such as finding a sentence containing both “Alice” and “rabbit” illustrate the need to combine signals across tokens and heads (see Figure 1, left vs. right).

2. Multi-Token Attention Architecture

• Key-Query Convolution: Applies a 2D convolution over the attention logits (pre- or post-softmax) to mix neighboring queries and keys (kernel sizes c_q, c_k).
• Head Mixing Convolution: Divides heads into groups and applies convolution across heads (kernel size c_h) to share information between attention maps.
• Group Normalization with Depth Scaling: Normalizes each head’s output independently, improving gradient flow as depth increases.

3. Toy-Task Demonstration

• Task: identify the block of N random letters containing L query letters and output all/first/last tokens.
• Standard Transformer fails (error rates up to 78%; Table 1), whereas MTA achieves near-zero error by setting c_q = L and c_k = 2N - 1 to cover entire blocks.

4. Large-Scale Language Modeling

• Models (880 M parameters) are pretrained on 105 B tokens of SlimPajama and finetuned to 4 k context length.
• MTA yields lower validation perplexity than both standard and Differential Transformer baselines (Table 2), with only a 0.001% parameter increase.
• Zero-shot evaluation on BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC, OpenBookQA, and MMLU shows average improvements (Table 3).

5. Long-Range Dependency Benchmarks

• On LAMBADA, MTA reduces perplexity significantly (Table 4).
• In the Needle-In-A-Haystack task, MTA achieves higher retrieval accuracy across multiple inserted “needles” and depths (Table 5).
• BabiLong QA1–5 tasks demonstrate consistent gains, especially with extensive distraction context (Figure 4, left).

6. Kernel Analysis and Ablations

• Kernel Patterns: Many key-query kernels learn identity or diagonal structures for pattern matching; head kernels often perform contrastive amplification (Figure 3).
• Ablations: Performance improves with as few as two MTA layers (Figure 4, right). GroupNorm and depth scaling each contribute to lower perplexity; variations in kernel size and pre/post-softmax positioning have minor effects (Table 6).

7. Conclusion and Future Directions

• MTA overcomes single-token attention limitations by leveraging local convolutions across queries, keys, and heads.
• It provides enhanced context retrieval in toy and real-world tasks, though current implementations incur higher memory and FLOPS (Table 10).
• Future work may optimize CUDA kernels for efficiency and explore larger or more global interaction kernels.