PaperlayerView exact revision on arXiv

Source-linked AI summary

Leaner Transformers Can Easily Learn to Cluster

Charlotte Park, Kenneth L. Clarkson, Lior Horesh, Takuya Ito, Parikshit Ram

arXiv:2610.09760v1cs.LGstat.ML

TL;DR

Clustering is a discrete, non-differentiable, NP-hard problem. The paper builds a smaller transformer that exactly executes Lloyd’s k-means algorithm and studies learned clustering transformers, which improve clustering iteratively and execute clustering on different-from-training tasks, with exceptions.

  • Problem

    The clustering objective is discrete and non-differentiable, and the problem is NP-hard even for k = 2.

    Source: §1 · p. 1

  • Method

    The paper presents a smaller transformer that exactly executes Lloyd’s k-means algorithm and trains transformers on clustering-task distributions while analyzing convergence and in-distribution generalization.

    Source: §1 · p. 2

  • Results

    The learned transformers progressively improve clustering under iterative application and execute clustering on different-from-training tasks, although naive iteration fails and exceptions exist.

    Source: §5 · p. 12

  • Takeaways & Limitations

    Learned transformers can handle clustering beyond their training tasks, but the reported exceptions preclude a universal generalization claim.

    Source: §5 · p. 12

    The guarantee is weak when this gap is large relative to the other terms.

    Source: §3 · p. 6

Abstract

from arXiv · show

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.

1 Transformers Can Execute Clustering Algorithms In-Context

Clustering is a hard, broadly useful task, and prior work showed transformers can execute learning algorithms in-context. This section frames the open question of when such algorithms can be learned and generalized, then outlines the paper’s contributions to clustering.

1 Transformers Can Execute Clustering Algorithms In-Context

  • Discrete k-means has a non-differentiable combinatorial objective and is NP-hard even for k=2 [MacQueen, 1967, Lloyd, 1982].Source: §1 · p. 1
  • The paper asks how and when transformers that execute algorithms in-context can learn those algorithms and what algorithmic generalization to expect.Prior work established this form of algorithmic expressivity, leaving learnability and generalization as central questions.Source: §1 · p. 2
  • The paper introduces an exact Lloyd-executing transformer with embedding size d + ⌈log2 k⌉ and studies training and in-distribution generalization across clustering tasks.It theoretically characterizes and empirically validates factors affecting convergence and generalization, including token embeddings and smoothing.Source: §1 · p. 2

2 Expressing a Clustering Algorithm with a Leaner Transformer

A binary encoding of cluster indices lets a transformer execute Lloyd’s algorithm exactly with fewer embedding dimensions than the one-hot construction. The reduced dimension lowers parameter dependence on k, with runtime and memory benefits most apparent at larger problem sizes.

2 Expressing a Clustering Algorithm with a Leaner Transformer

  • Binary encoding represents each cluster index with ⌈log2 k⌉ bits rather than the k dimensions required by one-hot encoding.The point-token update encodes each point’s closest-center assignment; subsequent updates modify the center-token embeddings.Source: §2 · p. 3
  • The binary-embedded transformer exactly matches Lloyd’s algorithm after t iterations for every t ≥ 1, using embedding dimension d + ⌈log2 k⌉.The theorem uses inverse softmax temperature γ = ∞.Source: §2 · p. 4
  • The architecture uses O((d + log2 k)^2) parameters, reducing the dependence on k from Ω(k^2) to O((log2 k)^2).When k ≈ d, the cited projection-matrix comparison gives an asymptotic parameter ratio approaching 0.25.Source: §2 · pp. 2, 4
  • With k sufficiently larger than d, the binary model achieves up to 12×/9× forward/backward speedups and almost 70% lower memory usage.Runtime gains are significant when d and k are large enough; Figure 1 reports memory gains across its tested settings.Source: §2 · pp. 4–5, Figure 1 · p. 4

3 Understanding Learning Convergence and Generalization

