Attention with Trained Embeddings Provably Selects Important Tokens

We thank the reviewer for the comments and for the positive evaluation of our work. We address all weaknesses and questions below.

W1. Simplified model.

As for the simplified attention model, we remark that dropping the attention weights does not qualitatively change the problem. In fact, the product of the attention weights and $p$ plays the same role as $p$ alone without the attention weights, which is to control the attention probability of each token after softmax. Thus, even if we add attention weights $W_{KQ}$, we still expect all the results for $p$ proved in our work to hold for $W_{KQ} p$. Besides, as discussed in [37, Lemma 1], one can recover the dynamics of $W_{KQ}$ given the dynamics of $p$.

We also note that we have conducted experiments on two-layer networks which show a similar phenomenology as the single-layer model, see Figure 1 of our submission.

W2. Frequent yet irrelevant features; semantic concepts that interact.

The metric we choose to capture the importance of features is equal to the frequency of occurrence with the +1 label minus the frequency of occurrence with the -1 label. This is different from just taking the frequency of the token. In fact, frequent yet irrelevant features will likely appear a similar number of times associated with the +1 label and with the -1 label, and our metric will be small in such cases.

The reviewer makes a good point in mentioning interactions between semantic concepts and we leave a detailed investigation of those to future work. Let us just mention that, while we agree that labels should depend on the interaction of semantic concepts, our results indicate that the first-order statistics already give important information about the labels.

W3. Limited set of synthetic experiments.

During the rebuttal, we have performed additional experiments and ablations:

(1) We have shown that the variance of token embeddings in the real data experiments is reduced by training the context embeddings just for one step (as opposed to training them until convergence). See our response to Q3 or reviewer vwUg for details.

(2) We have shown that a qualitatively similar picture emerges for multi-class classification. See our response to Q2 of reviewer E1M4 for details.

(3) We have analyzed the impact of MLP layers. See our response to Q4 of reviewer E1M4 for details.

(4) We have provided evidence that the trained embeddings improve performance on the classification tasks (as opposed to keeping the embeddings fixed to their initialization). See our response to W1 of reviewer q6jg for details.

(5) We have shown that a qualitatively similar picture emerges for a two-layer model with trained attention weight matrices $W_{KQ}^{(1)}$ and $W_{KQ}^{(2)}$. See our response to Q2 of reviewer q6jg for details.

W4. Only sentiment analysis with simple classification.

As mentioned above, during the rebuttal, we have considered a variety of other settings, including multi-class classification for which we provide both additional experiments and theoretical results. See our response to Q2 of reviewer E1M4 for more details.

As for contrastive learning, we discuss the connection in our response to Q3 below.

Q1. More realistic attention scenarios?

First, we note that adding attention weights $W_{KQ}$ in the second layer does not make qualitative difference, since now $W_{KQ} p$ has the same role as $p,$ and as discussed in [37, Lemma 1], we can recover the dynamics of $W_{KQ}$ given the dynamics of $p.$

To support this view, we have conducted an additional experiment on a two-layer model having attention weights both in the first ($W_{KQ}^{(1)}$) and in the second ($W_{KQ}^{(2)}$) layer. We have found that $E_s^\top v$ and $E_s^\top W_{KQ}^{(2)} p$ have the same behaviour as in Figure 1. See our response to Q2 of reviewer q6jg for the exact data in a tabular form.

Thus, the experiment confirms that dropping the attention weight matrix is a reasonable simplification of the model that does not make a qualitative difference. In the revision, we will add plots for this experiment, as well as a discussion (due to this year’s NeurIPS policy, we are not allowed to post links to plots in our response).

Q2. Practical relevance, e.g. for model diagnostics?

This is an interesting point. The correlations $E_s^\top v, E_s^\top p$ could provide a useful criterion to detect whether the model is performing well. In fact, if such correlations are always small, regardless of the importance of the token, then we expect the model to perform poorly. This is in fact the case if we leave the embeddings untrained, as in an ablation we performed to respond to W1 raised by reviewer q6jg.

