Title: Attention with Trained Embeddings Provably Selects Important Tokens
Venue: NeurIPS 2025
Accepted: yes
Keywords: attention layer

token embedding

gradient flow

max-margin
Abstract: Token embeddings play a crucial role in language modeling but, despite this practical relevance, their theoretical understanding is limited. Our paper addresses the gap by characterizing the structure of embeddings obtained via gradient descent. Specifically, we consider a one-layer softmax attention model with a linear head for binary classification, i.e., $\mathrm{Softmax}( p^\top E_X^\top ) E_X v =  \frac{ \sum_{i=1}^T \exp(p^\top E_{x_i})  E_{x_i}^\top v}{\sum_{j=1}^T \exp(p^\top E_{x_{j}}) }$, where $E_X = [ E_{x_1} , \dots, E_{x_T} ]^\top$ contains the embeddings of the input sequence, $p$ is the embedding of the $\mathrm{\langle cls \rangle}$ token and $v$ the output vector. First, we show that, already after a single step of gradient training with the standard logistic loss, the embeddings $E_X$ capture the importance of tokens in the dataset by aligning with the output vector $v$ proportionally to the corresponding average signed frequency that captures the relevance of tokens to the labels. Then, after training $p$ via gradient flow until convergence, the softmax selects the important tokens in the sentence (i.e., those that are predictive of the label), and the resulting $\mathrm{\langle cls \rangle}$ embedding maximizes the margin for such a selection. Experiments on real-world datasets (IMDB, Yelp) exhibit a phenomenology close to that unveiled by our theory.


Peer reviews:
comments: SUMMARY:
The paper studies the effect of gradient training on the encoding of text tokens in a simplified attention mechanism. They show that even a single gradient step aligns the feature weights such that high frequency words get up weighted.

STRENGTHS AND WEAKNESSES:
Strengths:
- Paper provides theoretical analysis of a simplified scenario for how attention mechanisms learn to adapt to token frequency.
- Theoretical claims are backed-up by experimental evidence.
- The mathemical analysis seems solid, although I could not follow all steps.

Weaknesses:
- Only a simplified model of attention is studied and only a single layer. 
- Importance of features is equalized (at least in writing) with importance - but there maybe frequent yet irrelevant features. Also, there may be other semantic concepts that interact, e.g. negating concepts. 
- Only limited set of synthetic experiments is performed.
- Only sentiment analysis with a simple classification setting is studied, what about other tasks like contrastive learning.

QUESTIONS:
- How can the proposed method be extended to more relatistc attention scenarios?
- Can the authors speculate about the practical relevance of their findings, e.g. can this be used for model diagnostics?
- Can other tasks then sentiment analysis be studied with the same framework, like contrastive learning?

