Title: Attention-based clustering
Venue: NeurIPS 2025
Accepted: yes
Keywords: Theory of neural networks

Clustering

Transformers

Attention-based models

Mixture models

Optimization

Unsupervised learning
Abstract: Transformers have emerged as a powerful neural network architecture capable of tackling a wide range of learning tasks. In this work, we provide a theoretical analysis of their ability to automatically extract structure from data in an unsupervised setting. In particular, we demonstrate their suitability for clustering when the input data is generated from a Gaussian mixture model. To this end, we study a simplified two-head attention layer and define a population risk whose minimization with unlabeled data drives the head parameters to align with the true mixture centroids.

Attention-based clustering

Abstract

Transformers have emerged as a powerful neural network architecture capable of tackling a wide range of learning tasks. In this work, we provide a theoretical analysis of their ability to automatically extract structure from data in an unsupervised setting. In particular, we demonstrate their suitability for clustering when the input data is generated from a Gaussian mixture model. To this end, we study a simplified two-head attention layer and define a population risk whose minimization with unlabeled data drives the head parameters to align with the true mixture centroids. This phenomenon highlights the ability of attention-based layers to capture underlying distributional structure. We further examine an attention layer with key, query, and value matrices fixed to the identity, and show that, even without any trainable parameters, it can perform in-context quantization, revealing the surprising capacity of transformer-based methods to adapt dynamically to input-specific distributions.

1 Introduction

Attention-based models (Bahdanau et al., 2015), in particular Transformers (Vaswani et al., 2017), have achieved state-of-the-art performance across a wide range of learning tasks. These include applications in natural language processing (Devlin et al., 2018; Bubeck et al., 2023; Luong et al., 2015; Bahdanau et al., 2016) and computer vision (Dosovitskiy et al., 2020; Liu et al., 2021; Ramachandran et al., 2019). The success of the attention mechanism has been linked to its ability to capture long-range relationships in input sequences (Bahdanau et al., 2015; Vaswani et al., 2017). They do this by computing pairwise dependencies between tokens based on learned projections, without regard to the tokens' positions in the sequence.

Onthe theoretical side, a full understanding of attention-based mechanisms has not yet been developed. This is due to the complexity of the architectures and the diversity of relevant tasks they manage to achieve. A promising research direction to bridge this gap involves identifying essential features from real-world problems and constructing minimal yet representative tasks that retain the essential difficulty-paired with provable models that solve them using attention-based mechanisms. Notable recent efforts in this vein include Ahn et al. (2023); von Oswald et al. (2023); Yang et al. (2025); Zhang et al. (2024); Li et al. (2024, 2023). However, the existing literature mainly focuses on supervised learning aspects, and in particular in-context learning (von Oswald et al., 2023; Zhang et al., 2024; Garg et al., 2023; Li et al., 2023; Furuya et al., 2024). The goal of in-context learning is to predict the output corresponding to a new query, given a prompt consisting of input/output pairs.

Beyond the standard supervised setting, Transformer models are often (pre)trained in practice with semi-supervised objectives such as masked language modeling (Phuong and Hutter, 2022). This raises important questions about their statistical behavior and training dynamics in unsupervised regimes. In this work, we examine Transformers through the lens of clustering, thereby revealing the inherent capacity of attention mechanisms to perform unsupervised representation learning. To the best of our knowledge, the only prior theoretical work that explicitly explores clustering with Transformers is He et al. (2025), who demonstrate that attention layers can mimic the EM algorithm (Lloyd, 1982), albeit assuming known cluster labels during training. In contrast, our analysis focuses on the fully unsupervised setting and further provides insight into the functional roles of individual attention heads in the context of model-based clustering.

Contributions. In this paper, we investigate the behavior of attention layers in an unsupervised learning setting, where input data is drawn from mixture distributions. We focus on a two-component Gaussian mixture model. Within this classical clustering framework, we introduce a two-headed linear attention layer designed to capture cluster membership through attention scores, while remaining analytically and computationally tractable. To assess the quality of the embeddings produced by the attention mechanism, we define a theoretical risk analogous to the classical quantization error in unsupervised learning. We analyze the training dynamics of the proposed predictor under projected gradient descent and prove that, with appropriate initialization, the algorithm can learn the true latent centroids of the mixture components, despite the non-convexity of the loss landscape and without access to cluster labels. To relax the initialization requirements in practice, we further propose a regularization scheme that promotes disentanglement between attention heads. Our theoretical findings are supported by numerical experiments under varying conditions, including different initialization regimes, mixture separability levels, and problem dimensionalities. Overall, we show that attention-based predictors can successfully adapt to mixture models by learning the underlying centroids through training. We also study their quantization properties, that is, how the layer summarizes the input distribution. In the oracle regime, where the attention parameters have converged to the true centroids, the problem can be viewed as an approximate denoising task, in the sense commonly used in the statistical physics literature, where each input can be interpreted as a noisy observation of an underlying centroid, and the transformed token corresponds, to some extent, to its denoised estimate. Finally, we focus on a particular attention layer in which the key, query, and value parameters are fixed to the identity matrix. Surprisingly, we show that this type of layer, despite having no trainable parameters, can still perform in-context quantization, meaning it adapts to the case where the distribution of each input sequence comes with its own mixing parameters. This further demonstrates the remarkable ability of transformer-based methods to adapt on the fly to the underlying data distribution, even when no attention parameters are trained.

Organization. Section 2 introduces the problem and outlines the proposed approach. In Section 3, we address the general problem with linear attention applied to a two-component Gaussian mixture model. In Section 4, we discuss the quantization properties of attention-based predictor with oracle parameters. In Section 5, we explore an in-context clustering framework and examine the quantization capabilities of a simple attention-based layer with no learned parameters. The proofs of all the theoretical results can be found in the appendices.

2 A starter on attention-based layers and clustering

Data distribution in model-based clustering. In attention-based learning, the key idea is to map a set of input tokens to a transformed set of output tokens. With this in mind, we consider an input sequence X ∈ R L × d composed of L tokens ( X 1 , . . . , X L ) , each token being a vector of R d , i.e., X = ( X 1 , . . . , X L ) ⊤ . We assume that the tokens are i.i.d. drawn from a simple mixture model: for 1 ≤ ℓ ≤ L ,

with balanced components and where the centroids ( µ ⋆ 0 , µ ⋆ 1 ) ∈ ( S d -1 ) 2 (i.e., ∥ µ ⋆ 0 ∥ 2 = ∥ µ ⋆ 1 ∥ 2 = 1 ) are assumed to be orthogonal, i.e., such that ⟨ µ ⋆ 0 , µ ⋆ 1 ⟩ = 0 . Therefore, for each token, there exists an associated latent variable, denoted by Z ℓ , corresponding to a Bernoulli random variable of parameter 1 / 2 and encoding its corresponding cluster.

Attention-based predictors. An attention head made of a self-attention layer can be written as

where the softmax of temperature λ > 0 is applied row-wise, no skip connection is considered and the matrices K,Q,V ∈ R d × p are usually referred to as keys, queries and values. We adopt the convention that the values are identity matrices; thus, the attention head simply outputs combinations of the initial tokens weighted by the attention scores. While this simplification is certainly convenient for facilitating the mathematical analysis that follows, it is also supported by experimental studies showing comparable performance when the value matrices are removed (He and Hofmann, 2024). Furthermore, assume that the key and query matrices are equal to the same column matrix µ ∈ R d × 1 , we obtain

With such an architecture, the ℓ -th output vector is therefore given by

which corresponds to aggregating the X k 's when X k and X ℓ are simultaneously aligned with µ . This suggests that attention heads may act as effective learners in a clustering framework.

The softmax nonlinearity used in the attention head (1) introduces a coupling between tokens, which undoubtedly complicates the mathematical analysis. To address this difficulty, we propose to look at a simplified linear attention head, still parameterized by µ ∈ R d , and defined for 1 ≤ ℓ ≤ L , as

This head uses a linear activation function instead of the traditional softmax found in practical architectures, and has already received interest in several mathematical studies (see Zhang et al., 2024; von Oswald et al., 2023; Han et al., 2023; Katharopoulos et al., 2020). This model is a particular case of Gaussian sequence multi-index models (Cui, 2025), which were already used to study attention models (Cui et al., 2024; Arnaboldi et al., 2025; Troiani et al., 2025), albeit not in a clustering context.

Note that when µ is chosen to be µ ⋆ 0 , then for tokens X ℓ and X k whose corresponding latent variables Z ℓ and Z k are both equal to 0 (i.e., the samples belong to the same cluster centered at µ ⋆ 0 ), the vectors X ℓ and X k are likely to be aligned with µ ⋆ 0 . In this case, we have ( X ⊤ ℓ µµ ⊤ X k ) X k ≃ X k . Conversely, if X ℓ and X k are associated with different latent variables (e.g., Z ℓ = 0 and Z k = 1 ), then ( X ⊤ ℓ µµ ⊤ X k ) X k ≃ 0 . This behavior suggests that when setting µ = µ ⋆ 0 , and if X ℓ belongs to the cluster centered at µ ⋆ 0 , the sum ∑ L k =1 ( λX ⊤ ℓ µµ ⊤ X k ) X k effectively aggregates the X k 's from the same cluster, whose expected number is L/ 2 , motivating the renormalizing factor of 2 /L . Overall, H lin ,µ ⋆ 0 ( X ) ℓ can be seen as producing an empirical mean of the tokens belonging to the same cluster, serving as an estimator of the corresponding centroid.

Therefore, assuming that the number of clusters in the data is known, it is natural to consider an attention-based predictor composed of two attention heads, parameterized by µ 0 and µ 1 ∈ R d ,

Metric loss. As no label is available, we focus on minimizing the following theoretical loss:

where T is an arbitrary attention-based predictor. The distinctive feature of this risk lies in the fact that, if the predictor were able to return, for each token X ℓ , its associated centroid µ ⋆ Z ℓ , the risk would exactly correspond to a quantization error, characteristic of a standard clustering task. Note that, due to the independence of the tokens, we have L ( T ) = E ∥ X 1 -T ( X ) 1 ∥ 2 2 , so we can confine the following theoretical analysis on the minimization of the predictive error for the first token only.

PGD iterates. In this paper, we focus on the training dynamics of Transformer-based predictors when minimizing the theoretical risk L . While we acknowledge that, in practice, an empirical version of this risk is typically used, analyzing the optimization of the theoretical risk is already a non-trivial task, which offers valuable insights into the behavior observed in practice.

For a given predictor T lin ,µ 0 ,µ 1 made of two linear attention heads parameterized respectively by µ 0 and µ 1 , one can reinterpret the objective L as a function R of the parameters ( µ 0 , µ 1 ) , defined by

Note that the computation of the risk also depends on the choice of the underlying data distribution. However, as is common practice in the literature, we do not explicitly indicate the dependence of the risk on the data distribution. This should not hinder understanding, as the distributional assumptions and context will always be made clear. As we rely on the theoretical analysis on an expression of this risk restricted to the sphere, we consider as a gradient strategy, the Projected (Riemaniann) Gradient Descent (PGD) algorithm (Boumal, 2023). Given an initialization ( µ 0 0 , µ 0 1 ) ∈ ( S d -1 ) 2 and a step size γ > 0 , the PGD iterates ( µ k 0 , µ k 1 ) ∈ ( S d -1 ) 2 are recursively defined by:

In what follows, we analyze the convergence of these iterates to the oracle centroids, both theoretically and numerically.

3 Training dynamics: The centroids are learned as attention parameters

In this section, we analyze the training dynamics of an attention layer optimized by minimizing the risk R through the PGD iterations defined in ( PGD ). We focus on the case of clustering data generated from a Gaussian mixture model, as specified in ( P σ ), where the input tokens are i.i.d. samples

for 1 ≤ ℓ ≤ L , where ( µ ⋆ 0 , µ ⋆ 1 ) ∈ ( S d -1 ) 2 are orthogonal unit vectors. As a first step toward a mathematical understanding of Transformer-based layers in clustering settings, we have also considered the degenerate case σ 2 = 0 , in which the data are drawn from a mixture of Dirac masses. We perform a detailed analysis of the training dynamics under this simplified setting, which provides valuable intuition and serves as a foundation for the subsequent study of the non-degenerate mixture model. For completeness, this preliminary analysis is presented in Appendix A.

3.1 Theoretical analysis

Preliminary computations. We start by introducing the following quantities:

and derive a closed-form expression for the risk w.r.t. this reparameterization. Although the full expression is somewhat complex (see Appendix C.1), the following proposition highlights its key structural properties, as being a polynomial in these five variables.

