Transformers versus the EM Algorithm in Multi-class Clustering

TL;DR

The paper addresses the problem of understanding how Transformer architectures can perform unsupervised multi-class clustering for Gaussian Mixture Models by linking Softmax attention to EM/Lloyd's algorithm. It develops a constructive approximation theory showing that pre-trained Transformers can emulate Lloyd iterations with explicit bounds and derives generalization guarantees for ERM pretraining. The key contributions include universal approximation results for multi-head Softmax mappings, generalization bounds, and minimax-rate results under sufficient pretraining and initialization, corroborated by simulations. The work advances the theoretical foundation of Transformer-based in-context learning for unsupervised algorithmic tasks and suggests practical implications for leveraging LLMs in clustering and related inference problems.

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.

Figures (5)

• Figure 1: $4$-Class Clustering with Different Minimum Distance, Data Dimension, and Number of Training Data. We train a small Transformer (layer $=3$, head $=2$, embedding $= 64$) and iterate for $300$ steps for each different setting. Each point in the figure is evaluated on $512$ testing data. We report the $10$ runs averaged result with a shaded region representing the standard deviation. Each training sample is generated according to isotropic Gaussian with covariances $\sigma^2 \boldsymbol{I}$. (1) First Row: Minimum Distance. We set $\sigma^2 \sim \mathrm{Uniform}[10,40]$. (2) Second Row: Data Dimension. We set $\sigma^2 \sim \mathrm{Uniform}[10,20]$, minimum distance $=5$. (3) Three Row: Number of Training Data. We set $\sigma^2 \sim \mathrm{Uniform}[0.5,5]$, minimum distance $=5$.
• Figure 2: $4$-Class Clustering with Different Number of Class and Inbalance Ratio. We train a small Transformer (layer $=3$, head $=2$, embedding $= 64$) and train for $300$ steps for each different setting. Each point in the figure is evaluated on $512$ testing data. We report the $10$ runs averaged result with a shaded region representing the standard deviation. Each training sample is generated according to isotropic Gaussian with covariances $\sigma^2 \boldsymbol{I}$. (1) First Row: Number of Class. We set $\sigma^2 \sim \mathrm{Uniform}[10,20]$, minimum distance $=5$. (2) Second Row: Inbalance Ratio. Two clusters each contain 50 data points, while the other two contain $50 \times \mathrm{ratio}$ and $50 \times \mathrm{1-ratio}$ respectively. We set $\sigma^2 \sim \mathrm{Uniform}[10,20]$, minimum distance $=5$.
• Figure 3: Comparision between Transformer and Lloyd's Algorithm. We compare the effect of the number of layers in Transformers with the number of iterations $\tau$ in Lloyd's algorithm under the same dataset configuration. We use a $6$-class dataset, where each cluster contains $50$ data points in a $d=10$ dimensional space. Each training sample is generated according to isotropic Gaussian with covariances $\sigma^2 \boldsymbol{I}$, where $\sigma^2 \sim \mathrm{Uniform}[20,30]$, and the minimum cluster separation is set to $1$. (1) Left: Transformer. We train Transformers with fixed head $=2$, embedding $= 64$, but vary the number of layers from $3$ to $20$. Each model is trained for $500$ steps per layer. (2) Right: Lloyd's Algorithm. We use sklearnpedregosa2011scikit to run the Lloyd's algorithm, varying the maximum iteration count from $1$ to $6$. Early convergence is declared when the Frobenius norm of the difference between cluster centers in consecutive iterations falls below $10^{-4}$. Each point in the figure represents an evaluation of $512$ test samples. Results are averaged over $10$ runs, with the shaded region indicating the standard deviation.
• Figure 4: Comparison of Concatenated Attention and Averaged Attention on Synthetic Dataset.Top: Performance Comparision on Minimum Distance Task.Bottom: Performance Comparision on Number of Data Task. We observe similar trend of performance between concatenated multihead attention a averaged multihead attention across three tasks, two evaluation metrics, and the converged loss. All the experiment settings are the same as the experiments in \ref{['sect4']}.
• Figure 5: Comparison of Concatenated Attention and Averaged Attention on Synthetic Dataset.Top: Performance Comparision on Number of Classes Task.Bottom: Performance Comparision on Inbalance Ratio Task. Again, we observe a similar performance trend between concatenated multihead attention and averaged multihead attention across both tasks. All experimental settings remain the same as those in \ref{['sect4']}.

Theorems & Definitions (28)

• Definition 2.1: Softmax Attention
• Definition 2.2: Un-normalized Attention
• Remark 1
• Definition 2.3: FC Layer
• Definition 2.4: Transformer
• Definition 2.5: Transformer+
• Lemma 2.1: Lower Bound yu2015useful
• Remark 2
• Theorem 3.1
• Remark 3