Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers
Authors: Lei Chen
Joan Bruna
Alberto Bietti
Venue: ICLR 2025 Poster (2025)
Keywords: reasoning
in-context learning
associative memory
transformers
distribution shift
Primary area: foundation or frontier models, including LLMs
TL;DR: We provide empirical and theoretical descriptions of how distributional associations and in-context reasoning mechanisms are learned during training, and tend to be disentangled in feed-forward and attention layers.

Abstract: Large language models have been successful at tasks involving basic forms of in-context reasoning, such as generating coherent language, as well as storing vast amounts of knowledge. At the core of the Transformer architecture behind such models are feed-forward and attention layers, which are often associated to knowledge and reasoning, respectively. In this paper, we study this distinction empirically and theoretically in a controlled synthetic setting where certain next-token predictions involve both distributional and in-context information. We find that feed-forward layers tend to learn simple distributional associations such as bigrams, while attention layers focus on in-context reasoning. Our theoretical analysis identifies the noise in the gradients as a key factor behind this discrepancy. Finally, we illustrate how similar disparities emerge in pre-trained models through ablations on the Pythia model family on simple reasoning tasks.

Published as a conference paper at ICLR 2025
DISTRIBUTIONAL ASSOCIATIONS VS IN-CONTEXT REASONING: A STUDY OF FEED-FORWARD AND ATTENTION LAYERS
Lei Chen'& & Joan Bruna
Alberto Bietti Flatiron Institute
Courant Institute of Mathematical Sciences Center for Data Science New York University
ABSTRACT
Large language models have been successful at tasks involving basic forms of in-context reasoning, such as generating coherent language, as well as storing vast amounts of knowledge. At the core of the Transformer architecture behind such models are feed-forward and attention layers, which are often associated to knowledge and reasoning, respectively. In this paper, we study this distinction empirically and theoretically in a controlled synthetic setting where certain next- token predictions involve both distributional and in-context information. We find that feed-forward layers tend to learn simple distributional associations such as bigrams, while attention layers focus on in-context reasoning. Our theoretical analysis identifies the noise in the gradients as a key factor behind this discrepancy. Finally, we illustrate how similar disparities emerge in pre-trained models through ablations on the Pythia model family on simple reasoning tasks.
INTRODUCTION
Large language models (LLMs) have shown impressive capabilities on variety of tasks, from generating coherent and grammatically correct text, to language understanding and basic mathematical
reasoning (Brown et al. [220 Touvron et al., 2023). At the heart of this success is the Transformer architecture (Vaswani et al.] 2017). which relies on a sequence of self-attention and feed-forward layers to efficiently combine information from the input context and patterns learned from training data, Despite recent progress on interpreting the mechanisms Iearned by different layers (Meng et al. 2022. Wang al.) 2022), these models remain largely black boxes. A better understanding of the role of Transformer layers and how they are affected by the training process could enable new monitoring and editing techniques, better training data, and ultimately more reliable LLMs.
The task of next-token prediction in language modeling inherently involves different subtasks that may be at odds with each other, shown in Figure T For instance, given the context "John gave a book to",, the word "the" is a natural and grammatically correct next word to predict, and relying on global bigram statisties might be enough to predict it given the last word "to". Nonetheless, if another character is present in the context, say Mary, then the name "Mary" may be a better prediction, and this would require a more involved form of reasoning over the context to retrieve this name. In the context of Transformer language models, previous work on interpretability has found that circuits of attention heads seem responsible for such in-context predictions (Wang et al., 2022), while feed-forward layers may be storing more general statistics such as the bigram "to the" or factual knowledge (Geva et al. 2021: Meng et al.) 2022 Bietti et al., 2023). To further strengthen this observation, the recent work (Sharma et al. 2023) found that selectively replacing certain layer weights to their low-rank approximation, particularly late feed-forwar layers, may improve performance on various reasoning benchmarks, and observed that the truncated components were often responsible for predicting "generic" tokens such as the word "the".
In this paper, we provide finer understanding of these phenomena by studying how such mechanisms arise during training, in particular how simple distributionall associations, such as the bigram "to the".
"The implementation is httppp//gggtiib..c///eeeitieeeee 8 ATTTNN Vs_Mi.. MIP

---

Published as a conference paper at ICLR 2025
MLP leams desbblinn
Atontion loarns i-cornesct ressoning
the
distr butonal aessociation
Madrid is located im
Spain
in-oontex: reasontng
Truncsling ML promotes in context reasoning
X
Training sleps
Figure 1: Distributional association v.s. in-context reasoning, In this work, we decompose tasks of next-token prediction into the distributional and the in-context ones, finding that MLPs learn distributional associations before attention develops in-context reasoning capabilities. Furthermore, truncating MLPs promotes in-context reasoning by weakening distributional associations. See Figure 5 for an example of this on the Pythia model [Biderman et al. 2023).
tend to be localized in bedd-fowwrrd layers, while attention focuses on in-context reasoning. We first provide a fine-grained study of training dynamics on a synthetic task with two-layer transformers exhibiting similar properties, where the task in-context recall (Bietti al. 2023) with additional noise on in-context tokens consisting of a fixed "generic" token:
In two-layer model with feed-forward layers (FF), we show that the generic noise token is mainly learned in FF and the attention attends towards correct in-context targets. Removing the feed-forward layers then leads to clean in-context predictions. We provide some theoretical justification through early training steps.
In model without FF, show that the generic noise can be identified in a rank-one subspace of the value matrix in attention block. When the noise level is small, low-rank truncation can filter it out and predict clean outputs.
then investigate such separation between distributional association and in-context reasoning on pre-trained language models, namely the Pythia family, which has checkpoints available at different training steps (Biderman et al. 2023). Overall, we provide a useful description how distributional associations and in-context reasoning mechanisms learned during training, and tend to be disentangled in different parts of the model, such that selectively removing certain components may lead to better predictions in reasoning tasks.
Related work. Sharma et al. (203)) recently empirically observed that a low-rank approximation of some weights in some pre-trained LLMs can improve reasoning capabilities. Several interpretability works have looked at the role of attention versus feed-forward layers for different tasks. The prominence Bood-forward/MLL layers for storing "global" or "persistent" associations or facts has been observed in (Sukhbaatar et al. 2019 (eva et al. 2021 Meng et al. 2022 (iva et al.) 2023). In contrast, several works have investigated the role of attention heads for "reasoning" or computation over the context, e.g., for simple copying mechanisms with so-called induction heads (Elhage et
2021
Olsson et al. 2022 Bietti et al.] 2023), or for more complex tasks (Merrill et al. 2022 Wang
al.
2022 Zhang al. 2022
Liu et 2023 Sanford et al. 2024b).
Training dynamics of transformers and attention have been studied in various works (Snell et al. Jelassi al.] [2022] Li al. 20023) Oymak al., [2023] Tian al. (2023] Bietti et al. [2023
[aadyy]20244 Tian et al. 2024| Zhang et 2024: Nichani et al. 2024| Edelman et al. 20244) In
partieular, the two-layer model and copy task we consider are similar to Bietti et al. (2023) yet their
data model does not involve noise on in-context predictions, and they do not study learning of global associations. Chan et al. (2022); Reddy (2024) study in-context vs. in-weights learning empirically, on different tasks than ours. Cabannes al. (2024) study training dynamics of linear associative memories, but focuses on deterministic data while our setup has generic noise. Training dynamics also studied empirically for interpretability (Olsson al. 2022 Nanda et | [223]] Quirke et al.) 2023) Chen et al.) 2024).. [Edelman et al.) (2022); Bai et al. (222)) Aberrethy et al. (2024) studied sample complexity of self-attention and in-context learning. but did not consider training dynamics.
2

---

Published as a conference paper at ICLR 2025
2 PRELIMINARIES
In this section, we provide some background and motivation on reasoning tasks, and describe the weight truncation technique which we use for ablating weights.
2.1 REASONING FROM CONTEXT
Recent LLMs have shown promising results in more complex "reasoning" tasks which may involve multiple steps of logical or computational processing from context or prompt (Srivastava et al. 2022] Wei et al. 2022 Bubeck et al. 2023 Dziri et al., 2024), as opposed to simple pattern matching or memorization of training data, for instance using learned n-gram predictions.
While it is difficult to clearly separate reasoning from memorization, in this work we will make the simplifying distinction that in-context reasoning involves dependencies between multiple tokens po- tentially far away in the context, while we consider distributional associations simpler predictions that only depend on the last token, e.g., through a bigram model. Thus, due to the residual structure of Transformers, reasoning will typically require using attention operations in Transformers over context, while feed-forward layers should suffice for learning distributional associations. We note that our assumption of distributional associations depending only on the last token mainly for convenience of our analysis, and could be extended to depend on the last token's 's residual stream (Elhage et al.) 202T). which may contain additional information from the context. For instance, this could include previous tokens thanks to position-based attention heads (Voita et al.) 2019 |ilhage et al. 2021 Akyürek al., 2024). which allows capturing n- n-grams instead of just bigrams.
Under this definition, we list few simple examples of reasoning that we will consider below:
In-context recall: when the last token is a, we'd d like copy the token that follows previous occurrences of a in the context. This [.. b layer induction head mechanism (Elhage et al. 2021 Bietti et al. 2023 Sanford et al. 2024a):
a] b pattern typically requires two-
Indirect object identification (IOI): we consider contexts of the form "When Mary and John went to the store, John gave the ice cream to" where the prediction should be "Mary" (IO, the indirect object), instead of "John" (S, the subject). |Wang et al. (2022) found circuit of several attention heads that perform this task by copying the name which only occurs once in the context;
Factual recall: sentences of the form "Paul Citroen is native speaker of"" with target "Dutch" as in (Sharma et al. (023). While this may be seen retrieving distributional association, we will treat here as reasoning since involves combining the subject and relation from the context, while bigram model that only depends on the last token "of" might instead predict the generic word "the".
2.2 TRUNCATING WEIGHTS WITH LASER (SHARMA ET AL. 2023
In order to assess the importance of different weight components for certain predictions, we use the weight truncation technique introduced by Sharma et al. 2023). They observed that reducing the rank of MLP matrices in certain layers of L.MMs effectively brings better performance on several reasoning benchmarks. Their proposed method, Leyer-Selective Rank Reduction (LASER), replaces any matrix in the full model by its low-rank approximation with fraction i.e., a matrix W E pddd.ddn would be replaced by its rank- min dun dowr J approximation via Singular Value Decomposition (SVD). After searching for the best parameters of different models on different datasets, (hhamma et al. (022)) found that applying their method to weight matrices of MLPs on relatively deep layers can enhance in-context reasoning performance on various benchmarks, consistent with our findings. The optimal p is smaller than 0.2 for many datasets.
Another observation from Sharma et al. (2023) is that, when LASER improves the model's prediction on some samples, the full model often predicts "generic" words while the improved model is able to predict the ground-truth answer. For instance, given input "Madrid is located in". the full model predicts "the" while the truncated model predicts the target "Spain" " in Table T Here, the generic word is consistent with definition of distributional associations in Section [2. as it may naturally follow from a bigram distribution conditioned on "in". while the factual answer is more akin to reasoning from context. Thus, we would like to better understand how such modification of
3

---

Published as a conference paper at ICLR 2025
feed- forward layers improves the model from predicting generic words to inferring the answer from context, and how such gap appears during training.
3 TWO-LAYER TRANSFORMER ON NOISY IN-CONTEXT RECALL
In this section, we consider simple onc- or two-layer transformers on an in-context recall task with added generic token noise, which allows to study the trade-offs between MLPs and attention layers for storing in-context versus distributional associations, in a controlled setting. We empirically show how transformers solve this task by storing the generic noise token in feed-forward layers, while attention implements the in-context mechanism. We then provide theory showing that feed-forward layers are more likely store the distributional association (generic token) while attention learns to attend to in-context targets. Finally, we show that when the model has no feed-forward layers, the value matrix in attention stores both in-context and distributional information, in different subspaces.
Data and task. The data model we consider is similar to Bietti et al. (202)). with additional noise. Consider vocabulary V {1.2... 2, ......... N, N + 13. The token 7 _tt N + I the generic noise token. We fix a trigger token q [N], which governs in-context recall, and a context length T. Each sequence of tokens 21:7 [21, 22, ...,. 2T is generated as follows:
i. Sample a correct output token y uniformly in [N].
ii Sample 21:T-1 according to the following Markov process (T are distributions on [N] defined later): 21 ~ (.), and
1 a,
if = D, = otherwise.
(o(()2)), if + q. (?a,());
244 12
Pa,y(2)
otherwise,
0,
iii. Set 2y q, and sample the final output y = 2T+1 Pa.s(.)-
Note that the true ÿ varies across sequences, so that the model needs to infer it from context, e.g., using an induction head as (Bietti al. 2023). Predicting y may thus be seen basic "reasoning" task, yet when training with a > O, the noisy output also requires the model to learn a distributional trigger-noise association, similar to the "of/in the" bigram discussed Section 2
We also consider
using multiple trigger tokens in Appendix B.4 and Figure
Two-layer transformer. We consider a simplified two-layer transformer formulated below. The in- put is a sequence of tokens 21.T = ss.....,T] E [N + 1)", and the output is Ç. The embedding ma- trix WE (N+)xx and un-embedding matrix Wo (N++1)xd are fixed at random initializa- tion. The two attention layers have learnable weights W1 QK. Wt. W2 W E dxddd with (-)) the soft- max on a vector. The two feed-forward layers F1, F2 also learnable, and typically we set them two- layer MLPs with ReLU activation. We will discuss dif- ferent architectural choices of Fi, F22 in Appendix We use the cross-entropy loss to predict = 2T+1 from the logits Ê9 N+1
It WE(Z)) + Pr,
11 Σ [0( W (ort1:)) Wha,,
x7 2 xe +
hit + Fi + ht),
h 4 Σ
(x) T W2r Wi
S5t
x² - mt + 1 Fa + 12))
& Woo?
Table I: Probabilities of the top-5 next-tokens in Pythia- IB before and after LASER. The input prompt "Madrid is located in".. Probabilities of two generic words, i.e., "the" and "a", drop sharply after LASER, while probabilities of meaningful words increase, especially the target "Spain".
"the"
"Spain"
"a"
"southern"
"northern"
Full 0.499
0.079
0.069
0.023
0.021
LASER
0.027
0.300
0.002
0.044
0.046
4

---

Published as a conference paper at ICLR 2025
Layer
Attn-2:
Σ
(Wigh++ + Wy0
FF-2:
train
Layer
W Wa4e)
Wg(0)
W2(
Wg(r)
Wete)
Attn-1
Attn-1
Layer
P=-1
W2(g)
Ps.
WgCF)
Pry-
Wetne>
Ps
Wg(+)
Pr
Wg(g)
Sequence
q
sigper
sumpted
genaric
Figure 2: Noisy in-context recall. Purpose of design: understand mechanisms of attention and feed-forward layers for tasks with in-context reasoning (predict f) and distributional association (predict 7). Task: predict tokens v.s. T from a sentence [..., ........ 4, T, ....., q where is trigger, y is sampled target token for sentence, and T is fixed generic token across sentences. Our findings: in a two-layer transformer, the second-layer attention (Attn-2) only attends towards target tuples (g, 5] while the feed-forward layer (FF-2) learns to predict T.
Margic: max(N
0
p-
MLP2 p=
MLF22
Probablity
Poibabiltt
fall
Marai
ML.2 p ML.2
10"
1000
1990
1900
2000
Thinig Stop
Traninn stop
Trotning sep
Tranting stop
Figure 3: Left three: Average probability of predicting correct and noise tokens, and test loss on clean data (a = 0), with different fractions P of preserved rank in U.n of the second-layer MIP F2. The full model learns to predict noise with probability around
0.5, expected from training data. When F2 is dropped (p = 0), the model predicts the correct token with probability 0.98, Rightmost: the FF-2 margin of 7 v.s. all the other tokens with input as g, i.e., F2(WF (q))), max [a F2 (We))) It reveals that FF-2 learns trigger-noise association early steps.
Experimental observations. Following Bietti et al. (2023). we take Tu and Tr to be the unigram and brigram character-level distributions estimated from the tiny Shakespeare dataset with N 65. The model setup includes = 256 and two-layer MIPs with ReLU for both F,, Fz. The training setup includes batch size 512 and the context length T' 256. When evaluating trained models, consider LASER on the input weight Un of F2. We consider noise level = 0.5 for training data (though any other constant value would lead to similar observations), During test time, set to compute the test loss, aiming to measure how likely the (full or after- (ifec-trunaation model predicts the ground-truth y.
Experimental results are reported in Figure 3 and g The full model predicts noise with probability close to a, which is expected since is trained predict the noise token w.p. a. However, when dropping the second-laye MLP Fz, the truncated model predicts the ground-truth ÿ with an almost perfect probability association "[trigger] + [noise]". Another observation is that the full model first learns to predict the noise with high probability in very early steps, after which starts learning to predict the correct g. which resembles the dynamics observed for learning the "to/in the" bigram in Pythia models in Figure [5] This suggests that learning the (distributional) trigger- association is easier than predicting y, and we will study this theoretically in Section 3.1
0.98. This suggests that F2 responsible for storing the distributional
5

---

