Leaner Transformers Can Easily Learn to Cluster

Abstract

Transformers have in-context learning capabilities, where some known learning algorithms can be executed in the forward pass through the model. Recent work shows that transformers can exactly perform Lloyd's algorithm for $k$-means clustering with $n$ points in $d$ dimensions with an embedding size $d_{\textsf{emb}} = d+k$ (thus, requiring attention projection matrices of size $(d+k)^2$). In this work, we build upon this result in the following ways: First, we present an equally expressive but smaller transformer that executes Lloyd's algorithm with embedding size $d_{\textsf{emb}} = (d + \lceil \log_2 k \rceil)$. Next, we train these transformers to learn the clustering algorithms given a distribution of clustering tasks, and theoretically characterize and empirically validate the factors affecting the convergence and in-distribution generalization of learning algorithms based on stochastic gradients. Finally, we probe the general clustering abilities of these learned algorithms (in the form of transformers), and try to understand situations where they succeed and fail.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

More in Machine Learning

• arXiv preprint · 7 October 2026

  Efficient Best-of-N policy evaluation for inference-time alignment

  Jonas Schweisthal, Yuxin Wang, Athiya Deviyani and 2 others
• arXiv preprint · 8 October 2026

  Decision-Sufficient Posterior Approximation

  Sean Plummer
• arXiv preprint · 7 October 2026

  The Symbol of the Surrogate: Measuring Numerical Provenance in Neural PDE Solvers

  Ridham Patel
• arXiv preprint · 8 October 2026

  A General $\widetildeΩ(\sqrt{T γ_T})$ Lower Bound for Kernel Bandits

  Chenkai Ma and Jonathan Scarlett
• arXiv preprint · 7 October 2026

  The Geometry of Anisotropic Dilation for Optimal Regularization

  Carson Newman and Oscar Leong

All new papers in Machine Learning