Research map / Learning theory and algorithms / Clustering

Leaner transformers can easily learn to cluster

NeurIPS 2026 · Learning theory and algorithms › Clustering · approach: k-means++ seeding analysis

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows a leaner transformer (embedding d+ceil(log2 k), binary cluster indices) can exactly execute Lloyd's k-means, then trains it end-to-end with smoothed (softmax/sparsemax) min surrogates to learn single-step clustering, analyzing convergence, generalization, and cross-distribution transfer.

Problem

Transformers with embedding size d+k needed to express Lloyd's k-means have projection matrices quadratic in k; it is unclear how learnable such clustering algorithms are with a discrete, non-differentiable k-means objective.

Approach

Replace one-hot cluster-index embeddings with binary (ceil(log2 k)-bit) embeddings so a single attention layer with embedding dimension d+ceil(log2 k) still exactly emulates Lloyd's iteration, and train it with a smoothed min surrogate (softmax or sparsemax regularized min) whose gap to the true k-means loss controls generalization.

Key results

• Binary embedding reduces embedding dimension from d+k to d+ceil(log2 k); e.g. for d=32, k=1000 the forward pass speeds up about 12x and GPU memory drops about 69%.
• Learned single-layer transformers with L2 (sparsemax) min regularization outperform a single Lloyd's iteration on unseen tasks for normal-type families at tau=0.1 and tau=1.
• Models trained on cauchy tasks fail to generalize, even to single-step cauchy clustering, and diverge on multi-step tasks.

Surprising findings

• Naively iterating the trained single-step model does not improve clustering; only re-embedding inputs at each step yields gains over Lloyd's.
• The leaner binary-embedding model converges more slowly than the one-hot model for small d, despite fewer parameters.

Limitations

• Model is tied to a specific data dimension d and cluster count k; cannot generalize across k or d without retraining.
• Failure to train and generalize on cauchy (heavy-tailed) distributions is not explained.
• Sparsemax surrogate is not generally beta-smooth, creating a mismatch with the theory; softmax attention used instead of hardmax limit.

Experimental setup

Datasets: synthetic mixture-of-distributions clustering tasks (normal, cauchy, gumbel, laplace, lognormal)

Benchmarks: synthetic clustering tasks (in-distribution validation) multi-step clustering T=20 on shifted synthetic distributions

Metrics: k-means objective forward/backward time GPU memory usage

Baselines: Lloyd's algorithm (single step) one-hot embedding (d+k) no index embedding (d+0)

Builds on: Lloyd's algorithm transformers executing Lloyd's algorithm in-context (d+k embedding) sparsemax Adam

Closest papers by meaning

• Transformer Circuits Can Realize Clustering Algorithms (ICML 2026 · similarity 0.92)
• Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers (NeurIPS 2026 · similarity 0.85)
• Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.84)
• Clustering via Self-Supervised Diffusion (ICML 2025 · similarity 0.84)
• Attention-based clustering (NeurIPS 2025 · similarity 0.84)
• Transformer Learns Optimal Variable Selection in Group-Sparse Classification (ICLR 2025 · similarity 0.84)