How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability
Authors: Shawn Im
Changdae Oh
Zhen Fang
Sharon Li
Venue: ICLR 2026 Oral (2026)
Keywords: Semantic associations
Interpretability
LLM
Primary area: interpretability and explainable AI

Abstract: Semantic associations such as the link between "bird" and "flew" are foundational for language modeling as they enable models to go beyond memorization and instead generalize and generate coherent text. Understanding how these associations are learned and represented in language models is essential for connecting deep learning with linguistic theory and developing a mechanistic foundation for large language models. In this work, we analyze how these associations emerge from natural language data in attention-based language models through the lens of training dynamics. By leveraging a leading-term approximation of the gradients, we develop closed-form expressions for the weights at early stages of training that explain how semantic associations first take shape. Through our analysis, we reveal that each set of weights of the transformer has closed-form expressions as simple compositions of three basis functions--bigram, token-interchangeability, and context mappings--reflecting the statistics in the text corpus and uncover how each component of the transformer captures the semantic association based on these compositions. Experiments on real-world LLMs demonstrate that our theoretical weight characterizations closely match the learned weights, and qualitative analyses further guide us on how our theorem shines light on interpreting the learned association in transformers.

Published as a conference paper at ICLR 2026
How Do TRANSFORMERS LEARN TO ASSOCIATE TOKENS: GRADIENT LEADING TERMS BRING MECHANISTIC INTERPRETABILITY
Shawn Im', Changdae Zhen
Zhen Fang", Sharon Lit
University of Wisconsin-Madison University of Technology Sydney {shawnim, changdae, sharonli 8cs wisc. zhen. fang@uts edu, au
ABSTRACT
Semantic associations such the link between "bird" and "flew" are founda- tional for language modeling as they enable models to go beyond memorization and instead generalize and generate coherent text. Understanding how these ciations are learned and represented language models is essential for connecting deep learning with linguistic theory and developing mechanistic foundation for large language models. In this work, we analyze how these associations emerge from natural language data in attention based language models through the lens of training dynamics. By leveraging a leading-tent approximation of the gradients, we develop closed-form expressions for the weights at early stages of training that explain how semantic associations first take shape. Through our analysis, we reveal that each set of weights of the transformer has closed-form expressions simple compositions of three basis canctions-bigrane and context maapings-reflecting the statistics of the text corpus and uncovering how each component of the transformer captures semantic associations based on these compositions. Experiments on real-world LLMs demonstrate that our theo- retical weight characterizations closely match the learned weights, and qualitative analyses further show how our theorem shines light on interpreting the learned associations in transformers.
aoken-interchangaabiiie
INTRODUCTION
Large language models (LLMs) based on self-attention have shown strong capabilities in capturing both factual knowledge and qualitative aspects of the human world (Grattafiori et al., 2024; Yang et al., 2025; Team et al., 2024; Achiam et al., 2023). This progress has sparked growing interest in understanding why these models work well and, in particular, what kinds of internal structures emerge during training (Engels et al., 2024; Li et al., 2023a; Meng et al., 2022: Cunningham et al., 2023). Among these structures, semantic associations cling (Harris, 1954; Firth, 1957; Miller & Charles, 1991), as they enable models to connect words and concepts in ways that support generalization and coherent text generation. While recent studies have identified specific mechanisms such as induction heads (Olsson et al., 2022), linear seman- tic relations (Nanda al., 2023), and topic clustering (Li et al., 2023b), we still lack principled account of how semantic associations arise during the training of attention transformers.
especially foundational to language mod-
By semantic associationss, we mean the statistical and functional relationships between tokens that encode meaning-forr example, the link between "bird" and "flew", the interchangeability of "car" and "truck" in adjectival contexts, or the coupling of "country" and "capital". These associations have long been recognized in linguistics under the lens of distributional semantics (Harris, 1954). In modern transformers, such associations are not explicitly programmed but instead emerge through gradient-based optimization over large corpora. Understanding how these structures crystallize dur- ing training is therefore essential not only for connecting deep learning with linguistic theory but also for developing mechanistic foundation of representation learning in large language models.
In this work, we develop a theory for the emergence of semantic associations in attention- based language models trained natural language data, through the lens of training dynamics. A formal

---

Published as a conference paper at ICLR 2026
(1) Gradient leading-term reveals
(3)
Learned weights match theoretical characterizations!
three basis functions of association.
red
happy
flowar
ballon
excited
jparden
pretty
Interchangeabllity
Context
(2) Composition of basis functions characterize Transformer weights. X
Softmax (aakk
X
Query/Key Weights (,
X
Relative
X
Value Weights
Output Weights A(B)
Positional Bias
Figure I: To understand the emergence of associative features, we analyze the training dynamics of Transformers by focusing on the gradient leading terms for weights, which allows to identify in- terpretable basis functions that characterize each weight by their compositions. Empirical validation confirms that our weight characterizations match the actual ones learned in practical transformers.
analysis of training dynamics attractive as allows us to rigorously discuss how modern lan- guage models learn features and capabilities. Unfortunately, the training dynamics of transformers are highly complex, which has led prior work to adopt unrealistic assumptions that diverge from practice: (I) synthetic structured language (Li et al., 2023b; Yang al., 2024), (2) simplified model architectures without, egg. positional encoding or residual connections (Tian et al., 2023; Huang al., 2025), and (3) non-standard training, such as sequential component-wise training partially frozen weights (Bietti et al., 2023; Li et al., 2023b). While these prior works provide valuable the- oretical insights, their departures from realistic conditions raise concerns about the generalizability of their insights to LLMs used in practice. In contrast, we ground our study in more realistic set- ting by focusing on naturalistic text distributions and attention-baser transformers with positional encodings. optimized with standard training procedure (Brown et al., 2020). This is essential to minimize the gap between our theory and practical use.
Our key technical innovation is to analyze training dynamics of the transformer at an early stage, through the leading term of an expansion the gradients for each of weights. In particular, trans- formers are known to acquire many core behaviors early in training-including semantic relations- and persist through convergence (Olsson et al., 2022; Elhage et al., 2021; Nanda al., 2023). This makes the early phase not only empirically important but also analytically tractable. During this stage, gradient updates admit a closed-form approximation: the leading term dominates parame- ter updates before higher-order corrections accumulate. Leveraging this, show that the learned weight matrices (including the output matrix, value matrices, query-key matrices) can be expressed simple compositions of three basis functions: bigram mapping, which captures next token de- pendencies; an interchangeability mapping, which reflect functional similarity across tokens (c.g., synonyms shared grammatical roles); and context mapping, which encodes longer-range pre- fix suffix co-occurrence.
Through experiments on natural language dataset, we verify that the learned weights an attention-baus based transformer model closely match our theoretical closed-form expressions, and fur- ther demonstrate that this holds even beyond the early stage. We also show rich qualitative examples of how each weight component of the transformer captures the actual word-wise semantic asso- ciations characterized by our theorem. Furthermore, we verify that our theoretically characterized features correlated with the behavior of real-world language model. Figure depicts overview of our analysis, and we summarize our contributions as follows:
1. We present the first explicit characterization of weights in attention-based transformers trained on real-world text corpora under the next-token prediction loss;
2. We interpret the features learned in weights compositions of bi-gram, interchangeability.. and context mappings, and then show how these basis functions capture semantic associa- tion across words;
2

---

Published as a conference paper at ICLR 2026
3. We finally validate our theoretical interpretation both self- models and practical LLM, demonstrating the generality and relevance of our theorems.
RELATED WORKS
Understanding emergence of features in Transformers. Many works have considered the train- ing dynamics of transformers under controlled settings to interpret their feature learning (Tian et al,, 2023; Bietti et al., 2023; Nichani et al., 2024; Kim & Suzuki, 2024). A line of them investigates how low-level associative features, such as bigram structure (Bietti et al., 2023), cyclic structure (Huang al., 2025), and co-occurrence (Tian al., 2023; Yang et al., 2024), are learned from data. There are also multiple works that analyze how high-level capabilities, such as thain-of-thouu Suzuki, 2025), topic clustering (Li al., 2023b; Jiang al., 2024), reasoning or memorization (Yao et al., 2025), and in-context learning capability (Nichani et al., 2024; Bietti et al., 2023; Wang et al., 2024a; Kim & Suzuki, 2024; Edelman et al., 2024), are obtained during training. Although insight- ful, they often assume structured or abstract language data (Li et al., 2023b; Nichani et al., 2024; Yang et al., 2024), unrealistic model architecture (Tian et al., 2023; Cui et al., 2024; Troiani et al., 2025), and adjusted training strategies far from practice (Bietti et al., 2023; Kim & Suzuki, 2024; Huang et al., 2025), which depart from reality. In contrast, our theoretical analysis is grounded in natural language data, realistic architecture, and standard training strategy. As result, our theory substantially reduces the gap between formal analysis and practical use, which further corrobo- rated by our empirical validations.
(Kim &
Understanding feature learning beyond Transformers. Recent work has also explored how models learn data-dependent features through dynamics for non- non-transformer models as well (Dandi et al., 2023). However, this line of work similarly
et al., 2023; Ba et al., 2022; Mousavi- considers abstractions of language, such Gaussian data (Ba al., 2022), single or multi-index models (Damian et al., 2024; Dandi et al., 2023), or spiked models (Wang et al., 2024b; Mousavi- Hosseini et al., 2023), and considers measures of data complexity with Hermite expansions (Bietti et al., 2022; Damian et al., 2024; Lee et al., 2024). On the contrary, we adopt a realistic theoretical setup to analyze features in transformers, which remains the dominant architecture in practice.
PRELIMINARY
3.1 PROBLEM STATEMENT
Semantic associations are foundational for language models: they enable models to go beyond mem- orizing sequences and instead generalize across contexts (Hinton, 1984), infer latent structure (Wu et al., 2018, 1)1,, and generate coherent text. Despite their importance, the mechanisms by which trans- formers acquire these associations during training remain poorly understood. Towards mechanistic and
interpretation of LLMs in a more realistic setup, we pose the question:
How do semantic associations emerge during the training of attention-based language models on natural language data?
It is worth noting that we focus here on general natural language data, rather than synthetically structured or abstractive language, which has been considered in previous works (Yang et al., 2024; Nichani et al., 2024; Huang et al., 2025). This is essential to minimize the gap between our theory and practical use, since real-world text highly diverse and is not restricted to a specific structure. In addition, prior studies (Olsson et al., 2022; Elhage et al., 2021; Nanda et al., 2023) have shown that critical semantic and reasoning abilities, such as induction heads and linear semantic relations, can already emerge in the carly stage and be preserved through convergence. This makes the early stage of training a natural and necessary focus for theoretical analysis, which we now develop.
3.2 MODEL ARCHITECTURE
Prior works have analyzed the training dynamics of attention ditentioo-based models under simplifying sumptions, such as restricting attention to low rank (Cui et al., 2024), removing causal mask- ing (Tian et al., 2023; Yang et al., 2024), without positional encodings (Bietti et al., 2023) or residual
3

---

Published as a conference paper at ICLR 2026
streams (Huang al., 2025). In line with Nichani et (2024), we study attention architec- ture that retains these components: positional encodings, causal masking, and residual streams. To further align with practice, employ relative positional encoding scheme, in T5 (Raffel et al., 2020), rather than augmenting embeddings with absolute position vectors. We begin by introducing the necessary notation before formally defining the transformer computation.
Let V -o,,,.., Ci...., ( denote the set of vocabulary. For an input sequence of length T, we represent the input as matrix X ?TX' where each row of X is the one-hot encoding of the t-th token in the sequence. an L-layer transformer, the parameters associated with self-attention are given by {WO) P(t), V R==1 together with Wo, where W(1) E R|M is the key- cey-query matrix of layer I, v RI x|V| encoding, and Wo E M is the output matrix. The model with input X defined as follows.
is the value matrix, p()) E R RxT is the learned relative positional
Definition 3.1 Antention-Based Transformer). Given an imput matrix X E RTxM,, the L-layer attention-based transformer with parameters θ =(W() p V U (Wo} is defined as
(X ho
where h is defined by the recurrence relation, i.e.,
h ( = h( + SiMask(h (t1) (-1>T () (-(()))))))))))))))))))))))))))) and h = X, (2)
where S() represents the softmax function, DM(u) maps the ith element of to the (-i+ 1)th subdi- agonal, and Mask(-) denotes the operator of attention mask. This architecture is in line with Nichani al. (2024), and recent work shows that self attention-only models can match the performance of architectures with MLP layers (Wang et al., 2025).
3.3 TRAINING SETUP
Learning objective. To align with standard language modeling practice and ensure comparability with prior works (Huang al., 2025; Nichani et al., 2024), adopt the standard cross-entropy objective: given N input matrices XI,..., XN with sequence length and corresponding output matrices Y....., YN. where Y, E TVVV, the objective function is defined as
((0) NT
log (())(X
where M denotes the 1-th row of matrix M and Y!!! corresponds to the one-hot embedding for the & + 1-th token of the sequence corresponding to X,.
Gradient descent. We analyze the evolution of the parameters under full-batch gradient descent with a constant learning rate y). Under gradient descent, the parameters are updated as follows:
@(t) O(t - ve(()..
Due to the nonlinear complexities of the gradient, deriving an exact form for even one of the weight matrices after steps is challenging. We address this challenge by considering a leading- approximation technique, allowing for closed-form expression of the gradients and weights while yielding tight approximation of the full gradient.
THEORETICAL ANALYSIS
In Section 4.1, we provide theorems demonstrating that the weights of inrention-basedd transformers remain close to their gradient leading terms for O(1/1) steps under both zero and Gaussian initial- izations. Then, Section 4.2 uncovers how three basis functions, which are crucial express token associations and language structure, are encapsulated in those gradient leading terms, and how these three functions are compounded to shape the desiderata of the transformers' weight matrices.
4.1
MAIN THEOREMS
Under the setup descrihed in Sec. 3, obtain the following results for attention- based transformers.
4

---

Published as a conference paper at ICLR 2026
Legend
Output Matrix
W RJ
Value Matrix
V RJ
+ B
leal Aependency
bigram (urzt-token) maggnnng
Step
Step ligrum (greviow- taken) mupping
long range dependexcy
tolon
connet nikan
Query-Key Matrix W
Q
context (Sugfi) wapping
Step
interchongrabiity cuntest napping
consal masking
Step mect-t-guryy avcociation wkjf
Figure 2: Illustration of theoretical results. We characterize weight matrices of the attention-only transformer as compositions of three basis functions: bigram mapping, interchangcability mapping, and context mappings. We illustrate how these mappings are composed across weight matrices to learn semantic associations between given query token and its surrounding text.
Theorem 4.1. (Informal) Given an atlention- transformer (Def. 3.1) under sufficiently small Gaussian initialization, with I, S vT/4 after gradient descent steps with learning rate rj Z T if S n min( 8/T 121.) then for all layers! 1 = 1, :.....L, L.
1
Wo sy/BB]
S $30211²
(5)
VI
17 $
B
12s"m²"
(6)
W
3
8
+
2
8
11 Q
< 1,'
(7)
4
3
P
3
11
A
135
(8)
where || 1e is the Frobenius norm, B corresponds to bigram statistic, E corresponds to a context co-occurrence statistic, Q corresponds token-to-token correlation based composition of B and 4b, and A corresponds to relative position correlation based on the same feature as Q
The above Theorem shows that any finite-depth L-layer attention-based transformer (Def. 3.1) has the same characterization for its weights uniformly across all layers under zero-initialization (The- orem D.9) and small Gaussian initialization (Theorem 4. 4.1,, suggesting that all layers of the model capture common associative features from natural language as a starting point before evolving differ- ently as training progresses (Figure 6). As seen in Figure 2, compositions of these features form the leading terms of the output matrix (B), value matrix ($B )), and query-key matrix (Q). We walk through these matrices in Section 4.2.1 and how they form the weights of the model in Section 4.2.2. The formal theorem and proofs are in Appendix D.
4.2 INTERPRETATION OF THFOREMS
In the previous section, we showed that the model parameters can be approximated by key corpus statistics B, Φ Q and A. Now, we discuss the definitions of these statistics by first introducing three basis functions and explaining how their composition characterizes the model's behavior.
4.2.1 THREE BASIS FUNCTIONS SHAPING ASSOCIATIVE FEATURES
(1) Bigram mapping B. The (i.j)-th element in B, corresponds correlation between token and token e, based on how likely e; is to be directly followed by ej as a bigram. More precisely,
B₁ (e(o:)
P.(e:)/(V),
where (e,) is the relative frequency e; over all tokens in the dataset X X and (e |e,) is the relative frequency of e, given that the previous token e;, The product between P. (e,) and (e,),, forms an estimate of the likelihood of e; followed by ej appearing as a bigram and the second term (P((e))/V) simply acts as centering term such that each row sums to
5

---

Published as a conference paper at ICLR 2026
(2) Interchangeability mapping Σg. We study ΣB = B, the correlation matrix of B, which captures correlations between pairs of tokens based on a frequency-weighter similarity of their previous-token distributions. From Eq. (9), the (i.j)-th element of ΣB can be represented
[V]
(((e))
Σ Prlez (e.ee/e,PPP PP) 10,)
(10)
Frequency weighting
Previous token similarity
In essence, Eq. (10) shows that ΣB captures symmetric relationship between tokens based on how similar of function or role they play across different contexts. Specifically, in Eq. (10), can see that the corresponding row, which acts as a feature for token e; captures its associations with interchangeable tokens captured by the previous token similarity factor and frequent tokens captured by the frequency weights. Similarities in previous token distributions are an indicator of functional similarities interchangeability, being preceded by articles or adjectives and objects being preceded by common descriptors. This interchangeability map. ΣB. acts a building block of characterizations for the weights W and p()) as illustrated Figure 2. We depict a simple example word-wise correlation captured by EB in Figure 1.
this captures structural patterns such as nouns
(3) Context mapping क. The (ith) element of & corresponds to a correlation between token e, and ej based on how likely ej is to appear prefix of e;. This can be written
T +
R (the + -th token is es; the m -th token eş) - μīj
Σ
(11)
m=1
where μj centers the columns of Eb to be 0. Considering each row as an embedding for a token e; which represents an av- erage of the tokens that appear in its context, i.e., smoothed context.
the
pond
contains
fish
24
More precisely, the strength of the association from token e; to ej is determined by the average probability that e, appears in the context of e; over possible positions of e; and ey. This matrix can interpreted as assigning a representation to a to- ken based e, appears in. This allows for learning associations between words that capture richer semantic relationships than bigram features. For example, we could expect to see correlations be- tween animal and habitat, country and capital, emotions and facial expressions (See Figure 3). This context mapping T is and value (!) matrices shown in Figure 2.
example of & with arrows pointing to prefix tokens for "fish" with context summary scores edges. Larger values in- dicate the token appears more fre- quently in the context of "fish".
Figure
summary of the possible contexts that token
building block of the gradients for the query-key attention W())
4.2.2 COMPOSITION OF BASIS FUNCTIONS FOR SEMANTIC ASSOCIATION
We now show how these three basis functions, bigram mapping B, interchangeability mapping En- and context mapping 4, compounded characterize four classes of weight matrices of the transformer.
(1) Output matrix Wo- As shown in Eq. (5), B is the leading term of Wo, and thus the mapping from embedding vectors to output predictions can be understood by examining the matrix product B, for a token embedding e,, The j-th element of the resulting output vector is B and each B. includes a factor of (e.). This implies that tokens are scored according to how frequently they occur the average next-token distribution of c; and explain how models at early stages effectively leam bigram- Mike like patterns.
(2) Value matrix VO), The leading term of the value matrix ( can be expressed as as noted in Eq. (6), which acts as composition of context summary and bigram mapping. Because T captures longer-term dependencies and B captures only bigram statistics, the resulting em- bedding from v()) still endows the original token representations with semantic properties similar to those of T seen in Figure 2.