The section studies transformers trained on distributions of clustering tasks using a differentiable smoothed upper bound on the k-means loss. Theory and experiments examine how regularization, penalty, and token embeddings affect training convergence and in-distribution generalization.

3 Understanding Learning Convergence and Generalization

  • Training samples clustering tasks from a task distribution and uses batch gradients on a differentiable surrogate upper-bounding the discrete objective.The scheme also applies to other discrete clustering objectives when a reasonably smoothed upper bound is available.Source: §3 · p. 6
  • The smoothed-loss Lipschitz constant depends on τ, the regularizer, attention inverse temperature, and embedding scheme; one-hot has a better bound than binary, trading fewer parameters for convergence.The convergence difference is limited when d is large relative to k, and faster binary forward/backward passes do not necessarily mean faster training to convergence.Source: §3 · p. 7
  • For every tested τ, L2-regularized min gives lower training loss and a tighter upper bound; at τ=1, only L2 reduces unseen-task loss as training progresses.With NE regularization at τ=1, validation loss increases during training.Source: §3 · p. 8
  • At τ=0.1, both regularizers outperform one Lloyd step except with no token embedding; higher τ speeds convergence but can yield worse solutions.For fixed τ=0.25, increasing d removes the convergence difference between one-hot and binary embeddings.Source: §3 · p. 9
  • L2-regularized loss surfaces appear fairly unimodal; one-hot and binary surfaces differ at d=4 but are similar at d=32.The authors connect the dimensionality pattern to the relative differences in embedding-related terms in their theorem.Source: §3 · p. 9
  • Across five tested distribution families, learned models generalize except on Cauchy tasks, where validation is worse than one Lloyd iteration; the authors hypothesize a larger surrogate-to-k-means loss gap.The Cauchy failure remains an open question.Source: §3 · p. 9

4 Probing Algorithmic Generalization Beyond Training Distributions

The learned models can improve clustering across multiple steps and generalize across several changes to task distributions, but naive iteration and Cauchy-family tasks expose important exceptions.

4 Probing Algorithmic Generalization Beyond Training Distributions

  • Re-embedding the input at each step improves multi-step clustering, with OH/BN models outperforming Lloyd’s iterations while naive repeated updates diverge.The models were trained for single-step clustering; models without token embeddings improve with re-embedding but underperform Lloyd’s.Source: §4 · p. 10, Figure 6 · p. 10
  • Across changes to normal-mixture properties, OH/BN models significantly outperform Lloyd’s except at larger scale, where they trail initially but overtake with more steps.The models were trained on equally weighted isotropic normal mixtures and evaluated on unequal weights, anisotropy, and smaller or larger scales.Source: Figure 7 · p. 10, §4 · p. 11
  • Across non-Cauchy distribution-family evaluations, OH/BN models generalize well and significantly outperform Lloyd’s, but Cauchy training diverges and Cauchy testing remains noncompetitive.The paper reports that Cauchy tasks are dissimilar to the other families, while the other families are similar to one another; these task dissimilarities align with observed generalization.Source: §4 · p. 11, Figure 8 · p. 11

5 Discussion

The discussion situates the work among transformer-based algorithm learning and distinguishes its clustering focus from prior research on token dynamics. It summarizes the leaner expressivity result and learned algorithms’ generalization, while noting limits in dimensionality and Cauchy-task understanding.

5 Discussion

  • The paper studies discrete clustering rather than the supervised learning tasks common in transformer in-context learning, where the objective is not natively differentiable.It examines the effects of smoothing the clustering objective.Source: §5 · p. 11
  • The expressivity result relies on softmax attention at zero temperature, which the authors call practically infeasible; experiments instead use standard softmax attention.The discussion builds on Clarkson et al. [2026].Source: §5 · p. 12
  • Unlike work studying self-attention token dynamics and their stabilization, this paper trains a single transformer layer to learn clustering algorithms from task data.The cited line of work studies representations that self-cluster, potentially down to a single cluster.Source: §5 · p. 12
  • Figure 9 compares 20 tasks using pairwise Earth-mover distances between distributions of point-to-initial-center distances, arranged in family-based blocks.The four tasks per family span normal, Cauchy, Gumbel, Laplace, and lognormal distributions.Source: Figure 9 · p. 12
  • The conclusion identifies two limits: learned algorithms are restricted to a specific data dimensionality d and cluster count k, and Cauchy-task failures remain poorly understood.The paper contrasts this restriction with hand-crafted algorithms, which are more general.Source: §5 · p. 12

