Leaner Transformers Can Easily Learn to Cluster

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

cs.LGOct 7, 2026

Leaner Transformers Can Easily Learn to Cluster

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

Organizations: MIT · IBM Research

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.

Figures & tables

Figure 1 : Comparison of the proposed ( BN ) embedding against the existing ( OH ) embedding with varying values of $d\in\{4,8,16,32,64,128\}$ and $k\in\{10,16,25,40,64,100\}$ . Each heatmap reports the relative improvement (in %) of ( BN ) over ( OH ) for each of following metrics: (i) Forward time (train) : forward pass time during training, (ii) Backward time (train) : backward pass time, (iv) Encoding specific GPU memory usage : peak GPU memory usage at training (less the usage from the self-attention matrices which does not depend on ${d_{\mathsf{emb}}}$ ), and (iv) Forward pass (eval) : forward pass time with the computation graph disabled (via model.eval() ). Each cell in the $(6\times 6)$ grid in the heatmap lists the relative gain (or degradation), with actual metrics for ( OH ) also listed in the cell as (mean $\pm$ std) for completeness; darker greens denote larger gain values. All runs use 32 clustering problems with $n=1024$ points, and numbers are generated with 30 repetitions (10 warm-up rounds). In terms of memory usage, ( BN ) shows gains across the board, ranging from +1-15%. For the computation times (forward (train), backward, and forward (eval)), ( BN ) does not always show a positive gain for smaller values of $k$ , especially when the baseline ( OH ) runtimes is already quite small ( $<10$ ms). However, once $k$ is large enough, the gains can be quite significant, reaching over 50% in some cases.

$k$	$({d_{\mathsf{emb}}}-d)$	Forward pass (ms)	Backward pass (ms)	GPU memory usage (MB)
	OH / BN	OH / BN [speedup]	OH / BN [speedup]	OH / BN [reduction]
6	6 / 3	4.25 / 4.83 [0.88 $\times$ ]	9.53 / 10.66 [0.89 $\times$ ]	419.93 / 416.29 [0.9%]
25	25 / 5	5.24 / 4.99 [1.05 $\times$ ]	11.21 / 11.01 [1.02 $\times$ ]	441.08 / 423.23 [4.0%]
100	100 / 7	9.76 / 4.52 [2.16 $\times$ ]	20.21 / 9.93 [2.03 $\times$ ]	532.96 / 447.04 [16.1%]
400	400 / 9	38.68 / 7.89 [4.90 $\times$ ]	66.13 / 16.47 [4.02 $\times$ ]	1095.60 / 531.02 [51.5%]
1000	1000 / 10	193.07 / 15.90 [12.14 $\times$ ]	308.28 / 32.33 [9.54 $\times$ ]	3930.34 / 1207.75 [69.3%]

Table 1 : Comparison between proposed ( BN ) and existing ( OH ) under the same measurement setting as in Figure 1 but where

k

is varied in the range

[6,1000]

(column 1) for a fixed data dimensionality

d=32

. We note the actual forward (column 3) and backward pass (column 4) times (and the speedup of ( BN ) over ( OH )), and the peak GPU memory usage (column 5) as described in Figure 1 (and the reduction from using ( BN )). We also note the additional embedding dimensions

({d_{\mathsf{emb}}}-d)

for each value of

k

for the two embeddings (column 2), and the significant reduction in the additional embedding dimensions lead to the significant reductions in the runtimes once

k

is sufficiently larger than

d

with up to 12

\times

\times

speedup in forward/backward pass, and almost 70% reduction in memory usage.

Panel of Figure 2 — Figure 2 : Effect of data dimensionality $d$ , $\min$ -regularizer $\Omega$ with penalty $\tau$ , token embedding $(E_{C},E_{X})$ on training convergence and generalization. We show the convergence of (i) the smoothed training loss $\tilde{\mathcal{L}}_{\Omega}^{\tau}$ (top row in each figure) and (ii) the relative non-differentiable clustering loss $\mathcal{L}$ on held-out validation clustering tasks (bottom row) for multiple configurations; the dotted black line at $y=1$ (bottom row) corresponds to the performance of a single Lloyd’s iteration on the validation tasks. See Appendix F for details and ablations. In all cases, lower values on the vertical axes are better. Here, we use clustering tasks with $n=512$ and $k=6$ . The initial learning rate for Adam is 0.01 (ablation in Figure 11 ) with no gradient clipping. In each plot, all 6 curves correspond to the same set of clustering tasks used for training, and we study the behaviour of the NE-regularized $\min$ (softmax) vs the L2-regularized $\min$ (sparsemax), jointly with the effect of $(E_{C},E_{X})$ — ( OH ) vs ( BN ) vs no token embedding ( NA ). In Figure 2(a) , we fix the data dimensionality $d=4$ and vary the $\min$ -regularization penalty $\tau$ (ablation in Figure 10 ), and in Figure 2(b) , we fix $\tau=0.25$ and vary $d$ (ablation in Figure 14 ).

