# Transformers can optimally learn regression mixture models

arXiv: https://arxiv.org/abs/2311.08362 · HTML: https://mlsanity.com/paper/2311.08362/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 28 references that are in the map.
- First posted on arXiv 14 Nov 2023.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [A Survey on Visual Transformer](https://mlsanity.com/paper/2012.12556/) (2020)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [A Field Guide to Federated Optimization](https://mlsanity.com/paper/2107.06917/) (2021)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [In-Context Learning through the Bayesian Prism](https://mlsanity.com/paper/2306.04891/) (2023)
- [Meta-learning for mixed linear regression](https://mlsanity.com/paper/2002.08936/) (2020)
- [Learning Mixtures of Linear Regressions with Nearly Optimal Complexity](https://mlsanity.com/paper/1802.07895/) (2018)
- [Solving a Mixture of Many Random Linear Equations by Tensor Decomposition and Alternating Minimization](https://mlsanity.com/paper/1608.05749/) (2016)
- [Learning Mixtures of Linear Regressions in Subexponential Time via Fourier Moments](https://mlsanity.com/paper/1912.07629/) (2019)
- [Avoiding Imposters and Delinquents: Adversarial Crowdsourcing and Peer Prediction](https://mlsanity.com/paper/1606.05374/) (2016)
- [Linear Regression using Heterogeneous Data Batches](https://mlsanity.com/paper/2309.01973/) (2023)

## Built on by (most cited first, 15 of 15)

- [A Theoretical Understanding of Self-Correction through In-context Alignment](https://mlsanity.com/paper/2405.18634/) (2024)
- [In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization](https://mlsanity.com/paper/2402.14951/) (2024)
- [Linear Transformers are Versatile In-Context Learners](https://mlsanity.com/paper/2402.14180/) (2024)
- [Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time](https://mlsanity.com/paper/2510.01098/) (2025)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [A transfer learning framework for weak-to-strong generalization](https://mlsanity.com/paper/2405.16236/) (2024)
- [Transformer learns the cross-task prior and regularization for in-context learning](https://mlsanity.com/paper/2505.12138/) (2025)
- [On the Robustness of Transformers against Context Hijacking for Linear Classification](https://mlsanity.com/paper/2502.15609/) (2025)
- [Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures](https://mlsanity.com/paper/2505.11918/) (2025)
- [Integrating the Expected Future in Load Forecasts with Contextually Enhanced Transformer Models](https://mlsanity.com/paper/2409.05884/) (2024)
- [In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamics](https://mlsanity.com/paper/2410.14183/) (2024)
- [On the Existence of Universal Simulators of Attention](https://mlsanity.com/paper/2506.18739/) (2025)
- [Limitations of refinement methods for weak to strong generalization](https://mlsanity.com/paper/2508.17018/) (2025)
- [Transformers Learn Latent Mixture Models In-Context via Mirror Descent](https://mlsanity.com/paper/2604.10848/) (2026)
- [Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context](https://mlsanity.com/paper/2502.04580/) (2025)

## Abstract (arXiv)

Mixture models arise in many regression problems, but most methods have seen limited adoption partly due to these algorithms' highly-tailored and model-specific nature. On the other hand, transformers are flexible, neural sequence models that present the intriguing possibility of providing general-purpose prediction methods, even in this mixture setting. In this work, we investigate the hypothesis that transformers can learn an optimal predictor for mixtures of regressions. We construct a generative process for a mixture of linear regressions for which the decision-theoretic optimal procedure is given by data-driven exponential weights on a finite set of parameters. We observe that transformers achieve low mean-squared error on data generated via this process. By probing the transformer's output at inference time, we also show that transformers typically make predictions that are close to the optimal predictor. Our experiments also demonstrate that transformers can learn mixtures of regressions in a sample-efficient fashion and are somewhat robust to distribution shifts. We complement our experimental observations by proving constructively that the decision-theoretic optimal procedure is indeed implementable by a transformer.