---

Published as a conference paper at ICLR 2026
(3) Attention matrix W Theorem 4.1 characterizes the attention weight (a shared query-key matrix) as Q. which is constructed as a composition of B D the input matrix X,, the output ma- trix Y, etc. We note that this compound feature captures token-to-token correlation determined by how predictive one token is of the other's s next-token distribution based on the context and inter- changeability mappings. We walk through overview of the construction of Q in three steps (See Appendix A for details).
1. Input-outuut matching scoring in context. As a preliminary step, first define a composed feature Σ & by multiplying the interchangeability mapping ΣB with the transpose of the context mapping Φb This composition utilizes local interchangeability to map a token to class of similar tokens and utilizes the context mapping to capture longer-range semantic correlations shared by the set of similar tokens. Using this feature, for each sample, we assign scores between each input and output token.
2. Masking and centering. The auto-regressive constraint is enforced by masking future to- kens, keeping only scores from input tokens that precede the output token. Then, the re- sulting scores for each output token are centered and normalized based on its position.
3. Next-to-query shift and averaging. The scores between each input and output token are then shifted so that the same score assigned instead to be between the input token and the token directly preceding the output token. Then, the scores are averaged across all samples.
(4) Positional encoding P(), The closed-form characterization A of the positional encoding P(!) follows a very similar composition to Q with the main difference being that the correlations are mapped to positional differences rather than to the vocabulary-space differences (See Lemma D.I).
4.2.3 How THE WEIGHTS COOPERATE
To illustrate how the weights work together and provide further context on the role of each of the weights as functions, we consider the leading-term computation of single-layer iitention-based model. Dropping constant factors to focus on the interactions between features, the leading terms of the entire model computation can be written as
(S (Mask XOXT
T + DM(A))) B + +X)
(12)
We can further decompose this into XWo and the computation from the self-attention block is:
S (Mask XQX + DM(A))) X ΣB.
(13)
Q and A capture correlations between two tokens or two positions based on how predictive the first token/position is of the next-token distribution of the second token/position according to (Φ T "EB) T. Notice that the attended tokens are mapped to the output space by Σg. the same feature that determines the correlations for attention. As result, the self-attention block effectively attends to tokens that, under the value and output matrix projection, lead to better next-token prediction. Thus, find that while the residual stream XWo provides average prediction of the next token, Q enables the model to refine this prediction by selectively focusing on tokens most indicative of the next-token given its current parameters, those capturing corpus association statistics.
Implication. By considering an end-to-end analysis of the model under simultaneous train- ing of layers and by decomposing the weights, we obtain clear interpretation of how differ- ent components collaborate to form semantic representations and can rigorously contextual- ize the function of each component in the full computation of attention-based transformers. While these features only yield small changes in the actual text output, they provide impor- tant insight into how the model's 's behavior develops during training. For example, if early training already associates fish with pond (as in Figure 3), we expect such relationships to be useful anchor for later training, allowing the model to complete more complex sen- tences, e.g., "A pond in the garden was filled with colorful fish that sparkled in the sunlight", coherently with learned semantic associations.

---

Published as a conference paper at ICLR 2026
5 EXPERIMENTS
3-LAYER ATTENTION-BASED TRANSFORMER
We begin with an experimental setting designed to closely mirror our theory, enabling direct veri- fication of results and analysis of the semantic relationships embedded in the learned weights. For clearer interpretabilityy we use the TinyStories dataset (Eldan & Li, 2023), truncated to the 3,000 most frequently occurring words, which also defines the model's vocabulary. A 3-layer self-attention model defined in Definition 3.1 is then trained with sequence length T = 200.
Cosine Similarities Large LR
14
Table I: Minimum cosine similari- ties between theoretical and actually learned weights across all epochs. Re- sults from a 3-layer attention-base model trained on TtoyStories (small n)).
Cosine Similarity
0.4
Attention (range across layers)
Value (range layem)
Output trange across layers)
Weights
Min. Cosine
0.0
20
60
100
Attention
0.999496
Epoch
Value 0.999169
Figure Cosine similarity between theoreti- cal and learned weights. Results from 3-layer transformer model trained on TinyStories.
Output
0.998486
Verification of theory. To verify Theorem 4.1, measure the cosine similarity between the learned weights and their corresponding leading terms at checkpoints over the first 100 epochs of SGD using batch size of 2048 for computational tractability with a learning rate of 0.005. We also consider the cosine similarity between the learned weights and their leading terms when using a larger learning rate of 0.05 to understand how features evolve later stages with respect to the leading term gradients. We provide results for both settings in Table I and Figure 4. The results show that the learned weights maintain strong agreement with the theoretical predictions: even after 30 epochs, all weights achieve a cosine similarity of at least 0.9. Moreover, all parameter matrices have a cosine similarity above 0.7, even after 100 epochs, where the loss had dropped from 8.00 to 5.35. These findings suggest that the features predicted by the theorem not only characterize the model dynamics during the early stage, but also remain informative well beyond it. We provide results for BPE tokenization and for a causal analysis in Appendix B, and we elaborate experimental details for the Tiny Stories experiments in Appendix C.
the red to
they happy wanted
fish
flower
birds
park ball the
she happy saw
fish
beautiful
bird
little car play they sad had
big yellow
tree
bird dress go he excited wanted
small butterfly
up
ball balloon be
scared asked
pond hose
park
dog truck help one proud went lake
garden
nest
big blocks see lily angry loved
water
bloom
flowers
tree apples her timmy nice
catch
pretty tweety
man shirt make tom
curious looked sea took boat
daisy field
sky
box hat do her
surprised
flew
(a) Examples for B
(b) Examples for Σg
(c) Examples for &
Figure 5: Selected tokens from the top 30 correlated tokens under different basis features from TinyStories. The characterized features actually capture both grammatical and semantic structures.
Semantic structure. To validate our interpretation of associative features, collect for each token the top 30 most correlated tokens under each of the basis functions: the bigram mapping (B), interchangeability mapping (En), and context matrix (Ø), constructed from the TinyStories corpus. We provide examples of tokens where the expected semantic relationships can observed in Figure 5. Under B, we that the word "red" is correlated with common objects such "truck"
8

