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 dtextsfemb = 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 dtextsfemb = (d + ⌈ log2 k ⌉). 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.
Related
- One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought
- Learn-Then-Differentiate Gradient Estimation
- Wait to Learn: Timing the Gradient Return in Over-the-Air Split Learning over an Aging Single-Antenna Channel
- What can linear attention learn from nonlinear teachers in-context?
- Learn Your Own Thoughts: Abstract Token Curriculum
BibTeX
@misc{park2026leaner,
title = {{Leaner Transformers Can Easily Learn to Cluster}},
author = {Charlotte Park and Kenneth L. Clarkson and Lior Horesh and Takuya Ito and Parikshit Ram},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.09760},
eprint = {2610.09760},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.09760},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Leaner Transformers Can Easily Learn to Cluster
AU - Charlotte Park
AU - Kenneth L. Clarkson
AU - Lior Horesh
AU - Takuya Ito
AU - Parikshit Ram
PY - 2026
DA - 2026/10/16
PB - arXiv
DO - 10.48550/arxiv.2610.09760
UR - https://arxiv.org/abs/2610.09760
AB - 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_textsfemb = 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_textsfemb = (d + ⌈ log_2 k ⌉). 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.
ER -