Specifically, we have performed an additional experiment providing evidence that the trained embeddings indeed improve performance on the classification tasks. In particular, we have fixed context embeddings (randomly initialized) and we have only trained the CLS token embedding and the classifier weights. The resulting test accuracy is around 72% with test error around 0.6, while for the original baseline the test accuracy is around 95% with test error around 0.1. This clearly indicates the importance of training the embeddings. We will report the result of this experiment in the revision.

Q3. Other tasks, e.g. contrastive learning?

We have provided extra theoretical and experimental results on multi-class classification problems. Please refer to our response to Q2 of reviewer E1M4 for details.

Regarding contrastive learning, while there is some high-level connection, our framework is unlikely to be directly applicable. As an example, consider CLIP. Then, intuitively we expect each image to be relevant only to certain words (tokens) in the text. However, we highlight two technical difficulties.

(1) We need a proper definition of the importance of each token to each image from the dataset.

(2) Due to the unsupervised nature of contrastive learning and the difference in architectures used compared to our case, we need to understand what structure of the embeddings would be crucial to learning the datasets.

Resolving such technical difficulties provides an exciting avenue for future research.

We thank the reviewer for the comments and for pointing out the strengths of our work. The reviewer raises valid points (e.g., on untrained orthogonal token embeddings and on adding $W_{KQ}$ matrices in the model) that we addressed via additional experiments. We now elaborate on these points, addressing all questions and concerns below.

W1. Untrained embeddings.

We have performed an additional experiment providing evidence that the trained embeddings indeed improve performance on the classification tasks. In particular, we have fixed context embeddings (randomly initialized) and we have only trained the CLS token embedding and the classifier weights. The resulting test accuracy is around 72% with test error around 0.6, while for the original baseline the test accuracy is around 95% with test error around 0.1. This clearly indicates the importance of training the embeddings. We will report the result of this experiment in the revision.

W2. Assumption 1.

If the sequence contains multiple positive or negative tokens, the behaviour is expected to depend on the specific structure of the dataset. The difficulty is that the model may not necessarily select all relevant tokens. In fact, only selecting one relevant token from each sequence is in principle sufficient, and the negative directional gradient of $p$ along the max-margin direction that selects all the relevant tokens may not be larger than that of other selections. Nevertheless, two results follow from our analysis:

(1) If $p$ converges in direction, it must select at least one relevant token for each sequence, simply because the loss is non-increasing.

(2) If $p$ converges in direction and selects all relevant tokens, $p$ must be the max-margin direction of such selection. This can be obtained by following the approach obtained to derive our Theorem 4.3, which compares the directional gradient of $p$ along the max-margin direction and non-max-margin direction of the same selection.

The case with multiple partially positive or negative tokens is even harder, and it is likely to require additional assumptions on the training set, as well as completely different techniques.

We finally note that assuming only one relevant token is common in the literature [A, B].

[A] Attention layers provably solve single-location regression. ICLR, 2025.

[B] Max-margin token selection in attention mechanism. NeurIPS, 2023.

W3. Novelty compared to [1].

Let us elaborate on the technical novelty compared to [1]:

(1) Our results work for more general settings as compared to [1]. The directional convergence results in [1] require either only one sequence in the dataset [1, Theorem 2], or an initialization that is close enough to certain directions [1, Theorem 3]. In contrast, Theorem 4.3 in our work requires none of those restrictions to characterize all possible directions the gradient flow converges to.

(2) Our results require fundamentally different techniques as compared to [1]. Specifically, [1, Theorem 3] suggests that, for any locally optimal max-margin direction, we can find a regime such that the direction of $p$ will converge to that direction. Note that a locally optimal max-margin direction in [1, Definition 2] does not necessarily select all the relevant tokens. Our results, on the contrary, show that $p$ must necessarily converge to a max-margin direction that selects all relevant tokens. Thus, our focus is on proving the impossibility of converging to “bad” directions under random initialization and one-step of gradient descent on the embeddings, while [1] focuses on proving the convergence under a close-enough initialization to the target direction.

