Attention-based clustering

TL;DR

This paper analyzes attention-based layers through the lens of unsupervised clustering on Gaussian mixtures. It shows that a two-headed linear attention predictor can learn the true mixture centroids via population risk minimization and projected gradient updates on the unit sphere, under suitable initialization and temperature settings. The authors establish convergence guarantees, quantify the quantization properties of trained heads, and demonstrate that even an identity-key/value attention module can perform in-context quantization without learned parameters. They further explore regularization strategies to improve robustness to initialization and extend the analysis to multiple clusters, random initializations, and higher-dimensional settings, validating results with experiments. The work highlights the Transformer’s capacity for adaptive, distribution-aware representation learning beyond supervised tasks, and points to principled directions for extending these ideas to richer architectures and in-context learning scenarios.

Abstract

Transformers have emerged as a powerful neural network architecture capable of tackling a wide range of learning tasks. In this work, we provide a theoretical analysis of their ability to automatically extract structure from data in an unsupervised setting. In particular, we demonstrate their suitability for clustering when the input data is generated from a Gaussian mixture model. To this end, we study a simplified two-head attention layer and define a population risk whose minimization with unlabeled data drives the head parameters to align with the true mixture centroids. This phenomenon highlights the ability of attention-based layers to capture underlying distributional structure. We further examine an attention layer with key, query, and value matrices fixed to the identity, and show that, even without any trainable parameters, it can perform in-context quantization, revealing the surprising capacity of transformer-based methods to adapt dynamically to input-specific distributions.

Figures (15)

• Figure 1: Distance to centroids vs \ref{['PGDrho']} iterations for the minimization of $\mathcal{R}$ (therefore without regularization), with an initialization on the manifold $\mathcal{M}$. 10 runs, 95% percentile intervals are plotted.
• Figure 2: Performance of \ref{['PGDrho']}, when initializing on the unit sphere and minimizing the regularized risk \ref{['penalizationreal']}. 10 runs, 95% percentile intervals are plotted.
• Figure 3: Comparison between inputs and embeddings ($d=10, \sigma=0.3$, $L=500$, $\lambda=\frac{1}{1+2\sigma^2}$, $\mu_0^\star=e_{10}, \mu_1^\star=-e_1$, $e_j$ is the $j$-th vector of the canonical basis of $\mathbb{R}^{10}$). PCA fitted on input tokens was used to project both input and transformed tokens to 2D.
• Figure 4: Comparison between inputs and embeddings ($d=10, \sigma=0.3$, $L=500$, $\lambda=\frac{1}{1+2\sigma^2}$, $\mu_0^\star=e_{10}, \mu_1^\star=-e_1$). PCA was fitted on the input tokens and was used to project both input and transformed tokens to 2D.
• Figure 5: Distance to centroids vs \ref{['PGDrho']} iterations for the minimization of $\mathcal{R}$, with data drawn from the degenerate case \ref{['def:dirac_mixture_model']}. 10 runs, 95% percentile intervals are plotted.

Theorems & Definitions (58)

• Proposition 3.1
• Lemma 3.2
• Proposition 3.3
• Theorem 3.4
• Proposition 4.1
• Proposition 4.2
• Proposition 5.1
• Proposition 5.2
• Proposition A.1
• Remark A.2