Panel of Figure 4 — Figure 4 : Effect of data distribution. The above figures are as described in Figure 2 . We train on tasks with $n=512,d=32,k=10$ generated from the noted distributions. The validation loss is computed on clustering tasks generated from the training distribution (see ablation in Figure 16 ).

Panel of Figure 5 — Figure 5 : Training loss gap surfaces. We visualize the landscape of the loss gap $(\tilde{\mathcal{L}}_{\Omega}^{\tau}-\mathcal{L})$ between the smoothed $\tilde{\mathcal{L}}_{\Omega}^{\tau}$ and the $k$ -means loss $\mathcal{L}$ at $\tau=0.1$ .

Figure 6 : Multi-step clustering performance. Models trained with

n=512,d=32,k=10,\tau=0.1

and used for

T=20

step clustering. We consider all configurations (NE/L2 regularizers and OH / BN / NA embeddings). [Left] Multi-step clustering by applying Equation 4 iteratively for

T

steps. [Right] Multi-step clustering where the model input is re-embedded at each of the

T

steps, treating

T

-steps as

T

single-steps of clustering. See Appendix C for further details.

Figure 7 : Generalization to varying distributional properties. Models are trained with

n=512,d=32,k=10,\tau=0.1

, L2-regularized

\min

with OH / BN / NA token embeddings, and evaluated on

T=20

step clustering with tasks from mixtures of normal distributions. Training tasks are from equally weighted mixtures of isotropic normal distributions with a per-dimensional scale of 0.1. We evaluate generalization in 4 scenarios: [Left] unequally weighted mixture of isotropic distributions, [Left center] equally weighted mixture of anisotropic distributions, [Right center] equally weighted mixture of isotropic distributions with a per-dimension scale smaller (half) than that during training , and [Right] same as Right center but with a scale larger (double) than that during training .

Figure 8 : Generalization to different distribution families. The models are trained on one distribution family (each column maps to a training family), and tested for multi-step (

T=20

) clustering from a different distribution family (each row maps to a test family).

Figure 9 : Pairwise task dissimilarities. We use 4 tasks from each of the families and compute the pairwise Earth-mover distance (EMD) between their respective distribution of the distances between the points ${\bm{X}}$ and initial centers ${\bm{C}}$ , giving us a $20\times 20$ matrix. Each row and column corresponds to a single task, and the tasks are arranged into family-based blocks in the following order: normal , cauchy , gumbel , laplace , lognormal . Light shades show low EMD.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 10 — Figure 10 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the $\min$ -regularization penalty $\tau$ in the training loss $\tilde{\mathcal{L}}_{\Omega}^{\tau}$ (the smoothed upper-bound of the discrete non-differentiable $k$ -means objective $\mathcal{L}$ ). This is an extended version of Figure 2(a) . We consider $\tau\in\{1,0.25,0.14,0.1\}$ . These set of results correspond to a training with clustering tasks involving $n=512$ points and $k=6$ clusters with $d=4$ ( Figure 10(a) ) and $d=32$ ( Figure 10(b) ) dimensions. Here the initial learning rate $\eta=0.01$ for the Adam optimizer with $M=10000$ training steps with a task batch-size $B=32$ .

Panel of Figure 11 — Figure 11 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the initial learning rate $\eta$ for the Adam optimizer selected from $\{0.001,0.01,0.1,1.0\}$ with $M=5000$ training steps with a task batch-size $B=32$ . These set of results correspond to a training with clustering tasks involving $n=512$ points and $k=6$ clusters with $d=4$ ( Figure 11(a) ) and $d=32$ ( Figure 11(b) ) dimensions with a $\min$ -regularization penalty $\tau=0.1$ .

Figure 12 : “Length-generalization” of learned models. We consider models trained with

n=512,d=32,k=10,\tau=0.1

utilized for

T=20

step clustering on tasks with

\{128,256,1024,2048\}

points. We only consider L2-regularized

\min

and the three forms of token embeddings, OH , BN and NA as described in Figure 2 . The performance of Lloyd’s algorithm is plotted in black.

Panel of Figure 13 — Figure 13 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the number of points $n$ in the clustering tasks used for training the transformers, with $n\in\{128,256,512,1024\}$ . These set of results correspond to a training with clustering tasks involving $k=6$ clusters with $d=4$ ( Figure 13(a) ) and $d=32$ ( Figure 13(b) ) dimensions. Here we use the initial learning rate $\eta=0.01$ for the Adam optimizer with $M=10000$ training steps, a task batch-size $B=32$ , and the $\min$ -regularization penalty $\tau=0.1$ .

