Finding Clustering Algorithms in the Transformer Architecture

TL;DR

This paper proves that transformers can exactly implement the $k$-means clustering algorithm, providing a neural architecture that mirrors Lloyd's algorithm and enabling interpretable modifications for diverse clustering variants.

Contribution

The paper introduces the $k$-means transformer, a theoretically proven neural architecture that exactly implements Lloyd's $k$-means algorithm using standard transformer components.

Findings

Transformers can precisely implement Lloyd's $k$-means algorithm.

Numerical experiments confirm the exact correspondence between the architecture and Lloyd's algorithm.

Alterations to the architecture produce diverse clustering algorithm variants.

Abstract

The invention of the transformer architecture has revolutionized Artificial Intelligence (AI), yielding unprecedented success in areas such as natural language processing, computer vision, and multimodal reasoning. Despite these advances, it is unclear whether transformers are able to learn and implement precise algorithms. Here, we demonstrate that transformers can exactly implement a fundamental and widely used algorithm for $k$-means clustering: Lloyd's algorithm. First, we theoretically prove the existence of such a transformer architecture, which we term the $k$-means transformer, that exactly implements Lloyd's algorithm for $k$-means clustering using the standard ingredients of modern transformers: attention and residual connections. Next, we numerically implement this transformer and demonstrate in experiments the exact correspondence between our architecture and Lloyd's algorithm, providing a fully neural implementation of $k$-means clustering. Finally, we demonstrate that interpretable alterations (e.g., incorporating layer normalizations or multilayer perceptrons) to this architecture yields diverse and novel variants of clustering algorithms, such as soft $k$-means, spherical $k$-means, trimmed $k$-means, and more. Collectively, our findings demonstrate how transformer mechanisms can precisely map onto algorithmic procedures, offering a clear and interpretable perspective on implementing precise algorithms in transformers.