Attention with Trained Embeddings Provably Selects Important Tokens
Authors: Diyuan Wu
Aleksandr Shevchenko
Samet Oymak
Marco Mondelli
Venue: NeurIPS 2025 poster ()
Keywords: attention layer
token embedding
gradient flow
max-margin
Primary area: deep_learning
TL;DR: We characterize the structure of embeddings obtained via gradient descent, showing that the attention mechanism provably selects important tokens.

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.

Attention with Trained Embeddings Provably Selects Important Tokens
Diyuan Wul Aleksandr Shevchenko
Samet Oymak
Marco Mondelli
ISTA ETH Zürich2
University of Michigan
Abstract
Token embeddings play crucial role in language modeling but, despite this practi- cal relevance, their theoretical understanding remains 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., Softmax(p' Ex) Σ. exp(p" Ez, FF where
Σ exp(p" Esy)
Ex = [Eg..... Egy! T contains the embeddings of the input sequence, p is the embedding of the (cls) token and the output vector. First, show that, already after a single step of gradient training with the logistic loss, the embeddings Ex capture the importance of tokens in the dataset by aligning with the output vector v proportionally to the frequency with which the corresponding tokens appear in the dataset. 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 (cls) embedding maximizes the margin for such a selection. Experiments on real-world datasets (IMDB, Yelp) exhibit a phenomenology close that unveiled by our theory.
1 Introduction
The introduction of the attention mechanism [5 marked a paradigm shift in the design of frontier machine learning models, leading to significant advances such as ChatGPT [2]. Claude [3), AlphaFold 18], CLIP [30) and Dall-E [1I. This success prompted a surge of interest understanding the structure and function of attention layers, with their optimization dynamies and inductive biases being object of extensive theoretical research 7
(see also Section 2 Embeddings crucial component of the attention mechanism 451 especially for downstream adaptation
7
10
23
37
with some works 220, 45 specifically highlighting their importance. However, despite the importance of learning embeddings, the existing analyses of transformer-like architectures either ignore the properties of embeddings by resorting to orthogonal structures 44 or omit embeddings completely by considering unprocessed inputs (39].
13 16
Our paper fills this gap by studying directly the embedding training dynamics. Specifically, we aim to provide theoretical insight to the following questions:
What is the structure learnt by the embeddings during gradient descent training?
How is this structure related to the statistical properties of the data?
In Figure I we investigate these questions by analyzing the embeddings of two-layer transformer trained on sentiment analysis task on IMDB and Yelp reviews. The plots reveal a remarkable simplicity in the structure of the learned embeddings, which capture the frequency of appearance of tokens in the dataset. Specifically, the predictive mechanism (overlap with the regression coefficient
Equul contribution. Corresponding author: diyuan.uu@ist.... ww@ist.ac.at
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

10
20
Mean
10
10
(Ex,
-10
-10
-20
Mean
D.
P11 Ex.) P(O(E,.)
P(] P(O(Ex))
10
Takens
Mean
20
10
10
(Ex,
-10"
Tokens
-20
-10
-30
-1.0
-0.5
0.0
0.5
-1.0
-0.5 P(IJE,) P(OJE,)
0.0
P(IEE,)) (EIIE,)-P(OPES P(O)E,))
Figure 1: Dot-product of token embeddings with (cls) embedding p (left) and regression coefficients (right), as function of the token-wise difference in posterior probabilities, for IMDB (top row) and Yelp (bottom row) datasets. We consider the two-layer attention model in (18) with all parameters trained until convergence.
v) favors the tokens which appear more frequently in the corresponding positive/negative context. A similar pattern emerges at the selection stage of the attention mechanism (overlap with the (cls) embedding p), i.e., more frequent tokens have a higher attention score.
For the theoretical study this emergent structure, we focus on a one-layer softmax attention model. Namely, for an input sequence X [x1, ... ,rT], the output of the model is given by
f(X:p. E) = Softmax Ex) Σt exp(p' Ez. ET
Σ/=1 cxp(p (E))
where Ex [E2) E contains the embeddings of the input X, p is the embedding of the (cls) token and is the final regression vector. Our main results are summarized below:
We show that, already after a single step of gradient training with the standard logistic loss, the embeddings Ex capture the importance of tokens in the dataset by aligning with the output vector v proportionally to the corresponding empirical frequencies (Lemma 4.1
setting where each sequence contains single important token, the (cls) embedding obtained from gradient flow must select all important tokens. We further characterize all the possible directions that the (cls) embedding may converge to, which are the max-margin solutions associated to feasible token selections (Theorem 4.3
While in general the (cls) embedding may select irrelevant tokens, we identify sufficient conditions leading the selection only of important tokens (Lemmas 4.4 and 4.5
2 Related work
Implicit bias, margin maximization, attention. The implicit bias literature has been instrumental understanding the behavior of neural networks or overparameterized models optimized by gradient A key phenomenon is that gradient descent on separable data with logistic loss directionally converges the max-margin separator (35] More recently, a series of works
methods g
4
26
36
22
40
34
37 32 has established an equivalence between the optimization geometry

---

Empirical statisties of each token in the dataset. The goal of the paper is to characterize the structure of the embeddings E,p p obtained by optimizing the objective (2) via gradient descent, and we show that such structure related to the empirical statisties of the tokens in the dataset. of self- celf-attention and a hard-margin SVM problem selecting subset of tokens via linear constraints on the outer- suter-products of token pairs. Compared these works that mostly focus on the training of single-layer attention weights, we point out two differences. First, we study the role of embeddings and their joint training with the (cls) token. Second, under our data model, we establish benign properties of the solution reached at convergence (which may not hold for arbitrary datasets (37))
has explored whether attention architectures can extract causal structure from Markovian inputs. The mechanics of next-token prediction when training a single self-attention layer is characterized in [21]. Towards understanding how to utilize structural properties of the data, the behavior of transformers
Theory of attention. line of work 23 24
sparse token selection tasks is The study [14] provides a theoretical justification to the tendency of modern language models to generate repetitive text by showing that the underlying self-attention mechanism collapses into sampling only a limited subset of tokens. This stands in contrast to the slightly strategy. The works [10. 9 take a mean-fiell view to analyze the clustering behavior in transformer representations that emerges after successive applications of the attention block. Under random feature design, it is shown in 66 that softmax attention exhibits a sensitivity property which allows for sharp change in attention scores given the perturbation of single token. The role of the attention mechanism is also studied in (29 for prompt-tuning and in 12 for pest-time-trainnng
considered [3 43
different setup of (38 where the transformer model does not degrade to a twinner-takes-al
3 Problem setup
Data model. We focus on binary text classification problems. We consider a (context) vocabulary set S with size [S], together with a (cls) token for classification. Let (X,, )²= be the dataset containing context sequences, where Yi E (-1,1) and each context sequence X E Xr == (XI..... X,, contains T' tokens, i.e., X = [x1,
ØT| with z; E S. Without loss of generality, we let be the set of tokens that appears in Xr,, the embeddings of the remaining tokens are not trained and are not relevant for the problem at hand.
Architecture. We consider one-layer softmax attention model with a linear head for classification. First, we append a (cls) token at the end of the sequence X, and then we embed each token into a vector of dimension d. Namely, after the embedding layer, we have Ex = [Ex,..... Egy IT Txd where E. E Rd denotes the embedding of the token s. We let E R|Sixd be the embedding matrix of all context tokens and p E Rd the embedding of the (cls) token.
We focus on the architecture defined in I1 where, given vector RT, (aaft expla.)
T}. tice, it is common to include the WKQ matrix and consider a model with output f(X; WKQ: E) = Softmax(p T Wxq EX) Exv. Since p WKQ plays the same role as and one can easily reconstruct WkQ from p in each gradient update as discussed in (37, we the model in (T) for simplicity. We also note that our experiments on two-layer networks with a W WxQ matrix reported in Appendix F4 do not show any qualitatively different behavior, which further justifies our simplification.
Ej= for i E (L....,T), We remark that the same model is also studied in 33 32). In prac-
Optimization problem. The output vector is fixed and all the embedding vectors p, E are trained with the standard logistic loss:
C(E. p)
log(1 +
E,p))) = [lo(1 + exp(-ys(X E. E.))
(2)
n
where the notation È is a shorthand for the average over the dataset D = {(X,, yk) 12
Specifically, after training, the softmax attention learns to select tokens that are more correlated to the labels based on the dataset. To quantify the correlation between a token s and the label y, we define the average signed frequency of a token as:
T
0g =
nT
1
Σ
7
1
y
2
(3)
(XD)ED
1=
3

---

In words, a, is obtained by taking the number of occurrences of s in sequences with positive label, subtracting the number of occurrences of s sequences with negative label, and finally dividing by the total number of tokens n7'' As such, it provides an average of the signed frequency of where the sign comes from the label of the sequences which the token appears.
Definition 3.1 (Positive, negative and irrelevant tokens). We say that a token is (i) positive if o, > 0, (ï) negative if Os < O, and (iii) irrelevant if a = 0. Moreover, we say that a token is completely positive if it appears only in sequences with label 1, and completely negative if it appears only in sequences with label 11
In words, token is positive (negative) when it is more frequently associated to the positive (negative) label; tokens that appear the same number of times associated to positive and negative labels are irrelevant. The quantity as quantifies how positive/negative token is. Intuitively, if either a,
0 or < O₂' < 0 then the token s is more relevant than the token s' for the classification task and, therefore, we expect that this will be reflected into the structure of the corresponding embeddings.
4 Main theoretical results
We show that the trained attention model statistics of the data. First, in Section 4.1 we prove that training the context embeddings with single gradient step suffices to capture the empirical importance of tokens in the dataset; then, in Section 4.2 characterize the implicit bias of training the (cls) embedding p until convergence, having fixed the context embeddings after the first gradient step.
[ learns to select important tokens from the empirical
4.1
One step of gradient descent learns the importance of the tokens
We start by showing that the first step of gradient descent is already enough to give correlation between embeddings and the output vector v. Furthermore, this correlation is proportional to the
average signed frequency defined in (3. We initialize with any unit-norm vector and E p lid. ((0, 7) for all & E S. Then, we perform one step gradient descent with step size Ylu on all trainable embeddings:
p' p'=p' 110 V ,P),
E
110 V A(EE,(P)) ,P), for all S.
(4)
Lemma 4.1. For any 5 0, let
|S|2 2
d 2 max 256,
2 log 5
(5)
Then, after the first step of gradient descent in we have that, for any E
E = E 0.2 + err..
170
p'= + errp
(6)
2
where the error terms err,, errp are bounded with probability at least 8 as
max ||erralla |lerr, 117od
SES
Lemma 4.1 implies that after one step of training, the embedding vector E, of each token learns the empirical importance of the tokens by adding a vector in the direction of the output vector e with magnitude proportional to Os.-
The proof follows from the structure of the gradient update. In particular, it can be shown that
T
T
Vg, L(E pp) - yo(X,v)
Σ
(1=:-
-1,,,) up
p. 1.
1=1
jpt
(8)
with 9: 1= exp(1x TE )
Note that, for all xi. (T is of order 1/Vd, due to the independent Gaussian initialization. Thus, the characterization in S implies that the gradient roughly 382 a plus term that is vanishing in d. The full argument deferred to Appendix C.I
ΣJ_ exp (T2)TE

---

4.2 Gradient flow on performs max-margin token selection
Lemma 4.1 shows the informative overlap between the output vector and the context embedding vectors E, after the first step of gradient descent] However, (6) also implies that the overlap between the {cls) embedding vector p and Eg does not improve after the first step. Thus, next, we study the training dynamics of characterizing its implicit bias. Specifically, we fix the context embedding matrix to E' (obtained after the first gradient step) and train the (cls) embedding vector with gradient flow initialized at p' (obtained after the first gradient step);
dt Pt - V Pc).
We consider gradient flow for technical convenience, and all results in this section can be readily extended to gradient descent with small enough step size. For the rest of the section and in the related proofs appearing in the appendix, we will refer the embeddings in E as E, and not E., omitting the superscript to favor readability.
Max-margin token selection. of tokens in X selected by p as
Given the (els) embedding p and sequence X, we denote the set
Sx(p) {s arg max p Es),
(10)
sEX
and we define (p) = X S Intuitively, given a sequence X, the selected tokens in X have the largest softmax weight (proportional exp ). Note that, for + p, may have that (x = (p) for all X. Thus, we define the equivalence relation
p p - (x) = =xx(p') (p'),
for all X E X
(11)
Intuitively, two vectors P. p' are equivalent under the above relation if they select the same tokens for all the sequences. Given a vector po, we denote by Pp. its equivalence class, and we define the set of max-margin directions among vectors in Pr. as
{
P |1 pEPpo
Pp(p.)
= arg min l!pl|2
(12)
s.t. (E, E
AS Sx(p)), As E Sx(Pp.), VX E x.).
We first show in Lemma 4.2 (proved in Appendix C.2) that the max-margin problem in 112 always has unique solution, i.e., (p.) always singleton. Thus, later on, we will use p(p.) as the p(pe)
solution to (1).. and p.(pp (p.) = 1TA(= We drop the dependency on Po when there is no confusion.
Lemma 4.2. For any po * 0, the max margin problem in 12 has unique solution denoted
Implicit bias of gradient flow. While Lemma 4.1 holld for any dataset, we need an extra assumption on the data to analyze the gradient flow, due to the complex loss landscape caused by softmax attention.
Assumption 1. Each sequence in Xr contains either a single completely positive token or a single completely negative token, and all remaining tokens are irrelevant.
Assumption I implies that all sequences in the dataset contain precisely one relevant token, and the relevant token also aligns with the label. We remark that datasets containing only one relevant token have been also considered in prior work, see (36] Theorem 1] and [25]. We further denote by S, the set containing all completely positive and all completely negative tokens.
Theorem 4.3. Under Assumption for any 8 > 0, let
SS 2 (887)) + 111 + 2)2 [8| |S|
No 2 2 d max 256,
2 log
(13)
8
Let Pt be the solution of the gradient flow 99. Then, with probability at least 8, we have that Pe exists, the limiting direction poo lp ll2 + oc. Furthermore, assuming that Pos := lim Tel
satisfies the following properties with probability at least - 8:
In Appendix F.3 we provide additional numerical evidence that learning all embeddings is essential. This corroborated by the fact that the more restrictive one-layer model embeddings. fails to catch up with its more flexible counterpart.
[ with untrained orthogonal token
5

---

1. Peo selects all completely positive and completely negative tokens, i.e., Se S UUx (po).
2. Poo is the max-margin direction for such a selection, i.e., Poo (Poo))
Theorem 4.3
4.3 shows that, if pr converges in direction, it must converge to the max-margin direction that selects all the completely positive/negative token. A sketch of the argument is given below and the complete proof is in Appendix C.3
Proof sketch. We prove the three statements separately. First, we show that |l ll2 - oo (Lemma C.I). To do so, explicitly construct vector P such that pV pV'C(E', .p) < 0 for all p. This means that there is no stationary point with finite norm, which implies that the norm of the vector obtained via gradient flow diverges.
Next, we show that, if the directional limit Pso exists, then must select all the important tokens (Lemma C.2). To do so, we note that, after one step. the model approximately selects the important tokens. This implies that pa selects important tokens for all t, as gradient flow cannot increase the loss.
Finally, we show that Poo is the max-margin solution of feasible selection (Lemma C.6). To do so, we assume by contradiction that the directional limit is any vector p' that is not the max margin solution of a feasible selection. Then under this assumption, we prove that lim which is in fact the solution of [2), thus giving a contradiction.
p. (p'), l
Pt
Before proceeding with the characterization of the max-margin solution, we highlight some differ-
ences with respect to the related works (37] and [40] whose setup is similar to ours. Theorem 3 in (37) shows that gradient descent on p converges to locally optimal max margin solution when initialized in a regime close enough to such solution, and Theorem 4 in [37] shows that the regularization path can only converge locally max- max-margin solutions. However, these results do not exclude the possibility of the gradient flow converging to directions that are not locally optimal and not the max-margin direction. In contrast, we characterize all possible directions the gradient flow converges to, showing that these are max-margin directions that select all completely positive/negative tokens. Furthermore, we do so neither starting from an initialization that is close enough to such solution, nor with only one sequence in the dataset. This requires different proof strategy as compared to [37]. Finally, (37) does not consider any connections between the directional convergence and the importance of the tokens, which, in contrast, is our main focus.
Let us now comment on the comparison with 44). There are two key assumptions crucially needed in 40) that do not apply to our case. More precisely. [40] Assumption 1] requires that all the tokens non-relevant positions are nearly orthogonal to each other, which is critical in their proof. 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). Thus, the global convergence results in (40 cannot be applied. Furthermore, 440, Assumption 2] requires all the irrelevant tokens in the sequence to have exactly the same score 7 = E v. This assumption leads to the fact that for each sequence, only the term that involves the relevant tokens contributes to the directional gradient. In our setting, even though the score differences between different irrelevant tokens are of order 1/ Vd, these errors matter when the gradient norm is diverging to infinity, and the analysis in [40] cannot be applied.
Characterization of the max-margin solution. Theorem 4.3 gives that, if gradient flow converges in direction, the limiting direction is the max-margin one that selects all important tokens in each sentence. However, this does not exclude a-priori the possibility that gradient flow also selects some irrelevant tokens. To address this point, we now establish if and how many irrelevant tokens can be selected: in general, it not possible to select all irrelevant tokens (Lemma 4.4b and, under an additional assumption, no irrelevant token is selected (Lemma 4.5).
Lemma 4.4. Suppose selects all the tokens, i.e., Sx S. Then, Pxo + p.
Proof. If p selects all the tokens, then p Ex; = TT Ex, for all xi x, E X and for all X E Xn Thus, by Lemma A.2 p p E.p) = 0 for any p. which gives the desired result.
Lemma shows that the directional limit P (when exists) cannot select all tokens and, as selects all important ones, it must be biased towards them. As application, consider the case where

---

there is only one irrelevant token in the vocabulary. Then, the combination of Theorem 4.3 and Lemma 4.4 gives that only the completely positive/negatiie tokens are selected by gradient flow.
Going beyond the case where there is a single irrelevant token, the result below provides a sufficient condition for gradient flow to select only important tokens.
Lemma 4.5. Under Assumption [1 for any 5 >> O, assume that 13) holds. Let P be the solution of the max-margin problem (12 that only selects the completely positive/negative tokens, i.e.
p = arg min |pll
E.) s.t. - 2,)
A 1 (s
VX E Nn.
where 8 denotes the unique completely positive/negative token in the sequence X. Assume that
Pro = lim Toel Pt exists and that, for any z' solving (12) with different selection, llðl|2 < (1 -
112 (1-1)
some constant je that does not depend d. Then, by further taking
(2)16T(((+2 2
1816TTn(11 (1++200)
S 8 + 44n
d 2
2700 "V 2 log
1
have that Poo
with probability least 6.
PPlla
The proof of Lemma 4.5 deferred to Appendix C.4 The sufficient condition of the result above requires the max-margin direction that does not select irrelevant tokens to have larger margin than any other max-margin solution associated to a different token selection. We expect this to be the case e.g. for datasets where all the completely positive/negative tokens have the same ag, In fact, given the structure of the context embeddings in (6). the max- mar-in solution ý is expected to satisfy pT v s 0.p Fgg x T F. R 0 for all s E (x and s' E Sx()). Since the token embeddings initialization approximately orthogonal each other, f ~ E.esx(r) E. meaning
that ||ll|2 (p)( which implies that the sufficient condition holds.
4.3 Additional generalizations of the theoretical results
Multiclass classification. Consider K-class classification, let v the k-th classifier and perform one step of gradient descent on the token embeddings. Then, after some manipulations deferred to Appendix E.
we have
T Σ
E.
Up
no T
K 1 Ly=eL.2:=8
K U
(14)
where by using R, we hide quantities that have a vanishing order in d. The RHS of (14) generalization of the average signed frequency that indicates the importance of token with respect to class k. In particular, if we divide the RHS by (K - 1)), the first term is the co-occurrence of token with label k and the second term is the average co-occurrence of token with other labels. Thus, intuitively, if token occurs much more frequently with label & than other labels, this token is more important w.r.t. label & and the correlation with Uy is larger. We also conduct experiments on the Yelp datasets with 4 classes (see Appendix F.5). and we find that E & is still increasing with respect to the importance of the token for each class. for binary classification.
Training and Ex. If we train but fix the context embeddings, the implicit bias of p is not affected in the limit. To see this, compute
-V,C(p,)) = vo(X, E yol
R É
vo(X,v)
(x(X) E,,, (P)
Σ
Note that, if p converges in direction, it must select all relevant tokens. This means that y will be increasing s is relevant, and will not move a lot if is irrelevant, since appears exactly the same number of times in positive and negative sequences. While do not have theoretical characterization for jointly training E, p, v, in our experiments the token embeddings E are trained simultancously with p. Thus, we expect the token selection be the same as if we only trained p.
7

---

15
10
Mean
(Ex.p)
two-layer model (E.. v.
-10
Moan
-20
-0.5
P(1]Ex) - P(OJE,)
P()EEx) P(O)Ex)
Figure 2: Dot-product of token embeddings with (cls) embedding p (left) and regression coefficients v (right), as a function of the token-wise difference in posterior probabilities for synthetic data sampled according (16). We consider the one-layer attention model in ① with all parameters trained until convergence. The point cloud around zero corresponds to the tokens in the irrelevant set.
Numerical experiments
To support theoretical findings, showcase the correlation of the embeddings with the (cls) embedding and the output vector v, having trained all the parameters with gradient descent until convergence. We consider different datasets (synthetic data in Figure 2 IMDB/Yelp datasets in
Figures and 33 and different architectures (one-layer model ①T in Figures 2 and
S
(18) in Figure I]. Taken together, the experiments display an excellent agreement with our theory going beyond the one-layer architecture ( and also beyond the requirements on the data-generating process. Specifically, the trained embeddings capture the importance of the corresponding tokens: the dot-product with is proportional to how frequently the token appears in positive sequences rather than in negative ones, and the dot-product with p is proportional to the modulus such frequency. We detail below the experimental design.
Synthetic data. Let define the data-generating process for the synthetic experiments in Figure The data is generated according to a K -leevel model. Namely, the vocabulary set S is partitioned as
S = u
{
(15)
Here, S contains irrelevant tokens appearing in both positive and negative contexts with equal probability, while St and S.¹ (for & E .I, K}) contain tokens appearing mostly in positive and negative contexts, respectively. Formally, define the importance levels 5, 5, ōR 0. Then, given the sequence label E [-1,+1} and 8 E S, we sample the tokens from the vocabulary as
1-
8 E SEŠ,
18
p(sly)=
8(1
((1-6))
8 E S
(16)
Σ [S
58
ST
ΣR |s
where - denotes the binary inversion, i.c., -(+1) = = 1 and (-11) = +1. The law (16 implies the following posterior distribution:
1/2,
s s
p(yss)
1 sb.
e S
(17)
8k
S.
From (7). it is clear that (i) S contains irrelevant tokens as the posterior is uniform, and (ii) 8 quantifies the importance of the tokens in St#1 by skewing the posterior to be (6,, 1 8x. For
the experiments in Figure 2 we select the following hyper-parameters: [S] 2048, = 8 and
sequence length = 256; S = |S with |st = + 2 and S = 964; = 0.05 and
(6x 1K == = (0.45, 0.35, 0.3, 0.25,,20000 0.2, 0.1, 0.05, 0.02}.
Figure shows clear separation between positive and negative tokens (right plot with the dot- dot-product (Fzz v)), and the selection mechanism ((cls) token) assigns high weights tokens that have larger importance dy (left plot with the dot-produc P)),.)))

---

60
10
40
Mean
20
(Ex.
-10
Mean
-10"
0.0
PV1 E..) P(O(E,))
P(] - P(O(E,.)
10
40
Tokens
Moen
20
10"
0
-10"
-20
Tokens
-10
-1.0
-0.5
0.0
0.5
-1.0 -0.5
0.0
P(I P(OJE,))
P(1EE,) - P(O)E,)
Figure 3: Dot-product of token embeddings with (cls) embedding p and regression coefficients as function of the token-wise difference in posterior for IMDB dataset (top row) and Yelp dataset (bottom row). We consider the one-layer attention model in (I) with all parameters trained until convergence.
IMDB and Yelp datasets. The IMDB datase consists of 50000 reviews of average length 239 words per review, associated to either a positive or a negative sentiment. Yelp review provide much larger selection. To align the data size and sequence length with the IMDB dataset, we randomly subsample portion of the Yelp dataset constrained on the sequence length, i.c., we select reviews which have at least 1000 and not more than 1500 characters. In addition, Yelp reviews provide a five-star ranking, which we convert to the binary sentiment based on the following rule: 1/2 stars reviews are assigned label 1; 4/5 star reviews are assigned label +1; neutral reviews, i.e., 3-star score, are removed. We adhere to a typical preprocessing pipeline for both datasets: we start by cleaning the data from punctuation symbols and omitting the words, followed by an application of stemming: and use the Bert tokenizer from Hugging Face to tokenize sequences. Tokens that appear less than times are purged.
The numerical simulations for both datasets are reported in Figure 3 which displays a phenomenology similar to that obtained for synthetic data Figure 2 We have also 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 around 72% with test error around 0.6, while for the original baseline the test accuracy around 95% with test eтror around 0.1. This clearly indicates the importance of training the embeddings.
Two-layer model. We also consider the following two-layer model:
E'x LayerNorm(Softmax Ex E Ex Ex), f(X;p, E) Softmax(p) TippTTEE' (Ex) E
(18)
which includes both skip connection and the layer-norm. We note that, for both IMDB and Yelp data, the model in (8) achieves significantly smaller loss values at convergence (of the order of 10- in contrast to the order of 10 achieved by the model in (T)). However, even if this model is more
http:///hwwkkkkkaaoaaa/////////a of -00kk-off-00-kk-- mive-movie movie-aoe reviees

---

complex than the one analyzed in Section 4 the results in Figure ! are still remarkably similar to those in Figures and
Finally, we note that all plots consider on the r-axis the difference in posterior probabilities
E(X,)eD & E.) E.,)
Ei 1zz=s
(19)
E(XADED Ex= 8;=8
in place of the quantity as defined in (3) In fact, while the quantity in (3) appears naturally from the analysis of gradient descent, the difference in posterior probabilities provides better visuals for real data (IMDB and Yelp). The difference between (3) and (9)) lies in the normalization used:
the posterior difference in (19) is the discrepancy between counts of the token z; in positive and negative sentences normalized by the total number of occurrences of æ;, while the quantity in (19) normalizes the discrepancy by the total number of tokens n. the datasets. For synthetic data sampled according to (I6), due to the uniform nature of the sampling procedure, all tokens appear the same number of times. Thus, both quantities are the same up to a fixed scaling and, thus, they equivalent. Additional details on the hyperparameter settings for all the experimental setups are contained in Appendix D
we consider IMDB and Yelp datasets, train one full-gradient step on the embeddings (as in the theory) and then fix the context embeddings only training the (cls) embedding as well as the output vector. As result, the average standard deviation over the bins is reduced by roughly times for the inner product with e and by roughly times for the inner product with p. compared to training all context embeddings until convergence. This gives evidence that the first step of gradient descent aligns the model with the first-order statistics of the tokens. With more training, the alignment weakens, i.c., the standard deviation grows. This is due to the fact that the dataset also contains higher-order statistics that an optimal classifier must capture; these emerge in the parameters only after additional training.
Additional numerical results. In Appendix F.
We have then performed several additional experiments to establish whether the conclusions from our theoretical analysis continue to hold in more elaborate predictive models. In Appendix FF.2 we show that the presence of the MLP layer does not qualitatively change the token selection mechanism (overlap with p). Nevertheless, the monotonic behavior of (Ez, v) becomes less evident. In Appendix
F.4 we inspect two-layer model (88) with key-query matrices w-( and W (2(2)
KQ in place for both layers.
(2)
We identify similar trends in overlaps with the regression vector v and the (cls) token. However, the latter takes slightly modified form. In particular, instead of correlating with (Ez, p), the correct
quantity to access in this case is p T (2(2) Ez,- In Appendix F.5 we observe a similar phenomenology .Q
in the multi-class case for overlaps with each class regression vector 2i and certain token-wise "one-versus-ale statistic.
6 Conclusions and limitations
In this paper, we study how the embedding vectors trained via gradient methods capture the importance of different tokens in the dataset. We theoretically characterize (i) the context embedding E, after one gradient step, and (ūi) the implicit bias of the (cls) embedding p after training with gradient flow until convergence. We conduct experiments on synthetic and realistic datasets which demonstrate the generality of our findings.
A limitation of our work is that the characterization we put forward is only in terms of the first-order statisties of the tokens (i.e., the frequencies with which they occur in the dataset), and does not describe how the model learns the causal structure between tokens. In practice, both first-order statistics and causal structure are expected to be crucial for the model to "understand" text. While our theory assumes lne-layer attention model, the numerical results of Figure [ suggests that similar qualitative picture holds more generally. This prompts us to conjecture that in deeper attention models with multiple heads, the earlier layers form induction heads 28 which learn the causal structure between tokens, and later layers perform elassification based on the empirical statisties of the resulting k-tuples. We regard this investigation
exciting future direction.
10

---

Acknowledgements
DW and MM were funded in part by the Austrian Science Fund (FWF) 10.55776/COE12. For the purpose of open access, the authors have applied a CC BY public copyright license to any Author Accepted Manuscript version arising from this submission. AS and MM were partially funded by the European Union (ERC. INF2, project number 101 161364). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them. AS was also supported by the Swiss National Science Foundation under grant 204439. SO was supported by the Office of Naval Research grant N000142412289
References
[1] Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can transformers reason? The globality barrier and inductive scratchpad. NeurlPS, 2024.
[2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint, 2023.
[3] Anthropic. Claude language model. https://www.aathhoopi.cc.o 2025. Accessed via https://www. http://wnthopp.c.om com/
[4] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. NeurIPS, 2019.
[5] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. International Conference on Learning Representations, 2015.
[6] Simone Bombari and Marco Mondelli. Towards understanding the word sensitivity of attention layers: A study via random features. International Conference on Machine Learning, 2024.
[7] Vivien Cabannes, Elvis Dohmatob, and Alberto Bietti. Scaling laws for associative memories. International Conference on Learning Representations, 2024.
[8] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. Conference on Learning Theory, 2020.
[9] Borjan Geshkovski, Hugo Koubbi, Yury Polyanskiy, and Philippe Rigollet. Dynamic metasta- bility in the self-attention model. arXiv preprint, 2024.
[10] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. NeurIPS, 2023.
[11] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective transformers. arXiv preprint, 2023.
[12] Halil Alperen Gozeten, M Emrullah Iddiz, Xuechen Zhang, Mahdi Soltanolkotabi, Marco Mondelli, and Samet Oymak. Test-time training provably improves transformers as in-context learners. International Conference on Machine Learning, 2025.
[13] Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficieet transfer learning for NLP. International Conference on Machine Learning, 2019.
[14] M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention to markov models: Unveiling the dynamics of generative transformers. Interna- tional Conference on Machine Learning, 2024.
[15) Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. Conference Learning Theory, 2019.

---

[16] Albert Q Jiang, Alicja Ziarko, Bartosz Piotrowski, Wenda Li, Mateja Jamnik, and Piotr Mitos. Repurposing language models into embedding models: Finding the compute-optimal recipe. NeurlPS, 2024.
[17] Addison Kristanto Julistiono, Davoud Ataee Tarzanagh, and Navid Azizan. Optimizing attention with mirror descent: Generalized max-margin token selection. Workshop Contribution at NeurIPS, 2024.
[18] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 2021.
[19] Jannik Kossen, Mark Collier, Basil Mustafa, Xiao Wang. Xiaohua Zhai, Lucas Beyer, Andreas Steiner, Jesse Berent, Rodolphe Jenatton, and Effrosyni Kokiopoulou. Three towers: Flexible contrastive learning with pretrained image models. NeurIPS, 2023.
[20] Daniel Levine, Syed Asad Rizvi, Sacha Lévy, Nazreen Pallikkavaliyaveetl, David Zhang, Xingyu Chen, Sina Ghadermarzi, Ruiming Wu, Zihe Zheng, Ivan Vrkic, al. Cell2sentence: Teaching large language models the language of biology. International Conference on Machine Learning, 2024.
[21] Yingcong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics next token prediction with self-attention. International Conference on Artificial Intelligence and Statistics, 2024.
[22] Roey Magen, Shuning Shang, Zhiwei Xu, Spencer Frei, Wei Hu, and Gal Vardi. Benign overfitting in single-head attention. arXiv preprint, 2024.
[23] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A curious case of single-layer transformers. International Conference on Learning Representations, 2025.
[24] Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Hyeji Kim, Michael Gastpar, and Chanakya Ekbote. Local to global: Learning dynamics and effect of initialization for transformers. NeurlPS, 2024.
[25] Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer Attention layers provably solve single-location regression. International Conference on Learning Representations, 2025.
[26] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. Workshop Contribution at International Conference on Learning Representations, 2015.
[27] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. International Conference Machine Learning, 2024.
[28] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint, 2022.
[29] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning, International Conference on Machine Learning, 2023.
[30] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. International Conference on Machine Learning. 2021.
[31] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. International Conference on Machine Learning, 2021.
[32] Keitaro Sakamoto and Issei Sato. Benign or not-henign overfitting token selection of attention mechanism. arXiv preprint, 2024.
12

---

[33] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. NeurIPS, 2023.
[34] Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison Zhou. Implicit regularization of gradient flow one-layer softmax attention. Workshop Contribution at International Conference on Learning Representations, 2025.
[35] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 2018.
[36] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transform- ers as support vector machines. Workshop Contribution at NeurlPS, 2023.
[37] Davoud Atace Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. NewrIPS, 2023.
[38] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Shaolei Du. Scan and snap: Understand- ing training dynamics and token composition in I-layer transformer. NeurIPS, 2023.
[39] Lorenzo Tiberi, Francesca Mignacco, Kazuki Irie, and Haim Sompolinsky. Dissecting the interplay of attention paths in a statistical mechanics theory of transformers. NeurIPS, 2024.
[40] Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast conver- gence rates for self-attention. Transactions on Machine Learning Research, 2025.
[41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NeurIPS, 2017.
[42] James Vuckovie, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of attention. arXiv preprint, 2020.
[43] Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D) Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. International Conference on Machine Learning, 2024.
[44] Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, and Yingbin Liang. Training dynamics of transformers to recognize word co-occurrence via gradient flow analysis. NeurlPS, 2024.
[45] Tian Zhou, Peisong Niu, Liang Sun, Rong Jin, et al. One fits all: Power general time series analysis by pretrained Im. NeurlPS, 2023.
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:
We state the scope of the results in the abstract and introduction part.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
13

---

2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Section
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 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 on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on 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 is 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 be 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 favor of transparency play an 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 proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes)
Justification: We state the assumptions in Assumption ! and for each Theorem and Lemma, state the assumptions under which it holds.
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 a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main 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)?
14

---

Answer: [Yes]
Justification:
We state all the training hyperparameters in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is 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 a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way 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
appropriate to the research performed.
While NeurIPS 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 to access this model for reproducing the results or way to reproduce the model (e.g., with an open- 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 (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We do not find it necessary to release the code. Our experiments concern the training of rather standard architectures and they can be readily reproduced without needing to upload the code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc.
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (c.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! for more details.
//nips //npps.cc/puliiii wttp:ccpppppcccceeecee.ee..ccltt http://ublcgguudeeeeeeeeee http://eeeeeeeeeeeuuaeeeoooooooo
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
15

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code 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:
We provide all hyperparameters for the experiments in
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We give the scatterplots corresponding to all embeddings (together with the average correlations).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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 1-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 figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
show in tables or
If 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 execution) needed to reproduce the experiments?
Question:
Answer: [No]
16

---

Our experiments are on shallow attention layers, and they do not require large amount of computational resources.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't Iodn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics /sttps////ourrss..cpppll//////i
Answer: [Yes]
Justification:
We have read and conform to the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
This paper presents work whose goal is to advance the field Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Justification:
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., deploymeet of technologics 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 direct path to any negative applications, the authors should point it out. For example, it 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, 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 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.
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 ML).
17

---

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: Our paper does not contain safety issues.
Guidelines:
The answer NA means that the paper poses 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 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 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: We provide the link of the datasets use.
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, if 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.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- 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:
We do introduce any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
18

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper 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.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
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?
Question:
Answer: [NA]
Justification: This paper 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 as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core results in this paper does not involve LLMs as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standard components.
Please refer to our LLM policy http::////errpscccoooofffeeeeee http://c/ooooeeeee////22222/20 for what should or should not be described.
19

---

Additional notation.
Throughout the appendices, to simplify the notation, we write
exp(a,(X)) Ej- exp(aa(X))
ai(X) = p Exzss
(x)
(20)
so that f( E) E (X) We will drop the dependence on X in qx()),((()) when there is no confusion. We also denote
N (X, v) yL
(21)
dropping again the dependency on X, when there is no confusion. Finally, we define
g(X,)) = exp(uf((X;p E))
(22)
Properties of initialization. at initialization we have
By standard concentration inequalities, with probability at least 1 8,
max
max Sf
EE Egr
max sES
EE! u1, max #ES
EE, pl. lo ol S Vdl V
2
6
(23)
max
max
|12 l1ell2 S
min sES
E 122 2
sE8
For all results of the paper holding with probability at least 1 - 6, we will be implicitly conditioning on
23
A Technical lemmas
Lemma A.1. The gradients of the empirical loss are given by
T
x = = 2:0 )
VE, ((E = - v9(X,v)
I (x)=0 + D
1=1
jfi
i=1
T
(X) (X)) - Ez, E2, |
V Pp((E,p) p) = T
vg(X,v) y)
Σ
q
1=1
where we have defined g(X, v) (x)))
Proof. We start by taking the gradient of qs as
(x) -1,=s exp (EH,P) p Ej= exp ET,P)) (ΣiT 10,= exp EI, p) exp exp(EET,)
E exp (EX,)) ))
p D 1ay=s 1my=s) exp FTT exp (FT.P)
2 Ej exp EI,
T
11, Iuy=s)
i=1
Σ
1, 1r,=x)a.03
1
20

---

(oxp (FT Ez.)
(Σ}= exp (FTp))
Eÿ= exp FTT Eu, exp (ET,))
V (x):(X)=
( exp F
ΣJ= exp ET exp EE (Ez. Exs)
E exp ET,D)
Σ 1=1
qi93 (Per; -Faz))
Σ (Eg,(Ex, - -Eu,).
Next, we look the gradient of f fXx:, E):
T
Ve,J(X:p, E)
V E, (z E +
1:
T
Σ [1 2, 12
Eg,upt +
1 mq
9 4:
==1
jti
1=1
vpf(X:p. -
Σ
Σ
9:93 (Ez, - -Ess) ET,0
i=1
This allows to conclude that
V FE, p) = E
-y
V [e,((Xpp. E))
exp(yf E))
T
=
-
Σ
Σ
11 11-=0
+ expl E))
T
2 (10)0
+
Σ
1 1
1=1
V vpC(E,p) E
pp E)
I exp (yx E))
E F
=
y 1+ expl exp(ly (X: E))
Σ
Σ J#i
(X) - Ezs)
thus concluding the proof.
Lemma A.2. For any vector P. we have
ΣΣ (211)) -33 (X, wi) )
T
D () = 9(x,v)
i=1 j21
where à pT Ez, for all i E (1.....,T} T}.
Proof. From Lemma A.I
we have
V vpC(E,p) -
V9(X,w)
Σ 4i (X) Fmss)
-
9(X, y)
Σ (X)) EEs)) 76(x v)
[9(X, 3) Ex (Diaglax) - (xxxx x(x,1)]
21

---

where qx = 191 (X), 90 (X]] ,xX,)) 8) = (x, y).... (.(.)).......................... v)]" and Diag(1x) denotes the
diagonal matrix with [ 1la (xX))
Thus, letting ā [āi,.... (......:000 @r] R with ā. p0' Exs, have
4 v P = E [g( u)p Ex (Diag(ax) (xax)
E [g(X, y) (Diag(ax) -xxx)7(X v)]
T Σ ā.qi(1 a:m
E
x(1,))
ΣΣ
i=1 s#i
Σ 0.9:9 (% -70)))))) }
E
9(. v)
(use - 91 Σ 9j)
T
az9:a; - 788) )
E
9 X v)
0 - -75)
2
i=1
j=1 iz1
T
E
9Xx 3)
(a. - -aa-ayyqqa 2,)0 73))
2
1=1
T
-75) 755) )
E
g(X, 9)
Σ
aas - 3g (::00 -
Lemma A.3 (Convergence lemma). Let ||pe|l2 - oc and suppose there exists p such that, for any 6 > 0, there is a f(e) ensuring
D
P. (.)(E
V
2
v
for all t 1 T(e).
(24)
1pl
pt 1
Then, flimo-oo TPella exists, we have
Pt
lim 1-000
||
l2
| 2
Proof. By the definition of the gradient flow, (4 is equivalent
p dp
2 c) Pr dp. llpella df
1l1 dt
We note that
Pr dp:
d
Ipellz dt at |lelll2-
Thus, by integrating both sides from [E(c), t), we have:
TT
(pt pi()) 2 (1 - c) [PT()l2))
1p
which gives
π Pt
2 (1 (11 |l TT PTPR)))
|
pell2
1|pell2
11@1221lee122
Since Pi(c)) P have finite norm for fixed by taking the limit on both sides, we have
Pt 2 1-c.
lim inf
2-000
22

