PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Stream Attention Residuals

Xiaohan Qin, Xiangdong Zhang, Tuo Dai, Sunan Zou, Huaijin Wu, Yebin Yang, Zhuo Xia, Lin Yao, Yuliang Liu, Junchi Yan

arXiv:2610.13630v1cs.CLcs.CV

TL;DR

Depth scaling can be limited by diluted early-layer features, while existing approaches do not provide diverse depth-attention heads without compromising hidden-state integrity. MSAR uses multi-scale causal convolutions to create contextual streams for independent depth retrieval and lightweight mixing. It consistently outperforms prior residual variants across evaluated scales and reports a 1.58× iso-loss compute speedup over the standard residual baseline.

  • Problem

    Depth-wise information dilution can make additional layers less efficient, and constructing diverse depth-attention heads from one layer output without fragmenting its representation remains a design challenge.

    Source: §1 · pp. 1–2

  • Method

    MSAR uses multi-scale causal depthwise convolutions to create contextual streams, retrieves over depth independently per stream, and fuses retrieved states with lightweight scalar mixing.

    Source: §1 · pp. 2–3

  • Results

    MSAR outperforms prior residual variants on all 12 evaluated benchmarks at both model scales and achieves a 1.58× iso-loss compute speedup over the standard residual baseline.

    Source: §4.2 · p. 9, §5 · p. 11

  • Takeaways & Limitations

    The results support MSAR as an effective residual architecture for large-scale pre-training, with single-stream source storage and on-demand expansion.

    Source: §5 · p. 11

    The source-view construction relies on autoregressive causal masking, under which earlier same-layer hidden states encode available context without exposing future information.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

Attention Residuals (AttnRes) enables content-dependent retrieval over network depth, which can be viewed as single-stream softmax attention along the depth axis. Inspired by the success of multi-head attention along the sequence axis, we introduce Multi-Stream Attention Residuals (MSAR), an extension of AttnRes that generalizes this mechanism to multi-head depth attention, yielding substantial performance gains with minimal additional overhead. MSAR draws stream diversity from sequence context through multi-scale causal convolutions, allowing different heads to retrieve historical representations through distinct contextual views while preserving the integrity of the hidden representation. The retrieved streams are combined through lightweight adaptive mixing to generate the single-stream layer input. By retaining single-stream sources and expanding them on demand, MSAR supports multi-stream computation without requiring persistent multi-stream storage or communication. Across all 12 downstream benchmarks at the 18B and 28B MoE scales, MSAR consistently outperforms prior residual variants, including standard residual connections, manifold-constrained Hyper-Connections (mHC), and AttnRes. Scaling-law experiments yield an iso-loss compute speedup of $1.58\times$ over the standard residual baseline, compared with $1.35\times$ for AttnRes and $1.30\times$ for mHC. These results make MSAR a promising residual architecture for large-scale pre-training.

1. Introduction

Depth-wise residual designs address dilution of early-layer information, but existing multi-stream or single-head retrieval mechanisms leave a challenge: obtaining diverse depth-attention heads without fragmenting hidden representations. MSAR uses sequence-context views to create those heads, and the supplied results report improved downstream performance and compute efficiency.

1. Introduction

  • 1.58× iso-loss compute speedup over Vanilla, versus 1.35× for AttnRes and 1.30× for mHC, accompanies MSAR’s lower 18B final training loss.MSAR reaches 1.750 loss, compared with Vanilla at 1.799, mHC at 1.776, and AttnRes at 1.769.Source: Figure 1 · p. 2, §1 · p. 3
  • MSAR generalizes AttnRes to multi-head depth attention using multi-scale causal depthwise convolutions, per-stream retrieval, and lightweight scalar mixing.The streams provide distinct sequence-context views while keeping the hidden representation full-dimensional.Source: §1 · pp. 2–3
  • MSAR stores single-stream summaries and expands them on demand, reducing communication and prefill activation storage from O(mNd) to O(Nd).The paper reports that this matches AttnRes storage costs with minimal training FLOPs and parameter overhead.Source: §1 · p. 3
  • MSAR delivers consistent downstream gains in both 18B and 28B MoE settings, with scaling-law and ablation studies supporting the design.The introduction attributes the gains to the proposed inductive bias rather than a model-size-specific artifact.Source: §1 · p. 3

