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.