Proposition 3.1. Under the Gaussian mixture model ( P σ ) , consider the attention-based predictor T lin ,µ 0 ,µ 1 composed of two linear heads parameterized by ( µ 0 , µ 1 ) ∈ ( S d -1 ) 2 . Then, there exists a function R < : [ -1 , 1] 5 ↦→ R such that R ( µ 0 , µ 1 ) = R < ( κ 0 , κ 1 , η 0 , η 1 , ξ ) and

We remark that when η 0 , η 1 , and ξ are fixed to 0 , most of the monomials vanish, yielding a fully explicit formula for the risk (see Lemma C.1 in the appendices). Far from being a mere rewrite, this step provides the algebraic foundation for all the exact calculations and insights that follow.

Optimality conditions. Given the complexity of the theoretical analysis, we focus on a simplified setting by restricting our study to parameters ( µ 0 , µ 1 ) lying on a specific manifold 2 :

In terms of notation, it is equivalent to assume that η 0 = 0 , η 1 = 0 , ξ = 0 . Therefore on this manifold, we adopt the shorthand notation R < ( κ 0 , κ 1 ) def = R < ( κ 0 , κ 1 , 0 , 0 , 0) .

Lemma 3.2. Under the Gaussian mixture model ( P σ ) , the risk R < restricted to M has the form

for A,B,C,D non-negatives constants, dependent on σ and L , made explicit in Lemma C.1.

Proposition 3.3. Consider R < ( κ 0 , κ 1 ) , there exists λ ⋆ ( σ, L ) > 0 such that the points ( ± 1 , ± 1) are global minima of R < ( κ 0 , κ 1 ) .

This result demonstrates that, under a suitable condition on the temperature parameter-specifically, when λ = λ ⋆ ( σ, L ) -the points ± µ ⋆ 0 and ± µ ⋆ 1 are global minimizers of the risk. The explicit form of λ ⋆ ( σ, L ) is provided in Proposition C.2 in the appendices. Moreover, it is worth noting that as the variance σ 2 of the Gaussian components tends to zero, λ ⋆ ( σ, L ) approaches the value λ ⋆ 0 = L +1 L +3 ,which coincides with the value previously identified in the degenerate case ( P 0 ). On the other hand, when σ is fixed and L grows large, λ ⋆ ( σ, L ) tends toward λ ⋆ ∞ = 1+4 σ 2 1+5 σ 2 +6 σ 4 , which will guide us to properly choose λ in our numerical experiments.

Convergence analysis. In what follows, we show that the ( PGD ) iterates can indeed converge to global minimizers, provided they are suitably initialized on the manifold M .

Theorem 3.4. Under the Gaussian mixture model ( P σ ) , consider the attention-based predictor T lin ,µ 0 ,µ 1 composed of two linear heads. Take λ ∈ ]0 , λ ⋆ ( σ, L )] , with λ ⋆ ( σ, L ) defined as in Proposition 3.3. Then there exists ¯ γ > 0 such that for any stepsize 0 < γ < ¯ γ , and for a generic initialization ( µ 0 0 , µ 0 1 ) ∈ M , the iterates ( µ k 0 , µ k 1 ) generated by ( PGD ) converge to the centroids (up to a sign), i.e.

Theorem 3.4 underlines the capabilities of linear attention-based predictors in a clustering framework. With appropriate initialization, the attention heads align with the true underlying centroids even when trained without access to labels. This result shows that attention layers can uncover and encode the latent structure of the input distribution in a fully unsupervised setting through their parameters.

3.2 Experimental verification of the theoretical results

Setting. To better reflect practical algorithmic behavior, we implement Projected Stochastic Gradient Descent (PSGD; see Appendix F.1), which serves as an empirical counterpart to the ( PGD ) iterates by relying on sample-based estimations.

In what follows, we use the metric referred to as distance to the centroids (up to a sign) , given by

where S 2 is the permutation group of two elements, µ ⋆ 0 , µ ⋆ 1 denote the true centroids, respectively, while ˆ µ 0 , ˆ µ 1 are the parameters returned by ( PSGD ). Note that this distance is invariant under relabeling and sign flips of the head parameters. More implementation details related to the following experiments can be found in Appendix F.2.

It is worth noting that the assumption of orthogonality of the centroids on the unit sphere always results in a constant distance between the centroids, namely ∥ µ ⋆ 0 -µ ⋆ 1 ∥ 2 = √ 2 . In this setting, to characterize

the separation between the two modes of the mixture, one can introduce a notion of interference that depends solely on the variance level of each mode and which is defined as I( σ ) = P ( X σ > √ 2 2 ) , where X σ ∼ N (0 , σ 2 ) . Remark that this function is increasing with supremum 0 . 5 . This motivates the choice of two contrasting scenarios for the numerical experiments: a low-interference regime with σ = 0 . 3 , where I(0 . 3) ≈ 0 . 01 , and a high-interference regime with σ = 1 , where I(1) ≈ 0 . 24 . More implementation details of the following numerical experiments can be found in Appendix F.2.

Results. When initialization is done on the manifold, the training analysis depicted in Figure 1 demonstrates linear convergence of the head parameters toward the centroids (up to a sign) during the first 10 3 iterations, which is in line with the obtained theoretical results. The error then plateaus at around 10 -2 in the lowinterference setting ( σ = 0 . 3 ), and around 10 -1 in the high-interference setting ( σ = 1 ). This saturation phenomenon is attributed to the stochasticity introduced by ( PSGD ) in place of ( PGD ) in the simulations.

3.3 Generalizations

We consider several generalizations of our theoretical setting, to give insight on the role of our assumptions.

0

2000

4000

6000

8000

10000

Iterations

10

2

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

=0.3

=1

Random initialization on the unit sphere (outside of the manifold). When the initialization is performed outside the manifold, PGD iterates only partially align with the underlying centroids. A way to handle arbitrary initializations (suggested by our analysis in the degenerate case, see Appendix A), is to introduce a regularized risk minimization problem:

for ρ > 0 and the regularization term defined by r ( µ 0 , µ 1 ) = E [ ⟨ µ 0 , X 1 ⟩ 2 ⟨ µ 1 , X 1 ⟩ 2 ] . The role of this term is to encourage the orthogonality conditions on µ 0 , µ 1 , thereby compensating for initializations that may fall outside the manifold M . Numerical results show that the centroids can be recovered with an appropriate level of regularization (see Figure 2a). Note that, as the strength of the regularization increases, it gradually overrides the original objective and impairs the alignment of the head parameters with the true centroids -an effect that becomes more pronounced at higher noise levels. In Figure 2b, we fix the regularization strength and observe linear convergence towards the centroids over the course of 4 × 10 3 iterations. The error eventually plateaus near 10 -3 for σ = 0 . 3 and near 10 -1 for σ = 1 . This shows that the regularization strategy inspired by the analysis of the simplified Dirac mixture case remains effective in the more realistic setting of Gaussian mixtures. In this context as well, it enhances the interpretability of the attention parameters by encouraging their alignment with the unknown components of the underlying mixture.

0

1

2

3

Regularization

10

2

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

=0.3

=1

0

2000

4000

6000

8000

10000

Iterations

10

2

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

=0.3

=1

In Appendix G.1, we present additional experiments in higher-dimensional settings, highlighting the impact of dimensionality on the training dynamics.

Orthogonality of the clusters. The framework studied in this paper assumes that the centroids of the clusters are orthogonal. Relaxing this assumption to allow for potential overlap between centroid directions significantly complicates the theoretical analysis by introducing additional terms. That said, the analysis in the orthogonal case remains highly informative for the non-orthogonal setting, particularly in high-dimensional regimes. Indeed, when the true centroids are randomly sampled on the sphere, they become nearly orthogonal as the dimension d increases. In such cases, learning the attention heads via the regularized risk minimization P ρ remains meaningful and yields compelling empirical results (see Appendix G.2).

Euclidean gradient instead of Riemannian gradient. The use of projected Riemannian gradient descent is both natural and mathematically justified in our setting, as we are able to compute the risk function only for parameters constrained to lie on the unit sphere. This constraint, however, is not a practical limitation: our numerical experiments indicate that standard projected SGD heuristics exhibit similar behavior to the theoretically grounded Riemannian updates, see Appendix G.3. We emphasize that the projection step is essential for the algorithm's proper functioning.

More clusters. A natural extension is to consider the case of K centroids and attention heads. We perform an experiment in the case K = 3 , which shows recovery of the centroids, see details in Appendix G.4. The main difference with the case K = 2 is the regularization term which now writes

to promote pairwise orthogonality between the parameters. This approach should further generalize to the case of K orthonormal centroids with K < d .

4 Attention-based layers as approximate quantizers

We have seen that attention-based predictors can adapt to mixture models by learning the underlying centroids through training. In this section, we investigate the quantization properties of an attentionbased predictor whose parameters have converged to the true centroids. To guide our analysis, we introduce the optimal quantizer T ⋆ as a statistical benchmark within a standard clustering framework. This oracle predictor returns the true centroid of each token, that is, for 1 ≤ ℓ ≤ L , T ⋆ ( X ) ℓ = µ ⋆ Z ℓ where Z ℓ is encoding the latent cluster of the token X ℓ . One can immediately note that the risk of the optimal quantizer is given by

Returning to the attention-based predictor T lin ,µ ⋆ 0 ,µ ⋆ 1 with oracle parameters, the first key observation is that it closely resembles an optimal quantizer: its ℓ -th output aligns, on average, with the centroid of the cluster to which the ℓ -th token belongs, as shown by the next lemma.

Proposition 4.1. Under the Gaussian mixture model ( P σ ) , it holds that

Therefore, choosing λ = L ( L +1)+2( L +3) σ 2 leads to unbiased approximate quantization, i.e.,

One can next characterize the asymptotic risk and variance of the oracle attention-based predictor.

Proposition 4.2. Under the Gaussian mixture model ( P σ ) , fix the temperature λ = 1+4 σ 2 +4 σ 4 1+6 σ 2 +12 σ 4 +8 σ 6 . Then, the risk of the attention-based predictor T lin ,µ ⋆ 0 ,µ ⋆ 1 with oracle parameters µ ⋆ 0 and µ ⋆ 1 satisfies

Moreover, for an arbitrary value of λ , when L →∞ , we get

Strikingly, as the input sequence length L increases, we find that

This result shows that the attention-based predictor asymptotically achieves a lower risk than the optimal quantizer. This phenomenon can be partly explained by the fact that the comparison is not entirely fair: the predictors T lin ,µ ⋆ 0 ,µ ⋆ 1 and T ⋆ do not belong to the same class of functions. Indeed, the optimal quantizer T ⋆ is only allowed to return two fixed vectors and relies on a single input token to predict the associated centroid (albeit with access to the latent label).

On the other hand, the image of attention-based encoder T lin ,µ ⋆ 0 ,µ ⋆ 1 is not discrete, and moreover this estimator aggregates a growing sequence of random variables, all drawn from the same mixture. The aggregation of multiple inputs can be seen as a variance reduction mechanism, which reduces the risk (this is also evident in the proof of Proposition 4.2, where the risk is shown to decrease as L increases). Note that the gap vanishes in a high-dimensional setting d →∞ .

Another insightful comparison between the predictors T lin ,µ ⋆ 0 ,µ ⋆ 1 and T ⋆ is through their conditional variances Var[ T lin ,µ ⋆ 0 ,µ ⋆ 1 ( X ) 1 | Z 1 = c ] . While the conditional variance of the optimal quantizer is null by definition, it is positive for the linear attention layer, and asymptotically independent of d as shown in Proposition 4.2. This once again highlights the fact that these two quantifiers belong to function classes of different complexity. These properties are illustrated in Figure 3: we observe that the attention-

1.5

1.0

0.5

0.0

0.5

1.0

1.5

PCA Component 1

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

PCA Component 2

Input tokens

Transformed tokens

0

,

1

based embeddings are approximate projections of the inputs on the line between the two centroids. In particular, the variance of the embedded point cloud is lower than the variance of the inputs, illustrating that (for λ independent of d and d large enough),

5 In-context clustering

5.1 Setting

So far, we have considered the traditional setting for model-based clustering, with a mixture of Gaussian made of two components: each token was assumed to be distributed as follows

with fixed centroids µ ⋆ 0 and µ ⋆ 1 of unit-norm and orthogonal. We have shown that despite the nonconvexity of the problem, attention-based predictors including two heads could perform approximate quantization and discover the underlying centroids encoded in their parameters. Note that for an input sequence of tokens X = ( X 1 | . . . | X L ) ⊤ ∈ R L × d , the first output of the type of attention-based predictor considered in this paper, parameterized by µ 0 and µ 1 , can be rewritten as

Therefore, when we consider two linear heads parameterized by row vectors for the queries and keys, and constrain them to be orthogonal, the setup can be interpreted as a single attention head with a query/key matrix of rank 2. Interestingly, we are able to effectively train this rank-2 query/key head by leveraging the non-convex optimization of two simple, row-structured heads.