2. Related Work

Prior residual designs either maintain multiple persistent states or provide access to earlier layers; MSAR combines explicit cross-layer retrieval with multi-stream contextual views.

2. Related Work

  • Multi-State Residual Recurrence methods maintain multiple residual streams, with HC-family variants differing in mixing, stability, and write-back design.The cited examples include HC [50], mHC [39], Single-Pass, xHC [44], and Gated Residual [28].Source: §2 · p. 3
  • Cross-layer access methods retrieve or aggregate earlier representations, ranging from learned depth weights to input-dependent attention and softmax retrieval.Examples include DenseFormer, MUDDFormer [38], MoDA [51], and AttnRes [35].Source: §2 · pp. 3–4
  • MSAR preserves AttnRes’s explicit cross-layer access while widening the depth-attention primitive through multiple sequence-context views rather than persistent residual recurrence.This places MSAR at the intersection of the two residual-design axes described in the paper.Source: §2 · p. 4

3. Method

The method frames residual connections as depth mixing, then extends depth retrieval with contextual streams and block-level sources. Its on-demand expansion and two-phase execution retain single-stream source storage while supporting multi-stream attention.

3.1. Preliminaries: Residuals as Depth Mixing

  • Standard residuals accumulate branch outputs into one running state, whereas HC methods maintain multiple persistent streams and AttnRes retrieves over earlier sources with normalized depth attention.The paper identifies compression into one state and progressively smaller relative updates as properties of standard residual accumulation.Source: §3.1 · pp. 4–5

3.2. MSAR: Multi-Stream Attention Residuals

  • MSAR expands each source into contextual streams with causal depthwise convolutions, performs depth attention independently per stream, then fuses retrieved streams with lightweight scalar weights.The main setting retains the original source and adds three convolutional streams with kernel sizes {4, 8, 12}.Source: §3.2 · p. 6

3.3. Block MSAR

  • Block MSAR uses completed block summaries and a current partial as retrieval sources, expanding the single-stream representations for per-stream attention before fusion.At a block’s first layer, the current partial source is absent.Source: §3.3 · pp. 7–8, Figure 2 · p. 7

3.4. Efficiency Analysis

  • On-demand expansion keeps persistent storage and pipeline communication at O(Nd), matching Block AttnRes; two-phase retrieval batches source reads and fused expansion avoids intermediate-stream writes.Retrieved multi-stream outputs still incur additional memory traffic for Hpre aggregation.Source: §3.4 · p. 8

4. Experiments

Experiments compare MSAR with prior residual variants across MoE scales and evaluate downstream performance, scaling laws, and design ablations. MSAR leads the downstream benchmarks and improves the loss–compute trade-off.

4.1. Experimental Setup

  • Experiments compare MSAR with mHC, AttnRes, and Vanilla on 18B and 28B MoE models, using a 10B MoE model for ablations across 12 downstream benchmarks.The downstream evaluation spans general capabilities, code and mathematics, and Chinese language understanding.Source: §4.1 · pp. 8–9

4.2. Main Results

  • MSAR outperforms mHC, AttnRes, and Vanilla on all 12 benchmarks at both scales; its 18B average score is 45.7 versus 41.8 for AttnRes.The 28B average is 52.4 for MSAR versus 49.4 for AttnRes.Source: Table 1 · p. 9, §4.2 · p. 9

4.3. Scaling Laws

  • Averaged across four scales, MSAR's iso-loss compute speedup is 1.58× relative to Vanilla, compared with 1.35× for AttnRes and 1.30× for mHC.MSAR's fitted curve remains below the three baselines over their shared compute range.Source: Figure 3 · p. 10, §4.3 · p. 10

4.4. Ablation Studies

  • On the 10B ablation, MSAR reaches Pile loss 1.847 versus 1.898 for Vanilla and 1.861 for AttnRes + mHC, suggesting its gains are not simple state addition.With m=4 and S=4, MSAR lowers loss by a further 0.014 relative to the hybrid.Source: §4.4 · pp. 10–11
  • Removing contextual convolutions raises loss from 1.847 to 1.862, while mean aggregation yields 1.851; Hpre also contributes to training stability.The comparison isolates contextual source expansion from multiple retrieval heads alone.Source: §4.4 · p. 11
  • Two, four, and eight streams yield losses 1.850, 1.847, and 1.846; S=2, 4, and 8 yield 1.845, 1.847, and 1.852 versus 1.843 for Full MSAR.Approximately eight blocks recover most of the benefit of full depth access, with diminishing gains from finer partitioning.Source: §4.4 · p. 11