Panel of Figure 14 — Figure 14 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the dimensionality $d$ of the clustering tasks used for training the transformers, with $d\in\{4,8,16,32\}$ . This is an extended version of Figure 2(b) . These set of results correspond to a training with clustering tasks involving $n=512$ points, $k=6$ clusters, trained with the smoothed loss with a loss $\min$ -regularizer penalty $\tau=0.25$ ( Figure 14(a) ) and $\tau=0.1$ ( Figure 14(b) ). Here the initial learning rate $\eta=0.01$ for the Adam optimizer with $M=10000$ training steps, each with a task batch-size $B=32$ .

Panel of Figure 15 — Figure 15 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the number of clusters $k$ in the clustering tasks used for training the transformers, with $k\in\{6,10,16,25\}$ . These set of results correspond to a training with clustering tasks with $n=512$ points. To highlight the interplay of the dimensionality $d$ and the $\min$ -regularization penalty $\tau$ , we consider four cases: $d=4,\tau=1$ ( Figure 15(a) ), $d=4,\tau=0.1$ ( Figure 15(b) ), $d=32,\tau=1$ ( Figure 15(c) ), and $d=32,\tau=0.1$ ( Figure 15(d) ). In all cases, the initial learning rate $\eta=0.01$ for the Adam optimizer with $M=10000$ training steps each with a task batch-size $B=32$ .

Panel of Figure 16 — Figure 16 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying the distribution of the points to the clustered. These set of results correspond to a training with clustering tasks with $n=512$ points in $d=32$ dimensions, and is an extended version of Figure 4 . To highlight the interplay of the number of clusters $k$ and the loss temperature $\tau$ , we consider four cases: $k=6,\tau=1$ ( Figure 16(a) ), $k=10,\tau=1$ ( Figure 16(b) ), $k=6,\tau=0.1$ ( Figure 16(c) ), and $k=10,\tau=0.1$ ( Figure 16(d) ). In all cases, the initial learning rate $\eta=0.01$ for the Adam optimizer with $M=10000$ training steps and a per-step task batch-size $B=32$ .

Panel of Figure 17 — Figure 17 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) in terms of actual training time for varying values of $n\in\{256,512,1024\}$ . In one case, we use $k=16,d=16$ with $\tau=1$ ( Figure 17(a) ), while in another, we use $k=25,d=32$ , again with $\tau=1$ ( Figure 17(b) ). Here we limit ourselves to the L2-regularized $\min$ , and compare the OH and BN token embeddings, giving us only the BN:L2 and the OH:L2 curves.

Figure 18 : Generalization to different distribution families. Larger version of Figure 8 . The models are trained on one distribution family (each column maps to a training family), and tested for multi-step (

T=20

) clustering from a different distribution family (each row maps to a test family). We consider

5

distributions — cauchy , gumbel , laplace , lognormal and normal — giving us a

5\times 5

pairwise evaluations. Note that the pairwise evaluations on the “diagonal” correspond to training and evaluations tasks being sampled from the same distribution (though the tasks themselves are different, and the training only considers single-step clustering while we are evaluating multi-step clustering). We only consider models trained with the L2-regularized

\min

and the three forms of token embeddings, OH , BN and NA as described in Figure 2 . Each curve corresponds to the inter-quartile range aggregated over 10 trials. The performance of Lloyd’s is in black.

Term	Description
$m$	Number of query tokens
$m^{\prime}$	Number of key tokens
$\rho$	Upper-bound on the query token norms
$\rho^{\prime}$	Upper-bound on the key token norms
$\Gamma_{1}$	Upper-bound on the spectral norm of the query/key projection matrix
$\Gamma_{2}$	Upper-bound on the spectral norm of the value projection matrix

Appendix

Table 2 : Different quantities affecting the various bounds in Lemma I.4 , Lemma I.3 , Lemma I.5 , Lemma I.6 , Lemma I.7 and Lemma I.8 . First, we provide a description of the terms, and then we present the effect on the bound when the terms of interest are increased. In almost all cases, the bounds worsen with increase in the term. The only quantities that lead to a decrease in the terms are the embedding dimension

{d_{\mathsf{emb}}}

and the number of key tokens

m^{\prime}

. We

	${{\mathcal{A}_{2}^{\gamma}}}(\bar{{\bm{X}}}^{(0)},\bar{{\bm{C}}}^{(0)};(\cdot,1))$	${{\mathcal{A}_{2}^{\gamma}}}(\bar{{\bm{X}}}^{(0)},\bar{{\bm{X}}}^{(0)};(\cdot,2))$	${{\mathcal{A}_{2}^{\gamma}}}(\bar{{\bm{C}}}^{(0)},\bar{{\bm{X}}}^{(1)};(\cdot,3))$	${{\mathcal{A}_{2}^{\gamma}}}(\bar{{\bm{C}}}^{(0)},\bar{{\bm{C}}}^{(0)};(\cdot,4))$