Published as a conference paper at ICLR 2025
Train: sccond-layer attention scores
Finetune: second-layer attention scores
0.8
Training step
Training stop
trigger
75
noise
0.4
correct
150
175
0.2
225
225
250
250
0.0
Token position
Token position
Figure 4: The second-layer attention scores of models trained with noise (left), fine-tuned with noise (right, initialized model pre-trained without noise), given the same input. It turns out both models learn to attend to the informative structure "[trigger]+ý" instead of "[trigger]+noiss" This implies that the attention in these models is only responsible to predict y. although the training input and output have noise with probability ((1). The fine-tuning setting is in Appendix B.
After the distributional noise association is learned, we observe slower learning of induction head mechanism, with similar dynamics to Bietti et al. (202)). Compared to Bietti et al. (022)), we notice that the induction head (i.c., the second layer attention head) filters out the noise tokens and only attends non-noisy output tokens following the trigger, corresponding to the correct y. as shown Figure 4 We present theoretical understanding for this mechanism Section 3.2 Figure 2 and Appendix B.2 summarize the roles of all components of the two-layer transformer in this task.
Simplified architecture and data for theoretical analysis. Understanding the full dynamics of the model used in experiments is out of the scope of the present paper, due the many moving parts and the complexity of non-linear MLPs. Instead, we focus on a simpler model involving one linear feed-forward layer and one attention layer, and look at the gradient dynamics near initialization. We consider the following simplified I-layer model. The input 2t E Idd at position
t is defined 2, W ((x) + W (W((2-1), where 2, [N + 1] is the token position t, WE (2x) its embedding and W e (2t-1) is a different embedding of the previous token to a different direction, as in the previous token head construction Biett et al. (2023), where the value matrix remaps the previous token to a different subspace. (Assumption D.I.. which requires large enough d, and holds in the flfmite-widii limit with random embeddings. This model allows us to simplify our analysis by considering a single attention layer with positional embeddings, while capturing the difficulty of long-range range interactions. We note that such simplification is standard the in-context learning literature (e.g., Akyüüek et al. 2023 Mahankali et al. 2024 Zhang et al.) 2024). For data generation, Tu and л. uniform distributions on [N]. Given sequence of inputs, Fi.T E IR Txd the output of model is g & Satn + Err
assume all embeddings to orthogonal
St A WE(2:) + W WE(#-1) Rd la (oT Wok] Wyx, E Rd
IST
(2)
(am(()1:)) Woo(er, ERN-, En(TIT)
WoF(zT) WoWFT N+1
where Wy (N+1))xd is the unembedding matrix, d(s t) the attention module with query s and context f, and F(-) is linear feed-forward layer. This architecture is similar to one-layer transformer, but already highlights the difference between feed-forward and attention layers in a way that expect to still hold for more layers. In the above parametrization. the learnable matrices are W QK. Wr, r. Wr f E ddxd At initialization, we set W Wo.. Wn p. Wy- = 0, noting that random initialization in high dimension would lead to similar behaviors thanks near-orthogonality

---

Published as a conference paper at ICLR 2025
3.1 FEED-FORWARD LAYERS STORE THE GENERIC NOISE
As we saw in Figure 3 and 8 the model very quickly learns to predict the noise token after a few steps. Then the gap between P 0 and 1 in Figure S suggests that the feed-forward layer Fg is responsible for storing the distributional association about noise, which is verified in Figure (miudle)) We now provide theoretical justification for this behavior. In particular, we will show that, at initialization, the gradients over the feed-forwart parameters are much more informative than the attention gradient, which is dominated by noise unless the sample size is very large. This shows that the feed-forward layer is much more likely to capture the distributional association.
We now look the first gradient step from initialization, which has commonly been used to understand feature learning and sample complexity in neural networks (Damian et al. 2022 Ba et al. 2022
Dandi et al. 2023 Oymak al. 2023 Bietti et al. 2023). Note that W OK has no gradient at
iniutaization, that the gradient of W v is most relevant initially (see also Snell et al. 2021 Li ct al.
2023 Oymak et al., 2023; Bietti et al. 2023).
Theorem (Logits after one gradient step). Assume N,T 1,a 1 (().. For the model in Eq
consider one gradient step update from aero-initialization on m i.i.d. samples of 21-T with separate learning rates F5 for W r and nu for Wy v (note that the gradient W qu is zero). With probability
1 - 5, the resulting logits for the feed-forward and attention blocks satisfy, for any test sequence Z
In 2(N-1)
((E(((T))) - TI al S 515.0
171
2(N+1)
In 2(N+1)
A A(Eam (11:))
7o N
I :7 0
TN
+ No
m
m
where A(E) EN+1 max E is the margin of predicting the generic noise token and å (a²9 a(1 i)), where q EisT 1 N + } is the fraction of noise tokens 21:T-
The margin AEE) reflects how much signal there in the logits for predicting the noise token, and the theorem provides concentration bounds on the contributions of the updates on W. and Wy to the margin. Note that q < 1 w.h.p. for large N, T, â = a. We make the following observations:
i. When m (1). there is enough signal in WF to predict the noise, say with my 1, and choice of 7/ ((1) will lead to small but controlled contribution to the prediction from Wy.
ii. When ((N). Wv can also reliably predict the noise by setting Tu ((N) (i.e., with small deviation on the r.h.s.), at the cost of many more samples.
Our result shows that in the initial phase of training, ford-forward layers are more likely to pick up the noise token, leading to a structure of the form WF F ze Wo + +1 1) W (q) while attention will be slower due additional noise and possibly smaller step-sizes. We may then expect the attention layers to focus instead on in-context reasoning, as we observe empirically and discuss next.
3.2 ATTENTION ATTENDS TO IN-CONTEXT TARGETS AND AVOIDS NOISE
When the feed-forward weight learns to predict the noise shown in Theorem 1 Figure 4 reveals that the socond-laye attention in the two-layer model attends only towards the correct tokens. In contrast, model pre-trained without noise has second-layer attention attend towards all tokens just after the triggers (Bietti et al., 2023), as observed in the attention pattern at the first step in Figure rright). Then, after being fine-tuned on noise data, the attention becomes only focused on the correct tokens. Understanding this mechanism requires the analysis of the dynamics of W WQK.
Following the simplified model and data distribution in Eq we take a step towards understanding how attention "avoids" the noise tokens. Concretely, this mechanism appears because, after the initial training phase when FF learns noise association much faster than the attention, Wy has structure of
EKSNA (k) =k)) + W s =k:)), similar to the non-noise setting in Bietti et al. (2023). After such a Wy is learned, the trigger- association provides a stronger gradient signal on W Wor than the trigger-noise association. We show this in the following theorem.
7

---

Published as a conference paper at ICLR 2025
Theorem 2 (Attention attends to in-context targets), Assume N, N.T. > 1 and Assumption F.1 hold. Consider the simplified model in Eq with infinite samples as m + co. After W learns the noise association in Theorem 7 in step the attention weight W QK leamss to attend to positions t E [T] where the correct label follows a trigger word, i.., 2t-1 = 9, 2, y.
More concretely, W QK has the following structure
クタIt - Ek-11 = 2(N > 0. i
4. V,.l,
(3)
Sq C- - -QNN 0, VjSN,
where G-- A WE(1) WoK Wox(Wel) (1) + We denotes the attention logit for different combina- tions of 2t-1 i, 2t = J. with L,j S N N+L
Note that set of logits induces probability distribution via differences between them as exp(E:)/ Σ, exp(E)) 1/ E, cxp(E; Ex). Therefore, the above theorem reveals that the at- tention has two patterns: Eq. (3) shows that W oK prefers attending to locations just after a trigger q- i.e., such that 2 = similar to Bietti et al. (2023), and Eq. (4) shows that among all positions that follow a trigger q. W oK places less attention on the noise token, iee, 2t N + 1, compared to correct tokens 2 I N. Such key difference for attention between noisy and non-noise tasks verifies our experimental observations in Figure
3.3 No FEED-FORWARD LAYERS: VALUE MATRIX STORES GENERIC NOISE ASSOCIATION
In the above discussion, we've seen separate roles of attention and feed-forward layers play to conduct noisy in-context learning. A natural question is, when there is no feed- fordffo layer, how the attention layer stores both in-context and distributional information. Figure 13 indicates that the value matrix stores the noise association in subspace with smaller singular values. In this section, we propose a setting of linear associative memory with noise understand this mechanism.
Unlike Theorem D and showing the separate roles of attention and FF, the attention in a non-FF model has to handle both noise and in-context information once the model is sufficiently trained to reach a global minimum. Due to symmetry from uniformly random sampling ÿ from N tokens, we consider passing the output x E R2 of the attention to the value matrix Wy y and output matrix Wu to predict next-token probability y E RN++ given 21T E IN + 1J" with noise probability of a as follows
x/y, 21.T
W E v + W ( _ W W N+1 N+1),
(5)
Pa (ylg) (1 a)
where W(Z1:) is an aggregate embedding independent of y. When T + co. W(2rT) converges to a fixed embedding W independent of J. so that we may consider simplified model ælğ & W E A Wx N++ with W E [N since W only contributes a fixed offset in all logits that can be easily canceled in the softmax predictions. Therefore, we investigate the following linear associative memory with noise.
Model and data. Consider learnable weight matrix W E Rdxd with d > N. Consider embeddings for N input tokens as {ei) f= C Rd and embeddings for (N + 1) output tokens {u, N++ C Rd. Given any pair of input and output tokens, the associative memory model takes the form fi,jj; W) A (uj. We,), A i [N] [N + 11, as logits to approximate Pa in (5). When k < d, we denote the rank-k approximation of f as (() by replacing W with ((k), where (()) is its rank-k approximation.
Experiments. During training, the dataset D, is generated with non-zero noise probability a > 0. At test time, the dataset Do is without noise as a = 0, so the computed loss is called pure-label loss. The full model is trained with Gradient Descent (GD) subjected to cross-entropy loss. The results are reported in Figure 18 with more discussions in Appendix G.I
Low-rank subspace stores noise. In Figure 18, the rank-1 1 subspace corresponding to the smallest non-zero singular value is responsible to store the noise. We prove this mechanism as follows. Note that, here N = 2 is for simplicity, which is easy to extend to any N > 2.
Theorem 3. Assume Assumptions I.n G.2
G.2 hold, considering N 2 and a E (0.2,0.4), we train the full model f(;; W) with gradient flow. Denote P(i.j: W) as the model's predicted probability
8

---

Published as a conference paper at ICLR 2025
for output j conditioned on input i. Then.forr + oo and i E (1,2), we have
P(i.j: W) (1 a0) 1 = 1) a _ N + 1), P(i,j: w (1-
1 = 0 ) = N
The above theorem implies, the full model always predicts noise w.p. a, while the rank-1 1 model eventually predicts correctly without noise, although training only on the full model with noise. Actually when N 2, the noise is stored in rank-1 subspace and the correct correspondence is stored in rank-(NV - 1) space. Therefore, this explains how the value matrix stores both in-context and noise information when the model is without FF.
4 EXPERIMENTS
In this section, we empirically investigate how LLMs process distributional vs in-context associations, and how this evolves during training. Meanwhile, we provide numerical results of how much low-rank truncation improves complex reasoning on real-world reasoning benchmark, GSM8K. Appendix provides another synethetic IOI dataset that requires counting tokens.
4.1 AN INVESTIGATION ON GPT-2 SMALL AND PYTHIA MODELS
We consider GPT-2 small and Pythia models on the indirect object identification (IOI) and factual recall tasks described in Section 2.1
Quick demonstration: 10I on GPT2 Small. Different from Wang et al. (2022). we would like to consider whether a model proposes an output beyond the input x. A quick demonstration is to consider the 1OI task with input "When Mary and John went to store, John gave drink to The top 4 predicted tokens for GPT-2 Small (Radford et al. 2019) on x are ["Mary", "them", "the", "John")] Although GPT-2 Small successfully predicts Mary (the 10 target) instead of John (S), the other two top candidate tokens, i.e., "them" and "the", do not even appear in the context. This prominence of such "generie" words is similar the factual recall example from Section 2.2 and plausibly follows from a distributional associative mechanism conditioned on the preposition "to".
Comprehensive experiment: IOI on Pythia- Now we would like verify this observation on more models and, more comprehensively, track the behavior of these models along training. We choose to conduct the IOI experiments on Pythia (Biderman al. 2023 a family of models ranging in sizes from 14M to 12B trained on web data, with hundreds of training checkpoints for each size. We generate an IOI dataset of 100 sentences with random names for [IO] and [S] in each sample. Figure reports the test results of Pythia- 1B along training. Here LASER is conducted on MLP weights, with parameters given in Appendix C.2 LASER boosts the probability ratio of [IO] over "the" from 2.3x to 12.3x at 14K steps.
Factual recall on Pythia- IB. As Table 1 we verify factual recall with input as "Madrid located in". The full model of Pythia- generates "Madrid is located in the north of Spain", while the model after LASER generates "Madrid is located in Spain". We track the probability of predicting "Spain" and "the" along training in Figure 5 LASER turns out to boost the probability ratio of "Spain" over "the" from 0.16x to 11.3x at 14K steps. We note that better prompting could avoid the need for LASER in this case (e.g.. "Madrid is located in the country of" predicts "Spain"), but increases the context length and thus the inference cost, though this is outside the scope of this paper.
Training dynamics on Pythia. The behavior of the Pythia models on the 1OI and factual recall tasks during their pre-training process displays several phases, as shown in Figure 5 For IOI, we observe:
i. Initialization: all tokens have similar logits since the weights are random initialized. ii Between 10 and 1000 steps: the models consistently output "the". They cannot solve IOI task at all, as long they have almost the same prediction for [IO] and [S]. After 500 steps, [IO] starts the growth towards one of the top predictions.
ii.. After 2000 steps: Pythia starts to be able to solve IOI task by preferring [IO] than [S] and "the". Meanwhile, the benefit of LASER appears enhancing the leading position of [IO].
Note that here we use "a" store instead of "the" store in the original example of Wang et al. 2022 The reason is rule the word "the" from the input context.

---

Published as a conference paper at ICLR 2025
IOI: average probability
Factual: average probability
10
10
prohahiitt
10
prohability
10
[1O] LAAIIR
10
10
LASER
10
Toa LASER
Tar 141
Spiia (ill Spiial LASTR
Texl LASTR
10"
10'
10"
10
10
10
10'
10'
10'
10'
training sleps
training steps
Figure 5: Left: average probability of tokens [1O], [S] and "the" in 100-sentence IOI task in the prediction by Pythia-IB along training. Right: average probability of tokens "Spain" and "the" in factual task predicted by Pythia- IB along training, with input as "Madrid is located in".. In both tasks, the full model learns to predict "the" with high probability starting from ~0 steps. and then learns to solve the tasks. LASER boosts the probability of correct answers against "the" in both tasks: the average probability ratio of correct answers against "the" improves from 2.3x to 12.3x (in IOI) and from 0.16x to 11 x x (in factual) at 14K steps.
Therefore, the training process reveals the capacity of predicting "the" is learnt much earlier than predicting [IO]. The reason might be that predicting "the" requires simpler grammar structure, while predicting [IO] requires a complicated architecture of attention heads different roles across layer (Wang al., 2002). Then note that the IOI task always has "to" before the masked [IO), which means "to" may be an indicator for the model to predict "the" with non-negligible probability, Similarly, for factual recall see early learning of the "generic" answer, while the factual answer learned later. Conceptually, if LLMs are able to write natural text or have been trained sufficiently with natural texts, it surprising for the model to predict "the" with high probability after seeing "to". This is verified in Appendix
4.2 THE EFFECT OF TRUNCATING SEED-FORWARD LAYERS ON GSM8K
As our previous examples of in-context reasoning tasks are too simple for real-world reasoning, we verify whether truncating MLPs improves reasoning on the GSM8K benchmark (Cobbe et al. 2021). As shown in Table 2 LASER improves the few-shot Chain-of- Thought (Wei et al. (0222) reasoning performance GSM8K when only using 1 or 2 shots, although the performance is worse in the standard 8-shot setting. This suggests that truncating MLPs may help promote in in-context reasoning in more complex settings, perhaps by removing spurious distributional associations.
Table 2: Few-shot accuracy (%) of pretrained and finetuned language models on GSM8K. Truncating MLPs (LASER) improves reasoning performances in few-shot CoT settings while it has worse performance in the standard 8-shot setting. The LASER
in Appendix C.2
1-shot
2-shot
4-shot
8-shot (standard)
Phi-3 (Abdin al. 2024)
56.0
72.2
78.2
82.7
Phi-3 LASER
66.1
74.4
77.0
82.3
Llama-3.1-8B AI@@ Meta
2024)
44.7
50.0
57.6
56.0
Llama-3.1-8B -LASER
46.1
50.7
55.9
53.8
-lmma-3.1-8B-In-e. (AI@ Meta 2024
72.6
74.7
78.5
79.7
-amma-3.1-8B-Insee LASER
73.6
75.6
77.7
77.0
DISCUSSION AND LIMITATIONS
In this paper, we studied the questions of how transformer language models learn to process dis- tributional associations differently than in-context inputs, and how truncating specific weights or layers, particularly feed-forward layers, can help in-context reasoning. While our work provides some initial theoretical understanding of how this may arise simple controlled settings, it would be interesting to study how these ideas may extend to more complex tasks where in in-context reasoning and distributional knowledge interact in more intricate ways.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We are grateful to Yifang Chen, Ekin Akyürek and Denny Wu for helpful discussions. This work was done part while AB was visiting the Simons Institute for the Theory of Computing. LC and JB were supported by the Alfred P. Sloan Foundation, and awards NSF RI-1816753, NSF CAREER CIF 1845360, NSF CHS-1901 1901091 and NSF DMS-MoDL 2134216. This work was supported part through the NYU IT High Performance Computing resources, services, and staff expertise. We thank anonymous reviewers for feedbacks.
REFERENCES
Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, et al. Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404. 14219, 2024.
Jacob Abernethy, Alekh Agarwal, Teodor Vanislavov Marinov, and Manfred Warmuth. mecha- nism for sample-efficient in in-context learning for sparse retrieval tasks. In International Conference on Algorithmic Learning Theory, 2024.
AI@ @ Meta. Llama 3 model card. 2024.
1lama3/
#lob/main/M00D35 CARD md md
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations (ICLR), 2023.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitec- tures and algorithms. arXiv preprint vrXiv:2401.12973 2024.
AI Anthropic. The claude model family: Opus, sonnet, haiku. Claude-3 Model Card, 2024.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 2022.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 2023.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony. Herbie Bradley, Kyle O' 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, pages 2397-2444 2430. PMLR, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 2023.
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. In Advances in Neural Information Processing Systems, 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint ar /2712, 2023.
Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent. arXiv preprint arXiv:2402. 18724, 2024.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. 2022.
11

---

Published as a conference paper at ICLR 2025
Angelica Chen, Ravid Schwartz-Ziv, Kyunghyun Cho, Matthew L Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in mlms. In International Conference Learning Representations, 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, arXiv preprint arXiv:21 10. 14168, 2021.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two-layer neural networks, one (giant) step at time. arXiv preprint arXiv:2305. 18270, 2023.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 2024.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, 2022.
Benjamin Edelman, Ezra Edelman. Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint arXiv:2402.11004, 2024.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfiel-Doode Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer key-value memories. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2021. 2021.
layers are
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. In Conference on Empiricul Methods Natural Language Processing (EMNLP), 2023.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems, 2022.
Diederik Kingma. Adam: method for stochastic optimization. arXiv preprint arXiv: 141.6698.. 2014.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference Machine Learning, 2023.
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts automata. In International Conference Learning Representations, 2023.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent provably the optimal in-context learner with one layer of linear self-attention. In International Conference on Learning Representations (ICLR), 2024.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 2022.
William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, (0:843-856. 2022.
12

---

Published as a conference paper at ICLR 2025
Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. International Conference on Learning Representations, 2023.
Eshaan Nichani, Alex Damian, and Jason Lee. How transformers learn causal structure with gradient descent. In International Conference on Learning Representations, 2024.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfiel-Doddss, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In International Conference Machine Learning, 2023.
Lucia Quirke, Lovis Heindrich, Wes Gurnee, and Neel Nanda. Training dynamics contextual n-grams in language models. arXiv preprint arXiv:2311.00865 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. Technical report, OpenAI, 2019.
Gautam Reddy, The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In International Conference on Learning Representations, 2024.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail to solve the induction heads task. arXiv preprint arXiv:2408. 14332, 2024a.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth. arXiv preprint arXiv:2402.00268, 2024b.
Pratyusha Sharma, Jordan Ash, and Dipendra Misra. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. arXiv preprint arXiv:2312.13558, 2023.
Charlic Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXhv:2103.07601, 2021.
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià Carriga-Alonse et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206.04611 2022.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory. arXiv preprint vrXiv:1907.000 01470, 2019.
Yuandong Tian, Yiping Wang. Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems, 2023.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint 230.:23077 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguisties, 2019.
13

---

Published as a conference paper at ICLR 2025
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593, 2022.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 2022.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context.
Journal of Machine Learning Research, 25(49):1-55, 2024.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301, 2022.
14

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2
Preliminaries
2.1
Reasoning from Context
2.2
Truncating Weights with LASER (Sharma et al. 2023
3
T Two-layer Transformer Noisy In-context Recall
3.1
Fordaforward layers store the generic noise
3.2
Attention attends to in-context targets and avoids noise
3.3
No fordarorward Layers: value matrix stores generic noise association
Experiments
An Investigation GPT-2 Small and Pythia Models
4.2
The effect of truncating feed-forward layers on GSM8K
10
5
Discussion and Limitations
10
Contributions and Implications
17
How Does the Two-layer Model Solve Noisy In-context Recall?
17
B.
Training settings
17
B.2
Summarizing: roles of key components in the two-layer transformer
18
How does attention attend less towards the noise token?
18
B.4
4
Multiple Triggers
22
Architectural Choices
22
B.6
Training Details about Experiments
23
More Experiments on Pythia
23
Learning Association with Prepositions
23
LASER Parameters for Evaluated LLMs
24
Other Pythia models on 10l and More Examples of Faccual Recal
Proof of Theorem]
25
D.I
Gradient for the Feed-forward Matrix W
25
D.2
Gradient for the Value Matrix Wv
28
D.3 Completing the Proof of Theorem
37
Proof for First and Second moments in Lemmal D.2
38
When y
39
When y = 9
45
15

---

Published as a conference paper at ICLR 2025
F
Proof of Theorem 2; Training Dynamics of the Attention Layer
53
Linear Associative Memory
57
G.1 Experiments and Discussions
57
G.2 Proof of Theorem
59
Useful Lemmas
60
Input Examples for LLMs
62
Examples for Prepositions
62
I.2
More Examples of Factual Recal
62
Synthetic 10I Task
63
16

---

Published as a conference paper at ICLR 2025
A CONTRIBUTIONS AND IMPLICATIONS
Our contribution focuses on understanding the different roles of attention and FF weights in disentan- gling distributional vs in-context associations, both empirically and theoretically. The application of low-rank truncation is simply way to verify our claims, and consistent with the findings in the LASER paper that truncating some FF layers may improve performance on some reasoning tasks.
Nevertheless, our perspective based on distributional associations versus in-context reasoning may be helpful in thinking about how to allocate parameters to feed-forward versus attention layers: for instance, in Figure 6 on our synthetic task, we found that for a fixed total parameter budget, models with fewer MLP parameters achieve higher loss on distributional predictions (e.g.. non- contextual bigrams) compared to models with more MLP parameters (and fewer attention parameters). These notions may also provide different way to reason about circuit discovery in mechanistic interpretabilityy from the perspective of training dynamics and properties of the training data. Finally, this disentanglement may inform more effective ways to fine-tune models, e.g., by selectively choosing which layers to fine-tune.
Bigram Loss after 100 Steps
3.50
#Params
0:8
3.25
1:7
4:4
3.00
7:1
bigram loss
2.75
2.50
2,25
2,00
1.75
10"
10
#parameters in total
Figure 6: The training loss of approximating the global bigram Ta with various allocations of parameters in MLP and Attentions. For each configuration of total parameters and ratios, the corresponding best learning rate after search to train 100 steps.
B HoW DOES THE TWO-LAYER MODEL SOLVE NOISY IN-CONTEXT RECALL?
B. TRAINING SETTINGS
In most parts of this work, we consistently train the model with fixed level of a > However, also present numerical results of fine-tur fine-tuning in Figure and 4 to show the mechanism of avoiding generic noise token in the sayend-layer attention. The details of such fine-tuning setting as follows.
Fine-tuning: there are two phases of training as
phase 1 (pre-training) starting from a model with random initialized weights, we train the model data generated with 0. This is exactly the same as Bietti al. (2023). At the end of this phase, the second-laye -lyer attention is expected to attend all tokens after the trigger token, i.e., f. T such that - q matter what 2 is.
17

---

Published as a conference paper at ICLR 2025
phase 2 (fine-tuning); starting from model after phase 1, train all weights in the model on data generated with a > 0. At the end of this phase, the second- attention learns to avoid the generic noise token, i.e., t S T such that 2t N q. shown in Figure
B.2 SUMMARIZING: ROLES OF KEY COMPONENTS IN THE TWO-LAYER TRANSFORMER
Recall the architecture of two-layer transformers in Section 3 as
Jt W
WE(Z) Pt,
1
[o( T W WoR Wi..,
+
hig
+ (x + ht);
1 A Σ
[a W W2 W6
$St
2 A n + h + F2(2) + h2))
&, - W Wo=7.
When the task is without noise, ie., 0. Bietti et al.
2023) point out the first-layer attention attends Therefore, when 2 ÿ with 2 = q. the output of the first layer z! W [0) W W E (g). Then they show that the second- attention
to the previous token through W QK Σ Li= Pt Pt-1P:
IP:
matches such =t with 2T = q by W Wax- (Wy W x (()) (g)) T, through which the information
of ÿ in ii is copied to last token as h R W W (ū). Finally W = Σ=E(M (o(=))x(()
helps output the correct label of y.
In our work with noise a > 0, the key difference is that there is fixed probability for noise token N ++ 1 to appear after each trigger This requires W QK to not only match the trigger but also avoid the noise token after trigger. Let's first summarize the whole pipeline of this model for our task.
Roles of key components. The first layer will be basically the same as Bietti et al. (0033 where
W QK = - 1P,1P, attends to the previous token. Consider two positions 11 12 with 2t,-1
2tp-1 = 9, 2t, J. 2tg N + 1, then outputs of the first layer at these two positions are s!, R WF(U) + W.WWa(p) sta WE(N + 1) + W (z))). Then the second-layer attention W Ook- (W,W WWWWe(() - e WE(N + W) with some positive e makes the attention attend to & and avoid l2 simultaneously, matching with the last token 2y = Therefore, the output of the second-layer attention at T is basically 12g s W (Ø). Similar to the noiseless helps output the correct label of j. Meanwhile, note that
case, W = EEEIN Wo(=)W
#1 actually contains (g through ar, so F2 is able to predict the noise N + 1 when seeing fixed W (g). As result, combining the two streams from h2 and (2 ) the full model is able to predict any w.p. 1 a and predict the noise N + 1 w.p.
Evidence. Figure 4 illustrates that the second-layer attention learns to attend to 2,, = ý and avoid F1 N + 1, with Appendix [B3 presenting primitive exploration on how the avoidance is learnt in a simplified setting. Figure (left) shows the attention pattern from Wax of attending to the previous token. Figure 7 (middle) shows the memory recall of W + 1)" 1)TFF(W F2 ((9)) to predict
the noise. Figure (right) illustrates the memory recall of Wo(() W W e (i) to predict the correct token.
B.3 HOW DOES ATTENTION ATTEND LESS TOWARDS THE NOISE TOKEN?
We use the same simplified model as in Section 3.1 to understand how the second-layer -lyer attention learns to avoid the noise. When using the same learning rate 1) = 77 = My. Theorem I [implies that the feed-forward Wr F makes the most contribution for predicting the noise after the first-step update. Denote the logits for the noise of the model at time as & The arguments in this section make the following assumptions, which hold at least after the first-step update:
i. Wr F dominates the logits & of predicting the noise token, compared with Wy.
18

---

Published as a conference paper at ICLR 2025
First-layer attention scores
Memory recall: noise
Memory recall: signal
current position
input token
20
input token
20
30
10-
40-
40
15
50-
50-
60-
60
20 40
40
60
context position
output token
output token
Figure 7: Left: first-layer attention attending the previous token from the current token. Middle: logits to predict noise from (#2 W B()). Wo(i)) with input i E [VV + 1] and output j E [N + 1], where the output channel 2 set as the noise channel. It turns out, for all input i, the logits on output 2 are large, which matches our construction that, at least for trigger
input, the output 2 has large logits. Right: logits to predict singal from (W W (1), Wu(()) for input i [N + 1] and output
j E IN + 1]. It matches our construction that i = has large logits. Meanwhile, j = 2 does not have large logits since is the noise channel.
Pre-train: prediction diversity
Fine-tuning: prediction diversity
1.0
1.0
0.5
mm
pred: Non-noise
prod: Noise
0.2
0.0
0.0
training steps
training
Figure 8: Fractions of predicting the noise token and the other non-noise tokens with a 0.5. (Left) pretraining steps on noisy data; (right) finetuning steps on noisy data, after pretraining on clean data with a In both cases, the models learn to predict noise with probability nearly 0.5. In the first few (~~ 5) steps, the models quickly learn to predict noise with probability close to 1. The fine-tuning setting is in Appendix B.I
ii. Logits for predicting any h < N is close to 0, which means the predicted probability is exp(Ex)
approximately Pt s N+exp(E.)
iii. The predicted probability Pt < a.
iv. The attention matrix WoR is approximately inducing uniform attention.
v. The dataset has T.N > and - Jo, the gradient is from population loss.
The first assumption holds after the first step from Theorem T with ny Te-
Then, since W
L)W (4) O(+) - W + (V L)W =<g)] for any k: < N
in Lemma D.I. the second assumption holds. Meanwhile, the projection of V ww, L onto any direction
in Lemma D.2 is also smaller than W + +1) (Vw, (q)) by factor of O(/N).
Let's check the condition of the third assumption. In the proof of Lemma D.I the gradient of W
has the form of
Wo(N W-
x()) Pt-
19

---

Published as a conference paper at ICLR 2025
This update induces & increase by 77a pe). This implies
= he-1 + (m(a-- aa-aaaa exp(E:)
N + exp(Ei) ) 4A 1/21. 1.
This sequence ffr has stationary point 5* = log N + log( -.)). Denoting & _ Er - with ε = -5" < O, the iteration becomes
exp(Et) I-a + exp(Ei)
t+ = Ee Eetetn tn
we would like to have & not hit the positive region by controlling 7,, it suffices to bound 77 with any
Ç < 0,
11 Σ
exp(E) T-o +exp(()
where RHS is continuous and decreasing & < 0 when a 0.5. Hence, have v a(I-oj evaluated by L' L'Hospital rule. This bound of v is very strong, since v O(log N) can still have & < 0 after one step.
The fourth assumption is basically from what we will show at the end of this section, the second observation.
Then consider the dynamics of Wv. which is much slower than Wr. From the proof of Lemma D.2 the gradient of Wv satisfies
Σ x )7
W. E
(pw (k/z)
11 -))=))))olll Wk)
T
1
(-V,
(a N7
1s N 11)
N
21
c 1 < - 1ik = N N -
(6)
where the projection WE + 1) always positive and smaller than that on other directions when < a. Projections onto other directions W (u))W Y S N, are smaller as O 22). Finally, let's consider the dynamics of W QQ.. At initialization, WoK due to zero initialization of Wv. After one-step, Wy has such a structure in Eq. Then, with I1T $ Eisier Ty from uniform attention, the gradient of W Wox satisfies
and VWex
- E
Σ (y - (k)
(k())
(oic))vwaa Wvx.) (x. F1. F1.r) W
N
T
Pt
E
N
k=
A
(a -Pt) E
Σ
(Wo(N WW) (W) -(x क.. Wr(g))
2B
20

---

Published as a conference paper at ICLR 2025
Then, have
1
W WW +1) WE(N + 1) (xs mnn)
WE + 1) BW BWE(1)
E
T
1
E
N 1}) cc cla -
WE(N + +1) (æ+ TIT)
T
T
(b)
E
N (c(a- - 1) = 11) WE(N + +1) (1 T1.T)
T
N daa 1) 0 ( 0
1 N2 <.
where (a) is from Eq. 6 (b) is due to 31.T
+ E, 24 and note that = x
Similarly, we also have
T
WE(N + 1 AW (q)) =
( WV (xe - Z14T)
Wc(k)
T
T
= -0( 1
= E
O
= N 1 +1) (xr - BITT)
T
N2
N3
1=1
For any h S N, we have
1
WE(N) BW BW =
W Wo(N + w W WE(k) (x, T1r)
T
=E
(cla 1) 1/z = k}) WE(N +1) (z+ 11T)
1 ) 0 \タ ) 00
N caa (o(--)(--
and
T
WE(R) AWe(g) = E
(W)) WE(k) (x. - Tnr))
T
1
= E
θ
1fa = N N W)))) (z, B1TT) N2
0 N
T
Combining the above four esimation of projections of A and B with Eq. 7 we have
WE(N +
(()) = 0
N22 0,
V (-V
O N3 0.
Then have three observations
i. Wor in this phase avoids the noise token + 1 and uniformly attends to all tokens k S N. is O N ), while the update of Wrr is ((1) in Lemma D.1 and that of Wy is O N in Lemma D.2 These three levels of updating speed also coincide with the assumptions that W. F dommates first and then Wy has micro structure that induces the
ii. The update of W QQ Wok
evolving of W QQ.
iii. The current proof for Wok strongly depends on the fact that the noise token appears less than other token by factor in expectation. The proof will have the opposite result if the noise token is made to appear more by manipulating the data distribution. Therefore, we leave a new proof that is robust to such assumption in data distribution future work.
21

---

Published as a conference paper at ICLR 2025
B.4 MULTIPLE TRIGGERS
In Section 3 we assume there is only one fixed trigger q E [N] for simplicity. Actually the case of multiple triggers has the same mechanism. As discussed by Bietti al. (2023) and Appendix [B.2] for one trigger, the second-layer attention has large logits in (W Wel W QK W We()) only for
i = q. For multiple triggers, basically (W/WE W QK W We(j)) only have large values
when E Q. This is verified in Figure 9
One trigger: attention
Five triggers: atfention
10
10
input token
input token
50
60
0
20
40
60
0
20
40
60
output token
output token
Figure 9: Logits of (W; W W2x We(i)) for input and output j when there is one trigger (left, 4 = 1) and five triggers (right, E Q = 39, 43, 53, 58)). both cases, the logits only have large values when i g, verifies the matching mechanism in Appendix B.2
B.5 ARCHITECTURAL CHOICES
In Section 3 and Appendix B.2 we were focused on experiments with both Fi, F2 being two-layer ReLU MLPs. Meanwhile, have also tried other choices of F1, F2 and then search for the best truncation method for each architecture. In this section, we would like to summarize our experimental results for better understanding of all modules in the two-layer transformer.
Generally, the feed-forward layer can be two-layer ReLU MLPs, one-layer Linear or "None", where None stands for there is no feed-forward layer so that the value matrices in attention layers the only weight matrices that transform features.
Both Fi, F22 are two-layer MLPs. This our main setting. The best truncation method is fully drop F2. We also try to fully drop Fi, as reported in Figure 10 It turns out fully dropping Fi makes the model predict the noise with high probability.
Fi is MLPss and F2 is Linear. Figure I1 reports the results. Dropping Fi and F2 both improve the correct prediction, and dropping Fi is better with lower test loss. Note that, when test accuracies are near 100%, lower test loss is a better measurement of the prediction quality, because accuracies are taken by argmax over the output logits while test loss are about the exactly predicted probability.
Fj is Linear and F2 is MLPs. Figure 12 prediction while dropping Fi makes the model predict noise more.
reports the results. Dropping Fz improves the correct
Both Fj and F2 are None. Figure 13 reports the results. While there is no feed-forward layer any more, low-rank truncating a part W of the first-layer matrix improves the model's prediction a little. This implies that, when there is not feed-forward layers, the noise association is possible stored in the first-layer value matrix attention. Note that the improvement such low-rank truncation is clearly smaller fully dropping one of fred-rdrward layers in the previous cases. Meanwhile, smaller p = 0.01 destroys the model's 's performance. This implies fully dropping is not the optimal choice for low-rank truncation of the value matrix, and there is low-rank subspace in it that is useful for predicting the correct tokens. Our discussion of the role of W in Appendix B.2 is a possible answer this phenomena.
22

---

Published as a conference paper at ICLR 2025
Both MLPs: Test accuracy
Both MLPs: Test loss
1.0
Tost accuracy
0.6
drop F.
Loss
10"
drop
drop F 2
drop
0.4
fill
full
0.2
10"
200
400
600
0
200
400
800
Training stop
Truining step
Figure 10: Test performance of fully dropping Fi, F22 when both F1, F2 are two-layer MLPs. turns out, while dropping F/2 makes the model predict correctly w.p. near 1, dropping F has the model predict noise with high probability.
ML.Ps Linear: Test accuracy
MLPs Linear: Test
1.0
0.8
Test accuracy
drop F_
10a
drop F_
drop F
Loss
drop P_
0.4
fill
full
0.2
10"
2000
4000
6000
8000
0
2000
4000
6000
8000
Training step
Training stop
Figure 11: Test performance of fully dropping Fi, F2 when both Fi is MLPs and F2 Linear. Both dropping methods turn out to help predict more correctly than the full model. Meanwhile, dropping the MLP Fj is better with lower test loss.
B.6 TRAINING DETAILS ABOUT EXPERIMENTS
of the training is with SGD optimization with learning rate {0.001, 0.033. The batch size is 512. The dimension is 256. The context length is 256. All results in the experiments are stable for any learning rate between 0.001 and 0.03. Each run of experiments is on single Nvidia Tesla V100 GPU. It takes 3 hours to finish each run for 2K steps, which probably can be optimized lot since are tracking lot measurement along training, not limited to hundreds of possible truncations at each test time.
C MORE EXPERIMENTS ON PYTHIA
C.I LEARNING ASSOCIATION WITH PREPOSITIONS
We would like to verify our guess about the structure of "to + the" in Pythia in Section 4.1 To make the argument generalizable than IOI dataset, we consider structure of "[preposition] + the", where [preposition] has a pool of 30 prepositions in English, including "to". The input is raw "[preposition]" random sentence ending with "[preposition]" with some examples in Appendix [.]]] For both kinds of inputs, Pythia-1660 160M/410M/BB turns out to learn the structure of "[preposition] + the" around 10 steps, shown in Figure 14
23

