cs.LGOct 7, 2026
SaveLeaner Transformers Can Easily Learn to Cluster
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 -means clustering with points in dimensions with an embedding size (thus, requiring attention projection matrices of size ). 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 . 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 and . 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 ), and (iv) Forward pass (eval) : forward pass time with the computation graph disabled (via model.eval() ). Each cell in the grid in the heatmap lists the relative gain (or degradation), with actual metrics for ( OH ) also listed in the cell as (mean std) for completeness; darker greens denote larger gain values. All runs use 32 clustering problems with 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 , especially when the baseline ( OH ) runtimes is already quite small ( ms). However, once is large enough, the gains can be quite significant, reaching over 50% in some cases.
| 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 ] | 9.53 / 10.66 [0.89 ] | 419.93 / 416.29 [0.9%] |
| 25 | 25 / 5 | 5.24 / 4.99 [1.05 ] | 11.21 / 11.01 [1.02 ] | 441.08 / 423.23 [4.0%] |
| 100 | 100 / 7 | 9.76 / 4.52 [2.16 ] | 20.21 / 9.93 [2.03 ] | 532.96 / 447.04 [16.1%] |
| 400 | 400 / 9 | 38.68 / 7.89 [4.90 ] | 66.13 / 16.47 [4.02 ] | 1095.60 / 531.02 [51.5%] |
| 1000 | 1000 / 10 | 193.07 / 15.90 [12.14 ] | 308.28 / 32.33 [9.54 ] | 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 is varied in the range (column 1) for a fixed data dimensionality . 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 for each value of for the two embeddings (column 2), and the significant reduction in the additional embedding dimensions lead to the significant reductions in the runtimes once is sufficiently larger than with up to 12 /9 speedup in forward/backward pass, and almost 70% reduction in memory usage.
Figure 2 : Effect of data dimensionality , -regularizer with penalty , token embedding on training convergence and generalization. We show the convergence of (i) the smoothed training loss (top row in each figure) and (ii) the relative non-differentiable clustering loss on held-out validation clustering tasks (bottom row) for multiple configurations; the dotted black line at (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 and . 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 (softmax) vs the L2-regularized (sparsemax), jointly with the effect of — ( OH ) vs ( BN ) vs no token embedding ( NA ). In Figure 2(a) , we fix the data dimensionality and vary the -regularization penalty (ablation in Figure 10 ), and in Figure 2(b) , we fix and vary (ablation in Figure 14 ).
Figure 3 : Loss surfaces for different dimensionalities ( vs ), token embeddings ( OH ) vs ( BN ) with the L2-regularized . We fix and vary to visualize its effect. We use the procedure in Li et al. (2018) with filterwise normalization and 2 random vectors in the parameter space.
Figure 4 : Effect of data distribution. The above figures are as described in Figure 2 . We train on tasks with generated from the noted distributions. The validation loss is computed on clustering tasks generated from the training distribution (see ablation in Figure 16 ).
Figure 5 : Training loss gap surfaces. We visualize the landscape of the loss gap between the smoothed and the -means loss at .
Figure 6 : Multi-step clustering performance. Models trained with and used for step clustering. We consider all configurations (NE/L2 regularizers and OH / BN / NA embeddings). [Left] Multi-step clustering by applying Equation 4 iteratively for steps. [Right] Multi-step clustering where the model input is re-embedded at each of the steps, treating -steps as single-steps of clustering. See Appendix C for further details.
Figure 7 : Generalization to varying distributional properties. Models are trained with , L2-regularized with OH / BN / NA token embeddings, and evaluated on 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 ( ) 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 and initial centers , giving us a 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
Figure 10 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the -regularization penalty in the training loss (the smoothed upper-bound of the discrete non-differentiable -means objective ). This is an extended version of Figure 2(a) . We consider . These set of results correspond to a training with clustering tasks involving points and clusters with ( Figure 10(a) ) and ( Figure 10(b) ) dimensions. Here the initial learning rate for the Adam optimizer with training steps with a task batch-size .
Figure 11 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the initial learning rate for the Adam optimizer selected from with training steps with a task batch-size . These set of results correspond to a training with clustering tasks involving points and clusters with ( Figure 11(a) ) and ( Figure 11(b) ) dimensions with a -regularization penalty .
Figure 12 : “Length-generalization” of learned models. We consider models trained with utilized for step clustering on tasks with points. We only consider L2-regularized 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.
Figure 13 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the number of points in the clustering tasks used for training the transformers, with . These set of results correspond to a training with clustering tasks involving clusters with ( Figure 13(a) ) and ( Figure 13(b) ) dimensions. Here we use the initial learning rate for the Adam optimizer with training steps, a task batch-size , and the -regularization penalty .
Figure 14 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the dimensionality of the clustering tasks used for training the transformers, with . This is an extended version of Figure 2(b) . These set of results correspond to a training with clustering tasks involving points, clusters, trained with the smoothed loss with a loss -regularizer penalty ( Figure 14(a) ) and ( Figure 14(b) ). Here the initial learning rate for the Adam optimizer with training steps, each with a task batch-size .
Figure 15 : Training loss convergence (top row) and relative validation loss from Equation 20 (bottom row) for varying values of the number of clusters in the clustering tasks used for training the transformers, with . These set of results correspond to a training with clustering tasks with points. To highlight the interplay of the dimensionality and the -regularization penalty , we consider four cases: ( Figure 15(a) ), ( Figure 15(b) ), ( Figure 15(c) ), and ( Figure 15(d) ). In all cases, the initial learning rate for the Adam optimizer with training steps each with a task batch-size .
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 points in dimensions, and is an extended version of Figure 4 . To highlight the interplay of the number of clusters and the loss temperature , we consider four cases: ( Figure 16(a) ), ( Figure 16(b) ), ( Figure 16(c) ), and ( Figure 16(d) ). In all cases, the initial learning rate for the Adam optimizer with training steps and a per-step task batch-size .
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 . In one case, we use with ( Figure 17(a) ), while in another, we use , again with ( Figure 17(b) ). Here we limit ourselves to the L2-regularized , 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 ( ) clustering from a different distribution family (each row maps to a test family). We consider distributions — cauchy , gumbel , laplace , lognormal and normal — giving us a 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 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 |
|---|---|
| Number of query tokens | |
| Number of key tokens | |
| Upper-bound on the query token norms | |
| Upper-bound on the key token norms | |
| Upper-bound on the spectral norm of the query/key projection matrix | |
| Upper-bound on the spectral norm of the value projection matrix |
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 and the number of key tokens . We
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 for query tokens , key tokens , where (i) the query token norms are bounded as for all , (ii) the key token norms are for all , and (iii) the corresponds to the appropriate query/key/value projection matrices. The bounds on the norms of the key and query tokens are obtained from Lemma I.1 .