# In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamics

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 51 references that are in the map.
- First posted on arXiv 18 Oct 2024.

## Builds on (most cited first, 30 of 35)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [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)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [cosFormer: Rethinking Softmax in Attention](https://mlsanity.com/paper/2202.08791/) (2022)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Statistical guarantees for the EM algorithm: From population to sample-based analysis](https://mlsanity.com/paper/1408.2156/) (2014)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Uncovering mesa-optimization algorithms in Transformers](https://mlsanity.com/paper/2309.05858/) (2023)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [Spectral Methods meet EM: A Provably Optimal Algorithm for Crowdsourcing](https://mlsanity.com/paper/1406.3824/) (2014)
- [A Study on ReLU and Softmax in Transformer](https://mlsanity.com/paper/2302.06461/) (2023)
- [Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding](https://mlsanity.com/paper/2106.12566/) (2021)
- [Meta-learning for mixed linear regression](https://mlsanity.com/paper/2002.08936/) (2020)
- [Ten Steps of EM Suffice for Mixtures of Two Gaussians](https://mlsanity.com/paper/1609.00368/) (2016)
- [Singularity, Misspecification, and the Convergence Rate of EM](https://mlsanity.com/paper/1810.00828/) (2018)
- [Local Maxima in the Likelihood of Gaussian Mixture Models: Structural Results and Algorithmic Consequences](https://mlsanity.com/paper/1609.00978/) (2016)
- [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)
- [Global analysis of Expectation Maximization for mixtures of two Gaussians](https://mlsanity.com/paper/1608.07630/) (2016)
- [EM Converges for a Mixture of Many Linear Regressions](https://mlsanity.com/paper/1905.12106/) (2019)
- [Randomly initialized EM algorithm for two-component Gaussian mixture achieves near optimality in O(√n) iterations](https://mlsanity.com/paper/1908.10935/) (2019)
- [Learning Mixtures of Linear Regressions in Subexponential Time via Fourier Moments](https://mlsanity.com/paper/1912.07629/) (2019)
- [Global Convergence of EM Algorithm for Mixtures of Two Component Linear Regression](https://mlsanity.com/paper/1810.05752/) (2018)
- [On the Minimax Optimality of the EM Algorithm for Learning Two-Component Mixed Linear Regression](https://mlsanity.com/paper/2006.02601/) (2020)
- [Spectral Experts for Estimating Mixtures of Linear Regressions](https://mlsanity.com/paper/1306.3729/) (2013)

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

- [Theory on Attention Dynamics for Out-of-Distribution In-Context Learning](https://mlsanity.com/paper/2609.36448/) (2026)

## Abstract (arXiv)

We investigate the in-context learning capabilities of transformers for the d-dimensional mixture of linear regression model, providing theoretical insights into their existence, generalization bounds, and training dynamics. Specifically, we prove that there exists a transformer capable of achieving a prediction error of order O(√d/n) with high probability, where n represents the training prompt size in the high signal-to-noise ratio (SNR) regime. Moreover, we derive in-context excess risk bounds of order O(L/√B) for the case of two mixtures, where B denotes the number of training prompts, and L represents the number of attention layers. The dependence of L on the SNR is explicitly characterized, differing between low and high SNR settings. We further analyze the training dynamics of transformers with single linear self-attention layers, demonstrating that, with appropriately initialized parameters, gradient flow optimization over the population mean square loss converges to a global optimum. Extensive simulations suggest that transformers perform well on this task, potentially outperforming other baselines, such as the Expectation-Maximization algorithm.