A Limitations

The paper identifies unresolved questions about why learning and generalization fail, especially for Cauchy tasks, and limits its architecture to fixed d and k. It also leaves open whether extra embedding dimensions are necessary and how several training choices affect learning.

A Limitations

  • The architecture is specific to d and k, although it can cluster varying numbers of points n for those fixed values.Zero-padding could allow use with smaller d, but training cost scales quadratically with the larger dimension, and strong generalization across d is uncertain.Source: A Limitations · p. 17
  • Whether the extra embedding dimensions are necessary for competitive learning remains theoretically unresolved, despite empirical results favoring models with them over NA models.The paper also leaves open whether training directly for multi-step clustering can efficiently yield a better model.Source: A Limitations · p. 17
  • The learned algorithm’s precise behavior remains unidentified, and the authors lack a clear explanation for learning and generalization failures in some settings.For Cauchy tasks, they hypothesize that outliers can degrade the regularized-min surrogate’s upper bound, but say the explanation and mitigation remain unclear.Source: A Limitations · pp. 17–18
  • The convergence analysis assumes β-smooth training losses, but sparsemax is only differentiable almost everywhere and is not generally β-smooth.The authors also leave adaptive or annealed attention temperature and regularized-min penalty, and the learned algorithm’s identity, for future work.Source: A Limitations · p. 18

B Leaner Transformers Can Emulate Lloyd’s Iteration

With inverse temperature γ = ∞ and no empty clusters during Lloyd iterations, the constructed transformer exactly reproduces Lloyd’s updates for any number of iterations. Its embedding uses d + ⌈log2 k⌉ dimensions to maintain both assignments and updated centers.

B Leaner Transformers Can Emulate Lloyd’s Iteration

  • The proof adapts Clarkson et al. [2026, Theorem 2] and uses averaging hard attention as softmax inverse temperature tends to infinity.The point-assignment update follows that prior construction, while the center update is different.Source: B Leaner Transformers Can Emulate… · pp. 18, 20
  • For any t, the constructed transformer exactly matches t Lloyd iterations when γ = ∞ and no iteration produces an empty cluster.The theorem uses d + ⌈log2 k⌉ embedding dimensions.Source: B Leaner Transformers Can Emulate… · p. 18
  • Hardmax attention assigns each point to its nearest center, while averaging hardmax forms updated centers from their assigned points.The center update assumes no empty clusters during Lloyd iterations.Source: B Leaner Transformers Can Emulate… · pp. 19–20
  • The leaner embedding maintains both Lloyd’s cluster assignments and center updates using d + ⌈log2 k⌉ dimensions.The center-update construction uses negative squared Euclidean distance attention rather than the linear attention in Clarkson et al. [2026].Source: B Leaner Transformers Can Emulate… · p. 20

C Single-Step and Multi-Step Clustering

The model is trained on one clustering update, then evaluated by repeatedly applying it to evolving data and centers. Iteration can improve clustering quality, but the learned token embeddings and attention updates are both needed for this behavior.

