Multi-Head Attention Residuals

TL;DR

Multi-Head Attention Residuals (MHAR) extends Moonshot’s Attention Residuals (AttnRes) — softmax attention across the depth dimension in place of unit-weight residual addition — by making the routing query itself multi-head. Instead of one shared query per layer that reads the depth history through a single softmax distribution, MHAR reshapes the query into H per-subspace heads, each with its own softmax over previous layers, so the depth-read becomes block-diagonal and different feature subspaces can attend to different layers. The reshape adds zero parameters and negligible compute, and H = 1 recovers standard AttnRes exactly. Trained from scratch on a Nemotron-based anneal corpus at 100M / 350M / 1B scales, MHAR improves validation loss by 0.061 / 0.149 / 0.140 nats over a standard Transformer — best of four methods at every scale — with a U-shaped optimum at H = 4 or H = 8 and consistent regression at H = 16. Fused Triton routing kernels lift attention-residual training throughput from 0.2–0.5× to 0.55–0.88× of a plain Transformer baseline; a delta-attention-residual variant enables identity-preserving 8B mid-training conversion that adds +3.2 GSM8K and +3.1 GPQA.

Key claims

Standard Transformers propagate information only through the most recent residual state; attention residuals let each sublayer softmax-attend over the full depth history, but the routing query is a single vector shared across the whole feature width — every subspace must read depth through one distribution [§Abstract].
The cost of this forced-single-distribution compromise grows with how much feature subspaces disagree about which prior layers to read, and disagreement grows with model width [§Abstract].
MHAR reshapes the depth-routing query into H per-subspace heads, each with its own softmax over the depth history; the depth read becomes block-diagonal, the reshape adds zero parameters and negligible compute, and H = 1 exactly recovers attention residuals [§Abstract].
At Nemotron-anneal from-scratch training on quality-filtered, STEM- and code-heavy data, MHAR is the best of four methods (standard Transformer, AttnRes-H1, MHAR, one other) at 100M, 350M, and 1B, with validation-loss deltas of −0.061 / −0.149 / −0.140 nats over the standard Transformer; the gain grows from 100M to the larger scales [§Abstract, Experiments].
The head count is a genuine design axis rather than a free knob: validation loss is U-shaped in H, with a flat optimum at H = 4 or H = 8 across scales; over-splitting to H = 16 consistently gives back part of the gain [§Abstract].
A direct probe of the trained queries confirms that learned subspace disagreement is the underlying driver of the H > 1 gain — the mechanism the paper proposes is the mechanism the training actually uses [§Abstract, Analysis].
Fused Triton routing kernels raise attention-residual training throughput from 0.2–0.5× to 0.55–0.88× of the plain Transformer baseline while maintaining near-baseline peak memory [§Abstract, Systems].
A delta attention residuals conversion is identity-preserving — it lets an existing standard Transformer be mid-trained into a MHAR model without loss discontinuity — and delivers +3.2 GSM8K and +3.1 GPQA at 8B scale [§Abstract, Mid-Training].
The authors adopt H = 8 as the default for large-scale runs and use H = 8 in the 8B mid-training experiment [§Abstract].

Method

The mechanism sits on top of Attention Residuals (AttnRes; Attention Residuals (AttnRes / Block AttnRes) — Technical Report). At each layer l, AttnRes replaces the unit-weight residual sum with h_l = Σ_i α_{i→l} · v_i where α_{i→l} = softmax(w_l^T RMSNorm(k_i)) is a per-layer softmax attention over prior layer outputs, with a single learned pseudo-query vector w_l ∈ R^d. MHAR observes that this vector is shared across the entire feature width — every d-dimensional subspace of the residual stream reads the depth history through the same distribution, forcing a single compromise across subspaces.

MHAR reshapes w_l into H per-subspace heads, giving H queries w_l^{(1)}, …, w_l^{(H)}, each of which runs its own softmax over the depth history to produce its own attention weights. The depth-read of a d-dimensional residual value is then block-partitioned into H slices, each mixed with its head’s weights — a block-diagonal routing where head h reads layers according to distribution α^{(h)}. Because the total number of query parameters is unchanged (a single length-d vector, viewed as H chunks of length d/H), MHAR adds zero parameters over AttnRes. Compute overhead is negligible; the only extra work is H softmaxes instead of one, over the same depth-history length.