---

Published as a conference paper at ICLR 2026
that would be described by the word "red". Under D, we see that the word "fish" is correlated with common settings where fish would appear, such as "pond" or "lake".
5.2 TRANSFORMERS IN PRACTICE
Setup. To evaluate how well our theoretical results extend to practical LLMs, we analyze token relationships learned from WebnWetText (Gokaslan et al., 2019), real-world large-scale dataset with text from millions of webpages, in Pythia- 1.44 (Biderman et al., 2023) and compare them with our theoretical predictions, examining how these relationships evolve across layers on datasets and models reflecting real-world complexities, We choose the Pythia model family, as they are open- sourced and uniquely provide access to intermediate checkpoints, enabling fine- ane-ggained analysis of training dynamics and interpretability (Marks et al., 2024; Gallego- Feliciano et al., 2025). Unlike our theoretical setting, Pythia includes additional components such as MLP and multi-head attention, making it impossible to directly read off average token correlations from the weights. In order to interpret the fayer-wise representations in terms of token- token correlations, we perform the analysis through the following steps:
Gallego-Feliciano
1. We pass in each token as the input to the transformer.
2. For each token and from each layer I, collect the following embeddings: the input to layer hppes the output of the (-th layer ho,t, and the output of the 1-th layer without the MLP component h,,,attn.
3. The embeddings hal,pre form the rows of pre E R Vixd which represents mapping from the input embeddings of layer to tokens. Similarly, the embeddings hi,,.ost and hi,Lottn form the rows of Er.po E RDV and Er,ttnn E RDV)xd respectively.
Attention correlations. To analyze the correlations captured by the attention weights at each layer, compute the product of the key and query mappings for each head and average these products, which we will call Am, E Rdd We then multiply the mapping Er. on both sides of Ar.emb to convert the average attention mapping into token-basis attention weight matrix Ai,tok- Finally, consider token correlations captured by At,o by using its covariance matrix, which we compare with the covariance matrix of Q, the leading-orde order attention mapping term from our theorem.
Embedding correlations. To analyze the correlations captured by the value mapping and the MLP,, consider the token-token correlations captured by the output of each layer. Utilizing the covariance matrix of Ej.post allows for direct comparison with the covariance matrix of the leading value matrix term T B since the matrices themselves have different dimensions. Furthermore, this enables control for shifts in the embedding space.
Comparison methodology. We compute the leading term matrices using 100K samples from WebTeet To control for differences in model architecture, we normalize each row of the leading term weights to have unit norm. Then, we compute cosine similarities between the corre- sponding covariance matrices across layers and across checkpoints. We perform the same analysis on the Web (Penedo et al., 2024) dataset and provide results Appendix B. More details on the experimental setup are in Appendix C.
Cosine Similarity Across Checkpoints
Attention Mapping Cosine Simiarity
No MLP Embedding Cosine Similarity
Embedding Mapping Cosine Similanity
Gosine Smilarty
Gosnee Smilaht
Dumb
Layer Numbos
Layer Num
Tranning Step
Training Step
Training Step
Figure 6: Cosine similarity between covariance matrices for Pythia- attention weights and em- beddings and the corresponding leading term features based on WebTet.
We remind the reader of each layer's 's structure in the Pythia model. The input is normalized and then passed into the attention block and the MI. block in parallel. Then the outputs of each added to the original input.

---

Published as a conference paper at ICLR 2026
Results. We provide visualization of results in Figure 6, where of training, there is very strong agreement between the Pythia embeddings and our teading-temn features. We can that for the embedding mapping. the token representations strongly match our theoretical analysis across all layers, and similarly for the attention weights, excluding only the first layer. We fixed associative features to represent richer knowledge beyond association, starting with the earlier layers. However, it still maintains these features to a large extent for relatively longer steps. This suggests that our analysis on attention- models generalizes with the addition of multi-head attention or MLP and acts as a starting point for finer- grei-ggaindd analysis of full training dynamics.
that, at the early stage
that as the model continues training, the weights gradually drift from
MLP ablation. We perform an ablation at each layer by performing the embedding correlation analysis using Eq,ttnn which is based on only the output of the attention block and excludes the MLP component. The results for this analysis can be seen in the middle plot of Figure 6. We can see that the correlations captured by embeddings with and without the MLP are similar except at the first layer. This suggests that at the first layer, the MLP maps tokens embeddings with structures similar to that of the leading-orde value matrix term and maintains a similar structure at later layers. Based these initial results, one possible hypothesis is that the MLP at early stages functions similarly to the leading-term value mapping.
Individual attention heads. In order capture fine-grained understanding of the attention block, we perform the analysis on attention correlations using individual attention heads. We per- form this analysis an early (Layer 2), middle (Layer 13), and late layer (Layer 24) to also under- stand how heads may evolve differently at different stages of the model. In Figure 7, we find that different layers evolve differently with respect the gradient leading-term for attention mappings. The earlier layers learn the leading-term features at slower rate, as seen by the high similarity (red) appearing later steps, especially for layer We can also see that layer exhibits faster specialization of attention heads than the other layers, seen by the high variance in each column later steps for layer 13. This provides insight into the rate specialization attention heads and suggests that intermediate layers are where specialization initially occurs.
Per-Head Attention Cosine Similarity Across Training Steps
Layer 2
Layer 13
Layar
Sim Iimiionity
Hend Vrdaa
Hereet undees
Head
312
Training
Training Step
Traning
Figure 7: Cosine similarity between covariance matrices for Pythia- 1.4B individual attention head
weights and the corresponding leading term features based Open WebText.
CONCLUSION
We present new theoretical results the emergence of semantic associations in self-attention mod- els learned from a natural language dataset. Our gradient leading term analysis for each model weight illuminates how the core basis functions that shape the associative features. i.c., bigram mapping, interchangeability mapping, and context mapping, develop from the training corpus. We show that transformer weights have closed- foomfoom expressions as compositions of those basis func- tions to represent semantic associations across natural language tokens. The extensive analyses on the weight matrices' characterizations grounded by empirical supports from toy transformers and real-world LLMs contribute to the theoretical foundations of representation learning in transformers while also opening pathways for interpretability research: discovering common factors that allow weight matrices across components to be decomposed into simple functions of those shared factors; leveraging theory to formulate broad hypotheses about how concepts arise in models, extending beyond individual mechanisms specifie behaviors to complex characteristics.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
We provide novel theorem that characterizes the roles of weights in the transformer model, which a de facto standard building block modern LLMs. We try to uphold high standards scien- tific excellence by making minimal assumptions for theoretical analysis while providing practical implications on mechanistic interpretability. The new insights on emerging features we presented contribute better understanding and diagnosis of the representation learning of transformers, which makes a big step towards transparent and reliable AL. As our study considers a setup of train- ing from scratch on public datasets, there is no direct privacy issue and harm. The authors also acknowledge and respect the ethics of confidentiality and faimess, and confirmed that there are no identified violations of them.
REPRODUCIBILITY STATEMENT
All the theoretical analyses in this work are accompanied by full proofs with detailed step-by-step explanations (in the Appendix) for verification, reproduction, and reuse. We elaborate on the details of our setup in the main body of the paper for all empirical validations, and the code is available here. In addition, our choice of models pursues maximum reproducibility and accessibility, given its simple and fully open-sourced configurations.
ACKNOWLEDGMENTS
The authors would like to thank James Oldfiel, Wendi Li, and Jimmy Di for their valuable feed- back. The work is supported in part by the AFOSR Young Investigator Program under award num- ber FA9550-25 1------------- 1-0184, National Science Foundation under awards I1S-2237037 and IIS-2331 Office of Naval Research under grant number N00014-23- 1-2643, Schmidt Sciences Foundation, Open Philanthropy. Alfred Sloan Fellowship, and gifts from Google and Amazon. Zhen Fang was funded by the Australian Government through the Australian Research Council (ARC) under grant number DE250100363. Shawn Im is also supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. 2137424. Any opinions, findings, and conclusions or recommendations expressed reflect the views of the National Science Foundation. Support also provided by the Graduate School and the Office of the Vice Chancellor for Research at the University of Wisconsin-Madiion with funding from the Wisconsin Alumni Research Foundation.
this material are those of the author(s) and do not necessarily
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303. 087774 2023.
Jimmy Ba. Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2390-2430. PMLR, 2023.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in neural information processing systems, 35:9768-9783, 2022.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023.
11

---

Published as a conference paper at ICLR 2026
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. phase transition between positional and semantic learning in solvable model of dot-product attention. Advances in Neural Information Processing Systems, 37:36342-36388 2024.
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen- coders find highly interpretable features in language models. arXiv preprint arXiv:2309.0880 08600, 2023.
Alex Damian, Loucas Pillaud-Vivien, Jason Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2003.05529 7, 2024.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2305. 18270, 2023.
Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. Advances Neural Information Processing Systems, 37:64273 64311, 2024.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.0773: 2023.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021.
Joshua Engels, Eric J Michaud, Isaac Liao, Wes Gurnee, and Max Tegmark. Not all language model features are one-dimensionally linear. arXiv preprint arXiv:2405. 14860 2024.
John Rupert Firth. Papers in Linguistics 1934-1951. Oxford University Press, London, 1957.
Jorge Fallego-Feliciano S Aaron McClendon, Juan Morinelli, Stavros Zervoudakis, and Antonios Saravanos. Hidden dynamics of massive activations in transformer training. arXiv preprint arXiv:2508.036111 2025.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http: //Skylion007 wtt.hhub...//peeeeeeeeeeeee
2019.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
Zellig S Harris. Distributional structure. Word, 10(2-3):146-162, 1954.
Geoffrey E Hinton. Distributed representations. 1984.
Ruiquan Huang, Yingbin Liang, and Jing Yang. aon-asymptotic convergence of training transform-
ers for next-token prediction. Advances in Neural Information Processing Systems, 37:80634- 80673, 2025.
Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, Bryon Aragam, and Victor Veitch. On the origins of linear representations in large language models. In Proceedings of the 41st International Conference on Machine Learning, pp. 11879-21911, 2024.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: nonconvex mean-field dynamics on the attention landscape. In Proceedings of the 41st International Conference on Machine Learning. pp. 24527-24561, 2024.
Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought. In The Thirleenth International Conference Learning Representations, 2025.
12