C Single-Step and Multi-Step Clustering

  • The authors explore multi-step clustering because iterative refinement is natural for clustering algorithms, while one-shot output only reduces a surrogate of the k-means objective.Lloyd’s algorithm and PAM are cited as iterative examples.Source: C Single-Step and Multi-Step Clustering · p. 20
  • Exact theoretical equivalence for repeated updates requires constructive-proof parameters and γ → ∞, unlike the trained model’s finite γ.The paper distinguishes this condition from its finite-temperature trained model.Source: C Single-Step and Multi-Step Clustering · p. 21
  • Repeatedly applying the trained model to evolving centers progressively improves clustering quality, whereas recursively applying attention updates alone does not.The comparison indicates that both token embeddings and attention updates are essential for successful iterative clustering.Source: C Single-Step and Multi-Step Clustering · p. 21
  • During training, each task receives one model update and the loss is computed from C(1), without recursion.At evaluation, the same model is reapplied to the data and current centers for T steps.Source: C Single-Step and Multi-Step Clustering · p. 21

D Clustering Task Generation

Training tasks are generated from mixtures of distributions, with sampled centers and optional anisotropic scales and unequal weights. Data and initial centers are scaled into the unit cube, which is compatible with k-means translation and scaling invariance.

D Clustering Task Generation

  • Tasks sample mixture centers uniformly from [0, 1]^d, with scales and mixture weights varied when using anisotropic or unequally weighted mixtures.The method samples n points from the mixture, applies MinMaxScaler, and chooses k sampled data points as initial centers.Source: D Clustering Task Generation · pp. 21–22
  • The generated data and initial centers lie in [0, 1], aligning the task distribution with the theoretical analyses.The paper relies on translation and scaling invariance to map general k-means tasks into this range.Source: D Clustering Task Generation · p. 22
  • The task-generation procedure shifts all points and centers by the combined per-dimension minimum, then scales by their largest shifted value.This preserves the k-means problem under translation and scaling; the all-zero case is trivial.Source: D Clustering Task Generation · p. 22

E Computational Resources and Implementation

The experiments use 10 random seeds and aggregate results across trials, with substantial model-training compute and specified CPU and GPU hardware.

Computational resources and implementation

  • 3,000 models were trained across 300 configurations and 10 repetitions, at about 15–20 minutes per model and roughly 1,000 total training hours.Loss-surface generation for Figure 3 took about 1 hour total, and inference evaluations took about 150 minutes total.Source: E Computational Resources and… · p. 22
  • Experiments ran on a 16-thread Intel i7 Core CPU with 64 GB memory and an NVIDIA V100 GPU with 32 GB memory.Timed runs for Figure 1, Table 1, and Figure 17 instead used an NVIDIA RTX PRO 3000 Blackwell Generation Laptop GPU.Source: E Computational Resources and… · p. 22
  • The implementation is in PyTorch, and results are aggregated across 10 random seeds.The paper states that the implementation is available in a GitHub repository.Source: E Computational Resources and… · p. 22

F Ablative Experiments

The ablations vary the training loss, optimizer learning rate, task size, dimension, cluster count, and data distribution. They show that validation performance depends on these choices, while learned models can generalize across task lengths and, in some settings, outperform a single Lloyd iteration.

Varying the min-regularization penalty

  • As τ decreases, the smoothed upper-bound tightens and validation performance can match or exceed one-step Lloyd, although higher τ yields faster training-loss convergence without guaranteeing strong validation performance.The L2-regularized min gives tighter bounds and better validation performance than the NE-regularized min, especially at larger τ and d = 4.Source: F Ablative Experiments · p. 23, Figure 10 · p. 23

Varying the learning rate

  • Learning rate 0.001 converges monotonically but more slowly than 0.01, while rates 0.1 and 1.0 can produce nonmonotonic losses and worse local minima.For d = 32, rates 0.1 and 1.0 converge to worse local minima; for d = 4, this is reported for rate 1.0. Where training converges favorably, L2-regularized min outperforms NE-regularized min on training and validation loss.Source: F Ablative Experiments · pp. 23–24, Figure 11 · p. 24

Length generalization

  • Models trained on n = 512 points generalize to tasks with 128, 256, 1024, and 2048 points, and OH- and BN-embedded models outperform Lloyd’s iterations at each tested size.The evaluation uses T = 20 clustering steps on tasks from the same distribution family as training.Source: F Ablative Experiments · pp. 24–25, Figure 12 · p. 24

