Research map / Learning theory and algorithms / Prediction and decision losses
Transformers can optimally learn regression mixture 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.
Shows that decoder-only transformers trained on prompts from a discrete mixture of linear regressions approach the Bayes-optimal posterior-mean predictor, and proves constructively that a transformer can implement that predictor.
Problem
Learning mixtures of linear regressions from batched in-context prompts, where each batch comes from one unknown subpopulation; existing algorithms are model-specific and need problem parameters.
Approach
Train autoregressive transformers on prompts drawn from a discrete mixture of linear regressions; show empirically their predictions track the decision-theoretic posterior-mean (exponential-weights) predictor, and give an explicit construction of a transformer that computes that posterior mean exactly.
Key results
- Transformers with 12 layers, hidden size 256 achieve prediction error close to oracle posterior mean and argmin procedures across mixtures with 5 and 20 components.
- In inference-time probing, the posterior mean with oracle weights is the closest algorithm to the transformer's predictions in all simulated settings.
- Transformers handle small covariate scaling (kappa in {0.33, 0.5, 2}) and small additive weight shifts (eps=0.25) but are sensitive to weight scaling.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use a trained transformer instead of OLS | mixtures of linear regressions · mixture of linear regressions | normalized MSE | Helped |
| use a trained transformer instead of the oracle posterior mean | mixtures of linear regressions · mixture of linear regressions | normalized MSE | No clear effect |
| use a trained transformer instead of the oracle argmin procedure | mixtures of linear regressions · mixture of linear regressions | normalized MSE | No clear effect |
| train a transformer with a fixed training set size rather than using model-specific procedures | mixtures of linear regressions · fixed training set sizes | mean-squared error | No clear effect |
| compare transformer predictions with the posterior mean using oracle weights | trained transformer · simulated mixtures with noise level σ = 1.0 | squared prediction distance | Helped |
| evaluate with covariates scaled by κ ∈ {0.33, 0.5, 2} | trained transformer · covariate scaling shift | normalized MSE | No clear effect |
| evaluate with covariate shifts much larger than κ ∈ {0.33, 0.5, 2} | trained transformer · covariate scaling shift | normalized MSE | Hurt |
| scale mixture weights away from their training values | trained transformer · weight scaling shift | normalized MSE | Hurt |
| apply a small additive weight shift (ε = 0.25) | trained transformer · additive weight shift | normalized MSE | No clear effect |
Surprising findings
- Transformers trained with fixed training set sizes match or beat EM- and subspace-algorithm plug-in procedures in sample efficiency.
Limitations
- Only discrete mixtures of linear regressions with Gaussian noise studied; nonlinear components not examined.
- Transformer is sensitive to weight scaling shifts and large covariate shifts.
- Posterior mean in the in-context setting with a distribution over mixtures would require high-dimensional integration, not addressed.
Experimental setup
Datasets: synthetic discrete mixture of linear regressions prompts
Benchmarks: synthetic mixture of linear regressions prompts
Metrics: normalized mean squared error squared distance between predictors d_k^sq MSE under covariate and weight shifts
Baselines: ordinary least squares argmin oracle posterior mean with oracle weights posterior mean with EM weights posterior mean with SA weights
Builds on: in-context learning of linear regression by transformers batch expectation-maximization subspace algorithm for batched mixed linear regression RAW operator transformer construction
Closest papers by meaning
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.87)
- Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.86)
- On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures (ICML 2025 · similarity 0.86)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.86)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.86)
- Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context (ICLR 2025 · similarity 0.86)
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.