---

Published as a conference paper at ICLR 2025
Lincar MLPs: Test accuracy
Linear + MLPs: Test loss
1.0
Tost accuracy
10"
0.6
drop
drop
Loss
drop
drop F
0.4
full
10"
0.0
2000
4000
6000
8000
0
2000
4000
6000
8000
Training stop
Truining step
Figure 12: Test performance of fully dropping F1, F2 when both Fi is Linear and MLPs. Only dropping P/2 helps predict more correctly. Dropping Fi makes the model predicting noise more.
No Feed-forward: Test accuracy
No forwardr Test loss
Tost accuracy
0.6
W p 0.05
Loss
W& p 0.01
0.4
10
full
W 0.05
0.2
wb:p=0.01
full
0.0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Training stop
Training stop
Figure 13: Test performance of low-rank truncating of W when there is no Fi, F2. Here p is the fraction of preserved rank W where actually we re-parametrize attention as W W. E RRdxd It turns out the best p == 0.05 improves the model's ss prediction little. Mcanwhile, smaller p destroys the model's performance.
first-layer value matrix in
C.2 LASER PARAMETERS FOR EVALUATED LLMS
Following the definition of LASER in Section 2.2 we search for the optimal layer, p and target weights in Pythia models and GPT-2 Small for each dataset.
I0I on Pythia-410M. The model has 24 layers. The truncation is the input matrix of MLPs the 22-th layer with = 0.02.
I0I on Pythia-11 1B. The model has 16 layers. The truncation is on the input matrix of MLPs on the 1 11-th layer with 0.008.
Factual recall Pythia-1B, 1B. The truncation is the input matrix of MLPs the 16-th layer with 0.0125.
Factual recall on Pythia-1.4B. 1.44. The model has 24 layers. The truncation is the input matrix of MI.ss on the 24-th layer with p = 0.025.
Factual recall on Pythia-2.8B. The model has 32 layers. The truncation is the input matrix of MLPs on the 32-th layer with p = 0.04.
IOI on GPT2 Small. Related parameters have been contained in Section
Phi-3 on GSMSK. The model has 32 layers. The truncation on the output matrix of MLPs on the 28-th layer with p 0.02.
24

