Transformers learn clustering algorithms with smaller models easily
Leaner Transformers Can Easily Learn to Cluster
Machine Learning
Summary
Clustering is a way to group similar things together automatically. The authors show that a type of AI called transformers can do clustering using fewer resources than before by using smaller internal sizes. They also trained these transformers to perform clustering tasks and studied what helps them learn well and how they generalize to new tasks. Finally, they explored when these AI-based clustering methods work best or might struggle.
What this means in practice
- •For machine learning engineers: Build smaller transformer models that perform clustering efficiently with fewer parameters for resource-constrained environments.
- •For data platform teams: Integrate transformer-based learned clustering algorithms into pipeline components to automatically group data with adaptable models.
Authors
Charlotte Park, Kenneth L. Clarkson, Lior Horesh, Takuya Ito, Parikshit Ram
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.