Multi-Head Attention Residuals: Per-Subspace Depth Routing with Zero Added Parameters

Authors: Cheng Luo, Zefan Cai, Junjie Hu

arXiv ID: 2607.27230

Problem: Transformers propagate information across depth through a single additive residual stream: every sublayer reads only the most recent state. Attention-residual variants relax this by letting each sublayer attend over the depth history through a learned softmax, but that read uses a single query shared across the entire width, so every feature subspace must read depth history through one distribution. The cost of that forced compromise grows with how much subspaces disagree about which layers to read, and disagreement grows with model width.

Key Methodology:

Multi-Head Attention Residuals (MHAR): reshape the routing query into H per-subspace heads, each with its own softmax over the depth history, making the read block-diagonal
The reshape adds zero parameters and negligible compute; H = 1 recovers attention residuals exactly
Trained from scratch on a deduplicated, quality-filtered, STEM- and code-heavy Nemotron-based anneal corpus at 100M, 350M, and 1B, compared against standard Transformer and three alternative depth-routing methods
Head-count study across H values to map the design axis
Identity-preserving conversion using delta attention residuals for 8B mid-training
Fused Triton routing kernels for training throughput

Key Results:

MHAR improves validation loss over a standard Transformer at 100M, 350M, and 1B by -0.061, -0.149, and -0.140, the best result among four methods in every setting, with the gain increasing from 100M to larger scales
Validation loss is U-shaped in H with a flat optimum at H = 4 or H = 8 across scales; H = 8 chosen for large-scale models, and over-splitting to H = 16 consistently gives back part of the gain
A direct probe of the trained queries confirms learned subspace disagreement is the underlying driver
Fused Triton routing kernels raise attention-residual training throughput from 0.2-0.5x to 0.55-0.88x of baseline while maintaining near-baseline peak memory
The identity-preserving 8B mid-training conversion yields +3.2 on GSM8K and +3.1 on GPQA

What it means for developers: A parameter-free architecture axis with measured gains that grow with scale: per-subspace depth-history routing is a drop-in reshape on top of attention residuals, and the H = 4/8 optimum makes the knob safe to default. The whole gain comes from restructuring the read, not from adding capacity, so it is relevant to efficiency-minded training stacks, mid-training conversions, and small-model annealing pipelines rather than only large-lab pretraining. It also serves as a counterexample to the reflex that architecture improvements require new parameters.

Paper: arXiv:2607.27230