5. Conclusion and Outlook

MSAR extends residual retrieval to multi-head depth attention using sequence-context-derived streams and adaptive mixing. Experiments report downstream gains and improved loss–compute efficiency, while on-demand stream expansion preserves single-stream storage.

Conclusion and Outlook

  • MSAR extends single-stream residual retrieval to multi-head depth attention, using multi-scale causal convolutions to create contextual source views and adaptive scalar mixing to form the layer input.The streams retrieve independently over depth.Source: §5 · p. 11
  • Block MSAR stores single-stream summaries and expands them on demand, matching Block AttnRes's persistent source storage and pipeline communication costs.It also supports cross-stage caching and two-phase computation.Source: §5 · p. 11
  • Experiments report consistent downstream improvements on 18B and 28B MoE models and a 1.58× iso-loss compute speedup relative to the standard residual baseline.The paper concludes that MSAR improves language-model pre-training quality and compute efficiency.Source: §5 · p. 11

A. Experimental Details

Experiments use a 10B MoE model for ablations and 18B and 28B MoE models for main comparisons. The evaluation suite combines likelihood-based and generation-based benchmarks.

Model Architectures and Training Settings

  • Each Transformer layer contains an attention and an MLP residual sublayer.The reported layer counts refer to Transformer layers.Source: A. Experimental Details · p. 16
  • The 10B MoE model is used for ablations, while the 18B and 28B MoE models support the main results.The backbones follow a DeepSeekMoE-style architecture with a leading dense Transformer layer, MoE layers, grouped-query attention, and SwiGLU feed-forward networks.Source: Table 3 · p. 16, A. Experimental Details · p. 16

Evaluation Protocols

  • Evaluation uses likelihood-based scoring for MMLU, MMLU-Redux, CMMLU, C-Eval, ARC-C, and C3, with task-specific shot counts.Generation-based evaluation covers MMLU-Pro, BBH, GSM8K, MATH, HumanEval, and LCBench, using task-specific answer scoring or execution-based pass@1.Source: A. Experimental Details · pp. 16–17

B. Scaling-Law Experiment Details

Scaling-law comparisons use matched backbones and training-token budgets across model scales. The experiment estimates iso-loss compute speedup by comparing fitted Vanilla compute with each method’s normalized training compute at observed losses.

Scaling-law configurations

Iso-loss compute speedup

  • Iso-loss compute speedup is averaged across four scales as the ratio of fitted Vanilla compute to each method’s normalized training compute at corresponding observed losses.The fit uses final training losses, fixes E=0.72, and estimates A and β by ordinary least squares regression of log2(L − E) on log10 C.Source: B. Scaling-Law Experiment Details · p. 17

C. Parameter, Compute, and Memory Costs

MSAR’s parameter increments are lower than mHC’s and higher than AttnRes’s at the 18B and 28B scales. The section also specifies how residual-module training FLOPs and memory-access costs are estimated or compared.

Parameter Overhead

  • MSAR’s added parameters comprise query and mixer parameters per sublayer plus parameters for its convolution banks.The parameter accounting includes biases and learned scales relative to the standard residual backbone; L counts residual sublayers and d is the hidden dimension.Source: C. Parameter, Compute, and Memory Costs · p. 17
  • MSAR adds 5.204M and 7.209M parameters at the 18B and 28B scales, versus mHC’s 11.356M and 15.730M and AttnRes’s 0.120M and 0.166M.With m=4 and kernel sizes {4, 8, 12}, the MSAR parameter expression becomes L(44d + 5), compared with L(96d + 27) for mHC.Source: C. Parameter, Compute, and Memory Costs · p. 18

Training FLOPs

  • Training-FLOPs overhead estimates include stream expansion, per-stream depth attention, and stream fusion for MSAR.The estimates count a multiply–add as two FLOPs and approximate training cost as three times forward cost, excluding activation recomputation and optimizer updates.Source: C. Parameter, Compute, and Memory Costs · p. 18, Table 5 · p. 18

