PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Head Attention Residuals

Cheng Luo, Zefan Cai, Junjie Hu

arXiv:2607.27230v2cs.AI

TL;DR

Standard residual streams and single-query attention residuals limit how independently feature subspaces can retrieve information from depth history. MHAR gives each subspace its own routing head and achieves the best validation loss among four methods from 100M to 1B, with gains increasing at larger scales.

  • Problem

    Residual streams expose only the latest state, while attention residuals force every feature dimension to retrieve depth history using one shared distribution.

  • Method

    MHAR splits the routing query into independent per-subspace heads over depth, adding no parameters and recovering attention residuals exactly at H=1.

  • Results

    MHAR is best among four methods at 100M–1B, improving validation loss over the baseline by −0.061/−0.149/−0.140, with gains growing at larger scales.

  • Takeaways & Limitations

    Validation loss is U-shaped in head count, with a flat optimum at H=4 or H=8 across scales; H=8 is adopted for large-scale models.

  • Takeaways & Limitations

    At 1B, single validation passes have approximately ±0.07 sampling noise, and single-cell tables report one fixed seed.

Abstract

from arXiv · show

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 disagreement grows with model width. We introduce Multi-Head Attention Residuals (MHAR): the routing query is reshaped into H per-subspace heads, each with its own softmax over the depth history. The read becomes block-diagonal, the reshape adds zero parameters and negligible compute, and H = 1 recovers attention residuals exactly. Trained from scratch on a deduplicated Nemotron-based anneal corpus that is quality-filtered and STEM- and code-heavy, MHAR improves validation loss over a standard Transformer at 100M, 350M, and 1B (-0.061, -0.149, and -0.140). It achieves the best result among four methods in every setting, with the gain increasing from 100M to the larger scales. The head count is a real design axis rather than a free knob: validation loss is U-shaped with respect to H, with a flat optimum at H = 4 or H = 8 across scales. We adopt H = 8 for large-scale models; over-splitting beyond this point (H = 16) consistently gives back part of the gain. A direct probe of the trained queries confirms that learned subspace disagreement is the underlying driver. Fused Triton routing kernels increase attention-residual training throughput from 0.2-0.5x to 0.55-0.88x of the baseline while maintaining near-baseline peak memory. An identity-preserving conversion using delta attention residuals supports 8B mid-training, yielding improvements of +3.2 on GSM8K and +3.1 on GPQA.

1 Introduction

MHAR replaces the single depth-history query of attention residuals with H per-subspace queries, allowing independent routing while preserving the original method at H = 1. The introduction motivates this change by subspace disagreement, and reports scale gains, a U-shaped head-count optimum, and improved routing efficiency.

  • Motivation: A single attention-residual query forces every feature dimension to use the same softmax distribution over prior layers, collapsing differing subspace depth preferences.The paper identifies this forced compromise as the central limitation of single-head depth routing.
  • Method: MHAR splits the routing query into H per-subspace heads, each independently attending over depth; reshaping adds no parameters or FLOPs, and H = 1 exactly recovers attention residuals.The resulting depth read is block-diagonal rather than shared across all feature dimensions.
  • Systems: 0.55–0.88× baseline throughput is achieved by fused Triton routing kernels, improving attention-residual training from ∼0.2–0.5× while maintaining near-baseline peak memory.These kernels make repeated depth-history reads practical at scale.
  • Results: −0.061/−0.149/−0.140 validation-loss gains over the baseline at 100M, 350M, and 1B make MHAR the best of four methods at every scale.The gains increase from 100M to the larger scales.

2 Method: Multi-Head Attention Residuals

Multi-Head Attention Residuals let different feature subspaces independently route over depth history, removing the single-distribution constraint of attention residuals. The method preserves parameter and FLOP matching while recovering attention residuals exactly at H = 1.

  • Motivation: Standard residual streams make earlier sublayer outputs present but not individually addressable for selective rereading.Each sublayer reads only the most recent running state.
  • Multi-head routing: H = 1 recovers attention residuals exactly, because the single-head route kernel is the only change from the original method.The routed slices are concatenated back to width d after independent depth softmaxes.
  • Multi-head routing: MHAR partitions coordinates into H heads, assigns each head its own query, and independently routes each slice over depth.Different subspaces can therefore attend to different layers, while the forward pass remains identical to attention residuals.
  • Properties: d parameters: MHAR’s H queries form a tensor with exactly d parameters, matching single-head routing.The head count is a reshape rather than a widening, so multi-head routing adds zero parameters.
  • Fused routing kernels: 2× over torch.compile and 6.4× over eager: fused deterministic Triton kernels cut per-microbatch routing time.Depth routing is memory-bound because every sublayer rereads the depth history.

