Multi-Token Attention — Paper Graph

🎯 Pitch

Transformer attention is fundamentally bottlenecked because each attention weight can only compare a single query to a single key — but MTA fixes this with convolutions that let queries and keys “talk” to each other before deciding where to attend, directly enabling multi-constraint lookups like finding a sentence that mentions both “Alice” and “rabbit” without cramming everything into one vector.

1. Executive Summary

This paper introduces Multi-Token Attention (MTA), a modification to the standard attention mechanism that enables attention weights to condition on multiple query and key vectors simultaneously—overcoming the "single token" bottleneck where each attention weight depends on only one query-key pair. Through convolution operations applied over queries, keys, and attention heads (key-query convolution and head mixing convolution), MTA allows nearby queries and keys to influence each other's attention weights, making it possible to locate relevant context using richer information than a single vector can encode—for example, locating a sentence containing both "Alice" and "rabbit" by first finding each independently, then combining those attention maps. Evaluated on 880M-parameter models pretrained on 105B tokens from SlimPajama, MTA improves validation perplexity by 0.34 points over the standard Transformer (10.91 vs. 11.25) while adding only ~0.001% more parameters, and achieves particularly large gains on long-context retrieval tasks—notably, multi-needle retrieval accuracy with 6 needles in 4K context reaches roughly 70% for MTA versus ~35% for the standard Transformer. The method's benefits are especially pronounced on tasks requiring precise location of information within long contexts, establishing that attention conditioned on multi-token interactions provides substantial advantages when the distinguishing information for relevant context exceeds the capacity of single token vectors.

2. Context and Motivation

The Core Problem: Attention Cannot Express Multi-Condition Queries

The fundamental limitation this paper identifies is that standard multi-head attention conditions each attention weight on exactly one query vector and one key vector. When deciding how much attention token $i$ should pay to token $j$ , the mechanism computes a single dot product $\mathbf{q}_i \cdot \mathbf{k}_j / \sqrt{d}$ — that's it. The entire decision of whether token $j$ is relevant to token $i$ must be compressed into the similarity of those two vectors.

This is a bottleneck — not in the information-theoretic sense that single vectors cannot encode complex queries (they can, given sufficient dimensionality), but in the practical sense that forcing all distinguishing information into a single vector pair places a heavy burden on the model's representational capacity. The paper's core example crystallizes this: suppose the model needs to find a sentence that mentions both "Alice" and "rabbit." The query token (the question mark, or "both," or some aggregation) must encode both concepts into one query vector, and the key vectors at each context position must similarly respond to both concepts simultaneously. This is possible in principle — a sufficiently high-dimensional query can encode arbitrarily many constraints — but it requires the model to dedicate substantial capacity to what is fundamentally a combinatorial search over concept conjunctions.

The paper argues this isn't merely a theoretical concern. It manifests empirically as suboptimal performance in long-context retrieval, where the model must locate information satisfying multiple criteria buried among distractors. The authors explicitly cite prior work (Kamradt, 2023; Kuratov et al., 2025) showing that "standard attention can suffer from suboptimal performance in this setting," and later work (Liu et al., 2024; 2025) demonstrating that "Transformers struggle to find relevant information especially in the middle of long context."

Why This Matters: Beyond One-Hop Retrieval

The practical significance extends beyond the synthetic "Alice and rabbit" example. Real-world language understanding routinely requires simultaneous attention to multiple constraints:

Multi-fact reasoning: "Where was the football before Mary took it to the garden?" requires locating facts about the football's location, Mary's actions, and their temporal ordering — not just any one of these in isolation.
Disambiguation through context: The word "bank" might be attended to highly based on similarity to a query about finance, but only a specific mention of "bank" near "river" and "loan" in the same sentence identifies the relevant passage. A single query-key comparison captures "bank ≈ finance," but not "bank near river AND loan."
Needle-in-haystack retrieval: When multiple needles are embedded in distracting text, the model must locate all of them and potentially identify which needle corresponds to which query — e.g., "The magic number of San Francisco is 8" and "The magic number of Boston is 12" require distinguishing which city matches which number.

In all these cases, the distinguishing information cannot be reduced to a single token match. Standard attention can handle some of these cases by using different attention heads for different constraints and then combining their outputs in the feedforward network. But as the paper points out, this is an indirect workaround: "Looking up 'Alice' with one attention head and using another head for 'rabbit' could find their mentions separately, but is not sufficient to pinpoint where both are mentioned together." The attention mechanism itself has no way to multiply (intersect) two attention maps — heads are combined only after the attention weights are already computed, at the output value stage.