(3) [1] does not consider any connections between the directional convergence and the importance of the tokens, which, in contrast, is our main focus.

Q1. $v$ and $E_X$ trained simultaneously?

If we train $v$ but fix the context embeddings, the implicit bias of $p$ is not affected in the limit. To see this, we compute

$-\nabla_v \mathcal L_n(p,v) = \hat{\mathbb E}[ y g(X,y) \sum_{i=1}^T q_i(X) E_{x_i} ] \approx \hat{\mathbb E}[ y g(X,y) \sum_{x_i \in \mathcal S_X(p)} q_i(X) E_{x_i} ].$

Note that, if $p$ converges in direction, it must select all relevant tokens. This means that $y E_s^\top v$ will be increasing if $s$ is relevant, and $E_s^\top v$ will not move a lot if $s$ is irrelevant, since it appears exactly the same number of times in positive and negative sequences.

While we do not have a rigorous theoretical characterization for training all parameters $E,p,v$ together, in our experiments of Figures 1-3 the token embeddings $E$ are trained simultaneously with $p.$ Thus, we expect the token selection to be the same as if we only trained $p$.

Q2. Only $p$ trained; adding $W_{KQ}^{(1)}$ and $W_{KQ}^{(2)}$.

First, we note that in the two-layer model experiments, $p,v, E_x$ are trained simultaneously.

Second, we note that adding $W_{KQ}^{(2)}$ in the second layer does not make a qualitative difference, since now $W_{KQ}^{(2)} p$ plays the same role as $p$ and, as discussed in [37, Lemma 1], we can recover the dynamics of $W_{KQ}^{(2)}$ given the dynamics of $p.$

Finally, we have conducted an additional experiment on a two-layer model with attention weights both in the first ($W_{KQ}^{(1)}$) and the second ($W_{KQ}^{(2)}$) layer. We have found that $E_s^\top v$ and $E_s^\top W_{KQ}^{(2)} p$ have the same behaviour as in Figure 1. The results are reported in the table below.

The experiment confirms that dropping the attention weight matrix is a reasonable simplification of the model that does not make a qualitative difference. In the revision, we will add plots for this experiment, as well as a discussion (due to this year’s NeurIPS policy, we are not allowed to post links to plots in our response).

We thank the reviewer for the comments and for appreciating the perspective brought forward by our work. The reviewer raises valid points (e.g., on multi-class classification and the addition of MLP layers) that we addressed via additional experiments. We also clarified our assumption on the embedding dimension $d$. We now elaborate on these points, addressing all questions and concerns below.

W1. Strong Assumptions.

While we agree that labels depend on token combinations, we argue that our setting, albeit simplified, is a first step towards a full understanding of the structure of trained embeddings. We bring the following evidence in favor of this argument:

(1) The setting with a single relevant token was also considered in previous work [1,2].

(2) Generalization to multi-class classification is rather direct and presents a similar phenomenology. See our response to Q2 below for details.

(3) Fixing the context embeddings and only training the CLS embedding as well as the output vector significantly reduces the variance in our scatterplots, indicating that the outlier points mentioned by the reviewer are likely due to the longer training of the embeddings. See our response to Q3 of Reviewer vwUg for details.

W2. Only focus on first-order statistics.

While we agree that labels should depend on higher-order relations, our results indicate that the first-order statistics still give important information about the labels. In particular, in the experiments, we see that tokens with different average signed frequency play a different role in the classification (tokens with large $|E_s^\top v|$ and $E_s^\top p$ play a more important role). In addition, the connections between token embeddings and first-order statistics were not revealed in previous works to our best knowledge.

W3. Generalization concern.

We will discuss the role of MLPs, see our response to Q4 below.

Q1. Why binary classification?