---

Published as a conference paper at ICLR 2026
Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selec- tion. Annals of staristics, pp. 1302-1338, 2000.
Jason D Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theoretie limit. Advances Neural Information Pro- cessing Systems, 37:58716-53 16-16615601 1-77666, 2024.
Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Inference-time intervention: Eliciting truthful answers from language model. Advances in Neural Information Processing Systems, 11044151-41555 2023a.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 19689-19725) PMLR, 23-29 Jul 2023b.
Samuel Marks, Can Rager, Eric J Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse feature circuits: Discovering and editing interpretable causal graphs language models. arXiv preprint arXiv:2403. 0244403.9977, 2024.
Kevin Meng, Arnb Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. Mass-cditing memory in transformer. arXiv preprint arXiv:2210.07229, 2022.
George A. Miller and Walter G. Charles. Contextual correlates of semantic similarity. Language and Cognitive Processes, 6(1):1-28, 1991.
Alireza Mousavi- Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient- feature learning under structured data. Advances in Neural Information Processing Systems, 36:7 1449- 71485, 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
Eshaan Nichani, Alex Damian, and Jason D) Lee. How transformers learn causal structure with gradient descent. In Proceedings of the 41st International Conference on Machine Learning, pp. 2024.
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 arXiv:2209.1188 2022.
Guilherme Penedo, Hynek Kydliček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https:/ aeperreveewnnne/forroee wtt./ 1d=n6SCkn2QaG.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, (1(110):1-61 2020.
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models a practical size. arXiv preprint arXiv:2408.001 18, 2024.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in neural information process- ing systems, 36: 36:71911111 1-7199777 2023.
Emanuele Troiani, Hugo Cui, Yatin Dandi. Florent Krzakala, and Lenka Zdeborová. Fundamental limits of learning in sequence multi-inde models and deep attention networks: High High-dimensional asymptotics and sharp thresholds. arXiv preprint arXiv:2502.000 00901, 2025.
13

---

Published as a conference paper at ICLR 2026
Mingze Wang. Ruoxi Yu, Lei Wu, et al. How transformers implement induction heads: Approxima- tion and optimization analysis. arXiv preprint arXiv:2410.1477,, 2024a.
Peng Wang, Yifu Lu, Yaodong Yu, Druv Pai, Qing Qu, and Yi Ma. Attention-only transformers via unrolled subspace denoising. arXiv preprint arXiv:2506.032 037990 2025.
Zhichao Wang, Denny Wu, and Zhou Fan. Nonlinear spiked covariance matrices and signal propa- gation in deep neural networks. In The Thirty Seventh Annual Conference on Learning Theory, pp. 4891 4957. PMLR, 2024b.
Yiling Wu, Shuhui Wang, and Qingming Huang. Learning semantic structure prrettrre-eeeeeeeed embeddings for cross- modal retrieval. In Proceedings of the 26th ACM International Conference Multime- dia, MM '18 pp. 825-833, New York, NY, USA, 2018. Association for Computing Machinery. ISBN 9781450356657.
An Yang. Anfeng Li, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388. 2025.
Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, Yinghin Liang, et Training dynamics of transformers to recognize word co-occurrence via gradient flow analysis. Advances in Neural Information Processing Systems, 37:46047-46117, 2024.
Junjie Yao, Zhongwang Zhang, and Zhi-Qin John Xu. An analysis for reasoning bias of language models with small initialization. arXiv preprint 2rX1::2502.04435. 2025.
14

---

Published as a conference paper at ICLR 2026
A DETAILED DESCRIPTION ON WEIGHT CHARACTERIZATION
The token-to-token correlation captured by Q is determined by how strongly correlated one token is with the other's next-token distribution. These correlations are captured by Q; where each element Qie of Q2 measures for X, the correlation between the token at position and the token at position k + 1. This correlation between the token at position j and position k + gets mapped back to correlation between the tokens at positions j and k through X Q:X; let Q given in Eq. (16) be the per-example correlation matrix computed from
token pairs in ith input,
N
Q
1
Σ Q.Q. 1=1
(14)
NT
We walk through an overview of the construction of Q four steps and provide the detailed com- putation in Appendix A.
Feature composition. As a preliminary step, we first define a composed feature Σg & by multiply- ing the interchangeability mapping EB with the context mapping Φ Each entry corresponds to the average product of path weights from token e; to Cj with one step on Σn and one step on $. This composition utilizes local interchangeability map token to its utilizes the context-summary to capture longer-range semantic correlations shared by tokens in the functional class. We will refer the resulting feature as the composed feature for simplicity in the remaining steps.
general functional class and
Scoring input-output pairs. For each input X; and its corresponding output Y,, utilize the composed feature, Eg & to compute correlation scores between input and output tokens as seen in Figure 2.
(Y, - U,)
(15)
where Uo baseline matrix with elements set 1/lV/. This assigns correlation score to each input-outpan token pair according the composed feature.
Masking and centering, The auto-regress regressive constraint is enforced by masking future tokens, keep- ing only scores from input tokens that precede the output token. The resulting scores for each output token are centered and normalized based on its position. matrix is then centered so that the scores for each output token sum to zero, yielding the per-example matrix Q
Q₂ = eint,k, +k (J,, (Y, Uo E
(16)
where J, is the masking operator and ein denotes an Einstein summation.
Next to Query Mapping. Lastly, the scores between each input and output token are then mapped to be the correlation between the input token and the token preceding the output token. In this way, the model learns to attend to the input token when expects the next token to be the output token.
Aggregation across the dataset. Finally, we map per- per-example correlations back to the vocabulary space and average over all inputs and T tokens per input:
N Σ X
Q=
(17)
NT 1=1
In this way, each token is associated with the average correlations other tokens across the dataset.
B ADDITIONAL EXPERIMENTS
BPE tokenization. We train 3-layer attention based model on TinyStories as Section 5.1 using BPE tokenization with vocabulary size of 10,000. We train the model for 10 epochs with learning rate of 0.005 and measure the cosine similarity between the theoretical and actual weights. We report the minimum the 10 epochs in Table 2.
15

---

Published as a conference paper at ICLR 2026
Weights
Min. Cosine
Attention
0.999914
Value
0.998800
Output
0.997891
Table 2: Minimum cosine similarities between theoretical and actually learned weights across all epochs. Results from a 3-layer attention-based model trained on TinyStories and with BPE tok- enization.
Causal intervention. We aim to understand how the model output changes when removing the leading terms from each of the weights. We perform this analysis on the 3-layer attention- transformers trained TinyStories with learning rate of 0.05. Unlike most causal intervention settings, the features considered have a general function rather than a specific function applicable narrower setting, and therefore, expect removing the leading terms result in performance degradation across the dataset. As result, we choose to focus on the extent to which the output distribution changes when the leading term component is removed for each weight matrix. For each weight matrix, we remove the projection of the weight matrix onto its corresponding leading term. After removing this projection, we compute the loss of the resulting model on the dataset. We provide the results of this intervention in Table 3. We can see that the output layer has the largest effect on the loss, while the attention weights have the least. This behavior is predicted by the theory
as the output layer has the largest order update, while the attention weights have the smallest order updates.
Weights
Loss
Original
5.349
Attention Layer
5.350
Attention Layer
5.352
Attention Layer
5.361
Value Layer O
6.192
Value Layer 1
6.526
Value Layer 2
6.520
Output
8.287
Table 3: Loss of the attention-basee model on TinyStories after the leading term component from each weight matrix is removed. The first row corresponds to the original model.
Validation on additional dataset. We perform the analysis in Section 5.2 on the token-token correlations captured by embeddings in Pythia-I.4B except instead of using Open WebText, use Fine Web (Penedo et al., 2024). We provide the results in Figure 8 where we see very similar results as with Textt....x.
WebText.
Cosine Similarity Across Checkpoints (FineWeb)
Attention Mapping Similarity
Embedding Cosine Similarity
Embedding Cosine Similanity
Layer Number
Cottii Smilett
Layer Number
Codnie Smili ft
Number
Tranning
Step
Training
Step
Praining Step
Figure 8: Cosine similarity between covariance matrices for Pythia- 1.4B attention weights and em-
beddings and the corresponding leading features based on Feb.
16

---

Published as a conference paper at ICLR 2026
C EXPERIMENTAL DETAILS
TinyStories Experiments We collect the vocabulary from TinyStories treating each word, punc- tuation mark, or number a token and use the 3000 most common tokens. We then filter out samples that include tokens outside of the set of 3000. For training, we use 65536 of the filtered samples with sequence length at least 201 and truncate all sequences to 201 tokens for training and computing theoretical leading terms. For the BPE tokenization, we tokenize the dataset using a vocabulary size of 10,000, and for training, we use samples with sequence length at least 201 and truncate all sequences to 201 tokens for training and computing theoretical leading terms. We compute the theoretical matrices using the first batch.
Pythia Experiments We use the first 100k samples of Open WebText/FineWeh with length at least 512 characters to perform the analysis.
We utilize A100 GPUs with 80GB of memory. These experiments can be performed with less compute by reducing batch size or sequence length.
D PROOFS
||-I] will be the operator norm unless denoted otherwise.
D.1 PROOF OF 1-LAYER THEBRRM
Lemma D.1 (General Gradient Form), Under the setting described, we have that
Oc
-1
N
h
(1)
R.
(18)
aWo i=1
NT
OC
-1
Σ
N
X (I)T
R
Wo
(19)
(V(I)
NT
Oc -1
X
(J,, v ( x,
(20)
ΦW())
NT
OC = NT einisk OP(1)
D, Σ.Eo ein einGjk,kk VUU xx,,,
(21)
S(Fo(X.)), J; E RTXIXT with Ju.- Diag( (Ag (1)10)) (DIT A (1)1 being the Jacobian of the softmax function for the uth token in the sequence, E XTXXIXI with Dr being matrix with ones along the (-t+ 11th sub aab-diagonal and zeros elsewhere, and ein is used to denote an Einstein summation.
where (1)) S(Mask( ((, XT +
Proof. We start by considering the derivative of the loss with respect to Fo(X,) which is
(X,)
(22)
and derivative of (X.))) with respect to Wo is h
(1)44
Then it follows that
OC = NT
N
(1)
R
(23)
aWo
1=1
Now, we consider the gradient with respect to V() using the chain rule which gives
dc = NT Σ
-1 N
X4A
(1)T
R
Wo
(24)
(V(1)
17

---

Published as a conference paper at ICLR 2026
Now, we consider the gradient with respect A, an intermediate step towards the gradient with respect W p(), Using the chain rule as before, we have
dC
N
NT Σ ROTT
(25)
0A
Letting B = X, p(),, we have that the derivative of A with respect B(1)dd
is
J Diag (i) 1 A
(26)
Then, in order to get the gradient of the loss with respect B,",, we need to consider the contribu- tion from each resulting in the Einstein summation
ac AB!")
NT Σ
Ue R.W VO x,))
(27)
From the chain rule, can derive the gradient with respect to both W1))) and p())
OC
N
x
einggra (J R.W. V()T (OTTxX)Xx
(28)
OWa)
NT
N
OC
einy 5k.
D.
Σ
(29)
0P(1)
NT
where D; has ones along the (-++ + 1)th sub-diagonal and zeros elsewhere. This completes the proof.
n
Lemma D.2 (Softmax Jacobian Norm). The norm of the Jacobian of the softmax function applied to vector with i elements unmasked is at most 1/vi.
Proof. If v is the masked vector before softmax is applied, then the Jacobian S(v), S(v),) as the jthh element on the diagonal and ((v)) for the element in the jth row and kth column with j * k. The norm is maximized when the output is the uniform distribution and is less than 1/Vi.
Lemma D.3 (First Gradient Step). Under the setting described, after one gradient step, we have that
Wo = (B) (11), ((),, p())
(30)
(31)
where B a [V| x [V| matrix where the jth row is the average next-token distribution of the jth token in the vocabulary weighted by the relative frequency of token j across the dataset and centered to have the row sum be O.
Proof. From Lemma 1.1, as the parameters are initially zero, we can see that W(1), V()), p(1) all have gradients of zero and therefore remain as O. For Wo. the value matrix is initially zero, (1) = Xx and as Wo zero, the output distribution for every token is the uniform distribution. Let Uo E TVV represent the resulting output with each element 1/|V]. Then, we have that
ac -1 N awo NT
(Y-Y-UI)
(32)
2=1
We consider the sum of each of the terms X, Y2 and X, Uo. First, we consider the sum of X Y,. The jth row of X,TY Y, is a Vimmmennsoona vector with each element being the number times the corresponding token appears after each occurrence of the jth token in the vocabulary in X,. Then,
18

---

Published as a conference paper at ICLR 2026
summing all and dividing by NT results in each mapping to the average next-token distribution weighted by the frequency of the token corresponding to the row. We can write this as
a 02 P2
B= NT Σ Y
(33)
where 0j is the relative frequency of the jth token in the dataset and Pj is the average next-token distribution for token j. For the sum UO X; divided by NT, we simply get that every row is oj times the uniform distribution over the vocabulary, and we will denote this matrix by U, Then, have that
ac -B-
(34)
awo
and therefore after the first step,
Wo = 7 - U)
(35)
Then, as B = B - U, this completes the proof.
Lemma D.4 (Second Gradient Step). Under the setting described, after two gradient steps, we have that
112
IWo - -2nBB1 S
(36)
viV
20²²
[(a) -np² BT|
(37)
V
W(1) p())= =0
(38)
where B as defined in the previous lemma and 4 is given by
D = P(ex + e,) pagk
(39)
where P(Ek + ez) corresponds the empirical probability that ej is the current token and eg is in its prefix and park k is the value that sets each column sum to O.
Proof. First, as Vi1) remains at zero after the first step, we have that the gradients for W(1,, p(1) are zero and therefore, they remain zero after the second step. We now consider the forward pass after the first gradient step. As the value matrix remains zero, we have that
Fo nX,B
(40)
Then, by the Softmax Jacobian lemma, we have that
V/T VIV
[[8(F(X.)) - Volls S 7
X.B <
(41)
Then, we have that
qVT V
|R, - (Yr - Uo)l S
(42)
and by Lemma 1.1 and that ||X VT, have
0c
N
T
n
+ +B
5
(43)
awo
F
Then, it follows that after the second gradient step.
n
|Wo 2nB
F
(44)
V
19

---

Published as a conference paper at ICLR 2026
Now, consider the gradient with respect to V By Lemma 1.1, we have that
Oc
-1
N nx ATTTT (Y S(F((X.)))
(45)
(A())
NT
and since w
p()) = 0, A, = Ag where the tth row of Ag has the first t elements equal to 1/t and the rest equal to 0. Then, by equation 41, have that
T
|n,, Ao ((F((X))) - nx AJ (Y; Uo) 1 S VIV
|
(46)
Then, using the discrete Hardy's inequality with p 2. we have that ||
and
ac
S 2n²
nx, UG)
(47)
(v())
NT
V
Now, we will analyze
N -Y;
NT Σ xx
(48)
Since B is independent of i, we can move it outside the sum and we can analyze
NT Σ AY, UO)
(49)
We start by consider the form of X7 AJ. Since the tth row Ao has 1/t as the first t elements and zeros for all other elements, we have that the jth element of the tth column of X, A is
meee. t)
(50)
where ej represents the jth token in the vocabulary and Me(e,, t) is the number of occurrences of ej in the first tokens of X,, Letting क'
Ei. x.T Ag Y, we have that
NT
N T
b+11 ex) 7((e))))
$ jk
1(x
(51)
NT
t
Swapping the order of the sums, have
4/4 NT
1xx [+1] =x) ex)r))
(52)
Then, as ()(e)) = E m==1 1(x!m) = ej), we have that
1 T Σ 1=1
N
Φ jk
1
1(x
[+]]
ex, x! =e))
(53)
N
m=1
1 Then, as N En corresponds to an average over the dataset, the average over N corresponds to the empirical probability of having a sequence with the t + 1th token equal to and the mtth token equal Cj, which we will denote as P(r++1 Cky F' = ez). Then, this gives
P(#+1 Ex, Fm 4m = eG)
(54)
T
Then, have average m which results in the average probability that the t 1th token is eg and ej is the first / tokens. We will denote this as P(ey E Iis, It+1 = =x). This gives
T
1/4
P(e E Eist, Tz+1 = ex)
(55)
T
20

---

Published as a conference paper at ICLR 2026
This probability of e, being in the prefix of #++1 e₁ is averaged over the different positions of ex to get an average probability that ej is in the prefix given that εg is the current token, which we will denote as P(e, + Ex. and
8/x = P(e, - ex)
(56)
Now, we consider Up 1 E X AJ VO. Then, we have that
N T Ale,, t)
UP
(57)
NT Am
[V
Rearranging the sum and decomposing Yi we get
3p
1 T
Σ
P(am =es)
(58)
Then, if we consider the average over positions and t, we get the average probability that ej is in the first tokens over all f multiplied by 1/ We can notice that the sum of the jth row of Up and o' are the same. Then, setting
4' = 4' - -Up
(59)
have
dc @V()
2n²²
10'B"" BT
S
(60)
V
F
Then, it follows that after two gradient steps,
2/p²
p 12
S
(61)
V
Defining क- TLT we have
2xy²
- 11 2
B
S
(62)
F
V
This completes the proof.
Lemma D.5 (Third Gradient Step). Under the setting described, after three grudient steps with 1,. we have that
I| 31 F S 372
(63)
((1) -07-3-722 BT
S 217²
(64)
|w(( - -21-
2n
(65)
| - 2n7 -2n0-222-11 2 1/1
S 2nnT
(66)
Proof. First, we start with bounding the norm of Wo, V(1), Ap. We have that
1/
Wollr 2 +
(67)
V
|
27/²
S
2n²²
(68)
VIV
|| 2
(69)
Now, we consider the deviation of the output from the uniform distribution. We start by bounding the norm of X; + A,X,
Xx++ Aox.vay|
< ((1+5n²) VT
(70)
21

---

Published as a conference paper at ICLR 2026
Then, upper bound the norm of FX()
57]
[|Fo(X:))l== S 2 51+ VT S 4nVT
(71)
and then using the Softmax Jacobian lemma, have that
[()(F((X))) Uoll= S 4n/T
(72)
VIV
Then, follows that
dc
+B
+
55)²T
-21
(73)
aWo
NT
VV
and
n
Wo 3n
S
+ 21j² - 31/²
(74)
V
Now, we consider the gradient with respect to ((1), Since
|| ((Fo(X))) - V2T, we have that
dc
4n
5n
n
22n²²
av() + 21/ BT
2
S
S n
(75)
2
IV
Then, we have that after the third step.
-y(1-31)) - BT
[v()
S 2T
(76)
Now, we consider the gradient with respect to W(1), p(1) which according to Lemma 1.1 are
Oc
-1
x
einejk,t (J,, R R.W. V x,)
(77)
aw()
NT
OC -1 ji
N D. D (J, RWW XT)
(78)
OP(I)
NT
We start by analyzing R,W, ( X,T First, We will use that I| Uoll S VT and
| S - max |X||| Aoll |Y, Toll s -
= T=
(79)
We know by the previous lemma and the bound on the deviation of the output that
- xTTT - UO) B |
25n'T
5T v
4n
3T
(80)
VIV
We start by considering the structure of Y,B B We first consider the simpler multiplication of BT B First, we define Σg = B B which has Emm = ΣM af Pmm where aj is the relative frequency of token and Pt is the average next-token distribution token e centered at O.
This corresponds to similarity measure of the previous tokens of em and e, with common tokens more heavily weighted. Then, we have that
VI Σ 912=1
ej DEI
EB P(em - Ck)
(81)
where P(em ex) is the probability that em is in the prefix of eu centered at 0. We can then interpret the each element (Eg 4) j) as a measure assocation between token j and based on a two-step chain of (intercaangeability mapping. suffix token mapping). Essentially, how often does token Eg succeed token e, and similar tokens. We will let G ΣB Now, we can consider 2np YY,,T This results in T T matrix where the jk-th element is P( -2 x(++1l and will denote this Jigk- Then, 2Y, Ux) will have elements centered to have row
22

---

Published as a conference paper at ICLR 2026
sums of 0 and we will let the centered elements be Pijk- Then, consider the einsum of J and 2ny² Go) The tth column of the resulting matrix will the be product of Je and the tth column of 2ny² (Y; Uo) This results in the tth row having the form
S:1e
Pg.
21y 3
(82)
Dije The tth row is Dije for < j I t weighted by 1/t and centered to have a row sum of 0 and we will refer to the centered and weighted elements 21yeqize
qujt and the resulting matrix Qi- Letting
Q = NT Em X, Q,X,, we have that
1
Oc owa)
27
< 2T
(83)
where we have used that the squared Frobenius norm of the einsum is the sum of the norms of each column and that ||Jal| 1/t. Then, it follows that
w(() 2n'² S 27
(84)
Now, we consider the gradient with respect to each element of p(1)
ac
217
N
Tr(D
2nn'VT
(85)
OP!
NT
Since the trace is a linear function, we let Am = Tr(D m NT Σm Qi) and let A be the vector
1
consisting of Ams and we have
OC 0Pa)
+ 2n7 A
2n
(86)
and it follows that
p(i) - 2n (01)-2n22/1 2n
S 2n
(87)
This completes the proof.
Theorem D.6 (Early Stage Features). Under the setting described, for 8 S 17
3, [V| 2 500, we have that after gradient descent steps with learning rate 87373, for T 2
Wo - snBB1
S s3221
(88)
11
B
(89)
2
W(1)
11 Q
S 16
(90)
4
3
F
p(1)
+
4
A
S 65 ST
(91)
3
Proof. We will prove the result by induction. The previous lemmas form the base case. The first phase will be bounding the deviation of the output from the uniform distribution after s gradient steps. To start, we bound the norm of A,, the resulting attention mapping for X,. We know from Lemma D.2, that for each row the attention mapping
(AA; Ao) <
[X,WX) XT + p()):
(92)
23

---

Published as a conference paper at ICLR 2026
Using the fact that there are t elements in (X, + it, t] with magnitude at most maxkm [Well + max, | and applying the inductive hypothesis for p() and Wa), we have that
: )
(max 12kml + max Am )) + 12 km
I(4 - Ao)lt, < 3 +
1
(93)
m
vt
Since the maximum magnitude of an element of B is EK-
km az S and each column of &" T has elements with magnitudes that sum to at most 1, we know that each element of has magnitude at most 1. Then, it follows that maxkm [Q., Then, we have
< 1 and therefore maxkm Qkml Is max Amm 1.
[I4, (elll - (o (0(2) +
17 17
(94)
Vi
Then, summing the upper bounds on the squared norms of each row, and using that > 1/9 S 1 + log we have
A
AollF S
6
+
(5)
n V + 12 / logT
(95)
4
3
Then, we have that
| 2 6
(6(5)
9
(3) 11 1'V + 125 ST / + log'
+
(96)
4
Then, upper bounding (() + 4 (3) by 2s we have
A | s 2 22+25 12 TT V + log T
(97)
Now, using that sy) S
873/5, we have that
3
25 2/ VT - 128 /1 + log T S 2811
(98)
and
|
< 2 + 2sm
(99)
Now, we bound the norm of (() which by the inductive hypothesis we have is at most
8 2
n2222 + 4
(100)
which is at most
22 + +40373
(101)
and since sn S
1 3
|v00/| S 1522n²
(102)
Then, we have that
|x, + 4l x S SV (1 16)2n22)
(103)
Since sn S. min ø7 and by the inductive hypothesis we have that
[(+(( < 2VT (sm 3)
(104)
Since 81) < 873 we have
[F & 4SHVT
(105)
Then, by Lemma D.2, we have
T [(FF(())) Vollr S 4sn V
(106)
24

---

Published as a conference paper at ICLR 2026
Now, we will utilize the bound on the deviation of the output from the uniform distribution well to perform the inductive step for Wo. We have based on the bound that after the (s + 1)th step
dc
8Sn
+
S
1252n1² S 68T)
(107)
aWo
F
Then, after (s + 1) steps, we have that
|Wo (s + 1nbB S 31/²3² 6sn S 317 1)²
(108)
Now, we perform the inductive step for v(1) using the bound on the deviation of the attention pattern and the output deviation. We have that after the (s + 1)th step
ac OV(1) S T BT
< [R: (Y; -o))ll
r | | IWol
|I(2; Uo)ll |4,00 Aoll ||
+ - - Aoll ||al Wo
Blr
(109)
Applying upper bounds and using that [V] 2 500, we have
|| ac lovas + syy BT
S -
24s Tys + 47 + $ $125212
(110)
Then, after (s ++11) 1) steps, we have that
y( po) s (at 1
2 B
$ 48 + 12521j² S 4(s 1)8,7²
(111)
Now, we perform the inductive step for W utilizing the earlier bounds on the output and attention
pattern deviations. We start by bounding the deviation between s² 11 E and W V())T By the inductive hypothesis and S VIV.. we have that
z3 82
W VI
17
S 86 +3 s 1/
(112)
2
Then, for each R.W V X, since I > 500, have that
-² (Y, (Y, Uo)>iEXX F VIVI
20
R.WJ X 2
13s" "T 14s" "T
(113)
Now, order to consider the deviation of the einsum of J and 30 (Y, - Un UO) x we
need to first bound the deviation of the Jacobian of the current attention pattern from J. We do so by considering the deviation for each Jr. As proven earlier, we have that
[IA,; - At) :11 < (6 + +
12 VE
(114)
Then, we have that for the current Jacobian for the sample X, corresponding to the tth row which will call Jei
Jr, - J, - Diag( Aolt, :)) Aolt, [1 - -AU--ADD--A
(A. Ab(4.:)) - (115)
and it follows then that for /2 2
2 + Vt
|| Jella < [| Aolt,
|A Aott, + ||A., A (116)
Then, as
|| [A[ Ao[t.:]]lac $ [|4.114)] -
(117)
25

---

Published as a conference paper at ICLR 2026
have that
|| Jell2 -
(6(2))
(1)))
4
8
17 +
12
3
Vt
1 +
2
+
6
8
"1
12
(118)
Vi
3
vt
Since J1, is always all zeros, we can ignore this term and for 1 2 2, we have that as 8n S
87375
|| Jell2 S 1/
(119)
Then, we have that
8² 5²
17
2
53
x2
|| J1ll
2 (Y, - Uo)
(120)
نو 3
| R.Wo
X,T
(Y; U
2
Then, as || = 1. 1111112 S + 152 VT S 2, and st
<
87375. we have that
11 2Q; ein sty (Ji, VOVT xT))
<5V25117+288''17 s
"0
F
(121)
Then, we have that
ac
83 -² p²
@W(1)
+
2
S <305"107
(122)
F
Then, we have that after (s + 1) steps,
WCO (3
+ (s(x+))
8 ++1 1) 1)
4 + 2
A²
66"'n"T + 30s s 6(s + 1)
(123)
3
F
Finally, have the bound on the deviation from Qi, we have that for P(1)
Qc OP(1)
83 3- 82
1 A
< 30'"'1'T
(124)
and that after (s + 1) steps,
p(t) - 3 8 ++1 1
8 ++1
)) 7'4
+ 2
30s +0 < 6(s + 1)5,
(125)
3
D.2 PROOF OF MULTI-LAYER THEOREM
Lemma D.7 (General Gradient Form). Under the setting described, defining
5(00 einijk, (J(C), ((D)(()) h (T
(126)
d(-1) = G( A 1 c. v(OT + +s( h (1-1 ( sS! h (011)
(127)
with
G = R.W.
(128)
we have that
ac = NT Σ
-1 N
n T R
(129)
awo
1=1
26

---

Published as a conference paper at ICLR 2026
ac
-1
N
(1-1)T
G.D)
(130)
OV
NT
1=1
dc
N
h
[111)
2 1(0) (1-1)
(131)
OW
NT
Oc
= NT finekk,j-i D. 5. 5!!
-1
N
5(1)
(132)
OP()
where A = (MMask(h (1-1) WO + p())). RY -(FoX)(F((X..)) J: E RTXIXT with J. Diag(A - A (1UT A being the Jacobian of the softmax function at the lth attention layer for the tth token in the sequence, D E Tx with De being a matrix with ones along the -tth sub-diagonal and zeros elsewhere, and ein denotes an Einstein summation.
Proof. We begin by considering the derivative of the loss with respect to Fo(X.)) which
OC
Y.
((F)( -R!
(133)
O IX
Since
(X, oWo
h (4)1
(134)
follows that
Oc -1
N NT
h
R
(135)
awo
We now consider the gradient through each attention layer in terms of the current and previous layer
embeddings i (((-1), !0 Let = h (t-1)
(0) G = G aC/OU!" V where
T:(0) V. We have
dC
(Ah) G = A
(136)
av
Summing over i and normalizing yields
Oc 0V
-1 N
Σ h (-1TT A 1 d
(137)
NT
i=1
The gradient for M = Ah is 6m = GV(TT _5A = 5yhh7 ((
AT AT das- Next, through the
row-wise softmax, each row Jacobian is
J = Diag(A (01l)) - (111) (DITT
(138)
Stacking these gives tensor
J." Applying it roww to ò.t gives
s einijk. tk-1 (J G ( a )
(139)
Finally, back back-propagating through = hw + p()) gives
@w = h T (:(0 h,
Ac
(140)
ap() ein - D, (()) s((0)
ac
(141)
Summing over and normalizing,
àc =
-1 N
Σ (111T (()) (1-1)
(142)
aww) NT
1=1
27

---

Published as a conference paper at ICLR 2026
QP()
ac
NT
-1
ein jjk
jk
D. Σ 5/
N
(143)
1=1
Collecting all contributions hg gives
- + T G. V(UTT + 101 +
(144)
Since the last layer,
G L 0c
d
az, Wo R.Wo
(145)
oh: (1)
we can inductively apply the recurrence and collecting the per-layer parameter derivatives gives the desired expressions for ac/oWo. OC/av") ac/aw0, and @L/OP().
Lemma D.8 (First Step. Multi-Layer Zero-Initialin Initialization). Under the setting described, after one gradient step, we have that
Wo = (B
(146)
W v(0 (/())p(' p(() =0
(147)
for S. S.S S. L where B is x matrix where the jth row is the average next-token distribution of the jth token in the vocabulary weighted by the relative frequency of token j across the dataset and centered to have the row sum be
Proof. By Lemma D.7,
Oc -1 NT
Σ
(LIT R
-1
Σ
N
x
(148)
awo
NT
where and U are defined the same as in the one-layer case. A single gradient step gives
dC
Wo = - awo naw = nB
(149)
At initialization Wo = 0, so by Lemma D.7 the upstream gradient from layer is
c (1) R.W.
(150)
Using the recurrence (Lemma D.7),
d!-1) G T
G(u + S "n"nn"")" (w" + + ((""ww
(151)
Since (. = 0 and ( (0) for all I, we inductively get G for every
Now, the layerwise gradients are (Lemma D.7)
ac -1
Σ K 1-1)T
(152)
0V
NT
1=1
and with s(1) einijk. ik (I 1OT
also have S.O
(because G
or (l.0) 0), hence
0c 0W
m NT
N
1-11T I)T
s
(10) (1-1) =
(153)
1=1
dc -1
D. D. ((l)
einijkjkk+
(154)
aP()
NT
1=1
Therefore single gradient step leaves (1)), V(I,, p()) 0 for 1 < 515 I.).
28

---

Published as a conference paper at ICLR 2026
Theorem D.9 (Barly Stage Features, Multi-Layer). Fix depth S L and assume zero initial-
5 ization for all parameters. Under the setting described, for s S n min [2L SVT with T 2 60 and [ > 500, after gradient descent steps with learning rale n we have, uniformly for every layer < 1 < L.
Wo snB
S $3821²
(155)
VO
8
1 B
S 12
(156)
2
W
3
11 Q
S <135
(157)
4
3
F
p())
+
"
A
S 135
(158)
3
F
where B, क. Q and A are as in the one-layer analysis (row-centered bigram matrix, centered prefix-statistics operator, and the third-step structures, respectively).
Proof. We prove the bounds simultaneously for all layers with induction.
By the previous lemma, with zero initialization and one step, Wo =nB, w() = V() 0, p() = 0 for 1 < & L). This gives the base case.
Now prove the inductive step. Assume the four bounds hold after steps, with (s 1)n I min 12L: ) We derive bounds on the deviations in the attention patterns, activations, and
outputs in the forward pass after steps.
Now, we start with a bound on the deviations in the activations from X, each row. Since, each row of A sums 1, we have that
|| s | "
[1,:] - x.((.:]) :]]|
||a5 |105-"0 (t.
|oo
(159)
and
|. |l S (1 | |
(160)
Then, by the inductive hypothesis and sn S T2E we have that || || ||
1
I s Using this and
that h (0) = =X, which has unit norm, we have that across all layers and rows
< (1
3 s, 2
(161)
and as sn S 12T and as (1 + c/L 1 2c for c S 1, we have that
t, [t
$ + 87 4
(162)
Using this and again that n, X,s we have that for all rows and layers,
| - x.(t,:|| :1| s (1+ sn 3 Sn < 252 S
2
81 6
(163)
again using that sy) < T2E.
1
Let Ao be the uniform causal attention with the t-th row having the first t elements equal 1/t and the remaining elements being 0. For each row, of A(D we have that
A (M [1,]] ( (1-1)T (M((P"))
(164)
29

---

Published as a conference paper at ICLR 2026
Decomposing 1 14, :]
[1, (n, (1,,) X:(1, =)) we get by the inductive hypothesis and
that maXkmm
Q km maxm Aml S 1 as shown in the one-layer case that
MASK (1:
(M |1.-1)
5 1/ v (Cw + Cp)s +
(165)
+ 2
h 1 [t, x,[t.
6
+ 4
8
n n'VT
4
3
M u. [1.1 X.
1/" 1/VT
3
By our earlier bounds, we have then
(t.
(1-1) 1)T
+ DM((())
(166)
5212²
S 1/ 265
26
87) "
A +
8 s VTT
36
which we can upper bound by
MA (I- (4.: DM 1)/
(1-IT
S / +
21
(167)
Then, by Lemma D.2, we have
(A - 1e
S 1551/ +
11
3 11 S
75211² 2
<
ST
(168)
2/8
VT
Then, also have that
| - Aol
S 7 S sq
(169)
From the deviation bounds on the activations and the inductive control of Wo.
[Fx()))))> = |1n! Wollr s |14,") s (1+ " >T (sn + 35²12) s 2
(170)
Applying Lemma D.2 gives
T [8())(X)) Voll < V
(171)
V
Then in the one-layer case but using equation state from X
171 and accounting for deviations in the hidden
dc
4s V
B
+ 4sm S 587)
(172)
aWo
F
Then, after s + steps, we have
I - (s (s+1)mm 1 < 3 + +5s7² 3 + 1)
(173)
From Lemma D.7,
h I1I)T 1 R.WO
(174)
0V
NT
Considering the deviation from each of the terms, we have
0V(1) + s7p
B
< 76
(175)
Then, we have that after + 1 steps,
1(0)
ss 1
1
+ B
< 1 + 36' < 12(s 1)373
2
P
30

---

Published as a conference paper at ICLR 2026
From Lemma D.7,
Oc
1
Σ h
1-1
(176)
aw NT
ac
D, D.E 5!!)
eintjk, jk-st
(177)
OP
NT
with 5/0 ein(J; J G V(U ). As in the one-layer bound, can the bound on the
attention pattern control 1 We have that for t 2
2
| Jell2 -
1 vi) [|A:[c:]] Ao(4, :ll12 + [A Ao(4,
(178)
Then, as we have that
87
[A : Aoft, :1l12 I
(179)
VT
follows that Since J1, is always all zeros, we can ignore this term and for 2, we have that as stj S 8/T"
10sn)
Jr Jell2 -
(180)
VT
10s7)
Je - Jell2 =
(181)
VT
Now, we bound the deviation of G.!) from n BO) Starting from layer L, we have
| (074)-s(((e snn(Y, UU)
T V (2sn) + VT(3s21)) S 46 1s2 || S 20m
(182)
We will let the bound on the deviation at layer be Da., Now, we consider the bound for each layer
1
sn - - UO) S Da, + 4s217²
c
2 |/|s (25T) (183)
Since we also need the norm of s(D) to iterate through layers, we bound the norm of S:
(()))
|
5
G
G (22)
2
(184)
$ 8s2n²²
|o
VT
< 55s |
Using this upper bound back the recurrence for DG.,, i. we have
| -sn( - Uo)BT
$ Da. + 4s2n7 | + 40s 77 6
(185)
Using S BV7, we have
|ff -V UO)
< Da. + 14 |c.0|
(186)
We can now write recurrence for |0.00 as ( S sy (V: Uo) BT + Da. we have
|a < (1 + 14s"
- - ||
Da.) < (1 + 14s2
Do.) (187)
Utilizing this with the recurrence for D we can then write a recurrence only in terms of Da and find that for all I, Da. S 1262n2 VT as L < +
vT and sn < min TWL 8VT
5
Then, we also
31

---

Published as a conference paper at ICLR 2026
have that for all 1, |G| I SMVT + 12 2snVT.. Then, we have that ||Jell2
|J:1l2 S 2 10=21p² < 2, and 8/ V min 12L 8VT
12 s( (00)
I 644' 11
(188)
2
This produces
dL
83 - 62 n²Q
< <'4'"11-T
(189)
OW
2
and similarly,
0c OP())
83
82
11 A
S 1445"1²
(190)
F
and hence after (s + 1) steps
w
3
(5 1
(+1s
3 1
11 n
S 1 +1)"18T
(191)
p(0)
(3
8
1+1
1
+
8 8811
1
)
A/A
$ 13(s + 1
(192)
4
Lemma D.10 (Gaussian Initialization Operator Norm). Under the setting described and with all .²
parameters initialized from (0, [ ) for & 2 0 and T s [V]. we have that with probability at least 1 - (3L + 1) exp
(- (-1V11- 1 4 ) for all 1 SISL L.
Woll, |/ol||||w00 ||0||wol|| pm))
3x
(193)
|V|1/2
We start with Wo. Using a concentration bound on Gaussian random matrices, we have that Proof.
P (1w/v1/+5Woll 2 2 -1-1/2
(194)
Then, setting t 1/44, we have that
P [ 2
32
[
I exp
(195)
2
Then, with probability at least 1 1 exp
exp
3v
||ol| S
(196)
VV/1/2
We can apply the same argument for each of V(),WW ( to derive the same bound. Since p()) is smaller than the other matrices and has the same initialization, we can also apply the same bound. Applying union bound on the probability of failures for each of the weights, we have that with
probability at least (3L + 1) exp
-
all Wo, V W(), p(t) have operator norm at 2
3r 1V//2.
most
Lemma D.11 (Gaussian Initialization Frobenius Norm), Under the setting described and with all parameters initialized from NO, ( ) for & 2 0 and T < [V]. we have that with probability at
least 1 - (3L + exp (- - for all 1 115
Wolle.
|
|///
p(0)
S 2v
(197)
p
F
F
32

---

Published as a conference paper at ICLR 2026
Proof. We start with Wo. Using Lemma I from Laurent & Massart (2000), have that
(Iwoll
2²
P
|1|2 + 2t)
se
(198)
Then, setting 1 have that
|
(1Woll? 2 23 exp
(199)
(-IMP+)))
Then, with probability least 1 - exp
|Woll S 2v
(200)
We can apply the same argument for each of ViO, W(() to derive the same bound. For P()), we have
(1 p())
22
T + 2VTT+ +
20)
S se
(201)
|V
Then, setting / = 1/
and using that T
s VV. we have that
(|
2
11/2-2
2 30²
S exp
(202)
F
(- [_
Then, with probability at least 1 - exp
p(t)
2v
(203)
Applying union bound on the probability of failures for each of the weights, have that with probability at least 1 - (3L + 1) exp (- [ all of Wo, V(,,, w p()) have Frobenius norm most 2v.
Theorem D.12 (Gaussian Initialization (Multi-Layer)). Assume the setting of D.9 with depth I. S I [V], and learning rate vT all parameters initialized i.i.d. from N(O,, ((pzzz)) with v I xf 72.
4
n T-1. Then, with probability least
|2|2+25
V 2422
1 (3I. + 1) exp
+ exp
2
for 8 S 1 min
with T 2 60 and VV 2 500, after s gradient descent steps with
12L 8VT
learning rate 1j we have, uniformly for every layer < s
Wo B S 3)²
(204)
V
8
11
P
B
S 12.3'12
(205)
2
W
+
2
17 QQ
< 5T
(206)
3
p())
+
13s
5
"T
(207)
3
F
where B. T. Q. and A in the one-layer analysis (row-centered bigram matrix, centered prefix-statistics operator, and the third-step structures, respectively).
Proof. We start by noting that as long the proof holds when 27 and show that the first gradient step satisfies the inductive hypothesis used in Theorem D.9, then the proof will be complete.
We will prove that the first gradient step satisfies the inductive hypothesis with 72.
33

---

Published as a conference paper at ICLR 2026
We will condition the event that the results of Lemmas D.10 and D.11 holds. Then, our results will hold with probability at least
|V|2+26
V 1/
1 - (3L + 1)
exp
+ +exp
2
4
and we have that at initialization for all
S L
Woll r4 |
| p(t)
31p²
- T V 1/2
(208)
and
[olll p v
[w0)
|
p())
201²
(209)
r
T
Now, we start with a bound on the deviations in the activations from X; each row. Since, each row of A sums to I, we have that
t. x.(e.:] ||| s |0 [e, x,[t, :||| 11155-1-11 " (t.
|oo|
(210)
and
| ||| S (1+ | (t.
(211)
Using that || layers and rows
3n2 S Toiraes and that
n X; which has unit norm, we have that across all
3rp² L
| 16:1 S 1 T2|V
(212)
and as 77 S T27. and as (1 + c/ $1+ + 2c for S 1, we have that
| [4:1 < 1 7/22 2
(213)
Using this and again that o X;, we have that for all rows and layers,
1t lt. - X a|| s L (1 (1+ (1+ 2 T2
1///2
3rd
777/2
S
(214)
3
again using that sn I 12L
Let Ao be the uniform causal attention with the t-th row having the first t elements equal to 1/t and the remaining elements being For each row, of A!") we have that
A (Mask ( (0) (1-1) DMM(P)))..
(215)
By our earlier bounds, we have then
MASK(N) 11)
[t.,,](0A ((-1) + DM(( [t.:)|
1||
7//22
31p²
3np
6/7T/2
(216)
VT
S
2
T
T/IV11/2
VT
where we have use
1 7 < 77 and
1
Then, by Lemma D.2, we have
12r.
7/2
|(4.00 Ab)
(217)
VTt
Then, we also have that
6n7
Ao
<
v + log s 7/
(218)
F
VT
34

---

Published as a conference paper at ICLR 2026
From the deviation bounds the activations and the initial bound on Wo.
1/2/
3r²²
[o(x)) ||,
(1) Wolle z Ie s (1+ 2
VT
12/21/22 S 4n
(219)
Applying Lemma D.2 gives
4nf
[S(F(XX)) - Vollr S VIV
(220)
Then following the argument in the zero-initialization case,
dc aWo
4n
5/2
+B
- 2
(221)
V
VT
Then, after the first step,
3r7²
Wo - BB S
+ 2n1 S 31 S 3rF
(222)
T2|112/2
From Lemma D.7,
Oc
Σ h
IT
TH R.W
(223)
0V(()
NT
Considering the deviation from each of the terms, we have
OC
15n/
T2IV I 157
(224)
@V(
Then, we have that after the first step
2n 2
||
S
+ 15n 11/2 E 33n S $111
F
T2
From Lemma D.7,
dC
1
h (1=1T s(00 (1(-1)
(225)
OW() NT
JC
D. -S
OP()
NT
jin
(226)
with 5:º = ein( J G ( vot h
(-)))) As in the zero-initialization case, we can use the bound
on the attention pattern to control J((). We have that for
11 Jella $
1+ V
[|A,[t,]] Ault. :lll + || - Aoft, :112
(227)
Then, as we have that
7/7
[|4:(1.] Ao(4,
(228)
/T
it follows that Since J1, is always all zeros, we can ignore this term and for t. 2, we have that,
15/77/2
[| Jell2 S
(229)
VT
|| J12 S 15/
(230)
Now, we bound the norm of (1)) Starting from layer L, we have
| S /2T
37²
& 5n)
(231)
T2/V
35

---

Published as a conference paper at ICLR 2026
We will let the bound on the deviation at layer be DG Now, we consider the bound for each layer L.
|
F -
5 D
Da.
T2 1/22 +2|lst"1
37/²
2T
30j²
S (1)+8)/2)) 0/2 Da. + 9n
||/0
(232)
Since we also need the norm of $ to iterate through layers, we bound the norm of 5!!)
1 F -
I|
2
G.I
T/IVI
3nd
V2T & 8n)" DG.
(233)
Using this upper bound back in the recurrence for Da,,, we have
c
" I| s (1+ 8n/ + 72n DG.
(234)
Then for all 1, Da. S 6n² as L S VT and 1 min
5
Then, also have that for
121. 8vT)
all L, || || I 617". Then, we have that as |elll
1
[| I 2 15n7/2 I 2, and syj I
min 127. SVT
(((0)
S 48np
(235)
This produces
ac (W()
S 4811
(236)
and similarly,
0c OP()
I 48np
(237)
and hence after the first step
317²
|| S 48h2
S ST
T2
31²
|p(0
S 48 +
< 4n
72
36


Reviews:
[{"review_id": "bzPL1RfpL1", "reviewer": "Reviewer_bKRV", "summary": "This paper addresses a critical gap in transformer interpretability: how semantic associations (e.g., \"bird\"-\"flew\", \"country\"-\"capital\") emerge during training of attention-based language models (LLMs) on natural language data. Unlike prior work that relies on synthetic data, simplified architectures, or non-standard training, the authors ground their analysis in realistic settings (natural text distributions, standard transformers with positional encoding, and standard next-token prediction loss). Their key technical innovation is leveraging a gradient leading-term approximation to derive closed-form expressions for transformer weights in the early training stage.\n\nThe authors show that all transformer weights (output, value, query-key, positional encoding) can be characterized as compositions of three interpretable basis functions: (1) bigram mapping (captures next-token dependencies), (2) token-interchangeability mapping (reflects functional similarity, e.g., synonyms), and (3) context mapping (encodes long-range prefix-suffix co-occurrence). They validate this theory empirically: on a 3-layer transformer trained on TinyStories, learned weights maintain cosine similarity ≥0.9 with theoretical predictions even beyond early training; on real-world LLMs (Pythia-1.4B trained on OpenWebText), early-stage activations and attention weights strongly align with the proposed basis functions.", "questions": "1. Drift in later training stages: You note that weights drift from leading terms as training progresses (Figure 6). Can you characterize this drift? For example, does it correspond to learning higher-level semantics (e.g., syntax, pragmatics) that build on the initial basis functions, or does it reflect overfitting to idiosyncrasies in the data?\n\n2. Multi-head attention specialization: In Pythia, you average attention heads to compute token correlations, but prior work shows heads specialize in different tasks. Do individual heads align with specific basis functions (e.g., some heads focus on bigram mappings, others on context)? If so, how does this specialization emerge?\n\n3. Causal validation: Your analysis uses correlational measures (cosine similarity) to link learned weights to theoretical terms. Have you tried causal interventions (e.g., modifying the bigram mapping component of the output matrix and measuring changes in next-token prediction accuracy for bigram-dependent pairs like \"bird\"-\"flew\")? Such tests would strengthen the claim that the basis functions are functional, not just correlational.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. Realism of the theoretical setup: By retaining critical components of practical transformers (positional encoding, causal masking, residual streams) and using natural text, the paper avoids the over-simplifications that limit the generalizability of prior work.\n\n2. Mechanistic interpretability: The basis functions provide a causal explanation for semantic association (e.g., the value matrix combines context and bigram mappings to encode long-range semantics), rather than just correlational observations.\n\n3. Strong empirical validation: The experiments are comprehensive, testing the theory on both small controlled models (TinyStories) and large real-world LLMs (Pythia-1.4B), with quantitative (cosine similarity) and qualitative (token correlation examples, Figure 5) evidence.\n\n4. Foundational value for future work: The closed-form expressions and basis functions can serve as a starting point for further research (e.g., diagnosing bias in LLMs by analyzing deviations from the theoretical bigram/context mappings, or designing more interpretable architectures).", "weaknesses": "1. Limited analysis of later training stages: While the paper notes that theoretical features persist beyond early training (e.g., cosine similarity ≥0.7 after 100 epochs in TinyStories), it does not explore why or how weights drift from the leading terms. Understanding this drift (e.g., whether it corresponds to higher-level semantic learning) would strengthen the theory’s completeness.\n\n2. Multi-head attention and MLP layers are understudied: The Pythia experiment adapts the analysis to multi-head attention by averaging attention heads, but it does not explore how individual heads or MLP layers interact with the proposed basis functions. For example, do some heads specialize in bigram mappings while others focus on context?\n\n3. Lack of causal interventions: The empirical validation relies on correlational measures (cosine similarity, covariance), not causal tests (e.g., ablating the bigram mapping component of weights to see if next-token prediction degrades). Causal interventions would more strongly confirm that the basis functions are necessary for semantic associations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uEuBMcb9Jq", "reviewer": "Reviewer_qEwF", "summary": "The paper studies how semantic associations between tokens emerge during the early training of language models. Focusing on a self-attention-only architecture, the authors derive the leading-order gradient updates and show that the learned weights can be decomposed using three quantities that reflect the statistics of the training corpus: bigram mapping (correlations between consecutive tokens), intercheangiability mapping (correlations induced by similarity of tokens’ previous-token context distributions), and context mapping (average contextual features of tokens across their occurrences). They validate the theory on small attention-only models - where the predicted leading terms match the learned parameters in the early steps - and then analyze embeddings in Pythia models, finding qualitative alignment with the theory in the initial training phase.", "questions": "- Th. 4.1 provides formal bounds on the error between the learned weights and their leading-term approximations. How tight are these error bounds in practice? For instance, in the TinyStories experiment.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents an interesting and, to my knowledge, novel blend of mathematical analysis, numerical experiments, linguistic insight, and mechanistic interpretability. The bottom-up approach of deriving representations from first principles of training dynamics is significant.\n- The paper analytically derives a leading-order approximation of the weight evolution and provides clear, intuitive interpretations for each component.\n- The theory does not rely on synthetic data models. Rather, the basis functions are derived directly from the statistics of the language corpora.\n- The paper is well-written and its assumptions are clearly outlined.", "weaknesses": "- The primary weakness is that the theory still relies on strong architectural assumptions. It is derived only for self-attention-only transformers. This analysis does not account for the impact of MLP layers, which are important components of present models. While this is a limitation, it is an understandable simplification necessary to make the analysis tractable, and future work might build upon this to relax these assumptions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kaeGRpVJ0K", "reviewer": "Reviewer_p4Xp", "summary": "This paper studies how semantic associations first emerge in attention-based language models by analyzing early training dynamics. Using a leading‑term expansion of the gradients, the authors derive closed‑form approximations for all major weight classes in a self attention only transformer with causal masking and learned relative positional bias. \nshows very high cosine similarity between learned weights and their leading‑term predictions, and an analysis of Pythia‑1.4B on OpenWebText shows strong early training alignment that gradually drifts laterr. Qualitative examples illustrate that the basis functions recover intuitive relations like \"fish\" <--> \"pond/lake\".", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The three basis functions (bigram, interchangeability provide an intuitive, corpus‑linked explanation for what each weight class is learning and how attention and values cooperate early in training\n\n\n* The writing is very clear\n\n* Near perfect cosine agreement on TinyStories over many steps and reasonable agreement in a non-toy sized model (pythia) support external validity. heatmaps/diagrams are clear", "weaknesses": "* Quantitative validation centers on cosine similarity and selected qualitative token lists. Some broader behavioral evaluations or ablations (e.g., other corpora, tokenizations, or stronger baselines) would further strengthen the claims\n\n* Unless I'm missing something, this doesn't hold well for later stages of training because of drift. That's fine, but I figured I'd raise it. I am not entirely sure how much this is a problem for the scope of the paper. Are there analyses on later stages of training that would be interesting but are blocked by this constraint?\n\n* I would like to see some interventional results showing some practicality: E.g., show that the leading term predictors enable training diagnostics or interventions (e.g., early phase monitoring that forecasts later perplexity or feature formation).\n\n* The theoretical result is clean and realistic for early training, and the empirical results are suggestive. I think the thing holding this paper back right now are scope", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p50fsAp5Nx", "reviewer": "Reviewer_T8ho", "summary": "This paper presents a theoretical and empirical analysis of how semantic associations emerge in transformer-based language models during early training. By analyzing the leading-term approximation of training gradients, the authors derive closed-form expressions for the model's weight matrices. They show that these weights can be characterized by the composition of three interpretable basis functions—bigram, interchangeability, and context mappings—which collectively explain how transformers acquire semantic structure from data. The theoretical claims are rigorously validated through experiments on both the controlled TinyStories benchmark and large-scale models like Pythia-1.4B, demonstrating strong agreement between the theoretically predicted features and the empirically learned weights.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper starts from the perspective of gradient flow and utilizes the contextual distribution in the text to explain the implied semantic associations within the parameters. This perspective is very novel and provides an in-depth way of understanding the distributional characteristics of the parameters of transformer-based language models.\n2. The paper is trained on real text, and the experimental results are highly consistent with the theoretical analysis, demonstrating the reasonableness of its results.", "weaknesses": "1. The paper is validated only on the TinyStories dataset, which I consider insufficient for comprehensive verification.\n2. The training loss on TinyStories remains very high (greater than 5), indicating that the model's learning on this dataset is inadequate and arguably unsuccessful. The value of discussing parameter characteristics under such inadequate fitting conditions is highly questionable.\n3. The description of how theoretical and experimental results are compared in Section 5.2 is difficult to understand. The authors should revise this section and provide more detailed methodological descriptions in the appendix, preferably in an algorithmic format.\n4. The paper uses one-hot encoding instead of the more commonly used embedding encoding. The authors should discuss how this distinction affects their theoretical results, particularly whether the presence of an embedding layer under zero initialization would invalidate their theoretical findings.\n5. Paper [1] also discusses from a gradient flow perspective how the embedding layer is influenced by semantic distributions under small initialization, and its findings share some similarities with certain results in this paper. I believe the authors should discuss the distinctions between their work and [1].\n\n[1] An Analysis for Reasoning Bias of Language Models with Small Initialization, Forty-second International Conference on Machine Learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4NSOiefasA", "reviewer": "Reviewer_XVTW", "summary": "This paper study how transformers learn semantic associations during training on natural language data. The authors develop a theory based on training dynamics, specifically using a leading-term approximation of the gradients to analyze the model's weights in the early phase of training. Their central finding is that the weight matrices, like the output $W_O$, the value $V^{(\\ell)}$, and query-key $W^{(\\ell)}$ can be expressed as compositions of three basis functions derived directly from corpus statistics. These basis functions are bigram maping, that captures the next token co-occurrence ($\\bar{B}$); the interchangeability maping, that fnds tokens with similar preceding token distributions (like synonyms or words in the same grammatical class) ($\\sum_{\\bar{B}}$); and context mapping, that encodes long range co-occurrance between a token and the tokens in its prefix ($\\bar{\\Phi}$).\nThe authors validate this theory first on a 3 layer attention only model trained on TinyStories, and they show that the learned weights have a very high cosine similarity (often >0.99) to their theoretical characterizations. Then, they extend the analysis to the Pythia-1.4B checkpoints, showing that these foundational associations are also learned in larger models, especially during early training and in the later layers.", "questions": "Your leading-term approximation brilliantly explains how associations emerge in the very early stages of training. The empirical results with Pythia suggest these structures persist for some time but also evolve, especially in the earlier layers. What are your thoughts on the mechanisms that drive this evolution away from the initial basis functions? Is it simply the accumulation of higher-order gradient terms, or do you suspect qualitatively new structures, like logical reasoning circuits, are being built on top of this initial associative foundation?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The leading-term gradient approximation is a clever technical innovation that makes an otherwise intractable analysis feasible while maintaining interpretability. The decomposition into three basis functions provides clear mechanistic insight into how different components capture semantic structure.\n- Unlike much prior theoretical work that relies on synthetic data or heavily simplified models, this analysis is grounded in a more realistic setup. \n- The paper provides closed-form expressions for all major weight matrices uniformly across layers, showing how they compose the basis functions differently.", "weaknesses": "- The theory applies to attention-only transformers, but validation on Pythia requires indirect analysis through covariance matrices of activations rather than direct weight comparison. This introduces an additional layer of approximation and makes it unclear how MLPs and multi-head attention affect the theoretical predictions.\n- The theoretical guarantees only hold for $\\mathcal{O}(1/\\eta)$ steps in early training. While empirical results show features persist longer, the paper lacks analysis of how and why these features evolve during later training, limiting practical applicability for understanding fully trained models.\n- To achieve theoretical tractability, the analysis relies on certain assumptions, such as \"sufficiently small\" initialization and bounds on the learning rate and number of steps (as seen in Theorem 4.1). The error bounds for the approximation, while derived, grow over time. This means the exact closed-form expressions are an idealization, and their accuracy naturally decreases as training progresses and the model's state becomes more complex.\n- Experiments use truncated vocabularies (3K for TinyStories, 20K for Pythia) which may not capture the full complexity of semantic associations in real LLMs with vocabularies of 50K+ tokens. The computational tractability argument doesn't fully justify this limitation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)