---

Published as a conference paper at ICLR 2025
average ranking of' the': input prepositions
average ranking of' the': input sentence prepositions
10"
10
10
rank
10"
160M
1M
10
10"
1B
10"
16'
10
10'
10
10
10'
10'
100
10"
10"
10"
training steps
training steps
Figure 14: Average ranking of tokens "the" in the prediction by Pythia- 160M/4100/IBB along training. The inputs are 30 preposition words (left) and 40 sentences ending with prepositions. It turns out "the" becomes one of top predictions around 10 steps.
Mlama3.1-8B-instruuee on GSM8K. The models have 32 layers. The truncation is on the output of MLPs on the 27-th layer with p = 0.02.
C.3 OTHER PYTHIA MODELS ON IOI AND MORE EXAMPLES OF FACTUAL RECALL
IOL In the same setting of Figure 5 (left), plot the prediction distributions of Pythia-4 10M and IB on the 100 IOI inputs in Figure T5 The model checkpoints are the final ones after training. LASER turns out to decrease the probability of "the" while keeping that of the correct [IO] high.
More examples Factual Recall. In additional to the factual query "Madrid is located in" in Figure 5 (right), consider more such examples in Table 5 We plot the prediction distributions of Pythia- 1B, 1.4B and 2.8B on these inputs in Figure 16 where LASER significantly lowers the probability of predicting "the" vesus the correct outputs.
D PROOF OF THEOREM
In this section, we will present the expectations and variances of V i and Vw with Wy W p at initialization. The targets are show:
1. gap between limm-xo Ww L and lim Vw, É that step of GD with large learning rates is enough to learn the noise in Wr, and
2. sample complexity of Vw, and Vw, f. based on expectations and variances.
Assumption D.1 (Orthonormal embeddings). The embeddings uy Rd are assumed be orthonor- mal, i.e., Uj = 1ii = j}. Meanwhile, if a matrix W E R is random initialized, it holds us Wu, = 0.
D.1 GRADIENT FOR THE FEED-FORWARD MATRIX WF
Lemma D.1. Consider zero initialization, Wy W. Wok = and N > 1. Then with probability 5, for any j. [N 1), holds
We (k) T (Vw, L) W) (()) ((+)
(8)
S 402(k) (In(N + 1) + in( m
4R(k) (In(N + 1) In(
m
where o( R(k) are expectation, variance and range for different choices of E [N] as follows:
25

---

Published as a conference paper at ICLR 2025
Pre-trained Pythia-410M evaludated on 101
1.0
0.8
probubility
0.6
0.4
0.2
0.0
[[o]] full
[10]: LASER
[S]: full
[S] LASER
'the': full
'the': LASER
Pre-trained Pythia- 1B evaludated on 101
1.0
0.8
probability
0.6
0.4
0.2
0.0
[10]: full
[10]: LASER
[S]: full
[S]: LASER
'the': full
'the': LASER
Figure 15: The prediction distributions of Pythia-410M and IB on the IOI task. The setting is the same as in 3 The evaluated models are the final checkpoints after training. LASER turns decrease the probability of "the" while keeping that of the correct [IO] high.
M(N + 1) p(k) N
2 + = a),
R(N + + maxfa, a), R(k) = 1.
VKSN:
- N
(2(k) L-a
N
Proof. Due to initialization, i.e. Wv WF F 0, the current predicted probability is pw(kla,) N+I for all i [m] and & E [N + 1). Therefore, from Lemma H.l we have
N+1 Σ
Vw, i. 1
16, - kb] Woo))a7
N+1
where F.T E Rd W W2(2r) + Pr is the input embedding with input token 21, at position T in sequence i, together with positional encoding pr for position T. Since z,, is set be the trigger in the data generation process and pr is assumed to orthogonal to any other vector in W g in Assumption D.1 we have the following projections for V L: A l E [N + 1].
Wo(k) (V =
L = k3) N W
1=1
From the data generation process, it is obvious get
1
1-a
E E(g.))
1fy 11y= kb
N+1 a.
1fk
N +1)
1 N).
(9)
N+1
N
26

---

Published as a conference paper at ICLR 2025
Pre-trained Pythia-IB on Factual Recall
0.8
probahility
0.2
0.0
Correct: full
Correct LASER
'the': full
'the': LASER
Pre-trained Pythia-I on Factual Recall
1.0
probuhility
0.4
0.2
0.0
Correct: Correct:
'the': full
'the': LASER
Pre-trained Pythia-2,88 on Factual Recall
probability
0.4
0.2
0.0
Correct:
Correct:
'the':
'the': LASER
Figure 16: The prediction distributions of Pythia- IB, 1.4B 2.8B on more examples of factual recall. Compared with the setting in Figure 5 (right), here we use 20 examples in Table 5 LASER turns out to significantly decrease probability "the" against the correct tokens.
Since O(1) much larger than N when N > due law large numbers, we have the population gradient Vw, satisfying
Wo + +1) L)We(() = O(1),
::NN:
(()) < 0, with absolute value in O(1/n).
27

---

Published as a conference paper at ICLR 2025
Layer
Layer
Layer
Healeod gobabiity
10
04
02
04
8.5
P
Figure 17: Predicted probability for c E ("Mary", "them", "the", "John"). LASER is conducted on input matrices of MLP layers on the layer = 9, 10. 12 of GPT-2 Small. The input is "When Mary and John went a store, John gave a drink to". The horizontal is the fraction perserved rank. o E [0, 1), where 1 stands for the full model. It turns out LASER clearly decreases probability of "the" and "them" when p [0.1, 0.8) for layer / 9, 10, 11, compared with the full model.
The variance of the gradient projection onto W + 1)W (q) of single data point follows that of Bernoulli distribution with parameter a, which means
1 + 1y==
Var
a(1 a.)
(10)
Similarly, for any h N, the variance of the gradient projection onto Wo(N ( of single data point follows that of Bernoulli distribution with parameter 1-a which means
N
1 13} 1-a (1 (1
a = 0(/N).
Var
(11)
N +1
The ranges of the gradient projections' deviation from the expectation are
N + = N +
N +1 a)
S &maxfa -
(12)
1 - -
1
V
k] N+1
$
N +11
N
For each choice of & E [N + 1] individually, after having the expectation (a(e), variance o²(k) and range R(k), by applying Bernstein's inequality, then: for each [N + 1], with probability 1 5, holds
40 (k) In(2)
4R(A) In(
Wo(k) (Vw,))((W - m(k) S
m
77
Then by the union bound in probability, we need (N + 1) events above to hold at the same time, so we can substitute 5 with 5
N+1 to have: with probability 1 5, for any E [N + 1). it holds
40²(k) (In(N 1) + (n)
4R(k) (In(N + 1) + in()))
Wo(k) (V.. u(K)| sV
(13)
D.2 GRADIENT FOR THE VALUE MATRIX Wv
Lemma D.2. Consider zero initialization, Wv WF= WoK = 0. Then with probability for any j, E [N 1), it holds
Wo()) (Vwv W)) m()
(14)
402(j. (2 1) + ln)
4R(jk) (2In(N - + In(3))
28

---

Published as a conference paper at ICLR 2025
Table 3: ju[j,) (j. R(j,k) for different choices of (j,k) in Lemma D.2
j
14
N +1
N +1
2
2
N
TN
N
N +1
7N
No
N 11
[M] (2)
TN
N
4
N+1
20-1
N
TN
No
20-1
d'-a" -
+ a4-
9
o/2
22TN2
[M] (1)
(2- 0.
N
TN N
[N] \ {o}
N+1
(2 a)
TN + N3
[N] {o}
(2 a)
TN + N3
N {q}
30-1
(2(1-a)2 a)
[+(1-a]2
N
TN2
N {g}
[N]] \ {q.5)
(2 a)
(TN2 + N3
1
N2
where k(jkk, 2 (), R(j, k) expectation, variance and range for different choices of (j, k) at listed Table 3
Proof. Due to zero initialization, ie., Wv WF = 0, the current predicted probability Pw (k[z,) N+I for all [m]] and [N + 1]. Mcanwhile, the attention score is uniform as + for context positions due to WK K = O. Therefore, from Lemma
we have
N+1 Vw = Σ
)
Liv = k} (w((k)
T Tz.t
N+1
k=
where Zi E Rd W ( Pe is the input embedding with input token 21, at position t in sequence i, together with positional encoding pt for position i. With the assumption of orthonormality in Assumption D.1 we have the projection of V Vw L: j, E [N + 1].
12 k3] )
(V.
1
(
16: =s)
Σ
N+1
T
1=1
Since each sample is drawn ii.d., it suffices to discuss the expectation and variance
A.(,A)= A
N+1 14,T41 = =st)
1
(+ Σ 1 [z,1 k))
T.G.A).
where we use the fact y: = 21,T+1
Recall that, for each sample in the data generation process, the trigger q is fixed while the correct next token ÿ ~ Uniform ([V]. Hence, conditioning on 2;,7 q. it has probability a for 2,,41 = N + and probability
a for 2, = y. This leads the necessity of discussing whether not ÿ = Meanwhile, a corner case of
is also necessary to consider, as this implies an event that increases
the counting ΣT. 1( o} than the case of q.
Therefore, generally there are 10 cases due to different choices (j, k) as follows:
1. J. = =
2. = -
3. == N - 1, (NT (g).
29

---

Published as a conference paper at ICLR 2025
4. 9. N 1,
5. j -jg..k- 9,
6. j = g, [N] \ {g).
7. JE (M (s), k = N +1
8. jE [M] (g) 4
9. e (M] (g),k= j,
10. j E [M] (0), A E [N] (g.5)-
For each r;(j,) k) individually, if we have its expectation μ(j,, k), variance o² k:) and range R(j, k), by applying Bernstein's 's inequality, then: for each j, [N + 1}, with probability 6, holds
(02 In (3
k)(j,))ln(3)
(G.A) - m(..R) S
Then by the union bound in probability, we need (N + 1)² events above hold the same time, so we can substitute 8 with (N+1)= to have: with probability 1 8, for any j. [N + 1], it holds
(G.A) - 402(1,A) (2In(N + ln())) 4R(1,)) (2
+ +1 In(3))
m
m
(15)
As a final step of the proof, now we elaborate the expectation, variance and range o for these 10 cases.
Case 1: j N 1, 1
There is probability I for ş q and probability N for ý # q. Hence, we have
1
B[T,(i k)] = E (TTT(C,A))) = a) + ETT, (i. k)ly # d)
N
N
1 -4 AT
N
E ET. N
E[T k) + gl-
N
From Lemma B.2 and the independence between 1 12,,T41 N + 1} and E1sT 1, kj. we have
1
ET. o
N²
1
1
TN
N2
where the second is from
-()))))
1
2
1) 2 (1 a
E
N +1
1 N
o)
N
W+1
+
N
NN11
Similarly, from Lemma E.5
we have
-a a
N
TN
N2
30

---