This is both a theoretical limitation (the attention operation's expressiveness is constrained) and a practical one (it wastes model capacity and hurts performance on tasks where multi-constraint retrieval is essential).

Prior Approaches and Where They Fall Short

The paper situates itself against several lines of prior work, each of which addresses related problems but misses the core insight:

1. Making softmax sharper or sparser. Methods like sparsemax (Martins and Astudillo, 2016), adaptive temperature scaling (Veličković et al., 2024), and scalable-softmax (Nakanishi, 2025) all modify the softmax operation to produce more focused attention distributions. While these sharpen attention, they do not change the fundamental computational unit: each attention weight still depends on a single query-key pair. Sharpening helps the model attend more decisively to whatever it already finds, but doesn't help it find things that require multi-token evidence.

2. Removing or filtering irrelevant tokens. Sukhbaatar et al. (2021) propose mechanisms to drop irrelevant tokens from memory entirely. This reduces distraction but doesn't enhance the model's ability to identify relevance based on multi-condition criteria — it's a filtering mechanism, not a detection mechanism.

3. Context-aware position encodings. Golovneva et al. (2024) and Desrochers et al. (2024) incorporate contextual information into position encodings, allowing positions to be amplified based on content. This is a step toward making attention sensitive to broader context, but operates through position biases rather than direct multi-query-key interaction.

4. Noise-canceling attention mechanisms. Several recent methods modify attention to better distinguish signal from noise:

Talking-heads attention (Shazeer et al., 2020) adds linear projections across the head dimension before and after softmax. This allows heads to share information — head $h_1$ 's attention weights can influence head $h_2$ 's — but the mixing is linear and global across all tokens. It doesn't perform the localized key-query interaction that MTA does. The paper explicitly uses Talking-heads as a baseline, and MTA outperforms it (Table 3: 44.9 vs. 44.4 average benchmark score).
Differential Transformer (DIFF Transformer) (Ye et al., 2024) computes attention as the difference between two separate softmax attention maps, inspired by differential amplifiers in electronics. This effectively cancels common-mode noise, helping the model focus on distinctive rather than ubiquitous features. The paper notes this "is related to our head mixing step" because both involve combining attention maps across heads, but the mechanism is fundamentally different: subtraction of two maps vs. learned convolution over groups of heads.

5. Convolution in attention for language. Several prior works have incorporated convolution into attention mechanisms:

Liu et al. (2018) use convolution to compress keys and values by a factor of 3, extending context length — this is about efficiency, not expressiveness.
Gulati et al. (2020) propose Conformer, where a convolution module is applied to the output of multi-head self-attention (post-value-aggregation), used primarily in speech recognition.
Zheng et al. (2024) apply convolution on attention weights in the key dimension only to enhance length extrapolation — this is a single-axis convolution for positional purposes.
Xu et al. (2024) modify attention so keys and values can shift by one time step — the paper notes this "can be viewed as a special case of MTA where the convolution performs shifting in the key dimension."

Crucially, none of these prior methods apply convolutions across queries, keys, and heads simultaneously to the attention weights themselves. The paper's contribution is not convolution-in-attention per se, but the specific three-dimensional convolution (queries × keys × heads) that enables attention weights to condition on multi-token interactions. This is what allows the mechanism described in the Alice-and-rabbit example: finding mentions of each concept independently via standard attention-like operations, then combining those attention maps via convolution to identify locations where both are active.

6. The "just use more heads" argument. A natural counterargument is that standard attention can already handle multi-condition queries by dedicating different heads to different constraints and combining their outputs. The paper implicitly addresses this: combining head outputs (value-aggregated representations) in the feedforward network is not the same as combining head attention weights themselves. Output combination says "head 1 found Alice-related content, head 2 found rabbit-related content, now let the FFN figure out if they're related." Attention-weight combination says "head 1 found Alice here, head 2 found rabbit nearby — therefore this specific position deserves higher combined attention." The latter is more direct and, the paper argues, more effective.

How This Paper Positions Itself

The paper frames its contribution not as competing with these prior approaches but as addressing a gap none of them fill: the inability of attention weights to condition on multiple query-key vector pairs through direct, learned interaction in a local neighborhood. The key phrase from the introduction captures this precisely:

"We argue that the dependency on single token vector similarity brings a fundamental limitation to the attention mechanism. In many cases, the relevant part of the context cannot be identified by a single token."

The conceptual framework is straightforward: attention should behave more like a pattern-matching operation over sequences rather than a similarity check between individual vectors. By applying learned convolution kernels over the attention logits (or weights) along query and key dimensions, MTA allows attention at position $(i, j)$ to be influenced by neighboring queries ( $i-1$ , $i-2$ , etc.) and neighboring keys ( $j-1$ , $j+1$ , etc.). By applying convolution across heads, it allows different heads' attention patterns to interact directly. Together, these operations create what is effectively a small spatial pattern detector operating over the attention map, capable of recognizing conjunctions, sequences, and other multi-token patterns that single-vector similarity cannot express.

The paper's empirical strategy reinforces this positioning. It doesn't just show MTA improves perplexity — it designs experiments that specifically isolate the limitation it claims to fix. The toy task (Section 4.1) is the clearest example: it constructs a problem where the correct answer can only be identified by attending to a block that contains all of a set of query letters. Standard attention fails dramatically (51.6% error for $N=5$ ), while MTA with appropriately sized convolution kernels achieves near-zero error (0.1%). This is a direct demonstration of the claimed limitation and the proposed solution, in minimal form.

The paper also positions MTA as practical and lightweight — a modification that can be added to existing Transformer architectures with minimal parameter overhead. The 880M-parameter model adds only ~30K parameters (0.001%) from the convolution kernels, and the key-query convolution is applied to only 1/4 of layers (every 4th layer). This addresses the implicit criticism that any solution to the single-token bottleneck would require major architectural changes or computational overhead. By contrast, MTA is presented as a drop-in replacement for standard attention that requires no changes to training infrastructure, tokenization, or model parallelism.

Finally, the paper positions MTA within the broader trajectory of attention mechanism improvements — from the original scaled dot-product attention (Vaswani et al., 2017), through differential attention and talking-heads, toward mechanisms that allow richer interaction patterns. The baseline comparisons (DIFF Transformer, Talking-heads) are chosen specifically because they represent the current frontier of attention modifications that MTA builds upon and extends. The group normalization and gating mechanism are adapted from DIFF Transformer's approach, while the head mixing convolution generalizes Talking-heads' linear head interaction to a learned convolutional one. MTA is thus presented as the natural next step: moving from intra-head multi-token attention (key-query convolution) through inter-head attention sharing (head convolution) in a unified framework.

3. Technical Approach

3.1 Reader Orientation (Approachable Technical Breakdown)

This paper proposes a modified attention mechanism — a drop-in replacement for the standard dot-product attention in Transformer models. The system solves the problem that standard attention can only compare ONE query vector to ONE key vector at a time, which is like trying to find "the sentence that mentions both Alice AND rabbit" using only a single word-level comparison — MTA lets attention weights be influenced by MULTIPLE neighboring queries, keys, and even other attention heads simultaneously, enabling richer pattern-matching within the attention operation itself rather than relying on later layers to recombine information.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three major components, all operating as modifications to standard multi-head attention:

Key-query convolution — a learned 2D convolution kernel applied over the attention logits (or weights) along the query and key dimensions, allowing attention at position $(i,j)$ to depend on attention from nearby queries like $(i-1, j)$ and nearby keys like $(i, j-1)$ .
Head mixing convolution — a learned convolution applied across groups of attention heads, enabling heads within a group to share and recombine their attention patterns (e.g., head 1's "Alice" attention and head 2's "rabbit" attention can be combined to produce stronger attention where both are present).
Group normalization with gating — a normalization mechanism applied to the output of the MTA convolution operations before concatenation, with a learned sigmoid gate per head that lets the model selectively amplify or suppress heads.

Information flows as follows: standard Q, K, V projections produce key, query, and value vectors → attention logits are computed via scaled dot-product → MTA convolution applies the key-query and head convolutions (either before or after softmax) → the modified attention weights produce the value-weighted output → group normalization with sigmoid gating normalizes the concatenated output → the result feeds into the standard feedforward network.

3.3 Roadmap for the Deep Dive

First, the standard attention formulation — I'll establish the precise notation and operations that MTA modifies, so we can see exactly where the changes are inserted.
Second, key-query convolution — the central mechanism that lets attention at one (query, key) position be affected by neighboring positions, including the pre-softmax and post-softmax variants, the convolution equation, and why this handles multi-condition queries.
Third, head mixing convolution — how attention patterns from different heads are combined within groups, the linear mixing equations, and why this generalizes prior methods like Talking-heads attention.
Fourth, the full MTA module — how key-query and head convolutions are composed together (pre-softmax, post-softmax, or mixed), the three-dimensional convolution view, and the specific configurations used in experiments.
Fifth, group normalization with gating — the output normalization step, why it matters for competing with the residual stream, and the choice of sigmoid gating over layer-dependent depth scaling.
Sixth, practical design choices — which layers get key-query convolution (every 4th), kernel sizes used ( $c_q=6, c_k=11, c_h=16$ ), initialization strategy (identity), and the rationale behind these choices.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architectural innovation paper whose core idea is that attention weights should be computed using information from multiple query-key pairs and multiple heads simultaneously, rather than from isolated single-vector comparisons. The mechanism achieves this through learned convolution operations applied directly to attention logits or weights across three dimensions: queries, keys, and heads.

Standard Multi-Head Attention (What MTA Modifies)

The paper builds directly on the standard scaled dot-product attention formulation from Vaswani et al. (2017) in decoder-only Transformers. A sequence of tokens $[x_1, ..., x_T]$ of length $T$ is transformed into hidden states $\mathbf{H} = [\mathbf{h}_1, \dots, \mathbf{h}_T]^\top \in \mathbb{R}^{T \times D}$ through an embedding layer and repeated Transformer layers.

Each Transformer layer contains a multi-head attention submodule with $M$ heads of dimension $d = D/M$ . Each head applies learned linear projections $\mathbf{W}_k, \mathbf{W}_v, \mathbf{W}_q \in \mathbb{R}^{D \times d}$ to produce key, value, and query vectors:

$\mathbf{K} = \mathbf{H}\mathbf{W}_k, \quad \mathbf{V} = \mathbf{H}\mathbf{W}_v, \quad \mathbf{Q} = \mathbf{H}\mathbf{W}_q$

where $\mathbf{K}, \mathbf{V}, \mathbf{Q} \in \mathbb{R}^{T \times d}$ are the key, value, and query matrices for a single head.

The attention logits $\hat{\mathbf{A}} \in \mathbb{R}^{T \times T}$ and attention weights $\mathbf{A}$ (probabilities) are computed as:

$\hat{\mathbf{A}} = \mathbf{Q}\mathbf{K}^\top / \sqrt{d}$

$\mathbf{A} = \text{Softmax}\left(\text{Mask}_{-\infty}(\hat{\mathbf{A}})\right)$

where $\hat{\mathbf{A}}$ contains raw similarity scores between each query and each key, the division by $\sqrt{d}$ prevents dot products from growing too large with dimension, the $\text{Mask}_{-\infty}$ function replaces values at positions $(i,j)$ with $-\infty$ for $i < j$ to prevent attending to future tokens in autoregressive generation, and the softmax (applied over the key dimension) converts these masked logits into a probability distribution where each query's attention weights over keys sum to 1.

What it computes: for each query position $i$ and key position $j$ , the attention weight $a_{ij}$ is $\text{softmax}_j(\mathbf{q}_i \cdot \mathbf{k}_j / \sqrt{d})$ — the normalized similarity between the $i$ -th query vector and the $j$ -th key vector. This produces a matrix $\mathbf{A}$ where each row $i$ is a probability distribution over all past key positions $j \leq i$ .

Why this form: the dot product measures vector similarity — keys similar to the query get higher weights, so their value vectors dominate the output. The scaling by $\sqrt{d}$ keeps the softmax from saturating (entering near-one-hot regimes) when dimensionality is large, which would kill gradients. The causal mask ensures the model cannot cheat by looking at future tokens during training.

The attention output for the head is $\mathbf{A}\mathbf{V} \in \mathbb{R}^{T \times d}$ , and outputs from all $M$ heads are concatenated and multiplied by an output projection $\mathbf{W}_o \in \mathbb{R}^{D \times D}$ to produce the final multi-head attention output, which flows into normalization and the feedforward network.

The limitation MTA addresses: each element $a_{ij}$ depends ONLY on $\mathbf{q}_i$ and $\mathbf{k}_j$ . There is no mechanism for $a_{ij}$ to be influenced by $a_{i-1, j}$ (the attention from the previous query to the same key), or by $a_{i, j+1}$ (attention to a neighboring key), or by attention patterns in other heads. All information needed to decide "is token $j$ relevant to token $i$ ?" must be encoded in those two vectors alone.

Key-Query Convolution

The core innovation of MTA is applying a learned 2D convolution operation over the attention logits (or weights) along the query and key dimensions. This allows the attention weight $a_{ij}$ to incorporate information from a neighborhood of nearby query positions and key positions.

Pre-softmax convolution (the default version): The convolution is applied to the raw attention logits before the softmax normalization:

$\mathbf{A} = \text{Softmax}\left(\text{Conv2d}_\theta(\hat{\mathbf{A}})\right)$

where $\text{Conv2d}_\theta$ is a 2-dimensional convolution with learned kernel weights $\theta$ and kernel sizes $(c_q, c_k)$ in the query and key dimensions respectively. The batch and head dimensions remain independent — each head has its own convolution parameters.

Expanding the convolution operation, the attention weight $a_{ij}$ from query $i$ to key $j$ is computed as:

$a_{ij} = \text{Softmax}\left( \sum_{i'=0}^{c_q-1} \sum_{j'=-\lfloor c_k/2 \rfloor}^{\lceil c_k/2 \rceil - 1} \mathbf{1}_{i \geq j - j'} \, \theta_{i',j'} \, \mathbf{q}_{i-i'} \mathbf{k}_{j-j'}^\top / \sqrt{d} \right)$

where $c_q$ is the query kernel size (how many past query positions to include), $c_k$ is the key kernel size (how many neighboring key positions to include), $\theta_{i',j'}$ are the learned convolution weights at relative position $(i', j')$ within the kernel, $\mathbf{q}_{i-i'}$ is the query vector at relative position $i-i'$ (looking backward in the query sequence), and $\mathbf{k}_{j-j'}$ is the key vector at relative position $j-j'$ (looking in both directions along the key sequence). The indicator $\mathbf{1}_{i \geq j - j'}$ ensures we never incorporate information from key positions that correspond to future tokens (enforcing causality).

What it computes: instead of $a_{ij}$ being just $\mathbf{q}_i \cdot \mathbf{k}_j$ , it becomes a learned weighted sum of dot products between queries in a $c_q$ -length window ending at position $i$ and keys in a $c_k$ -length window centered near position $j$ . Each weight $\theta_{i',j'}$ is learned and can be positive (amplifying), negative (suppressing), or near-zero (ignoring). The softmax then normalizes across the key dimension to produce valid attention probabilities.

Why this form: this allows the model to recognize patterns across attention pairs. For example, consider the "Alice and rabbit" scenario. If $\mathbf{q}_i$ is the query at token "rabbit" and $\mathbf{q}_{i-2}$ is the query at token "Alice" (2 positions back), then a kernel with $c_q = 4$ and learned weights that match the pattern "query at offset 0 has high dot product with keys here AND query at offset -2 has high dot product with keys in this same region" can amplify attention to positions where both conditions are satisfied. The kernel $\theta$ learns what patterns of multi-token similarity are predictive of relevance.

Practical masking simplification: implementing the exact causal masking in Equation 4 (the indicator function $\mathbf{1}_{i \geq j - j'}$ ) requires modifying CUDA convolution kernels, which is complex. The paper instead proposes a simpler two-pass masking approach:

$\mathbf{A} = \text{Softmax}\left(\text{Mask}_{-\infty}\left(\text{Conv2d}_\theta\left(\text{Mask}_0(\hat{\mathbf{A}})\right)\right)\right)$

Here, the first mask uses 0 (instead of $-\infty$ ) on future positions so those zeroed-out logits don't contribute to the convolution output, and the second mask uses $-\infty$ (before softmax) to ensure no information leakage from future tokens through the convolution operation itself. This masks slightly more aggressively than strictly necessary — it zeros out some past-key interactions that would be legitimate — but is simpler to implement and "prevents information leakage, so we use it as our default."

Post-softmax convolution variant: the convolution can also be applied after the softmax, on the attention weights (probabilities) rather than the logits:

$\mathbf{A} = \text{Mask}_0\left(\text{Conv2d}_\theta\left(\text{Softmax}\left(\text{Mask}_{-\infty}(\hat{\mathbf{A}})\right)\right)\right)$

Here, the interaction between attention weights becomes additive rather than multiplicative (in log-space). In the pre-softmax version, the convolution sums dot products in log-space, so the effect is multiplicative after exponentiation ( $\exp(\sum \dots) = \prod \exp(\dots)$ ). In the post-softmax version, the convolution operates directly on probabilities, so the effect is a linear combination: $a_{ij}^{\text{new}} = \sum_{i',j'} \theta_{i',j'} a_{i-i', j-j'}$ .

Why both exist: the pre-softmax version (multiplicative interaction) is the default and tends to work slightly better (the key-query kernel patterns shown in Figure 4 and Appendix H are most naturally interpreted as multiplicative interactions — e.g., amplifying attention when both a current query-key match and a nearby match are strong). However, the post-softmax version (additive interaction) is conceptually simpler and the paper experiments with both, finding only 0.01-0.04 perplexity differences between orderings (Table 5 bottom rows).

Kernel size interpretation: with $c_q = 6$ and $c_k = 11$ (the default configuration in the 880M experiments), each attention weight $a_{ij}$ is influenced by the query at position $i$ plus the 5 previous queries, and by the key at position $j$ plus the 5 keys on either side. This creates a $6 \times 11$ receptive field over the attention map. The kernel weights $\theta_{i',j'}$ are separate per head, so head 1 might learn to detect one pattern (e.g., a sequence of matching tokens) while head 2 learns another (e.g., isolated high matches on both sides of a gap).

Head Mixing Convolution

While the key-query convolution allows mixing within each head across time steps, the head mixing convolution allows mixing across heads — different heads' attention patterns can be combined to produce new attention patterns.

Operation: all $M$ heads in a layer are divided into groups of $c_h$ heads each (so there are $M / c_h$ groups total). Within each group, a non-overlapping convolution (or equivalently, a fully-connected linear combination) is applied across heads. For head kernel size $c_h = 2$ :

$\mathbf{A}^1_{\text{new}} = w_{11}\mathbf{A}^1 + w_{12}\mathbf{A}^2$

$\mathbf{A}^2_{\text{new}} = w_{21}\mathbf{A}^1 + w_{22}\mathbf{A}^2$

where $\mathbf{A}^1$ and $\mathbf{A}^2$ are the attention weight matrices from two different heads in the same group, and $w_{11}, w_{12}, w_{21}, w_{22}$ are learned mixing weights (the head convolution kernel).

What it computes: new attention patterns for each head that are learned linear combinations of the original attention patterns from all heads in the group. Head 1's new attention map might be $0.8 \cdot \text{head1} + (-0.5) \cdot \text{head2}$ , subtracting out the pattern that head 2 found from head 1's attention.

Why this form: this directly addresses the "Alice and rabbit" problem at the head level. Suppose head 1 learns to attend wherever "Alice" is mentioned (its query vectors detect Alice-like contexts), and head 2 learns to attend wherever "rabbit" is mentioned. Without head mixing, these two attention maps remain separate until they're combined at the output value stage in the feedforward network. With head mixing, the model can directly compute $\mathbf{A}^1_{\text{new}} = w_{11}\mathbf{A}^1_{\text{Alice}} + w_{12}\mathbf{A}^2_{\text{rabbit}}$ , and if $w_{11} > 0$ and $w_{12} > 0$ , the new attention map for head 1 amplifies positions where either Alice or the rabbit is mentioned — or with more sophisticated kernel patterns, positions where both co-occur.

The paper notes an important equivalence: head mixing can be viewed as increasing the effective rank of the value-output computation. As shown in Appendix B, for post-softmax mixing:

$\mathbf{O} = \mathbf{A}^1_{\text{new}}\mathbf{V}^1\mathbf{W}^1_o + \mathbf{A}^2_{\text{new}}\mathbf{V}^2\mathbf{W}^2_o = \mathbf{A}^1\left(w_{11}\mathbf{V}^1\mathbf{W}^1_o + w_{21}\mathbf{V}^2\mathbf{W}^2_o\right) + \mathbf{A}^2\left(w_{12}\mathbf{V}^1\mathbf{W}^1_o + w_{22}\mathbf{V}^2\mathbf{W}^2_o\right)$

This can be rewritten as $\mathbf{A}^1 \mathbf{H} \hat{\mathbf{W}}^1_v \hat{\mathbf{W}}^1_o + \mathbf{A}^2 \mathbf{H} \hat{\mathbf{W}}^2_v \hat{\mathbf{W}}^2_o$ where the new projections have dimension $2d$ instead of $d$ — effectively doubling the rank. However, the paper notes this is "not truly identical to $2\times$ rank because the parameters of the two heads are not independent" — the weight sharing imposed by the convolution structure is the key difference that makes this more parameter-efficient than simply doubling head dimensions.

Pre-softmax vs. post-softmax: like key-query convolution, head mixing can be applied before or after the softmax. In the pre-softmax version:

$\hat{\mathbf{A}}^1_{\text{new}} = w_{11}\hat{\mathbf{A}}^1 + w_{12}\hat{\mathbf{A}}^2$

this mixes attention logits (pre-softmax), which after softmax becomes a nonlinear combination. In the post-softmax version (Equation 7), the mixing is linear in probability space. Both are experimented with.

Head kernel size and related methods: the default configuration uses $c_h = 16$ — heads are divided into groups of 16, and within each group, a $16 \times 16$ linear transformation is learned. This generalizes several prior methods:

Talking-heads attention (Shazeer et al., 2020) applies linear projections across ALL heads before and after softmax, which is equivalent to MTA's head convolution with $c_h = M$ (a single group containing all heads). MTA's grouped version is more parameter-efficient and allows different mixing patterns in different head groups.
Differential Transformer (Ye et al., 2024) computes $\mathbf{A} = \text{softmax}(QK^\top/\sqrt{d}) - \lambda \cdot \text{softmax}(QK^\top/\sqrt{d})$ for two parallel attention maps. This can be viewed as a special case of head mixing with $c_h = 2$ where the weights are fixed to $[1, -\lambda]$ .

Observed kernel patterns: the paper shows (Appendix H, Figure 15) that common head convolution patterns include:

Identity with scaling: the head mostly passes through its own attention pattern, amplified or attenuated ( $w_{ii} \approx \alpha, w_{ij} \approx 0$ for $i \neq j$ ).
Contrasting: one head's attention is subtracted from another's ( $w_{12} \approx -1$ ), creating differential attention similar to the DIFF Transformer mechanism.
Scaling with depth: when group normalization is NOT used, the kernel weights increase in magnitude with layer depth (Appendix Figure 16), which the authors interpret as "competing with the residual stream, which gets larger with the model's depth." Group normalization eliminates this pattern by normalizing before the residual addition.

Putting Everything Together: The Full MTA Module

The key-query convolution and head mixing convolution can be composed in multiple ways depending on whether each is applied before or after the softmax.

Both pre-softmax (3D convolution view): when both operations are applied before the softmax, they can be implemented as a single 3-dimensional convolution over the attention logits, as illustrated in Figure 2. Two of the dimensions span query and key positions (as in the key-query convolution in Section 3.1), and the third dimension spans heads within groups of size $c_h$ . The same 3D convolution kernel $\theta_{i', j', h'}$ simultaneously captures:

Interactions across queries (dimension 1, size $c_q$ )
Interactions across keys (dimension 2, size $c_k$ )
Interactions across heads (dimension 3, size $c_h$ )

This unified 3D convolution is computationally appealing because it can be implemented as a single operation, but the paper's experiments apply key-query convolution on only 1/4 of layers while head convolution is applied to all layers, making the separate composition more relevant.

Both post-softmax: applying both convolutions after softmax on the attention probabilities. This is the additive-interaction version for both axes.

Mixed (key-query pre-softmax, head post-softmax): apply the key-query convolution before softmax (using Equation 5), then apply the head mixing convolution (Equation 7) on the resulting attention probabilities. This is the configuration used in the paper's default model (Table 5, top row: key-query conv pre-sm ✓, head conv post-sm ✓).

Why this composition matters: the paper's ablations (Table 5) show that all three configurations work (perplexity differences of only 0.01-0.04), but the mixed version (key-query pre-softmax, head post-softmax) with scalar gating achieves the best perplexity of 10.90. The authors interpret this as: multiplicative interaction in the time dimension (pre-softmax key-query) better captures conjunctive patterns ("both Alice AND rabbit are here"), while additive interaction in the head dimension (post-softmax head mixing) better captures differential/subtractive patterns (the common "subtract one head's attention from another" pattern observed in kernel visualizations).

The overall MTA flow (Figure 1, right):

Compute $\hat{\mathbf{A}} = \mathbf{Q}\mathbf{K}^\top / \sqrt{d}$ for all heads (standard).
Apply key-query convolution on $\hat{\mathbf{A}}$ (pre-softmax), potentially on a subset of layers.
Apply causal masking and softmax to get attention probabilities.
Apply head convolution to mix attention across groups of heads (post-softmax).
Repeat: apply another MTA convolution (both key-query and head) post-softmax (the "repeat this operation" arrow in Figure 1).
Compute attention output $\mathbf{A}\mathbf{V}$ per head, concatenate, apply output projection.
Apply group normalization with learned sigmoid gating.

The repetition of MTA convolution both before and after softmax is a design choice — the paper states "we apply the MTA convolution both pre-softmax and post-softmax" (Section 4), and the post-softmax repetition allows the additive interactions to further refine the attention patterns after the softmax normalization has already produced probabilities.

Group Normalization with Sigmoid Gating

The final component of MTA is a normalization mechanism applied to the concatenated attention outputs before they enter the residual stream and feedforward network.

Operation: after the MTA-modified attention produces per-head outputs $\mathbf{A}^h\mathbf{V}^h$ for each head $h$ , these are concatenated to a $D$ -dimensional vector. Group normalization is applied with a learned scaling factor, but instead of using layer-dependent depth scaling (as in DIFF Transformer), MTA uses a learned sigmoid gate per head:

$\text{output} = \text{GroupNorm}(\text{Concat}(\mathbf{A}^1\mathbf{V}^1, \dots, \mathbf{A}^M\mathbf{V}^M)) \odot \sigma(\mathbf{g})$

where $\mathbf{g} \in \mathbb{R}^M$ is a learned vector with one scalar gate per head, $\sigma(\cdot)$ is the sigmoid function (outputting values in $(0, 1)$ ), and $\odot$ is element-wise multiplication.

What it computes: group normalization normalizes the activations across the head dimension (treating each head's $d$ -dimensional output as a group), then each head's normalized output is multiplied by a learned gating value between 0 and 1. Heads with $\sigma(g_h) \approx 0$ are effectively turned off; heads with $\sigma(g_h) \approx 1$ pass through fully.

Why this form: the sigmoid gating addresses a specific problem observed in Transformer architectures — the residual stream accumulates activations across layers, creating a "competition" between the attention output and the already-accumulated signal. As shown in Appendix Figure 16, without normalization, MTA kernel weights increase in magnitude with depth to compete with the growing residual stream. Group normalization normalizes the attention output to a consistent scale regardless of depth, and the sigmoid gate gives each head a learned "volume knob" — the model can learn to selectively amplify important heads and suppress noisy ones.

The comparison to alternatives (Table 5):

Scalar gating (MTA default): 10.90 perplexity
Layer-dependent depth scaling (DIFF Transformer style): 10.92 perplexity
No scaling at all: 11.03 perplexity
Layer-norm scaling (Sun et al., 2025): 11.41 perplexity — notably worse, showing that not all normalization schemes work equally well for attention output

The paper uses group normalization (treating each $d$ -dimensional head as a group) rather than layer normalization because the head dimension is the natural grouping — heads compute independent attention patterns, and normalizing across them without mixing their statistics preserves their independence while controlling scale.

Practical Design Choices, Hyperparameters, and Training Configuration

The paper makes several specific design decisions that balance performance, parameter efficiency, and computational cost:

Selective key-query convolution: key-query convolution is applied only to every 4th layer of the Transformer (6 out of 24 layers for the 880M model), while head convolution is applied to all layers. The ablation in Figure 5 (right) shows that even 2 MTA layers outperform baselines, with 6 layers striking "a balance between performance and additional complexity." This is a crucial efficiency decision — key-query convolution increases computation more than head convolution (it operates on the full $T \times T$ attention map), so applying it sparingly controls overhead while still providing multi-token capability at key layers.

Kernel sizes: the default configuration uses $c_q = 6$ (query kernel span of 6 tokens backward), $c_k = 11$ (key kernel span of 11 tokens, $\pm 5$ around center), and $c_h = 16$ (head groups of 16). The ablation in Table 5 (middle rows) tests $(c_q=4, c_k=9)$ and $(c_q=8, c_k=13)$ , finding similar kernel patterns but slightly different final performance — $c_q=6, c_k=11$ achieves 10.90 with full MTA, while the smaller and larger kernels achieve 10.95 and 11.23 respectively (but these comparisons are confounded with different normalization choices, making direct comparison difficult). The choice of $c_h=16$ is ablated in Figure 6 (left), which shows monotonically improving perplexity with larger head kernel sizes (from 1.0 to 16.0 on the x-axis), motivating the choice of 16 as a practical maximum given the total head count.

Kernel initialization: the convolution kernels are initialized to identity — meaning the initial MTA model behaves exactly like a standard Transformer before training begins. Concretely, for key-query convolution, the kernel weight $\theta_{0, 0}$ (the center position) is initialized to 1, and all other positions to 0, so the convolution initially just passes through the original attention logit unchanged. For head convolution, the diagonal weights $w_{ii}$ are initialized to 1 and off-diagonals $w_{ij} (i \neq j)$ to 0. This identity initialization "leads to better convergence and final performance" compared to zero or constant (0.3) initialization, which reduces perplexity by 0.02 and 0.08 respectively.

Why identity initialization works: starting from identity means the model begins training with the standard attention mechanism it already knows how to optimize, and the convolution kernels gradually learn to deviate from identity to capture multi-token patterns. This is essentially a form of residual learning at the initialization level — the model learns corrections to standard attention rather than learning attention from scratch with a different mechanism.

Parameter overhead: the convolution kernels add very few parameters. For the 880M model (Table 8):

Standard Transformer: 876,553,728 parameters
MTA: 876,583,320 parameters — an increase of only ~29,592 parameters (0.0034%)

The key-query convolution on 6 layers with $c_q=6, c_k=11$ contributes $6 \times 16 \text{ heads} \times (6 \times 11) = 6,336$ parameters (assuming per-head kernels). The head convolution on all 24 layers with $c_h=16$ contributes $24 \times 16 \text{ groups} \times 16 \times 16 = 98,304$ kernel weights per kernel (and there are both pre and post-softmax head convolutions). The gating and group normalization parameters add another modest amount. The total is negligible compared to the model's existing 876M parameters.

Computational cost: the paper is transparent about the current implementation's inefficiency (Table 9). MTA training achieves only 5.7K tokens per second vs. 54.3K for the standard Transformer on 32 H200 GPUs — roughly 9.5× slower. However, the paper notes that the standard Transformer uses PyTorch's optimized scaled_dot_product_attention with CUDA kernels, while "our MTA implementation does not take advantage of such efficient kernels, which is the major reason behind its lower FLOPS." The theoretical additional computation from the convolutions is modest — a $c_q \times c_k$ convolution over an attention map is $O(T^2 \cdot c_q c_k)$ compared to the $O(T^2 d)$ of the attention matrix multiplication itself — so with optimized kernels, the overhead should be manageable.

Training configuration: the 880M models are trained on SlimPajama for 105B tokens with:

24 layers, 1536 hidden dimension, 16 heads
Context length 2048, RoPE theta 100,000
Batch size 262,144 tokens
Learning rate $1.5 \times 10^{-4}$ , weight decay 0.05
AdamW optimizer, 375 warmup steps
RMSNorm pre-normalization, SwiGLU activation, Rotary Embeddings

For the long-context finetuning (Section 4.3), the same models are trained for an additional 10.5B tokens with context extended to 4096, RoPE theta increased to 500,000, weight decay set to 0, and warmup reduced to 50 steps.

The theoretical intuition — why MTA solves the multi-condition problem: the core mechanism can be understood through the Alice-and-rabbit example from the paper. In standard attention, finding a sentence containing both requires a single query vector to encode both concepts and match against key vectors that also encode both — a quadratic capacity burden. In MTA:

Two separate query tokens ("Alice" at position $i-2$ and "rabbit" at position $i$ ) produce two independent attention maps through their respective dot products with keys.
The key-query convolution (with $c_q = 4$ ) combines these: $a_{ij}^{\text{conv}} = \theta_{0,j'}\mathbf{q}_i\mathbf{k}_{j-j'} + \theta_{2,j'}\mathbf{q}_{i-2}\mathbf{k}_{j-j'} + \dots$ . If $\theta_{0,0} \approx 1$ and $\theta_{2,0} \approx 1$ (both center positions get weight), the attention at position $j$ will be high only when both $\mathbf{q}_i \cdot \mathbf{k}_j$ is high (rabbit match at $j$ ) and $\mathbf{q}_{i-2} \cdot \mathbf{k}_j$ is high (Alice match at $j$ ) — exactly the conjunction pattern needed.

This is why Figure 4 shows a diagonal kernel pattern: it amplifies attention when a sequence of query tokens matches a sequence of keys, rather than just individual token matches. The convolution has effectively turned attention from a "single token lookup" into a "sequence pattern matching" operation.

4. Key Insights and Innovations

Innovation 1: The "Single Token Bottleneck" as a Diagnostic Concept for Attention Limitations

The paper's deepest conceptual contribution is not the MTA mechanism itself, but the diagnostic framing that gives it purpose: the identification and articulation of the single token bottleneck as a fundamental architectural limitation of standard attention. This framing moves beyond the vague sentiment that "attention could be better" to a precise, falsifiable claim about where and why standard attention fails.

What makes this distinctive: prior work on attention limitations has largely focused on capacity issues — attention is quadratic in context length (the efficiency problem), or softmax attention distributes weight too broadly (the focusing problem), or attention heads fail to specialize (the redundancy problem). The single token bottleneck identifies a combinatorial expressiveness issue: the fact that each attention weight conditions on exactly one query-key pair means that attention cannot directly express pattern matches over sets of tokens. This is a fundamentally different category of limitation — it's not about how attention distributes its weight, but about what information can inform that distribution in the first place.

The diagnostic power of this framing is demonstrated by the toy task (Section 4.1, Table 1). The task is deliberately designed to require exactly the capability that the single-token bottleneck claims is missing: locate a block containing a set of L letters, where no single letter suffices to identify the target. Standard attention with 4 layers and 256 hidden dimensions achieves 51.6% error on this task (with high variance: ±43.1%), meaning it sometimes learns a workaround and sometimes fails completely — exactly what you'd expect from a model forced to compress multi-condition information into single vectors via learned intermediate representations. MTA, with key-query convolution size c_q = 2 matching the number of query letters, achieves 0.1% error with near-zero variance. This is not just a performance improvement; it's a mechanistic validation of the diagnostic claim — give attention the ability to combine multiple token signals, and the bottleneck disappears; force it through single vectors, and the problem is genuinely hard even in this minimal setting.

Comparison to prior framings: the closest prior conceptual framing is the "lost in the middle" phenomenon (Liu et al., 2024), which diagnoses where attention fails (middle of long contexts) but not why in mechanistic terms. Methods like DIFF Transformer (Ye et al., 2024) and Talking-heads (Shazeer et al., 2020) improve attention by adding operations after attention weights are computed (differential subtraction, linear projection across heads), but they leave intact the fundamental computation unit — each raw attention logit still depends on a single query-key pair. The single-token bottleneck framing explains why these methods help but don't fully solve the problem: they make better use of attention patterns once computed, but don't expand what information can inform those patterns. This diagnostic distinction is what justifies MTA's architectural intervention — convolution before softmax on the attention logits — as a qualitatively different kind of change rather than an incremental refinement.

Significance beyond raw performance: the single-token bottleneck concept is portable. It provides a language for analyzing failure modes in any attention-based architecture: "is this task hard because it requires conditioning attention on multiple tokens simultaneously?" It also predicts where MTA should help most — tasks requiring multi-fact retrieval, conjunction queries, or sequence-pattern matching — and where it should help least — tasks where relevance can be determined by single-token similarity. The paper's results partially validate this prediction: MTA's largest gains come on Needle-in-Haystack with multiple needles (Figure 3, requiring multi-fact location) and BabiLong QA tasks requiring multiple supporting facts or argument relations (Figure 5 left), while its benchmark improvements on single-hop tasks like BoolQ and PIQA are more modest (Table 3, ~0.5-1% absolute gain). This diagnostic specificity — predicting differential improvement rather than uniform gains — is what elevates the framing from a post-hoc justification to a genuine conceptual advance.

Is this fundamental or incremental? The diagnosis is fundamental. It identifies a structural property of the attention operation that was previously overlooked — not a hyperparameter choice, not a training recipe issue, but an inherent limitation of conditioning attention on vector pairs. The fact that MTA's solution (multi-dimensional convolution over attention logits) is one of potentially many ways to address this bottleneck doesn't diminish the diagnostic contribution; it strengthens it by showing the bottleneck is real and that addressing it yields consistent improvements across scales (Figure 6 right, from 300M to 1B parameters).

Innovation 2: Convolution on Attention Logits as a General Mechanism for Multi-Token Interaction

The paper's architectural contribution — applying learned convolution kernels over attention logits across query, key, and head dimensions — is more than a specific implementation trick. It establishes convolution over the attention map as a general design pattern for enabling interactions among attention computations that were previously isolated.

What makes this distinctive: prior work has applied convolution in Transformers in several ways, but always outside the core attention computation:

Convolution on input sequences before attention (e.g., Conformer, Gulati et al. 2020, used in Llama 3's speech encoder) — this is preprocessing, not attention modification.
Convolution on keys and values for compression (Liu et al., 2018) — this is about efficiency, not expressiveness.
Convolution on attention weights in the key dimension only for length extrapolation (Zheng et al., 2024) — this is single-axis and positional, not about multi-token semantic interaction.
Convolution on the attention output (post-value aggregation) — again, after the attention computation is complete.

MTA's distinctive move is applying convolution to the attention logits themselves, before the softmax, and doing so across three axes (queries, keys, heads) simultaneously. This is not an incremental extension of prior convolution-in-attention work; it's a category shift. The previous approaches treat convolution as an adjunct to attention — preprocessing inputs or postprocessing outputs. MTA treats convolution as part of the attention mechanism, intervening at the exact point where similarity scores are computed.

Why this matters conceptually: the attention logit matrix Â_{ij} = q_i · k_j / √d is effectively a similarity map between all query positions and all key positions. Applying convolution to this map before softmax means the model learns spatial patterns in similarity space — it can recognize that a certain configuration of high-similarity pairs (e.g., a diagonal streak indicating a sequence match, or two nearby peaks indicating co-occurrence) is predictive of relevance, even if no single pair would trigger attention strongly. This transforms attention from a pointwise similarity operation into a local pattern detector over similarity maps. The fact that each head learns its own convolution kernel means different heads can learn to detect different patterns — some might detect sequences, others might detect co-occurrence, others might learn edge-detection-like filters that amplify the boundaries of high-similarity regions.

This interpretation is validated by the kernel visualizations in Appendix H. The paper shows that learned kernels exhibit interpretable patterns (Section 4.5): diagonal kernels that amplify attention when a query token sequence matches a key token sequence (Figure 4), "priming" kernels that amplify if the same key was attended by previous queries, and "edge detecting" kernels that amplify the first or last of multiple contiguous high-attention keys. These are not random learned weights — they correspond to meaningful spatial operations over the similarity map, and different heads converge to qualitatively different patterns (Figures 9-14 show striking diversity across heads and layers). This emergent specialization is exactly what you'd expect if the convolution mechanism is genuinely enabling new computational capabilities rather than just adding parameters.

Comparison to multi-head mechanisms: the closest prior architectural idea is Talking-heads attention (Shazeer et al., 2020), which applies linear projections across heads — essentially a 1 × 1 convolution in head-space with global receptive field. MTA generalizes this in two dimensions: it adds spatial (query and key) dimensions to the convolution, and it uses local receptive fields (kernel sizes c_q, c_k, c_h) rather than global linear projections. This locality is conceptually important because multi-token interactions are inherently local — "Alice" and "rabbit" influence each other's attention only if they appear within a few tokens of each other in the query or key sequence. A global linear projection across all head-key-query positions would be wildly overparameterized and would learn spurious long-range interactions. The convolutional inductive bias — that only nearby positions interact — is exactly the right prior for multi-token attention patterns.

The connection to DIFF Transformer (Ye et al., 2024) is similarly deepened by this framing. DIFF Transformer's differential attention softmax(Â₁) - λ · softmax(Â₂) can be seen as a fixed, hand-designed pattern in the head interaction space — subtract one head's attention map from another. MTA's head convolution learns these interaction patterns (and indeed, the paper observes that one common learned pattern is exactly this subtractive operation — Appendix Figure 15), while also learning other patterns (amplification, gating, selective combination) that a fixed differential mechanism cannot express. This is a case of a learned mechanism subsuming a hand-designed one, similar to how learned convolutions in CNNs replaced hand-designed edge detectors.

Significance beyond this paper: the idea of treating attention logits as a feature map to be processed by learned spatial filters opens a design space that this paper only partially explores. The paper uses small separable kernels (6×11 query-key, groups-of-16 head mixing), but the design space includes multi-scale convolution, dilated convolution for longer-range interactions, attention over the convolution outputs, and dynamic kernel generation conditioned on the input. By establishing convolution-on-logits as a viable and beneficial operation, MTA provides a template for a whole class of future attention mechanisms that treat the similarity map as a rich intermediate representation worthy of its own learned processing, rather than as a computation to be immediately consumed by softmax.

Is this fundamental or incremental? The specific convolution mechanism is incremental in implementation (adding a 2D/3D convolution to existing attention code), but the concept — attention logits as processable feature maps rather than ephemeral intermediates — is fundamental. It changes how one thinks about what happens between computing dot products and applying softmax: that intermediate space is now a locus of learned computation, not just a pass-through.

Innovation 3: The Complementary Roles of Pre-Softmax and Post-Softmax Mixing, and the Grouped Head Architecture

A more subtle but technically important contribution is the paper's empirical characterization of where different types of attention interaction should occur relative to the softmax, and the grouped approach to head interaction that balances expressiveness with parameter efficiency.

What makes this distinctive: the paper doesn't just propose one way to do multi-token attention — it systematically explores variations (pre-softmax vs. post-softmax for both key-query and head mixing) and finds that the optimal configuration is asymmetric: key-query convolution works better before softmax (multiplicative interaction in log-space), while head mixing works better after softmax (additive interaction in probability space). This asymmetry is not arbitrary — it reflects a deeper computational principle that the paper doesn't fully theorize but demonstrates empirically (Table 5: mixed configuration achieves 10.90 PPL vs. 10.95 for both-pre-softmax and 10.95 for both-post-softmax, though differences are small).

Why this asymmetry makes sense: pre-softmax key-query convolution enables multiplicative gating of attention patterns. When the convolution sums dot products before softmax, the exponentiation in softmax converts this sum into a product: softmax(∑ θ · qk) ∝ ∏ exp(θ · qk). This means attention is amplified only when all contributing query-key pairs have high similarity — a natural AND operation over multiple token matches. This is exactly what's needed for conjunction queries ("find where Alice AND rabbit are mentioned"). Post-softmax head mixing, by contrast, enables additive/subtractive combination of already-computed probability distributions: A_new = w₁A₁ + w₂A₂. This is a natural OR/NOT operation — "attend where head 1 attends OR where head 2 attends, but NOT where head 3 attends." The subtractive pattern observed in kernel visualizations (Appendix Figure 15, the contrasting pattern of w₁₂ ≈ -1) directly implements noise cancellation — subtracting out a head's attention pattern that represents ubiquitous or irrelevant features.

This division of labor — spatial (key-query) interaction in log-space for conjunction, cross-head interaction in probability-space for differential combination — is not an obvious design choice. A naive design would apply both operations symmetrically (both pre-softmax or both post-softmax). The paper's ablation (Table 5, bottom rows) shows that both symmetric configurations work (perplexity increases of only 0.01–0.04), but the asymmetric configuration is consistently best. This is a small empirical finding with large architectural implications: log-space and probability-space are suited for different types of attention combination, and a well-designed attention mechanism should use both.

The grouped head architecture: a second architectural insight is that head mixing should operate over local groups rather than globally across all heads. MTA divides the M heads into groups of c_h = 16 and learns a separate linear transformation within each group. This contrasts with Talking-heads attention, which learns global linear projections across all heads (equivalent to c_h = M). The grouped design has two advantages:

Parameter efficiency: (M/c_h) × c_h² = M · c_h parameters for grouping vs. M² for global mixing. For M=16, the difference is 256 vs. 256 — no savings at this scale — but for larger models with M=64 or M=128 heads, the grouped approach scales linearly while global mixing scales quadratically.
Functional specialization: different groups can learn different interaction patterns. The paper doesn't extensively analyze this, but the kernel diversity observed across layers (Figures 9-14) suggests that different layers' head groups are indeed learning different mixing strategies — some layers might specialize in noise cancellation (subtractive patterns) while others specialize in amplification (additive patterns).

Comparison to prior head interaction methods: prior work on cross-head interaction has oscillated between two extremes — either no interaction (standard multi-head attention, where heads are completely independent until output concatenation) or global interaction (Talking-heads, where all heads interact through learned projections). MTA's grouped convolution occupies a principled middle ground: heads interact locally in learned groups, similar to how convolutional networks use local receptive fields rather than fully-connected layers to capture spatial structure. This is analogous to the transition from fully-connected to convolutional architectures in computer vision — the inductive bias of locality matters for head interaction just as it matters for spatial interaction.

The normalization-gating interaction: the paper's finding that sigmoid gating outperforms layer-dependent depth scaling (10.90 vs. 10.92 PPL, Table 5) is another small but principled choice. Layer-dependent scaling (as in DIFF Transformer) assumes attention output magnitude should grow with depth to compete with the residual stream — a one-size-fits-all depth schedule. Sigmoid gating gives per-head learned scaling, allowing the model to independently adjust each head's contribution at each layer. The observation that kernel weights increase with depth when group normalization is absent (Appendix Figure 16) — but this pattern disappears with group normalization — reveals that the normalization isn't just stabilizing training; it's decoupling head importance from depth, allowing heads to have different effective "volumes" independently of their position in the network.

Is this fundamental or incremental? The pre/post-softmax asymmetry finding is incremental but practically important — it's a design rule that future attention modifications should incorporate. The grouped head architecture is an incremental refinement of Talking-heads but with a qualitatively different inductive bias (locality) that may prove important at larger scales. Together, these findings form a design pattern for attention mechanisms: treat log-space and probability-space as distinct computational regimes suitable for different operations, and use local rather than global head interaction.

Innovation 4: Identity Initialization as a Strategy for Architectural Modification of Pretrained Models

While much of the paper focuses on training from scratch, one of its most practically significant contributions is the demonstration that identity-initialized MTA convolutions can be added to already-pretrained models and fine-tuned with continued training, without requiring full retraining. This is established both through the identity initialization strategy (Section 4.6) and the preliminary continued-training experiments with both the authors' own 1.4B models and open-source Llama 3 models (Appendix I, Table 10).

What makes this distinctive: architectural modifications to Transformers typically require training from scratch — you can't just insert a new operation into a pretrained model and expect it to work, because the new operation's random initialization would corrupt the model's existing representations. The standard approach is to either (a) design the modification as a trainable adapter that starts near-zero and is fine-tuned while keeping the base model frozen, or (b) train the full modified architecture from scratch, accepting the cost.

MTA takes a third approach enabled by its convolutional form: initialize the convolution kernels to identity, so the initial MTA-augmented model produces exactly the same outputs as the original pretrained model. For key-query convolution, θ_{0,0} = 1 and all other kernel weights 0 means Conv2d(Â) = Â — the convolution is a no-op. For head convolution, w_{ii} = 1 and w_{ij (i≠j)} = 0 means heads pass through unchanged. Then, continued training allows the kernels to gradually deviate from identity, learning multi-token interaction patterns while the model's existing knowledge is preserved.

This is a deployment-friendly innovation because it means:

MTA can be retrofitted into existing pretrained models without architectural changes beyond inserting the convolution operations.
The continued training cost is much lower than pretraining from scratch (the appendix experiments use only 5.3–10.5B tokens of continued training).
The model's existing capabilities are not catastrophically disrupted during the transition.

What the continued-training experiments show (Appendix I, Table 10):

A 1.4B model pretrained with standard attention, then continued-trained with MTA for 10.5B tokens, achieves 10.61 PPL — better than the same model continued-trained with standard attention (10.69) and approaching the from-scratch MTA model (10.44).
Llama 3.2 1B, 3B, and Llama 3.1 8B all show consistent perplexity improvements when continued-trained with MTA compared to continued-trained with standard attention: MTA reduces perplexity by 0.07–0.08 points across all three model sizes.
These experiments use the same identity-initialized MTA insertion strategy and modest continued training budgets (5.3B tokens for Llama models).

Why this matters beyond MTA: the identity-initialization-as-adapter strategy is generalizable to any architectural modification that can be expressed as a learnable transformation with a natural identity parameterization. This provides a template for future Transformer modifications: design your new operation so that it reduces to identity at initialization, insert it into pretrained models, and fine-tune. The success of this approach with MTA (on models from 1B to 8B parameters, from two different model families) suggests it's a robust strategy, not a fluke.

Comparison to prior adapter methods: traditional adapters (e.g., LoRA, prompt tuning, (IA)³) add new parameters that start near-zero and learn to modify existing representations. MTA's identity-initialized convolutions are different in kind: they start as exact identity functions and learn to transform the computation. This means the adapter doesn't need to learn to preserve existing knowledge from a zero or random start — it starts already preserving everything, and only needs to learn beneficial deviations. This is a stronger form of the "do no harm" principle in continued training.

Is this fundamental or incremental? The identity initialization strategy for architectural modification is conceptually fundamental but experimentally preliminary in this paper. The continued-training results are in an appendix and use relatively small token budgets — it remains to be seen whether the approach scales to larger models (70B+) and longer continued training. But the principle is sound and the initial results are uniformly positive across model scales and families, making this one of the paper's most actionable contributions for practitioners who want to adopt MTA without training models from scratch.

5. Experimental Analysis

Evaluation Methodology

Dataset. Language model pretraining uses the SlimPajama dataset (Soboleva et al., 2023), a cleaned and deduplicated 627B-token version of RedPajama, with all models trained on 105B tokens of this corpus. Downstream evaluation spans standard academic benchmarks (BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC easy/challenge, OpenBookQA, and MMLU) plus long-range dependency tasks: LAMBADA (Paperno et al., 2016; Radford et al., 2019), Needle-In-A-Haystack (Kamradt, 2023), and BabiLong (Kuratov et al., 2025). The motivating toy task uses a synthetic dataset of 1M training sequences and 1K held-out test sequences, where each sequence consists of blocks of N random letters followed by L query letters, with the objective being to locate the block containing all query letters.
Base model(s). The primary experiments use 880M-parameter decoder-only Transformer models (24 layers, 1536 hidden dimension, 16 heads) following the LLaMA architecture (Touvron et al., 2023) with RMSNorm pre-normalization, SwiGLU activations, and Rotary Position Embeddings. Scaling law experiments span 300M, 550M, 880M, and 1B parameter models (Table 7 provides hyperparameters for each size), while continued-training experiments extend to 1.4B proprietary models and open-source Llama 3.2 1B, 3B, and Llama 3.1 8B models. PaLM 2-S*, mentioned in the query's reference example, is not used in this paper — the reference example was from a different paper and should not be applied here.
Metrics. Validation perplexity on SlimPajama subdomains (arxiv, book, c4, cc, github, stackexchange, wikipedia) serves as the primary metric during pretraining, with average perplexity across all domains reported. For downstream benchmarks, accuracy is used for BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC, OpenBookQA, and BabiLong; MMLU reports 5-shot accuracy; LAMBADA reports perplexity; and Needle-In-A-Haystack reports retrieval accuracy (fraction of queries for which the model correctly extracts the target information from the context). The toy task reports error rate as the fraction of test sequences where the model fails to output the correct target block tokens.
Baselines. Four architectures are compared: (1) Standard Transformer (Vaswani et al., 2017) — the vanilla multi-head attention baseline. (2) Differential Transformer (DIFF Transformer) (Ye et al., 2024) — computes attention scores as the difference between two separate softmax attention maps, using group normalization with depth-dependent scaling. (3) Talking-heads attention (Shazeer et al., 2020) — applies linear projections across the head dimension before and after softmax operations; the paper notes this is "related to our head convolution approach." (4) MTA — the proposed method with key-query convolution on every 4th layer, head convolution on all layers, and group normalization with sigmoid gating. The DIFF Transformer baseline uses half as many heads to account for doubled head dimension, ensuring the same total parameter count.
Generation budget / compute accounting. Test-time compute is measured in number of generated tokens during training (105B tokens for pretraining, plus 10.5B tokens for long-context finetuning), not inference-time budget as in the reference example. For fair architectural comparison, all models are trained with identical token budgets, batch sizes, context lengths, and optimizer settings. Parameter counts are carefully tracked (Table 8): the standard Transformer has 876,553,728 parameters, while MTA adds only 29,592 additional parameters (0.0034%) from the convolution kernels and gating parameters. Computational cost during training is reported in Table 9 as FLOPS (in units of 10^13) and tokens per second (TPS), though the paper acknowledges their unoptimized MTA implementation achieves only 2.6×10^13 FLOPS and 5.7K TPS versus 25.0×10^13 FLOPS and 54.3K TPS for the standard Transformer using PyTorch's optimized scaled_dot_product_attention.
Cross-validation / statistical protocol. For pretraining experiments, "each model, we conduct training twice and report average validation perplexity" (Section 4.2), providing two independent training runs per architecture to assess variance. For downstream benchmark evaluation (Table 3), "Results are averaged over two model training runs for each method." For the toy task (Table 1), experiments are run with three different random seeds and both mean error rates and standard deviations are reported. The long-context finetuning results (Table 2 bottom section) appear to report single-run results, as no averaging is mentioned.

Main Quantitative Results

Language Modeling Pretraining (Table 2, top section)

After 105B tokens of pretraining on SlimPajama with 2048 context length, MTA achieves an average validation perplexity of 10.91 across the seven SlimPajama subdomains, compared to 11.25 for the standard Transformer — a reduction of 0.34 perplexity points. This represents a meaningful improvement: the DIFF Transformer achieves 11.15 (0.10 better than Transformer baseline) and Talking-heads achieves 11.04 (0.21 better). The improvement is consistent across all seven subdomains, with the largest absolute gains on high-perplexity domains: book (13.09 vs. 13.47, −0.38), c4 (19.63 vs. 20.20, −0.57), and cc (14.00 vs. 14.41, −0.41). The smallest gain is on arxiv (4.54 vs. 4.65, −0.11), likely because lower absolute perplexity leaves less room for improvement. The standard deviation across two training runs for MTA is 0.01, indicating training stability.

Long-Context Finetuning (Table 2, bottom section)

After extending context length from 2048 to 4096 and finetuning for an additional 10.5B tokens, MTA maintains its advantage: 10.65 average perplexity versus 11.02 for the standard Transformer (a 0.37-point gap), 10.89 for DIFF Transformer, and 10.88 for Talking-heads. The gap between MTA and the standard Transformer slightly widens compared to pretraining (0.37 vs. 0.34), suggesting MTA's multi-token attention provides increasing benefit as context length grows. The improvement is again consistent across all subdomains, with the largest gaps on book (12.77 vs. 13.18, −0.41) and c4 (19.51 vs. 20.14, −0.63).

Standard Benchmark Evaluation (Table 3)

On nine standard zero-shot benchmarks, MTA achieves an average accuracy of 44.9% across the eight non-MMLU benchmarks (averaged with equal weight) and 25.8% on MMLU (5-shot), compared to the Transformer baseline's 43.7% average and 24.5% MMLU. This represents a +1.2 percentage point improvement on the aggregate benchmark score — a modest but consistent gain. Breaking this down: MTA outperforms the Transformer baseline on 7 of 9 benchmarks, with the largest absolute gains on BoolQ (62.1% vs. 56.2%, +5.9 points) and HellaSwag (39.7% vs. 38.5%, +1.2 points). It underperforms on ARC-Challenge (24.7% vs. 25.9%, −1.2 points) and ties on OpenBookQA (23.6%). Compared to the next-best baseline, Talking-heads (44.4% average), MTA's advantage narrows to +0.5 points, driven primarily by WinoGrande (57.2% vs. 54.5%, +2.7 points) and BoolQ (62.1% vs. 61.9%, +0.2 points). The small standard deviations reported for the Transformer baseline (±0.3 for the average, parenthetical in the "Avg ↑" column of Table 3) suggest these differences, while modest, are likely statistically meaningful, though formal error bars per benchmark are not provided for all methods.

The pattern of benchmark improvements is notable: MTA's gains are largest on tasks that plausibly require multi-fact reasoning or long-range dependency resolution (BoolQ — yes/no questions about passages; HellaSwag — sentence completion requiring discourse coherence; WinoGrande — pronoun resolution requiring contextual integration), while gains on tasks dominated by factual recall or short-range reasoning (PIQA, SIQA) are smaller (71.7% vs. 70.2%, 40.4% vs. 39.9%). This aligns with the paper's claim that MTA should particularly benefit tasks requiring multi-condition attention, though the benchmark suite does not isolate long-context retrieval, making this pattern suggestive rather than definitive.

Long-Range Dependency: LAMBADA (Table 4)

MTA achieves substantially lower perplexity on both LAMBADA variants: 13.2 on LAMBADA standard versus 17.6 for the Transformer baseline (a 4.4-point reduction, or 25% relative improvement), and 8.4 on LAMBADA OpenAI versus 9.5 for the Transformer (a 1.1-point reduction). The gap to the next-best baseline is also substantial: DIFF Transformer achieves 14.9 on LAMBADA standard, meaning MTA's improvement over DIFF is 1.7 points — larger than DIFF's improvement over the standard Transformer (2.7 points). LAMBADA is specifically designed so that the correct next word requires understanding the full preceding discourse, not just the local sentence — making it a direct test of the multi-token retrieval capability MTA targets. The large improvement here (the largest relative improvement on any metric in the paper aside from the toy task) provides strong support for the core motivation.

Long-Range Dependency: Needle-In-A-Haystack (Figure 3)

For multi-needle retrieval with 2K context (pretrained models, Figure 3 left), MTA achieves approximately 85% accuracy with 2 needles, decreasing to roughly 65% with 10 needles, while the standard Transformer drops much more sharply — from roughly 75% at 2 needles to approximately 35% at 10 needles. The performance gap widens with needle count: at 2 needles, MTA's advantage is roughly 10 percentage points; at 6 needles, the gap grows to approximately 25 points; at 10 needles, roughly 30 points. The DIFF Transformer and Talking-heads baselines perform similarly to the standard Transformer, suggesting their attention modifications (differential subtraction, global head projection) do not address the multi-needle challenge the way MTA's key-query convolution does.

For 4K context (finetuned models, Figure 3 right), MTA achieves roughly 70% accuracy with 4 needles versus the Transformer's approximately 30%. The gap between MTA and the standard Transformer is even larger at 4K than at 2K context — for example, at 6 needles the Transformer accuracy appears near or below 20% while MTA maintains roughly 45–50%. This suggests MTA's benefits compound with context length, consistent with the hypothesis that longer contexts impose greater demands on multi-condition attention for precise retrieval.

The paper notes that "each new sample, the needles are shuffled, thus removing bias the models might have toward extracting the needle that was inserted first or last" (Section 4.4), strengthening the claim that MTA's improvement reflects genuine multi-fact retrieval capability rather than positional heuristics.

Long-Range Dependency: BabiLong (Figure 5 left, Figure 7)

Averaged across QA1–5 tasks, MTA consistently outperforms all baselines at every distraction text length from 0K to 4K tokens. At 4K distraction — the most challenging setting — MTA achieves roughly 52% average accuracy versus the standard Transformer's roughly 37%, a 15 percentage point absolute improvement. The gap to the next-best baseline (Talking-heads, approximately 43% at 4K) is roughly 9 points.

The per-task breakdown (Appendix Figure 7) reveals several patterns:

QA1 (single supporting fact): All models perform well with 0K distraction (60–65% accuracy for baselines, approximately 65% for MTA). Adding 4K distraction causes the Transformer to drop to roughly 42% while MTA maintains roughly 57% — a 15-point gap that demonstrates even single-fact retrieval benefits from multi-token attention when distractor length is large.
QA2 (two supporting facts): At 4K distraction, MTA achieves approximately 25% versus the Transformer's roughly 12% — a 13-point advantage. This task requires locating AND combining two pieces of information, directly exercising the conjunction capability MTA is designed for.
QA3 (three supporting facts): All models perform poorly, with MTA achieving only roughly 15% at 4K distraction (random performance is 16.67%, since QA3 has multiple choice structure). The paper acknowledges this: "random performance on QA3 is 16.67%, thus all models perform poorly on QA3." This is an important failure case — it shows that even with multi-token attention, tasks requiring integration of three widely separated facts with heavy distraction remain extremely challenging for models of this scale.
QA4 (two argument relations): MTA achieves roughly 38% at 4K distraction versus the Transformer's roughly 27%. The gap to Talking-heads (roughly 32%) is roughly 6 points.
QA5 (three argument relations): MTA achieves roughly 55% at 4K distraction versus the Transformer's roughly 38% — a 17-point gap, and the largest absolute improvement on any QA task. This task requires tracking a chain of actions across multiple sentences, which aligns with MTA's sequence-pattern-matching capability (diagonal kernel patterns that amplify attention when query sequences match key sequences).

The systematic improvement across tasks with varying numbers of required facts (QA1–QA3) and argument relations (QA4–QA5) is particularly informative: MTA helps at every level of complexity, but the relative improvement is largest when the task demands integration of multiple pieces of information (QA2, QA5). This is consistent with the mechanism — key-query convolution directly enables combining multiple token-match signals into a single attention decision.

Motivating Toy Task (Table 1)

On the synthetic block-location task with N=5 block size and L=2 query letters (two variants: output all tokens or only first/last of the target block), the standard Transformer achieves the following error rates, averaged over 3 seeds with standard deviations:

Output all tokens (N=5): 51.6% error (±43.1%)
Output first token only (N=5): 78.2% error (±1.5%)
Output last token only (N=5): 1.3% error (±0.4%)

The extremely high standard deviation on the "all" variant (±43.1% on a 51.6% mean) indicates that across three seeds, some runs learned to partially solve the task while others failed completely — exactly what one would expect from a model attempting to solve a multi-condition task by compressing information into single vectors via learned intermediate representations, where success depends on whether stochastic optimization happens to find a viable representation.

MTA with key-query convolution (cq=2 to match the L=2 query letters, ck=2N−1=9 to cover a full block) achieves:

All variants, both N=5 and N=8: 0.0–0.1% error, with standard deviations of 0.0–0.1%

This near-perfect performance across all variants and block sizes is the strongest mechanistic validation in the paper: given the exact architectural capability that the single-token bottleneck framework says is missing (the ability to combine multiple token signals before softmax), the problem becomes trivially solvable. The fact that MTA achieves 0.0% error with standard deviation 0.0% (rather than, say, 5% with some variance) demonstrates that the convolution operation is not just helpful but sufficient — it directly implements the required computation rather than requiring the model to learn a complex workaround.

Scaling Laws (Figure 6, right)

Models at four scales (300M, 550M, 880M, 1B parameters) show consistent MTA advantage over baselines. The Y-axis reports "Perplexity gain, %" relative to a baseline (the zero line), with the Transformer baseline progressively improving (more negative gain) as scale increases. MTA's gain is approximately −1.5% at 300M, improving to roughly −2.5% at 550M, and reaching approximately −3.0% at 1B parameters. Compared to the next-best baseline (Talking-heads, roughly −1.5% at 1B), MTA's advantage grows with scale. The DIFF Transformer shows the smallest gains (approximately −0.5% at 300M, reaching roughly −1.0% at 1B). The monotonic improvement in MTA's relative gain with model size suggests the multi-token bottleneck becomes more severe (not less) as models scale — larger models have more capacity to encode multi-token information into single vectors, but evidently still benefit substantially from dedicated multi-token attention, possibly because larger models are also more capable of learning the complex multi-token patterns that MTA's convolutions detect.

Continued Training with MTA Insertion (Appendix I, Table 10)

On 1.4B proprietary models, continued training a standard Transformer for 10.5B tokens yields 10.69 average perplexity, while continued training the same model with MTA inserted yields 10.61 — an 0.08 improvement from MTA insertion alone. Training MTA from scratch for the full duration achieves 10.44, showing that continued training recovers roughly two-thirds of the from-scratch MTA benefit.

On open-source Llama 3 models continued-trained for 5.3B tokens:

Llama 3.2 1B: MTA reduces perplexity from 10.53 (standard continued training) to 10.46 (−0.07)
Llama 3.2 3B: MTA reduces perplexity from 8.93 to 8.89 (−0.04)
Llama 3.1 8B: MTA reduces perplexity from 8.53 to 8.48 (−0.05)

The consistent improvement across three model scales (1B, 3B, 8B) from two different model families demonstrates that MTA's benefits are not specific to the authors' training setup or architecture — the identity-initialized convolution insertion generalizes to independently trained models with different training recipes, tokenizers, and data distributions. The improvements are modest (0.04–0.07 perplexity) but achieved with only 5.3B tokens of continued training — a tiny fraction of these models' original pretraining budgets (Llama 3.1 8B was trained on 15T+ tokens). This suggests MTA can be practically adopted without full retraining, though the small continued-training budget means these gains may be a lower bound on what longer continued training could achieve.

Ablation Studies and Robustness Checks

Number of key-query convolution layers (Figure 5, right): Increasing the number of layers with key-query convolution from 0 to 12 (in a 24-layer model) monotonically improves average validation perplexity, with 2 layers (roughly 11.225 PPL) already outperforming both the Transformer (11.25) and DIFF Transformer (11.225) baselines. Six layers achieves approximately 11.175 PPL — the paper identifies this as striking "a balance between performance and additional complexity," though the curve continues downward to 12 layers, suggesting further gains are available at the cost of additional computation. Head convolution is applied to all layers regardless of this setting.
Head kernel size (Figure 6, left): Increasing head kernel size from 1 (no head mixing) to 16 monotonically improves perplexity from roughly 11.225 to approximately 11.05. The improvement is near-linear from kernel size 1 to 8, with diminishing returns from 8 to 16. The paper uses ch=16 as the default, which is the maximum tested — this matches the total head count (16 heads in the 880M model), meaning all heads are in a single mixing group at this size. Smaller kernel sizes (grouped interaction) are closer to Talking-heads attention's global projection but with a locality constraint.
Kernel initialization (Section 4.6): Identity initialization (kernel weights initialized so MTA initially computes standard attention) leads to best convergence and final performance. Zero initialization reduces average validation perplexity by 0.02 compared to identity, while constant initialization at 0.3 reduces it by 0.08. The paper interprets this as evidence that starting from the known-to-work standard attention mechanism and gradually learning multi-token patterns is more effective than learning attention from a perturbed starting point.
Pre-softmax vs. post-softmax convolution ordering (Table 5, top and bottom rows): The fully-enabled MTA configuration (key-query pre-softmax, head post-softmax, scalar gating) achieves 10.90 PPL, compared to 10.95 when both are pre-softmax and 10.95 when both are post-softmax. All configurations with group normalization outperform those without (11.03 for key-query pre-softmax, head post-softmax, no scaling — a 0.13 PPL degradation). Changing only the key-query convolution from pre-softmax to post-softmax (with head convolution post-softmax and depth scaling) increases perplexity from 10.92 to 11.11 (+0.19). Changing only head convolution from post-softmax to pre-softmax (with key-query pre-softmax and depth scaling) increases from 10.92 to 11.10 (+0.18). These are small but consistent gaps, suggesting the mixed configuration is genuinely optimal rather than a result of noise.
Normalization and gating strategy (Table 5): Sigmoid gating achieves 10.90 PPL, depth-dependent scaling (DIFF Transformer style) achieves 10.92, no scaling achieves 11.03, and layer-norm scaling (Sun et al., 2025) achieves 11.41 — substantially worse. The 0.51 PPL degradation from using layer-norm scaling rather than group normalization with sigmoid gating is the largest single ablation effect in the table, indicating that the normalization strategy is crucial and not all normalization approaches are interchangeable.
Kernel size variation (Table 5, middle rows): With key-query convolution only (cq=4, ck=9), perplexity is 11.23; with (cq=6, ck=11), 11.23; with (cq=8, ck=13), 11.31. The paper states these variants "display similar kernel patterns, while resulting in slightly different evaluation results" — the degradation at larger kernel sizes (11.31 vs. 11.23) suggests that larger receptive fields may overfit or add noise rather than capturing useful longer-range patterns. However, this comparison is confounded because these ablations were run without head convolution (indicated by the "×" markers in the head conv columns), making them weaker baselines than the fully-configured MTA.
Head convolution vs. higher-rank attention (Appendix B): The paper proves that post-softmax head convolution with ch=2 can be expressed as standard attention with twice the per-head rank (value and output projections of dimension 2d instead of d). However, the equivalence is not exact because "the parameters of the two heads are not independent" — the weight sharing imposed by the convolution structure makes it more parameter-efficient than simply doubling head dimensions. This provides theoretical grounding for why head convolution works, though no ablation directly compares head convolution to simply increasing head dimension at matched parameter count.
Negative result: QA3 on BabiLong (Figure 7): All models — including MTA — perform near random (16.67%) on the three-supporting-facts task with 4K distraction. MTA achieves roughly 15% versus the Transformer's roughly 12%, but none of the models meaningfully exceed random performance. This demonstrates a clear boundary: multi-token attention helps when the model can locate individual facts (as in QA1–QA2, where MTA shows large improvements), but when the underlying fact retrieval itself fails (as in QA3, where even single-fact retrieval may be unreliable given the three-fact query complexity), no amount of attention-map combination can recover the necessary information.

Critical Assessment

Claim 1: MTA overcomes the single-token bottleneck and improves language modeling.

Supported. The validation perplexity improvements in Table 2 (10.91 vs. 11.25, a 0.34-point gain or roughly 3% relative improvement) are consistent across all seven SlimPajama subdomains and two training runs, making this the most robust result in the paper. The improvement is modest in absolute terms but meaningful given that MTA adds essentially zero parameters and is applied to only 6 of 24 layers (key-query convolution). The downstream benchmark average improvement of 1.2 points (Table 3) is encouraging but less decisive — the individual benchmark standard deviations are not reported for most methods, making it impossible to assess statistical significance. A fairer characterization would be "consistent small improvements across most benchmarks, with some variation."

Weakness: The paper does not report ablation on whether key-query convolution alone (without head convolution) achieves most of the perplexity gain, or whether head convolution alone provides comparable benefit. The ablation in Table 5 partially addresses this (rows with key-query convolution only achieve 11.13–11.23 depending on normalization, compared to 10.90 for the full model), but head convolution without key-query convolution is tested only with depth scaling (11.11 PPL, fourth row from bottom) — relegating this important comparison to a single non-optimal normalization setting. A systematic 2×2 ablation (key-query conv on/off × head conv on/off) would clarify how much each component contributes to the overall gain.

Claim 2: MTA is especially beneficial for long-context retrieval tasks.

Strongly supported, with one important failure case. The Needle-In-A-Haystack results (Figure 3) show MTA maintaining high accuracy at difficulty levels where the Transformer collapses (6–10 needles, 4K context). The LAMBADA results (Table 4) show the largest relative improvement in the paper (4.4 perplexity points on LAMBADA standard). The BabiLong results (Figure 5 left) show consistent MTA advantage at 4K distraction across QA1–QA5. The convergence of evidence from three different long-context tasks, each measuring different aspects of retrieval capability, makes this the paper's strongest empirical contribution.

However, QA3 on BabiLong is a clear failure: with three supporting facts and 4K distraction, MTA's performance (roughly 15%) is not meaningfully different from random (16.67%) or from the Transformer baseline (roughly 12%). This bounds the claim: MTA helps when individual fact retrieval is partially successful (so multi-token combination can boost the correct answer above the noise floor), but when retrieval itself fails, no amount of attention combination can recover.

Missing experiment: The paper does not evaluate on the original single-needle Needle-In-A-Haystack task (the standard version with one needle at varying depths), which would help disentangle whether MTA's benefit comes from (a) better single-fact retrieval at depth, or (b) better multi-fact combination from multiple needles. The existing multi-needle results confound these two explanations. A single-needle baseline would also connect to the well-documented "lost in the middle" problem (Liu et al., 2024), which the paper cites as motivation.

Claim 3: The single-token bottleneck is the mechanism behind MTA's improvements.

Supported mechanistically by the toy task, inferentially by the long-context results. The toy task (Table 1) is the cleanest evidence: it directly constructs a problem requiring multi-condition attention, shows standard attention fails (51.6% error with high variance), and shows MTA solves it near-perfectly (0.1% error). This is an existence proof that the bottleneck is real and MTA addresses it.

However, the bridge to the large-scale experiments is inferential. The paper does not demonstrate that the specific kernel patterns learned during large-scale training (Figures 9–14) directly correspond to the multi-token conjunction operations hypothesized in the toy task. The diagonal kernel in Figure 4 is shown to amplify attention for a sequence match, which is suggestive, but this is a single hand-picked example from one layer — there is no systematic analysis of what fraction of kernels implement conjunction operations versus other patterns (priming, edge detection, etc.), or whether heads with conjunction-like kernels disproportionately contribute to performance on multi-fact retrieval tasks. Causal interventions (e.g., ablating specific kernel positions and measuring impact on Needle-In-A-Haystack performance) would strengthen the mechanistic claim substantially.

Claim 4: MTA outperforms existing attention modifications (DIFF Transformer, Talking-heads).

Supported, but the comparisons are not perfectly matched. Table 2 shows MTA consistently outperforms both baselines (10.91 vs. 11.04 for Talking-heads, 11.15 for DIFF Transformer). Table 3 shows MTA with 44.9 average benchmark score versus 44.4 for Talking-heads and 43.9 for DIFF Transformer. However, MTA has a slight parameter advantage: MTA's additional parameters (roughly 29,600) exceed DIFF Transformer's (roughly 13,800) and Talking-heads' (roughly 12,300), as shown in Table 8. While all additions are trivially small relative to 876M total parameters, the relative difference between modification methods is a factor of roughly 2× — MTA has about twice as many added parameters as the baselines. The paper does not discuss whether even this small parameter difference could account for the 0.13 PPL gap between MTA and Talking-heads, or whether an alternative configuration of Talking-heads (e.g., with group normalization added) would close the gap. The DIFF Transformer, in particular, is evaluated with half as many attention heads to match total parameters (Section 4.2), which may disadvantage it on tasks that benefit from more heads. A fairer comparison would match both total parameters AND total heads, perhaps by slightly adjusting hidden dimension.

Claim 5: MTA can be retrofitted into pretrained models via identity initialization.

Supported by preliminary experiments, but the evidence is limited in scale. The continued-training results (Appendix I, Table 10) consistently show 0.04–0.08 PPL improvements from adding MTA to already-pretrained models across multiple model families and scales. This is practically important and the results are uniformly positive. However, the experiments are preliminary in several key respects:

Token budget: The continued training uses only 5.3–10.5B tokens — modest relative to the 105B pretraining budget of the 880M experiments. It is possible that with longer continued training, the MTA kernels would learn richer patterns and achieve larger gains, or conversely, that the initial improvements are transient and would be matched by continued standard training given enough tokens.
Scale: The largest model tested is Llama 3.1 8B. Whether the approach scales to 70B+ models (where continued training is proportionally more expensive and catastrophic forgetting is a greater risk) is unknown.
Task evaluation: The continued-training experiments report only perplexity — they do not evaluate on downstream benchmarks or long-context tasks. The perplexity improvement might not translate to task performance, particularly if the continued training's small token budget leads to fragile kernel patterns that don't generalize.
Data distribution: The continued training uses SlimPajama (for the 1.4B model) and presumably SlimPajama or similar for the Llama models. The effectiveness may depend on the continued-training data matching the original pretraining distribution, which for Llama models is proprietary and may not match SlimPajama well.

Overall: What the experiments establish and what they leave open.

Established:

MTA provides consistent, small-to-moderate improvements in language modeling perplexity at negligible parameter cost.
MTA substantially improves performance on long-context retrieval tasks requiring multi-fact location, with gains that grow with context length and number of facts.
Convolution over attention logits is a viable mechanism that can be added to attention without destabilizing training, especially when initialized to identity.
The benefits scale with model size (Figure 6 right), suggesting MTA addresses a bottleneck that is not resolved by simply making models larger.

Left open:

Computational cost at deployment: The paper's implementation is 9.5× slower than standard attention during training (Table 9). With optimized CUDA kernels, the theoretical overhead of small convolutions over attention maps should be much smaller, but no optimized implementation is provided or benchmarked. The practical deployability of MTA at scale depends critically on closing this gap.
Single model family, single dataset: All pretraining experiments use LLaMA-style Transformers on SlimPajama. The continued-training results on Llama models partially address model-family generalization, but no experiments use non-Transformer architectures or substantially different pretraining data (e.g., code-heavy mixtures, multilingual corpora).
No combination with other attention improvements: MTA is compared against DIFF Transformer and Talking-heads as alternatives, but the combination of MTA with these methods is not tested. Since MTA's convolution operations are orthogonal to differential attention (subtraction) and global head projection, a model incorporating all three might outperform any individual method.
Causal evidence for the mechanism is limited to the toy task: At scale, the connection between learned kernel patterns, multi-token retrieval behavior, and task performance is correlational (kernel visualizations are interpreted post-hoc) rather than causally tested.
The threshold where MTA stops helping is unclear: The QA3 failure suggests MTA's benefit requires the base model to have non-trivial single-fact retrieval capability. Characterizing this threshold more precisely — in terms of base model pass rates, fact retrieval accuracy, or some other metric — would help practitioners decide when MTA is worth adopting.

6. Limitations and Trade-offs

The Computational Overhead in Practice Is Unknown and Currently Prohibitive

The paper reports that the unoptimized MTA implementation achieves only 5.7K tokens per second during training versus 54.3K for the standard Transformer using PyTorch's optimized scaled_dot_product_attention — a roughly 9.5× slowdown (Table 9). This is not a theoretical limitation but a pragmatic one: the paper does not provide an optimized implementation, making it impossible to assess what performance is achievable with proper CUDA kernels.

The authors are transparent about this:

"our MTA implementation does not take advantage of such efficient kernels, which is the major reason behind its lower FLOPS."

However, transparency does not resolve the uncertainty. The key-query convolution operates on the full T × T attention map at each layer where it is applied (6 of 24 layers), adding O(T² · c_q c_k) operations on top of the existing O(T² d) attention computation. For the default kernel sizes (c_q = 6, c_k = 11), this is roughly 66 multiply-adds per attention logit — modest compared to the d = 96 dot product operations per attention logit in the standard attention computation. The head convolution is even cheaper (a small linear transformation across heads). So the theoretical overhead should be manageable — perhaps 1.5–2× rather than 9.5×.

But this remains speculation without optimized kernels. The gap between theory and practice here is critical for deployment decisions: a 2× slowdown might be acceptable given MTA's perplexity and retrieval gains; a 9.5× slowdown almost certainly is not, especially for latency-sensitive applications. The paper acknowledges this as future work ("Optimizing the runtime performance was not the goal of this work, thus we leave further optimization of the MTA implementation for future research" — Appendix A), but the absence of even preliminary optimized-kernel results means the paper cannot address the most basic deployment question: what does MTA actually cost?

This limitation is partially mitigated by the fact that key-query convolution is applied to only 1/4 of layers, which limits the overhead. The ablation in Figure 5 (right) showing that even 2 layers with MTA outperform strong baselines suggests that further reducing the number of MTA layers could trade off performance for speed. But without optimized-kernel measurements at any layer count, these remain qualitative tradeoffs rather than actionable engineering decisions.

All Large-Scale Results Are on a Single Model Architecture and Dataset

Every pretraining experiment in the paper uses LLaMA-style decoder-only Transformers trained on SlimPajama (Soboleva et al., 2023). The benchmark evaluations (Table 3) use standard academic datasets, but the base models are all trained within this single architecture-data regime. The continued-training experiments in Appendix I extend to Llama 3 models (1B, 3B, 8B) with a different pretraining recipe, but these are evaluated only on perplexity after modest continued training — not on downstream tasks or long-context retrieval.

This matters for two reasons. First, the effectiveness of learned convolution kernels over attention logits may depend on properties of the attention distribution that vary across model architectures, training objectives, or data distributions. Second, SlimPajama is an English-dominated web text corpus. The multi-token patterns that MTA learns to detect — diagonal sequence matches, co-occurrence peaks, priming patterns — may be specific to the statistical structure of English web text. Languages with different word order or morphological structure might require different kernel sizes or produce different learned patterns.

The paper's demonstration that MTA insertion into Llama 3 models improves perplexity (0.04–0.07 points across all three scales, Table 10) partially addresses the architecture-generalization concern, but with two limitations: (1) the continued-training budgets are small (5.3B tokens), so these results demonstrate that MTA doesn't break existing models more than they demonstrate that MTA provides robust gains; and (2) no long-context retrieval evaluation is performed on the continued-trained Llama models, which is where MTA's largest gains appear (Figure 3, Figure 5 left). Without this, we cannot assess whether the perplexity improvements from continued training translate to the retrieval benefits that motivate MTA.

The paper does not acknowledge this as a limitation — the authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (implied in Section 4), but this is a claim about the base Transformer architecture, not about MTA's portability. The evidence for portability is preliminary.

MTA Does Not Help on the Hardest Multi-Fact Reasoning Tasks

The BabiLong QA3 task (three supporting facts required to answer a question) reveals a hard boundary on MTA's effectiveness. With 4K distraction text, all models — including MTA — perform near random: MTA achieves roughly 15% accuracy versus 16.67% random chance (Appendix Figure 7, QA3 panel). The standard Transformer achieves roughly 12%, meaning MTA's improvement over the baseline is approximately 3 percentage points, and neither model meaningfully exceeds the random baseline.

This is not an isolated failure. The difficulty-bin analysis from the paper's long-context experiments (implicit in the per-task BabiLong results, Figure 7) shows a consistent pattern: MTA's advantage over baselines is largest on tasks requiring one or two supporting facts (QA1, QA2) or simpler argument relations (QA4, QA5), and collapses when the underlying fact retrieval becomes too difficult. On QA3, the model cannot reliably locate individual facts from among 4K tokens of distraction, so the multi-token combination that MTA enables has nothing to combine — there are no reliable "Alice found here" and "rabbit found there" signals to intersect.

The consequence is a ceiling effect: MTA amplifies existing retrieval capability but does not create it. If the base model's probability of correctly attending to individual relevant facts is near zero, MTA's convolutions operate on noise and cannot recover the correct answer. This mirrors a pattern from the reference example paper (compute-optimal test-time scaling), where the hardest difficulty bin showed near-zero improvement regardless of method — test-time interventions help when the base model has non-trivial capability and fail when it does not.

The paper acknowledges this implicitly by reporting the QA3 results without attempting to explain or mitigate them, but does not discuss the implication: MTA helps precisely when standard attention has partial success, and fails when standard attention fails completely. Practitioners need to know whether their target tasks fall above or below this threshold. The paper provides no diagnostic for predicting when MTA will help versus when it will not, beyond the coarse observation that "more facts + more distraction = harder, and MTA helps until the task becomes impossible."

The mitigation is absent — the paper does not propose combining MTA with other mechanisms (retrieval augmentation, larger models, task decomposition) that might push the capability threshold upward for the hardest cases.

Identity Initialization for Continued Training Is Promising but Unvalidated at Scale

The continued-training experiments in Appendix I (Table 10) show that MTA can be inserted into pretrained models with identity-initialized kernels and fine-tuned to achieve perplexity improvements. This is a practically important result because it suggests MTA adoption does not require training models from scratch.

However, the evidence has specific limitations that constrain its applicability to real deployment scenarios:

Token budgets are tiny relative to pretraining. The continued training uses 5.3–10.5B tokens, which is 0.04–0.07% of Llama 3.1 8B's pretraining budget (15T+ tokens). The perplexity improvements (0.04–0.07 points) are positive but modest — they demonstrate that MTA doesn't catastrophically degrade performance, but don't establish that continued training with MTA would approach the from-scratch MTA performance if given a realistic budget.

No downstream task evaluation. The continued-training experiments report only validation perplexity. The paper's main results show that MTA's largest gains are on long-context retrieval tasks (Figures 3 and 5), not on perplexity alone. Without evaluating the continued-trained models on Needle-in-Haystack, BabiLong, or LAMBADA, we cannot assess whether the perplexity gain translates to the capabilities that motivate MTA adoption.

No study of forgetting or capability degradation. The paper does not evaluate whether continued training with MTA degrades performance on tasks the original model was good at — a standard concern in continued pretraining. The identity initialization ensures the model starts with unchanged outputs, but as the kernels deviate from identity during training, the attention patterns change, potentially disrupting capabilities the original model had developed through its much longer pretraining.

Scale is limited to 8B parameters. The largest model tested is Llama 3.1 8B. Whether the approach scales to 70B or 405B models — where continued training is proportionally more expensive, attention patterns are more specialized, and the risk of disrupting delicate capabilities is higher — is unknown.

The paper does not claim the continued-training results are definitive — they appear in an appendix and are described as "preliminary experiments" — but the prominence of identity initialization as a design choice (Section 4.6) and the practical importance of the retrofit capability mean that the gap between preliminary evidence and deployment readiness should be explicitly acknowledged.

The Mechanism Claims Are Correlational at Scale, Not Causally Validated

The paper's central mechanistic argument is that MTA helps because it overcomes the single-token bottleneck — enabling attention to condition on multiple query-key pairs rather than single vectors. This argument is validated causally in the toy task (Table 1), where the problem is deliberately constructed so that single-token attention cannot succeed, and MTA's convolution kernels are sized to exactly match the required multi-token interaction (c_q = 2 for L = 2 query letters).

At scale, the evidence for this mechanism is correlational and post-hoc. The paper shows that MTA improves performance on tasks that plausibly require multi-token attention (Needle-in-Haystack with multiple needles, BabiLong with multiple supporting facts) and visualizes learned kernel patterns that are interpretable as multi-token detectors (diagonal sequence matchers in Figure 4, priming patterns, edge detectors). But no experiment at scale demonstrates a causal link between specific kernel patterns and specific task improvements.

Key missing evidence includes:

Ablation of specific kernel positions: If the diagonal kernel in Figure 4 is responsible for sequence matching in Needle-in-Haystack, zeroing out its off-diagonal weights should degrade multi-needle retrieval while leaving single-needle retrieval intact. This experiment is not performed.
Correlation between kernel patterns and task improvements across heads/layers: The paper shows that kernels are diverse (Figures 9-14), but doesn't analyze whether heads with conjunction-like kernels contribute more to long-context task performance than heads with identity-like kernels.
Comparison to a "single-token MTA" baseline: An ablation where c_q = 1 and c_k = 1 (essentially standard attention but with head mixing and group normalization) would isolate how much of MTA's gain comes from multi-token interaction versus from head mixing and normalization alone. The head-kernel-size ablation (Figure 6 left) addresses the head mixing axis but not the key-query axis — the smallest query-key kernel tested is c_q = 4, c_k = 9 (Table 5), not c_q = 1, c_k = 1.

The consequence is that while the paper provides a compelling existence proof (the toy task) and suggestive evidence (kernel visualizations, long-context improvements), it does not provide a causal decomposition of MTA's benefits at scale. We cannot say with confidence that the multi-token interaction is the primary driver of improvement — as opposed to the regularization effect of the convolution, the increased representational capacity from head mixing, or beneficial optimization properties of the group normalization and gating. Multiple mechanisms likely contribute, but the paper's narrative attributes the gains primarily to multi-token attention without isolating that factor.

The mitigation is partial. The ablation in Table 5 shows that both key-query convolution and head convolution contribute (removing either degrades perplexity), but this shows they are both necessary for full performance, not that key-query convolution specifically helps through multi-token interaction rather than through some other mechanism.

No Evaluation on RULER-Style Aggregated Long-Context Benchmarks

The paper evaluates long-context performance on LAMBADA, Needle-in-Haystack, and BabiLong — three well-established tasks that each probe different aspects of long-range retrieval. However, these tasks have known limitations: Needle-in-Haystack is synthetic and tests pure retrieval without reasoning, BabiLong uses templated language that may not reflect natural text distributions, and LAMBADA is a next-word prediction task whose long-range dependency requirements have been debated.

The paper does not evaluate on more comprehensive long-context benchmarks such as RULER (Hsieh et al., 2024), which aggregates needle-in-haystack variants (multi-needle, multi-value, multi-query), variable tracking, and long-document QA into a unified score. Nor does it evaluate on LongBench (Bai et al., 2023) or ∞BENCH (Zhang et al., 2024), which test a broader range of long-context capabilities including summarization, multi-document QA, and code understanding. This matters because the paper's core claim — that MTA is "particularly beneficial" for context requiring precise location of relevant information — would be strengthened by demonstrating gains on a diverse set of realistic long-context tasks, not only on templated retrieval (BabiLong) and synthetic insertion (Needle-in-Haystack).

The BabiLong tasks, while informative, use artificially constructed distractors — irrelevant facts interspersed with relevant ones in a simple declarative sentence format. It is unclear whether MTA's improvements on this structured distraction transfer to naturalistic long documents where distractors are semantically coherent and thematically related to the target information. The Needle-in-Haystack task suffers from a similar artificiality: the "needle" (a sentence like "The magic number of San Francisco is 8") is semantically disjoint from the "haystack" (irrelevant text), making the retrieval problem primarily about locating a distinctive signal rather than distinguishing relevant from thematically similar irrelevant information.

This limitation is not a flaw in the experimental design — the chosen benchmarks are appropriate for an initial investigation — but it represents a gap between the paper's evidence and the strength of its claims about long-context benefits. The paper does not acknowledge this gap or discuss how the results might generalize to more naturalistic long-context tasks.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a new diagnostic category for attention limitations — the single-token bottleneck — and provides both a mechanistic demonstration of its reality (the toy task) and a practical architectural intervention that addresses it (multi-dimensional convolution over attention logits). The contribution is best understood not as a paradigm shift but as a reframing with architectural consequences: it changes what the field thinks attention can't do and provides a template for fixing it.

The reframing: from "attention distributes weight poorly" to "attention conditions on impoverished information." Prior work on attention limitations has focused almost exclusively on the distribution of attention weights — they're too diffuse (softmax temperature), too noisy (differential attention for noise cancellation), too evenly spread across heads (Talking-heads for head interaction). MTA's diagnostic reframing identifies a prior and more fundamental issue: before you can worry about how attention weights are distributed, you must ask whether the information available to compute those weights is sufficient. By showing that standard attention literally cannot solve a task requiring simultaneous attention to two letters without heroic representational compression (the toy task's 51.6% error rate with high variance, Table 1), the paper establishes that the single-vector conditioning is not just a theoretical limitation but a practical failure mode. This shifts attention research upstream — from post-hoc fixes applied to already-computed attention weights to interventions at the point where attention logits are formed.

The architectural template: attention logits as a processable intermediate representation. The paper's most transferable contribution may be methodological rather than algorithmic. By treating the attention logit matrix Â ∈ R^{T×T} as a feature map to which learned spatial convolutions can be applied, MTA opens a design space that the paper only partially explores. The specific mechanism — small 2D/3D convolutions over query, key, and head dimensions — is one instantiation of a broader idea: that the similarity scores between all query-key pairs contain structure (spatial patterns, temporal dependencies, cross-head correlations) that learned local operations can exploit, and that the softmax should operate on the processed similarity map rather than raw dot products. This is analogous to the transition in computer vision from applying classifiers directly to pixel values to applying them to convolutionally-processed feature maps. The paper demonstrates that even small kernels (6×11 query-key, groups of 16 heads) with identity initialization learn interpretable, specialized patterns (diagonal sequence matchers, priming detectors, edge detectors — Figures 4, 9-14), suggesting the design space is rich.

Resolution of a latent tension in the attention literature. The paper implicitly reconciles two seemingly contradictory findings about attention modification. On one hand, methods that modify attention after softmax (Talking-heads, DIFF Transformer) provide modest but consistent gains — they improve the utilization of computed attention patterns. On the other hand, the persistent difficulty of long-context retrieval (Kamradt, 2023; Liu et al., 2024) suggests these post-hoc fixes are not enough. MTA's results explain this: post-softmax modifications can't create information that wasn't present in the logits to begin with. When the distinguishing signal requires multi-token evidence (Alice AND rabbit), a post-softmax operation sees two separate attention maps — one for Alice, one for rabbit — and can at best combine them linearly (as Talking-heads does) or subtract them (as DIFF Transformer does). But it cannot multiply them to implement the AND operation that would identify positions where both are high. Pre-softmax key-query convolution enables exactly this multiplicative gating because it operates in log-space before the exponentiation. This explains why MTA substantially outperforms both Talking-heads and DIFF Transformer on multi-needle retrieval (roughly 70% vs. ~35% for the Transformer at 6 needles and 4K context, Figure 3 right) while providing more modest gains on standard benchmarks (Table 3). The pre-softmax intervention is qualitatively different from post-softmax modifications — it expands what information can inform attention, not just how computed attention is used.

Which research directions become more attractive. The paper makes several lines of investigation newly promising:

Learned processing of attention logits as an intermediate representation, including multi-scale convolution, attention over logit neighborhoods, and dynamic kernel generation conditioned on input content.
Mechanistic interpretability of attention patterns — the learned kernels provide a new window into what patterns different heads are detecting, and causal interventions on kernel weights could establish functional roles for specific attention computations.
Architectural modifications that can be inserted into pretrained models via identity initialization — the success of this strategy with MTA (Appendix I, Table 10) suggests a general approach for retrofitting capabilities into deployed models without full retraining.
Difficulty-aware application of attention modifications — the finding that MTA helps most when the base model has partial retrieval success (BabiLong QA1–QA2 but not QA3) suggests that routing mechanisms could apply MTA selectively to queries that benefit from multi-token attention.

Which directions become less attractive. The paper's results suggest diminished returns for:

Purely post-softmax attention modifications that don't address the information bottleneck in logit formation — Talking-heads and DIFF Transformer show consistent but ceiling-limited improvements compared to MTA's pre-softmax intervention.
Methods that focus solely on sharpening attention (sparsemax, adaptive temperature) without expanding the conditioning information — sharper attention to single-token matches doesn't solve the conjunction problem.
"Just make the model bigger" as a solution to attention limitations — the scaling laws (Figure 6 right) show MTA's relative gain increases with model size (roughly 1.5% perplexity gain at 300M vs. 3.0% at 1B), suggesting the single-token bottleneck is not resolved by additional capacity — larger models benefit more from multi-token attention, not less.

Follow-Up Research This Work Enables

1. Optimized CUDA kernels and rigorous overhead measurement. The most urgent follow-up is an optimized implementation that closes the 9.5× training speed gap between MTA and standard attention (Table 9). The theoretical overhead is modest — a 6×11 convolution over the T×T attention map adds roughly 66 multiply-adds per logit, compared to d=96 multiply-adds for the QK dot product — suggesting 1.5–2× slowdown is achievable with fused kernels that combine the QK computation with the subsequent convolution and avoid materializing the full attention matrix. A strong follow-up would implement MTA using Triton or custom CUDA, benchmark training throughput and memory at model scales from 1B to 70B, measure inference latency at realistic batch sizes, and compare wall-clock-time-matched perplexity (training the standard Transformer for longer to match MTA's wall-clock time) rather than token-matched perplexity. This would determine whether MTA's gains survive a fair compute-matched comparison — the question the current paper cannot answer.

2. Causal decomposition of MTA's benefits: which components matter for which tasks? The paper's ablation studies (Table 5, Figures 5-6) vary components individually but never perform a full factorial decomposition that isolates the interaction between key-query convolution, head convolution, and group normalization. A systematic follow-up would train matched 880M models with all 2×2×2 combinations of {key-query conv on/off, head conv on/off, group norm on/off} and evaluate on both standard benchmarks and long-context tasks. The crucial measurement is the interaction term: does key-query convolution provide additional benefit beyond head convolution alone, and is this additional benefit specifically concentrated on tasks requiring multi-token conjunction (Needle-in-Haystack with multiple needles, BabiLong QA2) rather than on tasks solvable by single-token matching? A negative result — finding that head convolution + group norm accounts for most of MTA's gain, with key-query convolution contributing minimally — would substantially revise the paper's mechanistic narrative.

3. Kernel ablation experiments to establish causal roles for learned attention patterns. The kernel visualizations (Figures 4, 9-14) are suggestive but purely correlational. A causal follow-up would intervene on specific kernel weights in trained MTA models and measure the impact on task performance. For example: (a) For each head's key-query kernel, zero out all off-diagonal weights (leaving only the center θ_{0,0}), which reduces that head to standard attention. Measure the degradation on multi-needle retrieval at different needle counts. The prediction from the paper's mechanistic story is that off-diagonal ablation should hurt multi-needle retrieval more than single-needle retrieval, and should disproportionately affect heads whose kernels show strong off-diagonal patterns (the diagonal sequence matchers in Figure 4) rather than heads with near-identity kernels. (b) For head convolution kernels, zero out cross-head weights (w_{ij} for i ≠ j), reducing head mixing to per-head scaling. Measure the degradation on tasks requiring noise cancellation (Needle-in-Haystack with distracting similar needles) versus tasks requiring amplification. (c) Compare the importance of different kernel positions — does θ_{2,0} (query offset 2, key center) matter more for tasks where query tokens are typically separated by 2 positions? This would connect the learned kernel structure to functional role with causal rather than correlational evidence.

4. Scaling MTA to larger models and evaluating on comprehensive long-context benchmarks. The current experiments are limited to ≤1B parameters (880M for full evaluation, 1.4B for continued training only) and evaluate long-context performance on three tasks (LAMBADA, Needle-in-Haystack, BabiLong) that, while informative, don't capture the diversity of real-world long-context scenarios. A strong follow-up would: (a) Train MTA models at 7B and 13B scale on the same SlimPajama data, measuring whether the scaling trend in Figure 6 (right) continues — does MTA's relative gain keep growing with model size, or does it saturate? (b) Evaluate on RULER (Hsieh et al., 2024) to test aggregated long-context capabilities across needle-in-haystack variants, variable tracking, and QA, and on LongBench (Bai et al., 2023) for multi-document QA, summarization, and code understanding. The paper's claim that MTA is "particularly beneficial" for long contexts would be substantially strengthened by gains on these naturalistic benchmarks, or substantially weakened if gains are limited to synthetic retrieval. (c) Measure performance on the "lost in the middle" phenomenon (Liu et al., 2024) — does MTA's multi-token attention reduce the U-shaped performance curve where models attend best to the beginning and end of contexts? The paper's depth-dependent Needle-in-Haystack results (Appendix G, Figure 8), which show MTA better at finding needles hidden deep in context, suggest this is likely but are reported only for pretrained models at 2K context, not the finetuned 4K models where the effect would be most relevant.

5. Combining MTA with retrieval-augmented generation and other attention improvements. MTA is evaluated in isolation against standard attention baselines, but its mechanism is orthogonal to several other performance-enhancing techniques. A natural follow-up would test MTA in combination with: (a) Retrieval-augmented generation (RAG), where the model must attend to retrieved chunks intermixed with the query — does MTA's multi-token attention help distinguish relevant from irrelevant retrieved passages when the distinction requires multi-fact matching? (b) DIFF Transformer's differential attention — since MTA's head convolution often learns subtractive patterns (Appendix Figure 15), adding explicit differential attention on top of MTA might provide complementary noise cancellation. (c) Sliding window or sparse attention, where MTA's key-query convolution would operate within local windows — does the convolution benefit survive when the attention map is sparse? The hypothesis from MTA's mechanism is that key-query convolution should help even within local windows (since multi-token conjunctions are typically local), but this needs empirical verification.

6. Failure mode characterization: when does MTA's benefit vanish, and can this be predicted? The BabiLong QA3 result — where MTA performs near random (15% vs. 16.67% random baseline, Figure 7) — establishes that MTA's benefit has a sharp threshold related to the base model's retrieval capability. A systematic follow-up would characterize this threshold quantitatively. Specifically: (a) For each BabiLong task, measure the base Transformer model's single-fact retrieval accuracy (can it locate individual relevant facts from the context?) and correlate this with MTA's improvement over the baseline. The prediction is that MTA's gain is near-zero when single-fact retrieval accuracy is near chance, grows as single-fact retrieval improves, and potentially saturates when single-fact retrieval is already near-perfect. (b) Construct synthetic multi-needle tasks with parametrically varying needle similarity to distractor text — when needles are highly distinctive (e.g., "The magic number of San Francisco is 8" in a haystack of philosophy text), MTA should help less because single-token matching suffices; when needles are semantically similar to distractors (e.g., numerical facts embedded in a financial report), MTA's multi-token attention should provide larger relative gains. (c) Train a lightweight "MTA-benefit predictor" that takes a query and context as input and predicts whether applying MTA (vs. standard attention) will change the model's answer — this would enable selective application of MTA, reserving its computational cost for queries where it actually helps.

Practical Applications and Downstream Use Cases

1. Long-context enterprise search and document QA. Organizations deploying LLMs for question-answering over long documents (contracts, financial reports, technical manuals, legal discovery) face a core challenge: the model must locate specific information satisfying multiple constraints within tens of thousands of tokens, often amidst thematically similar distractors. MTA's results on BabiLong with 4K distraction text are directly relevant — MTA achieves roughly 52% average accuracy on QA1–5 versus the standard Transformer's 37% (Figure 5 left), a 15-point absolute improvement. For a legal document review system processing 50-page contracts, this translates to finding relevant clauses roughly 40% more often without increasing model size. The continued-training results (Appendix I, Table 10) suggest that existing pretrained models (Llama 3.x series, 1B–8B) can be retrofitted with MTA through modest continued training (5.3B tokens), making this applicable to organizations that have already fine-tuned a base model on their domain and want to improve long-context retrieval without retraining from scratch.

2. Multi-hop reasoning over knowledge bases or research literature. Scientific literature search, competitive intelligence, and medical literature review require answering questions that depend on multiple facts distributed across different sections of a document or different documents. For example, "What drug interaction was reported between compound X and enzyme Y in patients with condition Z?" requires locating mentions of X, Y, and Z and identifying where they co-occur with interaction language. The BabiLong QA2 (two supporting facts) and QA5 (three argument relations) results are the closest proxy — MTA achieves roughly 25% on QA2 at 4K distraction vs. 12% for the standard Transformer, and roughly 55% on QA5 vs. 38%. These relative improvements (roughly 2× and 1.4×) on tasks explicitly requiring multi-fact integration suggest MTA-augmented models would reduce the failure rate on multi-hop queries by a practically meaningful margin. The key-query convolution's ability to detect diagonal patterns (Figure 4 — amplifying attention when query token sequences match key token sequences) is particularly relevant for tracking chains of reasoning across document sections.

3. On-device deployment where retrieval precision is critical but model size is constrained. For applications running on edge devices (smartphones, embedded systems, offline document readers), model size is severely constrained (typically ≤3B parameters), but users still expect accurate retrieval from long contexts (e.g., searching a downloaded ebook, querying cached emails). The paper's scaling results (Figure 6 right) show that MTA's relative gain increases as model size decreases — the 300M MTA model achieves a larger perplexity improvement over its Transformer baseline than the 1B MTA model does over its baseline. This suggests MTA is particularly well-suited to the small-model regime, where the capacity to encode multi-token information into single vectors is most limited. A 1B on-device model with MTA might approach the long-context retrieval accuracy of a 3B standard model at a fraction of the memory and inference cost. The parameter overhead (0.003%, Table 8) is negligible for on-device deployment, but the current computational overhead (9.5× slowdown, Table 9) would need to be addressed through optimized kernels — making the optimized-kernel follow-up (direction 1 above) a prerequisite for this application.

4. Training data generation and filtering pipelines. When LLMs are used to generate or filter training data (e.g., identifying all passages in a corpus that discuss a specific topic with specific constraints, or generating synthetic multi-hop QA pairs), the ability to precisely locate relevant information within long documents directly impacts yield and quality. A data filtering pipeline that must extract all paragraphs discussing "climate policy impacts on agricultural subsidies in Southeast Asia" from web-scale documents requires simultaneous attention to multiple concepts (climate, policy, agriculture, subsidies, Southeast Asia) that may not co-occur in a single sentence. MTA's multi-needle retrieval results (Figure 3) — maintaining roughly 65% accuracy with 10 needles at 2K context versus the standard Transformer's 35% — suggest MTA-augmented models would substantially reduce both false positives (retrieving paragraphs that mention some but not all concepts) and false negatives (missing paragraphs where concepts are distributed across non-adjacent sentences). For large-scale data curation pipelines processing billions of tokens, a 2× improvement in retrieval precision translates to meaningful reductions in downstream manual filtering and quality control costs.

When to Prefer This Method

The paper does not explicitly frame MTA as a choice against named alternatives with clear tradeoff conditions. It presents MTA as a drop-in improvement over standard attention that can be combined with other architectural modifications (it is evaluated against DIFF Transformer and Talking-heads as alternatives, but the paper does not discuss when one would choose MTA over these methods, or whether they can be combined). The continued-training experiments suggest MTA can be added to already-trained models, implying it is compatible with rather than exclusive of other attention modifications. Therefore, a decision matrix is not applicable here — the paper's position is that MTA is a broadly beneficial enhancement to attention, with particularly large gains on tasks requiring multi-token conjunction over long contexts, with the primary practical tradeoff being the current unoptimized computational overhead versus the demonstrated performance improvements.