NeurIPS 2025

Attention-based clustering
Rodrigo Maulen-Soto∗
Sorbonne Université, LPSM
Paris, France
Pierre Marion
EPFL, Institute of Mathematics
Lausanne, Switzerland
Claire Boyer∗
Laboratoire de Mathématiques d’Orsay
Université Paris Saclay
Institut Universitaire de France
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 anal-
ysis 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 underly-
ing 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.
1
Introduction
Attention-based models (Bahdanau et al., 2015), in particular Transformers (Vaswani et al., 2017),
have achieved state-of-the-art performance across a wide range of learning tasks. These include
applications in natural language processing (Devlin et al., 2018; Bubeck et al., 2023; Luong et al.,
2015; Bahdanau et al., 2016) and computer vision (Dosovitskiy et al., 2020; Liu et al., 2021;
Ramachandran et al., 2019). The success of the attention mechanism has been linked to its ability to
capture long-range relationships in input sequences (Bahdanau et al., 2015; Vaswani et al., 2017).
They do this by computing pairwise dependencies between tokens based on learned projections,
without regard to the t...