Multi-Token Attention
Authors: Olga Golovneva
Tianlu Wang
Jason E Weston
Sainbayar Sukhbaatar
Venue: COLM 2025 (2025)
Keywords: Deep learning architectures
Large Language Model (LLM)
Transformer
Attention
TL;DR: We present Multi-Token Attention: new method that allows LLMs to condition their attention weights on multiple query and key vectors simultaneously.

Abstract: Soft attention is a critical mechanism powering LLMs to locate relevant parts within a given context. However, individual attention weights are determined by the similarity of only a single query and key token vector. This “single token attention” bottlenecks the amount of information used in distinguishing a relevant part from the rest of the context. To address this issue, we propose a new attention method, Multi-Token Attention (MTA), which allows LLMs to condition their attention weights on multiple query and key vectors simultaneously. This is achieved by applying convolution operations over queries, keys and heads, allowing nearby queries and keys to affect each other’s attention weights for more precise attention. As a result, our method can locate relevant context using richer, more nuanced information that can exceed a single vector’s capacity. Through extensive evaluations, we demonstrate that MTA achieves enhanced performance on a range of popular benchmarks. Notably, it outperforms Transformer baseline models on standard language modeling tasks, and on tasks that require searching for information within long contexts, where our method’s ability to leverage richer information proves particularly beneficial.

Published as a conference paper at COLM 2025
Multi-Token Attention
Olga Golovneva
Tianlu Wang Jason Weston
Sainbayar Sukhbaatar
FAIR at Meta
Abstract
Soft attention is a critical mechanism powering LLMs to locate relevant parts within given context. However, individual attention weights determined by the similarity of only a single query and key token vector. This "single token attention" bottlenecks the amount of information used in distinguishing relevant part from the rest of the context. To address this issue, we propose new attention method, Multi-Token Attention (MTA), which allows LLMs to condition their attention weights on multiple query and key vectors simultaneously. This achieved by applying convolution operations over queries, keys and heads, allowing nearby queries and keys affect each other's 's attention weights for more precise attention. As a result, our method can locate relevant context using richer, more nuanced information that can exceed a single vector's 's capacity. Through extensive evaluations, we demonstrate that MTA achieves enhanced performance on a range of popular benchmarks. Notably, it outperforms Transformer baseline models on standard language modeling tasks, and on tasks that require searching for information within long contexts, where our method's ss ability to leverage richer information proves particularly beneficial'.
Introduction
The attention mechanism (Bahdanau et al., 2014; Vaswani et al., 2017) is a critical component of Large Language Models (LLMs) that enables them to retrieve and combine information from different parts the context. Attention is especially useful when the context contains large number of tokens, as focusing the relevant part while disregarding distractions becomes crucial. However, numerous works have shown that standard attention can suffer from suboptimal performance in this setting (Kamradt, 2023; Kuratov et al., 2025).
Standard multi-head attention works by comparing the similarity of the current query vector and the key vectors corresponding to the context tokens using their dot products. The keys similar to the query obtain higher attention weights, and subsequently their value vectors dominate the output vector. For example, query vector corresponding to a token "Alice" is capable of locating all mentions of " "Alice" in the context. However, each attention weight conditions only on single key and query vector (besides the normalization sum to one).
We argue that the dependency on single token vector similarity brings fundamental limitation the attention mechanism. In many cases, the relevant part of the context cannot be identified by single token. For example, looking up sentence that mentions both " "Alice" and "rabbit" requires the query vector to encode both tokens. 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. While it is possible to encode multiple tokens into single vector via the layers of the Transformer, this requires increased dimension, and the model to use lots of its capacity for this task.
In this paper, we propose novel attention mechanism that goes beyond this "single token" bottleneck, which we call Multi- Tolti-Token Attention (MTA). The high level goal is to make possible to the similarities of mulliple vector pairs to determine where attention must
Code is available https:/ http:////iiaaoom//cccoooreeeeaae RRAeeeh///MMM///e trre/mmann mmmain/rrojcts/mm/r

---

