Skip to main content
Back to timeline
arXivSource publication:

Beyond the residual stream, a depth memory: LDM beats PreLN, mHC and AttnRes perplexity at 400M and 1.5B

Synopsis

Transformers have long passed information across depth by adding layer outputs in a residual stream, and recent alternatives mHC and AttnRes make that depth aggregation more expressive. Depth Memory recasts cross-layer flow as streams that store and retrieve layer representations, and from it comes Linear Depth Memory (LDM), which treats retrieval across depth as a regression solved exactly by recursive least squares while carrying a residual stream, a matrix memory stream and an inverse covariance stream. Pretrained on 400M Llama-like and 1.5B Qwen-like backbones with 50B and 100B FineWeb-Edu tokens, LDM reaches lower validation, WikiText-2 and LAMBADA perplexity than tuned PreLN, mHC and AttnRes, with competitive zero-shot accuracy.

AI-generated editorial illustration: Depth Memory

Interpretation

Cross-layer information no longer has to travel only by adding layer outputs one after another; LDM lets each layer read earlier layers from a fixed-size memory using a content-dependent query. The standard PreLN residual stream gives every earlier layer the same weight, mHC keeps an expanded residual stream, and AttnRes explicitly retains and attends over previous layer outputs. On a 400M Llama-like backbone trained on 50B tokens and a 1.5B Qwen-like backbone trained on 100B tokens, the main LDM configuration has the lowest validation, WikiText-2 and LAMBADA perplexity; at 1.5B its LAMBADA perplexity is 8.42 against 9.05 for mHC, 8.84 for Block AttnRes and 10.49 for PreLN.

The memory update is posed as a regression across depth and solved exactly online by recursive least squares rather than approximated by an outer product or a delta rule. A plain outer-product update approximates the key Gram matrix by the identity, and the delta rule takes one gradient step per layer; LDM maintains the inverse covariance and keeps the exact solution. In 400M ablations, replacing RLS with the delta rule or a plain outer product still beats PreLN but neither matches RLS, and both variants that remove the separate residual stream fall below PreLN.

To make the layer-by-layer recurrence practical in training, the implementation does not store each layer's memory state but unwinds it in place from the final state during backpropagation, and defers writes within a window. A naive implementation stores a state per layer for backpropagation and moves the current state on every read and write, costing activation memory and memory traffic. Across Qwen3 models from 1.7B to 14B parameters, sequence and context lengths and batch sizes, LDM exceeds Full AttnRes throughput and beats mHC in all training and most inference settings, while staying broadly comparable to Block AttnRes.

Perspective

The demonstrated scope is language-model pretraining: a 400M Llama-like backbone on 50B FineWeb-Edu tokens and a 1.5B Qwen-like backbone on 100B tokens, plus aspect-ratio sweeps at 400M and 1.6B under the DepthBench protocol, with throughput measured on Qwen3 models from 1.7B to 14B. It suits readers working on Transformer architecture and training systems; larger scales and other domains are left to future work.

LDM can add peak memory during training, which other stable backward recurrences or checkpointing could reduce. At 400M on extreme aspect ratios of 42 and 70 layers, HC leads and LDM perplexity rises at 70 layers; at 1.6B, LDM has the lowest perplexity at every aspect ratio considered. Larger scales and other domains are not evaluated.