Attention-based clustering
Authors: Rodrigo Maulen-Soto
Pierre Marion
Claire Boyer
Venue: NeurIPS 2025 poster ()
Keywords: Theory of neural networks
Clustering
Transformers
Attention-based models
Mixture models
Optimization
Unsupervised learning
Primary area: theory

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
Rodrigo Maulen-Soto"
Pierre Marion
Sorbonne Université, LPSM Paris, France
EPFL, Institute of Mathematics
Lausanne, Switzerland
Claire Boyer'
Laboratoire de Mathématiques d' d'Orsay Université Paris Saclay
Institut Universitaire de France
Abstract
Transformers have emerged as a powerful neural network architecture capable of tackling a wide range learning tasks. In this work, we provide theoretical anal- ysis of their ability to automatically extract structure from data in an unsupervised setting. particular, we demonstrate their suitability for clustering when the input data is generated from Gaussian mixture model. To this end, we study simplified two-head attention layer and define 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 capture underly- ing distributional structure. We further examine an attention layer with key, query, and value matrices fixed the identity, and show that, even without any trainable parameters, it can perform in-context quantization, revealing the surprising capacity transformer-bant based methods to adapt dynamically to input-specific distributions.
Introduction
Attention- based models (Bahdanau et al. 2013), in particular Transformers /Vaswani et al. 2017). have achieved oftte-of-th-arr performance across 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 (Dosovitskly et al.) 2020 Liu al. 1021
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 al. 2017). They do this by computing pairwise dependencies between tokens based on learned projections, without regard to the tokens' positions in the sequence.
On the theoretical side, 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- -pal-edaideddd with provable models that solve them using attention based mechanisms. Notable recent efforts in this vein include Ahn al. (2023); yon Oswald et al. (2023): Yang et al. (202)): 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
ct al. 2024 Garg et al. 2023 Li et al. 2023 Furuya et al. (2224). The goal of in-context learning is to predict the output corresponding
new query, given prompt consisting of input/output pairs.
*Corresoondence: maulen@lpam pariipsampppriisaaaaae and claire. wwwire. www.rrrrbumieeeaterrrrraaaaaaaae
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Beyond the standard supervised setting, Transformer models are often (pre)trained in practice with sumi-iseee objectives such as masked language modeling (Phoogg 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 (Loyyd 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 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 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 aitention-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 noisy observation of an underlying centroid, and the transformed token corresponds, to some extent, to its denoised estimate. Finally, 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 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-basee methods to adapt on the fly the underlying data distribution, even when no attention parameters are trained.
focus on particular attention layer in which the key, query,
Organization. Section introduces the problem and outlines the proposed approach. In Section 3 we address the general problem with linear attention applied to wwo-component Gaussian mixture model. In Section 4 we discuss the quantization properties of attention- predictor with oracle parameters. In Section 5 we explore an in-context clustering framework and examine the quantization capabilities of 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 set of input tokens to a transformed set of output tokens. With this in mind, we consider an input sequence X E IR Lxd composed of L tokens (X1,... Xy)), each token being a vector of IR., i.e., x= (X...., XL) We assume that the tokens are i.i.d. drawn from simple mixture model: for 1 1se5 l LL
I X 80~2N(m N 02s) + 1 NN dal)).
I
(P,)
with balanced components and where the centroids (μā: ut) E (8d-1)2 (i.e., || = || = 1) are assumed to be orthogonal, i.e., such that (16 m4) = 0. Therefore, for each token, there exists an associated latent variable, denoted by Zz, corresponding to a Bernoulli random variable of parameter 1/2 and encoding its corresponding cluster.
2

---