Varying training-task size and dimension

  • Training-task point count n has no significant effect on convergence or in-distribution generalization, and changing training-task dimensionality shows no clear adverse effect from the model’s quadratic parameter growth with d.The study varies n across 128, 256, 512, and 1024, and d across 4, 8, 16, and 32; the authors suggest training with smaller n as computationally efficient while retaining use on larger tasks.Source: F Ablative Experiments · p. 25, Figure 13 · p. 25, Figure 14 · p. 26

Varying the number of clusters

  • For d = 4 and τ = 1, across k ∈ {6, 10, 16, 25}, L2-regularized min learns somewhat but never matches one Lloyd iteration, while NE-regularized min fails to learn; BN deteriorates as k rises, unlike OH.Source: F Ablative Experiments · p. 26
  • With d = 4 and τ = 0.1, OH with L2-regularized min beats Lloyd for all tested k, while BN performance slows as k grows; at d = 32, τ = 1, OH and BN with L2 beat Lloyd across k.At d = 32 and τ = 0.1, both regularizers with OH and BN beat Lloyd on validation tasks; OH has numerical issues in 3 of 10 trials at k = 25, though its median still beats Lloyd.Source: F Ablative Experiments · pp. 26–28

Varying the data distribution

  • For cauchy-distributed tasks, the proposed learning scheme fails to learn and generalize in the tested configurations; with τ = 0.1, both regularizers and OH or BN embeddings can beat Lloyd across the other tested distributions.At τ = 1, L2 generalizes better than NE, and NE fails to generalize on gumbel and laplace tasks.Source: F Ablative Experiments · pp. 28–29, Figure 16 · p. 28

Computational overhead and convergence

  • BN’s smaller embeddings reduce computational overhead, but the theoretical convergence guarantees indicate slower optimization than OH; in elapsed-time comparisons, the faster-converging embedding depends on the d, k setting.At d = k = 16, BN converges more slowly but both methods eventually perform close to each other and better than one-step Lloyd; at d = 32, k = 25, BN initially converges faster in runtime, and both match or outperform Lloyd by 10,000 steps.Source: F Ablative Experiments · p. 29, Figure 17 · p. 29

G Additional Figures

The additional figures examine cross-distribution generalization and pairwise task dissimilarities. They report that cauchy tasks are similar mainly to other cauchy tasks and are substantially dissimilar from other families.

Generalization to different distribution families

  • The 5 × 5 evaluation compares models trained on each of five distribution families against multi-step clustering tasks from each test family, with diagonal entries using the same family for training and testing.Training uses single-step clustering, evaluation uses T = 20 steps, and curves aggregate inter-quartile ranges over 10 trials; Lloyd’s performance is plotted in black.Source: Figure 18 · p. 30

Pairwise task dissimilarities

  • In the 80 × 80 task-distance matrix, cauchy tasks are similar mainly to other cauchy tasks and substantially dissimilar from other families; laplace appears least dissimilar to cauchy, but remains quite dissimilar.Rows and columns represent 16 tasks from each of five families, ordered normal, cauchy, gumbel, laplace, and lognormal; lighter shades indicate lower Earth-mover distances.Source: Figure 19 · p. 31
  • The pairwise matrix measures Earth-mover distances between task distributions of squared point-to-initial-center distances, with zero self-distance shown on the diagonal.Source: Figure 19 · p. 31

H Smoothed Clustering Objective via Regularized min

The appendix constructs differentiable upper-bounds for the discrete k-means objective by regularizing the minimum over center distances. It compares entropy-based softmin with L2-based sparsemax and discusses bound tightness and smoothness.

Regularized minimum

  • The discrete per-point k-means loss takes the minimum distance from a point to a center, and regularizing this minimization gives a differentiable surrogate for training.For each point, the vector r contains its distances to the k centers; the unregularized objective remains non-differentiable.Source: H Smoothed Clustering Objective via… · p. 32

