# Transformers versus the EM Algorithm in Multi-class Clustering

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 35 references that are in the map.
- First posted on arXiv 9 Feb 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Breaking the Curse of Dimensionality with Convex Neural Networks](https://mlsanity.com/paper/1412.8690/) (2014)
- [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)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Self-Attention Networks Can Process Bounded Hierarchical Languages](https://mlsanity.com/paper/2105.11115/) (2021)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [On the Computational Power of Transformers and its Implications in Sequence Modeling](https://mlsanity.com/paper/2006.09286/) (2020)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [Statistical and Computational Guarantees of Lloyd's Algorithm and its Variants](https://mlsanity.com/paper/1612.02099/) (2016)
- [Image Clustering with External Guidance](https://mlsanity.com/paper/2310.11989/) (2023)
- [Deep Transformation-Invariant Clustering](https://mlsanity.com/paper/2006.11132/) (2020)
- [How Well Can Transformers Emulate In-context Newton's Method?](https://mlsanity.com/paper/2403.03183/) (2024)
- [LSEnet: Lorentz Structural Entropy Neural Network for Deep Graph Clustering](https://mlsanity.com/paper/2405.11801/) (2024)
- [A Mechanism for Sample-Efficient In-Context Learning for Sparse Retrieval Tasks](https://mlsanity.com/paper/2305.17040/) (2023)
- [One-Layer Transformer Provably Learns One-Nearest Neighbor In Context](https://mlsanity.com/paper/2411.10830/) (2024)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)
- [Learning Spectral Methods by Transformers](https://mlsanity.com/paper/2501.01312/) (2025)

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

- [Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures](https://mlsanity.com/paper/2505.11918/) (2025)
- [Attention-based clustering](https://mlsanity.com/paper/2505.13112/) (2025)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent](https://mlsanity.com/paper/2605.06609/) (2026)
- [Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models](https://mlsanity.com/paper/2603.22801/) (2026)

## Abstract (arXiv)

LLMs demonstrate significant inference capacities in complicated machine learning tasks, using the Transformer model as its backbone. Motivated by the limited understanding of such models on the unsupervised learning problems, we study the learning guarantees of Transformers in performing multi-class clustering of the Gaussian Mixture Models. We develop a theory drawing strong connections between the Softmax Attention layers and the workflow of the EM algorithm on clustering the mixture of Gaussians. Our theory provides approximation bounds for the Expectation and Maximization steps by proving the universal approximation abilities of multivariate mappings by Softmax functions. In addition to the approximation guarantees, we also show that with a sufficient number of pre-training samples and an initialization, Transformers can achieve the minimax optimal rate for the problem considered. Our extensive simulations empirically verified our theory by revealing the strong learning capacities of Transformers even beyond the assumptions in the theory, shedding light on the powerful inference capacities of LLMs.