In binary classification, it is more intuitive to define the importance of the token w.r.t the labels (i.e., the average signed frequency we defined in the paper); in contrast, in other settings, e.g. generative tasks, quantifying such importance is less obvious.

As mentioned in the introduction, we first conducted experiments on the sentiment dataset, and then we provided a theoretical analysis to understand the experimental phenomenon, rather than picking a-posteriori a setting that aligns with the theoretical analysis.

Finally, we note that binary classification problems have been widely considered in prior work, see e.g. [2, 3].

Q2. Multi-class classification?

We provide below both theoretical and experimental evidence that our conclusions for binary classification extend to the multi-class setting.

Theoretical evidence. Consider $K$-class classification, let $v_k$ be the $k$-th classifier and perform one step of gradient descent on the token embeddings. Then, after some manipulations, we obtain that $E_s^\top v_k \approx \frac{\eta_0}{T} \hat{\mathbb E}\left[ \sum_{i=1}^T \frac{K-1}{K} 1_{y = e_k, x_i = s } - \frac{1}{K} 1_{y \neq e_k, x_i = s }\right].$

The RHS of the equation above can be viewed as a generalization of the average signed frequency that indicates the importance of a token with respect to class $k.$ In particular, if we divide the RHS by a factor of $K-1,$ the first term becomes the co-occurrence of token $s$ with label $k,$ and the second term is the average co-occurrence of token $s$ with other labels. Thus, intuitively, if a token occurs much more frequently with label $k$ than other labels, this token is more important w.r.t. label $k,$ and the correlation with $v_k$ is larger.

Experimental evidence. We conduct experiments on the Yelp datasets with 4 classes, and we find that $E_s^\top v_k$ is still increasing with respect to the importance of the token for each class $k$, as for binary classification. We display the results for class 2 and 4 in the tables below. For class 1 and 3, the results are similar.

In the revision, we will add the derivation of the formula above, plots for the experiments, as well as a discussion (due to this year’s NeurIPS policy, we are not allowed to update the PDF or post links to plots in our response).

Q3. Embedding dimension.

We highlight that the vocabulary in our paper only includes the tokens that appear in the training set, rather than the whole vocabulary of the tokenizer. We note that the embedding dimension $d$ grows roughly linearly in the vocabulary size in the numerical simulations: we pick embedding dimension 2048 (as common in practical models), and the vocabulary size is 7766 for IMDB and 5424 for Yelp.

Q4. Role of MLP.

We have performed additional experiments to study the impact of MLP layers. In particular, we have added one MLP layer after the softmax, with the same hidden dimension as the embedding. We find that $E_s^\top p$ is still increasing with $\alpha_s$, while the monotonicity of $E_s^\top v$ is less evident. This indicates that the model will still select important tokens using $E_s^\top p$, but the interaction between the embeddings and the output vectors is more complicated due to the presence of the MLP layer.

In the revision, we will add plots for the experiments, as well as a discussion (due to this year’s NeurIPS policy, we are not allowed to post links to plots in our response).

Q5. Vision tasks.

The embedding layer of a typical vision transformer (ViT) is a linear transformation over the patches, and we regard the frequency of patches as not very meaningful due to the continuous nature of the vision data. Since our results require a finite vocabulary, the study of vision tasks would require a different definition of the importance of patches, which we leave to future work.

References

[1] Attention layers provably solve single-location regression. ICLR, 2025.

[2] Max-margin token selection in attention mechanism. NeurIPS, 2023.

[3] Implicit bias and fast convergence rates for self-attention. TMLR, 2025.

We thank the reviewer for the comments. We appreciate that the reviewer found our setup to be clean and, while we acknowledge the simplifications it brings, we highlight that it still displays the same phenomenology as more realistic setups, with the important advantage of being mathematically tractable. We have also discussed in detail differences with [40] and performed the additional experiment suggested in the last question finding a conclusion aligned with the intuition of the reviewer. We now elaborate on these points, addressing all questions and concerns below.