H = 1 collapses back to AttnRes. H = d would give one softmax per feature dimension (independent per-channel routing) but is not run — the U-shape says the optimum is a few heads (4 or 8), not fine-grained per-channel routing, suggesting subspaces disagree in coarse groups rather than individually. The Triton kernels fuse the multi-head softmax with the depth-history gather-and-multiply so that MHAR’s block-diagonal structure translates into good arithmetic intensity on H100/H200.

The delta-attention-residual variant is what makes 8B mid-training practical: it parameterizes MHAR as a delta over the standard residual sum so that at initialization the network is exactly the pretrained Transformer, then trains the MHAR queries + routing kernels while the base weights adapt. This preserves the identity mapping at init (the AttnRes zero-init trick, generalized), and the 8B experiment reports +3.2 GSM8K / +3.1 GPQA from that mid-training pass.

Results

Scaling-from-scratch experiments run four methods on a Nemotron-based anneal corpus (deduplicated, quality-filtered, STEM- and code-heavy) at 100M, 350M, and 1B parameters. MHAR is the best method at every scale; validation-loss improvements over the standard Transformer are −0.061 at 100M, −0.149 at 350M, and −0.140 at 1B [§Abstract]. The gap between MHAR and the runner-up widens from 100M to the larger scales, matching the paper’s mechanism: subspace disagreement is a width-driven effect, so the payoff of splitting the query grows with model size.

The head-count sweep shows a flat U-shaped optimum at H = 4 or H = 8, with H = 16 consistently regressing part of the gain — a real design axis, not a free hyperparameter [§Abstract]. Query probes confirm the trained per-head query vectors disagree on which prior layers to attend to, and the disagreement grows with scale, giving a mechanism-level explanation for the improvement.

Systems: Fused Triton routing kernels move attention-residual training throughput from 0.2–0.5× of baseline (naive PyTorch) to 0.55–0.88× of baseline — a 2–4× speedup relative to the naive kernel — while keeping peak memory near baseline [§Abstract]. This is the same “small architectural tweak needs infrastructure to ship” story as the AttnRes systems section.

Mid-training: at 8B parameters, the delta-attention-residual conversion yields +3.2 on GSM8K and +3.1 on GPQA over the pretrained Transformer [§Abstract]. The identity-preserving initialization means these gains come from targeted mid-training rather than a full retrain, which is a smaller compute footprint than pretraining a fresh MHAR model at 8B.

Why it’s interesting

MHAR is the direct multi-head lift of Attention Residuals (AttnRes / Block AttnRes) — Technical Report — the same conceptual move Multi-Head Attention was over single-head attention over the sequence dimension, now applied to attention over the depth dimension. Where AttnRes framed itself as “softmax attention over depth, matching the linear-to-softmax transition attention did over sequence length,” MHAR completes the sequence-side parallel: the next design axis after softmax-routing is multi-head softmax-routing. This is a clean instance of a general pattern to track — depth-side architectural tricks lag their sequence-side cousins by roughly a decade, and the sequence-side history (single-head → multi-head → GQA/MLA → sparse) suggests the AttnRes lineage will continue to compound.

The finding that a shared routing query forces feature subspaces into a compromise is also a specific data point on the wider “attention micro-architecture surgery” line already mapped by Training stability at scale. Interleaved Head Attention argued the sequence-side attention block wastes capacity by producing only one pattern per head, and factorized Q/K/V across heads to induce more patterns; MHAR argues the depth-side attention block wastes capacity by producing only one distribution across width, and factorizes the query across width to induce more distributions. Same recipe (“find the load-bearing single-distribution bottleneck, factorize it, keep the kernel FlashAttention-compatible”), applied at the orthogonal axis. Neither paper has run the stacked ablation (MHAR + IHA + XSA), which is exactly the composition experiment the concept page flags as open. For Luma the practical hook is the delta-attention-residual mid-training: if MHAR can be added to a pretrained model without loss discontinuity for +3 points on reasoning benchmarks, it is a low-risk upgrade path for any existing dense stack, not a “retrain from scratch” investment.