Research map / Optimization / Implicit bias and generalization
Towards Learning High-Precision Least Squares Algorithms with Sequence Models
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.
Studies whether sequence models can learn numerically precise least-squares/gradient-descent iterates; shows softmax Transformers struggle with high-precision multiplication and proposes polynomial architectures (BaseConv, linear attention) trained with an adaptive-LR, EMA high-precision recipe to reach near float32 machine precision.
Problem
Standard Transformers trained in-context on least squares do not reach machine precision and do not generalize numerically out-of-distribution, so they do not learn numerical algorithms such as gradient descent.
Approach
Argues softmax attention cannot exactly express element-wise multiplication, so uses polynomial sequence mixers (BaseConv gated convolutions, non-causal linear attention) that can express GD exactly, and introduces a high-precision training recipe (cosine-similarity-driven adaptive LR plus EMA over Adam updates) that reduces minibatch gradient noise to reach near machine precision.
Key results
- Iterated 3-layer BaseConv trained with the recipe reaches about 1e-12 to 1e-13 MSE on GD iterates, about 100,000x lower MSE than the largest end-to-end Transformers.
- BaseConv shows about 10,000x smaller generalization gap on out-of-distribution targets than the Transformer.
- Learning up to k=4 GD iterates at once reaches MSE 3.1e-10 with BaseConv.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a high-precision training recipe with AdaptiveLR and EMA | BaseConv on explicit gradient updates · Least squares, explicit gradient task | MSE: $10^{-13}$ average MSE | Helped |
| Use BaseConv instead of an end-to-end Transformer and apply predicted gradients iteratively | Least squares solver · Least squares, iteratively applied to convergence | MSE: $100,000\times$ better MSE | Helped |
| Use iteratively applied BaseConv instead of the Transformer | Least squares solver · Out-of-distribution regression targets | generalization gap: $10,000\times$ smaller | Helped |
| Use iterative BaseConv rather than a 12-layer Transformer | In-context ODE solving · In-context ODE solving | MSE: up to $1,000,000 \times$ in MSE | Helped |
| Train the BaseConv k-th iterate task for larger k | BaseConv · Least squares, k-th GD iterate task | MSE: $5.0 \times 10^{-13}$; $2.5 \times 10^{-11}$; $2.5 \times 10^{-11}$; $3.1 \times 10^{-10}$ | Mixed |
| Replace softmax attention with BaseConv and train end-to-end (ablation) | Standard Transformer architecture on least squares · Least squares, end-to-end | MSE: $100\times$ worse | Hurt |
| Add EMA over Adam update vectors (ablation) | BaseConv on the explicit gradient task · Explicit gradient task | MSE: up to $100,000\times$ | Helped |
| Add MLPs and LayerNorms (ablation) | BaseConv on the explicit gradient task · Explicit gradient task | MSE: up to $1,000,000\times$ MSE | Hurt |
| Add LayerNorm (ablation) | BaseConv on small-k GD iterate tasks · Least squares, k-th GD iterate task | MSE: over $1,000\times$ | Hurt |
| Use linear attention instead of BaseConv (ablation) | Iterative least squares solver · Out-of-distribution regression targets | generalization gap | Hurt |
| Increase the initial learning rate from $10^{-3}$ to $10^{-2}$ with an aggressively decaying schedule (ablation) | Training on the Multiply primitive · Multiply primitive | MSE: $1000\times$ | Helped |
| Use longer training duration (ablation) | 1-layer Transformer and BaseConv on Multiply · Multiply primitive | MSE: 2-3 order of magnitude improvement | Helped |
Surprising findings
- Linear attention, despite near machine precision in-distribution, has poor numerical generality and its iterates diverge beyond sigma=4.
- LayerNorms and MLPs worsen precision by up to 1,000,000x on the explicit gradient task.
Limitations
- Deep networks (beyond ~4 iterates) are hard to train stably and precisely; LayerNorm reintroduction hurts precision.
- Experiments are on synthetic tasks with fixed small least squares sizes (20x5).
- Adaptive LR scheduler still needs careful hyperparameter choice; convergence to machine precision is slow.
Experimental setup
Datasets: synthetic in-context least squares problems synthetic Read/Linear/Multiply primitives synthetic in-context ODE problems
Benchmarks: in-context least squares (20x5) out-of-distribution regression targets in-context ODE solving
Metrics: MSE cosine similarity of minibatch gradients
Baselines: softmax Transformer (GPT-2 style) gradient descent Newton's method standard Adam with StepLR
Builds on: BaseConv GPT-2 Adam
Closest papers by meaning
- Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks (ICML 2018 · similarity 0.87)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.85)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.85)
- Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit (NeurIPS 2024 · similarity 0.85)
- Collage: Light-Weight Low-Precision Strategy for LLM Training (ICML 2024 · similarity 0.84)
- Deep linear networks for regression are implicitly regularized towards flat minima (NeurIPS 2024 · similarity 0.84)
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.