W1. Simplified model.

The problem we address in our work is to characterize the structure learnt by embeddings during gradient descent training. Dropping the attention weights does not qualitatively change this problem. In fact, the product of attention weights and $p$ plays the same role as $p$ alone without attention weights, which is to control the attention probability of each token after softmax. Thus, even if we add attention weights $W_{KQ}$, we still expect all the results for $p$ proved in our work to hold for $W_{KQ} p$. Besides, as discussed in [37, Lemma 1], one can recover the dynamics of $W_{KQ}$ given the dynamics of $p$.

W2. Comparison with [40].

There are two key differences between our setup and that of [40]:

(1) [40, Assumption 1] requires that all the tokens in non-relevant positions are nearly orthogonal to each other, which is crucial to their proof. In the notation of our paper, they require that, for any $X, X’,$ $E_{x_i}, E_{x’_j}$ are almost orthogonal to each other, for any $x_i, x_j’$ that are not relevant. This is not true in our case, since all the irrelevant tokens appear at least in two sequences (otherwise, the token would have non-zero average signed frequency). Specifically, for any irrelevant token $s,$ there exist $X,X’$ such that $x_i=x_j’=s.$ Since $s$ has non-zero embedding, [40, Assumption 1] cannot be true in our case, and the global convergence result cannot be applied.

(2) [40, Assumption 2] requires all the irrelevant tokens in the sequence to have exactly the same score $\gamma_s = y E_s^\top v.$ This assumption leads to the fact that for each sequence, only the term that involves the relevant tokens contributes to the directional gradient, see [40, Equation (5)]. Our setting does not satisfy this extra assumption. Even though the score differences between different irrelevant tokens are of order $O(1/\sqrt{d}),$ these errors matter when the gradient norm is diverging to infinity, and the analysis in [40] cannot be applied.

Thus, our focus is on proving the impossibility of converging to “bad” directions under random initialization and one-step of gradient descent on the embeddings. In contrast, [40] focuses on proving the convergence under extra assumptions on the orthogonality of embeddings and scores, which is a fundamentally different problem.

Q1. Why fix token embeddings?

Fixing the context token embeddings is indeed for technical simplification. However, we note that the behavior described by the theory is well aligned with our numerical experiments in which all token embeddings are trained until convergence, see Figures 1-3. Thus, we do expect that further training on the embeddings does not significantly change the structure learnt via gradient descent.

Q2. Clarification on synthetic data generation.

The synthetic data is generated exactly as in Equation (15), and the datasets in general do not satisfy Assumption 1. Thus, our results in the synthetic experiments hold for a more general setting. To recover a setting that satisfies Assumption 1, we can simply set $\tilde \delta = (|\mathcal S|-1)/|\mathcal S|,$ then with high probability the dataset satisfies Assumption 1. Overall, the goal of the synthetic experiments is to illustrate the connections between the importance of the tokens and the structure of the embeddings, which motivates our work. We will clarify these points in the revision.

Q3. High variance in IMDB and Yelp datasets.

This is an interesting question, and the reviewer has the right intuition. We have performed additional experiments, where we train one full-gradient step on the embeddings (as in the theory) and we then fix the context embeddings only training the CLS embedding as well as the output vector. As a result, the average standard deviation over the bins is reduced by roughly 8 times for the inner product with $v$ and by roughly 6 times for the inner product with $p$, compared to training all context embeddings until convergence. This gives evidence that, by training the context embeddings only for one step, the variance is indeed significantly lower.

We expect the large variance to be due to both the longer training of embeddings and the dataset itself. The dependence on the longer training of the embeddings is indicated in the experiments mentioned above, and the dependence on the dataset is apparent by comparing Figure 2 (synthetic datasets, small variance) with Figures 1 and 3 (real datasets, large variance).

In the revision, we will add plots for the experiments mentioned above (due to this year’s NeurIPS policy, we are not allowed to post a link to a plot in our response), as well as a discussion.