Let us now challenge the clustering setting, by assuming that each input sequence is still drawn from a 2-component mixture, but with its own centroids, i.e., for each input sequence X i = ( X i 1 , . . . , X iL ) ⊤ ∈ R L × d , i = 1 , . . . , n , we assume the tokens ( X iℓ ) ℓ to be i.i.d., such that the ℓ -th token is distributed as

for some random orthogonal centroids µ ⋆ i 0 and µ ⋆ i 1 of unit-norm.

If the prior distribution over the centroids is concentrated along preferred directions, say µ ⋆⋆ 0 and µ ⋆⋆ 1 , then the predictor T lin ,µ 0 ,µ 1 will likely perform well, as the Transformer's parameters will tend to align with these directions after training. Conversely, if the centroids are distributed in an isotropic way, T lin ,µ 0 ,µ 1 will struggle with in-context clustering due to its limited flexibility, as only two parameters µ 0 and µ 1 govern the embedding, whereas the centroids vary significantly from one input sequence to another. To address this issue, a natural idea is to give more degrees of freedom to the predictor by increasing the number of linear attention heads. Specifically, if we consider d attention heads whose parameters are constrained to be row vectors ( µ c ) c =1 ,...,d ∈ ( R d ) d , each of unit norm and mutually orthogonal, we obtain the following attention layer: for an input sequence X = ( X 1 , . . . , X L ) ⊤ ∈ R L × d , and for 1 ≤ ℓ ≤ L ,

so finally,

Somewhat surprisingly, using d simplified linear heads in parallel, while enforcing orthogonality among their parameters, ultimately amounts to employing an attention layer with no trainable parameters. In what follows, we discuss the properties of T ctx in an in-context clustering framework.

More formally, we refer to as the in-context clustering as a setting in which the input consists of a generic sequence of L tokens X 1 , . . . , X L , sampled from a Gaussian mixture model whose component means (centroids) are randomly drawn on the unit sphere:

Associated with each token X ℓ , we still consider a latent variable Z ℓ , corresponding to a Bernoulli random variable of parameter 1 / 2 and encoding its corresponding cluster, so that

5.2 Linear attention layers can perform in-context approximate quantization

We first observe, that similarly to the fixed centroids setting of Section 4, the attention-based encoder T ctx performs an approximate in-context quantization of the input distribution by aggregating sequences sampled from a Gaussian mixture. Specifically, as formalized below, this simple architecture effectively aligns its output with the correct centroid.

Proposition 5.1. For c ∈ { 0 , 1 } , one has

In what follows, we characterize the risk and variance of the attention-based embedding T ctx defined in (12), when the input sequence contains an infinite number of tokens.

Proposition 5.2. In the asymptotic regime where L →∞ , one has that

Choosing the temperature λ = 1+4 σ 2 +2 dσ 4 4 ( 2 ( σ 2 + 1 2 ) 3 +( d -2) σ 6 ) gives

Moreover, the conditional variance of the embedding satisfies when L →∞

Choosing λ = 1 1+2 σ 2 , we obtain an unbiased embedding with a conditional variance of

As in the fixed centroids setting (Proposition 4.2), we retrieve that for a suitable choice of λ , the loss satisfies

This result is all the more surprising given that it emerges in a more complex setting, yet with a simpler mechanism. Unlike the attention-based predictor T lin ,µ 0 ,µ 1 , which benefits from access to the true centroids and a trainable architecture, the incontext encoder T ctx achieves a smaller asymptotic risk without any learned parameters. The ability of T ctx to extract meaningful representations illustrates the relevance of the key-query-value structure in the attention mechanism, even stripped of any learnable parameter and of the softmax nonlinearity.

It is worth noting that for d ≥ 2 , taking λ = 1 1+2 σ 2 ,

Therefore, in the regime of infinite input sequences ( L →∞ ) and for the appropriate value of λ , the embedding becomes unbiased and exhibits a variance reduction effect compared to the input data points, as illustrated in Figure 4 and previously discussed in Section 4.

6 Conclusion

This work offers a mathematically grounded, principled perspective on the unsupervised learning capabilities of attention mechanisms within mixture models. By combining a classical clustering framework with simplified yet non-trivial attention architectures, we present theoretical and empirical evidence showing that, when properly trained, attention layers can effectively recover latent structure in data. Our analysis provides insight into the training dynamics, quantization behavior, and design choices, such as attention head regularization.

We further investigate an in-context setting, where attention-based models still perform efficient approximate quantization, and achieve lower risk than the optimal quantizer. Future directions include exploring richer attention architectures, closer to that used in practice, which may further challenge the theoretical analysis.

1.5

1.0

0.5

0.0

0.5

1.0

1.5

PCA Component 1

0.75

0.50

0.25

0.00

0.25

0.50

0.75

PCA Component 2

Input tokens

Transformed tokens

0

,

1

Acknowledgments

The authors would like to thank Gérard Biau and Lenka Zdeborová for insightful discussions. This work was supported by the Projet Émergence(s) of the City of Paris. P.M. is supported by a Google PhD Fellowship. Part of this work was done while P.M. was visiting the Simons Institute for the Theory of Computing.

Code availability

Our code is available at https://github.com/rodrigomaulen/Attention-based-clustering

References

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum?id= LziniAXEI9 .

Luca Arnaboldi, Bruno Loureiro, Ludovic Stephan, Florent Krzakala, and Lenka Zdeborova. Asymptotics of sgd in sequence-single index models and single-layer attention networks. arXiv preprint arXiv:2506.02651 , 2025.

Hedy Attouch. Viscosity solutions of minimization problems. SIAM Journal of Optimization , 3: 769-806, 1996.

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. 3rd International Conference on Learning Representations , 2015.

Dzmitry Bahdanau, Jan Chorowski, and Dmitriy Serdyuk. Neural machine translation by jointly learning to align and translate. 3rd International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 4945-4949, 2016.

Nicolas Boumal. An Introduction to Optimization on Smooth Manifolds . Cambridge University, 2023.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs. https://github.com/ google/jax , 2018.

Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Lee Peter, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, and et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv:2303.12712 , 2023.

Hugo Cui. High-dimensional learning of narrow neural networks. Journal of Statistical Mechanics: Theory and Experiment , 2025(2):023402, 2025.

Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning in a solvable model of dot-product attention. In Advances in Neural Information Processing Systems , volume 37. Curran Associates, Inc., 2024.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv:1810.04805 , 2018.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, and Dirk Weissenborn. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations , 2020.

Takashi Furuya, Maarten V. de Hoop, and Gabriel Peyré. Transformers are universal in-context learners, 2024. URL https://arxiv.org/abs/2408.01367 .

Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes, 2023. URL https://arxiv.org/abs/2208. 01066 .

Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang. Flatten transformer: Vision transformer using focused linear attention, 2023. URL https://arxiv.org/abs/2308.00442 .

Bobby He and Thomas Hofmann. Simplifying transformer blocks. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id= RtDok9eS3s .

Yihan He, Hong-Yu Chen, Yuan Cao, Jianqing Fan, and Han Liu. Transformers versus the em algorithm in multi-class clustering. arXiv:2502.06007v1 , 2025.

Isserlis. On a formula for the product-moment coefficient of any order of a normal frequency distribution. Biometrika , 12(1):134-139, 1918.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention, 2020. URL https://arxiv.org/ abs/2006.16236 .

Kenneth Lange. Optimization , volume 2 edition. Springer, New York, 2013.

Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning, 2023. URL https://arxiv.org/ abs/2301.07067 .

Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Klusowkski Jason, Jianqing Fan, and Mengdi Wang. One-layer transformer provably learns one-nearest neighbor in context. Advances in Neural Information Processing Systems , 2024.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Stuart Lloyd. Least squares quantization in PCM. IEEE transactions on information theory , 28(2): 129-137, 1982.

Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing , pages 1412-1421, 2015.

Pierre Marion and Raphael Berthier. Leveraging the two timescale regime to demonstrate convergence of neural networks. In Advances in Neural Information Processing Systems , volume 36, 2023.

Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer. Attention layers provably solve single-location regression, 2024. URL https://arxiv.org/abs/2410.01537 .

Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris J. Maddison, and Daniel M. Roy. The shaped transformer: Attention models in the infinite depth-and-width limit. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum?id=PqfPjS9JRX .

Mary Phuong and Marcus Hutter. Formal algorithms for transformers. arXiv:2207.09238 , 2022.

Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone self-attention in vision models. In 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada , 2019.

Michael Shub. Global Stability of Dynamical Systems . Springer, New York, 1987.

Emanuele Troiani, Hugo Cui, Yatin Dandi, Florent Krzakala, and Lenka Zdeborova. Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds. In Forty-second International Conference on Machine Learning , 2025.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30:6000-6010, 2017.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174, 2023. URL https: //arxiv.org/abs/2212.07677 .

Hongru Yang, Zhangyang Wang, Jason D. Lee, and Yingbin Liang. Transformers provably learn two-mixture of linear classification via gradient flow. In The Thirteenth International Conference on Learning Representations , 2025. URL https://openreview.net/forum?id=AuAj4vRPkv .

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research , 25(49):1-55, 2024.

A Training dynamics in the degenerate case

In this section, we discuss the training behavior of the Transformer-based predictor T lin ,µ 0 ,µ 1 in the context of clustering, assuming the data are drawn from the degenerate mixture model, where for 1 ≤ ℓ ≤ L ,

(the orthogonal centroids µ ⋆ 0 and µ ⋆ 1 still lie on the unit sphere). We emphasize that despite its apparent simplicity, this study framework is already sufficient to reveal some of the complexity inherent in the clustering task carried out by a self-attention layer.

A.1 Theoretical analysis

Since training is performed by minimizing the risk R , the first steps of our analysis focus on studying the critical points and extrema of this risk.

Critical points and minimizers. First, we reparameterize the problem using the quantities

The scalar products κ 0 and κ 1 measure the alignment of the parameters µ 0 and µ 1 with the true centroids µ ⋆ 0 and µ ⋆ 1 , respectively, while the scalar products η 0 and η 1 capture their orthogonality with the inverted centroids µ ⋆ 1 and µ ⋆ 0 . The theoretical risk w.r.t. κ 0 , κ 1 , η 0 and η 1 reads as follows. The proof of this result and the following are given in Appendix B.

Proposition A.1. Under the degenerate mixture model ( P 0 ) , considering the attention-based predictor T lin ,µ 0 ,µ 1 composed of two linear heads parameterized by µ 0 and µ 1 , the theoretical risk R can be re-expressed as a function R < : R 4 → R such that R ( µ 0 , µ 1 ) = R < ( κ 0 , κ 1 , η 0 , η 1 ) , where

In addition, if ( µ 0 , µ 1 ) ∈ ( S d -1 ) 2 are prescribed to the unit sphere, then dom( R < ) = [ -1 , 1] 4 .

Remark A.2. After a direct computation, we note that the critical points of the risk R correspond to those of its reparameterized version, R < , i.e.,

Proposition A.3 (Characterization of global minima) . Consider R < : R 4 → R defined as in Proposition A.1 with λ ⋆ 0 = L +1 L +3 , then a point ( κ 0 , κ 1 , η 0 , η 1 ) belongs to argmin( R < ) if and only if

While the characterization can be made for any value of λ , choosing λ ⋆ 0 = L +1 L +3 simplifies the system by setting the first two conditions equal to 1. Moreover, this specific value provides a critical upper bound on the temperature parameter that guarantees recovery of the underlying centroids via risk minimization, as highlighted in the theorem below, and discussed in Remark B.2 in the appendices.

The proof of this result in Appendix B also characterizes all critical points, beyond global minima.

Convergence analysis. From the characterization of the minima of R given in Proposition A.3, we observe that the points ( κ 0 , κ 1 , η 0 , η 1 ) that saturate the first two equations (i.e., satisfy κ 2 0 = κ 2 1 = 1 and η 2 0 = η 2 1 = 0 ) correspond to global minimizers of the risk that also recover the centroids (up to a sign). However, in general, other global minimizers may exist that do not exhibit this saturation behavior and are therefore disconnected from centroid recovery. In the next result, we show that under appropriate initialization conditions, the ( PGD ) algorithm converges to the desired global minimum, which aligns with the clustering objective. To this end, we introduce the following manifold

Theorem A.4. Under the Dirac mixture model ( P 0 ) , consider the attention-based predictor T lin ,µ 0 ,µ 1 composed of two linear heads. Take λ ∈ ]0 , L +1 L +3 ] . Then there exists ¯ γ > 0 such that for any stepsize 0 < γ < ¯ γ , and for a generic initialization ( µ 0 0 , µ 0 1 ) ∈ ˜ M , the sequence of iterates ( µ k 0 , µ k 1 ) , generated by ( PGD ) , converges to the centroids (up to a sign), i.e.,

