# Attention-based clustering

arXiv: https://arxiv.org/abs/2505.13112 · HTML: https://mlsanity.com/paper/2505.13112/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 28 references that are in the map.
- First posted on arXiv 19 May 2025.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Stand-Alone Self-Attention in Vision Models](https://mlsanity.com/paper/1906.05909/) (2019)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [FLatten Transformer: Vision Transformer using Focused Linear Attention](https://mlsanity.com/paper/2308.00442/) (2023)
- [Formal Algorithms for Transformers](https://mlsanity.com/paper/2207.09238/) (2022)
- [The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit](https://mlsanity.com/paper/2306.17759/) (2023)
- [Transformers are Universal In-context Learners](https://mlsanity.com/paper/2408.01367/) (2024)
- [Simplifying Transformer Blocks](https://mlsanity.com/paper/2311.01906/) (2023)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)
- [One-Layer Transformer Provably Learns One-Nearest Neighbor In Context](https://mlsanity.com/paper/2411.10830/) (2024)
- [Transformers versus the EM Algorithm in Multi-class Clustering](https://mlsanity.com/paper/2502.06007/) (2025)

## Built on by (most cited first, 3 of 3)

- [Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learning](https://mlsanity.com/paper/2602.04872/) (2026)
- [Dynamic Chunking for Diffusion Language Models](https://mlsanity.com/paper/2605.15676/) (2026)
- [Attention-based PCA](https://mlsanity.com/paper/2605.18315/) (2026)

## Abstract (arXiv)

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.