Published as a conference paper at ICLR 2025
Therefore, it holds
a
N -
a²
+
N
N
N
N
N
k(r,(i,)) B[T (J k) = 9) +
BIr,
1 +
N
N
1
N
-
a²
a²
a
N a
TN
N2
N
TN
N²
TN
N
a²
a
a
Var[T.(i - ([T B[T,(i,)]
TN
N2
The range of TT(j, k) is
[ kiT]( 1
2
and the extreme case is when half of the sequence is N + with the rest all being q.
Case 2: N q.
Similar to Case 1, we have 1 (2,,T+1 N 1} is independent of EiST 1 k).
From Lemma E.I we have
E|T, k)ly 9) a.
2
ET. o) aTN
a.
-1
a N2
From Lemma E.4 have
+
N'
ET
TN
NZ
Therefore, we have
1 BT](j,A]] 7 BI
N 1
B[T (j. k)
N
N
N
(])(( = =
1
N -
ETT.
a
+
N
N
TN
N2
a-a²
[ar[T.(j.))] -
TN
N2
The range of TT(J, k)) is
T.(i,k) k) F[T.(ii (i, 1,
and the extreme case is when ÿ q and the sequence is all q's.
Case 3:
[N] {g}.
Similar to Case 1, we have 1(2TT41- = N + 1} is independent of EnST 112,, = k}.
From Lemma 33 we have
1
BT, = d N'
1
1
E
TN
N2
31

---

Published as a conference paper at ICLR 2025
From Lemma E.7 have
1
N'
+
1
TN
N2
Therefore, we have
1
B[T((iAA]] R -a N'
1
B[T,(,,)))) a
+
TN
N2
a²
=
TN
N
The range of k)(j,k)in is
- -
and the extreme case is when all of the sequence except the last one is k.
Case 4: j= 9. = N +1
If ø & q. we always have 21,T+1 # q because 2,T+1 E {g. + 13. If conditioning ÿ= q. it has probability a for 2,,4+1 = q, independent of E1sT 1( = N + 1}.
From Lemma
E.5
we have
N + 1
N
a²
ET
E
++1
TN
N2
From Lemma E.2
we have
1
- a)
N
E - a))
1
1
TN
N2
Therefore, we have
1 [[],(;,)]] NV FTT =
N
2a-1 1
(J k)ly #
N
N2
E[T((,,)) N (T, +
N-1
E[T kj,
a²
a
+
1
N
TN2
N3
a² a
Var[T.(j,k)] ET. ET,
TN2
A3
The range of k)) is
1
- -T. &
2'
and the extreme case is when ÿ = q and half of the sequence is + with the rest all being
q.
Case 5: j=g.k= q.
Similar to Case 4, if ÿ A q, we always have 2,,441 # 9. If conditioning on ÿ = q. it has probability 1 for 21,T+1 q, independent of Etsr 1 (/ = 4).
32

---

Published as a conference paper at ICLR 2025
From Lemma E.4 have
1
N+1 N
1
B Bur
N+1
TN
N2
From Lemma E.
we have
ETT, (j. k)ly = g) (1-a) a.'
2
E[T =) a)
(-1
aTN
2²
Therefore, we have
N
20
= E /1,
N
N
aN2
N
B[T
BIT k,, - -
N
N
a
a +
a² a +1
TT2
2²N3
a
0.2
2
a + 1
Var[T,(j, = E|T ()) - ETT
TT
N²N3
The range of [i k) is
T.G k) ETT.GG,)))] - 1
and the extreme case is when ÿ q and all of the sequence are q.
Case 6: j= k E [N] \ {g)
Similar to Case 4, if ÿ q, we always have 2,, # q. If conditioning on ÿ = q, it has probability a for 2, q, independent of E1ST {=z,e = k).
Moreover, we need to consider whether ÿ = or not.
From Lemma E.6 have
- a
N+1 1 N
(2
(2 )
[[ (j, k) t,
N+1
TN
N2
From Lemma
E.7 7
we have
ET. (TT.(.))) a (a.
N +1
N'
1
1
1
T (1. k) # q N
(g.00)
+
N+1
TN
N2
From Lemma
E3
have
1
BTT a)
N'
1
1
E = a)
TN
N2
33

---

Published as a conference paper at ICLR 2025
Therefore, we have
B[]:(i.)]] = ( = x) (
N
+ N - 2
= [M] \ {0.0H1
N
23 a
N2
B[r,(i,A)] = [[ (j. - + 1 N
N
+ N 2
BIT, k) [N] \ {1.0))
N
(TN2 + 1 N
2 (2 a)
NN3
Var[T.(IAA]] = EIT - (CG6A))-- = (2 - a)
TN² +
N3
The range of k)) is
- - 2 1,
and the extreme case is when all of the sequence except the last one are k.
Case 7: j E [N]
If y # j. we always have 2, # j because 2,7+1 E {g. N + +1). If conditioning on ÿ = J. it has probability - for 21TH+ = j. independent of EIST 1 {2,, N 1).
Moreover, in the case of ÿ * j. we need to discuss whether or not q-
From Lemma
E.2 we have
1
N+1 1 N' -T. N+1
(N +
N2
From Lemma E.5
we have
1
a
~
N ++1 N
B 1)
1
α
a²
N+1
TN
N22
From Lemma E.5
we have
|) a) a
N'
a²
ED TN
+
N²
Therefore, we have
EJ], 1 ET.0,))0=
=l 1 -TT = 5 >
2
ET
N
N
a²
N2
= +1 N 2
B[T((;,)))] N
BT.(C,k)) 1- = 2)
N
N
TN2 + a²)
2 (2 a)
N3
a²
Var(T.(G,k)] ETT k)a
BIT k)1² (2
(2-
T
N3
34

---

Published as a conference paper at ICLR 2025
The range of (T(j, k) is
[T k) - ETT 1 3
and the extreme case is when y - j and one-third of the sequence are k, where the sequence has repeated pattern like [g.7, + 1, + I......
Case 8: j E [N] (a); k = 9-
Similar to Case 7, if y # j, we always have 2, # j. If conditioning on V = J, it has probability 1 for 21 j. independent of Ligr 1f2,, N + +}.
Moreover, in the case of y # j. we need to discuss whether not y q.
From Lemma
E.I we have
1
1
E[T
N
1
aN'
1
T
2 a
T
-1+
+
N+1
a.N
22
From Lemma E.4
we have
N W+1 N
1
1
N ++1
TN
N2
From Lemma E.4
we have
1
=
N'
1
1
=
TN
N2
Therefore, we have
E(T,(,,)) BTGG -
=]+ D a
N B
N
N -
ET,(J,A) ET. ET. = 9)
N k) = +
ET. サ
N
2 (2 - a)
TN2 + N3
= (2 - aa) (1 + 1
(j, =
ET; =[
TN2
N3
The range of (jj,k) is
[. - ET. 1
2
and the extreme case is when ÿ j j and half of the sequence are g-
Case 9: je [N] \ (g), l j.
Similar to Case 7, if g # j. we always have 21,T+ # j. If conditioning on ÿ = j. it has probability 1 for 2, j. independent of Eisr 1 N + 1).
Moreover, in the case of y * j. we need to discuss whether not ; q-
35

---

Published as a conference paper at ICLR 2025
From Lemma E33
have
1
EJT
N+1 N
N+1
TN
N
From Lemma E.7
we have
1 1
N+1 1
N
1
N+1
TN
N22
From Lemma E.6 we have
2-a
N
(2 2-6
(2 - a2)
" (1 a)
+
TNN
N2
Therefore, we have
[[]+(AA)]
ET ETT 7
1 -= =
N - (T((
N
N
a² + 3a-
-1
N2
E[T(Ú,,)
ET (j, k) 1=
- - E[T.(i.)) * $9000
1
1
N
N
1+ (1 (2
1+ a) az)²
TN2
N3
1+ 1+(1 a2 a) 1 (1 a) ax)²
Var[T. = ([TT(i,A))] BTT.,)))
TN2
N3
The range of (J) is
T. E[T (J)) 1 1,
and the extreme case is when ÿ j j and all of the sequence are k.
Case 10: j [N] (Ø), k [N] (q.5).
Similar to Case 7, if 5 # j. we always have 21,T41 # j. If conditioning on ÿ = j. it has probability 1 for 2 j. independent of Eisr 1 (2,,t N + 1).
Moreover, in the case of ğ # j. we need to discuss whether not q-
From Lemma
E33 we have
N
1
N
1
1
(T(
N+1
TN
N²
From Lemma E.7
we have
1
-T a)
N'
1
1
E
a)
TN
N2
36

---

Published as a conference paper at ICLR 2025
From Lemma E.6 have
2 2-a -
ET. k)ly
k] R
N ++1
N
2 a
a)²
E|T =
TN
N
From Lemma
E.7
we have
BT (U
N+1
B N
TN
N
Therefore, we have
[[](
LE - -TT -11.GA))
1
N
N 3
BTT
N
N2
([T,(,A)9]]
N E (j, - + + -T.
+ B kk]
N 3
+
BT, ()
N
~ (2 a)
+
TN
N2
1
Var[T.(;,)]] EIT.
EIT E[T (2 -
TNE + N3
The range of [j k) is
T. - E(T 1
and the extreme case is when 7 j and all of the sequence except the last are k.
D.3 COMPLETING THE PROOF OF THEOREM
Theorem 4 (Restatement of Theorem D. Assume N,T > 1,a=0(1 Consider a one gradient step update from zero-iiitialization on n i...d. samples of 21. with separate learning rates qi for W and no for Wv (note that the gradient on Wok is zero). For test sequence zuT, the resulting logits for the forward and attention blocks satisfy, with probability 1-
In 2(++1)
(AEG((IIT) - n15 -0
m
TN + No Im 2(++1)
In 2(++1)
A A(Sama(I1:T)) (a2g a(1 4))l STD Te N
m
m
where A(E) EN41 maxJe[N] &j is the margin of predicting the noise token and q
+ EisT 1(2, = N 11.
37

---

Published as a conference paper at ICLR 2025
Proof. For W p, since the input is always 2r 9. the logits will be Eala (k(k) We We k E [N + 1.. As WF is initialized from 0 and updated by GD with learning rate f/y, after one-step update, have
(o(k) ( V W
( (
E1r
W
By Lemma D.1 with probability
36, we have
2(N+1)
[knlN+1 - ---aal 50
0
Sn 0 Ss N, Kala Salk-nf- nf (1~0-- N N+1
x+11
In 2(N++) In 2(N+1)
Nm 771
and then triangle inequality finishes the proof for En-
For Wv, since the gradient WQk at initialization is zero, Wok being zero after the first step induces uniform attention the input sequence. Consider the input sequence {2i then the logits will be (Kamly (i) Tv) + Ers WE(a), 4 e [N 1].
Then considering the concentration bound of Wv after one-step update in Lemma D.2 denoting k(j = We(j) TW W (k),, we have
1 ErG [5amn] T IST
Σ k),
T
ASN+1
with concentration bound for each T(:') in Lemma D.2 From Table note that for all j = + 1, S N, the expectation and variances are the same, while = N + has slightly different expectation and variance (but still in the same order of the others). Hence, denoting
i = > 1( = N + 1] dependent of the test sample Z1.T, we have
(VLr (TN
2(N+1)
In 2(N-1)
((in())))))++1
Th N (a² a(1 (1)11-11)) 1)) S No.00 O
N2
m
m
Meanwhile, the terms in Table 3 for j + N always have much smaller mean and variance by factor 1/N, using the Bernstein's inequalites for these terms in Lemma D.2 finishes the proof for Wv.
E PROOF FOR FIRST AND SECOND MOMENTS IN LEMMA D.2
In this section, we will show the proof of the first and second moments of Σiser 1( k}}l for all cases. Note that we do not consider 27 = q, but including will not change the results, as T > and 2y is explicitly fixed g during data generation in Section Generally, there are three factors to classify the cases as follows:
1. The i.i.d. uniformly sampled correct token ÿ E [N]:
(a) y= = q.
(b) y 4-
2. The target token & [N + 1]:
(a) A
(b) k N + 1.
(c) k < N, g, k T.
(d) ( N
=
38

---

Published as a conference paper at ICLR 2025
3. A condition about the token 2u before the sequence { (212
(a) 20 9,
(b) 20 E I + 1 (4).
Note that when 20 will be implicitly or explicitly considered. When there no condition on the first token, which means ~ Uniform ([N), this belongs to Case 3bb, i.e., E [N + 1] (2). following the data generation process.
Table 4 summarizes all lemmas about the seven cases classified by the first two factors. The third factor about 2g explicitly presented in the proof of each corresponding lemma.
Table 4: All lemmas about the seven cases classified by ÿ and k.
2a
2b
(cc
2d
E.1
E.2
E.3
N/A
E,4
E.5
E.6
E.1
WHEN y q
Lemma E.1 (ğ = q, k Q) Following the data generation process, assuming N,T > and
= (1), if 9 and q. it holds
T
E
Σ
1(2 -
aN'
1S
(16)
E
Σ
1( =ki]
T
-1
2
T
aN
a
22N²²
1TT
Proof. For simplicity, we omit the condition of y = q, k = q this proof. Denote
Y(T)- & E
Σ 1f 20 =
1ST
((T) _ E
Σ 1( - 1)
E [N + 11, $
KT
Then the data generation process implies, VT 1,
Y(T) p(21 q q) (1 ((T - + - Y(T) p(21 = =/=0 * () (1+ L(T 1)) +() E (N] (0)l= X L(T-
1),
The iteration becomes
Y(T) = (1 - a) Y(T aa +
\(T) Y(T N N
-(T-1)++ 1 N
N
This gives
Y(T) -(T) (1 (1-cl a- -1 - N -) N'
1
N'
1
39

---

Published as a conference paper at ICLR 2025
Consider the initialization Y(0) (0) 0. This implies
Y(T) ((T)- = 1 #V
(1 1-(1----) )
a + #
N V(T)+ | T.
N
Then we obtain
1
N22 T aN + aN
Y(T) ~
(T aN)
+
aN -1
N 2
oN +1
oN +1
T
aN
-
-1
a²
Y(T)
ON
(aN
+ 132
aN +1
T
aN
Since the data generation process implicitly assumes + q, we have the desired expectation
] = -(TT)
Σ
- q, 1 = k3
T oN
EST
To obtain the expectation of the quadratic term, we similarly denote the following terms with different
Z0:
Σ 12 = ki) 1ST
Z(T) & E
20 =9
((T) & E
Σ 1 (2x = =k)) 1ST
50 [N 11, #9
Then the data generation process implies, V 2 1,
Z(T) p(=1 =0 = 9) ( (1 -)) + p
-
N 1/20 =Q Z(T 1),
((T) = p(21 = =0 + q)
+ # g/20 + 4) 2(TT -1),
where 2Y(T- 1) due to E1 + E2s1er 21 | + E Σ 12s5sT .
Then the iteration becomes
ZTT) = (1 a) (1 a -1)
=1 (1)a)
(1 +T - 11),
2(T) = N (1 + (T -1) (1) ZT N 2T 11)
N
N
Z(T - 1) N
2(T-1) 1))
(1 (T - 11). N
= N
This gives
1
Z(T) ((T)) a
N (2(T 1) - 2(T 1)) + (1 aa (1 (T -1)),
1
1
Z(T) + +z( = N Z 1) + - 1)) -
(1 + 2Y
N
N
40

---

Published as a conference paper at ICLR 2025
Considering the initialization Z(0) 2(0) 0, have
((T)-2(1)) = Σ (1 (1-a
1
(1+
N
KT 11
(1
a
N
1
1
T-
1+
aN
2t
2
a²
2
(ST
1-a a
2(1 -
a)
-1+
2²
a²
N
T
Z(T) + (T)
Σ
(())
N
N
1S1ST-1
T
2
t
11
N
N
aN
a²
1S1ST-1
T
2
T2
-1
+
N
aN2
Then we obtain
T
3
2
2
2 1). Q2
T
1-a a
Z(T)
+
+
a
a
a
a²
T
2
T
2(T)
-
+
aN
d²
Q2N?
Since the data generation process implicitly assumes 20 F q, have the desired expectation
1 (12==)]] k2) g 0= g.
)
T
2
T
E
Σ KT
((T)) ~ aN az
-1+
Q²N2
Lemma E.2 (ū = q, k- = N + 1). Following the data generation process, assuming N, T >> I and ((1), if q and N it holds
T
Σ 12 = g=
N
(17)
Σ 1iz k3) - N +1
)
T
T N2
E
+
N
KT
Proof. For simplicity, we omit the condition of ; = = k- = N + 1 in this proof. Denote
((T) & E
Σ 1fa = k) 20
((T) A
Σ ATT
1/2 = 1) 20 E [N + 10, $
Then the data generation process implies, YT 2 1,
Y(T) = p(21 = g/ q) Y(T - 1) + p = N + =) (1+ 1)),
\(T) p(21 ()) Y(T - p(a E [|] \ (0)1=0 # .(T -1--1).
41

---

Published as a conference paper at ICLR 2025
The iteration becomes
Y(T) (1 - a)- Y(T 1) a
Y(T) 1 N Y(T 1I+
N
Y(T 1).
N
This gives
1
((T) (1 a N | 8(T + a.
1 N Y(T) + 1 -T +) - N
N
Consider the initialization Y(0) Y() 0. This implies
1- (1 - a - N )
Y(T) - ((T) = a+ I
a
((T)+ N T.
Then we obtain
T Y(T) ~
+,
N
T
Y(T) ~ N
Since the data generation process implicitly assumes zo * q. we have the desired expectation as
E
Σ
1z== = o k N = (TT) N
FST
To obtain the expectation of the quadratic term, we similarly denote the following terms with different
=g:
2 Σ 1 1/ = k1)
Z(T)= _ E
20 - 9
1ST
2(T)) & IE
Σ 12 = ki)
20 [ +11, #
Then the data generation process implies, VT 2 1,
((T) p(=1 =0 = - - 1) p(21 = + 1/20 (
1)),
((T) = p(21 = 20 + g) Z(T - 1) p(21 g + q) 2(T - 1),
where 2 (T - 1) is due to IB[(1 + Σ2SisT ) 211+221 2E
2 -1 + EE ( (E2sesT
Then the iteration becomes
2(7)- = (1 - a) 2(T - 1) + (1+2 -1)
=(1 -a) 1) al
2(T) N (1) N -
2(T-1)
N
This gives
-(T) - (1 (Z 1) 2(T - 1)) + aa1 + 2 1)),
N
N ((T)+a2(1) N Z(T 1) aē(T - 1)
1
1
(1 2Ý (122((T-1)) 1)).
v
42

---

Published as a conference paper at ICLR 2025
Considering the initialization Z(0) 2(0) 0, we have
Z(T) - 2(T)= = Σ a(1 - a 1
) T- -- (1+22 (t)
N
KST
Σ
a(1 a N T-
1+
2t
2T
+1,
N
N Z(T) + aZ(
aT'
2a
Σ
4(1)
N
N
1S1ST-1
aT
20
2
N
N
N
151ST-1
aT
aT
N
N2'
Then we obtain
T
T
Z(T) ~ 30
+ 0
+
N
N
T
T
2(T)
N
N
Since the data generation process implicitly assumes zg * q, have the desired expectation
E
Σ KT
1(x 1()==)))
- N -
((T) ~
T
T
N
N2
Lemma E.3 (g- q. N, & q). Following the data generation process, assuming N, > 1 and a = 0(1) if = q and l E [M] {g). it holds
T
Σ 1T
1 1(2 =
{q}
N
(18)
E
Σ 12 = 1t)
0 9. 11)
\
2
T
T
N
N2
ET
Proof. For simplicity, omit the condition of ÿ = q, k [N] \ {q)
in this proof. Denote
Y(T) E
Σ
1(4 -k)
20
1ST
(T)) & E
1/ k} 20 E IN + 11, #9 q
IT
Then the data generation process implies, VT 1
Y(T) = p(21 = 2 q) Y(T 1)
1) p(2 = = 1/20 9) . -1),
8(7) = p(= = 2/20 # q) Y(T 1)
p(2 (N] (0)100 # 9) (p(2) k/2 Uniform ( ([N) \ (/)) Y(T 1)).
43

---

Published as a conference paper at ICLR 2025
The iteration becomes
Y(T) = (1 Y(T 1) -a ((T-11
- 1) N-1 1 N -1)-1))+ - (1
P(T) =
N
This gives
1)) N a N x(T)+ NY YT 1) + aT NN
Y(T) - Y(T) = (1 - a 1
N (X ( (
Consider the initialization Y(0) = Y()) = 0. This implies
- N
- (1 - a NT) )
- -(T)-
a+
N ((T)+ T.
Then we obtain
T Y(T) N
T
((7)
N
Since the data generation process implicitly assumes 2g * q. we have the desired expectation as
E
Σ
1
u- N = Y N
IT
To obtain the expectation of the quadratic term, we similarly denote the following terms with different
=o:
2
Z(T)= _ E
20 - 9 Σ 12 = -1)
1ST
2(T) & E
Σ 1 (2 =b(= 1<T
20 (N 11, 20 #
Then the data generation process implies, VT 2 1,
Z(T) p(21 q/20 q) Z(T 1) p(2 = N 1| 11-0-///1/ q) 2(T 1),
2(T) = -p(2i = -|=0 + q) Z(T - 1) + +p(a # g 2(T 1)
-p(21 k|z0 # q) (1 2Ý 1)),
where 2T 1) is due to E[(( + Σ E = 2E > E2SIST | + +E([ 22 )
Then the iteration becomes
((T)= (1 -a) 7(T 1) + a. ²i(T 1),
²(T) N 7(T - 1)
1
N 1
2(T 11) 1)+ + N
1
(1 + 2Ū(
N
This gives
1
1
Z(T) 7(T) = (1 a (Z(T 1) - 2(T 11)
1 (1 + 2 1));
N
1
N Z(T) +a2(T) =
N 2( - 1) + aZ( - 1)+ - (1+ (1 (T
44

---

N2 Published as a conference paper at ICLR 2025
Considering the initialization Z(0) Z(0) 0, we have
Z(T) - ((T)
Σ a-
1 (1+28(1)
N
SST
2
(1
11 a
1
T-1-t
1
2
N
N
N
1ST
2T
+11,
aN
N
N (
20
Σ
4(()
N
N
1SKTT-1
xT
20
Σ
N
151ST-1
AT TT
N N
Then we obtain
T T
Z(T) R
N
T
T
2(T)
N
N2
Since the data generation process implicitly assumes + we have the desired expectation
Σ M2 = (==))/p-000 (NI
T
T
I
((T)
N
N
E.2
WHEN ÿ =4
Lemma E.4 (g + 9. k = q) Following the data generation process, assuming N,T > and
((1), if and q. it holds
T
E
1 1f2 k k
~
N
(19)
E
Σ
1t = -k)
1 + g =
Rd
T
72
N
N2
IST
Proof. For simplicity, we omit the condition of z * q, k = in this proof. Denote
(()) & E
Σ 1(2 ((===k) k} 20 =
AST
(T)) & E
Σ 1( 12=
20 [N + 1), $
Then the data generation process implies, VT 2 1,
Y(T) = Y(T - 1); Ý(T) p(21 =/ 9) (1 Y(T 1)) p(a (N] (0)1 # 9) ( -1) 1).
45

---

T N Published as a conference paper at ICLR 2025
The iteration becomes
Y(T) = ((T-1 - 1),
1 Y(T)= N Y[T 1)
N-1
Y(T 1)+ 1
N
This gives
((T) - (X -)- N
N (
Y(T - -
1
N
Consider the initialization Y(0) - Ý()) 0. This implies
( 1- N )²)
Y(T) - ((T) 1+ I
N
N +(T)+((T) T.
Then we obtain
T Y(T) N'
T
Y(T) N
Since the data generation process implicitly assumes F q, have the desired expectation
kb o 1 = 8
E
Σ FST
12
To obtain the expectation of the quadratic term, we similarly denote the following terms with different
20:
1 (z 11) 20 -9
Z(T) A E
2
2(T) & E
Σ TT
1 1 k1)
20 (N 11, #
Then the data generation process implies, VT 2 1,
Z(T) = 2(T - 1),
2(T) p(= 20 -)+ ()) p(2 E [M] {g)1 + q) 2(T 1),
=) where 2Y (T 1) is due to E[(1 + E 1-1+2E 21 2 | EI )
Then the iteration becomes
Z(T) = 2( - 1), 1
N -- 2(T - 1 (1 2Y
2(T) 1 Z(T 11
This gives
1
1
Z(T) 2(T) N (7(T 1) 2(T - 1)) (1 + +T 1)),
1
2(T) + 2(T) N
1
Z(T - 1) 2(T
1)
1
1
1
(T
N
N
46

---

N2 Published as a conference paper at ICLR 2025
Considering the initialization Z(0) Z() 0, we have
Z(T)
1
Σ
T-1-t
(2)
N
N
(ST
1
1
T-1-t
1+
2t N
N
N
1
2T
N
N
( 2(T)
T
2
Σ
Y(t)
N
N
1-
2
+ N
N
N
15KTT-1
T T
N2
Then we obtain
T
T2
Z(T) ~
N
T
T
Ž(T) ~
N
N2
Since the data generation process implicitly assumes 20 F q, have the desired expectation
E
Σ 1
12 = k1)
2 g - go q. [N] (0) 2(T)
T
T
N
N2
Lemma E.5 (g & q, k = N + 1). Following the data generation process, assuming N, N.T. >> I and (1), i 7 and k N 1, it holds
aT
E
1[a k N+1
2
N
(20)
E
1 = k3
kg N+1
2
aT
2²T2
N
N2
KT
Proof. For simplicity, omit the condition of ÿ q, N + in this proof. Denote
Y(T) & E
Σ
1(a k) 20 = 9
1ST
(TT) - E
Σ 1(2 = k) 20 IN + 1), & q
1T
Then the data generation process implies, VT 2 1,
1/ = q),
(TT) = p(2 = 2|20 * q) Y - 1) + p(21 E [N] {9H120 # 9) Y(T-) - 1).
The iteration becomes
Y(T)= = (T 11)
1) + a
1 / Y(T 1) + Y(T)=
N 1
L(T-1)
N
47

---

Published as a conference paper at ICLR 2025
This gives
-(T)-((T) - = N - 1) - + a,
1
Y(T) Y(T) = N Y(T - 1) N
Consider the initialization Y(0) (()) 0. This implies
(1- (1- ---
Y(T) -0)--((7 Y(T) I+ N
N + = a T.
Then we obtain
Y(T) ~ aT
+
N
aT Y(T) ~
N
Since the data generation process implicitly assumes 20 o q. we have the desired expectation as
aT N
Σ KT
k===
=(
To obtain the expectation of the quadratic term, similarly denote the following terms with different
z0:
Σ 1 12 = k1)²
Z(T) & E
20 =9
15T
2(T) - E
Σ KT
1 -1)=)1
on [N 11,
Then the data generation process implies, VT > 1,
p (TT) - p(1 N 1/20 = q) (1+ + (T 1)).
(TT) = p(21 = g/ * q) 7T - 1) + p(z E [N] \ {g}120 + 2) 2(T - 1),
where (T 1) due to E[(( + E E2SI5T = 1+ 2E Dectr 1 (E )
Then the iteration becomes
Z(T) - a -
2(T) - Z(- - N- -
2(T 1). N
This gives
Z(T) ((T) 1
N (Z(T -(
N Z(T) + 2(T) N Z(T -1) -
1
(1+2((
48

---

Published as a conference paper at ICLR 2025
Considering the initialization Z(0) Z() 0, we have
1
- 2(T)-2(T)
Σ
---1-111
(1+
(1+
N
157-1
Σ
1
--
1
2at
N
N
1-T-1
2T
+a,
N
N
1
Z(T) + Z(T) = aT
N
+ N
2a
Σ
(()
1SIST-1
aT
2a
at
N
N
N
aT T2T
2
N N2
Then we obtain
Z(T) ~ T (2a² + a) +
T²
+
a.
N
N2
aT a²722
(T)) N N2
Since the data generation process implicitly assumes 20 + q, we have the desired expectation as
12 :aa-k))) 1(===)))
oT a²T²
E
Σ 1T
(M {p} 2(T)
N2
Lemma E.6 (ī + q, k D. Following the data generation process, assuming N,T > 1 and ((1) if q and y. it holds
E
1 = 8 q, g Σ 1ST
a)
T
N
(21)
2 = -1(==1))
, 1
(2 a)T
(2 - a) 27
E
Σ KT
1(
N
N2
Proof. For simplicity, we omit the condition of 5 + q, ÿ in this proof. Denote
Y(T) - E
Σ 1( = k) 20 =
(T) & E
Σ 1 k)
[N + 1], 20 # q
t
Then the data generation process implies, VT > 1,
Y(T) = YT 1) p(z) = م| - 4),
Y(T) = p(21 = 9 -/ + q) (T - 1) + p( E [N] fgHlen # +9) - 1) + p = 5 + g).
The iteration becomes
Y(T) ý(T 1) + (1-a),
1 N - Y(T)
1
Y(T - 1)+ N. N
49

---

Published as a conference paper at ICLR 2025
This gives
Y(T) N -VT--)-- a N 1 N +(T) = N Y(T - -))
1
1
2 2
N
Consider the initialization Y(0) (())= 0. This implies
1 a N
Y(T) (YT)
1-
1 N
N ((T)+((T)) = N
a T.
Then we obtain
T
Y(T) ~ (1 - a) a) (2 a)
N Y(T) ~ (2- a) N
T
Since the data generation process implicitly assumes 20 + q. we have the desired expectation as
T -(T) a) N
E
Σ 11 4
To obtain the expectation of the quadratic term, we similarly denote the following terms with different
20:
Σ 1 1(2 = k1) =1))
Z(T) & E
20 =9
T<T
2 1 1 k])
((T) & IE
Σ 1ST
in E [N + 1], #
Then the data generation process implies, YT 2 1,
Z(T)= (T -) 21 (1)-
2(T) = p(=1 = g/20 + 4) -) E [M] (0)2o + 4) 2(T -
+ p(2) u # q) 21+ 1)),
where 2T - 1) due to [1 + E3SIST = 21 DESEST + E1(E 2E(ESS1TT
Then the iteration becomes
Z(T)- = 2(T - + +1 -1
2(T) = T Z(T - 1) N 2(T - 1)+ 1
1
21
This gives
1
1
Z(T) 2(T))
N (7(T 1) 2(T -1) + -a a
N )1 [1+ 2T
1
1
2 2
N Z(T) 2(T) N ZT - 1)
- -
N (1 2Ý
50

---

Published as a conference paper at ICLR 2025
Considering the initialization Z(0) 2(0) 0, we have
Z(T) - = (1- =11)
1
(1+
N
N
SST
2(2 att
(1
T-1-t
1
N
N
N
1ST
2(2 T)T
2 0)
1
N
1 (
(2 - T
2(2 - a)
Σ Y())
N
N
ISITT-
T
2 a)
22 a)
+
N
N
N
12 a)T
(2 - 27
N
N2
Then we obtain
T
(2-a) 2T
Z(T)= (2 (3) 2a)
(1-a),
N
(2 a T
(2
TT
2(T)
N
N
Since the data generation process implicitly assumes + we have the desired expectation
M [N - (0) 2(T) (2 a)T (2 212
k1 11)) - h.
E
Σ 1(2
N2
KT
Lemma E.7 ( 9. N, + q. + D. Following the data generation process, assuming N,T and a = ((1), if # 4 and k: E [M] \ (ū. q). it holds
T
E
Σ S
12 k} 1 + 9. [N] {ū.g}
N
(22)
5 g ) [
kt) 2
T
T
E
1
N
N2
ITT
Proof. For simplicity, we omit the condition of q + 4, k E [N] \ (ū.4) in this proof. Denote
Y(T) & E
Σ
1[ k} 20 =
1ST
Y(T) - E
Σ
1 = 1} 2 E IN + 1), +
KT
Then the data generation process implies, A > 1,
Y(T) = Y(T - 1), ((T) = p(=1 =/ # n) Y(T 1) + p(a [N] \ (g)1=0 #9) - 1) p(z1
k|=0 # 9).
51

---

T N Published as a conference paper at ICLR 2025
The iteration becomes
Y(T) = -(T (1-o)) 1 ((T) = N Y(T 11) N-1
Y(T 1) 1)+
Note that these two equations exactly the same those in Lemma E.4 with same initialization as Y(0) (())- 0. Therefore, we have
T
~ Y(7)=
N
T
(TT)
N
Since the data generation process implicitly assumes 20 * q, we have the desired expectation
] = ((T)
E
Σ
1== - 1
CST
To obtain the expectation of the quadratic term, we similarly denote the following terms with different
Z(T) & E
Σ 12 = k1) 20 =
12T
2
2(T) & IE
Σ 12 = k1)
20
[N 11, #
ST
Then the data generation process implies, VT 1
Z(T)- = (T - 1),
2(T) p(21 = 20 # q) Z(T + p(21 [M] (g]120 $ 2(T 1)
+ p = k|20 + q) (1+ (Ý - 1)),
where (T 1) due to EE((1 + EσS1er ) = 2F 2 22SITT | + [E (25EET) )
Then the iteration becomes
Z(T) = ²TT 1),
1 ²(T) N Z(T 11) -(T
N - 2(T 1)+ N (Ý -T-
Again note that, since Y(T) ~ (T). these two equations are the same as those in Lemma E.4 Therefore, have
T
T
Z(T) ~
N
N
T N
T
2(T)
N
Since the data generation process implicitly assumes o we have the desired expectation
b-a (
T
T2
E
Σ tTT
1(
2(T) =
N
N
52

---

Published as a conference paper at ICLR 2025
F PROOF OF THEOREM 2 TRAINING DYNAMICS OF THE ATTENTION LAYER
We consider the following simplified I-layer model for the noisy in-context recall task.
2t Wa(#)) + WE(#-1) Rd
d(xT, T1:7)
(य. W WoRZE))), Wya, Rd
(23)
(F1.T)
Woo(ZT,M1r) ERN-1
(n(TI:T)
WuF(xz) Wo e RN+1
With zero initialization of W QK, Wv, Wr F, analyze the training dynamics of these three matrices in three phases:
1. WF learns the noise association in O time,
2. Wy learns be identity for all tokens k [NV 1].
3. W QK attends any position such that ž q and ūn = y.
Assumption F.1. In this section, make the following assumptions
1. (orthonormal embedding) WE(1) WE( = We()) = 1i6 = j} and WE(1) WE(i) for any i,j [N 1
2. (Feed-forward learns noise association) After phase 1, the prediction for noise always satisfies P(N + 1/21T) learn the noise association in a more quick speed than the other weights, so that it fair to assume p = for computing gradients of these weights.
for any Z1.T E [N + +1]0r, If deviates from a, W will
3. (Infinite samples)
the training loss L is population loss.
4. a S 1.5 - v5/2 0.38. This is to ensure the sign any j = k < N in 25
> 0 foror
Phase In this phase, the impact of W (E(2T-1) on WF and Wv is negligible compared with that of W (2r) because - is close to uniform in [N + 1] while 2T q fixed.
Lemma D.1 gives
(0(1) if = N+ ( &
(-W
Lemma gives
Lww = ) i N (0 N
(24)
Note that the entries of the above projection have the following signs, with details -14 in Table 3
> O, = N + 1) or (j = k) (j= q, = + 1),
(25)
K, otherwise.
The arguments in Appendix B.3 show
+ if +1, 0 ( fjsN.
(-
(26)
Therefore, during this phase, Wpr learss the noise association with effective graident norm of ((1) as (o +) (-V,, L)W x(g) @(1). Meanwhile, Wa r moves in the other directions
53

---

Published as a conference paper at ICLR 2025
uniformly in ( Wo ( (- W,, (q) @ for any l S N, which in fact ensures ((k(211)) = 1_ for any k S N and 21:7 E [N + 1| N
After O)n steps in this phase, we have p(N + 2 = a and p(k) = 1-x for any & S N and 21:7-
N
Phase 2: Assume P(N + 11-) starting from the beginning of this phase as discussed above. Due to symmetry for the rest channels, we have (kkl-) 1-ax
N Note that the attention scores in ((:-))
are still close to uniform, i.e.. [o (17 W Wor whose impact on attention scores is also in O(N-2) through exp(x) RJ track the movement of Wv under these conditions.
77. since the update Wok is in O(N
1 + x for x R 0. Then we
Since m xc, taking - + Te pa - Eaa k| and jis & E
((kW) -Ply
[
since
p(k)z) al [k N + 13 + - Like N = p(ylk). Lemmal H.l gives
N+1
Lw Σ
ply k) (k) E[Ely k])
ply = k) (B) Bizly- -Baal A))
W ( -
AJ)
N
k=
W (EE(())) WE WE(K) WE)
N
where the second equality is due to E/a = N + 1] due to y = N + 1 is uniform for any correct token ū S N, and the last equality from
Blaly = k1 N
(WE(k) WE) WEE()) WE)
with W E = N WE(1), WE=N-IS i= expected number of the tuple (4, 9) in a context length T is O f by comparing Lemma E.6 and E.7
WE(1) because E[r] = E, E. T and the
Therefore, the gradient for Wy has the following structure
~ NZ -
Vi,k N,
N3
(27)
N2 [
N3
Denote steps of phase 1 and phase 2 as L and L2 Combined with the structure of W y in phase 1 as in (4 ignoring projections that are O(N-3) or negative, Wv has the following structure after phase 2
OnTtNN-1)
if N - 1, Vk. if k < N, if g. k = N
WoO) Wr Wy O(itiN N (0(TTNN-2)
(28)
We() W (k) O(122N-2) iff k
Phase 3: now assume Wy y has the structure in 28 The model still predicts Pw (k(2) = aifk=
N 1) 1 1(k N) because the above projections of Wy onto Wuli ; N) is o(
54

---

Published as a conference paper at ICLR 2025
Meanwhile, the attention scores uniform + Wak R 0. Therefore, the gradient of W QK is
N+1 ply - (W) Wvr) wr(a; - =]]
Vwax L - T
-E((W W wr( 3)ylyy =]]
1
(B((W(()) (W) ( wr(a - x]
TN
>=1STT
(((W(()) Wyz) oT( 3) ly =k]];
IST 2t and the last equality holds due to the condition of y = N + uniform for any correct token D S N. Then, considering the above structure of W v. we notice that (i( "W WVI x 3g f2 = = 32
where I =-T--E Eisr
j} with 31 = N + and 3g ntyN for any j. k S N. Here note that we ignore the projection of j = q, k N + 1 in 2q because ÿ = q is with probability 1/N = o(1 so that it will not influence much the following derivation.
Plug-in Wo()" Wy and we get
WE(g) VW - We(b)))
1-a TN
ΣΣ IST
B[A
K,b k] -
-[A (t)
k.b. Az.
(29)
where
(t)
b,bi.b2 =
[1 = bi3 Σxst 1/2s = b1}
+ 1(2-1 -b2 Enst 1(25-1 = =ba] b2)
Now we are to control AR,b1 EiST k.b, k] - kbb, bg for different choices of
b b2. Note that b and by co-exist by sum in b,b,,b,, so the additivity of expectation allows to discuss choices of b,, by separately and then combine the results. Denote
D (3111z k} + 1 == k}) Afa = b1) E.st 1f2s = b1} =bbb))
T [1 = -b2) EXSt M6s-11 b2)
(30)
C = 131 = k} + 1 =k} ki))
T
Controlling Eisr E k.b, yy = k] E
k,by
If bj = k, from Lemma E.6 and E.7
we have
T
E
3,1 f = k}1fzz = k}
y= = k
- E
Σ 1ST
=k}
= 3 a) a))
N'
T
Σ EaST 1 1(2s = k)
Ensr f/z = kk
E
Σ
31f2 = k}
Σ
B11 = kb
T
T
IST
=-E
3 Σ
1(2,
k1))
+E
3,T
1 1(2, = =b))²
SST
STT
T
T
aT
(2 (2-a) 2T
T
= -
B1
N
N2
N
N2
N
The terms involving 1 = k} are negligible as OTT/N2). Therefore, we have
T
Σ
B k.
k]
(1 a:) N -
E
(31)
55

---

Published as a conference paper at ICLR 2025
If b1 + k, all terms are O(T/N2) because
If bi I N. it holds p(z= - bi |1-11 = k) - 1/N with the expected number of & in context of length L being (TT/N) from lemmas in Appendix F
If b1 N 1, it holds p(=z .N + 1 1 k) O(1/M) 1( = q} and the
expected number of in context of length is (T/NN) from Lemma E. and E.4
EIE, 1 condition y = le or not.
k) #b///T
E[#k- #b1/T ( (T//N3 no matter it is with
Therefore, for any b1 F k, we have
Σ
E E R,b ly k] BE (T//N)
(32)
Controlling Eisr
C|C
E
k,bg
kj B[C (1)
k
If by q. Lemma E.4 gives
11 =k) (1(2-111)) =
#y
E
y k.
T
1ST
B1 B1112 = (1 = -
#9 T
E
SST
#g T
B1 T
- (1-p(0- x)) E
Σ
B11 = =k] 11 =
y k
+
N
EST
= 31 - a T
N
where the last equality is from p(=; kly = - n) = 11 All the other terms are negligible with the same reason above. Therefore, we have
Σ ET
= (1 N
(33)
If b2 k, similar to the above discussion about BR we have
CBC k] E|C(!1 B2 a) N
T
(34)
Note that the key difference is that here we use Be instead of 31. and B2 < 81. If b2 * g and b2 + k. similar to the discussion for Eq(32). we have
Σ
EICL2 by h]
h B
oT//N).
(35)
KT
Therefore, combining the above results in Eq and C' from Eq gives
34 35p. taking sums of the corresponding B
(1 a a)TN + +(1-0)NN-
if by k, b2 = g:
B1 - a)TN + +2(1-0)7N-1
if by = k, bz = k.
81 - -)
if b, = k, other b2.
B1 a1
if by * k, bz a:
B1 a-
if b, * k. b2 Ic, otherwise.
O(
To take the summation over all k < N in Eq 29 we discuss the following cases of by and by for
W We(g)T (- (-Va (](We(b)) + We(l))).
56

---

Published as a conference paper at ICLR 2025
bf S b1 # b2 qf
when & = b1, we take Ak. by under the condition of bj = k. b2 when k * b1, take Ab,b, dy under the condition of b1 # k, q. Note that there are (N - such k. Therefore, it holds
1-
We(o) () (W + WE(1)) TN 31 a11-a)7(+NN)) (36)
If b1 = b2 q:
-when k. b1. take bn, under the condition of b, k. b2 k to achieve lower bound of the gap later.
-when k + b1, we take Ak.b, by under the condition of b # k. b2 q. Note that there are (N (N-1) 1) such k. Therefore, it holds
1-a TN (8i(1- - a)) 3e a)TN-1).
(-Vwa
Wn(e)
(37)
bf = N 1,b2 q: for any k S N, it holds k + bi. we take A under the condition b1 * ki by = q. Therefore, holds
WE(g) T (-V un L) + +1) W We())) = 1
B1(1 a)T.
(38)
TN
If by *.. 4 by: To get an upper bound of the projection length, we take Ak,b, ,2 under the condition of b_k, b2 l b # k, b2 k. Therefore, it holds
(- (W (b,)) (WE(b)) + Wx(b2)) - TN (8: + 21 a)TN-1, (39)
Comparing the above four cases, for any 7 < N, the attention weight Wok to attend more to W (x(() + W (y) than to WF(N + 1) WF(g), with
(WE + Wx(g)) - Wz(g) (-V () + (- We(d))
(1 a)
B2.
N
Meanwhile, any other setting of b,, by has smaller projection in (-V W L).
In summary, W QK has the following patterns
1. it learns attend indices t such that 2,-1 is the trigger word,
2. when there are multiple 4 such that -mm it learns attend to those with ÿ more than 2, N +
G LINEAR ASSOCIATIVE MEMORY
G. EXPERIMENTS AND DISCUSSIONS
In Section 3
showed that fully truncating feed-forward layer be helpful for reasoning. We now present a setting where noisy associations are stored in rank-one subspace of a layer, so that intermediate levels of truncation more useful remove noise.
Model and data. We consider a simple associative memory setting where the goal is learn an fixed permutation from input tokens output tokens (w.l.o.g taken be the identity), with linear model similar to Cabannes et al. (2024). Consider a learnable weight matrix W E Rdxd Consider embeddings for input tokens as 1n Rd and embeddings for output tokens as
(w) Rdd In contrast to Cabannes et al. 2024 consider additional "common noise"
57

---

Published as a conference paper at ICLR 2025
output token + 1, which is chosen for any input with probability (0, 1). For any input x E [r], the target distribution (a (z) is defined by
Pa(y)z) = (1 - a) 1 - x + 1/ 1 =
(40)
In other words, the last channel (c) for output the common noise with probability a for any input. The training dataset D, consists of uniformly distributed inputs x E [m]. and outputs conditionally sampled as y/æ ~ Pa (-1z).
Given any pair of input and output tokens, the associative memory model takes the form f(i,;; W) (ug, We), Vi,j [n] x [c),
(41)
When l & d, we denote the rank-k approximation of as (() by replacing W with W()), where W(A) the rank-k approximation W.
Training. During training, the dataset Da is generated with non-zero noise probability > 0. At test time, the dataset Du is without noise as a = O, so the computed loss is called pure-label loss. The model is trained with Gradient Descent (GD) subjected to cross- cross-entropy loss.
Experiments with randomness. Assume both {e.) and {urii_ are i.i.d. uniformly drawn from sphere gd-1 Also assume the model initialized as W. ~ (O, 4). Due to randomness from embeddings and model initialization, let's first conduct 20 runs of experiments to obtain significant factors before moving the theoretical argument.
Note that only full models trained, and track loss for low-rank models by conducting SVD in each step without manipulating training. In Figure 5, we illustrate the pure-label loss v.s. training steps for models of different ranks, where 3, 0.03 and 8 or 12. It turns out, while the full model (rank2 3) has a constant pure-label loss (~ 0.03, dependent on a), the rank-2 model is very likely have significant loss than the full model. Meanwhile, the larger d has more stable results than small d.
Pure-label loss rank 1, 1,2,3,4 models: dim-12
Pure-label loss for rank-1,2,3, 1,2,3,4 models: dim-8
10
10"
10
10"
rank-1
Loss
Loss
rank-2
rank-3
10a
10:
rank 4
10"
10"
500
1000
1500
2000
900
1000
1500
2000
Training step
Training step
Figure 18: Pure-label loss for rank- 1,2.3,4 models with n 3, 0.03 and 12 (left) or (right). Only full models are trained, and we report low-rank results by conducting SVD in each step without manipulating the training. In both figures, the experiments are run for 20 times to examine the randomness. For each rank, plot curves of the median, 25% and 75% out of 20 runs. It turns out: i) rank-2 models are very likely to have significantly lower pure-label loss thant full models (rank> 3), and ii) the larger dimension d has more stable results.
Therefore, we can qualify the following important factors for this model:
i. d v.s. n. when d n. c, random drawn embeddings tend to be orthogonal each other, with inner product in O(1vz). If n. R(d), embeddings will be in strong correlations, making the problem extremely difficult to understand. Cabannes et (2024) also discussed about such particle interaction in associative memory.
ii Low-rank subspace storing the noise. In Figure 18 the rank-1 subspace between the full and rank-2 models is responsible store the noise, removing which will induce model ideally predicting the ground- -ouuh without noise. This is understandable if the embeddings are orthogonal, shown in Theorem 3
58

---

Published as a conference paper at ICLR 2025
iii. When is large, orthogonal embeddings still induces low-rank subspace storing the noise, but a decides whether the low-rank subspace corresponds to the smallest singular values of W. If not, it requires more careful manipulation of the spectrum instead of low-rank approximation of W.
G.2 PROOF OF THEOREM 3
Now we present a theoretical analysis of this problem with some assumptions.
Assumption G.1 (Orthonormality). Embeddings of input and output tokens are orthonormal, i.e. e/ e = L = =}}, 2 j.j and ui Uj = Lii j}. V i.j.
Assumption G.2 (Initialization). The learnable matrix W is initialized from O when / = O.
Theorem 5 (Restatement of Theorem 3).
Assume Assumptions G.1 and G.2 hold, considering 2, 3 and (0.2, 0.4), train the full model f(-.;; W) with gradient flow. Denote
P(i.j: j; W) as the model's predicted probability for output j condliiioned on input i. Then, for l - op and E (1, have
(1) a - 1(j=), P(LJ;; (L,;:W))= (1- 1 = ) c
Remark Note that here the assumption (0.2, 0.4) a technical choice. In experiments, any value a E (0, 0.4) still has the same result.
Proof. W.l.o.g.. we assume the embeddings are standard basis in Rd. For any W, the gradient V L can be decomposed
1 -1 + +
[
V 1 = 71
[ 0].
(42)
Since W initializes from zero, this implies W can always be decomposed with the same basis
W = B1
-1 0
[
0 B2
[1
0].
(43)
2
Then gradient flow gives the following ODE
exp(-3: + (2) exp(31 + +.32)
B1 -- exp(- 81 + B2) + exp(31 B2) + exp( -2,2)
+ +1-a
expl 2231)
exp( 2xp(-231) + exppp 382) ) +1 +1-a,
(44)
(23g)
B2 = -2 exp(- B2) - (3) + (B2) + exp( (-232)
-3x
3 3exp( 3/22)
exp( -2,1) + exp(- -1 382) + 3a.
Denoting a -23,, -B1 3,3,, the ODE becomes
2 - 2 (x)
20,
cxp(a) (a) + exp(b) +1
(45)
2 - eex((b)
exp(a) + exp(b) +1 2 10a.
Lemma H.
gives the solution as, when + oc,
a + - log(t) - log(1 [1 - (4 - -2x), b + log a
59

---

Published as a conference paper at ICLR 2025
For the full model, taking the scores W1,: of the first input token example, have W11 31 + 82, W12 = -81 + 82, W13 = -22g, so the margins are
Wn W1 = =2,3 a, Wn1 - W13 31 + 3/3 = -b.
For the rank-1 1 model (assuming 31 f2), the margins
w (1) w 23,, w w (1) 31.
The proof finishes by computing softmax on the margins.
H USEFUL LEMMAS
Lemma H.1. Let p be data distribution on (I, v) E Bd x [N]. Consider training data as m i..dd
samples D - {(Ei, yi) C. R2 [N from p. Consider the following classification problem, with fixed output embeddings Wo:
L(W)= = Σ Mu.. We Wo;)).
The gradients take the following form: denoting pw (k((1)) as the current predicted probability of class k in IN + 1] classes for input Fz.
N-1
VwL(W) = Σ
Σ
(kw()) - 1i = k (k)a
When 0o, the above equation becomes
N+1
VwL(W) Σ ply W) - pT
where ji E[xcyy ki and ju
Eyy! fw (1e p(y=k)
(Relar s|.
Remark 2. This lemma is from Lemma 2 in Bietti et al. (2023).
Proof. Recall the form of the cross-entropy loss for classification with classes:
E
I(y,e)
Σ 1yy k} log
Σj eti
Its derivatives take the form
01
(y. s(E)x 11y k]
cl where s(E) Σ, c5
The gradient of L is then given by
N+1 Øl
VwL(W)
(y, W. W W2x)
m
N+
(Dw (k(z))) - = k
When - 00, the above equation becomes
N+1
N+1 Σ
VWL(W) Σ
(o())
kml
Am 1
Σ
j.k
N+1
p(u (k) T
60

---

Published as a conference paper at ICLR 2025
Lemma H.2. Consider sequence (Stte1 with S, d' .t where a + 1. Then Eiseer a(l-a*) a² T.
(a
Proof. Denote X₁ EiSIst St. Then we have Xt EesiST+1 a' -t-
(t- 1). Hence, it holds
1)-1)x,== -- d' a T+1 T a(l-a²) a T. Therefore, we have
1-a
X, a(1 a
T
+ a
(a - 1)2
Lemma H.3. Consider the following ODE with with a(0) b(0) and (0.2,0.4).
2 exp(a)
cxp(a) + exp(b) +1 2 2a,
2 - exp((b)
exp(a) + exp(b) +1 10a.
Then, when 0o, have
a - log(() log(1 a4(4 -0x), b - log
Proof. The ODE can be re- re-written as
à = 2 (a 2) exp(a) + (a - 1) exp(() +
2D
exp(a) + exp(b) +1
exp(a) + exp(b) +1
b b= 10 (a- - ) exp(a) + (a 1) exp(b) +
10E
exp(a) + exp(b)
exp(a) exp(b) +1
At /. 0, it holds a(0) < 0, b(0) < 0 since D 3a to decrease from Let's show D and E will never be positive when (0.2, 0.4) by contradiction.
< 0. Hence, and start The ending of the decreasing happens when one of 1) and F. gets positive.
0, F= 3a
Assume time Ti is when one of and equals to for the first time. This means E 0, because, for any time i, it always holds D < E since exp(a) 0 for any
R. Then at Ti, we have à 0 = O, which means exp(a) (a) will decrease for any small time window At > and exp(b) stays unchanged. Together with a >> 0.2, this means it has < again at time T₁ + At. Therefore, it is possible for E be O, but E. will never be positive. Meanwhile, this also guarantees D will always be negative because D < E.
Then, we make observation that when D is always negative and E is always non-positive, the decreasing nature of a will have D s E when / + do by exp(a) R 0. This implies log T-a Then, by taking exp(a) B 7,
a.
7, the ODE gives
1
(2a 4)3 7
7
3
I-a
which gives 7 1,8 (1=))
Therefore, when t + oo, we have
1
a log
b + log a
(1 (1-a)(4- 2a)
61

---

Published as a conference paper at ICLR 2025
1 INPUT EXAMPLES FOR LLMS
1.1 EXAMPLES FOR PREPOSITIONS
For experiments in Appendix C.I we use two synthetic datasets: inputs are 30 prepositions, and inputs are 40 incomplete sentences ending with a preposition.
The 30 prepositions are:
"about", "above", "across", "after". "against", "along". "around", "at". "before", "behind", "below", "beneath", "beside", "between", "by", "during". "for", "from", "in", "inside", "into", "near", "of", "on", "over", "through", "to", "under", "with", "without".
Generated by Claude (Anthropic 2024), the 40 incomplete sentences are:
[ ("Inspired painter gazed at pristine canvas, envisioning next creation about", "Children's 's delighted squeals filled yard as they frolicked, stumbling across". "Singer inhaled deeply. calming nerves before gracing stage before", "Ominous storm clouds amassed, promising downpour that would soon roll in", " "Awestruck trekker admired breathtaking summit vista, looking over". "Rich aroma of freshly roasted beans permeated cozy cafe, enticing during", "With deft sleight of hand, illusionist made coin vanish, leaving spectators awe without", "Majestic oak stood tall, branches reaching skyward above". "Gentle waves caressed shoreline, soothing rhythm lulling along". "Meticulous investigator scoured crime scene, searching for any evidence left behind". "Radiant sunbeams filtered through sheer curtains, warming hardwood floor beneath", "Concert pianist's 's nimble fingers glided across ivory keys, room resonating with melody around", "Crickets" evening chorus filled silent field from nearby meadow during". "Jubilant laughter resounded down corridor as jovial group headed towards celebration without", "Struggling poet tapped pen restlessly, seeking words to capture elusive emotion beneath". "Soothing patter of raindrops danced on windowpane, inviting serene relaxation with". "Mouthwatering scent of fresh bread beckoned passersby into cozy bakery without". "Mighty waves thundered against jagged cliffs, echoing roar along rugged shoreline around", "Seasoned trekker carefully navigated winding trail, cautiously avoiding exposed roots and rocks beneath", "Graceful ballerina flowed across stage, movements blending seamlessly with melody during". "Crackling campfire cast dancing shadows across gathered faces around", "Vibrant brush strokes danced across canvas, bold hues bursting into life before". "Photographer framed breathtaking sunset, capturing fleeting beauty over glistening ocean without", "Stern librarian hushed raucous group, reminding them to stay quiet inside", ; "Ink flowed from author's 's pen, words brimming with raw passion as page filled during", "Earthy aroma of freshly steeped tea perfumed air, inviting moment of serenity along". "Masterful guitarist's fingers danced nimbly across strings, room alive with haunting melody around", "Meticulous chef artfully garnished plate, adding delicate finishing touches over", "Indomitable marathoner pushed through punishing final stretch, fortitude driving every stride before". "Engrossed scientist examined specimen's intricate structures through microscope beneath", "Nervous thespian steadied breathing, striding into dazzling spotlight, delivering flawless performance with". "Skilled artist's pencil glided gracefully, deftly capturing subject's essence without", "Weary hiker paused to catch breath, marveling at sweeping panorama from lofty peak above", "Deep in thought, writer drummed fingers, secking perfect phrasing to convey profound emotion without", "Lost in reverie, violinist swayed gently, fingers dancing across delicate strings during", "Painter's brushstrokes burst into radiant life, canvas ablaze with vivid sunset hues over", "Adept photographer framed picturesque scene, preserving landscape's beauty without", "World-renowwed chef meticulously garnished plate, each component strategically placed around", "Dedicated researcher scrutinized specimen under microscope, documenting minute details beneath", "Seasoned actor inhaled deeply, embodying character as bright lights engulfed stage with",
I.2
MORE EXAMPLES OF FACTUAL RECALL
We consider more examples of factual recall with pairs of input and output shown Table
62

---

Published as a conference paper at ICLR 2025
Table 5: Inputs and Outputs of Factual Knowledge
Input
Target output
The Great Wall is located in Mount Kilimanjaro is located The Nobel Prize is awarded in
China
Tanzania
Sweden
The Statue of Liberty stands in Vatican City is enclosed within The Acropolis is situated in The Sydney Opera House located on The Galápagos Islands belong to The Aurora Borealis can seen The Amazon River flows through The Andes Mountains extend through Machu Picchu is found in The Kremlin is located in Uluru is landmark found in Petra is an archaeological city in Angkor Wat is located in The city of Toronto is in
New York Harbor
Rome
Athens
Bennelong Point
Ecuador
Norway
Brazil
Chile
Peru
Moscow
Australia
Jordan
Cambodia
Canada
The city of Barcelona is in The city of Mumbai is in The Fiffel Tower is located
Spain
India
Paris
J SYNTHETIC IOI TASK
Data and task. Here we consider a synthetic data model similar to the IOI task Wang et al. (0022).
with additional noise. Consider vocabulary V 21,2,.... N, N + 1]. The token N + the generic noise token. We fix trigger token [N], which governs in-context recall, and context length 7'. Each sequence of tokens 21.7 [21, ......, 2T] is generated as follows:
i. Sample a correct output token ÿ and a different distractor token y uniformly in [N].
ii. Sample three indices 11, 12, i3 for non-overlappingg)
[T 2] such that their distances are no smaller than (This
iii. Set 2i2 = 2ig = Among the three indices i + 1, + 1, + 1, random select one of them with Zic+1 = ÿ with the other as Zit+1 = D
iv. Set 27 q and sample 21+1 P.. with
a,
if x = y, if oz otherwise.
(z. = a2
0,
Random fill with tokens from V {q) into the remaining positions in [T+1]) (4,, é &2 1, 13, 13 + 1,T,T +1..
The key difference between the above data noisy in-context recall in Section Iss that, in additional to detecting the tokens and o after the trigger g. this task also requires counting to decide which of y and y² appear more. This mechanism is exactly the definition of the correct 10 token in Wang et al. (2022).
Most the other settings are the same as that in Section 3 including the training procedure, the architecture of transformer layer, dimensionality and the vocabulary size.
shows the test performance for models with layers L = 3, 6, where the models are trained with SGD. Dropping the last-layer MLP consistently improves the test
Results. Figure 19
performance across all models. Figure 20 shows the test performance for = 3, 4, 5 trained with Adam (Kingma 2014). Truncating the last MLP's input weights with p 0.01 significantly
63

---

Published as a conference paper at ICLR 2025
Test Loss vs Model Layers
Test Accumcy vs Model Layers
2.0
0.8
1.5
Test Accurncy
0.6
1.0
0.4
0.0
L-Jayer model
L-layer modal
Full
Drop
Drop (L 1)-th MLP
Figure 19: Synthetic IOI trained with SGD: test loss and accuracy for transformers with different
layers. Dropping the last-layer MIP consistently improves the test accuracies across all models.
improves the performance for L 3, We also that the model fails converge for possibly because we do not use any normalization technique in the architecture, so the Adam training less stable for deep transformers.
64

---

Published as a conference paper at ICLR 2025
layers: Loss along Training
layers: Test Accuraey along Training
10
Test Accuracy
Toa Losa
16º
0.2
0.0
L5000
Tizining stops
Trnining steps
p 0.01
p
layers: Test Loss along Training
layers: Test Accuracy along Training
0.8
10
Accuracy
0.3
014
Text
0.3
0.1
0.0
L.500
Trining
Tmining
Full
p=
MLP: p
layers: Test Loss along Training
layers: Test Accuracy along Training
0.5
0.4
Avcurcy
Text Loxs
0.3
10
0.2
0.0
5000
10000
L5000
20000
5000
10000
15000
20000
Tiaining steps
Tmining steps
Fell
Last MLP:: p 0.01
Last MLP: p= 0
Figure 20: Synthetic IOI trained with Adam: test loss and accuracy for transformers with layers L 3,4,5. Truncating the last-layer MLP's input weights with 0.01 improves the test performances for = 3, 4, while the model fails to converge for L = 5.
65


Reviews:
[{"review_id": "GpsHOptZ2I", "reviewer": "Reviewer_U1qi", "summary": "This work theoretically examines a phenomenon regarding the purpose of the feedforward and attention layers of a Transformer model, relative to the reasoning and knowledge learning tasks. It attempts to provide an explanation for the difference through a generic noise token, demonstrating that the feedforward layers store noise while the induction heads in the attention layers are insensitive to the noise tokens.", "questions": "* I would suggest moving the proofs of Lemmas C.1 and C.2 up, before they are invoked in the proof of Theorem 1 / 4. This would make it easier to follow the progression of claims made. \n* In the abstract, the term “gradient noise” is used in line 21, but not referred to in the rest of the paper — please clarify (or remove, if not relevant) this term. \n* The use of the term 'sequel' on line 125 is unclear; this should be clarified or rephrased accordingly. \n* As noted in the weaknesses section, are there any architectural changes that may be proposed, conditioned on the results of the exploration? Is the key takeaway that one should perform low-rank truncation as in LASER on the feed-forward layers since they store noise? Including discussion on this would improve the actionability of the findings of this work and its impact.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper sets out to address a clear objective, and designs a 2-layer model that enables a theoretical analysis of the phenomenon of interest, then reduces to a 1-layer model for theoretical analysis, revealing important, mathematically-grounded supporting evidence for said phenomenon. \n* I appreciate that the tasks of interest, and the distinguishment between the two abilities of interest (distributional associations and in-context reasoning) is made clear upfront. \n* Figure 2 is especially helpful to follow the 2-layer setup for the in-context recall task with noise tokens — the finding that the induction head filters out the noisy token is an interesting insight, and reinforced by Theorem 2 and its proof in Appendix D. \n* While the paper is dense, it is also organized quite well and in a manner that emphasizes the key takeaways.", "weaknesses": "* The novelty presented the results section (Section 4) is unclear — are there any architectural insights informed by the theoretical analyses that could be applied that have been considered, beyond an application of LASER? Aside from this, this section appears to be a more in-depth, investigative analysis on the behavior of the chosen pre-trained models for the selected tasks; it is also unclear how / whether these insights scale as the models considered are small. \n* I have no major concerns about the theoretical results in the work, which reveal the means through which noise is stored both in the presence and absence of feedforward layers, under a simple 1-layer model. I must acknowledge that while I have largely studied and followed the proofs of Theorems 1 and 2 and the theory in Appendix A, I have not thoroughly verified the correctness of the theory in the rest of the appendix.", "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": "MruUF8ZzA0", "reviewer": "Reviewer_Gq5Y", "summary": "Large language models excel in in-context reasoning tasks like coherent language generation and knowledge storage. Their Transformer architecture relies on feed-forward and attention layers linked to knowledge and reasoning. In controlled experiments, the authors found feed-forward layers learn simple distributional associations, while attention layers focus on context reasoning. Gradient noise explains this difference. Pre-trained models, like Pythia, show similar disparities in simple reasoning tasks.", "questions": "1. In addition to feed-forward layers and attention layers, there are also wording embedding layers, residual connections, and LayNorm layers in the Transformer architecture. What special role do these layers play in helping the model understand the semantics of the natural language? Did the authors explore?\n\n2. Because the attention module is the interaction between all tokens in the input sequence, it intuitively plays a role in understanding context. It seems obvious that the attention layer mainly focuses on in-context reasoning, right?\n\n3. Are the given conclusions applicable to all models based on the Transformer architecture? Is this applicable to all LLM models?\n\n4. In terms of format:\n    Line 1935, Line 1938, Line 1952, Line 2001, Line 2069 exceed the margins and have incorrect formatting.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper uses a large number of experiments and theoretical proof to draw the following conclusions:\nFor the Transformer architecture of large language model LLM, feed-forward layers learn simple distributional associations, while attention layers focus on context reasoning.\n\n2. The experimental results and conclusions are relatively solid, and the theoretical proof is quite convincing.\n\n3. This work may be helpful for future researchers studying LLM.", "weaknesses": "1. This work is highly theoretical and lacks certain practical significance and applicability.", "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": "h7tvWV3SH6", "reviewer": "Reviewer_d4XH", "summary": "In this paper, the authors investigate the role of transformers' self-attention and feed-forward networks (FFNs) both experimentally, in a controlled environment with a simplified model, and theoretically, using an even simpler model that allows them to derive calculations. The authors thus highlight an important distinction: FFNs are biased towards learning distributional associations (i.e. co-occurrence cues), while self-attention focuses on in-context reasoning. \n\nConcretely, as it is common to have the word “the” after the word “in”, a model using distributional cues only would be tempted to complete the sentence “Madrid is located in” with “the”, whereas a model using in-context reasoning would complete it with “Spain”.\n\nSo the authors set up an experiment in which the model is trained on a text containing with equal probability the “a + b” and “a + noise” patterns (where “a” varies from sentence to sentence, while the “noise” token remains the same) and then evaluated on a text containing only the “a + b” pattern. The idea is then to measure in-context recall in such a scenario: indeed, a purely distributional model will perform at 50% accuracy where a perfectly in-context reasoning model will achieve 100%. Using LASER to reduce the rank of the second-layer MLP matrix, until it is completely suppressed (rho = 0), they note that distributional behavior disappears with MLP rank, suggesting that FFNs are biased in this direction. They also note that self-attention manages to focus only on “a + b” patterns, ignoring “a + noise” patterns. \n\nThe authors then propose a theoretical explanation for such behavior, analyzing in particular the first gradient descent step. Finally, the authors propose to test their findings in a more realistic setting, i.e. on real models and real-world reasoning benchmarks.", "questions": "1) Why did you use such a small model for the experiments? Would it be possible to perform the same kind of analysis with more layers?\n2) What is your interpretation of the poorer performance of truncated MLPs after 8-shot on GSM8K (Table 2)?\n3) Your results seem to indicate that reducing FFNs greatly improves in-context reasoning, which is often desirable. However, FFNs also play an important role. It would therefore be interesting to study what is lost when they are reduced/removed?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- It is very well written and clear.\n- Their experimental environment is very clearly established, from the creation of their dataset to the architecture of the model.\n- I think that experimenting with LLMs in a controlled setting is an excellent methodology.\n- The paper proposes a theoretical justification for their discovery that is more than welcome and very rare in the field.\n- The findings are very interesting and help to better understand the behavior of LLMs, in particular FFNs, which are often less studied than self-attention.", "weaknesses": "- The main experiments and theoretical results are based on highly simplified models. However, the observed behavior may not scale with larger models.\n- The theorems and their consequences are relatively complicated to understand for a non-initiated reader of this literature, perhaps try to explain the different terms a little more.\n- Experimental figures have no error bars.", "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": "XQms6owBHK", "reviewer": "Reviewer_P7Ta", "summary": "The authors studied how attention mechanisms and feed-forward layers handle distributional associations and in-context inputs differently. They considered simple two-layer transformers on an in-context recall task and observed the distinct focuses of feed-forward layers and attention mechanisms through analyzing gradient.\nThey further employed LLMs like GPT-2 and Pythia, concentrating on IOI dataset . Additionally, the authors validate the effectiveness of the LASER technique.\nOverall, the authors underscore the differing roles of attention layers and feed-forward layers in language models.", "questions": "1. Is there potential modifications or improvements for the authors to propose to LASER that could strengthen this paper’s contributions?\n2. In Section 4.2, why LASER struggles with more shots? \n3. Is it possible to make other modifications beyond those in [2], such as designing an entirely new dataset to better validate your claims?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors conducts an in-depth study on how attention and feed-forward layers differently process distributional associations and contextual inputs, yielding clear and significant conclusions.\n2. The flow of the paper is easy to follow. \n3. The authors employ clear theoretical proofs and a logical structure, supporting the authors' arguments.", "weaknesses": "1. The experiments in this paper are insufficient to demonstrate authors’ claim, they only evaluate with limited number of datasets and model structures, which might not be representative of the general behavior.\n\n2. The technique authors recommended, LASER[1], is not introduced by this paper, making this paper seems an extension or complement to LASER. Differentiation between this paper’s contributions and [1]’s should be claimed.\n\n3. The methods and datasets this work employed is interesting but mostly appear a simple extension of that in [2]. This makes the research seem a supplementary work.\n\n**Reference**\n\n[1] The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction\n\n[2] Birth of a Transformer: A Memory Viewpoint", "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 (Poster)