Proof. This is a direct consequence of Theorem 3.4 with σ = 0 .

This result demonstrates that, despite the non-convexity of the objective function, the key and query row matrices of a linear attention layer trained via ( PGD ) align with the centroids of the underlying Dirac mixture. Although the setting is simplified, it already highlights the representational role of key and query matrices in attention-based learning, and serves as a foundation for addressing the more general case of Gaussian mixtures.

Note that the convergence is up to a sign, a consequence of the symmetry inherent in H lin ,µ . Nonetheless, this sign ambiguity does not affect the output of the attention layer. To resolve this ambiguity and identify the true centroids, one could compare likelihoods or perform a hard assignment step, selecting the centroid that minimizes the total distance to all points within its assigned cluster. Besides, by generic initialization, we mean that the set of initializations ( µ 0 0 , µ 0 1 ) ∈ ˜ M for which ( PGD ) fails to recover the centroids is of Lebesgue measure zero with respect to ˜ M .

Initialization on the manifold ˜ M relies on prior knowledge of the centroids, which may be impractical. While a theoretical analysis under generic initialization on the unit sphere would be of interest, it remains analytically intractable due to the complexity of the resulting dynamical system derived from ( PGD ). In the following, however, we present numerical experiments incorporating a regularization term that proves effective in solving the problem without initialization constraints.

A.2 Numerical experiments

In this section, we study the empirical convergence of the ( PGD ) iterates when the data follows the degenerate mixture model ( P 0 ).

Results. Figure 5a clearly illustrates that when initialized on the manifold ˜ M , the ( PSGD ) iterates, over the objective function R , converge to the centroids, as established in Theorem A.4.

The situation differs outside the manifold, where numerical evidence shows that the Transformer parameters only partially align with the true centroids as shown in Figure 5b. In fact, we observe empirically that each parameter learns a mixture of both centroids. This indicates that the ( PSGD ) iterates may converge to optima that do not coincide with the underlying centroids. To mitigate this and better guide the learning process, we propose using a specific form of regularization:

0

2000

4000

6000

8000

10000

Iterations

10

16

10

11

10

6

10

1

Distance to centroid

(up to a sign)

(a) Initialization on the manifold

˜

M

0

2000

4000

6000

8000

10000

Iterations

10 0

4×10 1

6×10 1

Distance to centroid

(up to a sign)

(b) Initialization on the sphere

(

S

d

-

1

)

2

Therefore, we train the attention-based predictor H lin ,µ 0 ,µ 1 now by minimizing the regularized risk

where ρ > 0 denotes the strength of the regularization. It can be rigorously shown that as ρ approaches 0, the minimizers of R ρ converge to those of R , exhibiting the saturation phenomenon, desirable to bolster the interpretability of the attention heads. We refer to Appendix B.3 for more details.

0.0

0.1

0.2

0.3

Regularization

10

14

10

10

10

6

10

2

Distance to centroid

(up to a sign)

0

2000

4000

6000

8000

10000

Iterations

10

4

10

3

10

2

10

1

10 0

Distance to centroid

(up to a sign)

(a) Distance to centroids vs. regularization strength

(b) Distance to centroids vs iterations (

ρ

= 0

.

1

).

In Figure 6a, we observe that a relatively small regularization parameter (of the order of 10 -1 ) is sufficient to achieve centroid alignment, with numerical error below 10 -14 . In Figure 6b, we fix the regularization parameter and observe that over the course of 10 4 iterations, the attention head parameters exhibit linear convergence towards the true centroids. This numerical experiment highlights the effectiveness of this form of regularization in enhancing the interpretability of attention heads-by promoting their disentanglement-in the context of mixture models.

B Proofs of Section A (degenerate case)

In this section, we present the postponed proofs that support and elaborate on the arguments developed in the main text. We begin by characterizing the critical points of the Dirac mixture risk, then proceed to a discussion on the effects of a regularization term in the Dirac setting. Finally, we outline the proofs of Theorems A.4 and 3.4, which constitute the main theoretical results of our work.

B.1 Proof of Proposition A.1 (expression of the risk in the degenerate case)

To facilitate the analysis that follows, we introduce the notation e k ( µ ) def = λX ⊤ 1 µµ ⊤ X k , for 1 ≤ k ≤ L , which allows us to write

In what follows, we give an expression of the risk of an attention-based predictor, in the case where the data is distributed according to the Dirac mixture model ( P 0 ). Then, the risk of an attention-based predictor T lin ,µ 0 ,µ 1 can be written as

We compute ( I ) by conditioning on Z 1 , Z k ,

This leads to ( I ) = λ L -1 L ( κ 2 0 + κ 2 1 + η 2 0 + η 2 1 ) .

Similarly for ( II ) , conditioning on Z 1 , Z k ,

Which gives ( II ) = λ 2 L -1 L 2 ([ κ 2 0 + η 2 0 ] 2 +[ κ 2 1 + η 2 1 ] 2 ) + λ 2 2( L -1) L 2 ( κ 0 η 1 + κ 1 η 0 ) 2 .

Finally, to compute ( III ) , note that

leading to ( III ) = λ 2 ( L -1)( L -2) 2 L 2 [[ κ 2 0 + η 2 0 ] 2 +[ κ 2 1 + η 2 1 ] 2 +2( κ 0 η 1 + κ 1 η 0 ) 2 ]

In a similar fashion, we get that

Putting everything together, we obtain that the risk can be written in terms of κ 0 , κ 1 , η 0 , η 1 , i.e., R ( µ 0 , µ 1 ) = R < ( κ 0 , κ 1 , η 0 , η 1 ) , where:

B.2 Proof of Proposition A.3 (critical points of the risk in the degenerate case)

Proposition B.1. Consider R < : R 4 → R defined as in Proposition A.1 with λ = λ ⋆ 0 = L +1 L +3 , then we characterize its critical points by

The point (0 , 0 , 0 , 0) is a local maximum.

The points ( κ 0 , 0 , η 0 , 0) , where κ 2 0 + η 2 0 = 1 , and (0 , κ 1 , 0 , η 1 ) , where κ 2 1 + η 2 1 = 1 , are strict saddle points.

The points ( κ 0 , κ 1 , κ 1 , κ 0 ) and ( κ 0 , κ 1 , -κ 1 , -κ 0 ) , where κ 2 0 + κ 2 1 = L +3 2( L +1) , are strict saddle points.

( κ 0 , κ 1 , η 0 , η 1 ) belongs to argmin( R < ) if and only if:

Proof. Let us define ζ 0 = ( κ 0 , η 0 ) , ζ 1 = ( η 1 , κ 1 ) , then there exists a function R << : R 2 × R 2 → R , such that R < ( κ 0 , κ 1 , η 0 , η 1 ) = R << ( ζ 0 , ζ 1 ) , in fact, let us define A = ( L +1) 2 L ( L +3) , B = ( L +1) 2 2 L ( L +3) , C = ( L +1) 2 ( L -1) L ( L +3) 2 , then with the value of λ defined in the proposition, we obtain

To analyze its critical points, we take the partial derivatives,

̸

And also, we compute its Hessian, we define

Then the Hessian will be defined by

To find the critical points, we solve the following system of equations:

(0 , 0) is a local maximum. We see that a trivial solution to this system is ( ζ 0 , ζ 1 ) = (0 , 0) , and replacing into the Hessian matrix, we see directly that this point is a local maximum.

̸

(0 , ζ 1 ) , ( ζ 0 , 0) are strict saddle points. We check the case when ζ 0 = 0 , ζ 1 = 0 , then we need to solve

Since ζ 1 = 0 , this forces -2 A +4 B ∥ ζ 1 ∥ 2 = 0 . Replacing (0 , ζ 1 ) into the Hessian matrix gives us

And eig( ∇ 2 R << (0 , ζ 1 )) = eig( -2 AI 2 +2 Cζ 1 ζ ⊤ 1 ) ∪ eig(8 Bζ 1 ζ ⊤ 1 ) , where eig is the set of eigenvalues of a matrix. We have that

where we have used that 8 B = 4 A . We also note that C -A < 0 , then we conclude there are 2 negative eigenvalues and 1 positive eigenvalue, concluding that these points are strict saddle points, due to symmetry we conclude the same for the points of the form ( ζ 0 , 0) for ∥ ζ 0 ∥ 2 = 1 .

̸

̸

Non-trivial critical points. We will first show that the critical points that are of the form ( ζ 0 , ζ 1 ) for ζ 0 = 0 , ζ 1 = 0 necessarily satisfy ∥ ζ 0 ∥ = ∥ ζ 1 ∥ ̸ = 0 , multiplying the first equation of (20) by ζ 0 and the second by ζ 1 , then subtracting both resulting expressions we obtain 4 B ( ∥ ζ 1 ∥ 4 -∥ ζ 0 ∥ 4 ) = 2 A ( ∥ ζ 1 ∥ 2 -∥ ζ 0 ∥ 2 ) , and then

thus either ∥ ζ 0 ∥ = ∥ ζ 1 ∥ and we get the first claim, or ∥ ζ 0 ∥ 2 + ∥ ζ 1 ∥ 2 = 1 , in this second case we divide in two subcases:

Let us assume that ⟨ ζ 0 , ζ 1 ⟩ = 0 , then multiplying the first equation of (20) by ζ 0 and the second equation by ζ 1 , we get that ∥ ζ 0 ∥ 2 = ∥ ζ 1 ∥ 2 = 1 , which is a contradiction since we are in the case where ∥ ζ 0 ∥ 2 + ∥ ζ 1 ∥ 2 = 1 .

Therefore ⟨ ζ 0 , ζ 1 ⟩ ̸ = 0 , we multiply the first equation of (20) by ζ 1 and second by ζ 0 , after dividing by 2 ⟨ ζ 0 , ζ 1 ⟩ we get that

Substracting both equations we get that ∥ ζ 0 ∥ = ∥ ζ 1 ∥ .

So we get that necessarily ∥ ζ 0 ∥ = ∥ ζ 1 ∥ = r > 0 . Then the equation (20) becomes

( ζ 0 , ± ζ 0 ) are strict saddle points. In the case where ζ 0 = ± ζ 1 , by (21) we get that when ζ 0 = ζ 1 , then

and ∥ ζ 0 ∥ = r 1 , with r 2 1 = A A -C = L +3 4 . Replacing this point on the Hessian matrix ∇ 2 R << ( ζ 0 , ± ζ 0 ) ,

We can write the Hessians in block form with 2 × 2 diagonal blocks

Considering vectors of the form ( x, ± x ) , x ∈ R 2 , this reduces the eigenvalue problem to the 2 × 2 matrices M 0 ± M 1 . Substituting r 2 0 = A/ ( A + C ) or r 2 1 = A/ ( A -C ) shows that each matrix has one positive eigenvalue along ζ 0 (equal to 4 A ) and one negative eigenvalue orthogonal to ζ 0 (equal to -4 AC A ± C ), so both Hessians are indefinite.

̸

Characterization of global minima. If ζ 0 = ± ζ 1 and ∥ ζ 0 ∥ = ∥ ζ 1 ∥ = r > 0 , then both vectors are linearly independent, thus the first equation of (21) is only possible when r 2 = 1 and ⟨ ζ 0 , ζ 1 ⟩ = 0 , in which case we have to analyze the points ( ζ 0 , ζ 1 ) such that ⟨ ζ 0 , ζ 1 ⟩ = 0 and ∥ ζ 0 ∥ 2 = ∥ ζ 1 ∥ 2 = 1 , we replace these points on the Hessian matrix and this gives us

A direct computation of the eigenvalues with eigenvectors ( ζ 1 , ζ 0 ) and ( ζ 1 , -ζ 0 ) gives us that all the eigenvalues are positive in this case, since R << is coercive, these points are in fact global minima.

