Multi-Stream Attention Residuals
Xiaohan Qin, Xiangdong Zhang, Tuo Dai, Sunan Zou, Huaijin Wu, Yebin Yang, Zhuo Xia, Lin Yao, Yuliang Liu, Junchi Yan
Multi-Stream Attention Residuals (MSAR) generalizes depth-wise attention by using multi-scale convolutions to create diverse source views.
How can we improve Transformer depth-mixing by replacing single-stream residual connections with multi-stream attention across layers?
In deep Transformer architectures, the standard residual stream aggregates all preceding layer outputs into a single hidden state, causing early features to dilute and become difficult for later layers to retrieve. MSAR addresses this by expanding each residual source into multiple streams using lightweight causal convolutions, allowing each stream to perform independent, content-dependent depth attention over different contextual scales. This approach consistently outperforms standard residuals and prior depth-attention methods, achieving an iso-loss compute speedup of 1.58× over the standard baseline in large-scale pre-training.
Paper Primer
MSAR generalizes the Attention Residuals (AttnRes) mechanism by transforming residual retrieval into a multi-head depth-attention primitive. It constructs multiple full-dimensional streams from a single layer output using causal depthwise convolutions, which encode different local-context scales, and fuses them via an adaptive mixer to preserve the integrity of the hidden representation.
MSAR significantly improves compute efficiency and downstream performance compared to standard residual connections and prior depth-attention variants.
In 18B and 28B MoE pre-training, MSAR consistently achieved the lowest training loss and higher average scores across 12 benchmarks, including MMLU, ARC-C, and BBH. MSAR achieves an iso-loss compute speedup of 1.58× relative to the standard residual baseline, compared to 1.35× for AttnRes and 1.30× for mHC.
To maintain efficiency, MSAR uses a two-phase computation strategy and fused kernels that expand streams on-demand. This design avoids persistent multi-stream storage, reducing pipeline communication and prefill activation storage to O(Nd), matching the cost of single-stream AttnRes.
Why does MSAR use convolutions instead of simply splitting the hidden state into multiple heads?
Splitting the hidden state into low-dimensional fragments degrades performance because the retrieved value must remain a coherent, full-dimensional representation. MSAR uses convolutions to derive diverse contextual views from the full-dimensional source, preserving its integrity while enabling head diversity.
How does MSAR differ from Hyper-Connections (HC) in its approach to multi-stream residual design?
HC-family methods maintain multiple persistent residual streams that evolve through recurrent mixing. MSAR retains a single-stream residual state and only expands it into multiple views during the depth-attention computation, significantly reducing memory and communication overhead.
MSAR demonstrates that enriching depth-wise residual retrieval with multi-stream sequence-context views is a highly effective, low-overhead strategy for improving depth scaling in large language models.
Introduction and Motivation
We expose why single‑stream residuals limit depth scaling and preview a multi‑stream solution.
Standard residual connections force a single‑stream information flow across layers, which limits the model’s ability to retrieve and reuse early‑layer features as depth grows.
Attention Residuals (AttnRes) replaces the fixed additive update with a depth‑wise softmax retrieval, and Multi‑Head Connection (mHC) expands the residual stream into multiple parallel paths, but both either keep a single‑stream view or suffer from reduced diversity when naively split.
Multi‑Stream Attention Residuals (MSAR) generalizes AttnRes by generating multiple full‑dimensional streams via low‑cost multi‑scale causal convolutions, allowing each head to attend to depth with distinct sequence‑context views while preserving hidden‑state integrity.
The limitation of single‑stream residual paths is the key obstacle to efficient depth scaling.
Residuals and Prior Baselines
We define baseline depth‑mixing mechanisms and their intuition.
This section introduces two baseline depth‑mixing mechanisms that the paper later compares against: Attention Residuals (AttnRes) and Multi‑Head Connection (mHC). Both operate on the residual branches of a Transformer stack, but they differ in how they aggregate information across depth.
AttnRes replaces the simple additive accumulation of previous layers with a softmax‑weighted sum, letting each layer choose how much of each earlier output to reuse based on content similarity.
mHC expands the residual state into $m$ parallel streams, allowing each layer to read from and write to a richer set of depth histories while keeping signal propagation stable via doubly‑stochastic mixing.
The MSAR Mechanism
MSAR expands residual sources into multiple contextual streams and fuses them via learned attention.
Naïvely giving each depth‑attention head its own dense projection would blow up parameters and FLOPs. The paper therefore seeks a cheap way to generate multiple depth‑wise views of a residual source. By re‑using nearby hidden states with causal depthwise convolutions, MSAR builds diverse streams without extra cost.
Instead of a single residual vector, MSAR creates a handful of cheap contextual copies, lets each copy attend over depth independently, and then merges them with a learned scalar mixer.
Stream 0 is the identity: $\mathbf{v}^{(0)}_{t} = \mathbf{v}_{t}$.
Stream 1 uses a 4‑token causal kernel $[0.5,0.3,0.2,0]$: $\mathbf{v}^{(1)}_{3} = 0.5\mathbf{v}_{3}+0.3\mathbf{v}_{2}+0.2\mathbf{v}_{1}$.
Stream 2 uses an 8‑token kernel (zero‑padded for early positions) yielding $\mathbf{v}^{(2)}_{4} = 0.4\mathbf{v}_{4}+0.3\mathbf{v}_{3}+0.2\mathbf{v}_{2}+0.1\mathbf{v}_{1}$.
Stream 3 uses a 12‑token kernel, which for this short sequence collapses to the same weighted sum as Stream 2 but with a broader receptive field.
Depth‑attention weights $\alpha^{(r)}$ are learned scalars; assume they become $[0.4,0.2,0.3,0.1]$ after training.
Fusion computes $\mathbf{h}_\ell = \sum_{r} \alpha^{(r)} \mathbf{v}^{(r)}_\ell$, e.g. for $t=3$: $\mathbf{h}_3 = 0.4\mathbf{v}^{(0)}_{3}+0.2\mathbf{v}^{(1)}_{3}+0.3\mathbf{v}^{(2)}_{3}+0.1\mathbf{v}^{(3)}_{3}$.
Even with tiny kernels, the streams capture complementary context ranges, and the learned scalar mixer can emphasize the most useful view without adding expensive parameters.
How does MSAR differ from the AttnRes baseline that also uses depth attention?
AttnRes applies a single softmax over depth for the whole layer, sharing the same query across all positions. MSAR, by contrast, first creates $m$ distinct streams via causal convolutions, giving each stream its own pseudo‑query and attention weights. This lets the model attend to depth at multiple contextual scales and then fuse the results, whereas AttnRes can only attend uniformly.
**Figure 2.** **Comparison of residual architectures.** (a) Standard Residuals; (b) Block AttnRes; (c) Block MSAR. Block MSAR expands single-stream sources into contextual streams for per-stream depth attention, then fuses the retrieved streams into a single input for each residual branch.
Block-Wise MSAR Implementation
Block MSAR swaps per‑layer hidden states for compact block summaries to cut memory and communication.
Retaining every layer’s hidden state forces a memory footprint that scales with depth, which becomes prohibitive in large‑scale training. Block MSAR tackles this by collapsing each contiguous group of layers into a single summary that serves as the depth‑attention source.
Instead of keeping a separate hidden vector for every layer, we keep one summed vector per block and let each layer attend to that block‑level summary plus its own partial accumulation.
Block 1 summary: $\mathbf{b}_1 = 1+2+3 = 6$.
Block 2 summary: $\mathbf{b}_2 = 4+5+6 = 15$.
Partial sum before layer 2: $\mathbf{b}_{<2}^{\,1}=1$; before layer 3: $\mathbf{b}_{<3}^{\,1}=1+2=3$.
When processing layer 2, the attention source consists of completed summary $\mathbf{b}_1$ (which is still being built) and the partial $\mathbf{b}_{<2}^{\,1}=1$.
After layer 3 finishes, the block summary $\mathbf{b}_1$ becomes 6 and is used by all layers in block 2.
Block MSAR stores only one vector per block (6 and 15) instead of six separate hidden states, cutting the per‑token storage by a factor of $L/N = 3$ while still providing each layer with a fresh partial contribution.
Pre‑compute pseudo‑queries for every target layer inside the block.
Batch all queries against the completed block summaries $\{\mathbf{b}_j\}_{j< n}$ to obtain per‑head retrieval results.
For each layer, compute its current partial sum $\mathbf{b}_{<\ell}^{\,n}$ on the fly.
Merge the partial contribution with the pre‑computed retrieval using an online softmax (Eq. 17).
Apply the causal‑convolution expansion to the merged result, producing the multi‑stream depth‑attention inputs for the layer.
Because only the block summaries are kept in memory, the persistent source traffic per token drops from $\mathcal{O}(mNd)$ to $\mathcal{O}(Nd)$, matching the baseline Block AttnRes while still supporting full depth‑mixing.
Experimental Setup
MSAR achieves top downstream scores, notably 55.7% on MMLU (18B).
MSAR outperforms all baselines on the MMLU benchmark at the 18B scale.
Table 1 shows MSAR scores 55.7% versus 48.9% (Vanilla), 21.1% (mHC), and 46.4% (AttnRes).
Performance Evaluation
MSAR outperforms baselines on MoE models across scales.
MSAR consistently outperforms the mHC, AttnRes, and Vanilla baselines on all benchmarks.
Table 1 shows average scores of 45.7 (18B) and 52.4 (28B) versus lower scores for the other methods.
**Table 1.** Downstream performance comparison of MSAR with Vanilla, mHC, and AttnRes on 18B and 28B MoE models. Scores are reported in percent (higher is better), with the best result for each benchmark at each scale in bold.
**Figure 3.** Scaling-law comparison of MSAR, AttnRes, mHC, and Vanilla. MSAR substantially outperforms all three baselines in the loss-compute trade-off.
Adopting MSAR yields up to 8 points higher downstream performance with negligible parameter and compute overhead.
Scaling Laws and Ablations
MSAR consistently lowers loss and improves compute efficiency across scales and design choices.
We first examine how MSAR’s compute‑efficiency scales with model size and then dissect which architectural components drive its gains.
Training loss typically follows a power‑law decay with compute; a flatter curve means you get more performance per FLOP.
Why does a higher $\beta$ translate into a speedup rather than just a lower loss?
$\beta$ determines the slope of the loss‑vs‑compute curve; a steeper slope means you reach a target loss with fewer FLOPs, which is reported as an iso‑loss compute speedup.
MSAR reduces training loss by 2.7 % relative to Vanilla at the largest model scale.
Loss 1.843 vs 1.898 for Vanilla (2.7 % reduction).
MSAR achieves a 1.58× iso‑loss compute speedup over Vanilla.
Average speedup across four scales: 1.58× vs 1.35× for AttnRes and 1.30× for mHC.
Parameter and FLOPs Overhead
Quantifies parameter, FLOPs, and memory overheads of residual mechanisms.
This appendix enumerates the extra parameters, training FLOPs, and memory traffic introduced by each residual mechanism.
The per‑token memory‑access cost of a residual sublayer follows the leading‑order term (N/S + 4 m + 5) d, where N is the number of blocks, S the number of residual sublayers per block, m the stream count, and d the hidden dimension.
**Figure 4.** Token-averaged depth-attention weights of the 10B Block MSAR model. Rows show Pre-Attn (top) and Pre-MLP (bottom); columns show the four streams, numbered 1–4. Source 0 is the embedding, and sources 1–8 are blocks, with the current block represented by its nonempty partial sum. Hatched regions mark unavailable sources.
Learned Depth-Attention Patterns
We examine depth‑attention distribution in the 10B Block MSAR model and the memory cost of each residual mechanism.
The model learns to route information across the depth dimension, with different sublayers showing distinct access styles.
How does this depth‑attention differ from the standard residual connection?
Standard residual connections simply add the previous layer’s output, whereas depth‑attention can selectively pull from any earlier block, enabling content‑dependent reuse of distant representations.
Vanilla residual merge incurs the highest memory I/O per token.
Table 6 reports 34 d per token for the vanilla mechanism.
mHC reduces memory I/O to 5.5 d per token.
Table 6 shows 5.5 d for the mHC mechanism.
Block AttnRes requires 21.5 d per token.
Table 6 lists 21.5 d for Block AttnRes.
Block MSAR’s depth‑mixing costs (8 m + 2) d per token.
Table 6 gives (8 m + 2) d for the Block MSAR mechanism.
Residual Mechanism Comparison
We compare residual update mechanisms and highlight stream‑specific depth‑mixing patterns.
This section surveys residual update mechanisms and highlights how MSAR induces stream‑specific depth‑mixing.
**Table.** Comparison of different recurrence and cross-layer access methods. The table categorizes methods into "Single-state recurrence", "Multi-state recurrence", "Cross-layer access", and "Multi-stream cross-layer access", detailing the update rule, weight type (Fixed, Static, Dynamic), and source for each method.
MSAR learns stream‑specific source‑selection distributions: Stream 2 (Conv‑4) and Stream 3 (Conv‑8) heavily attend to the embedding, while Streams 1 and 4 rely more on identity or later blocks, producing distinct vertical attention patterns.
Residual Mechanism Formulations
We compare diverse residual mechanisms and formalize their update rules and mixing weights.
Table 7 enumerates representative residual mechanisms, showing each method’s update rule, mixing weight formulation, and the sources it can draw from.
Experimental Details
Details of model configurations, training hyperparameters, residual setups, and evaluation procedures.
We train three MoE models—10 B for ablations, and 18 B and 28 B for the main comparisons—using the DeepSeekMoE architecture. Each backbone consists of a leading dense Transformer layer followed by MoE layers, grouped‑query attention, and SwiGLU feed‑forward networks. Training hyperparameters include sequence length 8192, vocabulary size 152 064, AdamW optimizer ($\beta_1=0.9$, $\beta_2=0.95$, $\epsilon=10^{-15}$), base learning rate $3.5\times10^{-4}$ with a warmup‑stable‑decay schedule, weight decay 0.1, and gradient clipping 1.0.
The mHC baseline uses four residual streams ($m=4$), 20 Sinkhorn–Knopp iterations, and gating scales initialized to 0.01. For the 18 B and 28 B experiments, both AttnRes and MSAR group four Transformer layers per block, yielding $S=8$ residual sublayers and $N=7$ or $N=8$ blocks respectively, with the embedding counted separately. MSAR also employs four streams, initializes its pre‑gating scale $\alpha_{\text{pre}}$ to 0.01, adds three causal depthwise convolution branches with kernel sizes $\{4,8,12\}$ (weights zero‑initialized), and sets learned pseudo‑queries to zero.
We evaluate using the OpenCompass suite, applying likelihood‑based selection of the lowest‑perplexity answer on MMLU, MMLU‑Redux, CMMLU, C‑Eval, ARC‑C, and C3 benchmarks. Generation‑based metrics are used for MMLU‑Pro, BBH, GSM8K, MATH, HumanEval, and LCBench, with task‑specific scoring for reasoning and pass@1 for code. All evaluations follow the few‑shot settings reported in the original benchmarks.
**Table.** Residual Configurations.
Scaling-Law Experiment Details
Experimental configurations and fitting details for the scaling‑law study.
We evaluate Vanilla, mHC, AttnRes, and MSAR at four model scales with matched backbones and training‑token budgets.
All models use one leading dense Transformer layer, 144 routed experts with top‑8 routing, one shared expert, and a sequence length of $8192$; common optimizer settings follow Appendix A.
mHC and MSAR use $m=4$ streams, while AttnRes and MSAR share the block sizes listed in Table 4, yielding $8$, $8$, $8$, and $7$ blocks across the four scales, with the last block possibly shorter and the embedding excluded from these block counts.
The table presents architectural and training hyperparameters for four different model scales. Columns include Total Parameters, Activated Parameters, Layers, Hidden Dimension, Dense FFN, Expert FFN, Q/KV Heads, Training Budget, Global Batch Size, Learning Rate, and Layers per Block (AttnRes/MSAR). A footnote indicates that backbone parameter counts exclude input embeddings.
We fit Equation 18 independently for each method using its final training loss and the corresponding normalized compute budget across the four model scales.
With $E=0.72$ fixed for all methods, we estimate $A$ and $\beta$ by ordinary least‑squares regression on $\log_2 C$.
The iso‑loss compute speedup is defined as the ratio of the fitted Vanilla compute at each observed loss to the method’s normalized training compute, averaged over the four model scales.
Depth Mixing Visualization
Appendix E details how structured matrices implement depth mixing across residual sublayers.
Figure 5 shows depth mixing for $L=4$ residual sublayers with block size $S=2$, where rows are target inputs $h_1,\dots,h_4$ and columns are sources $v_0,\dots,v_3$. Vanilla attention assigns unit weight to every preceding source, while mHC unrolls the recurrence using the product $A\times a\!\to\!b$ and doubly stochastic transition matrices. The figure also illustrates how Block AttnRes and Block MSAR restructure the score computation.
Groups two adjacent source vectors into a single pseudo‑source, so the block contributes one combined score instead of separate scores per column.
Creates a separate score matrix for each stream, normalizes each independently, and then merges the retrieved streams.
Matrix Representation Details
Detailed matrix forms and normalization steps for depth‑mixing mechanisms.
The appendix begins by enumerating the four residual entries (r) 1–4 together with their associated embeddings $E_r[\cdot]$.
Connections from stream 1 to streams 2, 3, and 4 are explicitly listed as 1 → 2, 1 → 3, 1 → 4.
Summation indices are set to $m-1$ and $r=0$, establishing the range of streams and the base residual level for the calculations.
Finally, the per‑stream results are fused to produce the overall depth‑mixed representation used by the model.