$m$	$n$	$n$	$k$	$k$
$m^{\prime}$	$k$	$n$	$n$	$k$
$\rho$	$\sqrt{d}$	$\sqrt{d}$	$\sqrt{d_{E}}$	$\sqrt{d_{E}}$
$\rho^{\prime}$	$\sqrt{d_{E}}$	$\sqrt{d}$	$(1+\Gamma_{2})\sqrt{d}+\Gamma_{2}\sqrt{d_{E}}$	$\sqrt{d_{E}}$

Appendix

Table 3 : The value of the relevant terms from Table 2 in the different attention operations in Definition 2.2 . We write the attention operations compactly as

{{\mathcal{A}_{2}^{\gamma}}}({\bm{Z}},{\bm{Z}}^{\prime},(\cdot,\mu))

for query tokens

{\bm{Z}}\in\mathbb{R}^{{d_{\mathsf{emb}}}\times m}

, key tokens

{\bm{Z}}^{\prime}\in\mathbb{R}^{{d_{\mathsf{emb}}}\times m^{\prime}}

, where (i) the query token norms are bounded as

\left\|{\bm{Z}}_{:,i}\right\|\leq\rho

for all

i\in[m]

, (ii) the key token norms are

\left\|{\bm{Z}}^{\prime}_{:,j}\right\|\leq\rho^{\prime}

for all

j\in[m^{\prime}]

, and (iii) the

(\cdots,\mu)

corresponds to the appropriate

({\bm{Q}}_{\mu},{\bm{K}}_{\mu},{\bm{V}}_{\mu})

query/key/value projection matrices. The bounds on the norms of the key and query tokens are obtained from Lemma I.1 .

Explore similar work

May 7, 2026cs.LG

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent

Transformers have demonstrated remarkable in-context learning (ICL) capabilities. The strong ICL performance of transformers is commonly believed to arise from their ability to implicitly execute certain algorithms on the context, thereby enhancing prediction and generation. In this work, we investigate how transformers with softmax attention perform in-context learning on linear classification data. We first construct a class of multi-layer transformers that can perform in-context logistic regression, with each layer exactly performing one step of normalized gradient descent on an in-context loss. Then, we show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model. Training convergence guarantees of the self-attention layer and out-of-distribution generalization guarantees of the looped model are provided. Our results advance the theoretical understanding of ICL mechanism by showcasing how softmax transformers can effectively act as in-context learners.

Chenyang Zhang, Yuan Cao

School of Computing & Data Science, The University of Hong Kong

In-Context Learning Transformer Architectures In-Context Linear Attention Logistic Regression Gradient Descent

Ask PDF

May 8, 2026cs.LG

Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers

Transformers perform inference by iteratively transforming token representations across layers. This layerwise computation has been studied empirically, and recent mean-field theories of Transformer dynamics explain how attention can drive token distributions toward clustering. However, existing mean-field analyses largely treat model parameters as prescribed, leaving open how training reshapes this clustering picture. We study this question in a noisy mean-field Transformer in which only a parameter-linear FFN is trained under

L^2

regularization. We find and analyze a training-induced phase in the dynamics: after initially following attention-driven clustering, the token distribution can leave the clustered regime near the final layers. Our mathematical analysis is based on an entropy-regularized interaction energy that captures the clustering bias of attention. More broadly, our results point toward a training-aware mean-field theory of Transformer dynamics, in which training and inference dynamics are treated together.

Noboru Isobe, Daisuke Inoue, Masaaki Imaizumi

1RIKEN Center for Advanced Intelligence Project · 2Imperial College London · 3The University of Tokyo +1

Transformer Attention Transformer Architectures Mean-Field Limit Inference Dynamics Clustering Escape

Ask PDF

May 29, 2026cs.LG

Fixed Universal Transformers

We introduce \emph{universal transformers}: fixed transformers that can simulate any transformer in a given class via a suitable input embedding. Analogous to a universal Turing machine, the input embedding encodes a description of the target model while all internal parameters remain fixed. We provide explicit sparse constructions achieving universality when the embedding dimension is sufficiently large, and further show that universality is generic: randomly initialized transformers are universal almost surely, which aligns with recent empirical results of Zhong and Andreas (2024). We empirically validate our theory on the algorithmic tasks of parenthesis balancing and multi-hop reasoning. Our results suggest that much of a transformer's expressive power may reside in its input representation rather than its learned weights.

Jingwen Liu, Alexandr Andoni, Daniel Hsu

Columbia University New York, NY

Transformer Architectures Energy-Based Transformer Multi-Hop Reasoning

Ask PDF