Published as a conference paper at COLM 2025
Concar
huody
Softmax
Concat
Geled group norm
Mask
Maimul
MTA convolution
Matmul
Scale
MTA convalution
MaiMul
MatMul
K
K
Standard multi-head attention
Multi-token attention
Figure 1: Multi-Token Attention (MTA) (right), compared to standard attention (left). In MTA, within each head we apply key-query convolution the attention scores and head convolution across groups of heads. We repeat this operation after the softmax. Finally, apply group normalization with scalar gating before final concatenation.
focus. We achieve this by making Etraight-forwant modifications to the existing attention mechanism. In particular, we design convolution operations over attention weights that operate on three dimensions: keys, queries, and attention heads. This allows its attention weights to condition following our previous example, MTA can find mentions "Alice" and "rabbit" separately first, and then combine those attentions together to focus only on where both exist.
neighboring keys, previous queries, and other heads. Intuitively,
We first experiment with a motivating toy task that reveals the shortcoming of standard attention and demonstrate that MTA easily solve it. Next, test method at scale by pre-training 880M parameters models on 105B tokens on a standard language modeling task. There we see MTA bring improvements in terms in validation perplexity well standard benchmark tasks, while only increasing the number of parameters by 0.001% %. Further, we evaluate the resulting models on long-context tasks such as Weedleein-the--Hasttae and Babil ong where MTA outperforms the baselines by significant margin. Finally, we investigate scaling laws of the proposed model and compare them with baselines.
Background on multi-head attention
We first describe the standard multi-head attention (Vaswaní al., 2017) and define the notation we use. In decoder-only Transformer architectures, the model receives sequence of tokens Xy, of length as an input. These tokens then undergo a series of trans-
formations into hidden states H = [/4.....p hr]" E R TxD through an embedding layer and repeated transformer layers. Each layer consists of multi-head attention submodule, feedforward network, and normalization operations. Multi-head attention includes M heads with dimensions D/M working in parallel. Each head uses key, value and query projections Wg. W,, Wg E R Dxd to construct key, value and query vectors:
K V = HW, Q = HW,
The attention logits Å and weights A (i.e. probabilities) are then calculated as follows:
Å ==QT//V A = = Softtax(Mask
(1)
where the softmax operates over the key dimension, and the mask function replaces values at (i, with -co for < to prevent information leaking from future tokens. Finally, the attention output AV E Txd from all heads is then concatenated, and multiplied by the output projection W, DxD which is then passed to the normalization and feedforward components. This standard approach is summarized in Figure (left).
3 Multi-Token Attention
Each attention value in standard multi-head attention, see Equation 1, depends solely on single key and query vector. That means all the necessary information for finding and
2

---

Published as a conference paper at COLM 2025
Head Tensey
Key query lanad
queriex
queries
heads
heads
keys
keys
Figure 2: MTA applies key-query and head convolution over attention values.
attending to a relevant part of the context must be compressed into these single vectors. This might not be ideal if we are looking for sentence containing multiple elements. Consider for example the sentence "Where did Alice see the rabbit?". We could try to find instances of " "Alice" and "rabbit" independently and then check if there is sentence that has both. Let qu and 9r be query vectors encoding "Alice" and "rabbit" respectively (assuming word tokenizer), then their attention weights are computed as follows:
Aa Softmax(q.k LoiKK//ar /vd), a, = Softmax(q, /vd)
(2)
By doing normal attention with those queries, we can attend where " "Alice" and "rabbit" are mentioned in the context. All we have to do then is to check if both attention weights aa and ay have higher probabilities at the same nearby locations, e.g. in the same sentence, which will indicate that the sentence mentions both "Alice" and "rabbit". Unfortunately, normal attention lacks such interaction between attention maps, and instead only uses them to compute output values. Even if we use different attention heads to find " "Alice" and "rabbit", there is no mechanism to combine these attention weights. This motivates us to modify the attention mechanism to allow combining different attention maps from nearby locations (both in terms of query and key locations), or between different attention heads.
As shown Figure 1 (right), our proposed Multi-Token Attention consists of three im- portant components built top of multi-head attention: key-query convolution, head mixing convolution, and group normalization with gating mechanism. The overall MTA convolution applies the key-query convolution to combine multiple keys and queries within heads, and the head convolution to share knowledge between heads and amplify impor- tant information. Finally, we apply group normalization with scalar gating to push back against residual streams and improve gradient flow. In this section, we will describe each component of M'TA in detail.
3.1 Key-query convolution
Pre-softmax convolution combine information from multiple query and key tokens:
MTA applies convolution operation on attention logits to
A Softmax (Aonv2dy(A)),
where Conv2dg is a 2-dimensional convolution operation with kernel weights 9, and kernel sizes (Cg, Cx). Convolution applied to the key and query length dimensions, while the batch and head dimensions remain independent. More precisely, the attention weight aij
from query 9: to key kj is computed as follows:
Tc/2 1
a:j- Softmax
Σ
Σ
1izi-9 0zZ 9:-P k / Vd
7--10/21
Given that the query position is i, any information from position should not be used to prevent information leakage from the future. That is why only past queries are used in Equation 4. For keys, we use indicator function 1/zjー" zero out future keys. However such masking is too complex implement (it is necessary to modify the convolution cuda kernels), thus we propose simpler version that applies the existing causal masking twice:
A Softmax (Mask (Conv2dg Masko(AA)))
(5)
3

---

Published as a conference paper at COLM 2025
Here, the first masking uses 0 so that those values will not affect the output from the convolution. Although this version masks out a little more than necessary, it is simpler to implement and prevents information leakage, so use it as our default.
Post-softmax convolution Similarly, we can apply the convolution on top of the attention weights instead of the logits
A Masko (Conv2dg (Softmax (Mask-ooAA)])
(6)
This makes interaction between attention weights additive instead of multiplicative. We will experiment with both versions, but will use the pre-softmax version by default.
Each attention head has separate 0 parameters, so they can perform different convolution operations. The chosen kernel dimensions dictate how far away tokens can be combined together. In the above example, the question "Where did Alice see the rabbit?" will make qa and dr queries separated by two tokens (assuming a word tokenizer), so we need Cq = 4 to cover both queries. Similarly, the target sentence "Alice saw the white rabbit under the tree" for example produces keys k2 and k, separated by three tokens, so C₁ = 5 is sufficient for combining them. Before applying the convolution, we pad the input with appropriate amount of zeros so that each convolution operation has a valid input.
3.2 Head mixing convolution
Unlike the key-query convolution, which allows mixing of attention weights from different time steps, we further propose to use a head convolution over groups of heads, attention weights from different heads can be combined. In particular, for a head convolution kernel of size C₁, all heads are divided in M/C, c/ groups. Within each group, we apply non- overlapping convolution operation". This way, MTA allows not only to condition attention weights on multiple query and key vectors within each head, but also shares attention information across heads. For example, consider splitting all heads into groups of two, such that the kernel size is сN 2. Let us use superscript denote head indices, so A and A2 are attention weights from two different heads. Then the new attention weights are:
New = 111 + W12 A = W21 A + W22
(7)
where W11, W12, W21, W22 are kernel weights. Here the mixing occurred after the softmax, but can also mix logits before the softmax:
L Wn1 AL 1 + W22 A = W21 A + Wr2
Like the key-query convolution, we experiment with both versions.
3.3 Putting everything together
In the previous sections, we introduced two different ways of mixing attention weights, one across key-query time steps and another across different heads. These two can be implemented together in a single MTA module. Because each has pre and post-softmax versions, there are multiple different ways of combining them.
If both mixing methods are pre-softmax, then they can be implemented by single 3- dim convolution operation as shown in Figure 2. Two of these dimensions will span key and query dimensions as described in Section 3.1. The third dimension will be across attention heads, but on groups of с heads. The same 3-dim convolution can be applied after softmax to perform both mixing methods post-softmax. The third possibility is to apply the key-query convolution before softmax, and then head mixing after softmax. This also straightforward by applying Equation 7 on the attention weights obtained from Equation 5.
Finally, we follow Ye et al. (2024) who also apply normalization, but our case we replace the layer-dependent scaling with sigmoid gating mechanism. Sigmoid gating allows the model to switch heads on and off across layers and better adapt to different tasks. We also perform ablations across different normalization approaches.
2Since the kernel size t) and group size Gj are the same, it can also be viewed as fully-connected.
4

---

Published as a conference paper at COLM 2025
Block size N = 5
Block size N = 8
Model
All
First
Last
All
First
Last
Transformer
51.6 + 43.1
78.2 + 1.5
1.3 + 04
31.2 + 50.3
28.9 + 24.9
58.4 + 4688
MTA
0.1 0.1
0.0 0.0
0.0 + 0.0
0.1 4: 0.0
0.0 + 0.0
0.0 = 0.0
Table 1: Error rates (%) the motivating toy task where the model needs to locate block of letters containing the given = letters. The output should contain all, first, or last tokens of the target block. MTA perfectly solves this task while a standard Transformer struggles to learn it. We report average error rate over 3 seeds and their standard deviations.
4 Experiments
We conduct experiments with the MTA architecture, comparing to several baselines a set of standard and long-range dependency tasks, starting with toy task. Unless otherwise specified, we apply the MTA convolution both pre-softmax and post-softmax.
4.1 Motivating toy task
We start with simple toy task to demonstrate the effectiveness of our method over standard multi-head attention. In this task, the model is given sequence of blocks where each block consists of N random letters. This followed by L question letters (L < N). The objective is to find the block that contains all question letters in any order. Then the model must output all letters of the target block, or only its first or last token (as three separate task variants).
Despite its simplicity, this task poses challenge to standard attention because it requires pieces of information (i.e. question letters) to identify the target block. To succeed, standard soft attention must encode all L question letters into single query vector. In contrast, MTA can first find the locations of each question letter, then use its convolution operation to increase the attention to the locations where all L letters are found together.
For this task, we train a small Decoder-only Transformer model with layers, 2 heads and 256 hidden dimensions. The results are shown in Table 1. As expected, Transformer with standard multi-head attention struggles to solve this task, often completely failing to find the target blocks even though the questions have only L = 2 letters. This highlights the inherent limitation of standard attention conditioned on single token vectors. For MTA, set the query kernel size C = 2 to match the number of question letters, and the key kernel c = 2N - so it can cover a whole block on both sides. To simplify the experiments, the key-query convolution is only applied before softmax, and no head convolution is used. This way, is sufficient for each query vector to encode only a single question letter while the convolution operation can aggregate their locations find the target block. As result, MTA successfully solves all the versions of the task with near zero error rate. See Appendix for more training details.
4.2 Large language modeling
For language modeling experiments, we perform pre-training of 880M-size models and compare the Transformer model (Vaswani et al., 2017), Differential Transformer (DIFF Transformer) (Ye et al., 2024), Transformer with Talking heads attention (Shazeer et al., 2020), and Transformer with MTA. DIFF Transformer calculates attention scores as the difference between two separate softmax attention maps, while Talking heads attention applies linear projections across heads before and after softmax operations, thus being related to our head convolution approach, so we use them as baselines. All models are trained in the same setup on the SlimPajama (Soboleva et al., 2023) dataset for 105B tokens using the Lingua framework (Videau et al., 2024). Training details are provided in Appendix D. To improve training efficiency, convolution is applied on all layers. Kernel dimensions are fixed at c. = 6, = 11, and we mix groups of C₁ 16 heads within each layer. We ablate these parameters in Section 4.6.
apply the key-query convolution on every 4th layer, while head
For each model, we conduct training twice and report average validation perplexity in Table 2. We observe consistent improvements across all validation datasets for the model
5

---

Published as a conference paper at COLM 2025
Model
arxiv
book
c4
github
wiki
Avg PPL
Pretraining
Transformer
4.65
13.47
20.20
14.41
4.28
10.13
11.64
11.25 (0.00)
DIFF transformer
4.62
13.33
19.99
14.28
4.25
10.04
11.54
11.15 (0.02)
Talking heads
4.59
13.26
19.82
14.15
4.20
9,93
11.33
(0.00)
MTA
4.54
13.09
19.63
14.00
4.12
9.76
11.18
10.91 (0.01)
Long context finetuning
Transformer
4.32
13.18
20.14
14.08
3.96
9.84
11.63
11.02
DIFF transformer
4.28
13.01
19.87
13,93
3,90
9.72
11.49
10.89
Talking heads
4.29
13.25
19.95
13.90
3.86
9.66
11.29
10.88
MIA
4.21
12.77
19.51
3.79
9.46
11.14
10.65
Table 2: Validation perplexity for 880M Transformer model on SlimPajama dataset after training for 105B tokens, and finetuning with 2048 4096 context extension for another 10.5B tokens. Pretraining perplexity was averaged across two runs.
Model
BoolQ
PIQA
SIQA
HellaS WinoG
ARCe
ARCc OBQA
MMLU
Avg
Transformer
56.2
70.2
39.9
38.5
56.4
57.9
25.9
23.8
24.5
43.7 (0.3)
DIFF transformer
59.6
70.5
39.7
38.9
56.4
57.7
25.6
21.4
24.9
43.9 (0.5)
Talking heads
61.9
71.4
40.6
39.4
54.5
58.2
25.2
23.6
24.7
44.4
MTA
62.1
71.7
40.4
39.7
57.2
58.9
24.7
23.6
25.8
44.9
Table Pretrained models' evaluation results on standard benchmarks. Results are averaged over two model training runs for each method.
trained with MTA, even though the key-query convolution only applied to quarter of the layers, keeping the total number of parameters on-par with the Transformer baseline (see Appendix Table 8). We also note that group normalization with layer scaling is important component that drives superior performance for both the DIFF Transformer and MTA architectures.
We further evaluate our models on set of popular benchmarks in a zero-shot setup as shown in Table 3. The model trained with MíA outperforms the baselines on most of them and achieves higher average score, despite these not being long-context tasks.
4.3 Long context finetuning
We further finetune our models the same mix of datasets for another 10.5B tokens, but increase the context length from 2048 to 4096. We increase RoPE's theta to 500000, change the weight decay 0, and reduce warm up steps 50. Other parameters remain the same as during pretraining (see Appendix D). The resulting Transformer model with MTA similarly outperforms the new baselines in perplexity evaluations as shown in Table
4.4 Long-range dependency tasks
It was previously shown that Transformers struggle to find relevant information especially in the middle of long context (Liu al., 2024; 2025). To test MTA this setting, we evaluate trained models on three tasks: LAMBADA (Paperno et al., 2016; Radford et al., 2019), Needle- In-A-Haystack (Kamradt, 2023) and BabiLong (Kuratov et al., 2025). All these tasks require models to almost sharply attend to the rong-range tokens buried in the context.
LAMBADA collection of texts that test the model's ability search long-range de- pendencies in the context. In particular, this dataset is designed such that humans can correctly predict the next word only if they have the whole text, but not they only see the last sentence preceding the target word. We observe models trained with MTA are better at correctly guessing the next word (Table 4), significantly outperforming the baseline Transformer model.
Needle-In-A A-hasstak We probe our models by inserting 2, 4, and needles in 2k and 4k context windows at varying depths. We additionally ensure that for each new sample, the

---

Published as a conference paper at COLM 2025
Model
LAMBADA standard ! LAMBADA OpenAI
Transformer
17.6
9.5
DIFF transformer
14.9
9.3
Talking heads
15.1
8.9
MTA
Table 4: Perplexity evaluations on the LAMBADA standard (Paperno et al., 2016) and LAM- BADA OpenAl (Radford al., 2019) datasets.
Multi-needie retrieval accuracy, context
Multi-needie retrieval accuracy, context
100
80
Accurasy, 16
Accuracy, %
60
Transformer
Transformer
transformer
40
Talking heads
Talking heads
MTA
MTA
20
20
2
2
Number needles
Number needles
Figure 3: Multi-needle retrieval accuracy averaged over different needle insertion depths. Pretrained models are evaluated with 2K context (left), while finetuned models are evaluated with 4K context (right).
needles are shuffled, thus removing bias the models might have toward extracting the needle that was inserted first or last. Each setup is evaluated on 500 samples. Accuracy evaluations averaged across the depths of insertion. As reported in Figure 3, we observe significant improvement in needle extraction abilities for models trained with MTA across all needle counts and varying context lengths. Breakdown by depth is reported in Appendix G.
BabiLong This benchmark tries to to assess language models' ability to reason across facts scattered long documents. We focus on tasks QA1-5 (Weston al., 2015) where a correct response requires various numbers of facts or argument relations. Examples of input and target output can be found in Table 6. We present average accuracy in Figure 5(left) and per-task in Appendix Figure 7. We observe that the MTA model performs well compared to other models, especially when there more distraction text (4K tokens) in the input.
4.5 Kernel patterns
Key-query and head convolution kernels across different layers and heads are shown in Appendix H. We observe that number of of the key-query kernels are close to identity, with near-one weight learned for targeted query and key, and near-zero weights learned for all other positions. However, there are numerous kernels that are more intricate.
One such kernel shown in Figure 4(left), which has a diagonal structure. This kernel will amplify attention if sequence of query tokens match sequence of keys. Such ability useful in searching sentences that matches the current sentence. Indeed, we observe that this particular kernel focuses attention on the target needle (see Figure 4(right)) in the Needle-in-the. Haystack task, where we must locate the magic number of the matching city.
Other kernels, for example, can be viewed as priming, amplifying if the same key was attended by previous queries (head on Figure 10), or edge detecting, amplifying the first or last of multiple contiguous keys with high attention (heads 8 and 11 on Figure 12). We leave further exploration of key-query kernel maps and their meaning for future research.
Head kernel patterns, the contrary, are simpler due to their small size, as shown in Appendix Figure 15. Besides an identity with scaling, a common pattern is contrasting: subtracting one attention weight from another. We also observe that the kernel scales increase with layers when there is no group normalization (see Appendix Figure 16). This is probably to compete with the residual stream, which gets larger with the model's depth.
7

---

Published as a conference paper at COLM 2025
Francisco
number
pinch
magic
Keys
10
fatal
The
of
San
is
B
The
50
022
Queries
Francisco
San
number
magic
Figure/ 4: Kernel pattern (left) and corresponding attention map (right), which has the highest attention scores on the targeted needle "The magic number of San Francisco is 8". This kernel amplifies attention if query token sequence matches key sequence - useful for searching for related sentences that match the current sentence.
Average Performance on Task QA1-5
80
11.250
Transformer
Dffferental Transfrmer
70
Ta Takking heads MTA
11.225
Accuracy (%)
60
Perplexity
11.200
50
11.175
40
11.150-
Differential transformer
11.125-
12
0K
1K
2K
4K
Key-query convolution layers
Distraction
Figure 5: (left) Average accuracy QA1-5 tasks in BabiLong. The models are all pretrained with 2K context and then finetuned with 4K context. Distraction text length varies from OK (no distraction) to 4K tokens. MTA consistently outperforms the baseline models. (right) Ablation on the number of MTA layers with key-query convolutions (head convolution is applied all layers). We report average validation perplexity on SlimPajama.
However, this pattern is not present with group normalization because it will undo the effect of such scaling.
4.6 Ablations
Key-query convolution To understand how key-query convolution affects the model performance, we run ablation studies on the number of layers where key-query convolution is added to the attention. The results shown in Figure 5(right) demonstrate that when only layers are enhanced with MTA, the model can outperform strong baselines, while 6 layers with MTA strikes balance between performance and additional complexity (see Appendix for more about computational complexity).
Head convolution An ablation study on the head kernel size shows that increasing kernel size improves average validation perplexity, as shown in Figure 6 (left), allowing for better communication across heads.
Kernel initialization Kernel initialization can affect the convergence rate and stability of training. We evaluated MTA models initialized with zero, constant (0.3), and identity kernels. The latter corresponds to initialization with a regular Transformer without convolution in the attention. We found that identity initialization leads to better convergence and final performance, while models initialized with zero and 0.3 values reduce average validation perplexity by 0.02 and 0.08 correspondingly.
8

---

Published as a conference paper at COLM 2025
Transformer
Talking heads MTA
Difl transformer
-
11.225
0.0
11.200
-0.5
Perplexity
11.175
Perplexity gain.
-1.0
11.150
11.125
2.0
11.100
-2.5
11.075
-3.0
11.050
2.5
7.5
10.0
12.5
15.0
0.4
1.0
1.2
Head kemel size
Parameters (B)
Figure 6: Ablation on the head kernel size in MTA convolution (left), and scaling laws (right). We report average validation perplexity on Sliml'ajama.
Key-query conv
Head conv Pre-sm
Pre-sm
Post-sm
Post-sm
Group norm
PPL |
scalar gating
10.90
depth scaling
10.92
X
scalar gating
10.95
10.99
depth scaling
10.99
X
no scaling
11.03
X
x
A
depth scaling
11.09
AXMX
x
11.16
no scaling
11.13
layer-norm scaling
11.41
C = 4, & = 9
XXX
X
depth scaling
11.23
Gq 6, 45
X
depth scaling
11.23
= 8, C =
X
depth scaling
11.31
x
V
depth scaling
11.11
XXXXX
x
11.19
depth scaling
11.10
11.20
Table 5: Ablation on MTA components: validation perplexity over SlimPajama dataset.
Ablation on MTA components We further experiment with MTA kernels different sizes, changing the order of convolution operations with respect to softmax, and different normalizations, such as layer-norm scaling (Sun et al., 2025). Results are summarized in Table 5. We found that different kernel sizes display similar kernel patterns, while resulting in slightly different evaluation results (middle rows). Group normalization and exponential vanilla MTA (top
depth scaling are both important factors each bringing improvement rows), while square root layer-norm scaling surprisingly underperforms. Finally, changing the order of convolutions wrt softmax operations increases perplexity only by 0.01-0.04 points (bottom rows).
Scaling laws We experiment with models of various sizes: 300M, 550M, and 1B, with hyperparameters listed Table As shown in Figure (right), we observe a consistent pattern across models of different sizes, where MTA provides superior performance.
Related work
Focusing attention There has been a number of attempts to focus the soft attention mechanism on important context. For example, modifications have been proposed make

---

Published as a conference paper at COLM 2025
softmax sharper. Martins & Astudillo (2016) propose to replace softmax with sparsemax for sparse activations. Veličković et al. (2024); Nakanishi (2025) propose Adaptive temperature and Scalable-Softmaa that adjust exponential base in softmax attention. Irrelevant tokens can be altogether removed from memory (Sukhbaatar et al., 2021) to make attention easier. In Golovneva et al. (2024); Desrochers et al. (2024) the authors incorporate contextual information into position encodings to allow positions to be amplified by context. More recently, several noise canceling mechanisms have been proposed for attention. Talking heads attention (Shazeer et al., 2020) adds linear projections across the attention-heads dimension before and after the softmax operation. DIFF transformer (Ye al., 2024) uses differential amplifiers to focus attention to the relevant context, which is related to our head mixing step. Cang et al. (2025) further extends it by introducing normalization steps. OpAmp adaptation (Wu et al., 2025) propose adapters, which are fine-tuned with noisy context to enhance an LLMs' denoising capability, Xu et al. (2024) modified attention so that keys and values can shift by one time step. This can be viewed special case of MTA where the convolution performs shifting in the key dimension.
Convolution in Attention Convolution layers have been used with attention (Gehring et al., 2017; Wu et al., 2019), especially in vision Transformers, to decompose each image into sequence of tokens with fixed length, and then apply multiple standard Transformer layers (Xiao et al., 2021; Wu et al., 2021). There have been some attempts to include convolution in language modeling as well. For example, Liu et al. (2018) use convolution to compress the keys and values in the multi-headed attention by a factor of 3 thus allowing to process sequences 3x in length. Liu & Lapata (2019) later augment this architecture with the ability to encode multiple documents in a hierarchical manner, and further improve on long-contedd summarization tasks. Later Subramanian et al. (2020) propose three hierarchical models to learn different levels of abstraction in language by interchanging casual transformer layers with convolutions or pooling, or compressing representations for tokens further away in the past. Gulati et al. (2020) proposed Conformer architecture for speech recognition, where convolution module is applied to multi-head self-attention output. This architecture was later used to encode speech in the Llama-3 herd of models (Grattafiori et al., 2024). (Zheng et al., 2024) apply a convolution on the attention weights in the key dimension that enhances Transformer's 's extrapolation capabilities to long context. However, none of these methods apply convolutions across multiple queries, keys and heads to the attention weights.
Conclusion
In this paper, we focused on a limitation of the standard soft attention mechanism that stems from conditioning on the similarity of single vector pairs. This makes challenging for Transformers to precisely locate relevant information based on richer distinguishing information. As a remedy, we proposed a novel Multi-Token Attention mechanism that combines attention weights from multiple queries, keys and heads, making it possible to attend using more tine-grained information. From a simple motivating toy task to large- scale LLM experiments on a variety of popular benchmarks we demonstrate that models equipped with MTA achieve enhanced performance, especially when tested on tasks that require the precise location of relevant information within long context.
References
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint ar arXio:1409.0473 2014.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Yueyang Cang, Yuhang Liu, Xiaoteng Zhang, Erlu Zhao, and Li Shi. Dint transformer. arxio preprint ar Xix::250111144 2025.
10

---

Published as a conference paper at COLM 2025
Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/me questions. arXiv preprint ar Xii::1955.10044, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arvi preprint arXiv: 803.05457,, 2018.
Sarah Desrochers, James Wilson, and Matthew Beauchesne. Reducing hallucinations in large language models through contextual position encoding, 2024.
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolu- tional sequence to sequence learning. In International conference on machine learning, pp. 1243-1252. PMLR, 2017.
Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Contextual position encoding: Learning count what's important. arXiv preprint xxXio:2405.18719 2024.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint ar orXiv:2407.21723 2024.
Anmol Gulati, James Qin, Shung-Cheng, Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmente transformer for speech recognition. ar Xio preprint ar XrXi:20055088000 2020.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. aiv preprint arXiv:2009.03300 2020.
Gregory Kamradt. Needle in haystack pressure testing 1lms. https://gi http://ggtttuhcccco com/ Skamradt/LLMTest.NNee Mww..eeest.eeedlennaaatttttttt 2023.
Yury Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of Ilms with long context reasoning-in- a-haystack. Advances in Neural Information Processing Systems, 37:106519-106554, 2025. 2025.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024.
Peter J Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. Xrv preprint ar 1Xii::80111019: 2018.
Xiaoran Liu, Ruixiao Li, Mianqiu Huang, Zhigeng Liu, Yuerong Song, Qipeng Guo, Siyang He, Qiqi Wang, Linlin Li, Qun Liu, et al. Thus spake long-contex: large language model. ar Xiv preprint arXit:2502.1712) 2025.
Yang Liu and Mirella Lapata. Hierarchical transformers for multi-document summarization. Xrv preprint ar arXio:1905.1114, 2019.
Andre Martins and Ramon Astudillo. From softmax to sparsemax: A sparse model of attention and multi-label classification. In International conference on machine learning, pp- 1614-1623. PMLR, 2016.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. arx preprint arXiv:1809.022 022899 2018.
Ken M Nakanishi. Scalable-softmax is superior for attention. arv preprint arXiv:2501.1939, 2025.
11