3 Experiments

Across unified from-scratch experiments on a quality-filtered, STEM- and code-heavy anneal corpus, MHAR consistently improves validation loss over the standard Transformer and competing routing methods. The gains persist in cross-distribution benchmarks and 8B mid-training, while head disagreement and an intermediate head count explain when the method helps.

  • Method comparisons: Hyper-connections improve baseline loss by −0.032/−0.066/−0.086, but MHAR outperforms them by 0.029, 0.083, and 0.053 at 100M/350M/1B.The comparison uses the same sweep with n=4 hyper-connections.
  • Head-count ablations: H=16 gives back MHAR’s gain: validation deltas worsen from −0.149 to −0.102 at 350M and from −0.140 to −0.069 at 1B.H=16 is past the H=4–8 optimum and lands above single-head routing at all three scales.
  • Downstream evaluation: MHAR improves perplexity and LAMBADA at every scale, and improves HellaSwag at 350M and 1B in zero-shot cross-distribution evaluation.The benchmarks are web-derived and held out from the anneal training mix.
  • 8B mid-training: +3.2 GSM8K and +3.1 GPQA are the routing gains in schedule-matched 8B mid-training, while MMLU, MATH, and code remain statistically unchanged.The gains are concentrated on GSM8K and GPQA; GSM8K changes from 123 versus 81 discordant items and GPQA from 27 versus 13.
  • Efficiency and mechanism: MHAR’s extra parameters are only +0.02%, and the multi-head split itself adds zero parameters, compute, and memory over single-head routing.The split reshapes one query into H smaller queries and replaces one depth-axis softmax with H softmaxes.

4 Ablations

Ablations show that multi-head routing is most effective with a moderate head count: validation loss has a flat optimum at H=4–8, while H=16 overshoots. The benefit of multi-head routing also depends on convergence and data distribution, where single-head routing is not consistently robust.

  • Head-count ablation: H=4 (3.132) and H=8 (3.129) are essentially tied for the 1B, KV=8 optimum, while H=16 (3.173) rises beyond the basin.The H=16 result is 0.044 above the optimum, beyond the approximate ±0.02 tail-mean noise.
  • Design rationale: Splitting the routing query into H heads addresses the limitation that one learned query cannot serve every feature subspace simultaneously.The attention-residual bypass is structurally a self-attention query over depth, motivating a multi-head design analogous to multi-head attention.
  • Head-count ablation: Multi-head routing improves loss at every KV, with per-KV optima at H=4 or H=8 and H=16 overshooting in every row.The full kv×H grid shows H=16 remains below H=1 despite giving back part of the multi-head gain.
  • Convergence signature: At 100M after only 5K steps, KV=1 shows a monotone 0.028 loss drop from H=1 to H=8, exceeding the ∼0.006 single-seed noise floor.The higher-KV rows are within the noise floor in this deliberately under-trained grid.
  • Data-distribution robustness: On the web-corpus replication, H=1 helps at 100M (−0.039) but exceeds the baseline at larger scales (+0.038 at 350M).The passage identifies single-head routing as the original attention-residual formulation and reports that it is not robust as the model grows.

5 Conclusion · A Related Work

MHAR removes the single-query routing bottleneck by giving feature subspaces independent depth-routing heads, while preserving parameter efficiency and extending to practical large-scale training. Related work places MHAR among approaches that enrich cross-layer connectivity, routing, and interpretability without replacing the single residual stream.

  • 5 Conclusion: MHAR gives each feature subspace its own depth-routing head, removing the single-query compromise at zero additional parameter cost.The method was trained from scratch from 100M to 1B and improved over a standard Transformer at every scale.
  • 5 Conclusion: 0.55–0.88× baseline throughput is achieved by fused routing kernels while maintaining near-baseline peak memory.These systems improvements make multi-head depth routing practical beyond research scale.
  • A Related Work: Residual connections provide additive identity paths for deep training, whereas DenseNets concatenate previous feature maps for direct access to earlier representations.MHAR’s source list is dense in spirit, but its mechanism reads depth history rather than concatenating feature maps.
  • A Related Work: DenseFormer learns depth-weighted averages, MUDDFormer learns input-dependent dense layer connections, and RealFormer carries residual paths on attention scores.ReZero and DeepNet instead rescale the standard residual branch.
  • A Related Work: Attention residuals learn attention over layer-output history; MHAR builds on them by splitting the single routing query across feature subspaces.Unlike approaches that enrich connection topology through separate streams, MHAR enriches routing distributions while keeping a single stream.
  • A Related Work: MHAR outperforms hyper-connections at matched recipe across all three scales, while concurrent work routes over per-sublayer deltas to reduce source collinearity.The approaches target different factors of the forced-compromise cost.
  • A Related Work: Mixture-of-experts and mixture-of-depths route tokens or layers to allocate compute, whereas MHAR routes every token over already-computed sublayer outputs.MHAR changes what each sublayer reads, not which tokens are processed.
  • A Related Work: MHAR applies multi-head attention’s specialization intuition across layers, motivated by views of the residual stream as a shared communication channel.Mechanistic analyses and evidence that transformer layers act near-linearly on the stream provide interpretability motivation.

