Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures

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.

💡 Research Summary

The paper investigates whether transformer models, which have become the dominant architecture for large language models, can serve as general-purpose unsupervised learning algorithms. Focusing on the classic problem of Gaussian Mixture Model (GMM) parameter estimation, the authors introduce a meta‑learning framework called TGMM (Transformer for Gaussian Mixture Models). TGMM is designed to solve a family of GMM tasks with varying numbers of mixture components K using a single shared transformer backbone, thereby achieving strong parameter efficiency.

The methodology proceeds as follows. A synthetic task sampler generates GMM instances by randomly selecting the number of components K, the mixture weights π, the component means μ, and the sample size N. Each task is represented as a sequence of tokens: the raw data points X are concatenated with a learned embedding of K, forming a matrix H =