Everything K-culture — comebacks to K-beauty, straight to your inboxGet it in your inbox›
Multi-Head Attention Residuals
arXiv:2607.272302026-07-21
Instead of letting a transformer look back at earlier layers with one shared question, letting each feature group ask its own question makes it work better
Standard transformers just add each layer's output to a running sum, so only the immediately previous state is directly readable; an existing method called attention residuals lets every sublayer look back at all past layer outputs through one learned attention query. This paper argues that one shared query forces every feature subspace to read the past in the same way, and fixes this by splitting the query into H independent per-subspace heads (MHAR), a change that costs essentially no extra parameters or compute. Trained from scratch at 100M, 350M, and 1B parameters, MHAR beats a standard transformer at every scale, and grafting it onto an 8B model during continued training improved GSM8K and GPQA scores.
METAL MEDIA explanatory visual
How different methods read past layers
Evidence statusMeasured results reported
Standard residualeach layer just adds its output to the running sum, so only the immediately previous state is read
Attention residuals (H=1)one shared learned query scans all past layer outputs, forcing every feature to read the past in the same proportions
MHAR (multi-head)features are split into H groups, each with its own independent query reading past layers separately, at no extra parameter cost
8B continued trainingan output gate starts at zero so the pretrained model is untouched at first, then routing gradually opens during further training
An explanatory diagram made by METAL MEDIA, not a figure supplied by the paper's authors.
What they did
The existing 'attention residuals' method routes every feature dimension through one shared learned query over past layer outputs; this paper splits that query into H independent heads, one per feature subspace, calling the result Multi-Head Attention Residuals (MHAR).
The split is essentially a free reshape: it adds no new parameters and negligible extra compute, and setting H=1 recovers the original single-query method exactly.
Trained from scratch on FineWeb-Edu at 100M, 350M, and 1B scale, MHAR was compared against a standard transformer, hyper-connections, and the original single-head attention residuals.
Setting the number of routing heads equal to the number of KV (key-value) heads worked as a tuning-free default, and directly probing the trained queries confirmed that different feature subspaces genuinely prefer different layers.
For an 8B model, the authors built a conversion method that grafts MHAR onto an existing pretrained model without disturbing it, and also built custom Triton kernels to make the mechanism run faster.
Figure 1: The depth read is an attention, so it should be multi-head. Left: attention residuals (19) feed every sublayer a learned mixture of the depth history (embedding and every earlier attention/MLP output). Zoom left: one routing site is exactly single-head attention over that history—one query, one depth map shared by all d channels (Eq. 2). Zoom right: MHAR gives each of H subspaces its own map (Eq. 3), parameter-free; H=1 recovers attention residuals exactly.
Table 1: From-scratch validation loss (↓, tail-mean of the last eleven evaluations, steps 15–20k; §3) at the tuned peak learning rate 1×10−3; Δ is relative to the same-rate baseline; CEGFLOPs is the compute-equivalent gain of MHAR over the baseline (see text). MHAR is best at every scale; single-head routing regresses below baseline at 350M and 1B.
Scale
d/L/KV
Baseline
Hyper-conn.
Single-head
MHAR
CEGFLOPs
100M
512/12/4
3.336
3.322 (−0.014)
3.297 (−0.039)
3.287 (−0.049)
1.27×
350M
1024/24/8
3.290
3.260 (−0.030)
3.345 (+0.055)
3.210 (−0.080)
1.49×
1B
1280/36/8
3.175
3.172 (−0.002)
3.315 (+0.140)
3.111 (−0.063)
1.38×
Figure 2: What each method lets the current sublayer read across depth. Each column: forward stack (top) and the depth-read weight matrix (bottom; entry [i,c] is how strongly channel c reads depth row i; illustrative). Standard pre-norm: a one-hot read of the previous state. Attention residuals: one shared query reads all earlier sources—every column identical (rank-1). MHAR (ours): each subspace reads depth independently (block-wise columns).
Table 2: Zero-shot downstream evaluation at 100M, 350M, and 1B (single run). Each model is evaluated at its training context (100M: seq. 2048; 350M/1B: seq. 1024), so perplexity is comparable only within a scale. MHAR transfers its gain to held-out perplexity and LAMBADA at every scale, and HellaSwag at 350M (tied at 1B, marginally lower at 100M). 100M is the earlier 6×10−4 batch; 350M and 1B are the 5×10−4 batch.
Scale
Method
WikiText-2 PPL ↓
LAMBADA ↑
HellaSwag ↑
100M
Baseline
48.6
10.4%
33.0%
100M
MHAR
43.8
13.4%
32.5%
350M
Baseline
53.9
9.6%
33.5%
350M
MHAR
46.0
13.4%
35.5%
1B
Baseline
49.6
12.4%
34.0%
1B
MHAR
40.1
15.0%
34.0%
Figure 4: The H heads carry genuinely different depth-links. Each head’s token-averaged deviation w¯h−w¯ from the head-consensus routing at the first ten routing sites of the trained 1B MHAR model (dark cells: source not yet available). Deviations reach ±0.28 vs. 0.067 under a matched-norm random query, replicate on disjoint evaluation text (r=0.77), and are near-uncorrelated across heads.
Table 3: Downstream accuracy (↑) after 8B mid-training on anneal_pt_v3 (final checkpoints, EMA weights). The control and MHAR columns are schedule-matched: identical LR schedule, batch, data order, and ≈10 B-token budget.
Task
Base
Plain CPT
+MHAR
General
MMLU
0.554
0.643
0.645
GPQA
0.266
0.315
0.346
Math & Coding
GSM8K
0.190
0.470
0.502
MATH
0.053
0.191
0.191
HumanEval
0.122
0.409
0.415
MBPP
0.148
0.386
0.392
Figure 5: The conversion is exactly identity and introduces no optimization shock. Train loss of the schedule-matched 8B mid-training pair (identical seed and data order). Left: the two curves coincide. Right: the per-step |Δloss| starts at ∼10−5, rises only as the zero-initialized gate opens, and saturates 36× below step-to-step batch noise (dotted).
Table 4: Training speed and memory. Baseline Transformer vs. MHAR (torch.compiled reference kernels) vs. MHAR with our fused routing kernels, at the three model settings of Table 1. Throughput is relative to the same-scale baseline; MHAR is parameter-matched (+0.02%); single-head attention residuals are cost-identical to MHAR by construction. Protocol, absolute medians, and the full breakdown in Appendix F (Table 8).
Throughput (vs. base)
Peak mem. (GB)
Method
100M
350M
1B
100M
350M
1B
Baseline Transformer
1.00×
1.00×
1.00×
41.5
19.4
19.0
MHAR
0.54×
0.32×
0.23×
52.4
40.1
47.5
MHAR + fused kernels (ours)
0.88×
0.71×
0.55×
42.0
20.0
20.1
Figure 6: The bypass wants more heads. Under-trained regime (100M, 1×10−3, 5K steps): routing-head (H) × KV grid of validation loss (lower/yellow better; gold stars: per-KV optimum; red dashed line: the H= KV diagonal). The boxed KV= 1 row is the one supra-noise signal (see text).
Table 5: Routing-operation speedup, isolating the kernels from Table 4’s end-to-end numbers: wall-clock of all routing calls of one microbatch, forward+backward, bf16, single H100; the 8B column is the mid-training delta variant. End-to-end gains are smaller because routing is only part of a training step (batch settings and the eager-reference comparison in Appendix F).
Routing per microbatch (ms)
100M
350M
1B
8B (delta)
torch.compile
30.2
53.1
153.2
651
Fused Triton (ours)
5.7
11.4
42.4
323
Speedup
5.3×
4.7×
3.6×
2.0×
Figure 8: Validation loss of MHAR at 100M as a function of routing heads H (x) and KV heads (y). Lower is better. Gold stars mark the best H per KV; the red dashed line is the head-matched diagonal H= KV. At 100M the KV= 4 and KV= 8 optima sit on the H= KV diagonal. The 100M architecture of Table 1, trained at 5×10−4.
Table 6: Direct probe of the trained single-head routing queries (no new training; S=4 slices). Width-disagreement is the mean KL(as∥α) between each query slice’s preferred depth distribution and the shared α; a matched-norm random query gives the null, and the learned excess (trained−null) is the genuine learned signal. The middle column is a width-isolating control (d512→d768 at fixed L12/kv4/N): the learned disagreement rises while the null and collinearity stay flat, so width itself drives it. Across scale the learned excess grows monotonically 2.6× (→0.606) and is the primary, empirically-supported factor; source collinearity (mean pairwise cosine of the N=2L+1 sources)—a secondary, hypothesized factor we do not cleanly isolate—is non-monotonic (1.9× from 100M to 350M, then lower at 1B, partly a dimension effect), so the disagreement factor is the clean cross-scale signal. Caveat: slices are sub-parts of one trained query (a proxy for independent heads) rather than independently trained routers.
100M
124M (width ctrl)
350M
1B
Probe quantity (S=4)
d512/L12/kv4
d768/L12/kv4
d1024/L24/kv8
d1280/L36/kv8
Width-disagreement KL (trained)
0.273
0.311
0.570
0.701
random-query null
0.038
0.030
0.058
0.095
learned excess (trained−null)
0.235
0.281
0.512
0.606
Source collinearity (mean cos)
0.060
0.063
0.114
0.087
N sources (=2L+1)
25
25
49
73
Figure 11: Training loss at 100M (EMA-smoothed): MHAR stays below the standard baseline throughout training (inset: tail zoom, steps 12k–20k). The two runs are identical except for the routing mechanism (same node, software, data order, and global batch).
Table 7: Head-alignment ablation (validation loss ↓, final-step eval): routing over arbitrary subspaces (MHAR) vs. aligned to KV head groups (MHAR-HW), with matched head count. Aligning to attention heads shows no measurable benefit over arbitrary subspace routing; all differences fall within the ±0.07 eval-noise floor (single-seed, final-step). Separate matched-run batch (§3); read the within-table Δ.
Scale
MHAR
MHAR-HW
Δ(hw−mh)
100M
3.330
3.333
+0.003
350M
3.225
3.223
−0.002
1B
3.213
3.229
+0.016
Table 8: Training compute and memory at the three model settings of Table 1 (single H100; per-GPU microbatch as in training: batch 8 / sequence 2048 at 100M, 4/1024 at 350M, 2/1024 at 1B; identical node and software for all cells; 120-step runs on a verified-healthy GPU). Each cell is the median of three steady-state windows (within ±1%); peak memory reproduces to the decimal across sessions, and the 350M compiled/fused cells reproduce in a second independent session. Throughput is relative to the same-scale baseline (absolute baseline medians: 113.9k / 45.0k / 19.7k tokens/s). Parameters are matched to the baseline up to the O(d) routing queries (+0.02%). †Single-head attention residuals are cost-identical to MHAR by construction in every row—the head split is a parameter-, FLOP-, and memory-free reshape (§2)—so they are not listed separately.
100M
350M
1B
Method†
thr.
mem. (GB)
thr.
mem. (GB)
thr.
mem. (GB)
Baseline (additive residual)
1.00×
41.5
1.00×
19.4
1.00×
19.0
MHAR, reference kernels
0.31×
68.8
0.16×
70.9
OOM (>80 GB)
MHAR, torch.compile
0.54×
52.4
0.32×
40.1
0.23×
47.5
MHAR, fused Triton (ours)
0.88×
42.0
0.71×
20.0
0.55×
20.1
Table 9: Width-control (tail-mean validation loss, steps 15–20k). Holding depth (L=12) and KV (4) fixed and only widening d→d768: single-head routing’s edge over the baseline erodes (−→−0.033) while MHAR holds (−→−0.053), so MHAR’s advantage over single-head more than doubles (−→−0.020). This is the loss-level counterpart of the width-isolated disagreement in Table 6 (which rises →0.281 over the same widening). Single-seed; the single-head deficit shift (0.008) is near the per-seed noise floor, so the robust signal is the widening MHAR−single gap.
Width
Baseline
Single-head
MHAR
MHAR−single
100M (d512)
3.462
3.421 (−0.041)
3.412 (−0.050)
−0.009
124M (d768)
3.220
3.187 (−0.033)
3.167 (−0.053)
−0.020
Table 10: Best-vs-best learning rate: each method at its own optimum over peak LR ∈{1,5}×10−4, 10−3 (validation loss ↓, tail-mean of the last eleven evaluations; Δ vs. the same-scale best baseline; superscript is the selected LR). MHAR wins at every scale even when every method is given its own optimal learning rate; single-head, even at its kinder 5×10−4, still regresses at 350M and 1B.
Scale
Baseline
Hyper-conn.
Single-head
MHAR
100M
3.336 10−3
3.322 10−3 (−0.014)
3.297 10−3 (−0.039)
3.287 10−3 (−0.049)
350M
3.290 10−3
3.260 10−3 (−0.030)
3.328 5e–4 (+0.038)
3.210 10−3 (−0.080)
1B
3.175 10−3
3.172 10−3 (−0.002)
3.279 5e–4 (+0.105)
3.111 10−3 (−0.063)
Table 11: Seed robustness: paired per-seed delta vs. baseline (validation loss ↓, tail-mean of the last eleven evaluations; three seeds per method per scale, shared 5×10−4 recipe; ± is the per-seed standard deviation). Seed s fixes the data order across methods, so the comparison is paired. MHAR improves at every scale with |Δ|/SE≥10; single-head crosses from help to harm and is the highest-variance method.
Scale
Hyper-conn.
Single-head
MHAR
100M
−0.031±0.003
−0.043±0.003
−0.045±0.007
350M
−0.028±0.009
−0.005±0.015
−0.090±0.004
1B
−0.016±0.015
+0.093±0.049
−0.071±0.002
Table 12: Composition of the anneal_pt_v3 mid-training corpus: share of text bytes per source, measured over a uniform random sample of 48 of the 2,048 shards (187 GB of text; group rows sum their constituents, so rounded columns may differ in the last digit). The corpus is English-only; FinePDFs is restricted to its top three of twenty quality bins (verified from the per-document final_bucket metadata).
At 100M/350M/1B scale, MHAR improved validation loss over a standard transformer trained at the same learning rate by 0.049, 0.080, and 0.063 respectively.
The original single-head method helped at 100M (-0.039) but became roughly neutral at 350M (+0.055) and clearly hurt at 1B (+0.140); even when each method used its own best learning rate, single-head still regressed at 350M and 1B.
MHAR outperformed hyper-connections at every scale, with the margin widening as scale increased (0.035, 0.050, 0.061).
MHAR's gains carried over to zero-shot downstream evaluation on WikiText-2 perplexity, LAMBADA, and HellaSwag.
Grafted onto an 8B model during continued training under a schedule-matched comparison, MHAR improved GSM8K by 3.2 points and GPQA by 3.1 points (p=0.004 and p=0.038), and custom Triton kernels raised training throughput to 0.55-0.88x of the baseline's speed.
Where it can be used
Could be used as a near-free architectural upgrade to improve depth-wise information reuse in transformer pretraining pipelines.
Could be grafted onto existing pretrained large models during continued training to boost reasoning-related benchmarks without disrupting the base model.
Setting the number of routing heads to match the number of KV heads offers a simple rule of thumb for practitioners designing wider models where a single shared routing query becomes a bigger liability.
Limits and open work
Results are only reported up to 1B parameters trained from scratch and 8B parameters in continued training; effects at larger scales remain unverified.
Downstream evaluation covers a limited set of benchmarks (WikiText-2, LAMBADA, HellaSwag, GSM8K, GPQA); broader benchmark coverage is left as future work.
In the 8B experiment, MMLU, MATH, and code benchmarks showed no statistically significant change.
A direct equal-compute comparison against the plain baseline transformer has not yet been performed and is left for future work.
Why the optimum tends to sit near H equals KV heads, and exactly what individual heads learn to attend to, are not yet fully explained.
Why it matters
It shows that as transformers get wider, forcing every feature to read the depth history through one shared query becomes a real bottleneck, and offers a near-free architectural fix, which matters for anyone designing or scaling up large transformer models.
Terms in this paper
residual stream · the running sum that carries information from one transformer layer to the next by simple addition
attention residuals · an existing technique where, instead of just adding, each sublayer reads a learned weighted mixture of all previous layer outputs
head · in attention, a subset of feature dimensions processed with its own independent query/computation
KV head · in grouped-query attention, the (smaller) number of heads used for keys and values, shared across multiple query heads to save compute
validation loss · a measure of how well the model predicts held-out text; lower is better
Original abstract (English)
Transformers propagate information across depth through a single additive residual stream: every sublayer reads only the most recent state. Attention residuals relax this by letting each sublayer attend, through a learned softmax. However, that read uses a single query shared across the entire width, so every feature subspace must read the depth history through one distribution. The cost of this forced compromise grows with how much the subspaces disagree about which layers to read, and disagree
![Figure 2: What each method lets the current sublayer read across depth. Each column: forward stack (top) and the depth-read weight matrix (bottom; entry [i,c] is how strongly channel c reads depth row i; illustrative). Standard pre-norm: a one-hot read of the previous state. Attention residuals: one shared query reads all earlier sources—every column identical (rank-1). MHAR (ours): each subspace reads depth independently (block-wise columns).](https://media.metallab.ai/papers/2607.27230/f1.svg)
