Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures

2025-05-17

AI Summary

This work investigates the fundamental unsupervised learning capabilities of Transformers, presenting the first systematic statistical estimation–driven analysis of their ability to solve Gaussian Mixture Models (GMMs). We propose TGMM, an end-to-end differentiable framework that jointly optimizes GMM parameter estimation and latent variable inference using a shared Transformer backbone. Theoretically, we prove that Transformers can uniformly approximate both the Expectation-Maximization (EM) algorithm and the cubic tensor power iteration method. Empirically, TGMM significantly outperforms classical approaches: it is initialization-robust, avoids local optima, exhibits robustness to distributional shift, and generalizes to unseen numbers of mixture components and covariance structures. This work establishes the first paradigm for large-model–based unsupervised learning that simultaneously provides rigorous theoretical guarantees and strong empirical performance.

Technology Category

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

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

The transformer architecture has demonstrated remarkable capabilities in modern artificial intelligence, among which the capability of implicitly learning an internal model during inference time is widely believed to play a key role in the under standing of pre-trained large language models. However, most recent works have been focusing on studying supervised learning topics such as in-context learning, leaving the field of unsupervised learning largely unexplored. This paper investigates the capabilities of transformers in solving Gaussian Mixture Models (GMMs), a fundamental unsupervised learning problem through the lens of statistical estimation. We propose a transformer-based learning framework called TGMM that simultaneously learns to solve multiple GMM tasks using a shared transformer backbone. The learned models are empirically demonstrated to effectively mitigate the limitations of classical methods such as Expectation-Maximization (EM) or spectral algorithms, at the same time exhibit reasonable robustness to distribution shifts. Theoretically, we prove that transformers can approximate both the EM algorithm and a core component of spectral methods (cubic tensor power iterations). These results bridge the gap between practical success and theoretical understanding, positioning transformers as versatile tools for unsupervised learning.

Problem

• Investigating transformers' capabilities in solving Gaussian Mixture Models (GMMs)
• Proposing a transformer-based framework (TGMM) for unsupervised GMM learning
• Bridging theoretical and practical gaps in transformers for unsupervised tasks

Innovation

• Transformers solve Gaussian Mixture Models unsupervised
• TGMM framework learns multiple GMM tasks jointly
• Theoretical proof: transformers approximate EM and spectral methods