Leaner Transformers Can Easily Learn to Cluster

2026-10-07

AI Summary

This study addresses the embedding dimension redundancy and unclear learning mechanisms when Transformers perform K-means clustering. We propose compressing the embedding dimension to d+log k, constructing a minimal Transformer that efficiently executes Lloyd's algorithm. By integrating stochastic gradient optimization with probing techniques, we theoretically characterize the convergence properties and in-distribution generalization conditions of the learned algorithm. Our experiments validate the effectiveness of this streamlined model on clustering tasks and identify the critical factors governing its success and failure. Ultimately, this work provides both theoretical and empirical foundations for understanding the algorithmic learning mechanisms underlying Transformers.

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.

Problem

• Transformers
• in-context learning
• k-means clustering
• Lloyd's algorithm
• generalization

Innovation

• In-context learning
• Transformers
• k-means clustering
• Lloyd's algorithm
• Generalization