Residual connections propagate information through a network by additively accumulating layer outputs. Recently proposed alternatives, such as Manifold-Constrained Hyper-Connections (mHC) and Attention Residuals (AttnRes), instead use more expressive depth aggregation mechanisms that improve language modeling performance at scale. In this work, we consider the mixing of representations over depth as a problem of storing and retrieving information across layers. Formalizing this perspective, we introduce Depth Memory (DM), a unified framework that describes the standard pre-layer-norm (PreLN) transformer, mHC, and AttnRes in terms of streams that carry representations through the network. One such stream, for example, could be an explicit memory stream that is queried to enrich the input to every layer, and whose output is subsequently written back into the memory for use by later layers. For AttnRes, this memory stream is a growing cache of explicit layer representations. For mHC, it is an expanded residual stream. Building on this view, we propose Linear Depth Memory (LDM), which formulates information retrieval across depth as a regression problem solved using recursive least squares. The resulting architecture propagates three streams across depth: the standard residual stream, a linear associative memory stream, and an inverse covariance stream. For high-throughput training and inference, we propose a hardware-aware implementation of LDM that reduces peak activation memory and memory traffic. LDM consistently achieves the lowest perplexities among PreLN, mHC and AttnRes baselines at the 400M and 1.5B parameter scales, alongside competitive zero-shot accuracy.
Figures & tables
Figure 1: The axes of representation mixing: sequence (tokens), channel (features), and depth (layers).
Figure 2: ( 2(a) ) LDM carries a separate per-token depth memory alongside the residual stream. Read and Write are simplified for brevity (the inverse covariance stream and the write address generation are omitted). ( 2(b) ) AttnRes reads from a growing key–value memory Ml and writes the sublayer output back into it. ( 2(c) ) mHC generates its three mappings from the current expanded state and applies them around one transformer branch.
Method
Streams
State update ( Write )
Memory readout ( Read )
PreLN
(xl)
xl=xl−1+vl
xl
AttnRes
((K≤l,V≤l))
Append(kl,vl)
Attend(Ml,ql)
mHC
(Xl)
Xl=Xl−1Hlres(Xl−1)+vlHlpost(Xl−1)⊤
XlHl+1pre(Xl)
LDM
(xl,Ml,Pl)
(xl,Ml)=(xl−1+vl,Ml−1+ulκl⊤)
xl+Mlql
Table 1: Overview of depth mixing methods. Each method maintains streams that are accessed through Read and Write operators. LDM uses the RLS update derived in Section 3.2 , with pseudo-value ul and gain κl . Gates and the inverse covariance update for Pl are omitted for brevity.
Figure 3: Efficient implementation of LDM for L=5 layers and windows (denoted by the dashed boxes) of size E=3 . Each dot denotes the Write at layer l fused with the Read for layer l+1 . In the forward pass, each dot saves a small context cl , which is used for the backward pass. Some details are omitted for clarity (see Appendix F ).
Figure 4: Left: Training loss at 1.5B model size (100B tokens, learning rate 10−3 ). Right: the last 25B tokens. Arrows: tokens LDM needs to reach each baseline’s final loss.
Method
ARC-E
ARC-C
Hella
PIQA
Wino
OBQA
Lmb
Avg
Lmb
Wiki
Val
acc↑
accn↑
accn↑
acc↑
acc↑
accn↑
acc↑
acc↑
ppl↓
ppl↓
ppl↓
400M @ 50B
PreLN
65.66
31.23
46.80
68.93
55.25
37.00
38.46
49.05
25.32
22.59
12.73
AttnRes (Full)
64.94
32.08
49.34
70.08
55.64
36.40
41.35
49.98
19.62
21.46
12.35
AttnRes (Block)
67.63
32.85
49.55
70.08
56.04
37.20
40.91
50.61
20.71
20.86
12.12
mHC
65.11
31.66
49.26
70.02
55.09
35.00
40.69
49.55
20.23
21.20
12.23
LDM, Dk=4
66.79
32.34
48.71
69.42
54.70
37.20
42.27
50.20
19.59
20.84
12.14
Table 2: Pretraining evaluations for both 400M and 1.5B models, generated using the lm-eval-harness ( Gao et al., 2024 ) . Val is the pretraining validation perplexity.
Figure 5: Throughput relative to PreLN on B200 for Qwen-like 1.7B and 14B models, comparing LDM, mHC and AttnRes. The mHC results use an improved version of the official mHC kernels. (a) Training step with 2K and 4K sequence length. (b) Prefill of 2K prompts at different batch sizes. (c) KV-cached decode after 2K prefill at different batch sizes. More settings in Appendix C .
Baselines
(a) Residual
(b) Update
(c) Addressing
(d) Gain, gates, forgetting
PreLN
LDM
F(r)
x=δ
Delta
Outer
v=x
QK(δ)
Static q
−γ
g=1
ω=1
λ=1
Δ
0.78
0.00
1.88
0.87
0.06
0.07
0.33
0.07
0.19
0.04
0.07
0.01
− 0.01
Table 3: LDM ablations, 400M @ 15B. Δ denotes change in held-out perplexity ( ↓ ) relative to our main configuration. (a) F(r) : memory-only layer input (normalization omitted). x=δ : no post-layer residual addition. (b) Delta/Outer: replace RLS with delta rule/outer product update. (c) v=x : store the residual stream. QK(δ) : queries and keys from layer outputs. Static q : learned query per read with read gate fixed to one. (d) −γ , g=1 , ω=1 : fix read gain, read gate, or write weight to one, respectively. λ=1 : disable forgetting. Selected training curves in Figure 18 , full definitions and results in Sections D.2 and 8 .
Figure 6: Held-out loss versus group count (400M, 15B tokens).
Figure 7: Final held-out perplexity on DepthBench against the aspect ratio (width/depth) at a fixed budget of 400M (left) and 1.6B (right) parameters. Tick labels denote the width/depth of each shape and its aspect ratio.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Configuration
400M Llama-like
1.5B Qwen-like
Backbone parameters
405.4M
1.512B
Transformer blocks
24
28
Model width D
1,024
2,048
Attention type
MHA
GQA
Query heads
16
16
Key/value heads
16
8
Appendix
Table 4: Backbone configurations used in our experiments. The bottom rows list the parameters that each method adds to the backbone.
1.7B
4B
8B
14B
Transformer blocks
28
36
36
40
Model width D
2,048
2,560
4,096
5,120
Query heads
16
32
32
40
Key/value heads
8
8
8
8
Head dimension
128
128
128
128
FFN intermediate width
6,144
9,728
12,288
17,408
Appendix
Table 5: Qwen3 shapes of the efficiency measurements. All other settings are those of the 1.5B Qwen-like model in Table 4 .
1.7B
4B
8B
14B
Training
PreLN
80.2
82.6
107.6
97.3
LDM
94.0
92.8
123.2
107.4
AttnRes (Block)
80.3
82.8
107.7
97.4
AttnRes (Full)
90.5
94.0
126.7
111.0
mHC
80.3
81.1
106.2
95.3
Appendix
Table 6: Peak memory allocated per GPU (GiB). Training runs on 8 GPUs with sequences of 2K tokens; prefill and decode run on one GPU at batch 64 with a 2K token prompt or context. Recomputation does not apply at inference, so the mHC row without recompute is training only.
Figure 8: Training throughput as a percentage of PreLN.
Scale
Method
5×10−4
1×10−3
2×10−3
4×10−3
400M
PreLN
2.5627
2.5441
2.5465
–
AttnRes (Full)
–
2.5172
2.5135
2.5133
AttnRes (Block)
–
2.5064
2.4950
2.5026
mHC
2.5291
2.5043
2.5041
–
LDM (Ours)
2.5138
2.4913
2.4957
–
LDM, Dk=4
2.5178
2.4965
2.4991
–
Appendix
Table 7: Learning rate sweeps: final held-out FineWeb-Edu cross-entropy of each method, model scale and peak learning rate. The selected learning rate of each method is in bold.
Figure 15: Final held-out loss vs peak learning rate (the numbers of Table 7 ). The filled marker is each method’s lowest loss.
Group
Variant
Loss ↓
Δ
Baselines
PreLN
2.6554
+0.0567
LDM
2.5987
0.0000
(a) Residual stream
Fl(rl−1)
2.7296
+0.1310
xl=δl
2.6611
+0.0625
(b) Recurrence
Delta
2.6028
+0.0041
Outer
2.6037
+0.0050
Appendix
Table 8: Final held-out FineWeb-Edu cross-entropy of the 400M ablation runs (15B tokens). Δ is relative to LDM.
Figure 18: Training loss of the 400M ablation runs over 15B FineWeb-Edu tokens, grouped as in Table 8 .
Tier
Layers
Width D
D/L
Head width
FFN width
Backbone
LDM adds
400M
16
1,216
76.0
76
3,248
406.5M
2.97M (0.73%)
20
1,120
56.0
70
2,992
414.1M
3.42M (0.83%)
24
1,024
42.7
64
2,736
405.4M
3.78M (0.93%)
28
960
34.3
60
2,560
406.3M
4.16M (1.02%)
32
896
28.0
56
2,400
399.4M
4.47M (1.12%)
42
800
19.0
50
2,144
404.2M
5.30M (1.31%)
Appendix
Table 9: DepthBench shapes. The 400M shapes have 16 query and 16 key/value heads, the 1.6B shapes 16 query and 8 key/value heads.
Figure 19: Final held-out perplexity on DepthBench at 400M parameters vs the aspect ratio (width/depth), from 16 to 70 layers. Tick labels give the width/depth of each shape and aspect ratio.
Scaling depth is a key driver for large language models (LLMs). Yet, as LLMs become deeper, they often suffer from signal degradation: informative features formed in shallow layers are gradually diluted by repeated residual updates, making them harder to recover in deeper layers. We introduce mixture-of-depths attention (MoDA), a mechanism that allows each attention head to attend to sequence KV pairs at the current layer and depth KV pairs from preceding layers. We further describe a hardware-efficient algorithm for MoDA that resolves non-contiguous memory-access patterns, achieving 97.3% of FlashAttention-2's efficiency at a sequence length of 64K. Experiments on 1.5B-parameter models demonstrate that MoDA consistently outperforms strong baselines. Notably, it improves average perplexity by 0.2 across 10 validation benchmarks and increases average performance by 2.11% on 10 downstream tasks, with a negligible 3.7% FLOPs computational overhead. We also find that combining MoDA with post-norm yields better performance than using it with pre-norm. These results suggest that MoDA is a promising primitive for depth scaling. Code is released at https://github.com/hustvl/MoDA .
Lianghui Zhu, Yuxin Fang, Bencheng Liao +10
School of EIC, Huazhong University of Science & Technology · ByteDance Seed
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× over the standard residual baseline, compared with 1.35× for AttnRes and 1.30× for mHC. These results make MSAR a promising residual architecture for large-scale pre-training.
Xiaohan Qin, Xiangdong Zhang, Tuo Dai +7
School of AI, Shanghai Jiao Tong University · Shanghai Innovation Institute · Dots Studio, Xiaohongshu Inc. +2
Self-attention selects information freely across the sequence, but across depth, Transformers merely add each layer's output to the residual stream, so later layers cannot selectively reuse earlier-layer representations. Recent cross-layer methods improve this flow but operate on hidden states outside attention, adding state beyond the key-value cache at inference--a cost that becomes increasingly salient as modern LLMs compress the cache with grouped-query and multi-head latent attention. We introduce Depth-Attention, which performs this selection inside the attention module itself: before a layer attends over the sequence, its query attends over the keys of earlier layers at the same token position and mixes their values into the value that self-attention then reads. Because Depth-Attention reuses the standard attention queries, keys, and value-cache slots, storing depth-mixed values in place of the original values, it adds no parameters and introduces no persistent inference state beyond the standard key-value cache--the same cache size as a vanilla decoder and less than hidden-state-based cross-layer methods. On Qwen3-style decoders at 1.5B and 3B parameters, Depth-Attention attains the lowest perplexity and the highest average downstream accuracy, improving over the vanilla Transformer by up to 2.3 accuracy points and surpassing strong cross-layer baselines in perplexity and average accuracy, while adding under 0.01% extra arithmetic FLOPs and no additional persistent inference state. The gains hold from 360M to 3B parameters and extend to looped Transformers.
Boyi Zeng, Yiqin Hao, Zitong Wang +7
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2