Entropy regularization

  • Negative Shannon entropy produces softmax weights and a soft-min objective whose upper-bound reaches the discrete loss as τ approaches zero.The learning procedure uses this differentiable objective instead of the original non-differentiable loss.Source: H Smoothed Clustering Objective via… · p. 32

L2 regularization

  • L2 regularization produces sparsemax weights, which can assign exactly zero weight to some centers and yields a differentiable upper-bound.The weights are obtained by projecting the negative scaled distances onto the probability simplex.Source: H Smoothed Clustering Objective via… · pp. 32–33

Comparing upper-bounds

  • For the same τ > 0, the sparsemax-based upper-bound is tighter than the softmax-based upper-bound, and the L2 bound can equal the discrete loss at positive τ.The bound comparison uses the distance between the exact minimum weights and regularized-min weights; the sparsemax bound is reported as tighter than the softmax bound.Source: H Smoothed Clustering Objective via… · p. 33

Lipschitz continuity

  • The smoothed objective’s Lipschitz constant generally grows as τ decreases and also depends on the regularized-min operation’s input-stability parameter.For softmax, the stated input-stability parameter also increases as τ decreases.Source: H Smoothed Clustering Objective via… · pp. 33–34

I Lipschitz-ness of the Smoothed Upperbound of the Clustering Objective

Under bounded-input and bounded-parameter assumptions, this section derives stability bounds for transformer outputs and a Lipschitz bound for the smoothed clustering loss with respect to learnable parameters. The bounds are built from attention-module sensitivities and bounded-dispersion softmax behavior, with the loss-bound constant depending on the regularizer and penalty, attention inverse temperature, problem size, embedding dimension, and initial-center embedding norm.

I Lipschitz-ness of the Smoothed Upperbound of the Clustering Objective

  • Lemma I.2 establishes that the smoothed loss is stable to changes in the model output, under its stated assumptions and the conditions of Lemma I.1.Lemma I.1 supplies bounded norms for initial and intermediate embeddings and cluster centers under both the OH and BN embedding schemes.Source: I Lipschitz-ness of the Smoothed… · pp. 34–35
  • Under Assumption I.2, the specified attention scores have dispersion at most ρ′ for every i ∈ [m′]; when D(v), D(v′) ≤ δ, softmax is bounded and input-stable.Source: I Lipschitz-ness of the Smoothed… · p. 37
  • Attention stability bounds cover perturbations to value, query, and key projections, as well as changes to the keys; Table 2 summarizes factors affecting these bounds.The factors include query and key counts, norm bounds, and projection-matrix bounds.Source: I Lipschitz-ness of the Smoothed… · pp. 37, 38, 39, 41
  • Theorem I.1 bounds changes in the model output under parameter perturbations, assuming bounded inputs and bounded projection-matrix norms.The projection bounds apply to query, key, and value matrices.Source: I Lipschitz-ness of the Smoothed… · pp. 34, 41
  • For fixed k, increasing the embedding dimension demb and decreasing the initial-center embedding norm dE reduce the bound constant ϕ.The derivation identifies k√dE and dE√demb as recurring quantities when treating Γ1 and Γ2 as architecture- and problem-independent constants.Source: I Lipschitz-ness of the Smoothed… · pp. 46–47
  • Theorem I.2 establishes Lipschitz-ness of the smoothed clustering loss with respect to learnable parameters, assuming bounded points, initial centers, and parameters.The theorem combines the preceding bound with a lemma on the smoothed loss.Source: I Lipschitz-ness of the Smoothed… · p. 47
  • The Lipschitz bound highlights dependencies on the min-regularizer Ω, penalty τ, attention inverse temperature γ, point count n, cluster count k, demb, and initial-center embedding norm dE.Source: I Lipschitz-ness of the Smoothed… · p. 47
Loading 2610.09760v1…