Transformers versus the EM Algorithm in Multi-class Clustering

2025-02-09

AI Summary

This paper investigates the theoretical performance of Transformers in unsupervised clustering under multi-component Gaussian mixture models (GMMs), aiming to establish a rigorous connection with the Expectation-Maximization (EM) algorithm and derive statistical learning guarantees. Methodologically, by analyzing the structural properties of the Softmax attention layer, the authors prove—for the first time—that its forward pass is equivalent to a joint iteration of the E-step and M-step of EM. Leveraging functional approximation theory and statistical learning theory, they further establish that Softmax possesses universal approximation capability for multivariate mappings and achieves the minimax-optimal convergence rate for clustering. Empirical simulations demonstrate strong robustness under model misspecification and finite-sample regimes. The work provides the first interpretable, EM-principled theoretical framework for unsupervised learning with large-scale models.

Technology Category

• Machine Learning: Unsupervised & Self-Supervised Learning
• Natural Language Processing: Learning & Optimization for NLP
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Normalization, clustering, classification, and summarization of Web text

Abstract

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.

Problem

• Transformers in multi-class clustering
• comparison with EM algorithm
• universal approximation by Softmax

Innovation

• Transformers for multi-class clustering
• Softmax Attention parallels EM algorithm
• Transformers achieve minimax optimal rate