Remark B.2. Note that the specific characterization of the global minima of R < was valid only for λ = λ ⋆ 0 = L +1 L +3 . However, when restricting the analysis to the manifold ˜ M and considering λ ∈ ]0 , L +1 L +3 [ , the global minima lie outside the domain [ -1 , 1] 2 . As a result, due to the structure of the update rule in ( PGD ) , the extreme points ( ± 1 , ± 1) of [ -1 , 1] 2 become fixed points of the algorithm and serve as global minimizers of R < .

and off-diagonal blocks

B.3 Discussion on regularization

In order to solve the clustering problem in the degenerate case, we train the attention-based predictor H lin ,µ 0 ,µ 1 now by minimizing the regularized risk

where r ( µ 0 , µ 1 ) = E [ ⟨ µ 0 , X 1 ⟩ 2 ⟨ µ 1 , X 1 ⟩ 2 ] , and ρ > 0 denotes the strength of the regularization.

It is direct to check that there exists r < : R 4 → R , such that r ( µ 0 , µ 1 ) = r < ( κ 0 , κ 1 , η 0 , η 1 ) according to the notation defined in (13), and r < ( κ 0 , κ 1 , η 0 , η 1 ) = 1 2 ( κ 2 0 η 2 0 + κ 2 1 η 2 1 ) . We define the following optimization problem

where R < is defined in Proposition A.1. Since R < and r < are coercive, we apply Attouch (1996, Theorem 2.1) to conclude that if u ρ ∈ [ -1 , 1] 4 is a solution of ( ˜ P < ρ ), then every limit point ˆ u of u ρ , when ρ → 0 , satisfies that:

Due to the geometry of r < and the characterization of argmin R < we got in Proposition 15, we obtain that if ˆ u = ( ˆ κ 0 , ˆ κ 1 , ˆ η 0 , ˆ η 1 ) , then

Then the optimal solution for the regularized problem when ρ → 0 achieves a saturation effect, corresponding to global minimizers that recover the centroids. However, due to the non-convex nature of the problem, it is not guaranteed a priori that PGD on ( ˜ P < ρ ) will converge to the desired solution. A possible direction of analysis is to study the dynamics of PGD in the limit where ρ → 0 . We know from Proposition B.1 that the only global minimizers of the unregularized problem lie on a manifold, so we expect that PGD converges to this manifold, before evolving on the manifold due to the regularization term, to converge to the minimizers given by (22). Technically, this dynamics could be studied by using two-timescale tools, e.g. similar to Marion and Berthier (2023) and references therein. We leave this analysis for future work.

C Proofs of Section 3 (Gaussian mixture model)

C.1 Proof of Proposition 3.1 (expression of the risk in the non-degenerate case).

Recall the notation e k ( µ ) def = λX ⊤ 1 µµ ⊤ X k , for 1 ≤ k ≤ L , which allows us to write

Under the Gaussian mixture model, we are going to show that the risk

admits a closed-form representation in terms of elementary functions. It holds that

We now proceed to compute each of the six terms. To compute ( I 0 ) , we can use Lemma E.2, since

To compute ( II 0 ) , by defining p 0 ( µ 0 , µ 1 , µ ⋆ ) as in Lemma E.3, we get

Then,

To compute ( III 0 ) , by defining p 1 ( µ 0 , µ 1 , µ ⋆ Z 1 , µ ⋆ Z 2 ) as in Lemma E.4, we get

where

And then

Consequently,

To compute ( I ) , we can use Lemma E.5 to obtain:

Thus,

(

I

) =

λ

L

-

L

[(

κ

0

+

η

0

+

κ

1

+

η

1

)(1 + 4

σ

) + 4

σ

(

∥

µ

0

∥

+

∥

µ

1

∥

)]

.

To compute ( II ) , by defining p 2 ( µ 0 , µ 1 , µ ⋆ Z 1 , µ ⋆ Z 2 ) as in Lemma E.6, we obtain:

1

2

2

2

2

2

4

2

2

where

And then

So we obtain,

Finally, to compute ( III ) , by defining p 3 ( µ 0 , µ 1 , µ ⋆ Z 1 , µ ⋆ Z 2 , µ ⋆ Z 3 ) as in Lemma E.7, we get

where

Observing that

we get

Finally, putting everything together, recalling the notation introduced in (7), and inspecting the formulas given by Lemmas from E.2 to E.7 allows to conclude.

C.2 Proof of Lemma 3.2 (expression of the risk on the manifold, non-degenerate case)

We provide in Lemma C.1 a more precise version of Lemma 3.2, with explicit constants.

Lemma C.1. Define c 1 ( n ) = 1 + nσ 2 and c 2 ( n ) = 1 + σ 2 ( d + n ) , then the risk R < ( κ 0 , κ 1 ) restricted to M has the form

where

Proof of Lemma C.1. Using the decomposition obtained in the proof of Proposition 3.1, after simple algebraic manipulation we get that on this manifold:

We conclude by noting that the risk restricted to this manifold is

and properly factorizing the terms.

C.3 Proof of Proposition 3.3 (global minima of the risk, non-degenerate case).

In what follows we provide an extended version of Proposition 3.3 with explicit constant, together with its proof.

Proposition C.2. Let us define

and consider R < ( κ 0 , κ 1 ) with the following λ :

Then the points ( ± 1 , ± 1) are global minimum of R < ( κ 0 , κ 1 ) .

Remark C.3. In the case where σ = 0 , we get λ ⋆ (0 , L ) = λ ⋆ 0 = L +1 L +3 .

Proof. Imposing first order conditions on R < ( κ 0 , κ 1 ) from Lemma C.1, we obtain an explicit form of its critical points. From this expression, we note that the global minimum are the points ( ± 1 , ± 1) if and only if 4 A +2 B +2 C = 0 . The function λ ↦→ 2 A ( λ ) + B ( λ ) + C ( λ ) is a quadratic which is negative for 0 ≤ λ < λ ⋆ ( σ, L ) , and vanishes at λ = λ ⋆ ( σ, L ) .

C.4 Proof of Theorem 3.4

The proof of this result is built upon a series of intermediate results that progressively lead to the desired conclusion.

Lemma C.4. At a point ( κ 0 , κ 1 , η 0 , η 1 , ξ ) such that η 0 = η 1 = ξ = 0 , we have ∂ η 0 R < = ∂ η 1 R < = ∂ ξ R < = 0 .

Proof. According to Proposition 3.1, we can directly obtain that

Taking the partial derivative in η 0 , we get

At a point such that η 0 = η 1 = ξ = 0 , this gives ∂ η 0 R < ( κ 0 , κ 1 , 0 , 0 , 0) = -∂ η 0 R < ( κ 0 , κ 1 , 0 , 0 , 0) , therefore ∂ η 0 R < ( κ 0 , κ 1 , 0 , 0 , 0) = 0 , the proof for ∂ η 1 R < , ∂ ξ R < is analogous.

Lemma C.5. The manifold M is invariant under ( PGD ) dynamics, this is if ( µ k 0 , µ k 1 ) ∈ M , then ( µ k +1 0 , µ k +1 1 ) ∈ M .

Proof. We apply the chain rule and Lemma C.4 to get:

We then follow the same ideas as in Marion et al. (2024, Lemma 4), where our Lemma C.4 takes the role of Marion et al. (2024, Lemma 14).

More concretely, let us consider c 0 = ∥ µ k 0 -γ ( I d -µ k 0 ( µ k 0 ) ⊤ ) ∇ µ 0 R ( µ k 0 , µ k 1 ) ∥ 2 , and c 1 = ∥ µ k 1 -γ ( I d -µ k 1 ( µ k 1 ) ⊤ ) ∇ µ 1 R ( µ k 0 , µ k 1 ) ∥ 2 , then recalling ( PGD ) updates, we have that if ( µ k 0 , µ k 1 ) ∈ M , then

And

where the last term is zero since

Lemma C.6. When initialized on the manifold M , the iterations generated by ( PGD ) can be reformulated as follows:

where φ : [ -1 , 1] 2 → [ -1 , 1] 2 is given by

and R < ( κ 0 , κ 1 ) def = R < ( κ 0 , κ 1 , 0 , 0 , 0) as in Lemma C.1.

Proof. By definition of the iterates and (23),(24), we have

In the following propositions, we consider R < : [ -1 , 1] 2 → R + defined as in Lemma C.1 with λ ∈ ]0 , λ ∗ ( σ, L )] , where λ ∗ ( σ, L ) is defined in Proposition C.2.

Proposition C.7. Let ( µ 0 0 , µ 0 1 ) ∈ M , consider (25) with initial conditions κ 0 0 = ⟨ µ 0 0 , µ ⋆ 0 ⟩ , κ 0 1 = ⟨ µ 0 1 , µ ⋆ 1 ⟩ . Then there exists ¯ γ > 0 such that for every 0 < γ < ¯ γ , the risk R < is decreasing along the iterates of (25) . Besides, the distance between successive iterates tends to zero, and, if ( κ ⋆ 0 , κ ⋆ 1 ) is an accumulation point of the sequence of iterates ( κ k 0 , κ k 1 ) k ∈ N , then

Proof. The proof is identical to that of Marion et al. (2024, Proposition 8) and is therefore omitted.

Proposition C.8. The points ( κ 0 , κ 1 ) ∈ [ -1 , 1] 2 satisfying (26) belong to the set

Proof. We recall that by Lemma C.1, the risk R < restricted to the manifold M , has the following form

Then

And we can rewrite equations (26) as

Since each equation is a product of 3 terms, the general solution to this system occurs when at least in each equation is zero. By considering only the first two terms in each equation, we obtain the solution set { ( ± 1 , ± 1) , (0 , ± 1) , ( ± 1 , 0) , (0 , 0) } . Now we consider the case when 2 Aκ 2 0 + B + Cκ 2 1 = 0 , this implies that:

If κ 1 = 0 , then κ 2 0 = -B 2 A .

If κ 2 1 = 1 , then κ 2 0 = -B + C 2 A .

If 2 Aκ 2 1 + B + Cκ 2 0 = 0 , then 2 Aκ 2 0 + Cκ 2 1 = 2 Aκ 2 0 + Cκ 2 1 , thus (2 A -C )( κ 2 0 -κ 2 1 ) = 0 . By inspection, C < 2 A , hence we get κ 2 0 = κ 2 1 and κ 2 0 = -B 2 A + C .

We note the following relation

Further remark by inspecting the proof of Proposition C.2 that for λ ∈ ]0 , λ ∗ ( σ, L )] , we have

Thus the only possible solution when 2 Aκ 2 0 + B + Cκ 2 1 = 0 is κ 2 1 = 1 and κ 2 0 = -B + C 2 A = 1 Putting everything together, the solution set is precisely

Proposition C.9. The fixed points of the dynamic can be classified as follows:

The points ( κ 0 , κ 1 ) = ( ± 1 , ± 1) are global minima of R < on [ -1 , 1] 2 .

The points ( κ 0 , κ 1 ) = (0 , ± 1) and ( ± 1 , 0) are strict saddle points of R < on [ -1 , 1] 2 .

The point ( κ 0 , κ 1 ) = (0 , 0) is a global maxima of R < on [ -1 , 1] 2 .

