Do multi-stream attention residuals improve large-scale language model pre-training?

MSAR turns depth retrieval into multi-head attention

The paper starts from a known weakness of deep PreNorm Transformers: the residual stream carries every earlier computation along one hidden-state trajectory with fixed additive updates. As layers accumulate, features formed early can be diluted, so later layers struggle to retrieve them. Attention Residuals (AttnRes) addressed this by replacing fixed accumulation with learned softmax retrieval over earlier layer outputs, which the authors describe as single-head softmax attention along the depth axis.

MSAR keeps that retrieval but makes it multi-head. Each stream reads the same historical representations through a different sequence-context view, produced by multi-scale causal convolutions, and the retrieved streams are fused by lightweight adaptive mixing into one layer input. The design borrows the motivation of multi-head attention along the sequence axis: several heads can select different sources instead of one shared depth mixture.

The paper reports why two simpler routes failed. Splitting the hidden state into channels for separate heads degraded performance in AttnRes, because the retrieved value is the hidden state itself and slicing it breaks its coherence. Deriving several candidate streams from one layer output via (m)HC-style write-back keeps each candidate full-dimensional, but normalization largely removes the magnitude differences, leaving nearly identical keys and little source diversity.

Stream diversity comes from convolutions, not extra storage

The mechanism matters because it separates expressiveness from cost. MSAR does not keep several persistent residual streams. It retains single-stream sources and expands them on demand through the convolution banks, so multi-stream computation does not require persistent multi-stream storage or communication.

In the block variant, each source is expanded before scoring and value aggregation, each stream is normalized independently over the available depth sources, and the retrieved states are fused with stream weights. The paper states that Block MSAR matches Block AttnRes in persistent source storage and pipeline communication under the same block partition and pipeline schedule.

The overhead is not zero, and the paper quantifies it. At the 18B and 28B scales, MSAR adds 0.31% and 0.27% of Vanilla's activated parameters and 0.51% and 0.48% of its total training FLOPs. For comparison, mHC adds 0.67% and 0.58% of parameters and 0.66% and 0.60% of FLOPs, while AttnRes adds 0.01% of parameters and 0.10% of FLOPs at both scales.

All 12 downstream benchmarks favor MSAR

The main comparison uses 18B and 28B mixture-of-experts models with matched backbones and training-token budgets. The abstract reports that MSAR consistently outperforms standard residual connections, manifold-constrained Hyper-Connections (mHC), and AttnRes across all 12 downstream benchmarks at both scales.

The 18B training curves in the paper show the direction of the gap: final language-model loss is 1.799 for Vanilla, 1.776 for mHC, 1.769 for AttnRes, and 1.750 for MSAR. Average downstream score at that scale is 37.5 for Vanilla, 41.7 for mHC, 41.8 for AttnRes, and 45.7 for MSAR.

Those two numbers answer different questions. Training loss measures next-token prediction on the training objective, while the average downstream score aggregates task performance across the 12 benchmarks. The paper's benchmark suite mixes likelihood-based evaluation, where the candidate answer with the lowest conditional perplexity is selected, with generation-based evaluation using task-specific answer scoring and execution-based pass@1 for code. A gain on one does not mechanically imply the same gain on the other, which is why the paper reports both.

The 1.58× speedup is a fitted scaling-law ratio

The headline efficiency number comes from scaling-law fits, not from a single training run. The authors trained Vanilla, mHC, AttnRes, and MSAR at four model scales with matched backbones and token budgets per scale, then fit a loss-versus-compute curve for each method using its final training losses and normalized compute budgets.

The iso-loss compute speedup is defined as the ratio of the fitted Vanilla compute at each observed loss to the corresponding method's normalized training compute, averaged across the four model scales. On that definition, MSAR reaches 1.58× over the standard residual baseline, compared with 1.35× for AttnRes and 1.30× for mHC.

Readers should keep the reference point fixed: 1.58× is relative to Vanilla, not to AttnRes. The gap between MSAR and AttnRes on this metric is the difference between 1.58× and 1.35×, and the paper's own framing treats all three as improvements over the same standard residual baseline. The fit fixes one curve parameter at 0.72 for all methods and estimates the remaining two by regression, so the speedup is a model-based summary of four scales rather than a measured wall-clock result.

Streams learn different depth sources

The paper checks whether the multiple heads actually diversify. Visualizing depth-attention weights of the 10B Block MSAR model, which has 15 Transformer layers, four streams, and eight blocks with the embedding kept as an extra source, it finds that streams within the same sublayer type develop distinct source preferences.

In Pre-MLP reads beyond the first block, attention to the embedding concentrates in Stream 2 (Conv-4) and Stream 3 (Conv-8), while Stream 1 (identity) and Stream 4 (Conv-12) assign it relatively little weight. The streams also differ in which blocks they repeatedly access: Stream 2 and Stream 4 show strong vertical bands at block 6, whereas Stream 3 prefers block 7.

The paper also separates sublayer types. Pre-Attn reads spread weight more broadly across depth and retain substantial access to the embedding and earlier blocks, while Pre-MLP reads concentrate on fewer sources, producing sharper near-diagonal and vertical structures. This is evidence that MSAR learns diverse source-selection distributions rather than repeating one depth mixture, though it is a descriptive visualization averaged over tokens from several 2,048-token Pile sequences, not a causal test of which stream drives the gains.

Where the conclusion holds

The evidence covers pre-training of mixture-of-experts language models at 18B and 28B total-parameter scales, plus scaling-law runs at four smaller configurations from 2.21B to 17.40B total parameters. All scaling-law models use one leading dense Transformer layer, 144 routed experts with top-8 routing, one shared expert, and a sequence length of 8192; mHC and MSAR use four streams.

The paper does not report results outside these scales, architectures, or training settings, and the downstream comparison is stated at the 18B and 28B scales. The efficiency claim rests on fitted curves and estimated FLOPs rather than measured throughput, and the memory-access analysis is a leading-order cost model per token per residual sublayer.

Within that scope, the paper's own conclusion is that the benefits of multi-state representations and cross-layer access are not fully overlapping, which is why it argues for residual designs combining both. The practical priority it supports is specific: when choosing among these residual variants for large-scale pre-training, MSAR carries a larger parameter and FLOP overhead than AttnRes but returns a higher fitted iso-loss speedup and better downstream scores in the tested range.

Sources used in this answer