Storage, Communication, and Memory Access

  • Block MSAR matches Block AttnRes in persistent source storage and pipeline communication under the same block partition and schedule.With fused execution, expanded source streams are consumed immediately rather than materialized; the reported leading-order memory-access cost is (N/S + 4m + 5)d per token per residual sublayer.Source: C. Parameter, Compute, and Memory Costs · p. 18

D. Learned Depth-Attention Patterns

Learned depth-attention patterns differ between Pre-Attn and Pre-MLP reads, and different streams develop distinct preferences for historical sources. Figure 4 displays token-averaged weights for four streams across embedding and block sources.

D. Learned Depth-Attention Patterns

  • The illustrated weights come from a 10B Block MSAR model with 15 Transformer layers, four streams, and eight blocks, averaged over tokens from several 2,048-token Pile sequences.Source: D. Learned Depth-Attention Patterns · p. 18

Sublayer-Specific Access Modes

  • Pre-Attn reads distribute weight more broadly across depth, while Pre-MLP reads concentrate on fewer sources with sharper near-diagonal and vertical structures.The near-diagonal band marks the current partial or most recent completed block; vertical bands mark repeated access to a source across target layers.Source: D. Learned Depth-Attention Patterns · p. 18

Diverse Patterns across Streams

  • Pre-MLP streams differ in source preferences: Streams 2 and 3 emphasize the embedding, while Streams 2 and 4 repeatedly access block 6 and Stream 3 favors block 7.The figure shows token-averaged depth-attention weights for four streams, with Pre-Attn and Pre-MLP reads separated; the embedding is source 0 and blocks are sources 1–8.Source: Figure 4 · p. 19, D. Learned Depth-Attention Patterns · p. 20

E.1. Comparison of Residual Mechanisms

The comparison organizes residual mechanisms by update rules, mixing weights, and accessible sources. MSAR differs from DAR by deriving head diversity from multi-scale sequence-context views rather than maintaining two residual trajectories.

E.1. Comparison of Residual Mechanisms

  • ReZero uses a learned scalar, Highway an input-dependent gate, and DenseFormer learned input-independent coefficients.The comparison covers update rules, mixing weights, and accessible sources.Source: §E.1 · p. 20
  • Block variants access completed summaries and the current partial, while MSAR expands each available source into stream-specific views.The current partial is omitted at a block's first sublayer.Source: §E.1 · p. 21
  • In the DDL row, H_l is its matrix-valued state, k_l and v_l its update key and value, and beta_l controls the delta update.These symbols are local to that row.Source: §E.1 · p. 21
  • HC/mHC use a stream-read vector distinct from scalar coefficients used by ReZero, DenseFormer, and AttnRes; DAR also explores stream-width expansion for depth retrieval.For DAR, the described history states and current partial are two-stream representations.Source: §E.1 · p. 21
  • MSAR derives head diversity from neighboring-position context, whereas DAR maintains two residual trajectories at the same token position.Each MSAR head can select depth histories using a different contextual scale.Source: §E.1 · p. 21

E.2. Residual Connections as Structured Matrices

Figure 5 represents depth mixing as target-to-source relationships, including grouped block sources. Block MSAR instead uses a score matrix for each expanded stream, normalizing and fusing retrieved states across streams.

E.2. Residual Connections as Structured Matrices

  • Figure 5 rows are target inputs and columns are sources; Vanilla assigns unit weights to preceding sources, while mHC transition matrices are doubly stochastic.The figure uses L=4 residual sublayers and block size S=2.Source: §E.2 · p. 21
  • Block AttnRes groups a completed block into one source whose score spans its constituent columns, and normalization counts that grouped source once.Dashed cells mark the evolving partial source.Source: §E.2 · p. 21
  • Block MSAR expands each grouped source before scoring and value aggregation, normalizes separately within each stream, and fuses the retrieved states.Its panel shows per-stream score matrices rather than a scalar mixing matrix over unexpanded sources.Source: §E.2 · pp. 21–22, Figure 5 · p. 22

Related papers

Loading 2610.13630v1…