Proof. Since R < is smooth, its extrema on the square [ -1 , 1] 2 occur either at critical points, where the gradient vanishes, or on the boundary. For the chosen range λ ∈ ]0 , λ ⋆ ( σ, L )[ , the only interior critical point is (0 , 0) , which corresponds to a local maximum. Indeed, the Hessian of R < at (0 , 0) is

We easily check that both eigenvalues are 2 B , which is negative since B < 0 , therefore concluding that (0 , 0) is a maximum. The rest of the extrema of R < must lie on the boundary of the square, we observe that

Thus concluding that ( ± 1 , ± 1) are minimizers. Finally, we observe that

Analogously, we see that

This shows that (0 , ± 1) and ( ± 1 , 0) are strict saddle points, concluding with the proof. We note that when λ = λ ⋆ ( σ, L ) , Proposition C.2 implies that the critical points coincide with the minimizers already identified.

Proposition C.10. Consider the context of Proposition C.7, then there exists ¯ γ > 0 such that for any stepsize 0 < γ < ¯ γ , the iterates ( κ k 0 , κ k 1 ) k ∈ N generated by (25) converge to an element of C .

Proof. By Proposition C.7, the distance between successive iterates ( κ k 0 , κ k 1 ) k ∈ N , then the set of accumulation points of the sequence is connected (Lange, 2013, Proposition 12.4.1). Since we have a finite number of accumulation points by Proposition C.8, the sequence has a unique accumulation point. Besides, the sequence belongs to the compact set [ -1 , 1] 2 , then it converges and its limit is one of the nine fixed points.

Proposition C.11. Consider the context of Proposition C.7, then there exists ¯ γ > 0 such that for any stepsize 0 < γ < ¯ γ , the set of initializations such that the iterates ( κ k 0 , κ k 1 ) k ∈ N generated by (25) converge to (0 , ± 1) , ( ± 1 , 0) or (0 , 0) has Lebesgue measure zero (with respect to the Lebesgue measure on the manifold M ).

Proof. The point (0 , 0) is a maxima of the risk R < on [ -1 , 1] 2 and the value of the risk decreases along the iterates of ( PGD ) by Proposition C.7. We follow the ideas presented in the proof of Marion et al. (2024, Proposition 12), we can conclude that φ is differentiable on [ -1 , 1] 2 , and that its Jacobian is not degenerate, besides φ is a local diffeomorphism around (0 , ± 1) and ( ± 1 , 0) , whose Jacobian matrix in each point has one eigenvalue in [0 , 1[ and one eigenvalue in ]1 , ∞ [ . The result follows from the Center-Stable Manifold Theorem (Shub, 1987, Theorem III.7), we refer to Marion et al. (2024, Proposition 13) for a detailed and analogous proof.

D Proofs of Section 4

D.1 Proof of Proposition 4.1

We provide hereafter the proof of Proposition 4.1, which proof follows.

Proposition D.1. Consider ( X ℓ ) 1 ≤ ℓ ≤ L i.i.d. drawn from ( P σ ) . Consider also Z ℓ ∈ { 0 , 1 } the latent variable of X ℓ , i.e. X ℓ | Z ℓ ∼ N ( µ ⋆ Z ℓ , σ 2 I d ) , and

where e k ( µ ) def = λ ⟨ X 1 , µ ⟩⟨ X k , µ ⟩ . Then

Moreover, when λ = L ( L +1)+2( L +3) σ 2 , then the encoding is unbiased, this is

Proof. We decompose the following term as follows,

Due to the independence of the variables,

On the one hand we have

On the other hand,

Therefore,

And then,

Which let us conclude that for c ∈ { 0 , 1 } ,

Proposition D.2. Consider ( X ℓ ) 1 ≤ ℓ ≤ L i.i.d. drawn from ( P σ ) . Consider also Z ℓ ∈ { 0 , 1 } the latent variable of X ℓ , i.e. X ℓ | Z ℓ ∼ N ( µ ⋆ Z ℓ , σ 2 I d ) , and

And

When λ = L ( L +1)+2( L +3) σ 2 , the encoding is unbiased, with variance

Besides, when L →∞ we get that,

In general, if λ is not fixed and L →∞ , we get

Notably, by selecting λ to ensure an unbiased encoding, the variance becomes independent of the dimension d and equals 2 σ 2 . This shows a variance reduction effect whenever the dimension d is bigger than 2. More generally, λ can be chosen independently of d such that

In this regime, the encoding also asymptotically reduces the variance of X 1 , conditioned on its cluster assignment, as the number of components L →∞ .

Proof. We note that the needed computations were already stated in the proof of Proposition 3.1, we follow as in the proof of Lemma C.1, without loss of generality, we assume Z 1 = 0 , then we get that for µ 0 , µ 1 ∈ M :

Recalling that κ 0 = ⟨ µ 0 , µ ⋆ 0 ⟩ , κ 1 = ⟨ µ 1 , µ ⋆ 1 ⟩ , in order to compute E [ ∥ T lin ,µ ∗ 0 ,µ ∗ 1 ( X ) 1 ∥ 2 | Z 1 = 0] , we just need to replace κ 0 and κ 1 by 1 in the previous expression, as follows

The expression of the variance comes from subtracting to this term the square of the conditional expectation given in Proposition 4.1. The asymptotic expressions are then straightforward to derive.

D.2 Proof of Proposition 4.2

Proof. Assume that the tokens are i.i.d., such that for any ℓ , X ℓ ∼ 1 2 N ( µ ⋆ 0 , σ 2 I d ) + 1 2 N ( µ ⋆ 1 , σ 2 I d ) . The risk of the oracle predictor T lin ,µ ⋆ 0 ,µ ⋆ 1 can be decomposed as follows

where, from the proof of Lemma C.1, by taking κ = κ = 1

When L tends to ∞ , only the first term together with ( I ) and ( III ) contribute. Therefore, we obtain that

Choosing λ = 1+4 σ 2 +4 σ 4 1+6 σ 2 +12 σ 4 +8 σ 6 (its value being independent of L ) minimizes the equivalent bound obtained above. With such a choice, the equivalent becomes

Remark D.4. In the case of the degenerate case ( σ = 0 ), similar computations lead to

Optimizing this quantity w.r.t. λ leads to choose λ = L +1 L +3 , plugging this value for λ gives

In the degenerate case, we observe that as the sequence length tends to infinity, the risk of the attention-based predictor with oracle parameters converges to zero, matching that of the optimal quantizer.

E Technical results

This section gathers a series of technical results about Gaussian random variables, used to derive expression of the risk in the rest of the document.

Lemma E.1. (Isserlis, 1918). Consider G ∼ N (0 , σ 2 I d ) and µ a , µ b , µ c ∈ R d , then

Lemma E.2. Consider X ∼ N ( µ ⋆ , σ 2 I d ) where ∥ µ ⋆ ∥ = 1 and µ a ∈ R d , then

Proof. We decompose X as follows,

Then, by Lemma E.1 one obtains

Lemma E.3. Let X ∼ N ( µ ⋆ , σ 2 I d ) , where ∥ µ ⋆ ∥ = 1 and µ a , µ b ∈ R d , then

can be expressed as

Proof. Write X = µ ⋆ + ε with ε ∼ N (0 , σ 2 I d ) . Then

Expand the product and drop all odd-moment terms of ε (their expectation is 0 ). The surviving types of expectations are of the form of Lemma E.1, collecting all nonzero contributions after expansion and simplifying gives the stated formula for p 0 ( µ a , µ b , µ ⋆ ) .

The proofs of Lemmas E.4-E.7 follow the same approach as the proof of Lemma E.3: we rewrite X i | Z i as µ ⋆ Z i + ε i , where ( ε i ) i are i.i.d. random variables with ε i ∼ N (0 , σ 2 I d ) . We then expand the product, discard odd-moment terms, and apply Lemma E.1 to obtain the desired results.

Lemma E.4. Let Z 1 and Z 2 ∈ { 0 , 1 } be fixed. Consider two independent R d -valued random variables X 1 and X 2 , such that

where the unit vectors µ ⋆ a , µ ⋆ b (i.e., ∥ µ ⋆ a ∥ = ∥ µ ⋆ b ∥ = 1 ) are orthogonal. For µ a , µ b , µ c ∈ R d , define

This quantity satisfies

Moreover, we define

which satisfies

Lemma E.5. Let Z 1 , Z 2 ∈ { 0 , 1 } be fixed. Consider two independent R d -valued random variables X 1 and X 2 , such that

where the unit vectors µ ⋆ a , µ ⋆ b (i.e., ∥ µ ⋆ a ∥ = ∥ µ ⋆ b ∥ = 1 ) are orthogonal. For µ a , µ b ∈ R d , we get that

Lemma E.6. Let Z 1 , Z 2 ∈ { 0 , 1 } be fixed. Consider two independent R d -valued random variables X 1 and X 2 , such that

where the unit vectors µ ⋆ a , µ ⋆ b (i.e., ∥ µ ⋆ a ∥ = ∥ µ ⋆ b ∥ = 1 ) are orthogonal. For µ a , µ b ∈ R d , define also:

These quantities satisfy

Moreover, we define

which satisfies

Lemma E.7. Let Z 1 , Z 2 , Z 3 ∈ { 0 , 1 } be fixed. Consider three independent R d -valued random variables X 1 , X 2 , X 3 , where

such that µ ⋆ a , µ ⋆ b unit vectors (i.e., ∥ µ ⋆ a ∥ = ∥ µ ⋆ b ∥ = 1 ) are orthogonal. For µ a , µ b ∈ R d , define also:

These quantities satisfy

Moreover, we define p 3 ( µ a , µ b , µ ⋆ Z 1 , µ ⋆ Z 2 , µ ⋆ Z 3 ) def = E [ ⟨ X 1 , µ a ⟩⟨ X 2 , µ a ⟩⟨ X 1 , µ b ⟩⟨ X 3 , µ b ⟩⟨ X 2 , X 3 ⟩| Z 1 , Z 2 , Z 3 ] , which satisfies

F Experimental details

This section provides algorithmic details, choices of parameters, and settings used for the plots displayed in Sections A and 3.

F.1 Projected Stochastic Gradient Descent

We formally define the method Projected Stochastic Gradient Descent (PSGD), which we run for our numerical experiments.

PSGD iterates for linear attention heads. Given the objective function R ρ : ( S d -1 ) 2 → R defined in ( P ρ ), we define h : ( S d -1 ) 2 × R L × d as

where X i is the i -th row of the matrix X . Consequently we can write

where D is the distribution over R L × d where each row is i.i.d. according to

Then, given and an initialization ( µ 0 0 , µ 0 1 ) ∈ ( S d -1 ) 2 , a stepsize γ , we define ( µ k 0 , µ k 1 ) ∈ ( S d -1 ) 2 recursively by:

where M is called the batch size, and for each k ∈ N , ( ξ k i ) i = { 1 ,...,M } are M independents samples of D .

PSGD iterates for softmax attention heads. Given the objective function R soft ,ρ 0 : ( S d -1 ) 2 × R 2 → R defined in ( P ρ 0 ), for simplicity, we note that for an appropriate h 0 , we can write

where D is the distribution over R L × d where each row is i.i.d. according to

Then, given and an initialization ( µ 0 0 , µ 0 1 ) ∈ ( S d -1 ) 2 , ( ψ 0 , λ 0 ) = (2 , 3) , a stepsize γ , we define ( µ k 0 , µ k 1 ) ∈ ( S d -1 ) 2 and ( ψ k , λ k ) ∈ R 2 recursively by:

where M is called the batch size, and for each k ∈ N , ( ξ k i ) i = { 1 ,...,M } are M independents samples of D .

Remark F.1. Gradient computations in the numerical experiments were carried out using JAX (Bradbury et al., 2018).

F.2 Experimental details

In the following, we provide the experimental setup corresponding to Sections A and 3.

We use input sequences of length L = 30 of 5-dimensional tokens ( d = 5 ), and define the true centroids as µ ⋆ 0 = (0 , 0 , 0 , 0 , 1) and µ ⋆ 1 = ( -1 , 0 , 0 , 0 , 0) . We recall that the metric used to quantify the distance to the centroids (up to a sign) is defined in (11).

Experimental details of Section A. Regarding the experiment on the manifold, i.e., Figure 5a, we perform 10 4 ( PSGD ) iterations without regularization ( ρ = 0 ) with a learning rate of γ = 0 . 01 , λ = 0 . 6 , batch size M = 256 . The experiment is repeated across 10 independent runs, each initialized randomly on the manifold ˜ M .

For the rest of the experiments of this section, we adopt the same setup as before, with the exception that each run is randomly initialized on the unit sphere. In Figure 5b, we perform 10 4 iterations to observe that without adding a regularization term, we only get partial alignment of the Transformer parameters towards the true centroids.

Then, in Figure 6a we perform 5 × 10 3 iterations of ( PSGD ) to minimize the regularized risk R ρ for 15 values of the regularization strength ρ , linearly spaced in [0 , 0 . 3] . Finally, in Figure 6b we choose ρ = 0 . 1 and perform 10 4 ( PSGD ) iterations.

Experimental details of Section 3. Regarding the experiment on the manifold, i.e., Figure 1, we run the algorithm for 10 4 iterations without regularization ( ρ = 0 ), with a learning rate of γ = 0 . 01 , batch size M = 256 , and choosing λ = 0 . 6 for σ = 0 . 3 , and λ = 0 . 2 for σ = 1 . The experiment is repeated across 10 independent runs, each initialized randomly on the manifold M .

For the rest of the experiments of this section, we adopt the same setup as before, with the exception that each run is randomly initialized on the unit sphere. In Figure 2a we perform 5 × 10 3 iterations of ( PSGD ) to minimize the regularized risk R ρ for 30 values of the regularization strength ρ , linearly spaced in [0 , 3] . Finally, in Figure 2b we choose ρ = 0 . 2 and perform 10 4 ( PSGD ) iterations.

Remark F.2. All experiments in Section A and 3 can be run on a standard laptop. Most complete within a few minutes, with the exception of those in Figures 6a and 2a, which require approximately 20 minutes and up to an hour, respectively, due to repeated problem-solving across a grid of regularization strengths.

G Additional numerical experiments

In what follows, we first present numerical experiments in dimension 100. We then vary the dimension from 4 to 200. Results are shown only for the linear approach, as the softmax variant exhibits numerical instability in higher dimensions.

G.1 Influence of the dimension

Experiments in R 100 . We use input sequences of length L = 30 in R 100 , where we define two centroids: µ ⋆ 0 = (0 , . . . , 0 ︸ ︷︷ ︸ 99 times , 1) and µ ⋆ 1 = ( -1 , 0 , . . . , 0 ︸ ︷︷ ︸ 99 times ) . The model is trained using ( PSGD ) with an online batch sampling strategy, with a batch size of 256, and a learning rate of 0.01. Due to the big dimensionality of the problem, we modify the concept introduced as distance to the centroid up to a sign by the concept of minimal root mean squared error, which is nothing but the distance to the centroid (up to a sign) divided by the square root of the dimension, i.e.,

0

2000

4000

6000

8000

10000

Iterations

10

17

10

12

10

7

10

2

RMSE

(a)

σ

= 0

.

0

2000

4000

6000

8000

10000

Iterations

10

2

10

1

RMSE

(b)

σ

= 0

.

3

.

0

2000

4000

6000

8000

10000

Iterations

10

1

3×10 2

4×10 2

6×10 2

2×10 1

RMSE

(c)

σ

= 1

.

0

2000

4000

6000

8000

10000

Iterations

10

5

10

4

10

3

10

2

10

1

RMSE

(a)

σ

= 0

.

0

2000

4000

6000

8000

10000

Iterations

10

2

10

1

RMSE

(b)

σ

= 0

.

3

.

0

2000

4000

6000

8000

10000

Iterations

10

1

3×10 2

4×10 2

6×10 2

2×10 1

RMSE

(c)

σ

= 1

where S 2 is the symmetric group of order 2, µ ⋆ 0 , µ ⋆ 1 are the true centroids, and ˆ µ 0 , ˆ µ 1 are the returned parameters from ( PSGD ). In Figures 7, 8 we can observe the behavior of the RMSE over the iterations for different levels of noise σ . We remark that in Figure 7 we initialize on the manifold M , and there is no regularization term (i.e. ρ = 0 ), in Figure 8 we initialize randomly over the unit sphere and we set ρ = 0 . 2 . In both experiments we set λ = 0 . 6 for the case σ = 0 and σ = 0 . 3 , and λ = 0 . 2 for the case σ = 1 .

In each experiment, the RMSE is of the order 10 -2 , which can be interpreted as, on average per coordinate, the estimators ˆ µ 0 , ˆ µ 1 are missing the true parameters µ ⋆ 0 , µ ⋆ 1 by 10 -2 , suggesting a high level of accuracy in the estimation procedure.

Making d vary. We repeat the same experiment as before, just varying the dimension of the problem, the two centroids in R d are defined by µ ⋆ 0 = (0 , . . . , 0 ︸ ︷︷ ︸ d-1 times , 1) and µ ⋆ 1 = ( -1 , 0 , . . . , 0 ︸ ︷︷ ︸ d-1 times ) . For d ranging between 4 and 200 , we show in Figures 9 and 10, on the x-axis the dimension of the problem and on the y-axis the minimal RMSE after 5000 iterations. We remark that in Figures 9b and 9c, the y-axis displays only a single line due to the logarithmic scale; this may appear misleading, but it simply indicates that the next power of 10 is much larger and is not captured on the plot.

Regardless of the initialization regime, in the noiseless case ( σ = 0 ) the minimal RMSE decreases as the problem dimension d grows. By contrast, for any strictly positive noise level, the minimal RMSE increases slowly with d -for σ = 0 . 3 it remains of order 10 -3 , and for σ = 1 of order 10 -2 . This reflects the growing difficulty of the problem as both dimensionality and noise increase. We recall that the minimal RMSE can be interpreted as the average discrepancy per coordinate between the estimated parameters and the true centroids, suggesting a high level of accuracy in each experiment.

0

50

100

150

200

Dimension

10

10

10

9

10

8

RMSE

(a)

σ

= 0

.

0

50

100

150

200

Dimension

10

3

RMSE

(b)

σ

= 0

.

3

.

0

50

100

150

200

Dimension

10

2

RMSE

(c)

σ

= 1

.

0

50

100

150

200

Dimension

10

3

10

2

RMSE

(a)

σ

= 0

.

0

50

100

150

200

Dimension

10

3

2×10 3

3×10 3

4×10 3

RMSE

(b)

σ

= 0

.

3

.

0

50

100

150

200

Dimension

10

2

2×10 2

3×10 2

4×10 2

RMSE

(c)

σ

= 1

.

0

2000

4000

6000

8000

10000

Iterations

10

1

RMSE

Sigma

0.3

1.0

Regarding running times, the experiments run in Figures 7 and 8 take approximately two hours on a standard laptop, while that of Figures 9 and 10 may require up to 12 hours, to cover the dimension grid.

G.2 Relaxing the orthogonality assumption

We replicate the experiments and parameter-selection procedure from Section G.1, but this time initializing the centroids and the initial points uniformly at random on the sphere S d -1 in each run. Figure 11 illustrates the algorithm's convergence behavior over 10000 iterations in the case d = 50 . In contrast, Figure 12 shows the minimal RMSE of ( PSGD ) after 5000 iterations for dimensions d ranging from 4 to 100.

We observed the expected behavior: as the dimension increases, randomly initializing centroids on the sphere makes them more likely to be orthogonal, and thus training via the regularized theoretical risk yields better results at higher dimensions. This effect is stronger at lower noise levels and becomes noticeably clearer beyond 40 dimensions.

Regarding running times, the experiments run in Figure 11 take approximately one hour on a standard laptop, while those of Figure 12 may require up to 7 hours, to cover the dimension grid.

0

20

40

60

80

100

Dimension

10

2

10

1

RMSE

Sigma

0.3

1.0

G.3 Euclidean Projected SGD

Based on Section F.1 and the algorithm ( PSGD ) presented there, we define the projected Euclidean SGD as the following update rule, with initialization ( µ 0 0 , µ 0 1 ) ∈ ( S d -1 ) 2 , and stepsize γ :

where h is defined in (28).

Our empirical analysis reveals that it is a viable and simpler alternative to the theoretically grounded Riemannian gradient. This section provides the numerical evidence supporting this choice. To compare the performance of Euclidean SGD against the Riemannian SGD (both projected on the unit sphere), we repeated the experiments outlined in Section 3. The results are presented in Figure 13. This empirical equivalence confirms that the simple projection of the Euclidean gradient serves as an effective and computationally simpler proxy for the true Riemannian gradient.

0

2000

4000

6000

8000

10000

Iterations

10

2

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

=0.3

=1

0

2000

4000

6000

8000

10000

Iterations

10

3

10

2

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

=0.3

=1

(a) Distance to centroids vs Projected SGD iterations for the minimization of R , initialization on the manifold. (b) Distance to centroids vs Projected SGD iterations for the minimization of R ρ , initialization on the unit sphere, with regularization ρ = 0 . 2 .

G.4 Extension to Gaussian mixture model with 3 components

We propose an extension of our work to the case of three orthonormal centroids. We believe that the approach described below would further generalize to the case of K orthonormal centroids with K < d . Specifically, we assume that the tokens are i.i.d. drawn from the mixture model

where µ ⋆ 0 , µ ⋆ 1 , µ ⋆ 2 are orthonormal vectors. It is natural to consider an attention-based predictor composed of three attention heads, parameterized by µ 0 , µ 1 , µ 2 ∈ R d ,

The associated risk is R ( µ 0 , µ 1 , µ 2 ) = E [ ∥ X 1 -T lin ,µ 0 ,µ 1 ,µ 2 ( X ) 1 ∥ 2 2 ] . There are two natural generalizations of the regularization term to this case:

0

2500

5000

7500

10000

Iterations

10

2

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

0.3

1.0

(a) With regularization r (1) .

0

5000

10000

15000

20000

Iterations

10

1

10 0

Distance to centroid

(up to a sign)

Sigma

0.3

1.0

(b) With regularization r (2) .

The first one promotes pairwise orthogonality while the second one promotes mutual orthogonality.

We use input sequences of length L = 30 in R 6 , where we define the three centroids µ ⋆ 0 = (1 , 0 , 0 , 0 , 0 , 0) , µ ⋆ 1 = (0 , 0 , 0 , 1 , 0 , 0) , µ ⋆ 2 = (0 , 0 , 0 , 0 , 0 , 1) . The model is trained with an online batch sampling strategy (similar to PSGD , changing the data distribution and the regularization term), with a batch size of 256, and a learning rate of 0.01. We take λ = 0 . 6 for σ = 0 . 3 , and λ = 0 . 2 for σ = 1 . Since any parameter could learn any centroid up to a sign, we introduce the following distance to the centroid (up to a sign):

where S 3 is the symmetric group of order 3 and ˆ µ 0 , ˆ µ 1 , ˆ µ 2 are the parameters returned by the algorithm. We present the results in Figure 14. We observe that the regularization r (1) outperforms r (2) , since it explicitly includes all pairwise terms to enforce orthogonality. However, we note that the number of regularization terms grows quadratically with the number of centroids.

Regarding running times, the experiments run in Figure 14 take approximately 15 minutes on a standard laptop.

H Softmax attention layers and clustering

In this section, we assess the abilities of attention-based predictors involving a softmax activation in a clustering context.

H.1 Problem setting

An attention-based learner with softmax activation. We recall that an attention head made of a self-attention layer can be written as follows:

where the softmax of temperature λ > 0 is applied row-wise, and the matrices K,Q,V ∈ R d × p are usually referred to as keys, queries and values. As in Section 2, we assume that the values are taken as identity, meaning that the attention head simply outputs combinations of the initial tokens weighted by attention scores. Furthermore, we assume that the key and query matrices are equal to the same row matrix µ ⊤ ∈ R 1 × d , we obtain

With such an architecture, the ℓ -th output vector is therefore given by

which corresponds to aggregating the X k 's when X k and X ℓ are simultaneously aligned with µ . This head should be a good candidate to estimate a centroid of a mixture model. In the case where the mixture involves two components, one could train two attention heads:

where

Remark H.1 (The attention heads are biased) . As the tokens ( X ℓ ) ℓ are independent, we have that

We note that H soft ,µ ( X ) 1 = softmax λ ( ⟨ X 1 , µ ⟩ v ) X , where the vector v is L -dimensional with components v ℓ = ⟨ X ℓ , µ ⟩ . In the idealized case where σ 2 = 0 , then each token X ℓ is sampled according to a mixture of Dirac masses given by 1 2 δ µ ⋆ 0 + 1 2 δ µ ⋆ 1 . Therefore, if we evaluate H soft λ ,µ ( X ) 1 on µ = µ ⋆ 0 , we observe the following:

Conditionally to X 1 = µ ⋆ 1 , then

This implies that even in a completely misaligned set-up (i.e., µ = µ ⋆ 0 , X 1 = µ ⋆ 1 ), the proposed attention head will return, as the transformation of the first token X 1 , the empirical mean of the sequence tokens. This highlight the bias introduced by such an attention head, which should be handled through the use of centering techniques.

Conditionally to X 1 = µ ⋆ 0 , then

This suggests that in the perfectly aligned case (i.e., µ = µ ⋆ 0 , X 1 = µ ⋆ 0 ), selecting a sufficiently large softmax temperature λ will cause the model to assign negligible weight to the misaligned components -a desirable property.

Debiasing and disentangling heads. To handle the bias introduced by the attention heads, discussed in Remark H.1, we propose to consider centered heads instead, leading to the following modified version of the risk

Considering such a risk is equivalent to using heads where a term proportional to 1 L ∑ L k =1 X k is substracted. This type of head is known as shaped attention (Noci et al., 2023; He and Hofmann, 2024). For instance, initializing ψ = 2 debiases both attention heads independently, without considering their interaction. Using heads with oracle parameters, one would expect that a single head provides all the necessary information, making it sufficient to debias only that head (i.e. ψ = 1 ). In that case, one should obtain:

However, when using non-oracle parameters µ 0 and µ 1 within the debiased heads, the risk function may admit global minima where the heads align with zero, one, or both centroids, which is undesirable for the clustering purpose. Therefore, we must enforce a constraint ensuring that each head aligns with exactly one centroid. To achieve this, we introduce the regularization term:

leading to the following regularized optimization problem

where ρ 0 > 0 .

H.2 Numerical experiments

We run Projected Stochastic Gradient Descent (see Appendix F.1) to learn the centroids µ ⋆ 0 and µ ⋆ 1 as well as the weights ψ and λ .

In this experiment, we use input sequences of length L = 30 of 5-dimensional tokens ( d = 5 ), drawn from a 2-component Gaussian mixture of centroids µ ⋆ 0 = (0 , 0 , 0 , 0 , 1) and µ ⋆ 1 = ( -1 , 0 , 0 , 0 , 0) . The variance of each component is either set to σ = 0 . 3 (low interference) or to σ = 1 (high interference).

The model based on two softmax attention heads parameterized by µ 0 and µ 1 is trained using ( PSGD soft ) with an online batch sampling strategy, with a batch size of 256, a learning rate of γ = 0 . 01 , and running for a total of 3000 iterations. Additionally, we initialize with λ set to 3 and a centering value ψ of 2. Here we use the metric distance to the centroids , given by

and µ ⋆ 0 , µ ⋆ 1 denote the true centroids, respectively, while ˆ µ 0 , ˆ µ 1 are the parameters returned by ( PSGD soft ). We remark that this distance is finer than the one defined in 11, as it does not disregard sign flips. The results are visualized in Figure 15a, we observe that a regularization term substantially improves the accuracy of the recovered solutions. However, as the strength of the regularization increases, it gradually overrides the original objective and impairs the alignment of the head parameters with the true centroids -an effect that becomes more pronounced at higher noise level, an effect also noticed in Section 3.

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Regularization

10

1

10 0

Distance to centroid

Sigma

=0.3

=1

0

2000

4000

6000

8000

10000

Iterations

10

1

10 0

Distance to centroid

Sigma

=0.3

=1

(a) Distance to the centroids after 5000 PSGD iterations vs regularization strength ρ for the minimization of R soft ,ρ 0 , with an initialization on the unit sphere, with data drawn from the non-degenerate case ( P σ ). 10 runs, 95% percentile intervals are plotted (b) Distance to centroids vs PSGD iterations for the minimization of R soft ,ρ 0 , with an initialization on the unit sphere and regularization ρ 0 = 0 . 5 , with data drawn from the non-degenerate case ( P σ ). 10 runs, 95% percentile intervals are plotted.

In Figure 15b, we set the regularization parameter ρ 0 to 0.5, and run PSGD soft for 10 4 iterations. We observe that the model yields accurate solutions under low interference ( σ = 0 . 3 ); however, where

as the interference increases ( σ = 1 ), the ability of the softmax attention heads to align with the underlying centroids is progressively impaired. A similar loss in alignment accuracy is observed as the dimensionality increases.

The experiments in Figure 15b run in a few minutes on a standard laptop, whereas those in Figure 15a may take up to two hours to cover the grid in the regularization hyperparameter.

I Proofs of Section 5

Proof of Proposition 5.1. For c ∈ { 0 , 1 } , one has

Where we have computed the conditional expectations by expanding each term using X i = µ ⋆ Z i + ε i with independent Gaussian errors ε i , discarding odd-moment terms, and applying Isserlis' theorem (Isserlis, 1918).

We remark that choosing λ = L 2 1 1+( d +2) σ 2 +( L -1) ( 1 2 + σ 2 ) we get that the encoding is unbiased.

And

Besides,

Also,

On the other hand,

Recalling the expressions stated at the beginning of the proof of Proposition 5.2 for moments of Gaussian r.v. , we conclude that

And

When L →∞ ,

Choosing the λ = 1 1+2 σ 2 , we have an unbiased encoding with variance

Proof of Proposition 5.2. We have

Furthermore we have the following,

These identities follow by writing

and expanding the expressions in terms of µ ⋆ i and ε i . The expectations then reduce to Gaussian moments, which can be evaluated systematically using Isserlis' theorem (Isserlis, 1918). Since this involves only straightforward but lengthy computations, we omit the details here.

Since no expression depends on µ ⋆ 1 , µ ⋆ 0 , we have

And

When L →∞ , we obtain

And we can choose λ = 1+4 σ 2 +2 dσ 4 4 ( 2 ( σ 2 + 1 2 ) 3 +( d -2) σ 6 ) to get

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our main claim is stated in the abstract, and fully studied in Section 3, particularly Theorem 3.4.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The authors are transparent about the simplified nature of the architectures and mathematical models they study, and consistently highlight these choices as part of the scope and limitations of their work.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theoretical assumptions are explicitly stated. Complete proofs are provided in Appendix B and C, with supporting technical results in Appendix E.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The main algorithm used for our numerical experiments is defined in Section F.1. Detailed descriptions of the experimental setups are provided in Sections A.2, 3.2, Appendix G, and H.2. This information enables others to replicate our experiments without the need for code.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our primary contribution is theoretical, focusing on the analysis of a projected gradient descent algorithm. The experimental results serve to illustrate these theoretical findings. The data used are synthetic mixtures of Gaussian, and the algorithm is a standard projected gradient descent on the unit sphere, both of which are straightforward to implement. The code can be found in the supplementary materials.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Detailed accounts of the experimental configurations are presented in Sections A.2, 3.2, G, and H.2.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: To assess the statistical significance of our experimental results, we conducted 10 independent runs for each experiment. The outcomes are visualized using Seaborn's lineplot function, which displays the mean performance along with error bands representing variability across runs.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All experiments were conducted on a standard laptop. Specific details regarding the computational resources, such as execution times, are provided at the end of each numerical experimental section.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have read and understood the code of ethics; and have committed to conform.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This work is primarily theoretical and employs simplified, synthetic distributions. It does not have direct societal implications beyond contributing to a deeper understanding of certain aspects of attention-based mechanisms.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper is mainly theoretical and does not involve the release of any models or datasets that could pose a risk of misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Our code is self-contained and does not include any source code or binary files from external libraries. Therefore, there are no concerns regarding permissions or licensing. We did cite open-sourced libraries used in our paper, such as JAX.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not introduce or release any new assets. The numerical experiments conducted serve solely to support the theoretical results and are not central to the paper's contribution.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: We did not employ any large language models (LLMs) during any stage of the research process.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
comments: SUMMARY:
This work provides a theoretical analysis of the clustering capability of Transformers brought by the attention mechanism. By taking a two-head attention layer as an example, this work shows that Transformers could capture the clustering structure underlying the input data. Experimental results also support the theoretical conclusions.

STRENGTHS AND WEAKNESSES:
- Strengths:
	1. It is interesting to study the clustering capability of the Transformer or attention mechanism. While the conclusions are not surprising, the detailed theoretical analyses are still of good value.
	2. Beyond theoretical analysis, the authors also conducted numerical experiments to show that the conclusions hold under different initialization regimes, mixture separability levels, and problem dimensionalities.
- Weaknesses:
	1. I suggest that authors provide an illustration to demonstrate the relationship between inputs, cluster centers, and attention layers. This might help the readers to understand the theorem.
	2. The current proofs and derivations are based on the simplified example. Could the derivation extend to more complicated scenarios?
	3. Since the attention mechanism is proven to be effective in clustering, is it possible to evaluate its actual clustering performance?

QUESTIONS:
Please refer to the weaknesses section.

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper shows, learning the weights in a very specific two-headed attention-mechanism recovers the cluster means of a very specific two component Gaussian mixture (up to signs) both theoretically and empirically given specific initialization and/or regularization procedures, when we enforce the transformed instances to be close to the instances in terms of their squared distance. The Gaussian mixture centroids are assumed to lie on the unit sphere and to be orthogonal. The attention-mechanism is based on linear dot-product attention (i.e. there is no softmax applied) with keys and queries projected to only one dimension by a shared weight vector.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n- The investigation of how attention relates to mixture models is interesting.\n- Theoretical results are backed by proofs (which I have not checked) and experiments on synthetic data.\n- The relation of the attention mechanism under consideration and variance reduction is interesting (but not really explored/explained).\n\nWeaknesses:\n- The constraints on the mixture model (i.e. two components, orthogonal centroids with norm = 1) are very harsh, reducing the significance of the results strongly.\n- Due to one-dimensional keys and queries, the specific form of attention (Eq. 2) can be rewritten as a product of a scaling factor (lambda X_l^T mu) that only dependends on the query and a weighted sum of values (here = instances, since there is no projection) where the weighting is only dependent on the keys. As such, the specific simplification of attention results in something that does not involve pairwise dependencies between tokens (an important / essential property of attention, also stated in the introduction), but rather attention heads that are learned means of instances and query weights that weight the attention heads.\n- The above point is not discussed. In my understanding, the attention mechanism is thus simplified to the point where it naturally resembles a mixture model. This would make the theoretical results very much less interesting.\n- The suggested initialization requires knowledge of the solution and is thus not applicable in practice.\n- The regularization looks like it might enforce approximate orthogonality of attention-weights mu to data instances. This might have a strong effects on optimization and further hinders the generalization of the papers result to more general attention models. This is not discussed.\n- That the problem statement (and or proofs) seemingly requires Riemannian gradient descent complicates the analysis further.\n- The propositions are proofed in the paper, but proof-sketches and intuition in the main text would help readability (I have only skimmed over the Appendix).\n- There are some minor typos and miss-nomenclature, (e.g. Z_l = 0 and Z_k = 0 in line 92, mu_2 in line 626, naming projection weights keys and queries in line 73) that do not help since the paper is very math-heavy and concise.\n\nQUESTIONS:\n- Why do you need the constraints on the mixture model (see above) + Riemannian GD?\n- How can you argue that results on the simplified attention in Eq. 3 are relevant to more general (or: proper) attention models (see notes on factorization above)?\n- Can you elaborate on the results in Section 5, i.e. the relation to variance reduction? How much is this reliant on lambda?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-02"}

comments: SUMMARY:
This theoretical paper investigates the clustering capabilities of transformers. Specifically, a two-class Gaussian mixture model is analyzed. First, the degenerate case, where the Gaussians become Dirac functions, is examined, followed by regular Gaussians with σ larger than zero.
The paper presents a theoretical convergence analysis, which is subsequently empirically demonstrated through simulations.
By combining a classical clustering

STRENGTHS AND WEAKNESSES:
Strengths 
1. The paper provides a complete theoretical analysis of transformers as an unsupervised clustering mechanism. This offers new insights into the properties of transformers. 
2. The paper is well-written. Theorems are clearly stated, while their proofs are given in the appendices. 
3. The convergence curves, with and without regularization, are insightful.

Weaknesses
1. Only a simplified scenario is analyzed (two classes, orthogonal centroids).
2. Only a brief discussion on the potential applications of the findings of the paper.
3. Only linear activations are discussed in the main paper. Appendix G discusses the softmax activation function. It would be beneficial to add a brief description to the main article as well.

QUESTIONS:
1. It will be interesting to simulate a problem with more classes. I assume that the theoretical analysis will be very challenging; however, a short simulation study may further strengthen the contribution.
2. Can the authors comment on the relation of the proposed method and the standard expectation-maximization method?
3. How important are the assumptions on orthogonal centroids? Are they mandatory for the convergence of the method, or only for the analysis?
4. Can you determine the number of classes from the attention mechanism?

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 4
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a theoretical analysis of the ability of attention mechanisms to perform unsupervised clustering. The authors investigate whether a simplified Transformer layer can learn the latent structure of data drawn from a mixture model without access to any labels. To do this, they analyze a simplified two-head linear attention layer where the value matrices are fixed to the identity. The learning task is to recover the two centroids of a Gaussian Mixture Model (GMM) from which the input tokens are sampled. \nThe core theoretical contributions are:\n\nA proof that for a simplified Dirac mixture model, projected gradient descent (PGD) on the risk function drives the attention head parameters to align with the true centroids, provided they are initialized on a specific manifold.\n\nAn extension of this result to the more general Gaussian mixture model, again showing convergence to the true centroids under similar initialization constraints.\n\nA proposed regularization scheme to overcome the impractical initialization requirements, which is shown empirically to promote disentanglement between the heads and enable learning from a random initialization on the sphere.\n\nAn analysis of the oracle regime (where head parameters equal the true centroids), showing the attention layer acts as an approximate quantizer and can achieve a lower risk than a singular optimal quantizer due to a variance-reduction effect from aggregating multiple tokens.\n\n\nThe theoretical findings are supported by numerical experiments on synthetic data that validate the convergence dynamics and the effectiveness of the proposed regularizer\n\nSTRENGTHS AND WEAKNESSES:\n***Strengths***\n\n**Novel Theoretical Framework:** This work provides a novel and rigorous proof that a simplified attention mechanism can perform fully unsupervised clustering on a simple Gaussian Mixture Model. This is a significant and complex theoretical step in understanding the unsupervised capabilities of Transformers.\n\n**Practical Regularization Method:** The paper proposes an effective regularization term that empirically overcomes the restrictive initialization assumptions of the main theorems, enabling the model to learn the correct centroids from a random start.\n\n**Clarity:** Numerical experiments, together with cohesive presentation, make the work pleasant to read. \n\n***Weaknesses***\n\n**Impractical Initialization:** The main convergence theorems (3.4 and 4.4) are a central contribution, yet they rely on an impractical initialization on a manifold that depends on the unknown, true cluster centroids. This creates a significant gap between the formal theory and practical application.\n\n**Limiting Simplifying Assumptions:** The analysis relies on strong simplifications that limit its generality, including a highly simplified linear attention model and a two-component, orthogonal GMM. It is unclear how these findings translate to more complex data, even GMM with more components.\n\nQUESTIONS:\nThe choice of 2 heads is connected with the fact that the mixture has two clusters. However, usually architectures are restricted to a reasonably low number of heads. Do you think that a similar approach can be generalized to N clusters, but a smaller number of heads?\n\nMany theoretical papers on Transformers focuse on their ability to perform supervised, in-context learning, for instance, by implementing known algorithms. Does your analysis of unsupervised clustering reveal a fundamentally different capability of attention, or could it be seen as another form of in-context algorithm implementation where the \"algorithm\" is centroid estimation?\n\nLIMITATIONS:\nLimitations are addressed in the work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}