B Attention residuals as the single-head route

Attention residuals are the single-head special case of MHAR: one shared query produces one softmax over depth that all feature coordinates use. Setting H = 1 recovers this route exactly and defines the single-head baseline used in the tables.

  • Single-head route: One shared query produces one softmax over depth, which all d coordinates read through.This is the defining single-head attention-residual route.
  • Single-head route: H = 1 recovers attention residuals exactly as a drop-in replacement for the multi-head route.The forward pass remains unchanged when the multi-head route is replaced by the single-head route.
  • Single-head route: The single-head method is the baseline reported in every table.It is the reference method for MHAR-versus-single-head comparisons.

C Routing-head sweep at 100M

At 100M, the routing-head sweep favors matching H to KV at the tested KV values, with H=8 adopted at scale. An under-trained replication shows that additional heads can help before router convergence, especially for KV=1.

  • Noise interpretation: At 1B, H=8 is treated as within noise of optimal because grid differences are ∼0.02, below the ∼±0.07 single-pass evaluation noise.The sweeps are single-seed, so the strict grid minimum is not emphasized.
  • Routing-head sweep: At 100M, KV = 4 and KV = 8 achieve their optima at H=4 and H=8, respectively.The H=8 setting is adopted at scale.
  • Under-trained replication: A 5K-step replication finds a monotone 0.028 validation-loss drop from H=1 to H=8 for KV = 1.This exceeds the ∼0.006 single-seed noise floor; higher-KV rows remain within noise.

D Why single-head routing breaks: the forced-compromise mechanism

Single-head routing forces every feature subspace to share one depth distribution, creating a compromise when their preferred layers diverge. Trained-query probes show this disagreement is learned, grows with scale and width, and tracks the loss penalty that MHAR removes.

  • Forced compromise: A single query routes the entire width through one softmax, forcing all subspaces to share one depth distribution.This forced compromise becomes costly when subspaces prefer different layers.
  • Direct probe: 0.27–0.70 versus ∼0.03–0.10: trained queries show far more width-disagreement than matched-norm random queries.The learned excess grows from 0.235 to 0.606 across 100M–1B.
  • Scale dependence: 0.235 →0.606: learned width-disagreement increases monotonically across scale, while source collinearity is non-monotonic.The forced-compromise cost peaks at 1B, where single-head routing regresses below baseline.
  • Loss connection: −0.009 to −0.020: widening d512 →d768 increases MHAR’s validation-loss advantage over single-head routing as measured disagreement rises.The matched-recipe control isolates width at fixed L12 and kv4.
  • Routing structure: No measurable benefit: aligning routing to KV head groups performs no better than arbitrary subspace routing within the ±0.07 eval-noise floor.This supports subspace disagreement, rather than attention-head alignment, as the relevant routing structure.

E Fused routing kernels

Fused Triton routing kernels address depth routing’s memory-bound, quadratic traffic by sharing source storage, fusing normalization with online softmax, and accumulating backward gradients in place. The implementation reduces saved activations while preserving the reference computation up to reduction-order effects.

  • Motivation: Depth routing is memory-bound and moves O(L^2BTd) source data because each sublayer reads all prior states.The reference implementation also incurs repeated stacking, materialized normalized keys, separate kernels, and many backward launches.
  • Forward implementation: A shared [2L+1, B, T, d] source buffer writes each residual state once and eliminates per-sublayer source-list copies.Routing call s reads rows [0, Ns) directly from the shared buffer.
  • Forward implementation: Online softmax and RMSNorm are fused in the forward kernel, which reads each source row once and saves only the routing weights.The saved routing-weight tensor is a factor d/H smaller than the stacked sources retained by the reference path.
  • Backward implementation: The backward kernel makes two source-row passes, recomputes norms, and accumulates gradients into one shared fp32 buffer in place.This replaces the reference backward’s thousands of small tensor additions; runtime assertions rely on autograd’s reverse-order execution.
  • Verification: ≤2.5 × 10−6 maximum relative error: in fp32, fused and reference losses are identical and every parameter gradient agrees within this bound.In bf16, fused gradients are on average closer to fp32 ground truth, with a geometric-mean gradient-error ratio of approximately 0.5 across batches.

