Transformers traditionally propagate information through a single additive residual stream, where each sublayer reads only the most recent state. Attention residuals relax this constraint by allowing each sublayer to attend to the entire depth history via a learned softmax. However, this read operation uses a single query shared across the entire width, forcing all feature subspaces to use the same distribution over layers—a compromise that becomes increasingly costly as model width grows and subspaces disagree on which layers to read.
To address this, researchers at Hugging Face propose Multi-Head Attention Residuals (MHAR). The key idea is to reshape the routing query into H per-subspace heads, each with its own softmax over the depth history. This makes the read operation block-diagonal, adds zero parameters, and incurs negligible compute. Setting H=1 recovers standard attention residuals exactly.
In experiments training from scratch on a deduplicated, quality-filtered, STEM- and code-heavy corpus based on Nemotron, MHAR improved validation loss over a standard Transformer at 100M, 350M, and 1B scales, with gains of -0.061, -0.149, and -0.140 respectively. It outperformed four other methods in every setting, with the advantage growing at larger scales.
The head count H is a genuine design axis, not a free knob: validation loss is U-shaped with respect to H, with a flat optimum at H=4 or H=8 across scales. The authors adopt H=8 for large-scale models, noting that over-splitting to H=16 gives back part of the gain. A direct probe of trained queries confirms that learned subspace disagreement is the underlying driver.
Efficiency-wise, 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.