Research map / Architectures and efficiency / Attention and transformers
Attention-based clustering
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies linear (and softmax-free) attention heads for unsupervised clustering of Gaussian mixtures, showing projected gradient descent on a quantization-style risk recovers true centroids and that a parameter-free identity-QKV attention performs in-context quantization.
Problem
Understanding whether attention layers can extract cluster structure in an unsupervised setting, where tokens are drawn from a Gaussian mixture without labels, and how training dynamics behave.
Approach
Model a two-head linear attention layer whose key/query vectors act as cluster centroid estimators; train it by minimizing a quantization-style reconstruction risk with projected gradient descent on the sphere, and show that with identity values and parameter-free attention (in-context setting) the layer still performs approximate quantization.
Key results
- Projected gradient descent on the linear two-head attention risk converges to the true centroids (up to sign) from initializations on the manifold; experiments show linear convergence over the first 10^3 iterations, plateauing near 1e-2 for sigma=0.3 and 1e-1 for sigma=1.
- For oracle parameters with L to infinity, risk ratio to the optimal quantizer tends to 1-2/d, so the attention predictor achieves lower risk than the optimal quantizer, though the comparison is not like-for-like.
- The parameter-free in-context attention layer (identity QKV, d heads with orthogonal parameters) attains unbiased embedding with risk bounded by sigma^2(d-2) and conditional variance below that of the inputs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Initialize the two-head attention parameters on the manifold M rather than randomly outside it (ablation) | Two-head linear attention trained without regularization · Gaussian mixture | distance to centroids: around 10^{-2} (σ = 0.3); around 10^{-1} (σ = 1) | Helped |
| Add orthogonality regularization to training (ablation) | Two-head linear attention with random unit-sphere initialization · Gaussian mixture | distance to centroids: near 10^{-3} for σ = 0.3; near 10^{-1} for σ = 1 | Helped |
| Increase regularization strength until it overrides the original objective (ablation) | Regularized two-head linear attention · Gaussian mixture | alignment with true centroids | Hurt |
| Use random unit-sphere initialization outside the manifold instead of manifold initialization (ablation) | Two-head linear attention trained without regularization · Gaussian mixture | alignment with true centroids | Hurt |
| Use standard projected SGD heuristics instead of Riemannian updates (ablation) | Two-head linear attention · Gaussian mixture | training behavior | No clear effect |
| Remove the projection step (ablation) | Two-head linear attention optimization · Gaussian mixture | training behavior | Hurt |
| Replace learned two-head parameters with the identity-parameter in-context attention layer (ablation) | Attention-based quantization · In-context Gaussian mixture, L→∞ | risk: \left(1-\frac{2}{d}\right)\frac{1+2\sigma^2}{1+6\sigma^2+12\sigma^4+4d\sigma^6} | Helped |
| Extend the regularized attention model from two clusters to three heads and centroids (ablation) | Regularized linear attention · Three-centroid mixture | centroid recovery | Helped |
Surprising findings
- The oracle linear attention predictor attains lower asymptotic risk than the optimal quantizer, which the authors attribute to comparing functions of different complexity.
- A layer with no trainable parameters performs in-context quantization with lower risk than the optimal quantizer.
Limitations
- Analysis restricted to orthogonal centroids, two clusters, and linear (softmax-free) attention with identity values.
- Convergence guarantee requires initialization on a special manifold; generic sphere initialization needs regularization without theory.
Experimental setup
Datasets: synthetic Gaussian mixture data synthetic Dirac mixture data
Metrics: distance to centroids (up to sign and permutation) quantization risk conditional variance of embedding
Baselines: optimal quantizer T* PCA estimator
Builds on: in-context learning theory of linear attention Gaussian sequence multi-index models projected Riemannian gradient descent
Closest papers by meaning
- Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers (NeurIPS 2026 · similarity 0.85)
- Leaner transformers can easily learn to cluster (NeurIPS 2026 · similarity 0.84)
- Clustering in Causal Attention Masking (NeurIPS 2024 · similarity 0.84)
- Hierarchical Compact Clustering Attention (COCA) for Unsupervised Object-Centric Learning (CVPR 2025 · similarity 0.84)
- On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures (ICML 2025 · similarity 0.83)
- Learning Mixtures of Gaussian Processes through Random Projection (ICML 2024 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.