arXiv：cs.LG· Charlotte Park, Kenneth L. Clarkson, Lior Horesh, Takuya Ito, Parikshit Ram·· 4 小时前AI 评分34
更精简的 Transformer 也能轻松学会聚类
Leaner Transformers Can Easily Learn to Cluster
AI 导读
研究者提出一种嵌入维度为 d+⌈log₂k⌉ 的更小 Transformer，可在前向传播中精确执行 k-means 的 Lloyd 算法，此前方案需要 d+k 的嵌入维度。该工作还训练这些 Transformer 在聚类任务分布上学习聚类算法，并从理论与实验两方面刻画了随机梯度学习算法的收敛性与同分布泛化影响因素。该论文已被 NeurIPS 2026 接收。
正文
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.
| Comments: | NeurIPS 2026 accepted paper |
| Subjects: | Machine Learning (cs.LG); Machine Learning (stat.ML) |
| Cite as: | arXiv:2610.09760 [cs.LG] |
| (or arXiv:2610.09760v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.09760 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Parikshit Ram [view email]
[v1]
Wed, 7 Oct 2026 09:46:01 UTC (4,860 KB)
来源：arXiv：cs.LG · arxiv.org