Multi-Stream Attention Residuals

What it claims

MSAR turns AttnRes's single softmax over earlier layer outputs into four depth-attention heads. Each head sees the same layer output through a different causal-convolution window. In the 18B MoE run, MSAR reaches a final training loss of 1.750, against Vanilla 1.799, mHC 1.776 and AttnRes 1.769. Its average downstream score is 45.7, against 37.5, 41.7 and 41.8. In the 28B MoE setting, it adds only 0.5% training FLOPs and 0.3% activated parameters.

How it works

The design answers two approaches the authors say do not work:

• Channel-splitting the hidden state into heads was already tried in AttnRes and degraded performance, because the retrieved value is a whole hidden state.
• (m)HC-style post write-back produces candidates that are scalar reweightings of one vector. RMSNorm erases those differences, so the keys come out nearly identical.

MSAR instead takes diversity from the sequence axis:

1. Stream expansion. Each depth source v_i (the embedding or a branch output) becomes m = 4 full-dimensional streams. These are the identity plus three layer-specific causal depthwise 1D convolutions with kernel sizes {4, 8, 12}. Positions before the sequence start are zero-padded.
2. Per-stream depth attention. Each stream has its own learned pseudo-query per target layer. It runs its own softmax over all earlier sources, using the AttnRes kernel exp(qᵀ RMSNorm(k)) with key = the source itself.
3. Stream fusion. No md×d output projection is used. An mHC-style H_pre mixer produces sigmoid scalar weights from RMSNorm(flatten(Z)) via an md×m projection. The layer input is the weighted sum of the m retrieved states.
4. The last layer skips expansion and does a plain residual merge into the LM head.

flowchart LR
  S[Stored single-stream source] --> I[Identity]
  S --> C4[Conv k=4]
  S --> C8[Conv k=8]
  S --> C12[Conv k=12]
  I --> A0[Depth softmax 0]
  C4 --> A1[Depth softmax 1]
  C8 --> A2[Depth softmax 2]
  C12 --> A3[Depth softmax 3]
  A0 & A1 & A2 & A3 --> M[Sigmoid scalar mixer]
  M --> L[Layer input]

Systems angle. Every stream is deterministically recomputable from the single main stream, so only single-stream sources (or block summaries, in block mode) are stored. The streams are expanded on demand, and a fused kernel can consume them without writing them back. This cuts pipeline communication and prefill activation storage from O(mNd) to O(Nd), the same as AttnRes. MSAR is stated to be compatible with AttnRes block mode, cache-based pipeline communication and two-phase computation.

Evidence

• 18B MoE training loss at ~46k iterations: figures as above.
• 18B average over 12 benchmarks: MSAR 45.7, AttnRes 41.8, mHC 41.7, Vanilla 37.5. The benchmarks are MMLU, MMLU-Pro, MMLU-Redux, ARC-C, BBH, HumanEval, LCBench, GSM8K, MATH, CMMLU, C-Eval and C3. The authors say MSAR's lead holds throughout training.
• Iso-loss compute speedup from scaling-law runs: 1.58× for MSAR, 1.35× for AttnRes and 1.30× for mHC, each relative to standard residuals.

How to use it

Code is linked at https://github.com/jianke0604/MSAR. The fetched text gives no install commands, weights or licence. The reported configuration is the identity plus kernels {4, 8, 12}, with one pseudo-query per stream per layer.

Limits

• The fetched text stops at Section 3.3 (Block MSAR). The following do not appear in it:
• per-benchmark numbers for either scale
• any 28B results
• details of the scaling-law setup
• the ablations the authors cite as proof that the gains come from the inductive bias
• The authors acknowledge remaining overhead over AttnRes: memory traffic for writing and aggregating the m retrieved stream outputs. No throughput or wall-clock figures are given.
• The storage savings depend on a fused expand-and-consume kernel ("in practice"). Without one, the intermediate streams are materialised.
• The stream diversity comes from causal sequence convolutions, so the method assumes autoregressive causal-mask models. No other setting is discussed.
• All validation is on MoE pre-training at 18B and 28B. Dense models and fine-tuning are not covered in the fetched text.