---

As we assume that lim,->oc ppll
exist and the above argument holds for any we conclude
Pt
lim 200
1
1
Lemma A.4. Given a sequence X, model parameters E, and indices i.,j Is. E Sp(p). 2g
X \ Sx (p); the following results hold.
1. We have
1
T qi S 1.
2. If there exist >> O such that p' (Ez.. - Ez,) 2 T for all xs, E (p). then we have
S 1 exp(((
3. If there exist > 0 such that p² (Ez., - Exi S for all xs. E Sx (p), then we have
Ij T Txpp(T)
Proof. The upper bound 4. is trivial. For the lower bound:
exp(p' Ez..)
9. exp(pT Ex..) + jzs. exp(p T Ex, exp(p' Ezi.))
2 T Texp(pT (ri.))..
T
there exists such that
ExT)) 2 for all n; E (x(p), then we have p (EE,,
1
4j
1 Eijj explpx (Fiz, - Fay))
$ 1 + exp(pp (Bzi, E25))
$ + exp(T))
there exists > such that pT (Fa., - Forr) S for all xs, E (x then we have
4. 1 Eny, exp(p (En, Eny))
1
(by definition of Sx
1+ (T' - exx (Eu,, E2,))
T Texp(T)*
B Properties after the first gradient step
Lemma B.1 (Boundedness of the embeddings). For any 5 > 0, let
[S|2 2 d 2 max 256, 2log 8
then with probability at least 1 - 8,
E < 2(1 + 2np).
||'||2 2|2 < 2 11md
23

---

Proof. By using
23
we have that
max #ES
|| 122 S max
(IÆ,/l2 + Du
2 [olll2 + |errallz)
< max
+ 1/
+ (pppl-1
SES
2
< 2 + 41/0
and that
|l 1l2 S 11 + |etplll S 2 + 11nd 3.
(25)
Lemma B.2 (Upper bound on the loss), For any 8 8> 0, let
|S|?) 2 (887 + 1117 + +2)8
d M max 256,
22og
8
then with probability at least 1 - 6,
T
L(E p²) S E log + exp
1
Σ
1o
you;
+ 22no
T
2
1=1
Proof. We first lower bound (X E) for each pair X,y, y. After the first step, we have max
(E - E max (p ( (E" - En) y10 (as <p") 8,5"
+ e (E Ez,) (err, err,) T p²1
We bound each term separately:
max ((0 (T")T(EE) (E. E.)1 < 2 max (600) S 2d
8,6"
770
(a. (x") (p))T(] Tol - no mnl T w S no ;
2
|enr,] (E. E.-) ) - ||emTp ||a|| E E lla S 44npd $/2 (1 + 217o).
((err, -err.) ) '0 - 2|1p 1|2 max |lerr 1|2 S 22 2 + 11md 1)
where we have used (33 By picking d 2 (88n)) + 1111j + 2 we get max,, Kp') ) [(E, -E ) s d= 3, which implies that, for any X and any E {1,..., T):
1
2d 2d
1
2d !
S q (X) - T
+
T
T
T
Thus, we lower bound yf(X:p F) for each pair (X, y) as
vf(X:p. E)
Σ 9: (X)
T
T
2d t/2 110
T
no
Σ
Oz + Σ (X)( (Ea, + +err,,)
2
youz
T i=1
2
T 2
i=1
T
11
- - (1 + 2d bo ) T (Fb, + e
T i=1
2
710
dd b 3(1 11pp)d
2 you,;
T
no
T
2
yo
227p
24

---

which allows us to conclude that
L(E',
E [og(( + (oxp(-yf(X;; E)))
E log 1 + exp
1
Σ
110
y0 +
1
(26)
T i=1
2
221jo
C Proofs for Section
C.1 Proof of Lemma
For simplicity, in the proof we drop the time dependency all the variables. By picking
|
2 log d 2
from 23 have
max
max #ES
EE c/. max sE8
E pl, 1p To vl
S d
max
max sES
| E2|1, ||p||2 22
Thus, at initialization, we have that, for all 8,
exp (-d S exp(p S exp (d=
which implies that, for any sequence X and any position i,
1
S {x())= S
T + 2T (d+
1 + (T exp 2d-
1 (T 1) exp (-2d
T
2T
d
where we use the fact that for = E [-1, - |2| S exp(2) < 1 + |2/.
Furthermore, for d >> 256 and for any sequence (X, 2) have
1 4d
1 4
T
T-
S (x))
- 7 T
and
-Td-3
Td
-2d <
< uf
E) <
< 2d 2d
T 2
T
2Td
Then,
o(X,y) 5
1
2 +
1
1 + exp -2d
2 - 2d 1/2
and similarly
9(X.1) v) 2 d 1
25

---

Now we look at the gradient update of the first step. By Lemma A.1
we have
T
((E,) = E
y9(X, v)
Σ
Σ
11,=0
1 10,=0)9005
E,, up +
Σ
1 = 9:0
i=1
R.
1
1
1 7 v
E
+
11
Ui
+
Vg( y)
Σ
1 2; 125=0)q.9;
Ex up
jfi
T
+ 12
g(x, v)
= 9:00
2
T
9i Es,. - Esy EE ))
vv((E,p) = E
yg(Xx,)
Σ
Σ
note that
1
1
E
y/
2 0
27
and we bound the remaining error terms. We have that
2
T
1
z,=s
Ci
7
11
v
S d
1=1
and
B
ug
Σ
Σ
11, 10,=
T
+ (1(x,8)
Σ
1. 4:0
S 1od
Furthermore, we also have that
[VV S 8d
Thus, the desired claim follows.
C.2 Proof of Lemma 4.2
Proof. We first show that, if 12 is feasible, then the solution is unique. Indeed, assume by con- tradiction that P1, P2 two different solutions of (12). Clearly, P1 and pa have the same norm,
Ipalaaleall # 1. Then, any convex combination of P1, P2 gives feasible solution with strictly smaller norm, which is contradiction.
Next, we show that (12) is always feasible. To see this, by definition, there exists some such that
(E, - Ex)
Vs E (x(),), Vs' E (x())), VX E X
Then, Pu is feasible solution of 12 which concludes the proof.
26

---

C.3
Proof of Theorem 4.3
prove each part separately. We first show that + Ipell2 - 00.
Lemma C.1. Under Assumption 1 for any 5 > 0, by picking
(/12)
d 2 max 256, 2log
|1| + 3
with probability at least we have lime ||p: = oo.
Proof. It is sufficient show that there exists a non-zero finite-norm ý, such that for any finite norm
VpC(E', p) # 0.
Indeed, above condition means that there no stationary point for any finite-norm p. For gradient flow, we have that
lim vyC(E', PP:)
which by contradiction implies the desired result.
Now we construct such p. Since d |S|+ + 2, we have that with high-probability E is full row rank. Furthermore, E = E + A and each row of A is in the subspace generated by v and po. Thus, we can pick 5 1 v, Po. so that
E"
Furthermore, for any fixed u, Po; let u1, U2 be the orthogonormal basis of span(r, (o) Then, with probability 1 we also have E"
uind u uzuu] also has full row rank.
Without loss of generality, let æı be an important token in a positive sequence Xx, i.., (1xx)) 2 Then, we define IRSS such that and a; for all + 1. Let
no
Tn
E
un u u2u; p
As E (1 - Un uzug has full row rank, there exists non-zero f that solves the above equation. By letting ý (T u uzu,] D,, we know that Po, and EoD By Lemma A.2 have that, for any
(a: a,) (75) ]
-D V,C(E' .p) = v(X,u)
ΣΣ
Σ
i=1
=g(Na, uk)
(Xx))
> 0.
4nT
J>1
which concludes the proof.
Next, we show that, if the directional limit exists, then it must select all completely positive/negative tokens. Lemma C.2. Under Assumption for any 5 0, by picking
S
no 2 4n2T2,
2 max 256, 2 2og
(88n) + 111 + 2)
with probability at least 1 8, if lim P
exists, then Poo satisfies
xxx Sx for all X E Xas
where s.x denotes the unique completely positive/negative token in the sequence X.
27

---

Proof. We prove the lemma by contradiction. W.l.o.g,, assume by contradiction that there exists X Xn cointaining the important token 21 s.t. 21 $ (Pe).. We show that there exists such that, for all 1 2 I,
((E' > C(E'.
which contradicts the fact that the gradient flow always decreases the loss.
To this, first note that by the definition of Sx (Pro); there exists some 0 independent of I. such that
min Poo (E., Ez,) -.
17
W.l.o.g, we assume that 22 is the token that achieves the minimum.
As lim I|Pall2 oo and lim,-soc large enough such that
Iplllz
Poo: we have that, for any >> 0, R >> O, there exists
lIpel12 2 2R,
Ppoll
Pt
Poc
S pr
for all t 122
Thus, we have:
(
Pt 1r
Em
Emg) = (F2 - Egg +
||rr12 Poo (Frr - Eiaz)
S - + 2p( (4nu 2)²
where we have used the result of Lemma B.1 Thus, by picking μ +()n++))2+ we have
Pi (Err Ema) S - 2'
T
|p:|12
which implies that
pl (Ezz 5 -RR
S
Next, we upper bound pt, E.). We first note that
41
exp (Fa, - Faz)) S exp(-TRR)
92
which gives
@1 S exp(-TR)).
Note that
T
yf(X:Pt, E²) = Σo 98%
i=1
S exp( (TR)) + max 7
5#1
< exp(--RR
1/0 (
+ (1 llm
+ (1 l1
Thus, by picking R
logd 4r
we have
3
23
3 4
uf(X:Po, E") <
2 110
d
which implies a lower bound on the loss:
1
- - 2n
3
1
C(Et,p)) 2 log(1 exp(_yf(Xx E
1))
log 1 exp
(27)
4
where we used that d > 256 in the last passage. Under Assumption I by Lemma 4.1 yo,, > 1/(rT") if x; is either the completely positive or the completely negative token in X, and otherwise yo², 0. Hence, given that each sequence X contains completely positive negative token, we have that
we have that
T
y0, >
T
1=1
28

---

As 1o > Tn > /2nI² / /11, by applying Lemma B.2
we obtain
1 2ni
C(E', p1) < log 1+ exp(
-
log(1 - exp(-n)) S exp(-nn)
4n
which gives a contradiction and concludes the proof.
Finally, we show that for each possible selection, pt converges direction, must converge the max-maie max-margin solution. In particular, we first prove the following lemma which gives an approximation the directional gradient of the locally optimal selection. To so, we define the secondary selection set and the locally optimal selection as follows:
Definition C.3. Given a vector p, for each sequence X, denote by S (p) the secondary selection set given by
S (p) arg max p E,,,, & Sx()))
(28)
We also denote by S (p) the set of tokens that are not chosen in the first and in the second place, i.e.,
(x() = x (Sx(p) (r))
(29)
Definition C.4. Given vector p, say that p is locally optimal if for every (X, y) pair, have
Σ (,(X, 7, 2 p 0,
for all (p),,,
(ESx(p)
for constant y that does not depends on p.
In the definition above and for the rest of this appendix, to help readability, we will abuse notation by letting indices (e.g., i above) also denote the corresponding tokens (e.g., Zi, Ij above).
Lemma C.5. Let p be unit-norm vector and = Rp for some positive constant R. Suppose p is locally optimal direction as defined in Desinition C.4 with some μ that does not depends R. Moreover, suppose there exists a constant T1 that may depend on P. no, n, T, d but not on R, such that:
min(p
DinnpT(F..) (F.- - F.)), (x E (x 2 n
X
(30)
min (F., E ()) E (x)) 21
X
Then, for any e 0, for any p 2 p such that lløl|2 does not depend on R and
min (E. EsE (p) A E X (x) 2 T2,
there exists R large enough such that:
- -VVCEE p) [1+
Σ Σ
(a,(X) - az(X)) (X. v.P.))
iESx (p) jES
D V p) 2 (1 (1
Σ
Σ
(x,(X) - (,(X))),,,
jESx (n) JESt (p)
where a.X)) = En,, (x(X) D Eng and
14 h (X,y,)) g(X, y) (XXx (X,
29

---

Proof. By Lemma A.2
we can write the directional gradient as follows:
T
p V
pp((E'
p) = E
(a.()) - -x,x)),,(xx)),,XXxxx11x))
i=1 j>i
= Σ
Σ
Σ
(ā:(X)
(BO)
iESx( JES (p)
+ E
Σ Σ
- (x))))
(BI)
1ESx (p) jES (p)
+ E
Σ
(0,(X) (x))))
(B2)
(EEX)
The rest of the proof is to show that
C1 exp(-
BO
B1
I C1 exp( R)
BO
C2 Rt) BO
B2
S C2
( exp(-Ti R))
BO
for some C2, C2 > O that do not depend on R. Then, by taking R large enough, we obtain the desired result.
First, we simplify [B). Note that, for all i, i0 E Sx()), we have that ā. (X) = 0:o (X). Hence, by switching the order of j, we obtain
Σ Σ
(x)(X) - (X)), = Σ
(T (X) - -0y(x)
hig MX,,X, y)P) (ESx(1)
Σ
(ESx(P) (p) (p)
jESZ(P)
= g(X, y)
Σ
(x), - (x)
(X))g, (X)
Σ
(7a(X, y) - (X,,)))
jES$ (p)
iESx(r)
for any io E (p). Since is locally optimal direction, we have
Σ (x((, 75(X, 2 14, for j E (x).
iESx (p)
Now, we compare
BT] and BO
By the exact same reason above, we can rewrite
Σ Σ
(X, az
iESx (P) (p)
= xg(,,y)
Σ
(ain (X) - -, (X);(X) Σ ()(Xx,))77(X) (ESx())
(ES
for any in E (p) and we compare to B0 term-by-term. Namely, for any X, E (x) and
1 E (p), we have:
a (X) - (x)l
< |0 Ez, llz <
2|lilla max, ||a,|12 == Cz.
(31)
a (X) xx,())
gk (X) = explax(X) ax(x)) S exp(-TT)),
(32)
(X)
EiEsx(0) hx(x,,) - (x(,, v)1
< 2T max, 1ad
2T max, || C2
(33)
EiESx(M) (1;(X,)) - 7,(X, v))
P which implies that, for any X.j E 8% (p) and l E Sg(P),
10: - a (X) (X (X)
Σ
1x y) 7/k (X, xxv))
(E)x(n)
I
(X) - (x) (x)
Σ (x((,)) - (x.()))))
iESx(D)
30

---

Thus, we get that:
B1
80
Next, compare (B2 and (30) S (p). We compare
Take any /' E E (x(p), X Sx()), E Ex(p),
(a, ã x
with each term in (BI. We note that the bounds (x) (xx)) and Big (X)-
(x are
the same as those in and 33
Furthermore,
9k
I exp(-ITRR
9i
which gives that thus concluding the proof.
B2
S T exp CCCC
B0
Lemma C.6. Under Assumption 7 for any ô > O, by picking
|S|2
Yo 2 4n2T2 d 2 max 256, 2 log
(8818 + 11170 2)
with probability 2 1 over the initialization, Pso lime llrlla exists, then peo E P. (poc).
Proof. We prove the lemma by contradiction. We first assume that there exists poc such that Then, we show that there exists P Pa (Poo) such that, for
Poo & P. (px)) and Poo = lim,+ro ||
any > 0, there is E(e) ensuring
pT
Dt [pell
V
.P:) -1 -
VpC(E' Pc),
for all t 2 E(c).
2
As consequence, by Lemma A.3 we have that Pso l1 which gives a contradiction.
For the rest of the proof, we fix any e and denote R = |p. We define p; = 2. 1 and we cquivalently show that:
Tipel
P Pt) 2 -(1 -1 .P).
(34)
To prove this, we first note that since pso $ P. (poo), for all independent of R such that
$ Pi
E P.(po)), there exists TO
|| P 1001211p1211 Tp-
Thus, by the definition of P. (Poc). there exists Xi S X, such that for each sequence X E Xo. we can find pair of indices (i.j) with Sx(Pe)), E X Sx (Pxo) violating the margin, i.e.,
(11ilaapoo) (E.; Ex3)
3r,
for some that does not depend on R. With a slight abuse of notation, we define 7 as
-= min| min {1 (118llp000) T - E2,), Sx(px)), E S (px)),
Xext ((17))2250) (Ez, - Ez,), (Poo), (px)), Xex,
min
2 (Ez, - E,,),i E 8 (poo).s E (poo)) XeX,
min
This means that, for all X E Xr and for all (i, pairs such that E Sx (Poo) E S (pxo), we have
(11112250) T (Ez; Ex,) 2 3r;
31

---

for all pairs (i.5) such that i S ( E (po)), we have
(1011200)) (Ez, Ez,) 2 3r: and for all X E Xo, e (x(ps)) E Sx (Poo), we have
(11i12Pbo) (L'n; Ex2) 1113. 3T,
with some that does not depend on R.
Now, we compute the overlap with pt. For all X and (i,j), we have
Pi (Ez. - Ez,) (l (Ez. Ez,) + (P l10llppoo) (Ex, - Exy)
We upper bound
(pi 11 ( (Ex, E2,)1 1P
Pt
Poc
||Ez;
Exq 12,
||pe|l2
and since 11 Ex,
Eag are finite, we have
lim (pr - |1 T (Be. B2, =0.
Thus, we can pick 1,, such that for 2 11, we have
(pt 11112p50) (Fz, F,,)1s5
which implies that, for X E Nn and for all (i,j) pairs such that i E Sx(Px)) E (Pcc) we have
Pt T (Ea, Ea,) 2
for all (i,j) pairs such that E (poo E (p(co)), have
Pt (Ea, Exy) ZT;
and for all X Xo, Sx(pso), S (Poc), have:
Pt T (Ez. - Ex,) 11-T,
some that does not depend on R.
Next, we show that Pi is a locally optimal solution as per Definition C.4 By Lemma Poo selects all the completely positive/negative tokens. Thus, as pi Poos Pi also selects such tokens, the rest Hence, for any pair (X, y) and for any E X (x(p:) (p:), we have:
being irrelevant by Assumption |
Σ
-X, Z
To 4nT
VESx (pt)
by picking d large enough (as per the hypothesis of the lemma). By construction, & ~ Pt, llpll2 does not depends on R and, moreover, for any X,
p (Ez. 1 1, for all i Sx))), j E X XX8,5 SSx
By applying Lemma C.5 on both and pr, have that for any E1 there exist to s.t. for all t 2 max(ti, t2), we have
p VVVLL(E' P..) 2 E)
Σ
Σ
(āx(X) - - (X.pp..)
iESx(p) jES (p)
- pi VVL(EE Pb) $ (1+6))
Σ Σ
(a,(X) - (XX
iESx (p) JESt (p)
where T7(X), are defined analogously à, (X), (X) by replacing & with pr.
32

---

Now, we further show that, for any > 0, there exist to such that for all 2 t3.
Σ Σ 1ESx (p) (P)
(x) Bg
(x(X))),,(Xx
$ (1 - 62 ExEx
Σ
Σ
(a,(X) - (x)))) M.PE)
(ESx(P) JES (p)
To see this, we use the same idea as in the proof of Lemma C.5
We can write
E
Σ Σ
(x, - (X, y.))
xESx (p) ( (p)
= Exexo
Σ
Σ
(x) - (X
(A0)
iESx(p) JESt (p)
+ Exrex. (Ab
Σ
Σ
- - -X
(AI)
iES (E) (E)
and it is sufficient show that
A1
E2 A0
To prove this, we compare term-by-term. Let E No, X' E Nr No.1 E (p.)) E Sx (pc). and recall that:
Σ
(T, - aij (X)) (X,
(x))
Cho (X) TT(())))))) (x)
Σ
- -x,
=
(ESx(Rd
(xT(X') (X')))) (X',3/
(p:)
= =( (aa (X') (xy(('))) (X')
Σ (n(X', "X',
for any io E Sx (x(), Sx (p:). Note that
g(X', 1/)
maxx,y ((, v)
max(1 exp(xf(X))) S (1++ exp(no)) Cs. (35) X
g(X,y) minx,y g(X,y)
By using the same argument as in 31 and
33
we have
ac (X') apl (X)
0io (X) ag (X) Cz.
Eiesy (76( X (/,8)- - 7t3 (X', (X',1))
S Ca.
Eiesx (m.) (76 X,y) - 75 1,,X,10 v))
Finally, need to upper bound:
9:, (x')) (X')
46o (X (X)4 (X)
We note that
a, (X') ax(xx)) 2 R/ ai (X) - (X) < (1 (1-7) (11b))//1811..
33

---

where ax(X) P Eg.. Thus by Lemma A.4. we have:
1
9h (X)
T (x Texp((1 - T)R )/(l)(l2)) 91, 11, qr(X)) z
exp(11/ (1
which implies that
9a (1'))) (gx1(()
5 (xpp((-TT/lilll)
lig {gj((()
Thus, for each X X' Xn Ko.j E Sx (pt). (Pt). we have
Σ
(X', s/ pt) - Co exp(-TR/llilla)
Σ (x)x)-(((())n (X,
(x)
(ESx(p)
Thus by picking large enough i3 which gives large enough Rt we have:
A1
A0
This allows to conclude that
-p" V V E p+) 2 (T)
Σ
Σ
(ū,(X) (X)
(EEx (p) JES
2 e1
Σ Σ
(ā(()) ā (X))a (x))x,),,X
iESx (p) JES (p)
-ii V/C(E', pd) +z
Σ (a(() - (X.())
(E
Note that, for each X E Ko.
ax))
(X) - 31 11-T;
which gives that
V pc)
T VpL(E',, .P). (1+
6
Since é1, 62 can be arbitrarily small, the proof is complete.
C.4 Proof of Lemma 4.5
We first prove two auxiliary lemmas showing that for any max-margin solution, at least constraint must be tight.
Lemma C.7. Consider the max-margin problem and let P be the solution. Then, there exist
X Xn, E (p), E X Sx (p) such that:
=x)= =.
Proof. Assume by contradiction that for all X E Xn 8 E Sx()), X Sx(p),
T EE,) > 1.
Define
min X Xx
D (E, Exx);
and we know that > Then, pi' is also feasible solution to 12) with || contradicts to the fact that p is the solution to the max-margin problem.
[p]|z which
L
Lemma C.8. Under the same condition on in Lemma 4.5 p the unique solution of the max-margin problem in (2) that only selects the completely positive and negutive tokens. Then we have:
llpll2 < 4n.
34

---

Proof. By the definition of the max-margin problem, it is sufficient to construct feasible p with Illlz I 4n, which implies that:
1 - |ll|2 S 4n.
Since p only selects completely
positive
and negative
tokens,
we
know that
(Ux ((U Sx(i)) = 0. We let
p=2
Σ
E
xEUx Sx (+)
and by (23) we know that |pl2 S 4m, since there are at most n completely positive and negative tokens.
remains show that is feasible with high probability. To see this, for E Ux 8x (D), have
p =2+
E
(E2)
E
7o a.p TT +
Topeen err,
#'E8x
2 2 170
2 0k2
Vd
T
2 2lg
S
6
22 2nmnn
2 2 To n
log
S
22v 22//2mm
2
where the first inequality follows from
23 and Lemma 4.1 Similarly, for E Ux Sx
L E
70
2 log
S + 2nn/ooo 4
vd
Thus, by picking
+ 44v 2n ))²
S| 8
d 2
2no0
n 2 log
we have p (E,- E!) > 1, with E Ux Sx (P), 5' E Ux Sx()), which indicates the feasibility of p and finishes the proof.
Proof of Lemma 4.5 Let ş' be the max-margin solution of (2) with a different selection. By
Theorem 4.3 we have that, for all X, E (Ø'). We denote by iX the index of s. Assume
by contradiction Peo = o' T We will now show that this implies the following statement: for any 0, there is a t(c) ensuring
TT
(p c) PJ
Ipdlle vpC(E,p)),
for all t t(c).
(36)
Then, by Lemma
A.3
we have that Poo - which gives contradiction.
As in the proof of Lemma C.6
define Pt 1| |p Thus, (36 equivalent to
PpC(E 2 (1 e)p V,C(E,P))
First of all, since p, o' are two max max-margin solutions, by the definition of max- margin solution and we have:
Lemma C.7
O (E.x E.)21 1,
V | 5
( - F25)) =1
X Xn, E Sx(i'), E (x(())
which implies that
||ll p
(E, Ex)) = l|:1|2 1 1- 11 1
XE E Xn s' E X x Vs e
| 1
|| 1
35

---

For simplicity, we define (F') - X \\ (x') such that
(f)) (E. E.) = 1,
X &n,s' s' E exx(0), E Sx(0),
and note that ( can be empty for some X. Also define
min
(F) - Egr) 1.
f' a for any E1 E (0, 14) small enough, there exists a t ensuring the following 1 2
As lime-+ Pi
for all 22t, t1:
Pt (E, - Ex) -1- < (1 11-#40 10(1 + TU) E1
E Xn, 8 E Sx E (P')), VX Xn, E Sx E Sx (x))
Pi (E, Ex) 2 (1 +1 no) E1,
(37)
which implies that SX S if)
By applying Lemma C.5 to pt, we obtain that, for any 62 0, there exists l2 ensuring that, for all
t 2 12.
Di T VoC(E',PE .P)) sp) 5
(1
g(X,y) Σ
E
- ([ xx)))
(ES (pe)
S (1+ g(X.v)
Σ E
- ( - 7x)(
iESx (pe) jESt (px)
(CI)
Now we further approximate CCI. For simplicity, we define Xo Xr such that, for all X E No. S ( + 0. We show that for any 63 > 0, there exists ty ensuring that, for all 2 max tz),
(1+e)
Σ s(x,w)
Σ Σ
n
XEX
iE5x(p.) (pt)
To see this, for any E3 O, we show that for 2 max t2, t3) and any X E No, X' E Xx \ No.
9(X',)) 9)
Σ Σ
(x') (xx)) (x' (X') (X') 7/3 (X')))
iESx (px) (PL)
(38)
<- <33((X.)) v)
Σ
Σ
(axx (X) - (x(x)))) (x)g (X) (x))).
iESx (P. JES (P:)
Indeed, using the same methods as in
31
33
35
we have
ब. (X') aj (X')
ax(X) (X) S Ca. E
(X') - (x()))
Cs, Vi (x())
(X) - (x))
g(X', v/) S Cs.
9(X, v)
where Cs, Ca, Cs are constants that do not depend on R. It remains to upper bound (x
gx(xg,
9xy 9.y (X) as all i Sx(p:) has the same qi-
which is equivalent to
By Lemma A.4 together with (37j. we have:
4: (X') )g,(),('' (X') S exp( (000(((11)(())) 11( + 7o) R.)
1
4x (X)q (X) 2 72
exp( (pp(-(111 + R))
36

---

where for simplicity we denote R = IIpa 12. This implies
1: (X') (X²)
9:5% (X) (X) S T exp(-((1- 11yo 22)).t).
Since l|p:|12 oc, and C3 C4, Ct do not depend on R, we can pick to large enough such that:
C exp(- - po - 22) S e3.
which implies 38
Thus, have
Di ( .Pi)
s (1+
+ (3)
Σ
9(X,y)
Σ
(a,(X) -x)0 - -2 (x)1
n
XeX
SESx (pr) JES (nc)
1
(1+
Σ
9(X,y)
Σ
(0. (X) 0. (X)) (X)) -x)))
71
xex
+ +1++2)(1 Ez)
=
Σ
3)
Σ
Σ
-
(x)
XEX,
(r)
(39)
We then compute by Lemma A.2 that
- V pp) - g(X.y) Σ
0xx (X) - (x)
(X) (x)))
jESx (pr)
(DI)
9(x,))
Σ
Σ
(a(X) (XX)))( (,)
iESx (Pe). (P8
(D2)
E
g(X,v)
Σ
Σ
(ā.()) (1000)0
iESx (p:)
(p
(D3)
We show that by picking
2816Tn2(1 21j0) )
4
d
have
#
D2
D1
2
and for any e4 > 0, there exists t such that for all 2 t4
D3
6a D1
To show
D2
it suffices note that, for any j E Sx(p:).
(x)( dr (X)= > 1,
(xx(X) xx(X) < 16n(1 + 270),
ii,
(x() 7x(1) S 22md 1/4
Vi
Rg (X) (x)) Z 4nT
No
37

---

where the second inequality follows from Lemma C.8 and Lemma B.1 we use again the inequalities above and the same strategy as Lemma C.5
To show
D3
Ea
DI
Thus we obtain that:
V PpC(E, El, PP))
1
P 2
c4 E
g(X,))
Σ (xx(()
jESx(P.)
1
E4
1
Σ
g(X,y) v)
(ā (X)
(x))) -7g)))
2
XeNo
jES (p.)
where in the last inequality use that all the summand are non-negative.
Note that
4p (X) xx(x)
x s +
XA E No.
Thus, it remains to show that
g(X,y)
Σ
(x)
(Xx - (x()
4n
X
(ES)(
21 o2)
E3) =3)=
1 Σ
g(X,y)
Σ
Σ
n XEX:
( (p))(
(40)
We have that ae -x)) 2
ax()- - ax 1 11-++-
ViEex(p)). VX E No.
Furthermore,
Ixa(X) - [x(1 S 22mpd-1/
Vi
(x) > 4nT
2816TM"(1+21)) 4
(266)
As d
40
holds and the proof is complete.
14
1
D Details of numerical experiments
For all numerical simulations, we use the AdamW optimizer from torch. optim, and we reduce the learning rate in multiplicative fashion by factor ? 0.1 at epochs 100 and 200, i.e.,
LR LRoid 7/-
We adhere to the batch size of 128 and fix the embedding dimension 2048.
IMDB and Yelp datasets. The hyperparameters do not differ between the two-layer model and the one-layer model. We set the number of training epochs 500, the learning rate to 0.01, and the weight decay to 10-5
Synthetic data. We set the number of training epochs to 196, the learning rate to 10-4, and the weight decay to 10 4
38

---

E
One-step analysis for multi-class classification
We consider there are K classes. The sequence X is defined as usual, and the label y R is the one-hot vector indicating the corresponding class. The new model becomes
f(X) VXSoftmax(X p) R
where V = [ui, ux T ER
Rxdd
We consider the standard cross-entropy loss:
exp(y f(X))
(), log K
E exp f(x))
and compute the gradient as:
E exp( c/ (x -x)
fn E)
Σ k= exp(e) ((x))
T
T
VE,(((X)] = Σ
11 12:== - 11 125=
FT, ukP
Σ 110kk
jzt
K
Tn. E) = E
[XII VE, (F())
Eesr exp f(x))
exp xxx))
E
1 y=m=
1 Luter
k== exple? f(X))
k= exp(el/ f(x))
Σ (1z, 1 it
Ey, U&P Σ 11 1b0K
Thus, the initial step, we have:
K
(K -K-1)
T 12,=0
Ve.Co(D.E)
y=ex
+
ll Nutee
K
K
Consequently, after gradient step, obtain the correlation
E.
U2
No T
Σ K
Ly=ec.2.=0
K
K
Here, by using R, we hide quantities that have vanishing order in d, and these passages could be made rigorous with a more involved technical analysis.
F
Additional numerical experiments
F.1
Variance of one-step gradient descent
In this section, we numerically validate the optimization procedure, which is more in line with the theoretical model of "one-step" analysis presented in the main body. More precisely, we perform one full-gradient step on the embeddings (as in our theoretical analysis), and we then fix the context embeddings, only training the (cls) embedding as well as the output vector v till convergence. The results for the one-layer model (I) IMDB data shown in Figure 4 In this case, both overlap statistics display similar behavior to the one Figures and 3 However, the average standard deviation over the bins is reduced by roughly 8 times for the inner product with and by roughly times for the inner product with p.
We attribute the larger variance of full training dynamics (i.e., with all parameters optimized si- multaneously) to both the longer training time for context embeddings and the dataset itself. The dependence on the longer training of the embeddings is indicated in the experiments in Figure and the dependence on the dataset apparent by comparing Figure (synthetic datasets, small variance) with Figures I and 3 (real datascts, large variance).
39

---

10
Mean
(Exp
-5
Mean
-15
P(] P(O)E,))
P() P(O)Es.)
Figure 4: Dot-product of token embeddings with (cls) embedding p (left) and regression coefficients v (right), as function of the token-wise difference in posterior probabilities, for IMDB dataset under the "one-step" procedure.
Tokens
Mean
Ex
-5
Takens
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.5
1.0
P(])Es)) ((O
P())E,)) - P(O)Ex.)
Figure 5: Dot-product of token embeddings with (cls) embedding p (left) and regression coefficients (right), as a function of the token- wise difference in posterior probabilities, for IMDB dataset for one-layer attention with MLP head.
F.2 The effect of MLP layer
In this section, identify how the presence of the extra MLP head after the attention layer in (I) affects the phenomenology described in the main body, Namely, we consider the MLP block which processes input r IRd follows:
31 ReLU(W x + b1).
W1 E Rddd
b Rd
22 ReLU(W22 + b2).
W2 E Rddd
by Rdd
Then, the output of the self-attention
I i...
Sof Smfx E Rd.
treated as an input for the MLP block. Afterwards, the MLP output multiplied with the final regression vector c. The correlation analysis for the architecture above is displayed in Figure is clear that the behavior of E p is mostly unchanged, while the monotonicity of E is less evident. This indicates that the model will still select important tokens using E. p. but the interaction between the embeddings and the output vector is more complicated due to the presence of the MLP layer.
F.3 Frozen orthogonal embeddings
In this section, we address training embeddings in model (I) is indeed necessary. Namely, we consider fixing token embeddings to their (almost) orthogonal configuration at initialization and training only (cls) embedding and the output regression vector v. The resulting test accuracy around 72% with test error around 0.6. This is significant downgrade from the original baseline: the test accuracy is around 95% with test error around 0.1. This clearly indicates the importance of training the embeddings in the self-attention model (I
40

---

Mean
-W10 Exx
(E vl
-1
-20
-2
-30
Moan
P(]]Ex) - P(O(E,)
P(] - P(O)E,)
Figure 6: Dot-product of token embeddings with (cls) embedding and key-query matrix of the second attention layer W difference in posterior probabilities, for IMDB dataset and two-layer attention model (T8) with key-query matrices.
AQ (left) and regression coefficients v (right), as a function of the token-wise
200
Tokens
Tokens
Mean
Mean
100
(Ex v2)
0
-100
-100
-0.02
-0.01
0.00
0.01
0.02
-0.03
-0.02
-0.01
0.00
0.01
0.02
0.03
T2(Ex)
T.(Ex)
Figure 7: Dot-product token embeddings with the corresponding regression coefficients v (left: k 2; right: k = 4), function of the token-wise multi-class statistic for Yelp dataset. The token statistic T3 on w-axis is defined as per (41
F,4 Influence of key-query matrix
In this section, we explore the presence of key-query matrices in the self-attention mechanism changes the qualitative implications of our results. For experimental setup, consider layer model (18) with key-query matrices in both instances of softmax operator. We present our findings in Figure [6] First, observe that the behavior of the overlap with the regression vector
remains unchanged. Second, we note that adding W-(2) in the second layer should not make
qualitative difference. In particular, wr(2)
KQ P plays the same role as p, which corroborated by the
results in Figure 6 Besides, one can recover the dynamics key-query matrices given the dynamics of (cf., Lemma in 37
F.5 Multi-class simulations
Given the analysis in Section E. we consider a modified version the token-wise statistic, Namely, we focus on the following cone-versus-all" score:
T
K K
1 K ky=kort
(F. T
11 1g
(41)
The RHS (41) can be viewed as generalization of average signed frequency 3 that indicates the importance of token with respect to class k. In particular, if we divide the RHS of (3) by factor of the first term becomes the co-occurrence of token s with label k and the second term is the average co-occurrence of token with other labels. Thus, intuitively, if token occurs much
41

---

more frequently with label & than other labels, this token is more important with respect to label k: and the RHS of (4T) will be larger.
For the multi-class dataset, we consider Yelp reviews. We use the 5-star rating directly and disregard the middle-ground rating of (4 classes in total). We train the one-layer model (T this data. The results for two classes are presented in Figure 7 Namely, it is clear that E. vh is still increasing with respect the importance of the token for each class k, which aligns with the previous observations for binary classification.
42


Reviews:
[{"review_id": "uMt91wqpIh", "reviewer": "Reviewer_vwUg", "summary": "This 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.", "questions": "1. 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?", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "I 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.", "was_revised": true, "final_justification": "The rebuttal clarified the connections to some closely related works and added useful discussion and new empirical results, which made the contribution, scope, and limitations much clearer. The core message and setup design are well chosen, and the experiments support the claims; the new results also reveal the limitations, giving a more complete picture. Based on this, I raised my score to borderline accept. I’m keeping it at borderline as the analysis still relies on several simplifying assumptions about the model.", "rebuttal": "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.\n\n---\n\n**W1. Simplified model.**\n\nThe 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$.\n\n---\n\n**W2. Comparison with [40].**\n\nThere are **two key differences** between our setup and that of [40]: \n\n(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.\n\n(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. \n\nThus, 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.\n\n---\n\n**Q1. Why fix token embeddings?**\n\nFixing 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**.\n\n---\n\n\n**Q2. Clarification on synthetic data generation.**\n\nThe 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.\n\n---\n\n**Q3. High variance in IMDB and Yelp datasets.**\n\nThis 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**.\n\nWe 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).\n\nIn 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."}, {"review_id": "s5VdMLjHS9", "reviewer": "Reviewer_E1M4", "summary": "This 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.", "questions": "(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?", "limitations": "Please see above weaknesses and questions.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "After the rebuttal and discussion with authors, the results are summarized as follows:\n\n**Weaknesses:**\n\n**Weak1 Strong Assumption and Weak3 Generalization concerns:** The authors' reply is convincing to me, which solves my concerns;\n\n**Weak 2 Only focus on first-order statistics:** I acknowledge that the connection between token embeddings and first-order statistics were not revealed in previous works, but the higher-order relations cannot be neglected. \n\n**Questions:**\n\nThe authors solve all my questions especially for the multi-classification case and the role of MLP.\n\n**Summary:**\n\nOverall, I appreciate the contributions of this work on revealing the connections of token embeddings and labels in the training process even though there are some limitations like the only focus on first-order statistics. However, because this theory cannot be trivially proved by the real transformer structure, I will give it 4 for borderline acceptance.", "rebuttal": "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.\n \n---\n\n**W1. Strong Assumptions.**\n\nWhile 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:\n\n(1) The setting with a single relevant token was also considered in previous work [1,2]. \n\n(2) Generalization to multi-class classification is rather direct and presents a similar phenomenology. See our response to **Q2** below for details.\n\n(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.  \n\n---\n\n**W2. Only focus on first-order statistics.**\n\nWhile 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. \n\n---\n\n**W3. Generalization concern.**\n\nWe will discuss the role of MLPs, see our response to **Q4** below. \n\n---\n\n**Q1. Why binary classification?**\n\nIn 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.\n\nAs 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.\n\nFinally, we note that binary classification problems have been widely considered in prior work, see e.g. [2, 3].\n\n---\n\n**Q2. Multi-class classification?** \n\nWe provide below both **theoretical and experimental evidence that our conclusions for binary classification extend to the multi-class setting**.\n\n*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 \n$$ 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].$$ \n\nThe 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**.\n\n*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.\n\n\n\n|       |        |        |        |        |        |        |        |        |       |       |       |       |       |       |       |       |       |       |       |\n|-------|--------|--------|--------|--------|--------|--------|--------|--------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|\n| **$P(1\\|E_s) - P(0\\|E_s)$** | -0.026 | -0.020 | -0.017 | -0.014 | -0.011 | -0.008 | -0.005 | -0.002 | 0.001 | 0.005 | 0.008 | 0.011 | 0.014 | 0.017 | 0.020 | 0.023 | 0.026 | 0.029 | 0.032 |\n| **Average $E_s^\\top v_2$** | -42.375| -27.306| -51.164| -20.713| -40.923| -41.925| -41.884| -35.691| 18.255| 35.633| 30.524| 56.211| 44.184| 23.280| 45.341| 43.210| 62.935| 81.274| 53.548|\n\n\n|       |        |        |        |        |        |        |        |        |        |       |       |       |       |       |       |       |       |       |       |\n|-------|--------|--------|--------|--------|--------|--------|--------|--------|--------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|\n| **$P(1\\|E_s) - P(0\\|E_s)$** | -0.021 | -0.018 | -0.016 | -0.014 | -0.011 | -0.009 | -0.007 | -0.004 | -0.002 | 0.000 | 0.002 | 0.005 | 0.007 | 0.009 | 0.012 | 0.014 | 0.016 | 0.019 | 0.021 |\n| **Average $E_s^\\top v_4$** | -64.776| -66.074| -37.040| -32.499| -31.109| -30.686| -26.341| -23.153| -22.310| -0.648| 31.915| 27.601| 30.951| 29.300| 19.455| 29.693| 41.506| 34.151| 21.060|\n\n\nIn 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).\n\n---\n\n**Q3. Embedding dimension.**\n\nWe 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.\n\n---\n\n**Q4. Role of MLP.**\n\nWe 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.\n\n\n|       |       |       |       |       |       |       |       |       |       |       |      |      |      |      |      |      |      |      |      |      |\n|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|------|------|------|------|------|------|------|------|------|------|\n| **$P(1\\|E_s) - P(0\\|E_s)$** | -0.95 | -0.85 | -0.76 | -0.66 | -0.56 | -0.47 | -0.37 | -0.28 | -0.18 | -0.08 | 0.01 | 0.11 | 0.21 | 0.30 | 0.40 | 0.50 | 0.59 | 0.69 | 0.79 | 0.89 |\n| **Average $E_s^\\top v$** | 0.13  | 0.24  | 0.06  | 0.09  | 0.04  | 0.10  | 0.01  | 0.04  | -0.01 | 0.03  | 0.06 | 0.11 | 0.18 | 0.24 | 0.34 | 0.44 | 0.57 | 0.72 | 0.87 | 1.40 |\n\n\n|       |       |       |       |       |       |       |       |       |       |       |      |      |      |      |      |      |      |      |      |      |\n|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|------|------|------|------|------|------|------|------|------|------|\n| **$P(1\\|E_s) - P(0\\|E_s)$** | -0.95 | -0.85 | -0.76 | -0.66 | -0.56 | -0.47 | -0.37 | -0.28 | -0.18 | -0.08 | 0.01 | 0.11 | 0.21 | 0.30 | 0.40 | 0.50 | 0.59 | 0.69 | 0.79 | 0.89 |\n| **Average $E_s^\\top p$** | 4.76  | 4.61  | 3.26  | 3.08  | 2.80  | 2.06  | 1.60  | 1.43  | 0.89  | 0.67  | 0.59 | 0.71 | 1.11 | 1.61 | 2.02 | 2.30 | 3.28 | 3.39 | 4.84 | 6.10 |\n\nIn 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).\n\n---\n\n**Q5. Vision tasks.**\n\nThe 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.\n\n---\n\n**References**\n\n[1] Attention layers provably solve single-location regression. ICLR, 2025.\n\n[2] Max-margin token selection in attention mechanism. NeurIPS, 2023.\n\n[3] Implicit bias and fast convergence rates for self-attention. TMLR, 2025."}, {"review_id": "HeAYlXj7lP", "reviewer": "Reviewer_q6jg", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "I thank the authors for addressing my concerns. Although it still lacks a rigorous theoretical analysis for training all parameters together, some experiments help verify the results. I have accordingly increased my score.", "rebuttal": "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. \n\n\n---\n\n\n**W1. Untrained embeddings.**\n\nWe 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.\n\n---\n\n\n**W2. Assumption 1.**\n\nIf 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**:\n\n(1) If $p$ converges in direction, it **must select at least one relevant token for each sequence**, simply because the loss is non-increasing.\n\n(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.\n\nThe 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. \n\nWe finally note that assuming only one relevant token is common in the literature [A, B].\n\n[A] Attention layers provably solve single-location regression. ICLR, 2025.\n\n[B] Max-margin token selection in attention mechanism. NeurIPS, 2023.\n\n---\n\n**W3. Novelty compared to [1].**\n\nLet us elaborate on the technical novelty compared to [1]:\n\n(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.\n\n(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.\n\n(3) [1] does not consider any connections between the directional convergence and the importance of the tokens, which, in contrast, is our main focus.\n\n---\n\n**Q1. $v$ and $E_X$ trained simultaneously?**\n\nIf we train $v$ but fix the context embeddings, the implicit bias of $p$ is not affected in the limit. To see this, we compute\n\n$$-\\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} ].$$\n\nNote 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.\n\nWhile 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$.\n\n---\n\n**Q2. Only $p$ trained; adding $W_{KQ}^{(1)}$ and $W_{KQ}^{(2)}$.**\n\nFirst, we note that in the two-layer model experiments, $p,v, E_x$ are trained simultaneously.  \n\nSecond, 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.$ \n\nFinally, 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.\n\n|   |   |   |  |  |  |  |  |  |  |  |  |  |  | | | | | | | | \n|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|------|------|------|------|-------|\n| **$P(1\\|E_s) - P(0\\|E_s)$** | -0.95 | -0.85 | -0.76 | -0.66 | -0.56 | -0.47 | -0.37 | -0.28 | -0.18 | -0.08 | 0.01  | 0.11  | 0.21  | 0.30  | 0.40  | 0.50 | 0.59 | 0.69 | 0.79 | 0.89  |\n|**Average  $p^\\top  W_{KQ}^{(2)} E_s$** | 15.47 | 7.37  | 7.79  | 4.64  | 3.68  | 1.60  | 0.68  | -0.61 | -2.22 | -2.80 | -3.92 | -3.18 | -2.23 | -1.43 | -0.15 | 0.29 | 3.21 | 2.98 | 7.90 | 12.44 |\n\n\n\n| | | | | | | | | | | | | | | | | | | | | | \n|---------------------------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------| ------|------|------|------|------|------|------|------|------|------|\n| **$P(1\\|E_s) - P(0\\|E_s)$** | -0.95 | -0.85 | -0.76 | -0.66 | -0.56 | -0.47 | -0.37 | -0.28 | -0.18 | -0.08 | 0.01 | 0.11 | 0.21 | 0.30 | 0.40 | 0.50 | 0.59 | 0.69 | 0.79 | 0.89 |\n| **Average  $E_s^\\top v$** | -1.07 | -1.14 | -0.82 | -0.69 | -0.58 | -0.65 | -0.53 | -0.44 | -0.18 | -0.13 | 0.01 | 0.11 | 0.26 | 0.39 | 0.54 | 0.63 | 0.81 | 0.73 | 0.61 | 0.93 |\n\n\nThe 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)."}, {"review_id": "uZD61fiJBQ", "reviewer": "Reviewer_XscA", "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.", "questions": "- How can the proposed method be extended to more relatistc attention scenarios?\n- Can the authors speculate about the practical relevance of their findings, e.g. can this be used for model diagnostics?\n- Can other tasks then sentiment analysis be studied with the same framework, like contrastive learning?", "limitations": "Brief, but ok", "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- Paper provides theoretical analysis of a simplified scenario for how attention mechanisms learn to adapt to token frequency.\n- Theoretical claims are backed-up by experimental evidence.\n- The mathemical analysis seems solid, although I could not follow all steps.\n\nWeaknesses:\n- Only a simplified model of attention is studied and only a single layer. \n- 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. \n- Only limited set of synthetic experiments is performed.\n- Only sentiment analysis with a simple classification setting is studied, what about other tasks like contrastive learning.", "was_revised": true, "final_justification": "After considering the reviews by the other reviewers, the rebuttal and the careful explanations and additional experiments, this paper still seems a worthwhile study, and the authors did a good job of adressing the other reviewer's comments. Of course, theoretical analysis requires simplification, so I do think the approach taken by the authors is warranted.", "rebuttal": "We thank the reviewer for the comments and for the positive evaluation of our work. We address all weaknesses and questions below.\n\n---\n\n**W1. Simplified model.**\n\nAs 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$.\n\nWe 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.\n\n---\n\n**W2. Frequent yet irrelevant features; semantic concepts that interact.**\n\nThe 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. \n\nThe 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.  \n\n---\n\n**W3. Limited set of synthetic experiments.**\n\nDuring the rebuttal, we have performed **additional experiments and ablations**:\n\n(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.\n\n(2) We have shown that a qualitatively similar picture emerges for multi-class classification. See our response to **Q2** of *reviewer E1M4* for details. \n\n(3) We have analyzed the impact of MLP layers. See our response to **Q4** of *reviewer E1M4* for details. \n\n(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. \n\n(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. \n\n---\n\n**W4. Only sentiment analysis with simple classification.**\n\nAs 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. \n\nAs for contrastive learning, we discuss the connection in our response to **Q3** below.\n\n---\n\n**Q1. More realistic attention scenarios?**\n\nFirst, 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.$ \n\nTo 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. \n\nThus, 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).\n\n---\n\n**Q2. Practical relevance, e.g. for model diagnostics?**\n\nThis 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*. \n\nSpecifically, 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.\n\n---\n\n**Q3. Other tasks, e.g. contrastive learning?**\n\nWe have provided extra theoretical and experimental results on multi-class classification problems. Please refer to our response to **Q2** of *reviewer E1M4* for details.\n\nRegarding 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.\n\n(1) We need a proper definition of the importance of each token to each image from the dataset.\n\n(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.\n\nResolving such technical difficulties provides an exciting avenue for future research."}]

Author rebuttal:

Meta-review: The paper gives a theoretical account of how trained token embeddings in a simplified attention classifier encode first-order token-label statistics after one gradient step and how training the CLS vector drives max-margin token selection under a single-relevant-token data model. The main claims are supported by synthetic and IMDB/Yelp analyses, including a two‑layer variant that shows similar trends. The authors’ rebuttal strengthened the work with clarifications against closely related results, a multi‑class extension (derivation and experiments), MLP and two‑layer WKQ ablations, and an accuracy ablation showing trained embeddings are important addressing several reviewer concerns. The chief limitations remain the strong data/model assumptions and the focus on first‑order statistics.
Decision: Accept (poster)