Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures

Zhiheng Chen, Ruofan Wu, Guanhua Fang

TL;DR

This work investigates transformers as unsupervised learning tools for Gaussian Mixture Models (GMMs) by introducing TGMM, a framework that solves multiple GMM tasks with a single shared transformer. The approach demonstrates that TGMM can match or exceed traditional methods—often outperforming EM and approaching spectral performance—while maintaining robustness to distribution shifts and varying problem sizes. The paper provides theoretical foundations showing transformers can approximate both the EM algorithm and cubic tensor power iterations, thereby bridging practical success with formal understanding. This positions transformers as versatile, data-driven solvers for core unsupervised tasks and suggests broader applications in unsupervised learning settings.

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.

Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures

TL;DR

This work investigates transformers as unsupervised learning tools for Gaussian Mixture Models (GMMs) by introducing TGMM, a framework that solves multiple GMM tasks with a single shared transformer. The approach demonstrates that TGMM can match or exceed traditional methods—often outperforming EM and approaching spectral performance—while maintaining robustness to distribution shifts and varying problem sizes. The paper provides theoretical foundations showing transformers can approximate both the EM algorithm and cubic tensor power iterations, thereby bridging practical success with formal understanding. This positions transformers as versatile, data-driven solvers for core unsupervised tasks and suggests broader applications in unsupervised learning settings.

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.
Paper Structure (34 sections, 18 theorems, 113 equations, 18 figures, 5 algorithms)

This paper contains 34 sections, 18 theorems, 113 equations, 18 figures, 5 algorithms.

Table of Contents

  1. Introduction
  2. Methodology
  3. Preliminaries
  4. The TGMM architecture
  5. Meta training procedure
  6. Experiments
  7. Experimental Setup
  8. Results and findings
  9. Theoretical understandings
  10. Understanding TGMM
  11. Proof Ideas
  12. Conclusion and discussions
  13. Appendix
  14. Algorithm Details
  15. On the parameter efficiency of TGMM
  16. ...and 19 more sections

Key Result

Theorem 1

There exists a $2L$-layer transformer ${\rm TF}_{\boldsymbol{\Theta}}$ such that for any $d\leq d_0$, $K\leq K_0$ and task ${\mathcal{T}} = {\left( {\mathbf{X}}, {\boldsymbol \theta}, K \right)}$ satisfying some regular conditions, given suitable embeddings, ${\rm TF}_{\boldsymbol{\Theta}}$ approxim

Figures (18)

  • Figure 1: Illustration of the proposed TGMM architecture: TGMM utilizes a shared transformer backbone that supports solving $s$ different kind of GMM tasks via a task-specific $\operatorname{Readout}$ strategies.
  • Figure 2: Performance comparison between TGMM and two classical algorithms, reported in $\ell_2$-error.
  • Figure 3: Assessments of TGMM under test-time task distribution shifts I: A line with $N_0^\text{train} \rightarrow N^\text{test}$ draws the performance of a TGMM model trained over tasks with sample size randomly sampled in $[N_0^\text{train} / 2, N_0^\text{train}]$ and evaluated over tasks with sample size $N^\text{test}$. We can view the configuration $128\rightarrow 128$ as an in-distribution test and the rest as out-of-distribution tests.
  • Figure 4: Assessments of TGMM under test-time task distribution shifts II: $\ell_2$-error of estimation when the test-time tasks $\mathcal{T}^\text{test}$ are sampled using a mean vector sampling distribution $p_\mu^\text{test}$ different from the one used during training.
  • Figure 5: Performance comparisons between TGMM using transformer and Mamba2 as backbone, reported in $\ell_2$-error.
  • ...and 13 more figures

Theorems & Definitions (43)

  • Definition 1: Attention layer
  • Definition 2: MLP layer
  • Definition 3: Transformer
  • Theorem 1: Informal
  • Theorem 2: Informal
  • Remark 1
  • Remark 2
  • Theorem C.1
  • Remark C.1
  • Remark C.2
  • ...and 33 more