---

Published as a conference paper at COLM 2025
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arviv preprint (Xit:1606.06033 2016.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106,
2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. ar preprint arXin:1904.09728, 2019.
Noam Shazeer. Glu variants improve transformer. ar preprint ar arXt::2002.05202 2020.
Noam Shazeer, Zhenzhong Lan, Youlong Cheng, Nan Ding, and Le Hou. Talking-heads attention. aro preprint ar Xix::2033002436 2020.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel SlimPajaam:: A 627B token cleaned and deduplicated version of RedPajama. slimpajama 627----ke-------eeeee----e //dtassstt/eeerrer///iiiiiiiiiia
Hestness, and Nolan Dey.
https://cerebre ai/blog/ June
2023. URL
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocompptting,, 568:127063, 2024.
Sandeep Subramanian, Ronan Collobert, Marc' Aurelio Ranzato, and Y-Lan Boureau. Multi- scale transformer language models. ar Xrv preprint ar Xri::2005500581, 2020.
Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, and Angela Fan. Not all memories are created equal: Learning forget by expiring. In International Conference on Machine Learning, 2021. URL /ttpss://api semanticscholar. org/CorpusID::2 2346816161
Wenfang Sun, Xinyuan Song, Pengxiang Li, Lu Yin, Yefeng Zheng, and Shiwei Liu. The curse of depth in large language models. arXiv preprint ar 2rXt::2502.0579,, 2025.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXio preprint ar 13311:1302.13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL httpss://proceedings. wewpreddiinnnnuuriisss... 3f5ee243547dee91 1od0ee011c1111144444aaaaaaae pdf..
2017ff///
Petar Veličković, Christos Perivolaropoulos, Federico Barbero, and Razvan Pascanu. softmax is not enough (for sharp out-of-distriiutiion)) arXiv preprint arXiv:2410.01104, 2024.
Mathurin Videau, Badr Youbi Idrissi, Daniel Haziza, Luca Wehrstedt, Jade Copet, Olivier Teytaud, and David Lopez-P'az. Meta Lingua: A minimal Py Torch LLM training library, 2024. URL https: //gi hatuub/ccoo/accboooreeeerrrcc/c/
Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart Van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. ar Xrv preprint ar arXie:1502.00555 2015.
Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. aiviv preprint ar 1io: 990.10030,, 2019.
12