LIMITATIONS:
Brief, but ok
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper shows the role of token embeddings. It studies how the trained token embeddings capture the importance of tokens and selects the relevant tokens. The authors characterize the token embeddings after one gradient step and the implicit bias of \\<cls\\> embedding. In addition, they conduct synthetic and realistic experiments to support the theoretical findings.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The topic is interesting and may be the first theoretical study of the embedding training dynamics.\n2. The paper is well-written and easy to follow.\n3. The mathematical analysis and proofs are solid and clear.\n\nWeaknesses:\n\n1. The paper demonstrates that the trained token embeddings help select the relevant tokens by capturing the importance of different tokens. However, it does not provide evidence that the trained embeddings improve performance on the classification tasks. It is possible that even untrained orthogonal token embeddings could achieve similar classification accuracy, suggesting that the training may not be essential for the task.\n2. Assumption 1 appears quite strong, as it assumes that each sequence contains only one relevant token, which must be completely positive or completely negative. What if a sequence contains multiple relevant tokens, or if some relevant tokens are only partially positive or negative? Would the model still be able to select all relevant tokens under these conditions?\n3. The output vector $v$ is fixed, and the token embeddings are trained for only one step. It appears somewhat simplistic. It involves fixing all other parameters and only training the embedding $p$, which does not seem to show much novelty compared to [1].\n\n[1] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. NeurIPS, 2023.\n\nQUESTIONS:\n1. What may happen if the output vector $v$ and the token embedding $E_X$ are trained simultaneously?\n2. In line 275, the two-layer model has three parameter matrices $p,v,W_X$, with only $p$ being trained. What would happen if there are two trained matrices $W_{KQ}^{(1)}$ and $W_{KQ}^{(2)}$ in the two layers?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper theoretically and empirically demonstrates how token embeddings capture the correlations between tokens and labels during training, using a simple one-layer attention model and a binary classification task. Specifically, the authors first prove that after a single gradient step with the standard logistic loss, the token embeddings begin to encode the empirical importance of each token by aligning with the output vector in proportion to the corresponding empirical frequencies. They then show that, under the assumption that each sequence contains exactly one completely positive or completely negative token and all remaining tokens are irrelevant, gradient flow converges in direction to a classifier embedding that corresponds to the max-margin separator, which maximally distinguishes relevant tokens from irrelevant ones. Finally, the paper verifies these theoretical results with numerical experiments on synthetic and IMDB/Yelp datasets, demonstrating that the learned embeddings indeed reflect the empirical statistics of token-label associations for both one-layer and two-layer attention models.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengthes:**\n\n**(1):** This work offers a new perspective on understanding the function and properties of token embeddings by analyzing their behavior during gradient descent through both clear and solid theoretical analysis and numerical experiments. It reveals that token embeddings capture the empirical importance of tokens in the dataset, while the classifier embedding converges to the direction corresponding to max-margin solutions that optimally separate relevant tokens from irrelevant ones. These findings help bridge our understanding of how token embeddings in large language models evolve during training and how they contribute to the model’s predictions.\n \n**(2):** The numerical analyses strongly support the theoretical claims. In both the synthetic data experiments and the real-world IMDB/Yelp datasets, it is clearly demonstrated that training a simple attention classifier leads to embeddings that effectively capture how frequently and strongly each token is associated with the label, consistent with the theoretical predictions.\n\n**(3):** The writing is overall clear and easy to follow. The paper is well-organized, and the theoretical claims and proofs presented in the main text are sound and do not exhibit major issues.\n\n**Weaknesses:**\n\n**(1) Strong Assumptions:** The paper assumes that each sequence in the dataset contains exactly one completely positive token or one completely negative token, with all remaining tokens treated as irrelevant. However, in reality, classification labels typically depend on the relationships among combinations of tokens. Under these conditions, the proposed theoretical results do not straightforwardly generalize to a wide range of tasks or more realistic language models. Even for the binary classification tasks studied in this paper, Figures 1 and 3 reveal many outlier points where token embeddings deviate from the predicted trends. Therefore, additional analysis is needed to understand and bridge the gap caused by such outliers.\n \n**(2): Only focus on first-order statistics** As indicates by the paper, only the marginal frequency with which each token appears in positive versus negative examples is considered. However, in reality, higher-order relationships are common in natural language processing for the exiting language models and a purely frequency-based perspective cannot fully explain how modern language models learn rich representations from data.\n\n**(3): Generalization concern:** The paper focuses primarily on a one-layer model and a simple two-layer extension, while modern language models employ deep architectures with multiple layers and attention heads. Bridging the gap between the theory developed here and the behavior of large-scale models remains non-trivial, as suggested by the more variable empirical results. Additionally, this work does not explore the role of feedforward networks (MLPs) in shaping token embeddings, which may further limit the conclusions drawn solely from attention mechanisms.\n\nQUESTIONS:\n(1) Why do the authors choose to use binary classification tasks as the illustrative tasks? Is it because it is simple for analysis and the results align with the sentiment dataset?\n\n(2) If yes to Q1, how to make sure that the conclusions drawn form the binary-classification can extend to other scenarios like multi-classification tasks? Is there any empirical or theoretical thinkings about it?\n\n(3) The theoretical results require the embedding dimension $d$ to grow at least logarithmically in vocabulary size, which can lead to very large $d$ for realistic datasets. I am wondering how sensitive are your conclusions to smaller embedding dimensions commonly used in practice? \n\n(4) The analysis omits the influence of MLP on token embedding training. Could the author please explain what will the attention with MLP impact the conclusions brought by the attention-only structures? \n\n(5) The study mainly focus on the language tasks; however, as transformer or attention mechanism is also widely in vision tasks, could the authors please elaborate how the conclusions brought by this paper can be extended to the vision tasks?\n\nLIMITATIONS:\nPlease see above weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies a simplified one-layer attention model in a binary sequence classification setting, where the only trainable parameters are the input token embeddings. The authors analyze the training dynamics and show that, after one gradient step, the embeddings begin to align with the fixed output vector, with their magnitudes scaled by the *signed* frequency of each token. This effectively encodes first-order statistics of the input into the embeddings.\n\nNext, with the input embeddings frozen after the first step, the authors characterize the implicit bias of the `<cls>` token embeddings, showing it converges to the max-margin solution that separates positive/negative tokens (tokens that appear more frequently in sequences from one of the classes) from the irrelevant ones (those equally likely to appear in the sequences of any class). This result relies on the assumption that each sequence contains exactly one relevant token that determines the label.\n\nThe authors support their theoretical findings by experiments on both synthetic datasets and sentiment classification tasks (IMDB and Yelp). In both settings, they observe that, on average, tokens with higher signed frequency receive higher attention (measured by inner product of their embeddings with the `<cls>` token) and align more strongly with the output vector.\n\nSTRENGTHS AND WEAKNESSES:\nI found the setup for analyzing how first-order statistics get encoded into the input embeddings clean, and the experiments support this intuition to some extent. \n\nThe model, however, is simplified. The attention weights are dropped entirely, and attention effectively reduces to computing directly the inner products between the token embeddings and the `<cls>` token. \n\nThe max-margin convergence result overlaps with prior work, as also cited in the paper, e.g., [32, 37, 40], which also considers similar synthetic setups. The comparison with [37] is helpful, especially in clarifying the distinction between local and global convergence. However, [40] also presents global convergence analysis in a related setting. If there are meaningful differences in results, techniques, or assumptions compared to [40], it would be helpful to include a discussion.\n\nQUESTIONS:\n1. What’s the motivation for fixing the token embeddings $E$ and only training the `<cls>` token? Is this just for simplification, or is there an argument that further training on $E$ does not provide more insights?\n\n2. In the synthetic data generation, are tokens sampled independently as in equation (15)? If so, does the generated data actually satisfy Assumption 1? Or are the results in the synthetic experiments for a more general setting?\n\n3. I understand that the IMDB and Yelp datasets differ substantially from the synthetic setup, but the variance in token embeddings seems quite high. Since the one-step result holds for arbitrary data, would training for just one step lead to lower variance and better alignment with the theoretical predictions? In other words, is the large variance an artifact of longer training or more of other dataset-specific factors?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-22"}