Transformer Circuits Can Realize Clustering Algorithms

Although transformers are most commonly optimized as statistical sequence models, it is unclear to what extent they can implement and learn exact algorithmic computations. Here, we specify a transformer implementation from first principles that executes a fundamental and widely used method for $k$-means clustering: Lloyd's algorithm. We theoretically prove and empirically demonstrate that this implementation of a transformer architecture, which we term the _$k$-means transformer_, exactly implements Lloyd's algorithm for $k$-means clustering using the standard circuit mechanisms of modern transformers: attention block, residual connections, and feed-forward block. In learning experiments, we find that training this base architecture on $k$-means clustering yields a generalizable clustering algorithm that surpasses Lloyd's algorithm in terms of clustering quality. Finally, we demonstrate that interpretable alterations (e.g., inclusion of layer normalizations) to this architecture yields diverse and novel variants of clustering algorithms, including soft $k$-means, spherical $k$-means, trimmed $k$-means. Overall, our results show that transformer circuit mechanisms can instantiate exact algorithmic routines for clustering, while simultaneously providing an effective learnable model.

Transformer models are the computational core of much of GenAI, and are often assumed to have "in-context learning" capabilities, being able to learn and adapt on the fly given input examples without requiring any explicit model updates. This capability is often considered critical in the general purpose utility of these models. However, it is not well understood what human-interpretable learning mechanisms (AKA algorithms) can these transformers precisely express. We show that these transformers are well equipped to express the most popular and well-studied learning algorithm for clustering. Furthermore, standard modifications commonly considered in the transformer translate to other known clustering algorithms, and in some cases, even novel algorithms. Finally, this connection between transformers and clustering algorithms allows the transformer to learn data-dependent clustering algorithms that are distinct from known ones, and provide better clustering quality. Our study provides a concrete bridge between transformers and exact algorithmic computations, allowing us to better understand what algorithms and learning mechanisms transformers can represent and learn, shedding light on their in-context learning capabilities.