Attention-based clustering

We would like to thank you for your feedback. In the following, we provide answers to the weaknesses you underlined.

W1 (Illustration). We 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.

W2 (Towards more complicated scenarios). Our theoretical analysis is conducted under three main assumptions: (i) a two-component Gaussian mixture model, (ii) centroids constrained to lie on the unit sphere, and (iii) mutual orthogonality of these centroids.

Regarding (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.

Regarding (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.

Regarding (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.

In 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.

W3 (Using attention layers as clustering methods). This 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.

In the case where the attention parameters converge to the true centroids, one can compute the risk of the resulting clustering rule: $\mathbb{P}(X^\top \mu_0 < X^\top \mu_1 | X \in \text{cluster}_1) + \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)$ where $\kappa = (\mu_0-\mu_1)/\|\mu_0-\mu_1\|$. We will add a discussion on this topic to the next version.

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. More specifically, there is generally a tradeoff in theoretical works between analytical tractability and resemblance to practice. Despite 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.

More detailed responses to your questions follow below.

Q1. Our theoretical analysis is conducted under three main assumptions: (i) a two-component Gaussian mixture model, (ii) centroids constrained to lie on the unit sphere, and (iii) mutual orthogonality of these centroids.

Regarding (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.

Regarding (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.

Regarding (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.

In 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.

Using 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.

We thank you for having raised these concerns, and will discuss them in the revised version.

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. In 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.

[1] Trained Transformers Learn Linear Models In-Context, Zhang, Frei, Bartlett, JMLR 24.

[2] Linear attention is (maybe) all you need (to understand transformer optimization), Ahn et al., ICLR 24.

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.

Regarding 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.

Finally, 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, $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.$

Interestingly, 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.

We thank you for your feedback and for acknowledging the relevance of our work.

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. To address this limitation in practice, we introduce a regularization term that empirically promotes global convergence.

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.

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.

It 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: $T_\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$ where $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.

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.

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, $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.$

Interestingly, 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.

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.

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. 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.

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.

That 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.

We thank you for prompting us to formalize this comparison, and we will include this discussion in the revised version of the paper.

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. It 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.

The formal analysis in this more general case is an interesting direction for future work.

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.