MHAR splits transformer residual attention into per-head reads

Transformers move information across depth through a single additive residual stream: each sublayer only ever reads the immediately preceding state. A prior technique called attention residuals relaxes this by letting a sublayer attend, through a learned softmax, back over the full history of earlier layers rather than just the last one. The catch is that this attention uses one query shared across the entire width of the model, so every feature subspace has to read the depth history through the same distribution, even when different subspaces would prefer to look at different layers. The researchers argue that the cost of this forced compromise grows with how much the subspaces disagree about which layers matter, and that disagreement grows as models get wider.

To fix it, the paper introduces Multi-Head Attention Residuals (MHAR): the single routing query is reshaped into H separate per-subspace heads, each running its own softmax over the depth history, so the read becomes block-diagonal instead of monolithic. The reshape adds no new parameters and negligible extra compute, and setting H = 1 reduces MHAR exactly back to standard attention residuals.

The authors trained models from scratch on a deduplicated, quality-filtered, STEM- and code-heavy anneal corpus built on Nemotron data, at 100M, 350M, and 1B parameters. MHAR improved validation loss over a standard Transformer at all three sizes, by -0.061, -0.149, and -0.140 respectively, and came out best among four compared methods at every scale, with the advantage growing from 100M to the larger sizes. Head count is not a free hyperparameter: validation loss is U-shaped in H, with a flat optimum at H = 4 or H = 8 across scales; pushing further to H = 16 gives back part of the gain. A direct probe of the trained routing queries confirmed that the learned disagreement between subspaces is what actually drives the improvement.

On the systems side, fused Triton kernels for the routing step raised attention-residual training throughput from 0.2-0.5x of a standard Transformer baseline up to 0.55-0.88x, while keeping peak memory close to baseline. Finally, an identity-preserving conversion method based on delta attention residuals let the authors carry MHAR into mid-training of an 8B-parameter model, where it produced gains of +3.2 points on GSM8K and +3.1 points on GPQA.

Key facts

MHAR reshapes the single routing query of attention residuals into H per-subspace heads, each with its own softmax over the depth history, adding no parameters and negligible compute.
At 100M, 350M, and 1B parameters, MHAR beat a standard Transformer's validation loss by -0.061, -0.149, and -0.140, ranking best of four compared methods at every scale.
Validation loss is U-shaped in head count H, with a flat optimum at H = 4 or H = 8; H = 16 gives back part of the gain.
Fused Triton routing kernels raised attention-residual training throughput from 0.2-0.5x to 0.55-0.88x of baseline, with near-baseline peak memory.
An identity-preserving delta-attention-residual conversion applied MHAR to 8B mid-training, gaining +3.2 on GSM8K and +3.1 on GPQA.

Why it matters

Standard transformers pass depth information through a single additive stream, so each layer only sees the immediately preceding state. Attention residuals already improved on this by letting a sublayer attend across the whole depth history, but they still force every feature subspace to share one attention distribution over that history. MHAR shows that this shared-query bottleneck has a real, measurable cost, and that removing it, by giving each subspace its own head over the depth history, delivers a bigger validation-loss improvement as models get wider, essentially for free in parameter count and compute.

Who it affects

Primarily researchers and engineers who design or train transformer architectures, particularly at the pretraining stage, since MHAR is an architecture-level change rather than something applied after the fact. Anyone maintaining or extending attention-residual variants, or evaluating alternative residual-stream designs, is the direct audience. The 8B mid-training result also signals relevance to teams doing continued pretraining of larger models.

How to use it

MHAR is a drop-in reshaping of the routing query in attention residuals into H heads; H = 1 exactly recovers standard attention residuals, so it is a strict generalization of the earlier method. The paper recommends H = 8 for large-scale models, since going to H = 16 loses part of the gain. Applying it to an already-trained model rather than training from scratch is possible through the paper's identity-preserving delta-attention-residual conversion, used to inject MHAR into 8B mid-training without retraining from zero. Fused Triton kernels are provided to close most of the throughput gap against a standard Transformer baseline.

How solid is it

The method is tested at three pretraining scales (100M, 350M, 1B) against a standard Transformer and two other compared methods, with consistent wins and a growing advantage at larger scale, plus a direct probe of the learned queries that confirms the claimed mechanism, subspace disagreement, rather than reporting a loss number alone. The 8B mid-training result adds a larger data point with concrete downstream benchmark gains on GSM8K and GPQA rather than validation loss alone.

Risks and caveats

The core pretraining experiments stop at 1B parameters; the only larger-scale evidence is a mid-training conversion at 8B rather than full pretraining from scratch, so it is unclear whether the gains hold at the scale of today's largest deployed models. The paper's own throughput numbers show attention residuals, with or without MHAR, still trail a standard Transformer baseline in raw throughput (0.55-0.88x) even after kernel optimization, so the architecture change is not free at the systems level. This is architecture research on a specialized STEM- and code-heavy corpus; it does not evaluate broader task diversity or instruction-tuned, chat-style performance.