---

Published as a conference paper at COLM 2025
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cvt: Introducing convolutions to vision transformers. In Proceedings of the IEEE/CVF international conference computer vision, pp. 22-31, 2021.
Haoyuan Wu, Rui Ming, Haisheng Zheng, Zhuolun He, and Bei Yu. Efficient opamp adaptation for zoom attention to golden contexts. arXiv preprint arXit:2502.12502, 2025.
Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, and Ross Girshick. Early convolutions help transformers see better. Advances in neural information processing systems, 34:30392-30400, 2021.
Mingyu Xu, Wei Cheng, Bingning Wang, and Weipeng Chen. Kv shifting attention enhances language modeling. ArXiv, abs/2411.19977 2024. URL https://api. semanticscholar. wrg/CorpusID:27442200
Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arv preprint arXiv:2410.052558, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv:1905.0788) 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Chuanyang Zheng, Yihang Gao, Han Shi, Jing Xiong, Jiankai Sun, Jingyao Li, Minbin Huang, Xiaozhe Ren, Michael Ng, Xin Jiang, et al. Dape v2: Process attention score as feature map for length extrapolation. arx preprint arXio:2410.04798, 2024.
A Limitations
To the best of our knowledge, the Multi-Token Attention method is not currently compatible with the popular optimized attention kernels. Optimizing the runtime performance was not the goal of this work, thus we leave further optimization the MTA implementation for future research.
B MTA Architecture
Head convolution normal attention with higher rank Let us consider mixing of two heads after softmax. If W! and W? are the output projections corresponding to those heads, their output can written as
hhe V')W + (Ahew W2)
= (WIA Vl + A W- + (122) A 122 A
Aww W! + w21 1 v2w,) + A2( wx2 W??
= A (W!) + +wy + AWW + w HW?
A H(w (a[(w//www W!! + w w3w?) + A3 + w2222
we compare the first term to a normal attention output A H-WWWW we note that the only have two rank-d matrix additions instead of one. Actually can
difference is that rewrite
Wn W wo + 22 W.W? WPW
where W! E Rx2d and W! R Thus, post-softmax head convolution can be repli- cated by a normal attention head with twice the rank. This may help in understanding why head mixing is useful, it may increase the expressive power using higher rank. However, it is not truly identical to 2x rank because the parameters of the two heads are not independent.
13

---

Published as a conference paper at COLM 2025
SUPPORTING
TWO FACTS
SUPPOKINGG
Tordare Toklns hoads
32-
Tordare
231
Tikles
32
Turstree Iaka bood
32-
12
45
D atrarttinn
Detraction
Detrrction
NELATIONS
RELATIOON
brling
Accurecy
50
22-
20
Ditraction
Détractin
Figure 7: Accuracy (%) on QA1-5 tasks in BabiLong benchmark. Note the random perfor- mance on QA3 is 16.67%, thus all models perform poorly on QA3. On other tasks, MTA demonstrates strong performance compared baselines especially when there lengthy distraction text (4K).
Let us see if the same is true for the pre-softmax version. The attention logits after mixing is
A' wnnQ'KT KIT + 22202K2T
wn WW WT T + 1 HW2 HT H(wn WT T + w1 Wo2 H
Again, this can be rewritten like a normal attention logits
A HW H where W3 W E Dx2d As before, it can be replicated by normal attention with twice the rank in the key and query projections.
C Toy task details
The blocks are built by randomly choosing N E {5,8} lowercase alphabets and concatenat- ing them. We join up to 50 blocks into a single sequence separated by ". followed by "#" and L = 2 question letters. Here is an example sequence:
siint.afigt.ww.. whftb. bjtpq. .m.mmy blocks)... pxjvf. www.kk..kkkkkk poi
Here the 4th block "bjtpq" is the target because it contains all query letters "pb". During data generation, we make sure there is only one target block in each sequence. We use 1M such sequences for training and test on held-out 1K sequences.
We train small Transformer model with layers, 2 heads and 256 hidden dimensions. The batch size is 64 and the training continues for a total of 100K steps. We run each experiment three times with different random seeds and report the average performance.
D Language model training details
Training setup follows LLaMa architecture (Touvron et al., 2023), and includes RMSNorm pre-normalization (Zhang & Sennrich, 2019), SwiGLU activation function (Shazeer, 2020),
14

---

Published as a conference paper at COLM 2025
Task
Name
Input
Target
Question
John travelled the hallway.
Mary journeyed to the bathroom. Daniel went back the bathroom. John moved to the bedroom.
QA1
single supporting fact
bathroom
Where is Mary?
Mary journeyed to the bathroom.
Daniel grabbed the football there.
QA2
two supporting fact
kitchen
Where is the football?
Sandra grabbed the milk there.
Daniel went to the kitchen.
Daniel moved to the bathroom.
Mary got the football. Mary went back to the kitchen.
Where was the football before the garden?
QA3
three supporting fact
kitchen
Mary journeyed to the garden.
The hallway is east of the bathroom. The bedroom west of the bathroom.
What the bathroom of?
QA4
two arg relations
bedroom
Fred picked up the football there.
Fred gave the football to Jeff. Jeff gave the football to Fred. Fred gave the football to Jeff.
Who did Fred give the football to?
QA5
three arg relations
Jeff
Table 6: Examples of QA1-5 tasks in BabiLong benchmark when there is no distraction context.
and Rotary Embeddings (Su et al., 2024). Table 7 shows detailed hyperparameter break- down for Transformer model pretraining. All models were trained in the same setup, except Diff Transformer, which had twice less heads account for doubled dimension.
We evaluate on the following tasks: BoolQ (Clark et al., 2019), PIQA (Bisk et al., 2020), SIQA (Sap et al., 2019), HellaSwag (Zellers et al., 2019), WinoGrande (Sakaguchi et al., 2021), ARC easy and challenge (Clark et al., 2018), OpenBookQA (Mihaylov et al., 2018), We also report 5-shot performance on the aggregated MMLU benchmark (Hendrycks et al., 2020).
E BabiLong details
We include examples of input and output from BabiLong benchmark in Table 6. The details of per-task evaluation results on QA1-5 are shown in Figure 7.
F Complexity evaluation
Table 8 shows estimated additional number of parameters and their actual counts. We note that DIFF transformer (Ye et al., 2024) introduces four learnable vectors per layer
15

---

Published as a conference paper at COLM 2025
Parameter
300M model
550M model
880M model
model
Dimension, D
1024
1280
1536
2048
Layers
20
24
24
24
Heads, M
16
10
16
16
RoPE theta
100, 000
100, 000
100, 000
100, 000
Batch size (tokens)
262, 144
262, 144
262, 144
524, 288
Context length
2048
2048
2048
2048
Learning rate
1.5e-4 4
1.5e 4
1.5e - 4
1.5e 4
Weight decay
0.05
0.05
0.05
0.05
Warm up steps
375
375
375
187
Table 7: Hyperparameters for Transformer training.
Model
Additional parameters, estimates
Actual number of parameters
Transformer
876,553,728
DIFF transformer Talking heads
#+ L# * M L* (2+d+1)
876,567 ,552
876,566,016
MTA
2+ MiG
876,583,320
Table 8: Estimation of the parameter counts for each model architecture.
Model
Memory, GB
FLOPS, 10
TPS, 10
Transformer
17.5
25.0
54.3
DIFF transformer
21.6
8.4
17.6
Talking heads
63.0
6.9
14.6
MTA
73.8
2.6
5.7
Table 9: Average memory consumption and training speed for models trained on 32 NVIDIA H200 GPUs with our unoptimized code.
(Aqi, ARI, Ag2, A42), plus group normalization weights. In total, that results in more parame- ters than Multi-Token Attention with key-query convolution added the quarter of layers, even though all kernel weights in MTA models are different between all heads and layers.
Table 9 shows actual runtime and memory consumption recorded during 880M model train- ing. Note that the baseline and DIFF transformer utilize scaled dot product attention function implemented in Torch, that calls optimized CUDA kernels for improved performance. In contrast, our MTA implementation does not take advantage such efficient kernels, which is the major reason behind its lower FLOPS.
G Multi-needle evaluation results
We report on the detailed evaluation results in Figure 8. We observe that models trained with MTA are better at finding needles hidden deeper in the context.
H MTA kernel patterns.
Key-query convolution kernels across different layers and heads are reported in Figure Figure 14.
Head kernel patterns for the model with 2,=2 are shown Figure 15. They are simpler due to their small size. Besides an identity with scaling, a common pattern is contrasting: subtracting one attention weight from another. We also observe that the kernel scales
generated
dot product attention.htmm
16

---

Published as a conference paper at COLM 2025
2k context length, needies
context length, 4 neadles
10-
08-
0.8-
0.6
0.6-
0.4
0.4
0.0
25.0
50.0
75.0
100.0
25.0
50.0
75.0
190.0
Depth,
Depth.
Tonsformr
DVf transformer
Talking hoads
MTA
Transfomer
bransforfemmr
Taiking heads
MTA
context needles
context
needles
0.8
0.6-
Scairayy
044
0.4
0.0
100.0
0.0
Ponsformr
tranffoorr
Trandfomer
Figure Needle-in-a accuracy across different models. Models trained with MTA are better at predicting needles hidden deep in the context.
17

---

Published as a conference paper at COLM 2025
Layer 3, head
Layer 3, heac
Layer 3, head
Layer head
Sengs
601
keys
keys
Layer 3, head 5
Layer 3, heac
Layer 3, head
Layer head
keys
keys
Layer 3, head
Layor head
Layer head
Layer 3, hoad
keys
keys
head 13
Layar head
Layer head
Layer 3, hoad
keys
Figure 9: Key-query kernel patterns across heads 3rd layer of the Transformer model with MTA.
18

---

Published as a conference paper at COLM 2025
Layer
Layer
Layer 7,
Layar
Layar
Layer
kame
Layer
head
Layer head
Layer head 15
Layer head
keys
ReyS
keys
keys
Figure 10: Key-query kernel patterns across heads at 7th layer of the Transformer model with MTA.
19

---

Published as a conference paper at COLM 2025
Laver 11, head
Layer 11, head 2
Layer 11, head
Layer 11, head
keya
Says
Layer 11, head 5
Layer 11, head
Layer 11, head 7
Layer 11, head
Kitel
keyl
Laver 13, head
Layar heac 13
Layar head 11
Laver head 12
head
heae
head 15
head 1.6
Figure 11: Key-query kernel patterns across heads 11th layer of the Transformer model with MTA.
20

---

Published as a conference paper at COLM 2025
15,
Layer 15,
15,
15,
GS
0.5
kow
Laver 15, head
Layer 15, head 20
Laver 15, head 11
Layer 15, head :2
keys
kmys
Layer 15, sead 23
Layer 15, hoad 24
Layer 15, head 15
Layer 15, head :6
C2
bays
CR
Figure 12: Key-query kernel patterns across heads at 15th layer of the Transformer model with MTA.
21

---

Published as a conference paper at COLM 2025
Layer 19, head
Layer 19, head
Layer 19, head
Layer 19, head
QU
19, head
Layer 29, head
Layer 29, head
10, head
01
head
06
0.
Eager
Layer
Layer
Layyrr
Layer
06
QA
Figure 13: Key-query kernel patterns across heads at 19th layer of the Transformer model with MTA.
22

---

Published as a conference paper at COLM 2025
Layer 23, head
Laver 23, head 2
Layer 23, head 3
Layer 23, head
keyt
keyp
leys
Layer 23, hozd
Laver 23, head
Layer 23, head 7
Layer 23, hoad
keys
keys
keps
beys
Layar 23, Iead
Layer 23, load 10
Laysr 23, heac 11
Layer 25, head 12
06
keys
Layer 23, head 13
Layer 23, head 14
Layer 23, heac 15
Layer 23, head 16
06
keyz
Figure 14: Key-query kernel patterns across heads at 23rd layer of the Transformer model with MTA.
23

---

Published as a conference paper at COLM 2025
pain
Layer
Layer
Layer
8.00
A.M
-35
05
-05
65
-05
0.5
-05
855
Layer pair
Layer pair
Layer pai?
Layer pair
-23
-29
6.6
-00
-06
0.5
Foad
hoad
Lever heed patn
Layer heed pair 2
Layer heed per
Layer heed parn
-05
+05
-05
Layer head pair
Layor hoad pair
Layer head pair
Layer head pai
3.5
3.5
0.5
A.1
+3.%
5
5
1oad
TOA
Layer heed pain
Layer head pai
Layer head pai
Layer head pa e
000
-25
-25
-2.
-0.
621-
221-
Fand
201-1
haed
Figure 15: Head kernel patterns across first three layers of the Transformer model with MTA.
increase with layers when there is no group normalization (see Figure 17). This is probably to compete with the residual stream, which gets larger with the model's depth. However, this pattern is not present with group normalization because it will undo the effect of such scaling.
Finetuning with MTA
One natural question readers might ask is MTA could be integrated into models that were already trained with standard attention? Would this require complete retraining, or could it be added through some form of adaptation? Architectually, MTA can be added as additional layers to already trained models, and weights can be updated with continual
24

---

Published as a conference paper at COLM 2025
with group normaliziion
win group normallztion
025
0.6
0.20
0.5
deggoaal ares
0.3
013
0.00
0.0
Figure 16: Average of the diagonal kernel values across head pairs for Multi-Token Attention models with (left) and without (right) group normalization. Values for layers with key-query convolution are colored in orange.
NTA with group no momiizziion
MTA wio group normaliztion
0.0
2.5
Fum equand diggiral td
-0.5
-0.4
-14
-0.6
-0.9
100
35
10
20
bayers
Ryers
Figure 17: Average of the ratio between non-diagonal and diagonal kernel values across head pairs for Multi-Token Attention models with (left) and without (right) group normalization. Values for layers with key-query convolution colored in orange.
25

---

Published as a conference paper at COLM 2025
Pretraining
Cont training
arxiv
book
c4
cc
github
wiki
Avg PPL !
1.4B models
Transformer
Transformer
4.54
12.87
19.21
13.77
3.99
9.49
10.96
10.69
Transformer
MTA
4.51
12.78
19.09
13.67
3,95
9.41
10.87
10.61
MTA
MTA
4.47
12.57
18.82
13.47
3.89
9.25
10.65
10.44
Llama 3 hend
Llama 3.2 1B
Llama 3.2 1B
4.54
14.60
18.14
13.40
3,92
8.84
10.25
10.53
Llama 3.2 1B
MTA
4.52
14.49
18.04
13.32
3.89
8.79
10.19
10.46
Llama 3.2.3B
Llama 3.2
4.12
12.08
15.58
11.50
3.37
7.52
8.35
8.93
Ilama 3. 3B
MTA
4.11
12.03
15.51
11.46
7.48
8.31
8.89
Llama 3.1 8B
Ilama 3.1
4.00
11.04
15.04
10.99
3.29
7.33
8.00
8.53
Llama 3.1 8B
MTA
3.98
10.97
14.97
10.94
3.27
7.29
7.97
8.48
Table 10: Validation perplexity on SlimPajama dataset after continuous training for 10.5B tokens on our 1.4B models, and continuous training for 5.3B tokens for Llama 3 models.
training. Since the main component of MTA is a convolution, we can initialize it to identity and insert it in existing Transformer layer. Such modification will not change the output of the layer, so when we start training the model, it should maintain its performance. However, the added convolution starts deviating from the identity state, will allow the transformer to condition its attention on multiple keys and queries.
To understand how well pre-trained models can learn perform preliminary experiments by finetuning our 1.4B model, as well as opensourced Llama models (Grattafiori et al., 2024). In all these experiments we used the same finetuning setup similar Section 4.3, but kept the context length at 2048 tokens.
Validation perplexity results are reported in Table 10. We observe that all models finetuned with MTA were not only able to incorporate new kernels, but outperform baselines in terms of perplexity.
26


Reviews:
[{"review_id": "CPWOPtlrge", "reviewer": "Reviewer_ryV8", "summary": "This paper proposes Multi-Token Attention (MTA), which conditions attention weights on multiple query and key vectors via convolutions across tokens and heads. It extends standard full attention by incorporating a pre-softmax convolution, head-wise mixing, and an additional group normalization layer. Experiments on both short-context and long-context benchmarks demonstrate consistent improvements with the proposed attention variant.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper is clearly written and easy to follow.\n\n2. The proposed technique is easy to plug into standard attention.\n\n3. The proposed attention variant consistently improves performance across various benchmarks.", "reasons_to_reject": "1. The first concern is that the example used to motivate multi-token attention is not convincing enough. After the current token-mixing operation, the information from multiple tokens is naturally fused and can be processed by the subsequent FFN layers and the next attention layer. It is unclear why we need to explicitly operate on the attention map of the current layer. The example and explanation only hold when there is a single attention operator in the model.\n\n2. Adding convolution on top of Q, K, and V is common in previous language models, especially in linear attention models like GLA, DeltaNet, and Gated DeltaNet. The difference is that this work introduces convolution after the matmul operation in attention, yet no benchmark or explanation is provided to justify why this is a better solution.\n\n3. The effectiveness of head mixing is not ablated. Conceptually, the subsequent FFN layers can also learn to average across different heads. It is unclear why weighted averaging of two consecutive heads is beneficial.\n\n4. The efficiency of the proposed attention variant is not analyzed, especially considering that the introduced convolutions and group normalization may introduce additional overhead.", "questions_to_authors": "In addition to the attention map of one example, could you also show the difference in averaged attention maps?\n\nOther questions have been included in the above weakness section.", "extra_scores": {}}, {"review_id": "I6hS7Vwvr2", "reviewer": "Reviewer_AZNg", "summary": "The paper proposes a new attention mechanism called Multi-Token Attention (MTA), aiming to address the limitations of traditional single-token attention mechanisms in precisely locating relevant information within context. The experimental results demonstrate that MTA outperforms baseline models across multiple benchmarks, particularly in long-context tasks, highlighting the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The method presented in this paper has a clear motivation and aligns well with the context of graph neural networks. Currently, attention scores are calculated based only on single-hop (directly connected) query-key pairs. However, by using convolution or other methods, multi-hop features can be incorporated.\n2. The experimental results highlight two significant advantages: lower perplexity in language modeling and better performance in long-context understanding.\n3. The additional parameters introduced are negligible.", "reasons_to_reject": "1. While the motivation of this paper is clear, the technical innovation is somewhat limited. Beyond validation through ablation studies, there is a lack of substantial theoretical or analytical support. Although the effectiveness of multi-hop makes sense, the paper does not discuss whether CNN kernels and group norm are the only ways to achieve this, or if other methods could be equally effective.\n\n2. I observed that in the language modeling tasks, MTA with group norm achieves better performance (Tables 2, 3, 4), while MTA without group norm performs comparably to the baseline. However, in the long-context retrieval tasks, MTA without group norm shows better performance (Table 5, Fig. 7), and MTA with group norm performs comparably to the baseline. Could the authors explain the reasons behind these results?\n\n3. As an improvement work based on the Transformer architecture, this paper lacks evaluation regarding scalability. The baseline Diff. Transformer has demonstrated scalability from 830M to 13B parameters, whereas this paper discusses only a very limited range of model sizes.\n\n4. Although the authors discuss efficiency issues in the Limitations section (memory and computational efficiency being only 1/5 of the baseline), attributing it to the absence of efficient operators, I would like the authors to discuss whether MTA is compatible with online softmax optimization algorithms. If not, this would represent a significant drawback of the method, implying that under fair computational efficiency comparisons, MTA should be compared to a baseline that is five times larger in size.", "questions_to_authors": "1. In the toy task experiment presented in Table 1, does the Transformer architecture align with that used in other sections of the paper (i.e., is it a llama-like architecture rather than a naive Transformer-Decoder)?\n\n2. Why is the DIFF w/o group norm. baseline included in Table 2 but absent from other experiments?\n\n3. I suggest noting that this paper discusses the Transformer-Decoder model, and that the complete Transformer also includes the Encoder.", "extra_scores": {}}, {"review_id": "L8OTKaMglU", "reviewer": "Reviewer_3NEq", "summary": "This paper proposes Multi-Token Attention (MTA), a modification to the standard transformer attention mechanism that allows attention weights to condition on multiple query and key vectors simultaneously rather than just single token pairs. The authors achieve this by applying convolution operations over queries, keys, and heads, allowing nearby elements to influence each other's attention weights. They test MTA on toy tasks, standard language modeling (pretraining an 880M parameter model on SlimPajama), and long-context tasks like Needle-in-a-Haystack and BabiLong, where they find improvements over standard transformers. The key innovation is allowing the attention mechanism to leverage richer information beyond what a single token vector pair can provide.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The core idea is quite interesting and addresses a fundamental limitation in transformer attention. The authors clearly identify that standard attention can struggle when \"looking up a sentence that mentions both 'Alice' and 'rabbit'\" as it \"requires the query vector to encode both tokens\" (lines 33-35). Their convolution approach seems like an elegant way to add this capability.\n- The motivating toy task effectively demonstrates MTA's benefits, showing how it can easily solve a problem where standard transformers struggle. The task requires finding blocks containing specific letters - something standard attention has trouble with, but MTA solves with \"near zero error rate\" (lines 170-172).\n- The ablation studies on different MTA components provide good insights. Varying the number of layers with key-query convolution (Fig. 4) shows that even adding just a few MTA layers brings significant improvement. The visualizations of kernel patterns in the appendix are interesting\n- The paper has real improvements on needle-in-haystack and BabiLong tasks that explicitly test long-range dependency capabilities, which is a clear application area where this approach makes sense.", "reasons_to_reject": "- The model size used in experiments (880M) is quite small by current standards. Given that MTA is meant to add expressivity that might otherwise require \"increased dimension, and for the model to use lots of its capacity\" (lines 37-38), it would be more compelling to see if the benefits persist in larger models (e.g., 7B+) where capacity constraints might be less of an issue.\n- The paper doesn't fine-tune any existing pre-trained LLMs with MTA, making it hard to assess how easily this could be integrated into current workflows. As the authors themselves note in the limitations section, \"Multi-Token Attention method is not currently compatible with the popular optimized attention kernels\" (lines 428-429), which raises questions about practical adoption.\n- The computational overhead seems significant. Table 10 indicates MTA uses nearly 3.5x more memory than the baseline transformer and achieves only about 20% of the training throughput. The authors don't adequately address whether this tradeoff is worth it - is a 0.16 perplexity improvement (11.25 to 11.09 in Table 2) worth this much additional compute?\n- The paper lacks some important baselines, like sparse attention methods or sliding window approaches that might address similar problems with less overhead. While they compare to DIFF Transformer, more comparisons to other long-context solutions would strengthen the paper.", "questions_to_authors": "- Could you clarify how MTA would be integrated into models already trained with standard attention? Would this require complete retraining, or could it be added through some form of adaptation?", "extra_scores": {}}, {"review_id": "5wQxNIL6Av", "reviewer": "Reviewer_maCc", "summary": "The paper introduces multi-token attention (MTA) as a replacement for single-token attention in the original Transformers. This approach utilizes convolution operations over queries, keys, and heads. Across various evaluations, MTA outperforms the standard Transformer and DIFF Transformer in language modeling (PPL), popular benchmarks in zero-shot settings, and tasks involving long-range dependencies. An ablation study highlights that group normalization is an important component of this approach.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The concept of replacing single-token attention with multi-token attention is compelling to me, and the idea is simple and novel. \n2. There are various tasks, including language modeling (PPL), popular benchmarks in zero-shot settings, and tasks involving long-range dependencies.", "reasons_to_reject": "1. The improvement over the Transformer baseline is negligible. For instance, in Table 2, the average PPL decreases from 11.25 to 11.09, representing an improvement of less than 2%. Similarly, the average score on standard benchmarks in Table 3 increases from 43.7 to 44.4, which is also an improvement of less than 2%.\n2. The experiments are conducted solely on 880M-size models, so it remains uncertain whether the methods would be effective on larger models.\n3. There is no comparison of training and inference costs with the baselines.", "questions_to_authors": "Please refer to the reasons to reject.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a new attention mechanism to address the limitation of single token attention. 

Initially, the scores were two 4 and two 6 with some major concerns.

During the rebuttal, the authors successfully addressed the issues and the scores becam 7, 6, 6, and 5.

AC carefully read the paper, the reviewers' comments, and the authors' feedback.

The authors response seem to sufficiently address most major concerns raised by a negative reviewer.

So, AC recommends accepting this paper, and asks the authors to reflect the comments in the final version.
Decision: Accept