Research map / Architectures and efficiency / Attention and transformers
Linear Transformers Implicitly Discover Unified Numerical Algorithms
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.
Trains a linear attention-only transformer on masked-block low-rank matrix completion and extracts a parameter-free iterative solver (EAGLE) that works across centralized, distributed, and sketched regimes.
Problem
Whether transformers trained only to fill missing blocks of low-rank matrices implicitly discover numerical algorithms (Nyström extrapolation / least squares) that adapt to different computational constraints.
Approach
Train a linear, attention-only transformer on masked-block completion under three architectural visibility/rank constraints, then algebraically extract the same two-line parameter-free update (EAGLE), a Newton-Schulz-like conditioning loop wrapped around direct prediction, rather than interpreting it as preconditioned gradient descent.
Key results
- Extracted EAGLE update reproduces transformer layer-wise activations to within about 6e-4 error (squared Frobenius norm per entry) across the three regimes.
- EAGLE reaches eps=1e-2 in at most 15 iterations for M in {1,3,5,8} when worker subspaces are nearly orthogonal (alpha about 1), while GD requires 10-100x more rounds.
- In the sketched setting, per-iteration time falls from 21 ms (s=240) to 5 ms (s=30), and EAGLE reaches 1e-6 in roughly one third of RandSketch+CG wall-clock at kappa=1e4.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use EAGLE instead of Conjugate Gradient | centralized solver · Centralized setting, κ=10⁴ | iterations: ∼ 100× fewer iterations | Helped |
| Use EAGLE instead of Gradient Descent | distributed solver · Distributed setting | communication rounds: 10–100× more rounds | Helped |
| Add sketching to the EAGLE update | EAGLE · Computation-limited setting, κ=10⁴ | wall-clock time to ε=10⁻⁶: roughly one third the wall-clock | Helped |
| Reduce sketch size from s=240 to s=30 (ablation) | EAGLE with sketching · Computation-limited setting | per-iteration wall-clock time: 21, 16, 7 and 5 milliseconds | Helped |
| Reduce sketch rank r (ablation) | EAGLE with sketching · Computation-limited setting | iteration count: linearly with n/r | Hurt |
| Prune attention heads from eight to one (ablation) | unconstrained trained transformer · Matrix completion test set | test MSE | No clear effect |
| Freeze updates to Aᵘ and Bᵘ after per-machine conditioning converges (ablation) | distributed EAGLE · Distributed setting | wall-clock runtime | Helped |
| Freeze updates to Aᵘ and Bᵘ after per-machine conditioning converges (ablation) | distributed EAGLE · Distributed setting | iteration count | No clear effect |
| Train with modest data noise (σ²=0.01) (ablation) | linear attention transformer · Masked-block completion across computational settings | emergent algorithm behavior | No clear effect |
Surprising findings
- The learned iteration is closer to a Newton-Schulz conditioning loop than to preconditioned gradient descent, so the authors reject the GD++ interpretation.
- The transformer recovers sparse, roughly O(d) communication patterns in the distributed setting without any explicit communication constraint.
Limitations
- Numerical stability beyond condition number 1e8 is untested.
- Only one pre-training recipe is studied; extensions to non-linear models are open.
- Experiments are on synthetic data with a fixed Gaussian training distribution.
Experimental setup
Datasets: synthetic low-rank matrices with anisotropic Gaussian factors
Benchmarks: synthetic low-rank matrix completion (Nyström extrapolation) distributed M/alpha sweeps sketch-rank sweeps
Metrics: squared Frobenius error iterations to tolerance wall-clock time per iteration
Baselines: QR-based least squares (torch.linalg.lstsq) Conjugate Gradient gradient descent SGD RandSketch+CG GD++ (preconditioned GD)
Builds on: Newton-Schulz method in-context learning as gradient descent Nyström approximation linear attention transformer
Closest papers by meaning
- CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs (NeurIPS 2026 · similarity 0.88)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.88)
- From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers (NeurIPS 2025 · similarity 0.88)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.87)
- Learning Linear Attention in Polynomial Time (NeurIPS 2025 · similarity 0.87)
- Predictive Geometry of Hidden Trajectories in Transformers (NeurIPS 2026 · similarity 0.87)
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.