Research map / Architectures and efficiency / Attention and transformers
Residual Matrix Transformers: Scaling the Size of the Residual Stream
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Replaces the transformer residual stream with per-token outer-product memory matrices (Residual Matrix Transformer), which scales the residual stream nearly free in parameters and FLOPs and improves pretraining efficiency on OpenWebText.
Problem
Scaling the residual stream in standard transformers increases parameter and FLOP counts linearly, so residual stream bandwidth cannot be scaled independently of model size and compute.
Approach
Replace each token's residual stream vector with an outer-product memory matrix (residual matrix) written to and read from via learned key vectors and tensor contractions, so the residual stream size (D_k x D_v) can grow with negligible parameter or FLOP cost while the attention and FF weight matrices are replaced by key vectors.
Key results
- At a 405M-param transformer's best train loss of 3.03, the 305M RMT reaches it with 58% fewer FLOPs, 25% fewer parameters, and 41% fewer tokens (runtime 4% slower).
- On zero-shot downstream tasks, the RMT (305M) beats the 405M transformer on all reported perplexity and accuracy tasks, e.g., LAMBADA perplexity 25.21 vs 61.34 and Winogrande 52.5 vs 49.5.
- Scaling residual stream size (D_k 384 to 4096) at fixed parameters and FLOPs monotonically lowers dev loss; the 4096 model matches the 768 model's final loss with 23% fewer FLOPs and 25% fewer tokens.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the vector residual stream with an outer-product residual matrix | Transformer · OpenWebText after 6B tokens | train loss: 3.03 -> 2.91 | Helped |
| Replace the vector residual stream with an outer-product residual matrix | Transformer · OpenWebText train loss 3.03 | FLOPS: 1.18 \times 10^{19} -> 4.97 \times 10^{18} (-58%) | Helped |
| Replace the vector residual stream with an outer-product residual matrix | Transformer · OpenWebText train loss 3.03 | training tokens: 6.04 \times 10^{9} -> 3.55 \times 10^{9} (-41%) | Helped |
| Replace the vector residual stream with an outer-product residual matrix | Transformer · OpenWebText train loss 3.03 | time: 1.86 \times 10^{5} -> 1.94 \times 10^{5} (+4%) | Hurt |
| Expand the RMT residual stream while holding model size and compute budget fixed (ablation) | RMT · OpenWebText; residual stream sizes 384–4096 | dev loss | Helped |
| Expand residual size from 768 to 4096 (ablation) | RMT · OpenWebText | FLOPS to same final loss: 23% fewer FLOPS | Helped |
| Expand residual size from 768 to 4096 (ablation) | RMT · OpenWebText | training tokens to same final loss: 25% fewer tokens | Helped |
| Use an RMT with residual size 768 instead of an expanded residual stream (ablation) | Transformer comparison · OpenWebText | final train loss: 3.43 -> 3.42 | No clear effect |
| Replace the Transformer with RMT | Transformer · Lambada OpenAI | perplexity: 61.34 -> 25.21 | Helped |
| Replace the Transformer with RMT | Transformer · OpenWebText downstream | perplexity: 20.64 -> 18.36 | Helped |
| Replace the Transformer with RMT | Transformer · Arc C | accuracy: 19.7 -> 20.6 | Helped |
| Replace the Transformer with RMT | Transformer · Arc E | accuracy: 44.3 -> 46.1 | Helped |
Surprising findings
- The RMT's runtime is roughly even or slower than the transformer (4% slower per time to target loss, 43% slower over 6B tokens), despite better FLOP and token efficiency, attributed to unoptimized tensor contractions.
- Prior residual-stream variants (Hierarchical Aggregation, Depthwise LSTM, Highway Transformer) all use more parameters and FLOPs than the transformer, whereas the RMT uses fewer.
Limitations
- Runtime is slower than the transformer because the tensor contractions are not hardware-optimized.
- Scaling experiments were limited to models up to 405M parameters and 6B tokens due to resource constraints.
- Attention storage moment propagation is worse than the transformer's (variance ratio 0.4 vs 1 forward).
Experimental setup
Datasets: OpenWebText
Benchmarks: LAMBADA OpenWebText perplexity pile 10K WikiText ARC Challenge ARC Easy HellaSwag PIQA Winogrande
Metrics: train loss dev loss perplexity accuracy FLOPs to target loss tokens to target loss
Baselines: GPT2 transformer (standard) Hierarchical Aggregation Depthwise LSTM Highway Transformer
Builds on: Transformer Kohonen linear associative memory Mu-Parameterization (muTransfer) Highway Networks DenseNet
Closest papers by meaning
- The Quantization Benefits of Residual-Free Transformers (NeurIPS 2026 · similarity 0.86)
- CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs (NeurIPS 2026 · similarity 0.86)
- Linear Transformers Implicitly Discover Unified Numerical Algorithms (NeurIPS 2025 · similarity 0.85)
- MemoryFormer : Minimize Transformer Computation by Removing Fully-Connected Layers (NeurIPS 2024 · similarity 0.85)
- Predictive Geometry of Hidden Trajectories in Transformers (NeurIPS 2026 · similarity 0.85)
- A unified framework for establishing the universal approximation of transformer-type architectures (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.