F Compute and memory cost

The multi-head split adds no parameters, FLOPs, or memory, so MHAR and single-head attention residuals have identical routing cost by construction. Compute and memory were measured across three model sizes on a single H100 using controlled, repeated runs.

  • Routing cost: MHAR and single-head attention residuals have identical cost because reshaping the routing query is parameter-, FLOP-, and memory-free.H = 1 multi-head routing reproduces single-head attention residuals exactly, so any overhead is inherited from attention residuals rather than introduced by MHAR.
  • Measurement setup: Table 8 measures throughput and peak memory at 100M, 350M, and 1B on a single H100 with identical node and software settings.Each cell is the median of three steady-state windows, and the runs used verified-healthy GPUs.

G Width-control: isolating width in loss

Widening the model at fixed depth and KV increases learned subspace disagreement and changes validation loss in-distribution. The resulting loss comparison shows MHAR’s advantage over single-head routing more than doubles with width.

  • Width-control: 0.235→0.281 learned disagreement when widening d512→d768 at fixed depth and KV, while the random-query null and source collinearity stay flat.The probe isolates width as the driver of subspace disagreement.
  • Width-control: −0.041→−0.033 single-head routing’s edge over baseline when widening d512→d768, whereas MHAR’s edge changes −0.050→−0.053.Depth is held at L=12 and KV at 4.
  • Width-control: −0.009→−0.020 MHAR’s advantage over single-head routing as width increases from d512 to d768.This loss-level result corresponds to the width-isolated disagreement increase measured in the probe.

H Robustness of the from-scratch comparison

MHAR’s advantage remains robust on FineWeb-Edu: it stays below the standard baseline during training, wins when each method uses its own best learning rate, and improves across seeds at every scale.

  • Web-data comparison: MHAR stays below the standard baseline throughout 100M FineWeb-Edu training under identical conditions except routing mechanism.The runs use the same node, software, data order, and global batch.
  • Web-data comparison: At 100M, single-head routing helps, but at 350M–1B it rises above baseline while MHAR retains its full gain.This reversal occurs on FineWeb-Edu and is attributed to distribution dependence.
  • Robustness to learning rate: MHAR wins at every scale even when each method receives its own optimal learning rate.Single-head routing remains above baseline at 350M and 1B even at its kinder 5 × 10−4 rate.
  • Seed robustness: −0.045/−0.090/−0.071 nats: MHAR’s paired seed deltas versus baseline at 100M/350M/1B each exceed ten standard errors.Three seeds per method share a 5 × 10−4 recipe, with data order paired across methods.

I Training and reproducibility details

Training uses a shared code path with specified Qwen3-based architecture, optimization, data, and validation settings across methods and scales. The 8B mid-training setup additionally specifies distributed hardware, corpus construction, and reproducibility details.

  • Shared setup: All runs share a single code path, with settings common to every method and scale unless noted.
  • Mid-training and corpus: The 8B mid-training setup uses FSDP full-shard, bfloat16, activation checkpointing, 8×H100-80GB, sequence length 4096, weight-EMA β=0.999, logit z-loss 10−4, and weight decay 0.05.The anneal_pt_v3 corpus contains approximately 8 TB of text and 1.9 T tokens across 2,048 parquet shards, with measured source shares.
  • Architecture: The models use pre-norm Qwen3-style decoder blocks with RMSNorm, SwiGLU, grouped-query attention, RoPE, tied embeddings, and no dropout.Weights use N(0, 0.022) initialization; attention biases are disabled.
  • Optimization: AdamW training uses β1=0.9, β2=0.95, ϵ=10−8, global gradient clipping at 1.0, and a 20K-step warmup-plus-cosine learning-rate schedule.Training uses bfloat16 under distributed data parallelism through 1B parameters.
  • Batch and sequence length: 100M models use sequence length 2048 and global batch 64, while 350M and 1B models use sequence length 1024 and global batch 32.The base random seed is 42, offset per worker.
  • Validation protocol: Validation reuses the same corpus with a different shuffle and a +9999 seed offset rather than an explicit held-out split, evaluating every 500 steps.Token-level cross-entropy is averaged over evaluation batches and data-parallel workers.
Loading 2607.27230v2…