Attention-based predictors. An attention head made of a
layer can be written as Haofta (X) softmaxx (XQK x XV
where the softmax of temperature X > 0 is applied row-wise, no skip connection is considered and the matrices K, Q, V E Rdxp are usually referred 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 μ E Raxi we obtain
x (x) = softmax (Xup² X X.
With such architecture, the l-th output vector is therefore given by
H (X) Σ softmax, (x7 14 x))x Xs.
which corresponds to aggregating the Xg 's when Xg 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 ( introduces a coupling between tokens, which undoubtedly complicates the mathematical analysis. To address this difficulty, we propose to look at simplified linear attention head, still parameterized by μ Rd, and defined for s l < L, as
2 11 Hi (X)e I. EaxT XR)XK. (
This head linear activation function instead of the traditional softmax found in practical architectures, and has already received interest in several mathematical studies (see ZZhng et al. 2024] von Oswald al.
2023 Han et al. 2023 Katharopoulos et al. 2020). This model is particular case
of Gaussian sequence multi-index models (Cui 2025), which were already used to study attention Arnaboldi et al., 2025 Troiani al. 2025). albeit not clustering context.
models (Cui et al. 2024
Note that when ju is chosen to be μΑЛ then for tokens X, and X₂ whose corresponding latent variables Ze and Z₁ are both equal to 0 (i.e., the samples belong to the same cluster centered at pī). the vectors X, and Xs are likely to be aligned with μ²z In this case, we have (X uy Xx)
X Conversely, if Xg and Xx are associated with different latent variables (e.g., Z₁ = O and Z = 1), then (X, 144 X2 = 0. This behavior suggests that when setting HỘ and if X₂ belongs to the cluster centered at μő. the sum El= (AX7 yqp T Xxx)X,,, effectively aggregates the X₂ 's from the same cluster, whose expected number is L/2, motivating the renormalizing factor of 2/L. Overall, Hlln. (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 μg and μq Rd
Po-
X
Fli (X) + Hlin, (X).
Metric loss. As no label is available, we focus on minimizing the following theoretical loss:
C(T) del
EE E E [ +(
L
where T is an arbitrary iitentionbbased predictor. The distinctive feature of this risk lies in the fact that, if the predictor were able to return, for each token Xx. its associated centroid /zz,- the risk would exactly correspond to quantization error, characteristic of standard clustering task. Note that, due to the independence of the tokens, we have L(T) E| - T(X), 113 so we can confine the following thcoretical analysis on the minimization of the predictive error for the first token only.
3

---

PGD iterates. In this paper, we focus on the training dynamics of Transformer-base 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 non-trivial task, which offers valuable insights into the behavior observed in practice.
For given predictor pl. made of two linear attention heads parameterized respectively by μo and M1, one can reinterpret the objective L function R. of the parameters (mo, R1), defined by
R(vo 11) = L(T
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 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, nq) (Sd-1)2 and step size (S" 1)2 are recursively defined by:
> 0, the PGD iterates (Má. 14)
14 (Is (18) V R(p p²)
+
1/ 7(1 (1(()) T) T Mo R(U0 12)
(PGD)
a 7(4a - (M V R(u - v4)
A+1
|A -I(la (A R( (ap0.-i)) p1)
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 the case of clustering data generated from a Gaussian mixture model, as specified in P where the input tokens are i.i.d. samples
X 2 N( (Ia) + +2 NN
for S L, where (16. pi) E (Sa-1)2 are orthogonal unit vectors. As a first step toward mathematical understanding of Transformer-based layers in clustering settings, we have also considered the degenerate case d² 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
in which the data are drawn from mixture of Dirac masses.
3.1 Theoretical analysis
Preliminary computations.
We start by introducing the following quantities:
Kg (Pos
(1is roo def (1u, (6), ?h del (Ho. pT), & _ (po-
(po- p1),
(7)
and derive a closed-form expression for the risk w.r.t. this reparameterization. Although the full expression is somewhat complex (see Appendix C.T). the following proposition highlights its key structural properties, as being a polynomial in these five variables.
Proposition 3.1. Under the Gaussian mixture model (Pa) consider the aatention-basse predictor plin.po..i composed of two linear heads parameterized by (Ho, (1) E (gd-1y2. Then, there exists function R 1,-1,1]] 1]6 + IR such that R(p,, p1) = R (Ko, f61, Too Th, {) and
4 k4, o. n 2 iim n 1 ko7o817h R RK(K,, K1, No. () &span {ng No:
(8)
Kgp n N², m. BOnDE, R 2
We remark that when Noo 7, and are fixed to 0, most of the monomials vanish, yielding fully explicit formula for the risk (see Lemma C.T in the appendices). Far from being 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 (Ho. p1) lying on specific manifold
M = (() E (s (o) = 0, (1ō pi) = 0, (oo P1) = 0}-
(9)
In terms of notation, it is equivalent to assume that 7jo 0, n/1 0,5 = 0. Therefore on this manifold, we adopt the shorthand notation R (No, K1) d R (No: K1, 0, 0).
Lemma 3.2. Under the Gaussian mixture model P the risk R restricted to M has the form
R (Ko, x1) A(kō xt) + Brn + R2) Crani D,
(10)
for A, B, C, D non-negatives constants, dependent and L, made explicit in Lemma C.1 Proposition 3.3. Consider R2 (Bo,,1), there exists X² (o, L) > 0 such that the points (+1, +1) are global minima of RR (Roo K1).
This result demonstrates that, under suitable condition on the temperature parameter parameler--seciffcallyy when = (a Lhethe points =/1 and + are global minimizers of the risk. The explicit form of (a L) is provided in Proposition C.2 in the appendices. Moreover, it is worth noting that as the variance o² of the Gaussian components tends to zero, (a, L) approaches the value L+1 which coincides with the value previously identified in the degenerate case (Po)- On the
Aj = L+3 other hand, when a is fixed and I. grows large, X (a, I.) tends toward A guide us to properly choose À in our numerical experiments.
1460 6o which will
1+46²
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 Paa) consider the attention- predictor plip composed of two linear heads. Take X EJO, (* L)]. with (a L) defined as in Proposi- Then there exists 5 O such that for any stepsize O < 5, and for generic initialization
tion 3.3
(po. PT) E MM, the iterates (pl p() generated by (PGD) converge to the centroids (up to sign), i.e.
(ph p()
(= +1))
k-000
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 Gra- dient Descent (PSGD; see Appendix F.I). which serves as an empirical counterpart to the (GDD) 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
min +E32
min
Σ |l - -5147 13,
(11)
where S2 is the permutation group of two elements, Hō. μ² denote the true centroids, respectively. while Ao- fú 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 Appendix F22
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 ||ا - lill2 = /2. In this setting, to characterize
uup to relabeling the head parameters, since priori, μo (resp. μn) does not have to be automatically related μ (resp. pi).
5

---

the separation between the two modes of the mixture, one can introduce notion of interference that depends solely on the variance level each mode and which is defined as 1( P(X,, L )
where X. (O,0 o²). Remark that this function increasing with supremum 0.5. This motivates the choice of two contrasting scenarios for the numerical experiments: a low-intefference regime with o 0.3, where I(0.3) 0.01, and high-interference regime with 1, where [(1) 0.24. More implementation details of the following numerical experiments can be found in Appendix F2
Results. When initialization is done the manifold, the training analysis depicted in Figure demonstrates linear convergence of the head parameters toward the centroids (up to a sign) during the first 10 iterations, which is in line with the obtained theoretical results.
Distance to centraid
100
Sigma
to a sign)
0.3
10-1
0=
The error then plateaus at around 10-2 2 in the low-
interference setting (a 0.3), and around 10- in the
-0
high-interference setting (a = 1). This saturation phe- nomenon attributed to the stochasticity introduced
2000
4000 6000 8000 10000 Iterations
by (PSGD) in place of (PGD) in the simulations.
Figure 1: Distance to centroids vs PSGD) iterations for the minimization of R (there- fore without regularization), with an initial- ization on the manifold M. 10 runs, 95% percentile intervals are plotted.
3.3 Generalizations
We consider several generalizations of our theoretical setting, to give insight on the role of our assumptions.
Random initialization on the unit sphere (outside of the manifold). When the initialization 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:
min (² Mo. 684-1
with R" 11) R() + pr M1),
(Pp)
for p > 0 and the regularization term defined by r(po- (1) = (e (1) The role of this term is encourage the orthogonality conditions μ4o- M1, thereby compensating for initializations that may fall outside the manifold M. Numerical results show that the centroids can be recovered with appropriate level of regularization (see Figure (ap. Note that, 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 higher noise levels. In Figure 2b we fix the regularization strength and observe linear convergence towards the centroids over the course of 10 iterations. The error eventually plateaus near 10-3 for o 0.3 and near 10 for 1. This shows that the regularization strategy inspired by the analysis of the simplified Dirac mixture case remains effective the more realistic setting of Gaussian mixtures. In this context well, it enhances the interpretability of the attention parameters by encouraging their alignment with the unknown components of the underlying mixture.
Distance to centroid
100
Distance to centroid
100
Sigma
9 to a sign!
to a sign)
o= 0.3
Sigma
10
10-1
0= 1
a= 0
10-2
10022 -
2000
4000 6000 8000 10000 Iterations
Regularization
(a) Distance to centroids after 5000 iterations vs regu- (b) Distance to centroids vs (PSGD) iterations for the larization strength p for the minimization of RD. minimization of R2, with regularization, p = 0.2.
Figure 2: Performance of (PSGD). when initializing on the unit sphere and minimizing the regular-
ized risk (Pop. 10 runs, 95% percentile intervals are plotted.

---

In Appendix G.1 we present additional experiments in higher-dimensionn 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 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, function only for parameters constrained to lie on the unit sphere. This constraint, however, is not 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 's proper functioning.
are able to compute the risk
More clusters. A natural extension is to consider the case of K centroids and attention heads, We perform 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 the regularization term which now writes
r(po: Pl, 42)
Σ
(m., X1) (Mg,
0SKJSS
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 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 attention- based predictor whose parameters have converged to the true centroids. To guide our analysis, introduce the optimal quantizer 7* as a statistical benchmark within a standard clustering framework. This oracle predictor returns the true centroid of each token, that is, for 1 S S L, T (X), uzz where Ze is encoding the latent cluster of the token Xy. One can immediately note that the risk of the optimal quantizer is given by
= -
-2 12] do
Returning to the attention- based predictor Tlin./6 with oracle parameters, the first key observation that it closely resembles optimal quantizer: its 8-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 (Po). it holds that
p Z = c] = 1 [[L [1L+1) + + 3)02), for (0,1)}
Therefore, choosing X (I+ 2)(+))(()+))) leads unbiased approximate quantization, i.e.,
ET (X)) 1 Z1 c] k
One can next characterize the asymptotic risk and variance of the oracle attention- predictor.
Proposition 4.2. Under the Gaussian mixture model P fix the temperature À 1+66
144 44
Then, the risk of the attention- basedd predictor Thin.pi- with oracle parameters μ and μ satisfies
lim L(Tlii,, = o²(d (d-22). 2).
Moreover, for arbitrary value of X, when L + 00, get
lim Var Mar/plin.s (X) = - 2X"02 202)²
10000
7

---

Strikingly, as the input sequence length L increases, we find that
2 d'
limn 1,0000
L(T*)
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 not entirely fair: the predictors Ti and T do not belong to the same class of func- tions. Indeed, the optimal quantizer T* is only allowed to return two fixed vectors and relies on single input token to predict the associated centroid (albeit with access to the latent label). On the other hand, the image of attention-based en- coder T'in-/6.'i 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 seen as vari- ance 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 T. increases).
0.75
0.50
Compunert:2
0.23
0.50
Note that the gap vanishes in a high-dimensional setting d + oc.
Insut tocers Transfomeed lelers
Another insightful comparison between the predic- tors Pi and T is through their conditional variances Var[TTin.H.. (X), |Z1 c). While the conditional variance of the optimal quantizer 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- 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),
Component
Figure 3: Comparison between inputs and embeddings (d 10, a 0.3, L = 500, À 1+202. Pó C10. ur = -1, ej is the j- th vector of the canonical basis of 12))). PCA fitted on input tokens was used to project both input and transformed tokens to 2D.
0 lim Var(T c) = -2X00 + 202)3 < do² Var[X = c.. L
5 In-context clustering
5.1 Setting
So far, have considered the traditional setting for model-based clustering, with mixture of Gaussian made of two components: each token was assumed to be distributed as follows
Xx 2 N o² daIa) N o
with fixed centroids μD and μ of unit-norm and orthogonal. We have shown that despite the non- convexity of the problem, attention-basee 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 = (X1l... (Xt) E Lxd, the first output of the type of attention-based predictor considered in this paper, parameterized by po and μn, can be rewritten
L
po (x) = Σ
A(X H A(X Tnn X
2 T,
Σ
XXX (HOUT +
8m1
8

---

Therefore, when we consider two linear heads parameterized by row vectors for the queries and keys, and constrain them be orthogonal, the setup can be interpreted as single attention head with query/key matrix of rank 2. Interestingly, we are able to effectively train this rank-2 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; (Xil...,, ......,,,....... E IR_ = 1,..., n, we assume the tokens (Xat) to be i.i.d,, such that the
2-th token is distributed as
So P4i 2"d Vi 2 o
for some random orthogonal centroids Mio and p1 of unit-norm.
If the prior distribution over the centroids is concentrated along preferred directions, say va and Mi" then the predictor T'im..oo-/1 will likely perform well, as the Transformer's 's parameters will tend to align with these directions after training. Conversely, the centroids are distributed in an isotropic way, Tin. will struggle with in-context clustering due to its limited flexibility, only two parameters (l and μ11 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 linear attention heads. Specifically, if we consider attention heads whose parameters unit norm and mutually orthogonal, we obtain the following attention layer: for an input sequence
constrained to be row vectors (Me)a=1,..dd I...d E (IR4), each of
X
(X1,- , .z) I lxd, and for S l S
d
2x d
2x
Teix (
Σ c=1
H Mc (X)e
Σ
[ Mape Xx)Xx
X2
Σ Motol
Xx XX.
L
k=
ld
so finally,
T (X), 2A
x, X
(12)
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, discuss the properties of 7 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 XI,..., XL,, sampled from Gaussian mixture model whose component means (centroids) are randomly drawn on the unit sphere:
HO ~ U((=-1) and #T Mg arbitrarily distributed on gd-1 n (H6)*
X1 XulMo #1 ~ WN(H,, o²Ia) + (N
Associated with each token Xe, we still consider a latent variable Z,, corresponding to Bernuuli random variable of parameter 1/2 and encoding its corresponding cluster, that
Xelpō= pî, Ze N(H
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 Tetx performs an approximate in in-context quantization of the input distribution by aggregating se- quences sampled from Gaussian mixture. Specifically, as formalized below, this simple architecture effectively aligns its output with the correct centroid. Proposition 5.1. For (0, 1}, one has
2A
(L-)) ()+ Me.
E I (X) |4 , Mg Zi cl
(1 (d (d+
L
Choosing À = (4 11) yields an unbiased embedding. i.e..
E [ c (2)1lk7 nō. Zi = c] p..

---

In what follows, we characterize the risk and variance of the attention based embedding T defined in (I2), when the input sequence contains an infinite number of tokens.
Proposition 5.2. In the asymptotic regime where L - oc, one has that
lim 1-+00
2X(1 40 2do*) + 412
2
2 0²
1 +
+ + 2)08)
2)
Choosing the temperature X = +2102+2do
gives
1 ((2(02+2) 2 +(d-2)o0)
20²
lim C(TC(X) o² 2) 1 6o + 12o + 6 S o 2). 1.
Moreover, the conditional variance of the embedding satisfies when L + 00
lim Var
(x) Mû Z1 = =c = 22 + 40² + 2do"))
1-000
Choosing
I we obtain unbiased embedding with conditional variance of 4a² 2do
20
(1+ 2or2)
As in the fixed centroids setting (Proposition
4.2 we retrieve that for suitable choice of X, the loss
satisfies
[(Tetx) ((T*) = -11 (1-2 2 d
+ 20²
+ 12o + 4 I (1-22
2
lim L-+00
1 1+660 6
This result is all the more surprising given that emerges in a more complex setting, yet with sim- pler mechanism. Unlike the attention-based pre- dictor 7 No., which benefits from access the true centroids and trainable architecture, the in- context encoder T achieves smaller asymptotic risk without any learned parameters. The ability of 7etx to extract meaningful representations illustrates the relevance of the key-query-valle value structure in the attention mechanism, even stripped of any learnable parameter and of the softmax nonlinearity.
Tranfforedd Ar.ii
Sic Camprrent
PCA Companont
It is worth
noting that for 2 2, taking À T+2
Figure 4: Comparison between inputs and embeddings (d = 10, a 0.3, L 500,
T4267. n Eoo μ - PCA was fitted on the input tokens and was used to project both input and transformed tokens to 2D.
Jim Var [Texx (x) Ini Z1 =c 1-000
= 20 + 40 2do
(1 + 20
Soad= Var [ 71 c]
Therefore, in the regime of infinite input sequences (L oc) and for the appropriate value of X, the embedding becomes unbiased and exhibits a variance reduction effect compared to the input data points, as illustrated in Figure and previously discussed in Section
6 Conclusion
This work offers a mathematically grounded, principled perspective on the unsupervised learning capabilities of attention mechanisms within mixture models. By combining 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 practice, which may further challenge the theoretical analysis.
10

---

Acknowledgments
The authors would like to thank Gérard Biau and Lenka Zdeborová for insightful discussions. This work was supported by the Projet Emergence(s) of the City of Paris. P.M. is supported by 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 https://gi thub. wtt.bb/www/rrrttggaaaaaaatattttt basen-baadd--ee leeeed-custering
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. In 7hirty-seventhh Conference on Neural Information Processing Systems, 2023. URL http://pppaaeeeeee.ttt/toorllla
LziniAXEI9
Luca Arnaboldi, Bruno Loureiro, Ludovic Stephan, Florent Krzakala, and Lenka Zdeborova. Asymp- totics of sgd in sequence-single index models and single-layer attention networks. arXiv preprint (rXiv:2506.02633 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 Reprexentations, 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-494 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-Milnee 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. 2rX11:2303..2271, 2023.
Hugo Cui. High-dimensional learning of narrow neural networks. Journal of Statistical Mechanies: Theory and Experiment, 2025()):02340,, 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-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- Pre-training of deep bidirectional transformers for language understanding. wrXiv:/810.04800 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. httpsttps: //arxiv...v////////////a//aaaavl www./osg////o/ooooooot..oo...set abs/2408. 24308.013677777

---

Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes, 2023. JtLppttts//arxii.oorr wttp:/iii..rg/ab////000/ 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////aaxx org/ 2308. 00442
Bobby He and Thomas Hofmann. Simplifying transformer blocks. In The Twelfth International Conference on Learning Representations, 2024. URL httpp::s///p/eeeeeee// wwtp://eee/oeeeee REDok9eS3s
Yihan He, Hong- Chen, Yuan Cao, Jianqing Fan, and Han Liu. Transformers versus the em algorithm in multi-class clustering. arXin:2502.006077v,, 2025.
L. Isserlis. On formula for the product-moment coefficient of any order of normal frequency distribution. Biometrika, ((1):134-139) 1918.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention, 2020. https://://ar org/ abs/2006. 16236
Kenneth Lange. Optimization, volume 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://axxxt //wp:i//aaiivv//// org/
abs abs/2301.000 ////22211.00000
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.
Staart 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 iiention-basedd neural machine translation. Proceedings of the 2015 Conference on Empirical Methods 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 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. WRt.ttts://// httpsstt//x/xxi... wrw.s/ii.....voooooooo. htse/bb//200 2e10. 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 limit. In Teirth-seventh Conference Neural Information Processing Systems, 2023. URL
httpps://openreviex www.iiiwwweeeeeeeeee mett/forumidd-PofPSSS jS9JRX
Mary Phuong and Marcus Hutter. Formal algorithms for transformers. arXiv:2207.09233. 2022.
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone self-attention in vision models. In 33rd Conference 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.
12

---

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 .rg/ abs/ aws/ab12...000 221.222222777
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 htt.hhttss///c ptstps://openreviess net/forum - AooouuniddAuAAaaa AuAj4vRPky
Ruiqi Zhang, Spencer Frei, and Peter 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 Tlin. in the context of clustering, assuming the data are drawn from the degenerate mixture model, where for
1 S l $ L.
Xe ~ 2 Snu + 2 5
(Po)
(the orthogonal centroids μá and μ4 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
Kg def (16: Ho): K1 del (1).
I7u (M1,16),
def
Th def (po. MT)-
(13)
The scalar products Ko and 11 measure the alignment of the parameters fn and μr with the true centroids f and μL, respectively, while the scalar products no and 171 capture their orthogonality with the inverted centroids μ and μğ. The theoretical risk w.r.t. Ko, R1, 7o and m reads as follows. The proof of this result and the following are given in Appendix [
Proposition A.1. Under the degenerate mixture model (Po). considering the attention -raeed based predictor Thin composed of two linear heads parameterized by pn and μn, the theoretical risk R can be re-expressed as a function R R.d. >> RR such that R(Ho: p1) R (no: Ki, Too 77), where
L + 1
L L+3
R- (Ro, N1, no. m)) = X
(ND
+ K2 nō
X2
(IK + mo + |n + 1)1
L
2L
L -
(NOm + + x2 L
(14)
In addition, if (Ho, p1) 12 are prescribed the unit sphere, then dom ( 1-1,
Remark A.2. After direct computation, note that the critical points of the risk R. correspond to those of its reparameterized version, RS, i.e.,
(wo: M1) E crit(R) < (No, K1, Voo m1 E crit(R (R.
Proposition A.3 (Characterization of global minima). Consider RS IR. + I. defined as in Proposition A./ with \j
L+1 then a point (No, Ki,, no, 7)) belongs to argmin(R*) if and only if L.33
KO + n =
2
N² i =1, and Ro71 + K1700 0.
(15)
13

---

L+3 simplifies the system by setting the first two conditions equal to I. 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.
While the characterization can be made for any value of X, choosing =5= L+1
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 observe that the points (Kog Ki, fjo '7) that saturate the first two equations (i.e., satisfy x² KG and n/ = of = 0) correspond to global minimizers of the risk that also recover the centroids (up to 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
we
M {(Ho, 11) (2= (pi, 0, (Pú= 11) -0.
(16)
Theorem A.4. Under the Dirac mixture model Po). consider the attention- predictor T pl composed of two linear heads. Take X EJ0, 1. Then there exists 7 > 0 such that for any stepsize
0 < < 9, and for a generic initialization (n6= 1D) E MM, the sequence of iterates (16. pf), generated by PGD). converges the centroids (up a sign), i.e.
(mb p
(=1ō: +17)-
k-+00
Proof. This a direct consequence of Theorem 3.4 with o = 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 simplified, already highlights the representational role of key and query matrices in attention-bant -based learning, and serves as a foundation for addressing the more general case Gaussian mixtures.
Note that the convergence is up to a sign, consequence of the symmetry inherent in Hlin, 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 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 (p8- 14) E M for which (PGD) fails to recover the centroids is of Lebesgue measure zero with respect to M.
Initialization on the manifold M relies 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 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
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 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 specific form of regularization:
r(po, m1) de Bl X1) (11, 21)
(17)
14

---

Distance to centroid
1
Distance to centroid
sign)
(up to a sign)
10
10
200-1
66-10-1
-
210-16
4x
0
2000
4000 6000
8000 10000
0
2000
4000 6000 Iterations
8000 10000
Iterations
(a) Initialization on the manifold M
(b) Initialization on sphere (s"-1)2
Figure 5: Distance to centroids (PSG)) iterations for the minimization of R, with data drawn
from the degenerate case P 10 runs, 95% percentile intervals are plotted.
Therefore, we train the attention-based predictor H Mon.u now by minimizing the regularized risk
min (uo, m1) Po
with
RD p1) R(uoo p1) + pr(po, p1),
where > denotes the strength of the regularization, can be rigorously shown that as approaches 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.
Distance to centroid
Distance to centroid
10"
10
3 to si sign)
lup to a sign)
0010-
10
10-2
10-10
10-3
10
10-4
0.0
0.1
0.2
0.3
2000
4000 Iterations
6000
8000 10000
Regularization
(a) Distance centroids regularization strength
(b) Distance centroids vs iterations (p 0.1).
Figure 6: Convergence analysis of the (PSGD) iterates for the minimization of the regularized risk
RP, under the degenerate mixture case Po). 10 runs, 95% percentile intervals are plotted.
In Figure 6a| we observe that a relatively small regularization parameter (of the order of 10-1) sufficient to achieve centroid alignment, with numerical error below 10-14, 14, In Figure 6b we fix the regularization parameter and observe that over the course of 10 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 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 discussion on the effects of 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, introduce the notation (#) der AX; ah X for 1 5 k
L. which allows us to write
2
R(pa, m4) =
X1
(ex(wo) + (c((n)))
15

---

In what follows, we give expression of the risk of an attention- based predictor. in the case where the data is distributed according the Dirac mixture model (Fo) Then, the risk of an attention-based predictor pll.. can be written
L
2
R(yo, p1)
B
X1
(ex())) (x((n))
k=1
=B
[|
(xx:(ex (ex( + (x((m)))x))
L
(ex(w)) (mi)))) |2)
+
E
T2
L
Σ
13 (X (ex(po) + +xx(m))Xaa)
Blll(e((uo) (x(
8
+
Σ ((() + (e( + eg (mr () (Xx,
1
ISKCJSL
[l((e (HO) + e (14)))
E [ll(e) (11))))
L
L
(11b)
(((((())) + (e(()))((e((n)) (x), Xx)]
L
(1(11))
L
E [[X11 (x(ma))
E
[Il(e((1o) + Xxx
L
L
(1)
det
(m(
+ T2
Σ
[[(e)((o)) + =x(mi))(es(HH) + (edn))))))eee(ee)eeee))
(x)(m)))(Xx, X,)]
(1II)
We compute (I) by conditioning on Z1, Zt. [((e((p)) + (x())))( Xt)]
E [e<x (No) + (x(())) Xx) Zel]
>AE(((HZ, mo) (p2z no) + (1) H2.)l
~ N₂ + nd r/
This leads (I) = X (NO i n 1/t)
Similarly for (IT), conditioning on Z1, Zk.
E [[((e((p))) + (xx1)))xxx11 E (a())) | Zæ]]
+
= >32Ill( PO) (42,- 4o) + (4z H1) (Hz, 1z 1131
xa INB + nl + iN?
(No07) + x170)2
2
Which gives (II) = X² 12 (1kg + 2 + [n] + 1 + 2 2(1-1)
(1077) + k2
1
Finally, to compute (IIT), note that (((((p))) (#x()) Xx,,
16

---

- ((( + (x(Ha))(e)(H)) + (x, 2x. 2,1]
= mo) uo) no) + 1 (uz, (2))) Ho) + (1), M2))
(122 1z2,1l
= x2 INO + nol + [n]] + m
X² (NOV) + +1177)
leading to (ITT) = 22 (1-11
[ + n0812 + [M] + n + 2(n0m: + 9.
21
similar fashion, we get that
(To) (Kg x4 + 10 nn);
(11p) [N + 2 [w] 2
(ITT)) = [<] + 16 [w] + 1
212
Putting everything together, we obtain that the risk can be written in terms of Ko, K1, 70, m. i...
R(uo, ui) = R (KO, K1, 7/0, 7 where:
L L-1
R _ 1
(ng + n² + 11 + n/))
L
L
2(L
L-1
(L- 1-(1)(L 2)
X
(IK3 + - + +M + )
T22
1
212
1)
2)
2(L 1)
+
(Ni!!) + K41 k11
L
L
L + 1
L +
1 X
(ND +x + no + ni)
X2
1/1 + [n] + 1))
L
2L
L -
+ 2 (Nom + x170)
B.2 Proof of Proposition A.3 (critical points of the risk in the degenerate case)
Proposition B.1. Consider R² R + R defined as in Proposition A./ with = Aū = L.43*
L+1 then
we characterize its critical points by
1. The point 0,,0,0,0) is local maximum.
2. The points (Ko, No, 0), where xă + no 1, and (0, k1, 0, where ki + +n 1, are strict saddle points.
3. The points (No, K1, kn) and (No, Ki, -ki, - -ko), where No + n2 2[L+1) are strict saddle points.
I.t3
4. (Ko, K1, no. Th) belongs to argmin(RS)) if and only if:
2 + n
1,
K + r
1,
(18)
K0771 + K1770
0.
Proof. Let us define Ço = (No. n1)) Çi = (m 61), then there exists a function RCC IR:2 R.2 - R,
such that R (No, N1. Mo. m1) < (1), in fact, let define A (L+1)²
B (L+1)²
C
L(L+3)
2L(L+3)+
(L-1) then with the value of À defined in the proposition, we obtain ML-3)
<<= (1) = A(I[G 2 + ||G 12) +
+ + C(Go.G12
To analyze its critical points, we take the partial derivatives.
V R ( -2Ao + 4Bllcoll 2C((a, Ci)C1.
17

---

R (1) 2A4 + 4B + (C Si)Co-
And also, we compute its Hessian, define
V 2 R (60 (1) = -211 + 4B(26000 + 12) + 20G51
V coisi R (Go (1) 2C(GGG + (1T2);
Vi.5o R < & (1) = 20(G1 + +G. (222)
V S1.5,
RC (Lo, S1) 2AA2 + +4B(2ci + 12) 2CCoSO
Then the Hessian will be defined by
V2 So.Co (So (. R<
V Vioce R R= S1) Va R R (Co, Si)
V²RZ< (1o
(19)
((1,,
To find the critical points, solve the following system of equations:
-2Açe + 4Blk (o + 2C Gi)G 0,
(20)
-2C + +BI| [ 2C Co 0.
(0,0) local maximum. We see that trivial solution this system (Co, (i) (0, 0), and replacing into the Hessian matrix, we see directly that this point is a local maximum.
(0, Ci), (So, 0) are strict saddle points. solve
We check the case when Co 0, S1 O, then we need to
-AA + 4B||</1122 =0.
Since Ç1 # 0, this forces -2A + 4B [<a
0. Replacing (0, (1) into the Hessian matrix gives
(OR<<(0,G11) (-2A12 + 20G SBCISI
And eig( (VR<<< (T) = eig( -212 + 20GGT) (ig((8G(T)) where eig is the set of eigen- values of matrix. We have that
eig(-2A12 + - (-
eig(8B(icT) = {0, 8B][(11))))
where we have used that 8B 4A. We also note that C' - A < O, 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 (Co, 0) for 2 1.
Non-trivial critical points. We will first show that the critical points that are of the form (Co, C1)
for Co # 0, * 0 necessarily satisfy || = ||S1l| of O, multiplying the first equation of (20) by Co and the second by Sı, then subtracting both resulting expressions we obtain 4B([| [C
2AA[Isa || 12), and then
([Coll IK< [ + 1Kc.
||c.113) = 0,
thus either || = || and we get the first claim, |/o/12 + ||<a 1|? 1, in this second case we divide in two subcases:
Let us assume that (Co, (1) 0, then multiplying the first equation of 20 by Co and the
second equation by (r, we get that |oll || 1, which contradiction since we
are in the case where ||Soll + || 1.
18

---

Therefore (So. (1) # O, we multiply the first equation of (00) by &i and second by Co. after dividing by (So. <T> we get that
A + All + C(Co. Ci) 0
- A + Alll + C(So: (1) O.
Substracting both equations we get that |Coll = ||iil
So we get that necessarily l|sall = ||</l = r > 0. Then the equation 20
becomes
A(r² - 1)ço C(Co. S1) O, A(r² 1)<i + C(oo, <i>Co 0.
(21)
(Co, =u)) are strict saddle points. then
the case where Ço +(1, by (1 we get that when So = C1.
A 1)5o + Croso 0,
and ||olll = ro. with r
L+3 2(1+1))
Besides when Sa = - then
44C
A(r 1)Co + Cri 0,
and 1coll ri, with r = 4-C
A
=
L+3 4
Replacing this point on the Hessian matrix
RRR (Go, +GO,
V²R=<(GO. Co) = (A((r) 1)I2 + 2(2A + C)CoSo
2C + CoCo
2A(r3 1)/2 + 2(2A + C)OOT)
2C(ri + (oCT)
V²R -Co) = (2A(1)
-020(1 - -1)12 + 2 Caso -2C(r9 + CoSo
-20(f + (660))
2A(r7 1) + 2(2A + CCKoo
We can write the Hessians in block form with 2 x 2 diagonal blocks
Mo 2A(12 1)/2 + 2(2A +
and dia-diagonal blocks
Mi = 2C + SoCo
Considering vectors of the form (z, -x), 123 this reduces the eigenvalue problem the 2
matrices Mo + M₁. Substituting r A/(AA + C*) or or AA(A C*) shows that each matrix has one positive eigenvalue along Ço (equal to 4.4) and one negative eigenvalue orthogonal Ço (equal to 4AC AiG), ) so both Hessians are indefinite.
Characterization global minima. [o * +C1 and |Kulll ||<1|| = 0, then both vectors are
linearly independent, thus the first equation of (2I) is only possible when / = 1 and (Ço. G 0, in which case we have analyze the points (So= (1) such that (So. 0 and || l[s 11 replace these points on the Hessian matrix and this gives
we
(4ACoSO + 20CGTT 20G1
2C
)
VRR <RR<<
(Co. G1)
4AGiCT 2C
direct computation of the eigenvalues with eigenvectors (G1, Co) and (S1. gives us that all the eigenvalues are positive in this case, since R< is coercive, these points in fact global minima.
Remark B.2. Note that the specific characterization of the global minima of R was valid only for x A/ L+1 A E|0 1-1 the global minima lie outside the domain []1, 122 As result, due the structure of the update rule (PGD), the extreme points (11,+1) of [-1, become fixed points of the algorithm and serve global minimizers of R3.
L+3" However, when restricting the analysis to the manifold M and considering
19

---

B.3 Discussion on regularization
In order solve the clustering problem in the degenerate case, we train the attention- basodd predictor H Ho- now by minimizing the regularized risk
min Payai ES
R (HO, 11)
with
("(Ho: del R(Ho: p1) + pr(po: p1),
(pp
where r(Ho: P1) = El (1,, and p > 0 denotes the strength of the regularization.
It is direct to check that there exists r5 E+ R, such r μ1) r (Ko, K1, 7o. 7/) according to the notation defined in (13). and r (No: K1, 700 mh) = 1 (nov + xZnl). We define the following optimization problem
min AR C[-1,1)
R (Ko. K1, No, mi + pr (No. K1, Too m)),
(P/
where R is defined in Proposition A.1 Since R and x are coercive, we apply Attouch (9996
Theorem 2.1) to conclude that if up E
1, 1] a solution of (p< then every limit point up-
when P 0, satisfies that:
r (+) S (r)) for every argmin RS, E argmin R
Due to the geometry of r and the characterization of argmin R we got in Proposition 15 we obtain that if û (Ko, K1, no, fi), then
=1 = 0, or hm 1, = 1, R, O.
(22)
Then the optimal solution for the regularized problem when -> 0 achieves a saturation effect, corresponding global minimizers that recover the centroids. However, due to the non-convex nature of the problem, is not guaranteed 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 p -> 0. We know from Proposition B. 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 (Gaussian mixture model)
C.1 Proof of Proposition 3.1 (expression of the risk in the non-degenerate case).
Recall the notation (4) del XXX un
Xx, for S S L, which allows us to write
2 t
R(po, m1) E
X
Σ (en(uo) + (x)))
L
Under the Gaussian mixture model, we are going to show that the risk
R(Ho, p1) = E [[|X] - (H + H (x)+1)))
admits closed-form representation in terms of elementary functions. It holds that
2
R(Ho, 11) E
X1
(ex(p)) xx(
=E [|
E [[X (ex(mo) + (x()))
20

---

E
(ex(p)) + (x(I)))
1
(1 do?)
R (X1, (ex(mo) + =xx(1))xxa)
L
k=1
4
+
1
8
+
T?
ISkejSL.
4
4 L
(1 do²) [[((e + (e)
E [ll (No) (11))
def (11a)
([(e + (e()))(ex()))) + (xx(
1.2
(111)))
4
(X1, (en(mo) +x((m)))xe)
[I((e((o) + (x(M)))
L
(T)
-(1
Σ
-[(ex(m))= - =x(())(()(()))) +
+e)(m)) +(x)))(()))((( X,)
L
L<<jjL
(um
= -1+ - (1o) + (11) + (111)) (1) + (11)+ (111).
We now proceed to compute each of the six terms. To compute (To), we can use Lemma E.2
since
E(((X1,Ho)² + (X1, (1) >>) ||
= 2 (B[(X,, (o) X / 21 = 0] E[((XI, ||x 11 = 1])
1
+ (II((1,, m1 ||
| 11
= -| 11 -1)
2
[m)
n + + + (1 + o ()) (d 2)) 12 [|mm|12)
Then, (Jo) 2A [(n] no ot + n1) (d 4)) 203 40) + o² (d 2)) [1)]]
To compute (To),, by defining Po(Mo, /1, k²) as in Lemma E.3 ((( + (3
we get
BICK 2(X1, 2
X1, )1) + (X1, "||x, 11
EI(( + + 24) (X1, (X1, HI) || 1/ 1] 2
= 5 (poluo- No. m5) +
16) + Po(111 (1, 17))
2 (po (uo, no, ui) + 2po M1, u1) + po M1, v())
Then,
412
(11)) 2 B[((X,, 20)2 22 13
21

---

2X²
(poluo, po, 16 + Mi: + polun, M1; (G))
2X²
+ T,2
No. 14) + 2po (po, M1, (()) + po po MI (1))))
To compute (1TT)), by defining p1 (No, M1, 12, 122z) as in Lemma we get
E +
(xx,
no) + (X1, (i))
Σ
(2)
(E
where
T1(21, ([(((1,)0) + (x),, No) (X2, + (1,,1)((22111 Hz))
(X1, X2) 21 = 21, 7/2 = 22).
And then
Σ
T2()1,22)
4
(E
(pI (Ho: no: po. 15) + pr (pos M1, 15) pı (11: Po: Po: 16) + pı M1,
M1, Ho. mō))
(Mo, Mo, 1i, (²) + pr (No, p1, (i, + (11, Po, Mi, p5) pa (11, P1, Mi: H6)) 4
+ (pi (Ho. No, H6: pi) + (Po: Pr, Pos w1) + Pi (Pi, Po, P6- pi) + pa (1) M1, Mo: mf))
(po, No, pi: pT) + pı (No, P1, (1,4i, + pi(pi, po, pT. pr (11 Pz, Mi, vi)).
Consequently,
(1I1)) = 8x2 -
E[((X1,, (Xx,
I?
= 21 (L 11)
Σ
(2)
I?
(00112
To compute (I), we Lemma E.5 obtain:
E[((X,, po) (X2, Ho) + (XX
(o2, -, =, + [B((((X)) (X (o)
+
= 0,2,
+ [E[((,,, (4)) (X,, (4) XX = 1, -,,2=
+ [E(((X,,10)) (X2
(X1, (X, X2) 1, 1, 1
+
40
(K0
10
vi) + 011mol1 + || 11
4
Thus. = x [(N3
+
+
1) 4 + 40 (11 + || [13)]
To compute (IT), by defining P2 p1, Mzz mzzz) as in Lemma E,6 we obtain:
E ( mo) (22, no)
m1)) X₂
= Σ T2(=1,22),
(01.)))
22

---

where
x2) Bl((XI, + (M 21, Z2 22.
And then
Σ T2(21,
(#1.22)E(0.13
4 (p2(po Pus No: mā) + 2p2 P1, (ā) + P2(11, P1, Po: 15))
(pz(vo- po: Po- vi) 2p2(No: pr, P0- vi) p2 (1, 16 (4i)
(p2(vo, po, pi, 16) 2p2(No- ni ri, pa (M,, |n, /i, p6))
+ 4
(p2(Ho: po. pi, (T) + 2p2 P1, (T, + (11 (11, /1, M4, (T).
So obtain,
4X2 -
L2 (II)
[[[[ [0(( Ho) (x)
(I)
Σ
T2(21.22).
L
(0).)))
Finally, to compute (III), by defining Ps(Mo, H1, P2, HZzz 1zzz as Lemma E.7 we get
E[((X, no) (X2, Mo) + (X1, (x), (Xg, Mo) (X,, Xy)l
Σ
22,
where
T3 22, 23) = (o)( No) + (X1, (1),,
((X1, (X3, m0) + (X1, (x)( m11)
(X2, Xa) 21 = 21, Lo = 22, Us = 23
Observing that
Σ
T2(21,22,34
Σ
Pa(Has Hos Hes Ha: (c)),
(0..))
E
get
(III) 8X2(L. 1)(L -)
Σ
Ta(21, 22, 23)
212
((0.11
X2(L 1)(L 2))
Σ
Pa(Mas ub, ue: Hà: uC).
212
(e.b..d))E(0...p)
Finally, putting everything together, recalling the notation introduced in (7). and inspecting the formulas given by Lemmas from E.2 to E.7 allows conclude.
C.2 Proof of Lemma 3.2
(expression of the risk on the manifold, non-degenerate case)
We provide in Lemma C.I more precise version of Lemma 3.2 with explicit constants.
23

---

Lemma C.1. Define (n)) = 1 + no² and c2(n) 1 + od + n), then the risk R (Ko, R1)
restricted to M has the form
R (KO, A(k0 + x1) + B(g + xi) CRini D,
where
2X
L (2 + 2A
01 (5) + 12( c2(4) + X2 - -)(L 2) CI
XL(L-
A
L2
12
212
2X
161a02
(2(6) + 8X 1)
A(I -
B =
L (2( +
(1
(1(4)
L
12
7
2 1)
X (L 1)([. -2) 2)
(2(3)
(6).
12
L
4x (1
(I-
C
L
SλØ²
32
64X20'(L -1) 1)
D = (d)
022 (2) +
(2(
L
L
I
8AO SAo 1
8X
8 1- -
+
(2) +
L
L2
Proof of Lemmn C.7 Using the decomposition obtained in the proof of Proposition 3.1 after simple
algebraic manipulation we get that on this manifold:
(IIn) +|000:/0000000000444444400000010
(IH1) = 2Xa -,,-0 320
(I) +1)+
(II) x² ا (Kg x+) 4)) 40 xz) (d 80
d2(d+2)))
(III) X2 x2(1-1X4-2 (ng) xx) + 40 + 20 +x) 66?) + 160
We conclude by noting that the risk restricted to this manifold is
Rs (No: K1) = (1 + do²) (To) (I1) (III)) - (I) + (IT) + (III),
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
C3(o, L) 160"c2(6) + 802 1)c, (6) + 4L + OZ - [) ((LL-2)ci1 2)c1(6) + 4c2(8) + 4(L - 1)ci + - - + (L (4)) 1))
and consider (No: K1) with the following X:
() 2Le2 (4) L(L - (4)
C3(o,L) L)
Then the points (+1, 11) are global minimum Rf (No, ki).
Remark C.3. In the case where O, we get X1 (0, L) 2a 143.
24

---

Proof. Imposing first order conditions on R R (Fo. k1) 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,
if and only if 4.A + 2B + 2C = 0. The function À + 2A(A) + B()) C()) is quadratic which is negative for v (a, L.), and vanishes at X X (a, 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 (Ko, Ki,70, 71, E) such that 171 = = = 0, we have ONy R = On R 0.R5 0.
Proof. According to Proposition 3.1 we can directly obtain that
R (No: K1, no. m1 R (No. Ki, -770- 771, -E).
Taking the partial derivative in va, we get
pn R < (K0 N1, 70, 71.
Om R (No. K1, -100 7/1, -E).
At a point such that yo = 171 = < = 0, this gives 0, R- (150, 11, 0, 0, 0) -gp (~ K1, 0,0,0,0.. 0). therefore On R- (Ko, K1, 0, 0, 0) = 0, the proof for 0, R.S R. is analogous.
11
Lemma C.S. The manifold M invariant under (ph+). .11 (+1) M.
PGD dynamics, this is if (nă: pt) E M, then
Proof. We apply the chain rule and Lemma C.4 to get:
V R = Omo R S 10 + 0n R R ii + 0. R5 MN 3no R 16:
(23)
Pp R = 0x R R 1 + 8Tp R5 ND + 0g no ON R 14.
(24)
We then follow the same ideas as in Marion et (2024 Lemma 4), where our Lemma C.4 takes the role of Marion et al. (2024) Lemma 14).
More concretely, let us consider co |R 7(4 C || (Ia (x()) R(mh, plll then recalling (PGD) updates, we have that if (u6. nf) E M,, then
( V noo R( u4) ll,, and
(wi) u/ - (mi)) (Ta - (z()) ) R (mo x²) 0,
(ni) 1 A+1
(HO) ur - 7(6)) T(a - wi((tt T 10x, R R- ntt xt) O, (m6) 11 k+1
(Ta-
And
(u)+2)= MD
(ut) un
COCI
100m R- (n 15 (TT) - (+())) T)tt) n/
7(0mo RR(()) xf))((Ta (18(1 )H6) un
coc
12k0ka R (Ntj (0,, R (M (+a) (bt( ((G)) (La (nt) T 0.
CDc1
where the last term is zero since
(Laa (la--r(mb)) (((H)) ()6))) T(a (A(H)) = O.
Then (1b+1 k+1 .H 1 E M.
25

---

Lemma C.6. When initialized on the manifold M, the iterations generated by PGD can be reformulated as follows:
(not1,ntt)) ((()::))))
(25)
where :: --1,1, 1|1, is given by
Ko 7m0 R R (K50 No ()(1 xN)
N1 - 7(0,, R (1o. x1))(L nG))
(kKo.KK)= =
+1++7(00m)) R- (No,, Ki ))²(1 ki)
v + (On,, R (No. Ki (1])a(1 xi)
and R (No, k1) def R (No, R1, 0,0,0,000 0) as in Lemma C.1
Proof. By definition of the iterates and 23 24 we have
k+1 = (10 + 1 110 = 10 ~80o < (N5: ((6))))) (Ta - 1 (1())) m6))
Kg
v1 + 72 (2 R (1o0, 161 ))2 - 1/
n/ mmo R < (No.. 1z x1)(1= (1 (x
V1 200 R (No kz) (xz))))
x 70n, R < (N5O K51 (MT)) (Ia (nt)))ntt
.+1
1 111= (1(+1)
K1
v1 (0,x R (MONO, N1) (2))) (x )
N 7m, R (NO: K3) x1)((1 (xf)2)
V1 + 2k, , R (No, x=) x1)(1 (xr))2
In the following propositions, we consider R
1|, + R. defined as in Lemma e0, X (o, L)), where (o L) is defined in Proposition C.2
C.1 with
Proposition C.7. Let (1g) p M, consider 25 with initial conditions xa (pg, vō) xf
(4). Then there exists 7 > 0 such that for every 0 < < 5, the risk R is decreasing along the
iterates of (25). Besides, the distance between successive iterates tends to zero, and, if (Kā KQ) is accumulation point of the sequence of iterates (155, set )EE, then
(1 - (KT)²)0, R (NO: x3)) = 0, (1 - (xT)))0,,, R (Ng, x2]) = 0.
(26)
Proof. The proof is identical to that of Marion et al. (2024 Proposition 8) and is therefore omitted.
Proposition C.8. The points (Ko, K1) E [-1, satisfying (26 belong to the set
6 [(11, 11), (0, =1), (=1, 0), (0,0)}
Proof. We recall that by Lemma C.1 the risk R restricted to the manifold M, has the following form
R (= k1) = A(x + xT) + B( + x²t + Crg D.
Then
On R R = 4Ang + +Bro + 2C'
0z, R (Ko, K1) = 4Ant + +BB + 2C
And we can rewrite equations 26 as
Ko
12AN] + B B Ckil = 0,
x1 - [2Am? + B CK = 0.
Since each equation a product of 3 terms, the general solution to this system occurs when 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 + B + Cx² = 0, this implies that:
26

---

If K1 0, then Kg
B
2A'
If n then Kō
R+C. 2A
If 2Ax + Ckj- 0, then 2Axi + Cki = 2Akg + Cna thus -2A- Cx - -x 0.
By inspection, C' < 2A, hence we get knj = KÎ and Ko
2A+C
We note the following relation
B +
B
B
2A
2A
2A
Further remark by inspecting the proof of Proposition
C.2 that for À EJO, L)), have
B
+
24
Thus the only possible solution when 2A + B Cv. 0 is ni = and n0 everything together, the solution set is precisely
B+C
Putting
24
{(11, (=1,
Proposition C.9. The fixed points of the dynamic can be classified as follows:
1. The points (Ro, k1) = (11,=1) are global minima Rf on 1|2.
2. The points (Ro, K1) = (0, =1) and (+1,0) are strict saddle points of R on 1-1,
3. The point (Ko, K1) (0, 0) is global maxima Rf
R on 1-1
Proof. Since R is smooth, its extrema on the square [-1, 1J occur either at critical points, where the gradient vanishes, or on the boundary. For the chosen range À E|0, (o, L)L. the only interior critical point is (0, 0), which corresponds local maximum. Indeed, the Hessian R (0, is
2B 0
V S (0,0)
2B
We easily check that both eigenvalues are 2B, which is negative since B < 0, therefore concluding that (0, 0) is maximum. The rest of the extrema of R must lie on the boundary of the square, we observe that
R- (=11,#11 R R= (x, (z. y) 0([-1, =11,
Thus concluding that (#1, minimizers. Finally, observe that
R R5 +1)
< R < (0, +1) < R (0, K,n E] - 1,
Analogously, we see that
(+1,k
R c (+1, 0) < R (), K - 1.
This shows that (0. +1) and (+1, 0) are strict saddle points, concluding with the proof. note that when À (a, 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 7 > such that for any stepsize 0 < y the iterates (wl xc)) ke) generated by 25 converge element of 6.
Proof. By Proposition C.7 the distance between successive iterates (NO 16) kEN, 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 unique accumulation point. Besides, the sequence belongs to the compact set 1, 1)2, then it converges and its limit one of the nine fixed points.
27

---

Proposition C.11. Consider the context of Proposition C.7 then there exists 7. > O such that for any stepsize 0 y < 3, the set of initializations such that the iterates (Ky 16 keN generated by
25 converge to (0, +1). ()),0) or (0, 0) has Lebesgue measure zero (with respect to the Lebesgue measure on the manifold M).
Proof. The point (0, 0) a maxima of the risk R on [-1, 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 yp is differentiable on [-1, 1]2 and that its Jacobian not degenerate, besides op is 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, oof. The result follows
from the Center-Stable Manifold Theorem (Shub) 1987, Theorem III.7), we refer to Marion et al. 22024 Proposition 13) for 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 (Xe)isis i...d drawn from P Consider also Ze E {0, 1) the latent variable of Xt. i.e. Ze|Ze NN 0² Ia), and
T po (X)1 =
2 Σ
(ex(m)) + e
where (p) def X(x1, 11) (Kk. p1). Then
lin (x)) |21 = =c]
111 + - c= 10.
E fin (X)]] Lú 1 X
3)o2).
L
Moreover. when X [L+1] 24 then the encoding is unbiased, this is
E (())) = = M2: c {0, 1}.
Proof. We decompose the following term follows,
(Xx, = (X1, + Σ (x,, Nr)
Due to the independence of the variables,
2
lin. mi (X)= (E(( + (X1, (i),)i)))
+ (L 1) MD) (X1 (X2, x2)x
On the one hand have
Xi 122, + e~N(0,
(x), 121 (42, 26)2 + +2222 M0) (e, mb) (e, mi)?
(X),)5)² 2 X1|Zi = (U2,, 2 P2z + (42, M8) 2 + 2(12, m (e, Hol
+ +24U3 (6) mā)e + (e, 21² + (c, 2
E (X1,
21]
(43, 2 + E El 2c + B/2 H60 (E, M6)
28

---

+ E|2(M2, : (E) + (E),,MO) + Blie, B|(E,AA) 2))
202 142200
ORKE,
S1 1z, 133, + (12, +2)42,
On the other hand,
X;|Z: 11z, + Gis 61 oT)); 61 -1 E2.
1I (X2, X2 Z1, Σ2] = 1B[(X1, MJ)
B] X2 Zg]
(u2, 1o) 2 16) + (82, 11) (uz + 62]
- (uiz, Mo) (1 16) + (1 Bl2) + Bliez, 1 uz +o² ui)
(62 (422, H6) + o²Pi)
Therefore.
2x L
Elplin.nosai (x)
(122, (12, )uz, + 2000000 20²uu
-000 20000(2200 (122, 16) + (p2, mi)
(I. 1)(((z 10) /1 m²) + 11) (12,
(L. - ( (1)2, + (u2 mi)1))}
And then,
Mo (X)
21 11p22 M1) (1i)² (42 11) + +20222, 20 (422 + (1)m)))
+ (L - 1) ( +0² (1M2, ( mō) + vi),))
Which let us conclude that for {0, 1}.
Ho 1 = c]
ped
((I+ 2(T. + +3)0?)
L
(PD + μT)x ((L+1) 1) +2
E T H6 (X)1]
2L
Proposition D.2. Consider (Xt) i.i.d. drawn from P variable of Xt. i.e. Ze ~ N (1 IIa), and
Consider also Ze E (0, the latent
2
mi (x)1
Σ
(en(m6) (xT)))
L
where (x(p) del XX1, 1,, u)) (Xtx n)) Then for 1),
E Ell
12), c] L7 [1 +
412
o od 16) + So (d + 7) + 80 4)]
+ 212 (L - 1)
3 o 22 28) + 4 (d + 16) + 4 (d + 10)
122
+ A22 2 (L () 2)
1 + 60² 2 + 120 + 80
17
29

---

And
(X))1)
412
1 +
8o' (od &
+ 21² (1 1)
o2 28) 40 ( (d + 16) 4a (d + 10)
+ 22 (L
1)(L -2)
[1 + +60 + 120 80
I
X²
[L + 1) + + +300
2
When X
(4+))) the encoding is unbiased, with variance
L
Var lin. (x)
16) ()
[[L + 1) + 2(L 3)0212 2(L -
(d 28) + 40 (d+ 16) 40 (d 10)]
[(L + 1) 2(L + 3)02
(6--1)1 (L- 2)
[(L + 1) 2(L + 3)02 [1 60 + 12o +00 1.
Besides, when I. - get that,
Van (X)1 c] 20 A~ 1
1 26
In general, if X is not fixed and L + 0o, ge!
Var (X)1 Zi = c) ~ 2X² 2 (1
Remark D.3. We recall that
Var Z₁ = = o2d.
Notably, by selecting À ensure an unbiased encoding, the variance becomes independent of the dimension d and equals 20 This shows a variance reduction effect whenever the dimension is bigger than 2. More generally, X can be chosen independently of such that
21 + 20²)² <<
In this regime, the encoding also asymptotically reduces the variance of X2, conditioned on its cluster assignment, as the number of components L
x..
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. without loss of generality, we assume Z = 0, then we get that for po: #1 E M:
lin.pin (8) 0]
412
8)) + +o od 4))
L
412 L -
50 + 460
160
1
2x (T. (1-1)
[nd
(d
3)) 40 (1 + (d
1
x2 (L 1)(L 2)
[ + 4 20 66 86
Recalling that Ro = (Ho, 16)) K1 = (11, n4),, in order to compute Ell (x) = 0), we just need to replace No and K4 by in the previous expression, as follows
lil (x), 13 0]
30

---

412
+
80 (
T
1)
+412
1 40
160
1,2
+ +21 (T.
11+0% (
o
L
+ (L 1)(L
+ 46 +
80
L
4X2
+ ( 16) So (d 80 (d+
1
+ 21 (L - 1)
( + 40 (d + 16) 40 (d 10)]
L7
x L () 2)
[1 60 + 120 + +80
1
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 f, Xx [ Ta + 1
The risk of the oracle predictor :6 can be decomposed as follows
((T = + do²) (1o) + (110) (111)) - (1) (11) (111),
(27)
where, from the proof of Lemma C.I by taking Kō Ki 1
(4o) 1A
2))]
(TT) =
4x La
8d
80 80 4))]
(111) = 4xa 2 [1 10 + 24 + 160%)
(T) 2X - + 2 + 4
(II) 2 للاياا
(d 40 7) + 4 (d + 2)].
(III) X2 (
60 12o + 80
When L tends to co, only the first term together with (1) and (11I) contribute. Therefore, we obtain that
(1 + do² A/2 +02 +0 + 6a² + 12o + 80
L000
Choosing À 4120 (its value being independent of L) minimizes the equivalent bound obtained above. With such a choice, the equivalent becomes
(1 + 4o + 404
L(Tlin.to
(1 11 (1+do?)
1-000
1 60 420 + 8
(1+ 20
(1 + doz)
1,400
(1+ 203)3
(1 + da
20²
L+000
I'd - 2).
L+30
31

---

Remark D.A. In the case of the degenerate case fa = 0), similar computations lead to
41
X²
L 4-1
L -
L- 1
1,2 + X2 1)(L 2)
C(Tlinnii -1
41
2X
+ 2X
T.
T,2
1,2
T.
1.2
I ++1
(T.
+3)
= 2X
+ X²
L
L
Optimizing this quantity w.r.t. À leads to choose
L+1 L+3
plugging this value for À gives
(T. 1)2 +(L+3)
E(TlinAAoN =
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 zero, matching that of the optimal quantizer.
E
Technical results
This section gathers 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 ~ NNO0, La) and pas Pbx Me E Rdd, then
1. ||||]|]]] ddd.
2. Ep G = 0.
3. E( G)G) o² Pa.
4. (p Gy(us, G3] o² (Pas ub)-
5. (M,, () G)2l ot (l1 || 2 2))3))
6. (p Gl G)² (Me: G)] ('(luoll (Por He) 2(pres uo) pe)).
7. (M Gy G) |||?? o'(d 2)(Haa
8. (p (u G) dd () 26Haa
Lemma E.2. Consider X ~ N(,,,, o² Id) where ||~0|| and Ma E Rdd, then
,a) (o) (p' Ha) (1 + (21+o(
Proof. We decompose X as follows,
= p +E, E E~ N~((00=2aa))
(X, pa) 1 ((H",aa) (E. 2 2(e, n²) + [|=|1))
E| m.. 1 || = (p",ua) + (a) BB|le|| + 4/p² ma) Bl(e, pa) (e) k²)l
+ Ell,e, (olls,Ma)2 E ((, m..) ||=
Then, by Lemma E. one obtains
xx (u²)),,)(0 oz + 4)) + (1
Lemma E.3. Let X ~ N(u,, o2la). where || and Has Mub E Rdd. then
po(Hay pb, p²* der (XX(X, pay) (2, (b) (|((1
32

---

can be expressed as
PolHa, Pb: (b:
(y, Ma)² (p* 1n)
+ 0² (1r (b) a1laaall 112 4(x², Ho) ub) (Maa No) + (4", (u",p)) (l
+ o² o² 8)(1", pa)2 (10)²
+ o (11lall ||/b111 + +2(Mas Nb)² +
,b)² + l| ma) )))
+ 40 40 + ,)("",a))rr" pb) (Mas Mb) + 0d () + 2(Ma: 2).
Proof. Write X = 1² + with ~ (N( TT)) Then
(X, 2a) (b) ) ||111 = (1,", + (e, pa)) ((n", ((n")) (1)) + (e, ps)) + 2(., v²*) + ||e|13).
Expand the product and drop all odd-moment terms of E (their expectation is 0). The surviving types of expectations are of the form Lemma E.I collecting all nonzero contributions after expansion and simplifying gives the stated formula for (o (66, v1")
The proofs of Lemmas E.4 H_7 follow the same approach the proof of Lemma
E.3 rewrite
X; Z; as uzz + Ei. where (5i); are i.i.d. random variables with s; ~ (O, la). We then expand the product, discard odd-moment terms, and apply Lemma E.1 to obtain the desired results.
Lemma E.4. Let Zi and Z2 E {0,1} be fixed. Consider two independent -va-ue random variables X1 and X2, such that
X. N(22, Ia), foreach each i = {1,2},
where the unit vectors Par 16 (i.e. || |1 1) are orthogonal. For pas Mbs Me E Rd, define
P1.0( Pb. Pb, UZ, Mc) Ma) (X1, (X1, Zz.].
This quantity satisfies
Pio(Hay Hb: 12, Mc) = (uz, Haa (Hz, Hoi (122, He)
+ o² [[|mal] 1l (12 Ho) (12,, Mo) + 2(H22 Ma) (H22, (Mas mo) + (U22, Mb) (Ma: Mc))]
o² [122 (pbr Holl
(1llaall (140 Mc) 2(Ma: (4b) (Mas)e)))
Moreover, we define
Pi Mb, Uz, 1zzz)) del
X2) X2) Zz).
which satisfies
Pr(pas piz M2zz)= = Na) +
( Mb, ,b))
Lemma E.5. Let Z,, Zz E {0, be fixed. Consider two independent R-=valled random variables Xi and X2, such that
1, ~ N(22, 02T)) for each each = {1,2),
where the unit vectors pas M5 (i.e. || = || = 1) are orthogonal. For Has 4bb E Rdd we get that
E El((X,, ((11) Ma) (X2, Ho) + (X1, (X2, (X1) X2)) [2) Z1, 7z)
(UZ, 12,, M/12 Ma) (/2, Ma) + (122, (12) Hb)) o² (Maa 122z)²² 2 (Mn + (1o, p2, )2 (14, )2)
5 (1|Hall +11mo11
33

---

Lemma E.6. Let Z1. Z2 E {0, 1] be fixed. Consider two independent R vvalud random variables
X1 and X2, such that
X1IZ ~ N( I2), for each i = {1,2),
where the unit vectors Par pb (i.e., | || || 1) orthogonal. For pax Mbb R.. define also:
P2,0 Mb. 12z,) E| (x1,
def EK Ma (x2 Mo) Xx2 [22] P2,1 (Ma Mb: 12zz)
[22)
These quantities satisfy
P2.0 ubb uz,1 = (1z Ho) (1Z, Hb) + o² (Ha,
P2.1 (2.1 uZz) ) = (uzz (a) (1Z2 Ho) + o² + 4)( Ma) (MZ. mo) + (Mar mb)
0² (d + 2) 2) Hb)-
Moreover, we define
P2(Ha, Hb: Uzz (Z,) E ( (X), (ax
Zz).
which satisfies
p2(Ma: Mbo uz, uZ, .22,) = P2,0lua, Pz, P2,, (Pax Mbb 1zzz))
Lemma E.7. Let Zı, Zz, Za E {0, 1} be fixed. Consider three independent -aluul random
variables X1, X2, X3. where
X:12i ~ N(22, o Ia), for each (1,2,3)
such that Pa 1h unit vectors (i.e. || = | = 1) are orthogonal. For pas Pb E Rd,, define also:
Ps,u(Pas Mb p'z, def [B[(X1, Pa)) (o),
Ps,1 (Jas Pb: P'z) def [x Ma) (X3, (b6 (X2, Xz) 72, Zo).
These quantities satisfy
P3 Pb; 1z,z)
(1Z, Ma) (42, 14) + o² 022(11 (Mar Hb),
P3, (Ma Mb, u2,, 1122)) (u2, Ma) (42, (1z u22z)
0² (1220 Ha) (12,, mb) + (U22 Ma) (u2s, 0 (p)
Moreover, we define
Pa(Has pb: p2z, M2,)) E[(X1, Ma) (X), Ho. Zz, Zz],
which satisfies
Ps(Pa: Pb, uzz, 1Zz0 1z,)) Ps,u(Ha, Mzz )a,1 (Pay
(Pa, 1b, /z Hzzz)
F
Experimental details
This section provides algorithmic details, choices of parameters, and settings used for the plots displayed in Sections A and
F11
Projected Stochastic Gradient Descent
We formally define the method Projected Stochastie Gradient Descent (PSGD), which we run for our numerical experiments.
34

---

PSGD iterates for linear attention heads. Given the objective function R : =8d-1)2 + R defined in Pep we define h : (Su-1)2 x IR Lxd as
2 Σ
h(po: M1, X) = X1
Xx (HORT + Man ) x
piuo, (1) X1) (28)
k=1
where X, is the i-th row of the matrix X. Consequently we can write
(uo, v1) Ex~D μl, X)],
where the distribution over IR Lxd where each row i.id. according to
1
2 N Ta) + 2 Niui,
Then, given and an initialization recursively by:
(m), v9) E (Sd-1)2, stepsize y, we define (/) (8d-15a
M
1
no n(mu ph. (7)
M
1
9
V ne h( MM E. et,,,
M
(PSGD)
M 7(aa M (ut)) T Dg0
10
k+1
MM 7(1a P5 (et) )9e ll2
A ((a (M() Do
A A+I
11T >(La (x T)
where M is called the batch size, and for each l: E N. (E (( are M independents samples of DD.
PSGD iterates for softmax attention heads. Given the objective function Reoft.po (Sd-1)2
R2 defined in (po for simplicity, we note that for an appropriate hg, we can write
Reoft.po (po, 14, 15, A) = Ex-d (No AI, v, X, X)],
where D is the distribution over Lxdd where each row is i.i.d. according to
1 2 (( + 2 N oF
1
o²Ia).
Then, given and an initialization (ng: 2 E (Sd-1)3, (w) 22) = (2,3), stepsize y, we define (ph, 2f) E (22 and (v* 2²1) E R2 recursively by:
M V Vag holu SoolunsHt ui, wh A (h,
M
M
1
o
V ps ho ui, 2h XR e.)
M
111
r(a (1( '99
A+1
16
11
(Ja
u (vE)) )
(PSGD
7la
(M+)) Dg
6+1
7Ia
M
2 1+1
1**
V ha ui, wh AA .t).
M
1
*+1
xh
Σ
v Vaho(ug pt, 1ith A ct)
M
1=1
where M called the batch size, and for each & E N. (ct (=(...) are M independents samples of D.
35

---

Remark F.1. Gradient computations in the mumerical experiments were carried out using JAX (Bradbury al. 2018).
F.2 Experimental details
In the following, we provide the experimental setup corresponding to Sections A and
We use input sequences of length L = 30 of 5-dimensional tokens (d = 5), and define the true centroids as μo (0, 0, 0, 1) and μ (- 0). We recall that the metric used quantify the distance to the centroids (up to sign) is defined in
Experimental details of Section A Regarding the experiment on the manifold, i.c., Figure 5a we perform 10a (PSGD) iterations without regularization (p = with a learning rate of y 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, 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* iterations to observe that without adding regularization term, only get partial alignment of the Transformer parameters towards the true centroids.
Then, in Figure 6a perform 5 10 iterations of (PSGD) to minimize the regularized risk R for 15 values the regularization strength p. linearly spaced [0, 0.3]. Finally, in Figure 6b we choose 0.1 and perform 104 (PSGD) iterations.
Regarding the experiment on the manifold, i.c., Figure D run the algorithm for 10a iterations without regularization (p 0), with a learning rate of y = 0.01, batch size M = 256, and choosing À 0.6 for a 0.3, and À = 0.2 for a 1. The experiment is repeated across 10 independent runs, each initialized randomly on the manifold M.
Experimental details of Section 3
For the rest of the experiments of this section, we adopt the same setup as before, with the exception that each run randomly initialized on the unit sphere. In Figure 2a we perform 5 103 iterations of (PSGD) to minimize the regularized risk R2 for 30 values of the regularization strength linearly spaced in [0, 3). Finally, in Figure 2h we choose p = 0.2 and perform 10a (PSGD) iterations.
Remark F.2. All experiments in Section Section[Aand And be run on a standard laptop. Most complete within a few minutes, with the exception of those in Figures ba and which require approximately 20 minutes and up to an hour, respectively, due to repeated moblem--sllving 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, the softmax variant cxhibits numerical instability in higher dimensions.
G.1 Influence of the dimension
Experiments RR 100 We use input sequences of length = 30 in R 1000 where we define two 0.1)
centroids: 1 (0, 00.
1) and ii (-1,0,...,0) The model is trained using (PSGD) with
99 times
online batch sampling strategy, with batch size of 256. and 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 sign by the concept of minimal root mean squared eиror, which is nothing but the distance to the centroid (up to a sign) divided by the square root the dimension, i.c..
1 min min
Minimal RMSE
[1ā=()) - s:14 11?
Vd
#ES, sE{ --1,1
36

---

10
2x10
RMSS
10
RMSE
10"4
10
6x
10-31
10000
2000
0000
6000
0000
10000
(a) ===0= O.
(b) σ = 0.3.
(c) a=)=
Figure 7: Minimal RMSE Iterations, Initialization on the manifold in dimension 100.
2
10
IMSE
10-3
10
10-4
3x
2000
4000
8008
10000
2000
4000 5000
8000
10000
Iterations
Iterations
lterations
(a) a
(b) σ 0.3.
(c)
Figure 8: Minimal RMSE Iterations in dimension 100, Regularization 0.1 for p 0.2
for σ > 0, Initialization on the unit sphere.
where S2 is the symmetric group of order 2, No: μ4 are the true centroids, and Aoo Pr are the returned parameters from PPSDD). In Figures| 7 we can observe the behavior of the RMSE over the iterations for different levels of noise We remark that in Figure we initialize on the manifold M, and there no regularization term (i.e. p = 0), Figure we initialize randomly over the unit sphere and we set p 0.2. In both experiments we set X 0.6 for the case 0 and 0.3, and 0.2 the case a
In each experiment, the RMSE of the order 10-2, which can be interpreted as, on average per coordinate, the estimators Ao fi are missing the true parameters M pi by 10 suggesting a high level of accuracy in the estimation procedure.
Making d vary. We repeat the experiment before, just varying the dimension of the problem, 0,
the two centroids in Rd are defined by /á (O.....0.1) and n1 = (-1.0......0) 0). For d ranging
d-1 times
d-1 times
between and 200, we show Figures and TO on the x-axis dimension 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 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 the initialization regime, the noiseless case (a 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 remains of order 10-3 and for o 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 interpreted as the average discrepancy per coordinate between estimated parameters and the true centroids, suggesting high level of accuracy in each experiment.
BMSE
PANSE
RMASIE
10-10
10~3
100
200
100
150
100 Dimersion
Dimension
Dimension
(a)
(b)
0.3.
(c)
Figure 9: Minimal RMSE vs Dimensionality, Initialization on the manifold M.
37

---

4x10-1
10
RMSS
RMSE
AMSE
1x
10-
50
100
150
200
30
100 Dimension
150
200
10
50
100
150
200
Dimension
Dimension
a O.
(b) σ 0.3.
(c) =1
Figure 10: Minimal RMSE vs Dimensionality, Regularization p 0.1 for o = 0, p = 0.2 for σ > Initialization the unit sphere.
Sigma
0.3
RMSE
10
1.0
2000
4000 Iterations
6000
8000
10000
Figure 11: Minimal RMSE vs Iterations in dimension 50, Random initialization on the unit sphere of
the centroids and initial guesses, Regularization p = 0.2, 10 runs, 95% percentile intervals are plotted.
Regarding rumning times, the experiments run in Figures 7 and takee approximately two hours standard laptop, while that of Figures and 10 may require up to 12 hours, to cover the dimension grid.
G.2 Relaxing the orthogonality assumption
We replicate the experiments and initializing the centroids and the initial points uniformly random the sphere S in each run. Figure illustrates the algorithm's '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.
procedure from Section G.I but this time
We observed the expected behavior: the dimension increases, randomly initializing centroids on the sphere makes them more likely to orthogonal, and thus training via the regularized theoretical risk yields better results higher dimensions. This effeet is stronger lower noise levels and becomes noticeably clearer beyond 40 dimensions.
Regarding running times, the experiments run in Figure II take approximately one hour on standard laptop, while those of Figure 12 may require up to hours, cover the dimension grid.
Sigma 0.3
RMSE
1.0
40
60
100
Dimension
Figure 12: Minimal RMSE vs Dimensionality, Random initialization on the unit sphere of the centroids and of initial guesses, Regularization p = 0.2, 10 runs, 95% percentile intervals are plotted.
38

---

G.3
Euclidean Projected SGD
Based on Section El and the algorithm PSGD) presented there, we define the projected Euclidean SGD as the following update rule, with initialization (1): mq) ($d-1)2, and stepsize
M
1 M
no hooh
M
o4
V Ma h(wu: pr, k/:,
M
==1
(PSGD Euclidean)
1 700
k++
P4
196 1l
a 7
1
79 12
where & is defined in 28
Our empirical analysis reveals that it is 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 the unit sphere), we repeated the experiments outlined in Section 3 The results are presented Figure [3] This empirical equivalence confirms that the simple projection of the Euclidean gradient serves effective and computationally simpler proxy for the true Riemannian gradient.
Distance to centroid
10a
Sigma
Distance to centroid
100
Sigma a= 0.3
to a sign)
( to sign)
0.3
10-1
1
100
1
10-
10
10-3
2000
4000 6000
8000 10000
2000
4000 6000 8000 10000 Iterations
Iterations
(a) Distance centroids vs Projected SGD iterations (b) Distance to centroids vs Projected SGD iterations for the minimization of R. initialization on the mani- for the minimization of R", initialization on the unit sphere, with regularization 0.2.
fold.
Figure 13: Performance Projected SGD. runs, 95% percentile intervals are plotted.
G.4 Extension to Gaussian mixture model with 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
X ~ 3 N( IaI)) + N o²Ia) + N(2,, 2
(Pa)
where Hό: Mi, /12 are orthonormal vectors. It is natural to consider attention-based predictor composed of three attention heads, parameterized by po, (1, E Rdd,
T po M2 (X) lin. (X) Hlin, (X) M2i (X).
(29)
The associated risk is R(po, M1 u2) = B|J[X, - Mo.. (X [.
There are two natural gener-
alizations of the regularization term to this case:
Mi, u2)
Σ
(p,, (13;
39

---

Distance to centroid
Distance to centroid
(up to a sign)
10
Sigma
Ln to a sign)
10"
Sigma
0.3
0.3
1.0
1.0
10
10--
0
2500
5000
7500
10000
5000 10000 15000 Iterations
20000
Iterations
(a) With regularization p())
(b) With regularization y())
Figure 14: Distance to centroids vs number of iterations, with regularization strength = 0.2, initialization on the unit sphere. 10 runs, 95% percentile intervals plotted.
( P1, p22 = II (4,,X))²
The first promotes pairwise orthogonality while the second one promotes mutual orthogonality.
We use input sequences of length L 30 in R where we define the three centroids μá (1, O, 0, O, 0, 0), μ (0.0,0 0, O, 1, 0. 0). μu = (0. 0, 0, 0, O, 1). The model is trained with online batch sampling strategy (similar to PSGD] changing the data distribution and the regularization term), with a batch size of 256, and learning rate of 0.01. We take À = 0.6 for = 0.3, and À = 0.2 for a 1. Since any parameter could learn any centroid up to a sign, we introduce the following distance to the centroid (up sign):
Σ || 5: 13,
min HES (-1,1
min
where Sy is the symmetric group of order and fo fu P are the parameters returned by the algorithm. We present the results in Figure 14 We observe that the regularization y(1) outperforms 72)), since explicitly includes all pairwise terms enforce orthogonality. However, note that the number of regularization terms grows quadratically with the number of centroids.
Regarding running times, the experiments run in Figure standard laptop.
14 take approximately 15 minutes on
H Softmax attention layers and clustering
In this section, we assess the abilities attention- -aased predictors involving softmax activation in a clustering context.
H.1 Problem setting
An attention-baad based learner with softmax activation. self-attention layer can be written as follows:
We recall that attention head made of
H (x) = softmaxx (XQK XV
where the softmax of temperature À > 0 applied rise and the matrices K, Q, V E Rdxpp usually referred to keys, queries and values. As in Section 2 we assume that the values 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
E RR xx we obtain
ao (X) = softmaxA (Xpp²) Tx] X.
(30)
40

---

With such an architecture, the L-th output vector is therefore given by
H (x)e
Σ softmax (Xx uu xx Xx.
(31)
which corresponds to aggregating the Xx when X and X₂ are simultaneously aligned with This head should a good candidate estimate a centroid of a mixture model. In the case where the mixture involves two components, one could train two attention heads:
(jo. f1) E argmin p0.10 E84-1 R (Mo 111)
(32)
where
Reoft Ruo (10, =
| - (H + Hanft (XX) [1]]
(33)
1 L
E
Xe (Hooftx. Hoon (X))
Remark H.1 (The attention heads are biased). As the tokens (Xe)e are independent, we have that
Roof (1) = E|||x - (H + H (
We note that (x) = softmax, ((X1, u)))x, where the vector is L-dimensional with components (X), In the idealized case where o² O. then each token Xe is sampled according to mixlure of Dirac masses given by 10ng + 28m.. Therefore, if we evaluate Heols H (X)) on J Poj we observe the following:
Conditionally X1 p4. then
de (X)1
1 x.
This implies that even in a completely misaligned set-up (i.e., μ === M6. =()), the proposed attention head will return, as the transformation of the first token X 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 X1 = Ph. then
exp(A)μg pi exp(A) +1
pG (X)) (1)X
This suggests that in the perfectly aligned case (i.e., = 1 X1 Mj). selecting sufficiently large softmax temperature X will cause the model to assign negligible weight to the misaligned components desirable property.
Debiasing and disentangling heads. To handle the bias introduced by the attention heads, discussed in Remark H.I propose to consider centered heads instead, leading to the following modified version of the risk
27
Reoh (Mo. A, 1
X, Ho + m (x3) +
(34)
k=
Considering such a risk is equivalent to using heads where a term proportional to E X₁ is substracted. This type of head is known 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, would expect that single head provides all the necessary information, making it sufficient to debias only that head (i.e. d = 1). In that case, one should obtain:
Reoft (HŌ, Mi, X*.1) Ro min Rooft
(35)
41

---

However, when using non-oracle parameters do and μ사 within the debiased heads, the risk function may admit global minima where the heads align with zero, one, both centroids, which is undesirable for the clustering purpose. Therefore, we must enforce constraint ensuring that each head aligns with exactly one centroid. To achieve this, we introduce the regularization term:
ro(Ho: del [((oo X1) 1)2 ((11, X1) 1) (4),)1),
leading to the following regularized optimization problem
min No.11 2
poo (1o, μ1, A, 1/) Ro (uo,Hi,A, 2, + poro(uu.
(Ppo)
where Po >
H.2 Numerical experiments
We run Projected Stochastic Gradient Descent (see Appendix F.I to learn the centroids μó and μi as well as the weights v and
In this experiment, we use input sequences of length = 30 of 5-dimensional tokens (d = 5), drawn from a 2-component Gaussian mixture of centroids μ (0. 0, 0, 1) and иィ (-1, 0, 0, 0). The variance of each component either set to a = 0.3 (low interference) or to a = 1 (high interference).
The model based on two softmax attention heads parameterized by μo and μ4 is trained using (PStDsoft with an online batch sampling strategy, with a batch size of 256, a learning rate of 0.01, and running for total of 3000 iterations. Additionally, we initialize with À set to 3 and centering value vi Here we use the metric distance to the centroids, given by
v min {distı, dist2],
(36)
where
dist 11 Hōl12 2 11 1il1?
dist2 = llāo - nT|12 + 11 nōl1?
and Prg 14 denote the true centroids, respectively, while Ao jún are the parameters returned by PPSGD We remark that this distance is finer than the one defined in [1. as it does not disregard sign flips. The results visualized in Figure [5a] we observe that regularization term substantially improves the accuracy of the recovered solutions. However, as the strength of the regularization in- creases, 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
Distance to centroid
Distance to centroid
100
10
Sigma 0=
Sigma o= 0.3
0 1
10-1
a= 1
10-1
0.0
0.5
1.0
1.5 2.0 2.5
3.0
2000
4000 6000
8000 10000
Regularization
Iterations
(a) Distance to the centroids after 5000 PSGD itera-(b) Distance centroids PSGD iterations for the tions vs regularization strength p for minimization minimization of R Rooft, with an initialization of R io with an initialization on the unit sphere, unit sphere and regularization Po = 0.5. with data
with data drawn from the non-degenerate case (Pa). drawn from the non-degenerate case (Pa) 10 runs, 10 runs, 95% percentile intervals are plotted
95% percentile intervals plotted.
Figure 15: Performance of (PSGD soft. with data drawn from the the non-degenerate case (P 10
runs, 95% percentile intervals are plotted.
In Figure 15b we set the regularization parameter (o to 0.5, and run PSG soft for 10 iterations. We observe that the model yields accurate solutions under low interference (a 0.3); however,
42

---

as the interference increases (a = 1), the ability of the softmax attention heads to align with the underlying centroids is progressively impaired. similar loss in alignment accuracy is observed as the dimensionality increases.
The experiments in Figure 15b in few minutes 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
Proof of Proposition 5.7
For e E {0, 1], one has
21
E
(X1,
Xx pi 16: 21 =c
2A
2X(L -
| X1
11, Mū: 21 = c]
[| X2) u1 14: Z = c]
L
2X
|||| 14 Mb: 21 +(I 1) (X1, ut 16: 21 = =c]
(1+(d+2)o2)x
(1+02)1E
2A (1 (1+ [d+ (L-
( + 0² Jui
T,
Where we have computed the conditional expectations by expanding each term using X; = μzz + E; with independent Gaussian eπrors E₁, discarding odd-momen add-moment terms, and applying Isserlis' theorem [Isserlis 1918).
We remark that choosing À (1 1)
we get that the encoding is unbiased.
And
2X
412
1 (d(++ 2)0² +(L (L- (1-1) 1) (5 2 oz)
E
X1, Xz) Xx
Pi 21
T
T,2
t=1
Besides,
2X L
X
1 + (d 2)0² (( (1----1 1) -2+ (15 + +11
[1 Xe
Hi,
Also,
2\
L
21² L
1 + (d 2)0² + (L 1) ( oz )}
E
(1
Xe 14:
Xz)
L
On the other hand,
21 L
2
4X2
11
Xx) Xe
11, 16 71 = c
TEll X1| 1 1o, 21 =
L
I
t=1
12X²
+
1 (L
(X,, X2l 1/ 21 = =c]
8X2 (I. - (0_1))[L-1
(X Xa) Zi c).
T22 2
Recalling the expressions stated at the beginning of the proof of Proposition 5.2 for moments of Gaussian r.v., we conclude that
2) Σ X1,
E
X
pi, Mb: Zi =
1=1
43

---

412
T,2 (1 + 3(d + 1)0² + 3(d + 2) 4)0² + d(d + 2) + 4)00)
12x²
1,2 (L 1) (1 2
(d+
+
0² +3(d+ + 2)0² + d(d + 2)0²)
2
8A2 (L- - 1)(L -2) 2)
2 (2(02
1)
+
0²
+
+ (d 2)08
1.2
2
2
And
21 I
Var
Σ (X1, XI)
UT Hō: 21 =c
22 L
2x 7
2
=
E
L
(X1: Xe
ui H6: 21 = c
L
(X1,
Xe 21 = c
t=1
t=1
4X2
I (1 + 3(d + 4)0² + +3(d+
+ d(d + 2) + 4)00)
12X²
1
(d+8) 8)
(T. 11--1 1)
+
0² + 3(d + 2) + d(d + +)00
12
2
2
8X2 (L - 1)((L -2)
1
3
+
2
2
+
+ (d - 2)06
1
2
2
4X2
(
1
+ + ]²
2
1+ (d + 2)0² + (L - 1) (d+
1
2
When L - 00.
2x
Var
(x1, Xz) Xe "i, 21 =
T.
2
2 202 + 1)
2) 1) + (d - 2)06 (d-2))0
1 1)2 02+
~ 4X2
-222
2
2X20² + 40 + 2d
Choosing the À = 142az, we have an unbiased encoding with variance
20 + 40² + 2 4
(1+ 20022
Proof of Proposition
5.2
We have
21
E
X1
X1 X²)
Xe
41 I
=
EE[|X,1' 11 pt mál
41
-
H6ll
I
Σ
B
ビ=2
412
4x2
E|||X| /1
Mi
Ei]l(,,, Xx)xxxx.x))xx||111x
12
8=2
8X2
8A2
B|||X, (X,, xx) 11)1 +
Σ
XX),X,,
L2
L2
25tckEL
44

---

Furthermore we have the following,
[[1 X1 I14, = 1 + +
E Xx) 14
= + +20 + do
2
- 2 + d+8 0² + 3(d + 2)0* + d(d + 2)0²
(2) 112,
x [o] 2
2)0² + d(d+2)o' 2)0",
B[| X1 11 ||1, = +1+ + 4)0² + +(d 3(d+2)(d
2) + 4) + + 4)0²,
Xs) (X2, Xx),xxMMii.. Hol = ( 2
+ (d - 2)0º.
These identities follow by writing
X,In = 14 + Es.
6, ~ 50,000 02aa) i.d..
and expanding the expressions in terms of μrt and E,. The expectations then reduce to Gaussian moments, which can 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 pi, MÓ. we have
21
L
X |
2A
L (x
1| (iir
E
X1
(x
E X1
X
Xx
L
L
1=1
1=1
And
21
Xe 2
E
Xi
(1
T,
i=1
4)
(1 + 2(d +10++22+++++ =(d++2()+)) -/1 (L 1) (1
41
=1+
+ 20 + do²)
L
L
4X2
12 (1+ (1+3(d+ + 4)0² +(d++2)( 4)0² d(d + 2)(d + 4)00)
12X²
(1-
d+8 0² + (3d + +6) + d(d + 2)00)
+
(L-1)
1
2
8X2 (L (1]1([L-)) 2)
(2 2 02
(3(0²
1 3
+
+ + - 2)08
2 2
2
When I. - oc, we obtain
2x
(1,, Xx) Xx) /|]
E
X1
X1)
=(1 =(1+02d) - 2X(1 + +402 + 2
L
E=1
+ 412
2 02+ 2)
1)
+ (d (d-2))0))
And we can choose
À-
1+40 +4-10+2200
to get
1 +) 2)00
21 L
(X1, Xo)x /
(1 + 40 + 2)o
E
X1 L
= = 0dd)
4 ( (02 + 1))² + (d - 2)00)
1 20²
= - 2)
1+ 6 12o + 6
<d 2).
45

---

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 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 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 is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
The authors transparent about the simplified nature of the architectures and mathematical models they study, and consistently highlight these choices as part the scope and limitations of their work.
Justification:
Guidelines:
The answer NA means that the paper has 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 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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant 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 prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
46

---

Answer: [Yes]
Justification: All theoretical assumptions are explicitly stated. Complete proof' are provided in Appendix B and C
with supporting technical results in Appendix
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 cross- referenced.
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 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 properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The main algorithm used for our numerical experiments is defined in Section F.I Detailed descriptions of the experimental setups are provided in Sections [A.2 3.2
Appendix G and H.2 This information enables others 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, 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 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 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 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions 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 new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or 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 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 (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
47

---

reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, 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 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.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7 w/n.ps.cc/publcc/////ccccccc
for more details.
The authors should provide instructions on data access and preparation, including how 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 subset of experiments 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) 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, hyper- parameters, 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
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, 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 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:
48

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence 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 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).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All experiments were conducted 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 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 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 into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics htttp::/heuripp..ccppbb..a..
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 (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
49

---

Answer: [NA]
Justification: This work is primarily theoretical and employs simplified, synthetic dis- tributions. It does not have direct societal implications beyond contributing to a deeper understanding of certain aspects of attention-base mechanisms.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper is mainly theoretical and does not involve the release of any models or datasets that could pose risk of misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
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 licensing. We did cite open-sourced libraries used in our paper, such JAX.
50

---

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 used and, possible, include 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 use in the package should be provided. For popular datasets, paperswi thcode. com camdsatasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that 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 '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/cood/mmodel 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals 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?
51

---

Answer:
[NA]
Justification:
This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such 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, 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 non non standaadd components. for what should or should not be described.
involve LLMs as any important, original,
Please refer to our LLM policy
52


Reviews:
[{"review_id": "rMfqz6dzLx", "reviewer": "Reviewer_3mf1", "summary": "This 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", "questions": "The 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?", "limitations": "Limitations are addressed in the work.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "***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.", "was_revised": true, "final_justification": "The work has an interesting theoretical contribution.\n\nAuthors answered my questions and even expanded some of their results.", "rebuttal": "We thank you for your feedback and for acknowledging the relevance of our work. \n\n**Initialization.** We agree that restricting initialization to the manifold is a limitation. Relaxing this would require proving the stability of the manifold, so that initializations sufficiently close to it still lead the PGD iterates to converge toward the manifold. Unfortunately, we have not yet been able to establish this result theoretically.\nTo address this limitation in practice, we introduce a regularization term that empirically promotes global convergence.\n\n**Simplified architecture.** We acknowledge that using a linear head is a simplification compared to practice. This choice was guided by two considerations: (i) this simplified head appears to be sufficient for the structure of the data at hand, and (ii) the mathematical analysis of softmax attention heads (especially when combined with a quadratic loss) is complex and lacks a conclusive theoretical framework to date.\n\n**More mixing components.** Thank you for this question. Our framework could be extended to handle $K\\leq d$ clusters with orthogonal centroids. The analysis would proceed along similar lines by employing $K$ linear attention heads. While technically feasible, such an extension would significantly increase the complexity and reduce the readability of the mathematical exposition, which motivated our choice to focus on the two-cluster case for clarity. \n\nIt is also worth noting that stacking attention heads can be reinterpreted as a single attention head with matrix-valued parameters subject to a rank constraint. Specifically, the transformation applied to the $\\ell$-th token in a two-head architecture can be expressed as:\n$\nT_\\ell(\\mathbb{X}) \\propto_\\lambda \\sum_{k=1}^L X_\\ell^\\top (\\mu_0 \\mu_0^\\top + \\mu_1 \\mu_1^\\top) X_k X_k = \\sum_{k=1}^L (X_\\ell^\\top Q Q^\\top X_k) X_k\n$\nwhere $Q\\in \\mathbb{R}^{d\\times 2}$, so that $QQ^\\top$ is a rank-2 matrix. Therefore, our work can be seen as an analysis of a non-convex attention-based strategy for implementing spectral-like methods. This perspective also sheds light on what happens when using fewer heads than cluster components: the resulting performance should be comparable to that observed in spectral clustering methods under similar conditions. We leave a formal analysis of this phenomenon for future work.\n\n**In-context setting.** We thank you for raising this question, as it led us to new results for an extended version of the architecture, within an in-context learning framework. We consider prompts composed of tokens drawn from a mixture of two Gaussians, each parameterized by prompt-dependent centroids that are randomly sampled on the unit sphere and mutually orthogonal.\n\nWhile the original architecture studied in the paper appears too limited for this scenario, a straightforward solution is to increase the number of heads. Since the head parameters in our work are assumed to be orthogonal, this leads us to consider at most $d$ heads. In this setup, the architecture computes, for the $\\ell$-th token, \n$T_\\ell(\\mathbb{X}) = \\sum_{h=1}^d \\sum_{k=1}^L (X_\\ell^\\top \\mu_h \\mu_h^\\top X_k) X_k = \\sum_{h=1}^d [ X_\\ell^\\top (\\sum_{h=1}^d \\mu_h \\mu_h^\\top) X_k ] X_k = \\sum_{h=1}^d \\sum_{k=1}^L (X_\\ell^\\top X_k) X_k.$\n\nInterestingly, stacking $d$ orthogonal heads in this setting results in an architecture with no parameters to learn. Yet, it can be shown that this architecture enjoys quantization properties similar to those established in the standard clustering framework, and furthermore successfully adapts to the underlying mixture structure. We will include these new developments in the revised version."}, {"review_id": "RfKTXYbCYp", "reviewer": "Reviewer_tsYz", "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.\nThe paper presents a theoretical convergence analysis, which is subsequently empirically demonstrated through simulations.\nBy combining a classical clustering", "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.\n2. Can the authors comment on the relation of the proposed method and the standard expectation-maximization method?\n3. How important are the assumptions on orthogonal centroids? Are they mandatory for the convergence of the method, or only for the analysis?\n4. Can you determine the number of classes from the attention mechanism?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths \n1. The paper provides a complete theoretical analysis of transformers as an unsupervised clustering mechanism. This offers new insights into the properties of transformers. \n2. The paper is well-written. Theorems are clearly stated, while their proofs are given in the appendices. \n3. The convergence curves, with and without regularization, are insightful.\n\nWeaknesses\n1. Only a simplified scenario is analyzed (two classes, orthogonal centroids).\n2. Only a brief discussion on the potential applications of the findings of the paper.\n3. 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.", "was_revised": true, "final_justification": "The paper analyzes transformers as an unsupervised clustering mechanism.\nThe authors have addressed all my comments and suggestions.\nI will stay with my 'Accept' decision as the paper proposes a fresh perspective on clustering, but the method  has some limitations, as explained adequately by the authors.\nOverall, I think that the paper should be accepted.", "rebuttal": "Thank you for your time to review our paper, and for acknowledging the relevance of our work. We provide some answers to your questions below.\n\n**Q1.** Our framework could be extended to handle $K\\leq d$ clusters with orthogonal centroids. The analysis would proceed along similar lines by employing $K$ linear attention heads. While technically feasible, such an extension would significantly increase the complexity and reduce the readability of the mathematical exposition, which motivated our choice to focus on the two-cluster case for clarity. \nThat being said, we ran new numerical experiments with $K=3$ clusters, 3 heads and random initializations. We first use a natural extension of the previously introduced regularization. In this specific case, it consists in adding 3 (=$\\binom{K}{2}$)  regularization terms to enforce pairwise orthogonality between the head parameters. We also observe that using a single, simpler regularization term based on the product $\\mathbb{E}[\\prod_{k=1}^K \\langle X_1, \\mu_k \\rangle^2]$ achieves almost comparable performance while significantly reducing the number of regularization terms. We will add these experiments to the revised version.\n\n**Q2.** While the two methods are fundamentally different and originate from distinct principles, the comparison is insightful. The EM algorithm is a likelihood-based approach that heavily relies on the prior assumption of a Gaussian mixture model. In contrast, the attention layer considered here operates in a more geometric fashion, aligning tokens based on directional similarity.\n\nThat said, some interesting parallels can be drawn. In the EM algorithm for Gaussian mixtures, the steps are explicit, so that the centroids are estimated by aggregating data points weighted by soft-assignment scores, which reflect the estimated probability of membership in each cluster. While the attention-based model lacks an underlying probabilistic framework, it also performs a weighted aggregation of tokens, this time using alignment scores that quantify directional similarity with certain vectors.\n\nWe thank you for prompting us to formalize this comparison, and we will include this discussion in the revised version of the paper.\n\n**Q3.**  The attention-based model relies on directional similarity to perform clustering. If the clusters are aligned (or anti-aligned), the model is not able to distinguish them. For intermediate cases (neither orthogonal nor aligned), additional cross terms complicate the analysis. \nIt is nevertheless worth noting that as $d$ increases, even without explicitly enforcing orthogonality between the centroids, randomly choosing them on the sphere results in centroids that are nearly orthogonal. Thus, using the algorithms introduced in the paper, we obtain similar numerical results with centroids randomly chosen on the sphere, even though they are not strictly orthogonal. \n\nThe formal analysis in this more general case is an interesting direction for future work. \n\n**Q4.** We conducted new numerical experiments where the inputs were drawn from a two-component mixture, but the model used three attention heads, meaning there were more heads than necessary. In the low-noise regime, one head successfully recovers a centroid, while the remaining two heads align with each other, but the recovery of the second centroid is partial. However, as the noise level increases, the attention parameters struggle to align with the true centroids, suggesting the existence of bad local minima in the optimization landscape under such conditions. Notably, this setting offers a natural diagnostic: when too many heads are used, redundant heads tend to converge to the same parameters, effectively signaling the model's overcapacity. We leave further investigation for future work."}, {"review_id": "sx2UsAfcSr", "reviewer": "Reviewer_g8nw", "summary": "The 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.", "questions": "- 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?", "limitations": "Yes.", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strength:\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.", "was_revised": true, "final_justification": "I still think the paper simplifies attention to essentially a mixture model with two components which renders theoretical results rather uninteresting. However, the other reviewers convinced me that there is something in the paper and I don't object acceptance. I'll stick to my score though. (Let me know in case this is a problem).", "rebuttal": "Thank you for your time dedicated to review our paper. We would like to begin by addressing and putting into perspective some of your concerns about our work. \nMore specifically, there is generally a tradeoff in theoretical works between analytical tractability and resemblance to practice.\nDespite some simplifications, we believe our results provide valuable insights into both the diverse learning capabilities of transformers and a novel perspective on attention mechanisms for unsupervised learning.\n\nMore detailed responses to your questions follow below.\n\n**Q1.** Our theoretical analysis is conducted under three main assumptions:\n(i) a two-component Gaussian mixture model,\n(ii) centroids constrained to lie on the unit sphere, and\n(iii) mutual orthogonality of these centroids.\n\nRegarding (i): Our framework could be extended to handle $K\\leq d$ clusters with orthogonal centroids. The analysis would proceed along similar lines by employing $K$ linear attention heads. While technically feasible, such an extension would significantly increase the complexity and reduce the readability of the mathematical exposition, which motivated our choice to focus on the two-cluster case for clarity. That being said, we ran new numerical experiments with $K=3$ clusters, 3 heads and random initializations. We first use a natural extension of the previously introduced regularization. In this specific case, it consists in adding 3 (=$\\binom{K}{2}$) regularization terms to enforce pairwise orthogonality between the head parameters. We also observe that using a single, simpler regularization term based on the product $\\mathbb{E}[\\prod_{k=1}^K \\langle X_1, \\mu_k \\rangle^2]$ achieves almost comparable performance while significantly reducing the number of regularization terms. We will add these experiments to the revised version.\n\nRegarding (ii): Constraining centroids to lie on the sphere serves enables the use of linear attention heads in place of softmax-based mechanisms. This assumption effectively normalizes the queries and keys, a role typically played by explicit normalization layers or softmax operations in standard architectures.\n\nRegarding (iii): The attention-based model relies on directional similarity to perform clustering. If the clusters are aligned (or anti-aligned), the model is not able to distinguish them. For intermediate cases (neither orthogonal nor aligned), additional cross terms complicate the analysis. The orthogonality assumption remains a mild condition in high dimensions. We ran additional experiments with randomly initialized centroids on the unit sphere in dimension 10 to support this claim, obtaining similar results to the orthogonal case, and will add these results to the revised version. \n\nIn summary, while these assumptions may seem restrictive, they were chosen deliberately to keep the analysis tractable and transparent. We believe they still capture essential aspects of the mechanism and provide useful theoretical insights.\n\nUsing Riemannian gradient descent is the most natural and mathematically sound approach, given that we can only compute the risk for parameters that lie on the sphere. As such, we are not in a position to compute a standard Euclidean gradient and must instead project each update onto the tangent space of the sphere. Moreover, this choice does not appear to be restrictive in practice: numerical experiments using standard SGD strategies behave similarly to the theory. \n\nWe thank you for having raised these concerns, and will discuss them in the revised version. \n\n**Q2.** We emphasize that linear attention is a standard and widely studied setting (see among many others [1, 2] and references therein). Importantly, any linear attention head can be decomposed into a sum of heads with one-dimensional keys and queries. Therefore, we respectfully disagree with your remark regarding the factorization.\nIn this setting, the attention score between tokens $\\ell$ and $k$ is $X_\\ell^\\top \\mu \\mu^\\top X_k = \\langle X_k , \\mu \\rangle \\langle X_\\ell , \\mu \\rangle$. It measures a specific type of dependency between tokens, namely, their joint alignment with a fixed direction $\\mu$, determined by the attention head. Exploring extensions to nonlinear attention mechanisms is an interesting direction left for future work.\n\n[1] Trained Transformers Learn Linear Models In-Context, Zhang, Frei, Bartlett, JMLR 24.\n\n[2] Linear attention is (maybe) all you need (to understand transformer optimization), Ahn et al., ICLR 24.\n\n\n**Q3.** Thank you for raising this question. We used the phrase ``variance reduction'' to mean that the attention-based predictor achieves lower reconstruction error than predicting the mean of the cluster---for which the reconstruction error is exactly the variance of the cluster. We acknowledge that this formulation was unclear, and will clarify in the revised version.\n\nRegarding the influence of $\\lambda$: for simplicity, Proposition 5.2 reports the result for a specific value of $\\lambda$ (corresponding to the minimal reconstruction error). However, our computations are carried out for arbitrary values of $\\lambda$, and show that the variance reduction effect holds over a range of values and not just for a single point. For instance, when $\\sigma=1$, this effect occurs for $\\lambda$ approximately between 0.3 and 0.4.\n\nFinally, note that we have derived additional results on the approximate quantization phenomenon in a more challenging in-context setting, where mixture parameters vary for each prompt. We consider prompts composed of tokens drawn from a mixture of two Gaussians, each parameterized by prompt-dependent centroids that are randomly sampled on the unit sphere and mutually orthogonal. While the original architecture studied in the paper appears too limited for this scenario, a straightforward solution is to increase the number of heads. Since the head parameters in our work are assumed to be orthogonal, this leads us to consider at most $d$ heads. In this setup, the architecture computes, for the $\\ell$-th token, \n$T_\\ell(\\mathbb{X}) = \\sum_{h=1}^d \\sum_{k=1}^L (X_\\ell^\\top \\mu_h \\mu_h^\\top X_k) X_k = \\sum_{h=1}^d [ X_\\ell^\\top (\\sum_{h=1}^d \\mu_h \\mu_h^\\top) X_k ] X_k = \\sum_{h=1}^d \\sum_{k=1}^L (X_\\ell^\\top X_k) X_k.$\n\nInterestingly, stacking $d$ orthogonal heads in this setting results in an architecture with no parameters to learn. Yet, it can be shown that this architecture enjoys quantization properties similar to those established in the standard clustering framework, and furthermore successfully adapts to the underlying mixture structure. We will include these new developments in the revised version."}, {"review_id": "IyaMeTLUTB", "reviewer": "Reviewer_8aE9", "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.", "questions": "Please refer to the weaknesses section.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "- Strengths:\n\t1. 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.\n\t2. 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.\n- Weaknesses:\n\t1. 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.\n\t2. The current proofs and derivations are based on the simplified example. Could the derivation extend to more complicated scenarios?\n\t3. Since the attention mechanism is proven to be effective in clustering, is it possible to evaluate its actual clustering performance?", "was_revised": true, "final_justification": "The authors' response has addressed my concerns. I would like to keep my original score.", "rebuttal": "We would like to thank you for your feedback. In the following, we provide answers to the weaknesses you underlined.\n\n**W1 (Illustration).** \nWe appreciate the suggestion. In the revised version, we will include a 2-dimensional illustration showing the input data---two point clouds around their cluster centers---and the corresponding outputs of the trained attention layers, which form flattened point clouds, illustrating how the attention model performs approximate quantization.\n\n**W2 (Towards more complicated scenarios).** \nOur theoretical analysis is conducted under three main assumptions:\n(i) a two-component Gaussian mixture model,\n(ii) centroids constrained to lie on the unit sphere, and\n(iii) mutual orthogonality of these centroids.\n\nRegarding (i): Our framework could be extended to handle $K\\leq d$ clusters with orthogonal centroids. The analysis would proceed along similar lines by employing $K$ linear attention heads. While technically feasible, such an extension would significantly increase the complexity and reduce the readability of the mathematical exposition, which motivated our choice to focus on the two-cluster case for clarity. That being said, we ran new numerical experiments with $K=3$ clusters, 3 heads, and random initializations, which we will add to the revised version.\n\nRegarding (ii): Constraining centroids to lie on the sphere enables the use of linear attention heads in place of softmax-based mechanisms. This assumption effectively normalizes the queries and keys, a role typically played by explicit normalization layers or softmax operations in standard architectures.\n\nRegarding (iii): The attention-based model relies on directional similarity to perform clustering. If the clusters are aligned (or anti-aligned), the model is not able to distinguish them. For intermediate cases (neither orthogonal nor aligned), additional cross terms complicate the analysis. The orthogonality assumption remains a mild condition in high dimensions. We ran additional experiments with randomly initialized centroids on the unit sphere in dimension 10 to support this claim, obtaining similar results to the orthogonal case, and will add these results to the revised version. \n\nIn summary, while these assumptions may seem restrictive, they were chosen deliberately to keep the analysis tractable and transparent. We believe they still capture essential aspects of the mechanism and provide useful theoretical insights.\n\n**W3 (Using attention layers as clustering methods).**\nThis paper studies attention layers when the input data follows a probabilistic model commonly used in clustering contexts. We show that, when attention heads are trained in an unsupervised setting, they learn to encode the parameters of the underlying data distribution—here, the centroids. The transformation performed by the attention layer can further be interpreted as an estimation of the centroid associated with each input token. While this is not clustering in the strict sense (no discrete assignments are made), cluster labels can be derived post hoc. For instance, one may assign token $k$ to cluster $1$ if $X_k^\\top \\hat{\\mu}_1\\geq X_k^\\top \\hat{\\mu}_0$, where $\\hat{\\mu}_0$ and $\\hat{\\mu}_1$ are the learned centroids.\n\nIn the case where the attention parameters converge to the true centroids, one can compute the risk of the resulting clustering rule:\n$\n\\mathbb{P}(X^\\top \\mu_0 < X^\\top \\mu_1 | X \\in \\text{cluster}_1) +\n\\mathbb{P}(X^\\top \\mu_1 < X^\\top \\mu_0 | X \\in \\text{cluster}_0)= \\Phi\\left( \\mu_1^\\top \\kappa / \\sigma \\right) + \\Phi\\left( -\\mu_0^\\top \\kappa / \\sigma \\right)\n$\nwhere $\\kappa = (\\mu_0-\\mu_1)/\\|\\mu_0-\\mu_1\\|$.\nWe will add a discussion on this topic to the next version."}]

Author rebuttal:

Meta-review: This paper provides a theoretical analysis of attention layers as clustering mechanisms under Gaussian mixture models, complemented by a simple regularizer and supporting experiments. While the theory is developed in a simplified setting (two components, orthogonal centroids, linear attention), the work nonetheless offers new insights into the unsupervised capabilities of Transformers. Reviewers agree the paper is well written, technically solid, and addresses concerns adequately in the rebuttal, though they note limitations in generality and scope. Reviewer g8nw maintained a low rating (2), yet after discussion he indicated no objection to acceptance. Overall, the consensus is that the contribution is novel and meaningful despite its restricted setting, and I recommend acceptance.
Decision: Accept (poster)