Rethinking Associative Memory Mechanism in Induction Head
Authors: Shuo Wang
Issei Sato
Venue: COLM 2025 (2025)
Keywords: transformer
induction head
associative memory
positional encoding
TL;DR: This paper investigates how a two-layer transformer thoroughly captures in-context information and balances it with pretrained bigram knowledge in next token prediction, from the viewpoint of associative memory.

Abstract: Induction head mechanism is a part of the computational circuits for in-context learning (ICL) that enable large language models (LLMs) to adapt to new tasks without fine-tuning. Most existing work explains the training dynamics behind acquiring such a powerful mechanism. However, it is unclear how a transformer extract information from long contexts and then use it to coordinate with global knowledge acquired during pretraninig. This paper considers weight matrices as associative memory to investigate how an induction head functions over long contexts and balances in-context and global bigram knowledge in next token prediction. We theoretically analyze the representation of the learned associative memory in attention layers and the resulting logits when a transformer is given prompts generated by a bigram model. In the experiments, we design specific prompts to evaluate whether the outputs of the trained transformer align with the theoretical results.

Published as a conference paper at COLM 2025
Rethinking Associative Memory Mechanism in Induction Head
Shuo Wang
Issei Sato
The University of Tokyo wang-shuo311 12.. ecc. www. w.t.kko..ccc..
The University of Tokyo satoêg. ecc. u-tokyo. ac.
Abstract
Induction head mechanism is part of the computational circuits for in- context learning (ICL) that enable large language models (LLMs) to adapt to new tasks without fine-tuning. Most existing work explains the training dynamics behind acquiring such powerful mechanism. However, is unclear how a transformer extract information from long contexts and then use it to coordinate with global knowledge acquired during pretraninig. This paper considers weight matrices associative memory to investigate how an induction head functions over long contexts and balances in-context and global bigram knowledge in next token prediction. We theoretically analyze the representation of the learned associative memory in attention layers and the resulting logits when transformer is given prompts gener- ated by bigram model. In the experiments, we design specific prompts evaluate whether the outputs of the trained transformer align with the theoretical results.
1 Introduction
In recent years, transformer-bbsed models, such as BERT (Devlin et al., 2019) and GPT (Radford, 2018; Radford et al., 2019; Brown al., 2020), have achieved remarkable success in natural language processing. Especially, in-context learning (ICL) (Brown et al., 2020; Dong et al., 2024) has emerged as groundbreaking capability within large language models (LLMs), enabling them to adapt to new tasks without traditional fine-tuning. Instead, these models leverage patterns from a prompt or input sequence, effectively learning "in context" by interpreting examples or instructions provided in real-time (Liu et al., 2022; Wu al., 2023). During phase change, models acquire the ability to complete complex patterns through induction heads, suggesting that induction heads are key mechanistic basis of ICL (Olsson et al., 2022). is known that the induction head mechanism, which is a pattern of attention in transformer that enables the model to copy and reuse information from earlier in the context, can emerge in two-layer transformers (Élhage et al., 2021), and many theoretical studies have focused on analyzing two-layer architectures to understand how induction head develops (Edelman et al., 2024; Nichani et al., 2024; Chen al., 2024b). Among them, Bietti et al. (2024) demonstrated, both theoretically and empirically, that training a transformer leads to weight matrices behaving as associative memories, composing the induction head mechanism.
However, the effectiveness of the induction head mechanism tends to diminish for the latter part of sequence (Bietti et al., 2024), which does not align with the behavior observed in real-world L.M. In addition, while the learning dynamics of the induction head have been elucidated in the existing work, it remains unclear how in-context knowledge provided in the prompt and global knowledge acquired during pretraining influence the output of two-layer transformers. Understanding how a transformer predicts the next token is crucial, as it allows us to model phenomena such context hijacking (Jiang et al., 2024), in which altering the context disrupts factual recall, causing the model to produce incorrect outputs influenced by misleading in-context information.
In this paper, we theoretically and empirically show, from the perspective of associative memory, how two-layer transformer with relative positional encoding (RPE) can avoid

---

Published as a conference paper at COLM 2025
Theoretical analysis
Paper
Model
Objective
Meng et al. (2022)
GPT
knowledge editing
no
Cabannes et al. (2024a)
linear model
scaling laws
yes
Cabannes et al. (2024b)
linear model
learning dynamics
yes
Nichani al. (2024)
two-layer transformer
causal structure
yes
Friedman et al. (2023)
modified transformer
transformer programs
no
Jiang et al. (2024)
one-layer transformer
concept association
yes
Chen et al. (2024a)
linear model
noisy classification
yes
Bietti et al. (2024)
two-layer transtormer
learning dynamics length generalization & global v.s. in-context
yes
OURS
two-layer transformer
yes
Table 1: Comparisons among related work and ours in terms associative memory. The concept of associative memory is widely adopted for theoretical analyses of various phe- nomena.
failure in detecting patterns that appear the latter part of a sequence. We also investigate how the model prioritizes in-context knowledge versus global knowledge when predicting the next token, using training data generated by a bigram model with triggered transitions. To summarize, make the following contributions:
We theoretically show that the induction head learned by transformer using RPE can avoid neglect of patterns within the framework of associative memory.
We experimentally confirmed that the model with RPE can comprehensively use in-context knowledge for the bigram model. We also demonstrated its ability to effectively infer on sequences where the transformer using absolute positional encoding (APE) fails in next-token prediction.
We investigated how the model prioritizes in-context knowledge versus global knowledge when generating outputs, using both mental evaluation.
theoretical analysis and experi-
2 Related work
Associative memory Associative memory refers to the storage and retrieval of patterns, and its computational model was designed to retrieve complete memories from sufficient fragments of information in neuroscience (Amari, 1972; Hopfield, 1982; 1984). After the emergence of modern Hopfield networks (Krotov & Hopfield, 2016), which store patterns in an energy function and retrieve them using an updating rule. Moreover, many studies employ the concept of associative memory in various ways. For instance, Zhang et al. (2024a) propose memory unit that estimates conditional probability distribution with Gaussian kernel smoothing. Some work considers attention blocks memory because they read and write information flowing through the residual stream (Nichani et al., 2024; Friedman et al., 2023). Meng et al. (2022) locate factual association patterns in the weight matrix of an MLP. Recently, it has become increasingly popular to regard matrix as associative memory from a theoretical perspective, as shown in Tab. 1. Previous studies examine scaling laws in relation to model capacity (Cabannes et al., 2024a), learning dynamics of linear model in terms of particle system (Cabannes et al., 2024b), and the effect of low-rank approximation of weight matrix to classification involving noisy token (Chen et al., 2024a) for linear model. Additionally, Jiang et al. (2024) demonstrate that a one-layer transformer can recover latent concept from long enough context.
In-context learning and induction head Since Brown et al. (2020) introduced ICL, theoret- ical understanding of this phenomenon has been a major interest within the community. For example, Akyürek et al. (2023) showed that transformer can implement learning
2

---

Published as a conference paper at COLM 2025
algorithm that solves linear regression. Von Oswald et al. (2023) demonstrated that linear self-attention layer can emulate gradient descent. Similarly, many studies (Ahn et al., 2023; Mahankali et al., 2024; Zhang et al., 2024b; Dai et al., 2023) examined the theoretical con- nections between ICL and gradient descent. Mechanistic interpretability, which aims to understand the model behavior by reverse-engineering its internal components (Nanda et al., 2023; Wang et al., 2023a; Conmy et al., 2023), is another perspective for analyzing ICL (Olsson et al., 2022; Elhage al., 2021). In particular, Elhage et al. (2021) discovered that two-layer attention-only transformers understand the previous use of the token and look for similar situations in the context. In addition, Olsson et al. (2022) provided empirical evidence that induction heads are the mechanism behind ICL. Moreover, Bansal et al. (2023) identified that many induction heads are also among the attention heads deemed critical for ICL. The most relevant study to ours is by Bietti et al. (2024), who investigated the learning dynamiis of the induction head mechanism in two-layer transformers from the perspective of associative memory. In contrast, our study focuses on the functionality of the induction head mechanism and the influence knowledge acquired during training and that obtained in context the final output.
Preliminaries
3.1 Transformer architecture
In this work, we focus two-layer transformer described below.
Calculation Given sequence of tokens 21:7 of length T E N from the vocabulary set
V, each token 2, is embedded into a d-dimensional vector and the calculation proceeds as follows:
x, = wr(Zz) E R"
(0)
xi = i xix
(1)
(0,z) (o (( (0x)) T wox,
(0,))
(04] w +
(0,)) x(00) ER"
(1)
(2) = W&WG (1) o ((W2 (1) Wzx ) x E ER.,
(1)
x
(2)
W ())) + x
(3)
2111 arg mammmaxxWuuu
(4)
DEV
where WE dx an embedding matrix, : R' R' is the softmax function, x000),xx00) and x (0p) are the token embeddings, which may or may not include positional information, 4 E RR represents W/W, Wy = (. ()) E R V is an unembedding matrix, and represents the i-th vocabulary in V. The feed-forward layer consists of matrices W1 E d and W₂ RRVx.
Note that Bietti et al. (2024) incorporate absolute positional encoding:
x (0.)) = x, (0.4) (0.p) = wF (zz) + Pt,
x,
where Pr is the f-th absolute positional encoding. In contrast, we use simplified relative positional encoding:
wg(Z1n) + R1 t:0, and x, (04) = x1 (0,r) we (Zt),, x Lat
(0.))
where positional information R1-r0 (n t ro) E Rd and r.; represents the positional information of the i-th previous token 2, from the current one 2,-
3.2 Associative memory
Associative memory plays an important role in analyzing the behavior of memory recall at induction head in the next section. Before defining associative memory, first adopt
3

---

Published as a conference paper at COLM 2025
technical assumption commonly used in theoretical studies transformers, that is, that embeddings are high-dimensional random vectors, allowing them be nearly orthogonal. Mathematically, we impose the following assumption.
Assumption 1 (Near-orthogonality (Huang et al., 2024; Li et al., 2024; Bietti et al., 2024)). The embeddings (u;); are d-dimensional vectors with Gaussian entries, each having a unit norm ||uil| = 1 and i; Uj = 0 for + Also, Wou; forms new vector that has near-unit norm and near-orthogonal (u;); if Wo IR IR x is a Gaussian random matrix with appropriate entries.
With the assumption, we define associative memory.
Definition 1 (Associative Memory). Suppose (u:); and (v;); are the embeddings satisfying Assumption Then, A matrix W is called associative memory if can be expressed as a weighted sum of outer products of u; and Dy:
W
Σ Laijuio
where aij E the score for the pair (ui, Uj).
Thanks to orthogonality, we can derive the score &ij for the pair (u)) using the operation u, Woj. Henceforth, we consider the weight matrices in the attention layer and the feed- forward network of the transformer as matrices that implement associative memory.
3.3 Induction head
Induction heads are specific type of attention mechanism observed in transformer models, responsible for detecting patterns in token sequences and utilizing previous occurrences of the same token to improve predictions. The behavior of induction heads is particularly relevant for tasks that require understanding and leveraging repeated patterns, such as in language modeling.
In Fig. 2, we summarize the process by which the induction head identifies and outputs repetitive patterns in the input. When the prompt ABC A is given, the first-layer attention, which is called previous token head, copies information from the previous token the current token. Consequently, the token B retains the information of A in the form of Φ A. Then, in the second-layer attention, the associative memory W2 matches the pair 4 A with A, attending to the position of token B. As result, the information of B is stored in the last token A as W.B.B Finally, the output matrix W& transforms W? B into (B))), producing an output that follows the repetitive pattern.
3.4 Bigram model with triggered transitions
This section describes the bigram model with triggered tokens used in Bietti et al. (2024).
In this bigram language model over vocabulary of size V, sample trigger tokens q from and output tokens from Ni each time we generate training sequence.
a distribution Output tokens al always follow their corresponding trigger tokens in the generated sequences.
if é {9ch. [ = ox if = 4k-
(i (i
Motivation for the model In written texts centered on particular theme, it is common for identical patterns of tokens to recur multiple times. For instance, in texts about Harry Potter or machine learning, the token "Avada" is often followed by "Kedavra," and "neural" is typically followed by "network." The induction head is computational mechanism capable of detecting repeated patterns in context and facilitating the output of the tokens that complete these patterns when trigger token appears.

---

Published as a conference paper at COLM 2025
The next token following given token is determined by conditional probabilities in pure bigram model. As a result, bigram sequences often lack consistent repeated patterns, which prevents the induction head from emerging. In contrast, pairs of trigger and output tokens in a sequence enable two-layer transformers to learn a circuit that leverages repeated patterns, rather than relying solely on bigram conditionals.
Theoretical result
We analyze two key aspects of induction heads: the impact of positional encoding on the neglect of patterns in sequence, and the contributions of in-context information and pretrained knowledge during inference.
4.1 Neglect of in-context patterns
In this section, we address the theoretical aspects of overlooking patterns in sequence by the induction head mechanism. We argue that the learning of the previous token head fails for transformers with APE, and thus, the model cannot take patterns in the latter half of sequence into account for the prediction. Mathematically, we observe the representation of the weight matrix W when weight matrices are sequentially trained, in the order of W w2, and W through one step of gradient descent. We refer to Theorem from Bietti et al. (2024):
Theorem (Theorem 3, Bietti et al. (2024)). For any the learned matrix achieves the following associative memory;
PL-WWPP Maa' (T (T-1) P(t =t (1- +o )
(6)
T2t
where pr is the absolute positional encoding for token position f, a, a' IR are constants, and P (fg = =--1] denote the probability that the first trigger token appears at position f f- 1,
From Eq. 6, we can see that the score for the pair of p:-1 and pt is inversely proportional to f. Therefore, as increases meaning as we move toward the latter part of the sequence the function of the associative memory diminishes.
In contrast, we prove that transformers employing RPE consistently direct the previous token head to attend the preceding token, regardless of the input sequence length. We calculate how the transformer's 's weight matrices are represented by following the same learning procedure as when using a transformer with APE. We again present the training outcome of the matrix W as associative memory, which is the key component for the previous token head. We provide the complete statement and proof in Appendix E.
Theorem 2 (informal). Under the setup described in Sec. E.2, two-layer attention-only trans- former with relative positional encoding learns the associations by one step of gradient descent for any vocabulary v E V:
T (kWE(W) naa'
P( =t 0(1)
naa' I P(g = t) 12
Σ
{o
(0(t))
1 V
+0
}
where 1 IR is the learning rate, a,a' IR are constants, by and T are the first and second occurrence positions of trigger token, and V is the size of the vocabulary set V.
Theorem suggests that the association between r-1 and wy(v) (v) is independent of both the vocabulary and the token position t. Therefore, regardless of the input sequence length, the previous token head will always attend to previous tokens with consistent strength.
Remark 1. transformer using RPE learns to attend to the previous token regardless of its position within the sequence. However, if the input sequence contains out-of-distributiu tokens, the induction head mechanism does not activate. such cases, only transformer with APE can effectively attend to the previous token.

---

Published as a conference paper at COLM 2025
We have seen that positional encodings have significant effect the induction head's ability to generalize to longer sequences, as shown in the comparison between APE and RPE.
In the context length generalization in transformers, models without positional encodings (NoPE) are often discussed (Kazemnejad et al., 2024; Wang et al., 2024). To deepen our understanding of NoPE settings, we show that a three-layer transformer without positional encoding can implement the induction head mechanism.
Proposition A construction exists for a three-layer transformer without positional encoding that successfully achieves the induction head mechanism.
4.2 Global knowledge v.s. in-context knowledge
Here, we discuss how two-layer transformer uses statistical information obtained from the training data and the information provided in-context by calculating the logits.
Associative memory construction achieving induction head We first introduce the fol- lowing lemma, which states that the induction head mechanism is achievable by setting the appropriate weights for the matrices in two-layer transformer with APE.
Lemma 1 (Bietti et al. (2024)). Define Q as the support of TT. The induction head can be achieved by constructing matrices in the following manner.
we w3= I,
W
Σ PiPl
V W Σ mu())(((
w = Σ (k((k (2
RE
where I is the identity matrix, and the matrices Wo Wh, and W2 are a Gaussian random matrix.
straightforward to prove that a two-layer transformer with RPE also has associative memory construction that achieves induction head by modifying W = Exeq Wo wr(k)r Now, using the transformer with RPE architecture, we define an associative memory transformer: model that incorporates induction head and feed-forward network storing knowledge learned through pretraining.
Definition 2 (associative memory transformer). A two-layer transformer with RPE is called associative memory transformer if its weight matrices are set as Lemma 1 except for
w (v we (V2) T
W&=
WE (k) kc
Σ
W1
:
(v((v)
W2
Σ
v
log (u in (1)
log (1( v2)wu(() ... Σ
V
log TTp (o))ou()))
This construction reveals that the feed-forward layer in Eq. 3 of Sec. 3.1 functions explicitly as a key-value memory (Geva et al., 2021). Each row of W1 (key) acts as detector for an input pattern wr (D),,, and each column of W2 (value) encodes global knowledge, i.e., the distribution Ty over the output vocabulary. Appendix D discusses the difference from the architecture of Bietti et al. (2024). Note that when the probability TT v) approaches +0, its logarithm log 7t (u |v) diverges to -oo. If we admit log 7y (u v) -oo, the transformer never outputs regardless of context. This does not align with the actual behavior of the induction head. Also, due to the limitations floating- -oint representation in computers, is not possible to represent co directly. Thus, we consider threshold and set nip 7 v)

---

Published as a conference paper at COLM 2025
Associative memory transformer on limited input sequences. We present the following proposition, which focuses on two token patterns in the input sequence.
Proposition 2. Suppose two-layer transformer is an associative memory transformer. Given a length-T sequence Z1:T where the last token zT 9, let the f1-th token be 2₁, Vi following 21 10 9. and let the 12-th token be 2t2 02 following 2t2-1 q, where & 02 and (3 ) < l2 without loss of generality. Also assume that there exists only one occurrence for each of Vi and 02 in the context Z1.T. Then, the logits Sor and for can be expressed as follows:
1te
e'230-1
So Z + log (T( q).
6o2
log ( q),
Z
where Z is the normalization factor of the softmax function and e Euler's mumber.
We can observe a general trend from Eqs. 7. First, if the model contains global knowledge of both tokens i1 and 22, the output token is influenced by the knowledge, in addition to the in-context knowledge. In particular, the logits of tokens that rarely appear in the training dataset become significantly low. In our associative memory transformer model, global knowledge never promotes the output of any specific token; rather, it solely suppresses inappropriate predictions, because we always have log (- I 0. Second, if token patterns, such as q P1 and q 22, do not appear in the training data, global knowledge affects the token logits to the same extent, due to the way WF defined using the threshold e in Def.
is worth noting that an associative memory transformer exhibits a significant difference in scale when handling global knowledge and in-context knowledge. For instance, if ( 9) 0.1, then log (T 9) R 223. Even with the functionality of an induction head, it is unlikely to predict the token patterns from the context. This is avoidable by reformulating W3 as in Def. and adjusting the strength of induction head. We consider this setting in the next paragraph, along with more general input sequences.
If neither token appears during training, the difference between the two logits before applying the softmax function E.2 5% is calculated as follows:
e(4 - (2) + 1 +
(8) The result implies that the token 82 is more likely to be generated if f1 is large and the distance between the tokens dn and U2, iee, 12 f1 is small. This trend arises from the first layer of the attention block, where the attention is not fully concentrated on the previous token, resulting in diffusion of attention across other tokens well. In other words, the more tokens that precede the current token, the lower the attention score assigned to the current token.
2 5% (41 +
-e
Stronger associative memory on general input sequences. Next, we introduce a more sophisticated analysis of in-context and global knowledge by relaxing assumptions in the sequence ZI.T. The transformer considered here has associative memory with larger scores, characterized by weighted sum of the individual terms of the weight matrices Sec. 3.3, where each term is scaled by an appropriate coefficient. We consider special case in which each term has the same coefficient T1, T2, T3 E R.
Definition 3 (stronger associative memory transformer). We redefine the following matrices W&. W and W3 in Def. 2 for a stronger associative memory transformer
W T, Wt, W2 T
W6 = T3W?
with the parameters T1, T2, T3 E R4. The other weight matrices remain the same. Proposition 3. Suppose a two-layer transformer is stronger associative memory transformer. Given a length-T sequence žn. where the last token is 2T q, let f(v) be the number of token pattern" "qp" appearing in the sequence for vocabulary E V. For sufficiently large Tı and T2, the logits &o can be expressed as follows:
f(0) + +1f0 9)1 = 9}
E. & dognz(0 ()) + T3.
(10)
E (E2-((0')) + +1 9
7

---

Published as a conference paper at COLM 2025
B1 and B2 Prediction and Global Prediction Ratios
Memory Recall Comparison during Training between APE and RPE
In-context ratio
6.7
Global
B1 ratio
B2 ratio
Rabio
6.4
Miemory
0.2
ARE
50
21
52 62 71 82
Iteration
Number of B1 Prompt
(a) Even when trained with sequences of length 256, the attention block in the first layer of TPE fails to attend to the previous token at positions f 128, while TEREE attends to previous tokens regardless of the positions.
(b) Ratio of predicting B1 and B2 given the prompt A B A B1, A B2...., B2, A". The transformer shows an increasing ten- dency complete the context with frequency token pattern A Bj the context grows.
the
Figure 1: (a) Comparison of two-layer transformer with APE (Thm. 1) and RPE (Thm. 2) capturing previous token information. (b) Output behavior as the frequency of the token pattern A B, increases in the input sequence (Prop. 3).
can be observed that the transformer equipped with general associative memory is affected by the proportions of each token pattern present in the context as well as by the global knowledge. In this case, it can be concluded that the information regarding the positions where token patterns are observed does not influence the final logit. We can easily derive the following properties.
Corollary 1. Given a length-T sequence Zι:T where the last token is 2T = q and the only token patterns q Ui and q 12 appear f(1) and f(02) times, respectively. The transformer outputs one of the following: arg max maX,gV((p,,o)) [T 9), 01, or V2.
For vocabulary v &f Vi, U2, the logit &o determined only by log TT (v | and this means that the candidate for the next token is arg max 7T q). On the other hand, The logits for iy and v2 have the other term in Prop. 3, which makes them other candidates for the next token. Overall, Corollary 1 states that the model output depends on the bigram conditionals Fp 9) and the frequency of token patterns quı and qu2 the context.
Experiments
This section empirically examines how positional encoding affects the transformer's abil- ity to capture patterns in input sequences, as discussed in Thm. 1 and 2, and how the transformer's output changes depending on the proportion of pattern occurrences within the context, as stated in Prop. 3. Henceforth, we denote by TFape and TF two-layer transformer with APE, RPE.
5.1 Neglect of in-context knowledge
Training process of previous token head We investigated whether the training of trans- formers directs attention towards the previous token in the first attention layer. Specifically, TFrpe and TFape were trained using sequences of length 256 generated according to the sequence generation rule. Details regarding the hyperparameters and sequence generation process are provided in Appendix G.
8

---

Published as a conference paper at COLM 2025
For evaluation of the learned associative memory Liiee / we employed the memory recall metric (Dar et al., 2023; Geva et al., 2023; Bietti et al., 2024):
1
R(W)
Σ
aarrg max up Woj =i}.
[M]
(LjjEM
For Trre we want each vocabulary wF(k) to be paired with the information of the previous position --1- Hence, R(W)) is computed as
1
Σ
marrg W -1}.
|QI|
KFQ
/
Similarly, we have for TFape
T 1 1arg p Wkpt =t 11,
Σ
1'
where T is the maximum input length for the transformer.
In the RPE setting, the metric quantifies the amount of vocabulary tokens that are more strongly associated with the relative position embedding T-1 than any other position em- bedding in the learned associative memory. We can find Fig. la that the recall approaches 1.0 as training progresses. This indicates that nearly all vocabulary items are successfully paired with r-1, confirming the effectiveness of RPE for associative memory formation. In contrast, for APE, we observe that memory recall is substantially higher for positions 128 compared to the full positions, even though they are trained with sequences of length 256. This supports our theoretical claim that it is harder to embed the relationship between pt and P1-1 into associative memory when t is large, due to the diminishing alignment between position embeddings distant positions.
Discussion One of the worst-case scenarios for TF occurs when trigger token does not appear by chance in the range S 128 and only appears for the first time at > 128. Although such sequence is contrived, its probability of occurrence is non-zero. For such sequences, models with an incomplete previous token head fail to associate the trigger token's 's information with the output token, resulting in incorrect predictions without leveraging in-context information. In contrast, TFxpe, which links words to their relative positions, ensures that the in-context information is not overlooked, regardless of where the trigger token first appears, thereby enabling correct predictions.
5.2
Global knowledge v.s. in-context knowledge
To validate the results derived from the theory, we conducted experiments using real data.
Analogical reasoning task We utilize the analogical reasoning task, one of the most practical tasks for bigram analysis. Specifically, we focus on the "capital-world" analogy type of questions from the Google Analogy Dataset (Mikolov et al., 2013), which contains a total of 19,544 questions. This subset includes 232 unique vocabulary items. Denoting word A (e-g., Tokyo) and its corresponding analogy counterpart A' (e.g. Japan), the prompts were constructed in the form A A" B B* CC* ,..., with pairs separated by commas. Using these generated prompts, we trained TF Trppe- For details on the training and dataset setup, please refer to Appendix G.
Collision of context information As shown in Corollary 1, the model predicts either token based on global knowledge or the token that appears in the context. To confirm this, we randomly sampled from the capitals
the Google Analogy Dataset and constructed prompts using word pairs A B₁ and A B2 that were not learned as global knowledge. The resulting prompts were of the form A B, By...,, A B, B2,.... A B2. A, which were then used for predicting the next token. We generated 1,000 prompts and calculated the propor- tion cases where either B₁ or B2 was predicted, while controlling the frequency of Bj and A B2 in the context.

---

Published as a conference paper at COLM 2025
The result is shown in Fig. 1b. As the number of token pattern A A increases, the model predicts Bj as the next token more frequently. Furthermore, prediction trend reverses when the frequency of B₁ and B2 in the context is nearly equal. Additionally, we observe from global ratio that the model predicted some vocabulary from global knowledge. Especially, the figure illustrates that it is more likely to output global knowledge when the number of B1 and B2 are the same, or when the number of B₁ or B2 is almost maximum. Please refer Appendix H.2 for further discussion on these phenomena.
can be observed that the
6 Conclusion
We analyzed the influence of interaction between global bigram and in-context knowledge to the two-layer transformer through the lens of associative memory. We theoretically and empirically verified that relative positional encoding empowers transformers with the ability to capture information in longer sequences, and that how the next token prediction is conducted within the transformer that was trained to store global knowledge and to have induction head.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614-45650, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Represenlations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
S-I Amari. Learning patterns and pattern sequences by self-organizing nets of threshold elements. IEEE Transactions on computers, 100(11):1197-1206 1972.
Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the role of scale for in-context learning: An interpretability- based case study at 66 billion scale. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 11833-11856, Toronto, Canada, July 2023. As- sociation for Computational Linguistics. doi: 10.18653 htt///122222222 2t06://///1222.aaooog..666 acl-long.-6666 660. URL https://accle http:://aclantholog// org/ acl -ong... 660.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Vivien Cabannes, Elvis Dohmatob, and Alberto Bietti. Scaling laws for associative memories. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024a.
Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent. ar Xio preprint 4rX1042402.18724, 2024b.
10

---

Published as a conference paper at COLM 2025
Lei Chen, Joan Bruna, and Alberto Bietti. How truncating weights improves reasoning in language models. Xiv preprint ar wxXi:22466.03063 2024a.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. In Advances in Neural Information Processing Systems 38: Annual Conference Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 - 15, 2024, 2024b.
Arthur Conmy, Augustine Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alonso. Towards automated circuit discovery for mechanistic interpretability. Advances in Neural Information Processing Systems, 36: 36:16318-16352, 2023.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta- optimizers. In Anna Rogers, Jordan Boyd-Graber and Naoaki Ōkazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 4005 4019, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/ (1111863/11/// w1/l/2132222 a02....nniingsacll 247. URL https: / htt////caaaanna wwwoo/oooooo.oooooooooooooooooo. org/2023. findings-acl. acl. 247.
Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant. Analyzing transformers in embed- ding space. In Proceedings of the 61st Annual Meeting of the Association for Computational Lin- guistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-14, 2023, pp. 16124-16170. Association for Computational Linguistics, 2023. 2023.......111666////224444AAA2 /11/022AALLLLONG..55
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Do- ran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chap- ter of the Association for Computational Linguistics: Human Language Technologies, NAACL- HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171-4186 Association for Computational Linguistics, 2019. doi: 10.18633/11/N19-111
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 1107-1128, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. URL httpp::/aaee http::::///claahhllogyoo org/ 22tg.//222...22 emnlp-main. 64.
Ezra Edelman, Nikolaos Tsilivis, Benjamin L. Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. In Adoances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1: 2021.
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. In Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forwan layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pp. 5484-5495. Association for Computational Linguistics, 2021.
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pp. MAM.N..-MAIN.55 V1 FM1/223.EMNI..
Association for Computational Linguistics, 2023. doi: 10.18653/
11

---

Published as a conference paper at COLM 2025
Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: decodiing-enhanced bert with disentangled attention. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
John J Hopfield. Neural networks and physical systems with emergent collective computa- tional abilities. Proceedings of the national academy of sciences, 79(6)2555--2588 1982.
JohnJ J Hopfield. Neurons with graded response have collective computational properties like those of two-state neurons. Proceedings of the national academy of sciences, (100) 3088-3092, 1984.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
Zhiheng Huang, Davis Liang, Peng Xu, and Bing Xiang. Improve transformer models with better relative position embeddings. In Findings of the Association for Computational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020, volume EMNLP 2020 of Findings of ACL, Pp. 3327-3335. Association for Computational Linguistics, 2020.
Samy Jelassi, Stéphane Ascool, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint ar arXim:2306.15400 2023.
Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Do lims dream of elephants (when told not to)? latent concept association and associative memory in transformers. In Advances in Neural Information Processing Systems 38: Annual Conference Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
Andrej Karpathy. The unreasonable effee of recurrent neural networks, May 2015. URL httpttttt///:aattttett git.ub. 2tt.://00..0502222mmmfeeeeeeeee1 Accessed: 2024-12- 15.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Guolin Ke, Di He, and Tie-Yan Liu. Rethinking positional encoding in language pre-training, In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Shun Kiyono, Sosuke Kobayashi, Jun Suzuki, and Kentaro Inui. SHAPE : Shifted absolute position embedding for transformers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLPP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pp. 3309-3321. Association for Computational Linguistics, 2021.
Dmitry Krotov and John J. Hopfield. Dense associative memory for pattern recognition. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Anmual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp- 1172-1180, 2016.
Yingcong Li, Yixiao Huang, Muhammed Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pp. 685-693. PMLR, 2024.
Tatiana Likhomanenko, Qiantong Xu, Gabriel Synnaeve, Ronan Collobert, and Alex Ro- gozhnikov. Cape: Encoding relative positions with continuous augmented positional embeddings. Advances in Neural Information Processing Systems, 34:16079-11092. 2021.
12

---

Published as a conference paper at COLM 2025
Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for GPT-3? In Eneko Agirre, Marianna Apidianaki, and Ivan Vulié (eds.), Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pp. 100- 114, Dublin, Treland and Online, May 2022. Association for Computational Linguistics. doi:
http://v12222.deeeeeee
URL https:/ https:///caaatthooo.o.o hrtp:/oggg.rg///2/// org/2022. deelio-l.11. 1.10.
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359-17372, 2022.
Tomás Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. In Yoshua Bengio and Yann LeCun (eds.), 1st International Conference on Learning Representations, ICLR 2013, Scottsdale, Arizona, USA, May 2-4, 2013, Workshop Track Proceedings, 2013.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Confer- ence on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Masato Neishi and Naoki Yoshinaga. On the relation between position information and sentence length in neural machine translation. In Proceedings of the 23rd Conference on Computational Natural Language Learning (CoNLL), pp. 328-338, 2019.
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning, ICML 2024, Viemna, Austria, July 21-27, 2024, 2024.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11 11899, 2022.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint vrXiv:2108.12409 2021.
Alec Radford. Improving language understanding by generative pre-training. 2018
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, (1144)):1667 2020
2020.
Jie Ren, Qipeng Guo, Hang Yan, Dongrui Liu, Quanshi Zhang, Xipeng Qiu, and Dahua Lin. Identifying semantic induction heads to understand in-context learning. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics: ACL 2024, pp. 9916-6932, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/ /00///v1222////// w112024.findingssa. findings-acl412. URL https: //aclanthology. or 22244 rindings-acl.441
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), ACL 2023, Toronto, Canada, July 9-14, 2023, pp. 1889-1903. Association for Computational Linguistics, 2023.
13

---

Published as a conference paper at COLM 2025
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT, New Orleans, Louisiana, USA, June 1-6, 2018, Volume 2 (Short Papers), pp. 464 468. Association for Computational Linguistics, 2018
Koustuv Sinha, Amirhossein Kazemnejad, Siva Reddy, Joelle Pineau, Dieuwke Hupkes, and Adina Williams. The curious case of absolute position embeddings. In Findings of the Association for Computational Linguistics: EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 4449 4479. Association for Computational Linguistics, 2022.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, PP. 35151-35174. PMLR, 2023.
Benyou Wang, Donghao Zhao, Christina Lioma, Qiuchi Li, Peng Zhang, and Jakob Grue Simonsen. Encoding word order in complex embeddings. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
Jie Wang, Tao Ji, Yuanbin Wu, Hang Yan, Tao Gui, Qi Zhang, Xuanjing Huang, and Xiaoling Wang- Length generalization of causal transformers without position encoding. arv preprint Xio: 12224, 2024.
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023a.
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words are anchors: An information flow perspective for understanding in-context learning. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 9840-9853 Singapore, December 2023b. Association for Computational Linguistics. doi: 10. 10.18653/v1/200 mminp- URL httpt:///cclannni h/ap://////laahhogg// org/2023. mmnlp-main.. 609.
Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Zhiyong Wu, Yaoxiang Wang, Jiacheng Ye, and Lingpeng Kong. Self-adaptive in-context learning: An information compression perspective for in-context example selection and or- dering. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1423-1436, Toronto, Canada, July 2023. Association Computational Linguistics. doi: 10.18653/v1/2033..oo
79. URL https: /tc///llllll/ https::///aaaathhlllooo wrw.:y....rg/eeee org/2023..22.. acl acl-long. 79.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In The Tenth International Conference on Learning Representations, [CLR 2022, Virtual Event, April 25-29, 2022, 2022.
Jianyu Zhang, Niklas Nolte, Ranajoy Sadhukhan, Beidi Chen, and Léon Bottou. Memory mosaics, 2024a. URL
06394.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. J. Mach. Learn. Res., 25:49:1 25:49:1-4::55, 2024b.
14

---

Published as a conference paper at COLM 2025
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in- context learning learn? bayesian model averaging, parametenization, and generalization. arXiv preprint XfXi::2305.1940, 2023.
Liang Zhao, Xiachong Feng, Xiaocheng Feng, Weihong Zhong, Dongliang Xu, Qing Yang, Hongtao Liu, Bing Qin, and Ting Liu. Length extrapolation of transformers: A survey from the perspective of positional encoding. In Yaser Al-Onaizan, Mohit Bansal, and Yun- Nung Chen (eds.), Findings of the Association for Computational Linguisties: EMNLP 2024, Miami, Florida, USA, November 12-16, 2024, pp. 9959-9977. Association for Computational Linguistics, 2024.
15

---

Published as a conference paper at COLM 2025
Output Matrix (W
A
W
A
Attention Layer (Wi w
4
4 B
A
Attention Layer (W and W w
residual connection
Figure 2: The visualization of induction head mechanism. The first attention layer copies the previous token information. Then, the current token matches the copied information to find the probable next token.
A Limitations
Our study has certain limitations that should be acknowledged. First, the analysis presented in this work is conducted using two-layer transformer model. While this allows for controlled and interpretable exploration of the underlying mechanisms, the findings may not directly generalize to LLMs, which typically involve significantly more layers and complex interactions.
Second, while our focus induction heads provides valuable insights into ICL, it is important to note that ICL leverages additional computational circuits beyond induction heads, As such, our explanation captures only part of the broader mechanisms underlying ICI,, leaving room for further investigation into other contributing factors.
B Ethics statement
This work is purely theoretical, supported by controlled experiments designed to validate the proposed analyses. No personal or sensitive data were involved at any stage of this re- search. All experiments were conducted using datasets that publicly available, ensuring compliance with ethical standards for data usage.
C Additional related work
In-context learning and induction head Bayesian inference is also popular perspective for comprehending ICL. Xie et al. (2022) and Wies et al. (2024) investigated how a pretrained transformer implicitly discovers the underlying concept latent task of prompt by leveraging a mixture of distributions. An alternative line of work (Zhang et al., 2023) relates the Bayesian model averaging to the attention mechanism. Wang et al. (2023b) studies how information is aggregated to label words in ICL. In terms of two-layer transformer, Chen et al. (2024b) studied ICL. on n-gram Markov chain, proving that gradient flow optimization under the cross-entropy loss leads to a model that exhibits a generalized induction head mechanism. Edelman et al. (2024) showed two-layer attention-only transformers learn induction heads that estimate the next token's conditional probability based on prior occurrences, enabling near Bayes-optimal performance. Similarly, Nichani et al. (2024) showed that two-layer transformers learn causal structure via gradient descent, suggesting that induction head is a special case of this structure. Ren et al. (2024) discovered a more generalized version of induction heads, which are referred to as semantic induction heads.
16

---

Published as a conference paper at COLM 2025
Length generalization and positional encoding Although we point out the neglect of in-context information by a transformer with APE, this can be considered as the failure of length generalization. Since the introduction of the Transformer architecture (Vaswani, 2017), there has been a wealth of research on positional encoding (Zhao et al., 2024). Originally, Vaswani (2017) adopted sinusoidal positional encoding that transforms the tokens' absolute positions. Subsequently, Kiyono et al. (2021) and Likhomanenko et al. (2021) proposed shifting absolute sinusoidal positional encoding during training to achieve shift invariance. Wang et al. (2020) extends word embedding as continuous functions with respect to token position, leading to smooth word representations. Meanwhile, is superior to APE for longer sequence generalization in various tasks (Neishi & Yoshinaga, 2019; Likhomanenko et al., 2021; Huang et al., 2020; Jelassi et al., 2023). Additionally, Sinha et al. (2022) questioned whether models can learn the relative distance between words when equipped with APE. The first RPE was formulated by Shaw et al. (2018), where a trainable encoding is added to the key before computing the dot product with the query. The transformer in our work also adopts this type of positional encoding, while we fix the positional encodings to randomly initialized vectors. A variety of different RPE methods were presented (Huang et al., 2020; Ke et al., 2021) until now, but DeBERTa (He et al., 2021), RoPE (Su et al., 2024), T5-Bias (Raffel et al., 2020), and ALiBi (Press et al., 2021) are the most popular choices. Ruoss et al. (2023) suggested that randomizing absolute and relative positional encoding enhance the model's ability to generalize to input sequences of lengths not encountered during training, but they used the encoder-ony Transformer, which is different from ours.
demonstrated that RPE
Short summary of Bietti et al. (2024) The remarkable success LLMs has increased the need for deeper understanding of their internal mechanisms and for enhancing their reliability. Existing studies lack detailed insights into how reasoning abilities evolve during the learning process based on contextual information. Bietti et al. (2024) explored the dynamics of balancing in-context knowledge and global knowledge through a bigram model, analyzing the development of these abilities as part of the training dynamics. They conceptualized the weight matrices of transformers as associative memory and theoretically demonstrated that the induction head mechanism can be trained by proposing the storage of specific embedding pairs from training data. In their experiments, they trained two- layer transformer on bigram model estimated from the tiny Shakespeare dataset. The proximity of each weight matrix to the theoretically derived weights was measured using memory recall probe. They analyzed how these weights changed over the course of training. Although this work is the most relevant to our research, we note that while they primarily focus on the learning dynamics of induction heads in a two-layer transformer, our primary emphasis lies on length generalization and the reasoning process.
D MLP as a key-value memory
Geva et al. (2021) discovered that the feed-forward layers in transformer-based language models function as key-value memory. In a transformer, the feed-forward layer is typically computed as
F(x) = ()(ReLLU(,x))),
where x E Rdd1 is the input vector, and W1 E R4xd W2 E R4 According to Geva et al. (2021), the rows of W1 (keys) serve to detect specific patterns in the input sequence, while the columns of W2 (values) have been shown to represent distributions over the output vocabulary.
Bietti et al. (2024) uses a transformer architecture that uses a linear projection Wy instead of an MLP layer.
V
V log To ()
WF
Σ
Σ
Due to the near-orthogonality of the embedding vectors, this projection maps an input Wr to the output ΣY= log ( ()) This enables the matrix We to function
17

---

Published as a conference paper at COLM 2025
associative memory that stores the pairs of an embedding vector and corresponding global knowledge.
Now, suppose construct the matrices W1 and W2 follows:
(e(w1)
(E((2))
W1
(11)
wF(2v)
V
W2
log [u by (u)
log (2( 02) .... log TB vy oy)wu(())
(12)
Each row of W1 acts as detector for an input pattern (F((), such that only the correspond- ing entry activated when wv(v;) is provided as input. Each row of W2, in turn, encodes global knowledge, i.e., the distribution over the output vocabulary conditioned on the corresponding input token.
We can easily find that the resulting computation matches exactly that of the linear projection WF under the oraronothhonnaltyy condition.
Therefore, this construction confirms that the feed-forward layer operates as key-value memory.
Learning of ideal transformer
E.1
Notation table
We prepare tab. 2, which provides a concise summary of the notations used throughout the appendices.
Symbol
Description
Dimensionality of the positional encoding
V
Vocabulary set
V
Vocabulary size
wE(Z))
Embedding of token Z₁
rs-t
Relative positional encoding expressing the relation between s-th token and current one softmax function
d
product of two randomly initialized matrices W1 and W/t Index of the first trigger token
tq
Index of the output token corresponding to the first trigger token, to ty Index of the second trigger token, or the input length set of natural number up to N, (1,2,...,N) cross entropy loss function learning rate
T
[N]
I
constant, defined by E, 1/f] constant, VT constant, /VTT
x'
q
variable representing the trigger token uniform distribution from which z is sampled.
πg
a uniform distribution by which the trigger token is determined.
To
uniform distribution by which the output token is determined.
a uniform distribution from which 22:7 is sampled conditioned on the previous token.
Table 2: Table of Notations
18

---

Published as a conference paper at COLM 2025
E.2 Setup
Here we consider simplified setting to analyze the influence of relative positional encoding in the training dynamics of our two-layer attention only transformer. Specifically,
1. Input Sequence We consider an input sequence 21.T E V which has one trigger token If appearing twice, and ends with the trigger token. In other words, let fy be the first occurrence position. Then, we have 21, 27 = 9. From the bigram generation rule in Sec.3.4, 2t and ZT are the output token.
2. Probability Distribution Assumptions: The bigram is generated by uniform dis- tributions over [V] for any index i.e,, Nur Ny. no, and πT( f) are uniformly distributed.
3. Simplification of Loss Function: Consider the loss only for sequences of length where the final input token ZT is the second occurrence of the trigger token, and the label is 2T+1, the corresponding output token.
4. Simplification for Learning Focus: Our approach involves sequentially training W W and W from top to bottom. We employ zero-initialization and carry out a single gradient descent step.
5. Initialization and Freezing: To achieve our goal of showing that W W and W learn to be associative memory, we zero-initialize the three matrices. For other matrices such as W/, W/ and W. are randomly initialized from Gaussian distribution. We set W and W3 to identity matrix. All these matrices except for wg. W& and Wg are freezed during the training process.
E.3 Theoretical analysis of learning
Theorem 3 (formal). Under the setup described in Sec. E.2, a two-layer attention only transformer with relative positional encoding learns associations of the associative memory transformer, i.e., the memories W W& and W can be written as the sum of the terms that constitute the associative memory transformer.
w = (x) (WRY,
(13)
kE@
W
Σ (()))a(()) (W3)',
(14)
REQ
3 Σ (())
(W3),
(15)
where X. Φ V and N + IR denotes the learned score of each association, and (Wt)', (W2)' and (W2y y are the other associations.
Proof. We will demonstrate that associative memory can be learned by performing one gradient descent step in top-down manner, starting with W3, followed by Wł, and then W..
Step1: Training of W3
We begin with the dynamics of W3. The input to the second layer attention, xf"), is composed of the sum of the values from the first layer attention, x; (1,0) and the values from the residual connection, (1.1) Since W is initialized to zero, the attention of the first layer is evenly distributed across all the tokens Z and we have
(1.0)
x
19

---

Published as a conference paper at COLM 2025
where 4 WoWp Additionally, the residual stream carries the token embedding:
xj (1,1) (z(zz)-
Since the matrix W is also zero-initialized, the value of the second layer attention can be computed through
x = w? ( 4 (
Now the logit prediction is given by Wux + x with the residual connec- tion. We note that when d is large, thanks to the near-orthggonality Wuxy (1) Wu [I_ (wwr(zs)/ T + wF(ZT)) is negligible because xf" does not contain wy(w) for any v E V. This enables us to use Lemma 2 and one gradient step with learning-rate yields
w3
(2) (k(EEE[x
wu
y = k]
y
wu
@w x) ply = k) xx
V
1 wu (k)
xT
(2) y = k]
E E
(16)
V
This transformation comes from the assumption that every token sampled from a uniform distribution (p(y = k) = 1/V) and all the logits are when W? = Odxá, which means that Pw(k x) = 1/V.
Given that y k, holds that the first output token 2t, k. For z; + fo), the uniform distribution can possibly generate any token, independently of the condition y = k. Since have
x
(2)
1
W3
w wF(Z))
E p (ow (o)(())) s=1
1=1
the expression IE
XT
(2)
y = k
xx
leaves only the terms related to 2t. This yields
(2)
y k E x
E
W3 [E] W (Zt,)
ly = k]- EB (z+,)
+ T
W/ D (e(wE)))) 1y k
W/ 4 (F((()))))
(17)
T
w (w= (lwe((E))) W3 Dy Elwx(ZZ,)),,
(18)
where 7 T: Σ t, :] Since each token generated from a uniform distribution, the expected value of we (zt,), conditioned on nothing, given by:
1 V
EE [wg(Zt)]
(k))).
Combining eq. 16 and 17, the near-orthogonality of Gaussian vectors gives us
(k( T W3 W)wE(()
y 1(k = +0
(19)
VT n VT
V2T
(k) T 13
~
1(k = +
17
(20)
1/T
20

---

Published as a conference paper at COLM 2025
This is the same result as in the proof from Bietti et al. (2024). So far, the weight matrix W is trained to perform as an associative memory so that
wu(k) W3WFWE(Z)) ~ a1(z) = k],
(21)
with constant a = TV.
Step2: Training of W
Since the training of W3 is finished, we have Eq. 19. When W W Oaxd Lemma demonstrates that the next token prediction is distributed among all tokens, which means that f(v x) 1/V for any vocabulary v. According to Lemma 5, one step of gradient
descent yields
W =
TN E
Σ
wwr (k) T (1) (x.") - x(1)) x (x ly = k]
D22
k,t
"
E wu(k) Φ2x1" (1) (x.") - x()) (x
(22)
TN k
with x(1) + 1T=1
(0)
We simplify our transformer architecture by setting (1,1) as the queries and values, and Xi (1,0) as the keys. Furthermore, we add the condition that the trigger token 21, = f and rewrite eq. 22. This conditioning leverages the fact that each token in 21, is sampled from uniform distribution. Now we have
N
W TN
E
wu (k)
(2x((11) (1.1) x (x(110)
(x 1y = k, =
7
Σ
IE (k) 02x; (L1) x (x, (10) - (x Ztg = i].
TN ktt
N
where x(1,0)
7 E
x
(1,0)
Using the formula 21, we obtain
N
N
W
an
E
{z = k) x (x/100 x(100) x(10)) (wE(ZT) 1y y Ztg =1
TN2
N
0.7]
Σ IE
Σ
1(z== = x (xf110) - x(1.0) ((=0)) 21, =
TN2
=1k=
N
N
Ax,wwr(i
(23)
TN² j=
where 4xj is defined by
T
Σ 1lz k] (xf(.0) x(1,0) 1 -k, =
Auj
E
1 1(z =3 xxx x(10)) 21, = 1]
E
The sum in Arj can be partitioned into three distinct groups: (i) Ag,- where I is the index of the output token; (ii) 42 where corresponds to the indices of the trigger tokens; and (iii De,p which includes all other cases. Mathematically, let to be random variable, and
tq = to - 1, we write
42 = E 1[2.. = =x) (x(0) x(10)) ly = k, k, =1
21

---

Published as a conference paper at COLM 2025
E 112 = (x (xx] - x(100) 2tz = =i]
1/
E
Σ
((2 (x} (1xf x(10)) 2t =
1
1(z) = (xxxx(1) x(10))| Zzp =1|
E
E
(ETy
M
E
Σ
1 = =k (x, (x,1 x(100) y 2, =
x(1,0)
LITT,
x(200) | 21, = = 1(z - kx (x
E
Σ
IET,
where To (lg T} and T, (1 sisT to + 1g, # T}.
First, we will take look at 4//- Note that El E[1 = k} y k, [B[ k} 2,, = = 1/N, so find that
and
1 =
1 N
E
(1,)
x(1.0) y y=
N
1
1 I
Σ
Φ WE (zs)
2 PI
N
to
T
5=1
1=1
N
1
N
Σ
anjj
0 w (i),
(24)
N
where Akjj is the coefficient of d (i) by calculating
1 te to s=1
1 T 1 wc((z)
E
Σ
O (e (zs)
Σ T
(=1
s=1
The analysis of akjaa is done by the work (Bietti et al., 2024):
1
N
if こ
akii
1
otherwise.
Thus, we have the following approximation:
N Σ
O
if
(+IWE(I) T
Dij
(25)
n
ifi
Similarly, we have - - (()) and 4/.- = O (z) based on the discussion in Bietti et al.
(2024). By combining these analysis with Eq. 25, we multiply (f)) from left and
(r.(j) from right to Eq. 23 and we establish that
(PIWE(1) W2WE()) ( ~
an TN
a
(
1(1 = 10 +
T
(26)
N
Step3: Training of W So far, the gradient descent step as to W3 and W2 gives them the ability to behave as an associative memory:
wu(k) ((k)) ( = k],
(27)
22

---

Published as a conference paper at COLM 2025
(#)wF(z)) (WewF a'1(z; =2
(28)
From Lemma 3, we admit that the model predicts the next token almost randomly, i.e., p(v x) = 1/V for any vocabulary when Wg = Odxd- This enables us to employ Lemma 5
and W after one gradient descent step gives the following:
1 $2WE (ZI)
E
(k)
Σ
( (Iwr(Z))) (2 ly = k
V
T
i
E
wu (k) (2)0F(z))
1
E
d (zs))
(
V
T
E
(wu(k) T 12
1
Σ
1
Σ
(
(P]we(z)))
( T y= = k
T
t
7 +
V
Σ
E
(mu(k) (2211T)
T
Σ 7
1
Σ
(IwE(Z)))
with qs (wy (zs) r==1)) - ( + --) and (0((zz)) 1 E (z(z)) Now, fix an
arbitrary vocabulary E V. Given that
(i)(k) T (2WE(Z))= = a1 = k}
and that
( ( = a' a'1(zs = h
which equals a' when s = tq or s = T, we manipulate the expression:
wk
7
Σ
E
T
a1 22 = k} a'
; , (1(tg I t)gt,t + 1(f = ((gr)) y =k
V
7 "
Y
IE
T
Σ
a a1 = - (1st S f) + 1[t = (onr)waass)
7
E
1
Σ
T
1<
+
1
Σ
T
a
(1l1g S 1)9tt + 1(1 = (() T y = k
T
1=11
-=1
+
7
E
1
Σ
T
1(2
k)
1
k]
1 _ x
Σ
(1(tg S f991tt + 1(t - ((t=))99))
T
P-1
1=1
With the near-orthogonality of the token embeddings, we write
naa' V
T
Σ
(Aex -ax- - + Dfx)
VT k=11-1
yaa' V VT
Σ
(AT - BTA -
+ DTx)).
(29)
where, we define each term as follows:
A1 = E 1(z = k)1(4 <
yy= k).
1(2 = k] < f)9 = 8}}
B1 = E
C22 = E
Σ
1(4 = k)1((t S f9 = t} ly = k
T
-=1
23

---

Published as a conference paper at COLM 2025
T
T Σ 1 s 02,, = 8)
Dix E
ATA = E [1/z] = k) (} yy
BTT = E
[1 = ( m}],
1
T
Tr
E
Σ
1(2 = k) = 0} y k
=
T
1==1
DIT
E
1
Σ
T
1|z
k)q,,]122= = o)
T P-1
We will further split these four terms based on relative positional encoding and token embedding that are contained in Itg.. For example, we decompose ALe into A&x and A by defining
A E 1(2 14g 1 (14-1- 11 no) 1(21 = 0) y k).
Ai E 1(z+ k)1(t, < (wE(Zz)) - @E(ZIH)) 1/2z =-} y = =k]
(i) when v= k We begin with ABE
A
E -|==k)114 1 (r,,-1 71 10) 1(z) = t} 1 yy-kk k]
Σ P(t = s | y K) [1/z] = k} y 1 = s] (ra-1 h1+
= P((4 = -1 (-)) - -11-:00== + 1s = T} P(4 = (ro - -P1--0)
1(1=1 T}
Σ
P(tg- s)((s-i- - - -00
V
sclt 2)1(t}
= P(4 = t - 1) - P1-1:0) + 11t 1[t= T[t=T] = (ro-)ror--0a) (1-00) + + 1
Here, we used Lemma 6 and 7 for the calculation of the expectation IE |[1 = k} |.]. Similarly, we will examine f. C and D as well.
B = L 1(z 1]1 S f (re-- P1-10) 1(2+ m}}
Σ
P(ty = 5) [E =k] fq = -s((r--- - P1-10)
Σ
P(tg = 5)
(rs - P1-100
V
0
For the term CR ig we make use of Lemma 8 and 9,
CR
=
=E
Σ
1zz==
S/} (Ita - -i00 ) 1(zz = o} y = k
T P=11
1
P(g = -)(rs-t----o)) Po { 1lzr k} 1/2 = t} 1y = =s
[
Σ
=
T
24

---

Published as a conference paper at COLM 2025
P(1g =1 =1-
(r-1-P1- p1-)0) { (1+11) 1s1 (1+1(1=T) T)) + (T- 3) 1st + TH
T
V
2P(1y- =1)
2 Plg =s) 0
(
11
+
(ro - -;1-100 +
T
T
1
- ( -100 + 21g = 1) (ro 1.00-Fop - +0 +0 }
(1 1 T))
=
T
As for D the condition y = k does not exist, and z follows uniform distribution. Therefore, we obtain
D 1x
E
1
Σ
T
1/z = k] S f
( -1 1izz ta
T
1'=1
P(t s)
I Σ
Σ
(Ts t -- ;1::00BB
1(2 k]11zz = f} y - - =s
==1
1
0
Next, we will examine Aix Bi Ci2 and Di in this order. Using lemma 10 and 11, we obtain.
A
E 1(2z = k] (x - = 0) o0 kk]
E 1(|z = k}1(tt s f y k] -E [1(z] k) S f/ y y- k]
=E
= P(tt =t- =t-1)1(t- (k) + = -)](t = -T -T-1))
p(tp=tt 1)
1(t -(t==-10 (1) W (k) +0 (
(1 + 1 = T))
(k))
t
f
We also get
B = E 1(z 11g S (x1, - x1.)
= m1}
E [1z = k) s [)x1]]
- E [1 = k) S f]
= P(z = E1t, s f)2;,) 1zz= k] - P(zz = K) [1 S f) |z1 = k]
= 1
Now we move Ct From lemma 12 and 13, it follows that
1
T
Ci
E
Σ 1(zp =k]
1lt S (x(, - 11.) = t} y
1
T
=E
Σ
1 1zz =k]
1(4 S [1,x,11(,= ly=k k
T
P=1
1
1(z = k) ) 114 S 11 = k) ly=
E
1-1
{
(2((gg = 1 - 1)
1(1 [f=11 T) + 2P(1g 1)
= (k)
= TT- 13
T
P(1 =1 1) 4. 1(t = T] + +1
P(t, = 2- 1[ = T 1)
T
T
1
+ + 1. (
25

---

Published as a conference paper at COLM 2025
For Di we have
1
T
Di E
Σ
1[z =
=k] k)
1f4y S 51) f11)1(2t 2)
T
P(zz = F
Σ 1(z k) 114 -14551(2)) -111) 21
T
1
because of the term 1(z = f} under no condition. So far, we have calculated Aik, Biks Cis: and Dr for u = k. Now, we turn our attention to # k. (ii) when of
In this case, it is guaranteed that 1(z, = k} + 1f = 4,, and thus, we have A Bi 0. We only need to examine CLK and Di,- We use the fact that the token Z₁ has be sampled uniformly from potentially reduced vocabulary space with some words excluded depending on the position of I, we obtain
1 T
Ct = E
Σ
1 = k] S 132 = t} 1y ]
T
P(z = v y kB
-
1(4 = k)1(4 S 1)9tut y = k, = Σ
=0 1
and
T
DiA = E
Σ 1(z = k}
1lt S fx), - 8z)1((= = o}
P=11
P((
y K)IB
7
Σ
T
1(z =k]
1(4, 1 = v}
-=1
1 4
o
So far, have seen the content of AG,, Br Ct,, and Dig. Now, we will shift our focus to ATxx B CT and DT We recall that
9s (wF(2s) +Ts ((0F(11)) + F1-10).
By substituting T and T, we have
qT,7 (we (ZT) + ro) (we (ZETT + F1-T:0).
Also, 1(2r = d) implies that the trigger token is v. we use Lemma 14 and obtain the following:
AT1x
E [1 k] =4 y = k]
P(l. = - 1) (ro -1-1:0)= + P(t. = T - -1 2 + ++)
if k
0
otherwise.
Similarly to the argument of Bi and Dr,,, the trigger token is chosen uniformly, and thus, we have P(zr = v) = J. This gives
B E [11
k2r = m}]
26

---

Published as a conference paper at COLM 2025
10
(+)
f otherwise,
0
and
L11 = E
2 = (2g - ta
T
if k, otherwise, Finally, according to Lemma 15, it is established that
0
0
CT = E
E 1z = -kgrr11(2z = = y =k k
T
2P(4, (T(+)1-1)
(ro - 71-T0) + 2P(+,=T-11
T I WE (k) + - (t)
of
otherwise.
To sum up, we get the following associative memory behaviors by examining Eq. 29. We point out that F1 +o also contains 1-1 and ro-
(W) qaa'
T
P(1 =t 1)
1111111=7 (1-1)
Σ
T
t
T
naal
P(t =t)
7
1(t= T}
f2
naa'
v
P(q
T T-1)
1)
(1 1 7
VT2
and
ro (oW) = naa' T
P(ty =t)
(1[1= T} - 7) (1 (1-
Σ
T
naa'
T
P(g
= =t (1
1 1 1+1
12
naa'
P(tg = -T-1) 1 2
2)
VT2
T
Also, for any j E V, we have
WF wF (i) =0 1 (
The proof is finished.
F Other lemmas and proofs
Lemma (Lemma Bietti et al. (2024)). Given data distribution p over pairs (x.y) E Rd x [ and a fixed matrix Wu E Rdxd, mininizing the loss L(W) -E(xy)-p [e (y. Wu gives the following gradient:
VwL(W)- v w=1
=) No)T
(o)() v| and po = Ex p(y=o) x), and Pw x) is the probubility of generating v.
where pk =E(x
27

---

Published as a conference paper at COLM 2025
Lemma 3. Consider two-layer attention only transformer that meets the initialization condition
in Sec. E.2, and the imput sequence of length T ends with the trigger token 2T = 9 at its second occurrence. Even after the gradient descent step of W3 and WR, the prediction p(v 21:T) = 1/V for any E V.
Proof. From the way the transformer is initialized, we have W = -axxd Therefore, the output of the first attention block for any t is:
x wr(z)) + 1 Σ plwr(zs)
In the second layer, the key matrix is trained such that (DIWE(Z))) T (ZTE) = a'1(z zT}. Thus, we have
(WE(ZI) Diwe (21))
(WE(ZZ) + E ()we(z)) 2
((aft)) (wxx)" =
Di (WE(ZT-1) Σ7-
--1
(r((T)) + E
T
x
Σ KE
(K) (k) (2O)
1x
(wE(Z1) + 1 (21) T
(WE(Z2) 2 mpl T
ΦIWE(9)
WEE( + Ex ))an((z)
+ - d dywe (2:) T
(wE(z))
' TV
1/ (to - 1)
1/t
1/ L/(T-1) 1) 1/T
When V is large enough, the attention is spread across the whole sequence evenly, and obtain
Σ W3W3x (1) +x++ (1)
T
~
T
(30)
Σ wBWB (WE(Z)) + 1 Σ ()w((2)) ) + (E((T) + *T Σ Φ WE (zz)
1 T
T wr((z))
28

---

Published as a conference paper at COLM 2025
where focus on the terms represented by wir(o) for vocabulary by awy. Taking the expectation over 21,7 while keeping the trigger token 2T = 2 = 4 fixed yields
-|x
2
wu Σ
EE
T - 2
Σ wy(w) + 2w)
T
M<P<V D#
The coefficients of (u) are the same for any E V and this concludes the proof.
Lemma Consider two-layer attention only transformer that meets the initialization condition in Sec. E.2, and the imput sequence of length T ends with the trigger token Zr = 9 at its second occurrence. After the gradient descent step of W2, the prediction p(v Z1.T) 1/V for any e V.
Since W = xdd it is guaranteed that we have Eq. 30 as an equation. Thus, the result follows directly from Lemma
Proof.
Lemma 5. Suppose the loss function is given by
I(W) Exxy)
((X) WydD T W2xT),
(W)) with (W) = E 11 + R1 R1-10) Wx))s. )s and o (un:r): = + (1 + it T - Us) is the linearization of the softmax function around 0. Then the
where Z(W)
gradient at W- 0 has the following form.
VwL(W)= Σ
E
42 7 wu(
Σ
(1x))
W3X
ly m==
Σ
N
IE
T
wy
wy(i)
42211
1
Σ
(D)xs))
W2xx79x1x7
T
s=1
Σ
IE
(i)
12
1
Σ
T
(1xxs)
T W2XT9,,11 y =i
T
1=1
s=1
+
Σ
IE
(i)
T T2281:1 7
1
Σ
T
11x))
W2X
where 9,,t (I, + +r ry-)) ($1. P1-10) and 8111
$ Ex= Xs.
Proof. The result obtained by replacing PT with (x14 + R1-10) T and pt with x; in 2, from the proof of Lemma from Bietti et al. (2024).
It is worth noting that x; x (0) (z(z;) and that if we directly follows the transformer architecture the vector in the linearized version of softmax ē should be Z'(W) W2z+ where Z'(W) (2 (W) (W)) with z/(W) x; Es= o (xxxx(((xx1 + R to) Wx))s. Since W is learned so that it only cares the ordered pair of (0wpp(z and wF(2)), the loss function is simplified and ignores the other terms. This also applies to g(X), where
(x) Wy (W) W2=7)
the authentic form the output. Since we know W3 (and therefore (2) is already trained, can simplify the expression.
29

---

Published as a conference paper at COLM 2025
E1 Supporting lemmas
This section provides useful lemmas to show the learning of matrices, especially for W
Lemma 6. Given a sequence 21:7 that satisfies Sec. E.2. For any E [1, T), and any vocabulary k E V, we have
E[1(z| = k} yy = k, lg = - = 1.
Proof. Since the trigger token is 2t, 21-1, the next token 2₁ is the output token. The condition y = gives us 2, = k. Therefore, 1(2, = k} = and the equation holds.
Lemma 7. Given a sequence Z1:T that satisfies Sec. E.2. Fix some [1, C N. For any s sE [1, - 2 2] U {f}, and vocabulary k E V, we have
ifs 1 and 1 = T 1, S S 1 2 and + T, otherwise,
E[1 = k} ly k, =
(31)
0
Proof. First, consider the case = t. Under the condition = we have fy = t, and thus, 2, is the trigger token. The only possible case allowing the same token for trigger and output is when lg T - 1. This yields 27 1 2τ = k because zy not only works as trigger, but also output. Therefore, when s = f, we have
(1 i
E [
k}
k, tg s]
(32)
otherwise.
Next, we focus s < f 1 2. Suppose we given t T. The token z; is then the second trigger token, and 1[2: = k} = implies that 2t, = k. However, since the first trigger position is l < - 2 and the output token is also y = k, zr(= k) cannot be the second trigger occurrence. Hence, it suffices to consider + T. We now have s 1-22 < Assuming t, = k will also lead to contradiction because the trigger token k appears more than twice at positions s, + 1 and T. With this in mind, we obtain
[2 = k} y = k, lg = = P(z, = y = k, lq = s)
Σ
P(z+ = k, = 9 1))
Σ
P(z, = k ,,,,,,24 = 9) P(1 = 1.)
=
(EVi(K)
1 1
Σ
V-1V-1
(EV((k)
(33)
V 1
Here we use P(Z1, =q y =) = 1V - 1) because the trigger token is uniformly sampled from V (k). Also, we have P(z,
y = k = S, q) 1/(V (V-1) 1) because 2, for f * T' is neither the trigger token nor the output token, and thus, sampled from V {g). These give us the following when s < f- 2:
i otherwise.
E k} ly - -
(34)
Combining Eq. 32 and 34, we obtain the result.
Lemma 8. Given sequence 21:T that satisfies Sec. E.2. For any vocabulary k E V, we have
2
=t==T,
Σ 1zz = k] 1(z: = k) 1 = k, 1
7 (1 (1 + [2)
otherwise.
(35)
30

---

Published as a conference paper at COLM 2025
Proof. By the linearity of expectation, can decompose the expectation follows:
T Σ 1(z = k} 1(2 k}
E
E|1(21, = k |-]
+ + = = k} |-]
+ 1 + E1 k)1lzz k} [.]
+E
1(z k) 1(z1 = k} 1. (*
Σ
where we abbreviate the conditions of the expectations to enhance readability. Note that we always have 1(z; k} 1 when lg = 1 and y k. In the case of T, the trigger token appears at positions T - and T given ig f - 1. With the condition y = k, it follows that
1/z k}1lzz = k} = 1,
1 1 1 k) =
In the above, we point out that we have by + = T this situation. Additionally, we get
E
1/2 = k} 1(2, = k} = 0, H4g
Σ
where the final equation follows from the fact that Zy not a trigger token k.
For the case where t * T, the trigger token is not k because otherwise there will be more than two triggers in the input sequence positions 1q. lg + and T. Thus, we have
[[1 1)+ k} = 0,
E ( +1, = k) = k] = [[1 1) = k} |] = 0.
Following the same argument in Eq. 33, we also have
Σ
1(z k}
1(2 = [1)==}]1
=E
Σ
1(2p = k}
('
T-3 3
V
We prove the statement by combining these results.
Lemma 9. Given a sequence Zı.T that satisfies Sec. E.2. Fix some t [1, C N. For any s E [1, - 2] (t), and vocabulary k E V,
2
ifs t and f T - 1, ifs and # T, otherwise.
E
Σ 1 12z = k} 1(z = k7 ly = k,
=
V-1 + (V 1)2
T-4
(36)
Proof. First, we look at the scenario s Since tq s t, the input sequence Z will contain more than two trigger tokens unless f == T - 1, which does not satisfy the
condition of input sequence, we examine the case of f 1(2r-1 k} = 1(zT = k} = 1, and therefore, we can use the instance = T from Lemma 8.
1. This condition ensures that
E
Σ 1ze = =k} 1(2: = k} =
31

---

Published as a conference paper at COLM 2025
2. It is impossible to have f = T because this violates the input sequence condition, where the trigger token only appears twice. Now, by the linearity of expectation, can decompose the expectation as follows:
Next, require: f f- -
Σ 1zz =] 1(2: = k} 1-
- = k} [-]
E k}1lzz k} 1-]
E[1 = 1 k] |.]
E[1(ZT k)1[zz = k) |-
+E
Σ
1(z = k]
1(2 k]
Pttl LAT
Since 4 =s t - and f + Т, the token 2, is neither trigger nor output. This yields the following:
E|1 = k] = k) | -1]=0, 0, E [[1(2r = () = =
Given y equals k, we have 1 = k} 1. Therefore, from Lemma 7, we get E Z[1 {)+ k} 1
| k} [-]
1
V
=[1(2=k]11zz = k] | = E (2 = k} 1-]
1
Finally, for 1 + tgrto + 1, and T, the tokens zy and z, are independent. With Lemma 7, we obtain
T T 4
E
Σ
1[z k] = 1(2) [z===]]] k]
(V- 1)2
Integrating the equations above allows us to arrive at the conclusion.
L
Lemma 10. Given a sequence 21.T that satisfies Sec. E.2. For any vocabulary E V, and e
[1, C N, have
E 1(z S f) ly= k) P(tg () TTwr(k)
P = 1)1(t 1) + 11
v
Proof. We split 1(tg S t) into 1ft 1 1) + 1119 = f} + 1f4g S 2}. Then, the expectation can be decomposed as:
I [ k) S x) y=
E [[ 11 t 1)x,, ly k E 1(z+ k)1 = f)xt, y k
E 1/z, k) (t 23 x1, | y k
Note that when lg = 1, z₁ is the output token. In addition, if we have T, 21 works as the second output token as well, which means that z, and X1, we (k). If we find + T, the trigger token is not Thus, we derive
E 1(2 = k)1(4g = 1 1)xt, ly k = P(l, = 1) 1(zz = k)x,, 1y = k, tq = 1
32

---

Published as a conference paper at COLM 2025
P(t. - 1)1ft (k)
wE(i V V-1
+ P(g = -
E jEVi(R)
Next, we discuss fg f. Since ig is the first trigger position, always observe that by S T - 1. If two trigger tokens are adjacent to each other, i.e., by = = - 1, we can state that 2T-1 = ΣT= = k because of the equation y = k. Otherwise, the trigger token must belong to
V (k), and 1f2 = k} = O. This gives
E
1(z = k fe, ly= k| P(t. = = -
For 1n < f 2, 2, we again analyze the cases of = T and f + T. Provided that = T, 1/z; = k}- = implies that the trigger token is decided to be k, while the output token is also k. This leads to more than two occurrences of trigger token, which violates the input sequence condition. Now, we assume t < T - 1. This time, the trigger token Zi, is not k, and is not a trigger token, which yields
1(z = (tt S - 23xt, y = k)
=P(lg S 1- - 22 y = k, lg S 1 - (2) = k] 2t,, y = 4,4 S 2
Wr (9) E E[1(2+ = k} Zt, = 9,y = k, S - 2]
=((tg - 2) Σ
V
qE (EV)(E)
WE 1)2
= P(4 S 2)
Σ
(V
qE (R)
This concludes the proof.
Lemma 11. Given sequence z1.r that satisfies Sec. E.2. For any vocabulary k E V, and f e
[1,7] C N, we have
P(fg- = =t-1) 1)
E [1 = k]1(tt S &1 y=
1 (1+1(t= T} (())))
+ 1f = T 11 = wo (k) +0
(+)
(37)
t
Proof. As in the proof of Lemma 10, we analyze each scenario depending on the value of fy. Firstly, let us consider lg = f 1. This leads to the fact that 2, is the output token. Given
y k, we have = k.
E [1 = k4 = - 13 y =y=k k] = P(4 = 1 - 1)I [
y = k, 4 = 11- 1]
In the event that T is valid, Xt, XT wr (k) holds and the other T 2 tokens not k.
Thus, it gives
1 2k + -T Σ
wc(i) V - 1
E E = k] =
jEVi(()
When t + T holds, we can say that X1, # (k()), while Xt, 11 = (k)()) because of y = k. This gives us
WF (i) IE E[xt, I y = k, = =- Σ
(EV)(() V-1
E [[xx,11 k, t =t 1] w (k).
Also, we can derive the following equations for i [1, \\  (tq kq + 13:
E y = k, fq = t
Σ P(z: = ww (j) CCV
33

---

Published as a conference paper at COLM 2025
V
V
E
E
P(z = 1,2t, =9 (-)w((i)
E
Σ
P(zi 1, =q1 1))
jeVi(a}
Σ
P(zi = j.l 2t4 = 9, P((z, -))
(EV((K) icV((a)
wE(i)
E Σ (E) jeV
(V- 1)2
In the above, use the fact that P(x; j, q) 0f 9 or q k holds. To summarize, we have
2we
7 Ejerit V mF(g)
we(i)
if
T,
(E((K) + ENEVi(R)
E 81a y = ky =t
V-1
(38)
1
WEli) E Eger((K) (V
otherwise.
17
iEVilat
fwo(() +0 tav (k) + +0
1
ift=
otherwise.
Next, assuae 1. In this case, 1 = T implies lg T which a contradiction. Furthermore, ly = I < T 1 implies by < 19+1 < T. Under this condition, we cannot have 2, = k, because produces three appearances of the trigger token. Hence, it suffices consider If ly = T 1 is true, we can state that X; k, leading to the formulation below:
B[1(z+ k} Rn. y k,
[[ y fy = f]
wE(K) -T- Σ wF
V
1S/SV
T T-1 W (k) +0
(k(K)
where the second equation follows from the same argument in the instance T in Eq. 38, but this time X7 is not contained in X1:t-
Lastly, we focus on lg = s S I We argue that the trigger token 2,, = zr not k, due to the condition of the number of trigger token in the input sequence. This means that t * T and 2, does not serve as a trigger token. Now, we split E1t into } (xx, + Xt, 1 + x + E. xp) and obtain
E(1(z = x)xt, 1y = =
E (2)
P(Z: = k, w)
we (((9)
Σ
(V 1)22
(k)
E E k/xx,11 y k, k, s]
P(zı = k -)w (k)
Σ P(z: k, = 9 1)) (k)
(EV
WE (k) (V 112
Σ
(
34

---

Published as a conference paper at COLM 2025
w0 (k)
V
[[1 = k3xx = k = = P(zz = = = = (k)
we (k)
V
1'
E[1/2 = x) 1y= k, = s]
E
P(zi = k, = 1, - -)wx
VEVI
(E)
E
P(Zta = =91
(E)
Σ jEV\(n)
P(zr = 22 = 21g 9, 2 (i)
wE(1)
E V-1
Σ
(V-1)2
{x)
jEVi(g)
In short, we have
1
[[1 = 8 y =|
Summing up all the related terms leads to Eq. 37, which finishes the proof.
Lemma 12. Given a sequence 21.r that satisfies Sec. E.2. For any vocabulary k E v, and f e
[1, C N, we have
1 T
2 (k)
IE
Σ
1zz = k)1(4 S 1)2, = k} 1y y= k
1 = TTP((t = =f-1) 1)
T
(k
T 1ft=T-13P(4, =))
T
1
1
0
(39)
v
Proof. We consider the following cases: Case 1: fy =t 1.
If T holds, this gives Ztg = ZT y k and 1 = k] 1(2; = kk} 1. Thus, we have
E [ x),,, 1fzz = k}} y k, lg = 1 1 IE X1, I y
k, lg = t t- 1]
WF (k),
E [11z k)x,, 1fz1 = k} 1 y k, -t-1 IE Xip
I y
tg - t- 1
W (k).
Since zu for I' # by + 1 is not trigger token, it follows that 1[2p = k) = 0, and we have
IE
Σ
1(2 = ,| = k} 1- = 0.
11 1,
Otherwise, when 1 + T, 21, = Σς is not k, yielding
IB
|
= = 1 = k} | 0.
E 1(27 = k 1fz1 = k}} 1- O.
Also, Zt, is sampled from V (k), we get
E 1(2 = k) xx, 1,, = k} y = k,1 =E Xt, = k 111
35

---

Published as a conference paper at COLM 2025
wr
Σ
(E V-1
For the rest of t', we note that 21, # k and obtain
E
Σ
1(2p = x) 1 (21 [ k}} 1-]
Σ
P(ze = 2, = 12)0((()))
# LT
vEVi(K)
WE(9)
= (T-3)
Σ
(40)
(V - (V-
geVi
In summary, we find that
E
Σ 12p (1g 1- 13xx, 1(24 = =k -
ly k
P(tg =t t- 1-1111(--- T))==(()))) + P(ty= =t - 1)1ft + TJO (
(41)
Case 2: fg = f. For this scenario, 1(z; = k} = holds only when tug = = T - 1. Thus, we can state that
E[1 k) xx, 1 (z+ = k] y = k, = = EE [[ y = k, lq = t
wg(k),
[[1 ,, = k} y = k 1 EE Ex(t, y = k, kg = t]
(g(k).
For the other H, we have 1/z k} 0 because these tokens are not a trigger token. This means that
E
Σ
1[zz k3 1(= k} ly
k, = 0.
11,,T
In short, we have
1(1 T - 13 1)P(14 =1) 2()) (42)
E
T Σ 1(z = k)1ltt = f)x;, 1 = ly=k
T
Case 3: ly S f 2. We first let T. 1f = k] = holds, then the trigger token becomes k, but we find that 2tg = 21,+1 zT k, meaning the trigger token appears more than two times. Therefore, also fails to satisfy the required conditions
we only consider t + т. Given zt, (=27) for the sequence, hence follows that
Hi1 x),,, 1f = k} y k, lq s) 0,
E (z x 1(z - k) lyykk,, k, s] s] 0,
where 1 1.555 5- Similarly to Eq. 40, we can derive that
E E|1 x)x,, 1(z = k} y s]][__ 1(2; =k} 1y- k, = s]
= Σ WE(q)
(V - 1)22
&EV1(2)
E[1 x)x,, = =]
1 =|
k) | y k, =s]
36

---

Published as a conference paper at COLM 2025
1e w6(9) (V - 1)2
Σ
(E
In addition, we have
E
E
1/2 = 1 = k} ly = k, lg =
1(z
(+
f*t
E
P(Zz = k, 2,, = =q -
=
(K)
Σ
P(zy = k, = k, Zt, 9, ,.,) 9 ly = k, 1 = =9)
PEViRK
WE(9)
= 4T-
(V - (V-1)3*
QEVI(K)
These equations give us that for 1 < --2 2,
E
1( = xx, 1 = =k 1y = = O
(43)
12
Gathering all the equations of 1 S 14 S leads to Eq. 39 and have proven the statement.
Lemma 13. Given a sequence zn.r that satisfies Sec. E.2. For any vocabulary E V, and e
[1, C N, we have
E
Σ 122 = k] $ f] = k} yy=k
1-
4 1st = T} + 1(t +
2 2.1(t= = 11) =
P(ty 1)we(k) +
We (k)
Tt
Tt
[T)
+0
(44)
VT
Proof. Considering the linearity of expectation, we focus on each term contained in 8 We consider the following cases:
Case 1: ly = =1 1. We already have the result as to Xt, in Eq. 41:
1
E
Σ
1(2z k)1ltt =
=z kk] y =
T
P(t. = t - 1)
P(t. = t 1)
T 1ft # OJL (
T 1lt T3 T)2W (k) +
We can also use Lemma 8 for x; = k, iee,
E
Σ
= =
=k} ly yy= k
T
=E
E
T Σ 1(z = k] y = k, = 1
P(ty = -- (k)
37

---

Published as a conference paper at COLM 2025
- 1)
T 200
ift=
P(4=1-1)
(1+ --1) We (k)()
otherwise.
Now we deal with x; for 1 < <is 1 2. Note that 2; # k for these under the conditions t T l'and y
E
yy= k Σ = =)
[1(2) -] [-]
+ E 1(21, = x] k]
+ E = x2, k [-]
B x,, 1-1
Ejeri(k) V-I
wrlil + Ejcri [k) V-T
wE(
ift= T,
EjEviHAt V-1 + V + (T - 3) E (E #g(i)
wElil mg(k)
otherwise
(V-1)
(c)(())
(1)
Case 2: fg = f.
For this scenario, have seen that iq 1 is the only possible way for the expectation to take non-zero value. For Xt,, we can observe Eq. 42 and obtain
1(t T 1) = f)
E
T 1zz kllst = x 1
2 (k).
T
The other terms related to x; for 1 <is can be calculated
T
E
Σ
1[z = k]111(2 k) y 1,14
F [1 = k)x1((= k3 |.]
1
+ E [1{z] = k)11(z = k} [-]
T
1
+ T
Σ
= = l-
1
wr(i)
wr(i)
Σ
+ Σ
V V-1 0
T
V-1
(EV
jEV((K)
1
1
o
T
We note that 1(2) = k} = and 1(z = k} = 0 for 1' < t4 = because of the input sequence condition.
Case 3: fg = < t t- 2.
Again, we have the result for Xt, in Eq. 43.
1 T
1zz = k)1ft, stx, 1fz1 = k} ly= k 0 1
(12)
E
T
v2
38

---

Published as a conference paper at COLM 2025
We have x1
under y k, and Lemma tells that
E
= = k)1ft = s s)xt, = k} y = k
T
P(tg 5) T
=
E
Σ
1(z = k}1fzt = k} | y = k, =
WE (k)
p(t+=))
\2² + マ wF(K)
otherwise.
Next, we consider xt. Note that T forces Nr, to be wr(k) for the expectation to take non-zero value, which contradicts the input sequence condition. Hence we assume + To that 2t,
+ k.
1(z = k)x,1(2 = =k} ly = k,t =s
1/2,, = =) |-]
T
T
[| |-]
=
E
= 13 1-1
1
200
(F((k
E
(V-
T
V-1
/'
1
1
O
T
Lastly, regarding x; for i A t, tq by + 1, we also consider
* T and follows that
E
T 1(2p -,, k} = k =
T E [1z [\,,= k)
1
+ E [1T =k)
+E [2,
+ {) {11]
1
0+0+2 Σ 1 Σ
wF(i) Σ Σ
1
Σ wr(i)
T
V-1
(V 1)2
V-1 (q)
(V - 1)3
(EV)
(EV)
1
= -
Thus, we obtain Eq. 44.
Lemma 14. Given a sequence z1r that satisfies Sec. E.2. Then, we have
P(t T 1) (ro F1-T:0)
+ P(1 - 1- +0 -+)
E [17 1 {ZT = = y k]
if k, otherwise,
0
39

---

Published as a conference paper at COLM 2025
where 9T,7 (wFr(zT) ro) (10F(Z1:T) + --TT:)
Proof. The inequality & produces 1/2T = kzr = 0} = 0, and thus, we focus on v = k.
Since 2r is the second trigger token, have 214 ZT k. This achievable only when lg 1.
[1(zT = (2) ly =
[ 9T y
P(zz = y [TT,T =
(45)
T
Σ P(Zz = k = y = E) lqT, 1/ k, = k]
i-1
= P(4g = T 1 y = Pzr = y = k =T [9TT 1 = k, = k]
Σ (T-2) 2)
P(4g = T - 1) (k(k) + - 81 T:0 T
2wr(k) +
- wi(()
15JSV 1
T
= P((t = T 1)
WE (k) 71 T:0
T
(v)
Lemma 15. Given a sequence Zı: that satisfies Sec. E.2. Then, we have
2P(t, - TT 1) (ro - 51 r:)
2P T- 2 (k)
T
1lz = k) =) ly=k)
T
+
if k, otherwise.
1o
where 9T,7 = (wF((T) + ro) - (10 (ZIr)) + 7-
Proof. We first consider the case k. If the second trigger token zT output token is also k, it is impossible to have tg < T because it produces more than two trigger tokens in the input sequence. This means that we have, for 1' E [T 2], that
and the
E 11/2 1)1 k} y = k| 0.
We can use Lemma 14 for I' =
E Zzt = k/gT,11(27 = t} = k] = P(tg = T -1) (1o
T-2
wF (k) F1-T:0)
T
(
+
Finally, for =
T T- --1,, 1, we note that the two conditions y = k and 27 = immly ZT = k.
E 1 = =k]4r.1122 = ly y=
P(ZT k 1y [1
k)9r,7 =, k)
=P(zr k k =B= [qr, = k]
-E(1(2r k)qr,r1 = -) 1yy k]
PP(t = T 1 ( T-2
(kr(k) - F1 T:0 +0 +0 (0)) (
40

---

Published as a conference paper at COLM 2025
where the second-to-last expression follows from Eq. 45. Now, we look at the case v * k. Different from v = k, we have
E [zr-1 = - () = y = k] = 0,
[1 = :onn=k)gr..1t k) = - = 0.
Furthermore, the trigger token is uniformly distributed, we have
[[1 0) y kj = P(zr = y)
T
Σ
P(Zz 0 = = k)
T-2
Σ
P(zt v ly 1= i) =
1 y
1=1
P(tg
T-
V-1
1
V
This leads to the following expression for f E [T [
EE [1 = =kgr.11122 = n} 1 y = = [1 [2r = v) [1 = k]
y k]
= P(Zt = V = K) [|z = k}gT,7 ly = k = v)
= =0 1 (
This concludes the proof.
E2 global v.s. in-context
Proposition 4 (Restated). Suppose a twoo-ayyr transformer an associative memory transformer. Given a length-TT sequence 21.T where the last token is 2T q, let the t1-th token be Zt, 01 following 2t, = 9, and let the t2-th token be 2t2 = 02 following Zt--1 = 9, where V, # 82 and (3 5) < t2 without loss of generality. Also assume that there exist only one on and v2 in the context Z1:7. Then, the logit Sor and for can be expressed as follows:
e
Zo. = o
+ +log FIp (p 1))
1 h
1+e
Eoz =0
+ +lo (22 1))
12-e-1
where (f)) represents the softmax function.
Proof. We consider the output of the first layer of attention for z,- Since we use the associative
memory transformer, we substitute W] I and W/ "ieo we (k)r into Eq. 1, we have
(1) (0/0F(Z1a)e (wwr(2)))) + R10)
Σ (r-ww((k) Iwog(z)) ke@
+ wF(2))
(WE(ZI) + r1--1) T (we(22) + +(2-1)
= w
x (-wwE(zi) wE(Z))
+ #wE(Zi)
(We(Zr-1) + r--) T (WE(Zi) +ro)
41

---

Published as a conference paper at COLM 2025
$ 10 (2
+ wr(zi)
(46)
1
e( w
1 (21-1) -
(zi)
-wy(z1) +
Σ
(47)
+
11-1
where e is Euler's 's number. In the transformation above, Eq. 46 follows from the near- orthogonality among the relative positional encodings and the embedding vectors Next, we calculate the output of the second layer attention block. We first consider softmax values and then attention head WO W3. Similarly to the first attention layer, we use Eq. 47 and ideal matrices W3 and W = Exeq (e(()((1 to calculate softmax. For any I, we can
express )x T (W2) w2xx""
(1)
(WE(ZI) + D WF (21)) (w)((2)) + cl (()))
Σ
e+1
e+1
Σ KEQ
4 w (k)) 1x(")
(WE(Zr-1) + < (0 Σ D1wl(z)
wr 142-2
148-2
(FF(Z))
(( Ite-1
I+e-1
where the summation in the 1-th row is given by Eisigi Since xf7 WF(4)
e ((
T+e-1
111
ΣKKT (pw(() and thanks to the near-orthogonality, we have
-T-1 T+c-1
(xx))) ( (p)
where each entry of p E R is
ift < f1 - 1, if f1 - 1,
1/(11 ---)) 2)
e/(t+ 1)
ift= f1,
1/(i+-11) 1)
if l < 12 1,
pt
2/(12+e-22
(1+e)/(1+ee) (e-1)
if l
2/(i+e-1)
if </<tt
(/
ifl=T.
Back to Eq. 2, now write xf2 with w3 E (1(o)
(2)
Σ (pW3)(p))x + +
x
V I Σ [p(p) Σ (u(o)(W((0((1))
((
Σ Φ1wE(Zi)
x W3 wF(Z+)+ +
+x (1)
t+e- -
t e -
1/t 1
42

---

Published as a conference paper at COLM 2025
E (p()))u(z)) + +w=()) + e( ((0)wF(2r-1)
I
& (1))(())
Σ
T+e
T
15ist
1/T
(2)
Finally, the feed-forward network in the second layer is applied to xT and we obtain the following:
MLP((x (2)) + x
WE wr wr(22)
(r(()) + Σ (l(p)) (p)))ollz)) +c)
W2
ReLU
+xf
ww((v)
(2) + +xT
W2
ReLU
V
Σ log T (1)
Σ
log uu uy (v)
+ xx
(2)
21
log Try (u ()) + xf)
log ( ()) Σ (p))))u()))+++ wF(4) + C,
w(22r)121))) 1+r-1
where C
+ E Eici! 14T
T+e-I
We can compute the logits To for v E V after the unembedding layer. Since the row vectors of Wir are also gear-orthogonal (P1) ) and (w)). Since 71 and D22 appear only once at I = and I f2 respectively in the context Zi:T, the logits are calculated as follows:
each other, sufficient look the coefficients of
e
+ +1(0 U log (vr 1))
h 1
1te
Eoz
+1(02 log (02 1))
12+e-1
S two-layer transformer is stronger associative memory transformer as in Def. 3. Given a length-T sequence 21:7 where the last token is ZT = 9, let f(v) be the number of token paltern "qv" appearing in the sequence for vocabulary V. For large enough I and T2, the logits &o can be expressed as follows:
Proposition 5 (Restated). Suppose
Eo ~ dognz(0 () T3 f(v) -1v = 9}
(48)
EV ( + + 9}
43

---

Published as a conference paper at COLM 2025
Proof. The proof proceeds in the same manner as in Prop. 2. With the associative memory W the attention values for z is expressed
x ((WE(Z))))((W(((Z))
R1-1:))
Er-Iwe(()) (20E(z))) (g((2))
11
KEQ
(wg(Z1) (21) + +11-1) (20p. (22) + +12-+)
T1
r 1-1WE (zt) wE(())
+ wE((t)
(wE(Z 1) + (WE(Zi) + ro)
+ w (zI)
T1
22 Biwr(Z 1) + (e
Now the second layer attention has the form
(xx (1)
(W2) w3x(")
(WF(ZI) + PIWF(Z)) (WE(Z2) PIWE(Z1))
Σ kEQ
(1) wr (k)) Ix
2
(wF(2+-1) (D(FF(2(22))
(WE(Z+) + D1 (2:0E(2)-1))
We note that xx() (q) Df and consider the case T in the above equation. Then, we will have
((+ (WP)) T W w2x1")
(WE(21) + D (21)) (wr + 1 (21))
Σ kcQ
s ( w (K)))
Txff)
T2
(WE(ZT-1) + D (ZT-2) (WE(2T) + WET 1))
(WE(Z1) + W (21) (wr(z2) + w (21))
(WE(a) + wi (q))
T2
Φ)wE(9)
(wr(b) + w)
(WE(2T) D
Σ tze
ieloch
44

---

Published as a conference paper at COLM 2025
(1) where 1 {t { wg(d') + Φ (, 3v' E V} is the index set of x; containing (1) and 4 = (1} z = g and 4 = Ø otherwise. Also, e, denote the unit vector where only the i-th component is 1, and all other components are 0. Note that | U 41| =
f()) 1(21 = 9).
Going back to the output of the whole attention block including W3 and WG, we have from Eq. that
(2)
WO
Worl
Σ
x2)
(1) + +x7
iEt,U
Σ
WGWBx!" + +x
1 Ughh h
iet,uh
1
(3wn(o)(33xx(()) {f(D) (WB(WE(O)+ + (
14 1=1
+ +1/21 9 (w)(w((9) (1 + x
f(v) + 1(0 9) - 9 (zwe(() + wE(4) + @T
= Σ
E f(o)) + 1/21 =
Σ (b) w + # + Φ w (3T-11))
Lastly, through the feed-forward network, we obtain the outcome.
XT
= -L + (2)
Σ log (b( 11) (1)
Σ log ( oy (w)
x (2)
V
log nz (u| (w) + xx)
Σ log 1T ()) +
Σ -==
((()))ul( E()))a(()) + (F(() + DIwEE(ZT-) 1)
f(v) + 110 9 9
1lv log log (v 4)) t
(wu(v)
((*))())) + 1(21 = 4}
+ ( q (E
The proof concludes by examining the coefficients of (v)
E3 Induction head without positional encoding
In this section, we provide the construction of a three-layer transformer with no positional encoding that achieves the induction head mechanism.
We first introduce the following theorem (Kazemneja et al., 2024) stating that transformer can implement APE using one transformer block.
45

---

Published as a conference paper at COLM 2025
Theorem 4 (Theorem 1, (Kazemnejad et al., 2024)). Let 21:T+1 be an input sequence of length ++D 1, where 21 = (bos). Then, there exist transformer block consisting of an attention block and feed-froward network that recovers absolute positions [1,2,... T 1] and writes in the hidden state H()
In the proof, they prepare extra 3 dimensions for word embeddings to effectively reconstruct the positional information:
1
1
1
1
0
0
0
0
0
0
WE
1,2
1,3
0d
4,2
where (bos) is represented by the first column without loss of generality.
Theorem 4 states that the hidden state (()) has the following form after the appropriate transformer block:
1
0
T+1
H(1)
Ø1,1
A12
41,3
01,T+1
Ad,1
ad2
4433
dd
Now we show that three-layer transformer with no positional encoding can implement induction head mechanism.
Proposition 6 (Restated). There exists construction of a three-layer transformer without positional encoding that achieves induction head mechanism.
Proof. The first layer of the transformer block is given by the construction in Theorem 4, which gives us the hidden state:
0
0
0
T T+1
H()
Ø1,1
A12
4,3
R1,T+1
Ad,1
Ad2 443
4d,++1/
where each entry denoted by as,t is Gaussian random value to guarantee near- orthogonality.
Then, for the second transformer block, the following matrix weights.
W2
0
0
0
C
0
0
0
W2
46

---

Published as a conference paper at COLM 2025
0
W2
0
11
W
1,
20
0
0
0
0
W
0
1,1
W
1,4
wo
w 4,a
1,1
where Ce R, constant value, and w denotes the Gaussian entry to guarantee the near-orthogonality With these matrix, the query and key vectors are represented by
at = (1,0,...,0)) k; = (c.i,O,...00
Thus, the block outputs
Σ (33 W3WAn
(49)
where the approximation holds for sufficiently large C and ( the 1-th column of H(1), Here, we assume that the attention block uses strict causal attention, which attends only to positions in (1,2, 1} for the token position f.
Eq. 49 implies that the attention block attends only to the previous token, thus forming
previous token head. With the residual connection, the hidden state H(2) represented as:
0
0
0
0
T + 1
(22)
0
0
0
41,1
A1,
01,
01,T+1
+
0
2221
P222
Q2aT
4d,1
042
0d3
0 11/
where a; ((,,,)).,,44.)) and <2 the bottom right sub-matrix from WWGWW)aa443,0d-3
Finally, in the third transformer block, we set the following attention matrix weights:
0
0 0 0
3
1
0
3
0
1,1
1a
wkk
d
47

---

Published as a conference paper at COLM 2025
0
W
0
0
111
1a
0
wes
4
0
0
0
0
0
0
0
0
W
1,
w00
wo'
4,1
4
where w.,, denotes the Gaussian entry to guarantee the near-orthogonallty, and the sub- matrices in the bottom right of W and Wa are constructed the same way as W3 and W3 in Lemma 1. The result of this calculation is the same as in Lemma 1 except for the first three rows. Thus, the induction head mechanism is implemented the three-layer transformer with no positional encoding.
G Experimental setup
G.1 Example of Bigram Sequence Generation with Trigger-Output Pairs
To illustrate the sequence generation procedure in our modified bigram language model, provide concrete example with small vocabulary V A, B, C, D. The generation proceeds in discrete steps as follows:
Step 0: Sample trigger-output pair
new trigger-output pair (gk, 0z) is sampled from Ty and no. For example, suppose (q),0, (A,
Step 1: Sample the first token
The first token 21 is sampled from a predefined distribution over V, e.g., uniformly. Suppose B.
Step 2: Trigger check and transition
If z1 q1, then set 22 01. Otherwise, sample 22 TT( 21). Since 21 # 91, we sample
22
Step 3: Trigger condition applies
Since 22 = 91 = A, we set Z3 == 01 C.
Step 4: Normal bigram transition
Since 23 # 91, sample z4 Tb(- 23) and suppose z4 D. These steps result in the
sequence BACD.
Step 5: Re-sample trigger-output pair
Return to Step and sample new (g))0) from Tg and no to generate new sequence.
G.2 Neglect of in-context knowledge
We conducted our experiments using the following setup, carefully designed to evaluate the performance of a transformer model with APE and RPE.
48

---

Published as a conference paper at COLM 2025
Sequence generation Following Bietti al. (2024), we define Nh and πt, unigram and bigram character-level distributions estimated from the Tiny Shakespeare dataset (Karpathy, 2015). We sample triggers from 7 Tur and corresponding outputs of also sampled uniformly for each sequence generation. The vocabulary size [V] is 65.
Data Arguments The model was trained sequences with length of 256 tokens for the previous token head experiment and 128 tokens for the generalization experiment. The number of trigger tokens was set to 5, and we did not fix trigger tokens, i.e., all trigger tokens were sampled uniformly.
Model Configuration The transformer model used for the experiments had the following configuration:
Model dimension: 256
Vocabulary size: 65
Input sequences were capped at maximum length of 256 tokens.
The embedding layer Wt and unembedding layer Wu were frozen during training.
The model did not use WF in the second layer to focus on the role of induction head.
The model has 0.5M parameters.
Optimization Setup parameters:
The optimization process was conducted using the following hyper-
Optimizer: Stochastic Gradient Descent (SGD) with momentum
Batch size: 512
Learning rate: 0.2
Momentum: 0.9
Weight decay: 0.0001
Training Strategy
The model was trained for a total of 1000 iterations, which takes 0.5 hours of A100 GPU time.
Loss computation was restricted to output tokens.
G.3
Global vs in-context
We conducted our experiments using the following setup for the evaluation of performance of a transformer model with RPE trained on The Google analogy dataset.
Data Arguments tokens.
The model was trained on sequences with a maximum length of 256
Since The Google analogy dataset offers analogy pairs, such as (Tokyo, Japan), the bigram conditionals will mostly be in the form like Tt (Japan Tokyo) We constructed new bigram conditionals by adding some randomenss.
We first sample 10 fake target words, such as USA, China for each source word. Then, the conditionals will be calculated as follows:
1. count the number CA (B) of analogical pair (A, B).
2. sample PA E [0101,0.1] for each source word A.
3. consider the analogical pair (A, B,) appear (BC)(B) times, where B, is the fake target words for source word A.
49

---

Published as a conference paper at COLM 2025
4. calculate the bigram conditionals based on the number of appearance of analogical pairs.
Everytime generate input sequence, sample 5 trigger tokens from the set of all source words, and uniformly selects the corresponding output tokens from all vocabulary. start from source word that is sampled uniformly from the set of source words. Then, we transition with either the bigram conditionals or to the output token the current token the trigger token. Then we place comma, and sample uniformly new word from the set of source words for the next token, and repeat this process until the maximum length is achieved.
Model Configuration The transformer model used for the experiments had the following
architecture and configuration:
Model dimension: 512
Vocabulary size: 233
Input sequences were capped at a maximum length of 256 tokens.
The embedding and output layers were frozen during training.
The model has 2M parameters.
Optimization Setup
The optimization process was conducted using the following hyper-
parameters:
Optimizer: Stochastic Gradient Descent (SGD) with momentum
Batch size: 512
Learning rate: 0.1
Momentum: 0.9
Weight decay: 0.0001
Training Strategy
The model was trained for total of 100,000 iterations, which required 9 hours of A100 GPU time.
Loss computation was applied to all tokens except for commas.
G.4 About datasets
The Tiny Shakespeare dataset licensed under the Apache License, Version 2.0, and we can freely use the content long comply with the terms of the license. Similarly, the Google Analogy dataset is licensed under CC BY-NC 4.0, allowing use of the content for non-commercial purposes with appropriate attribution. Our of these datasets is consistent with their intended purposes. Furthermore, there is no explicit measure that should be taken check for personal information or offensive content.
H Other experimental results and analyses
H.1 Generalization to longer sequence
To evaluate the generalization ability of Taape and TFrpe we trained these two types of transformers with sequences length 128 generated from the bigram model. Details regarding the training procedure are provided in the Appendix G. We measured the accuracy of APE and TFRPE with respect to the second and subsequent occurrences of the output tokens. In addition, We present the memory recall to examine the generalization.
The results are shown in Tab. 3. terms of accuracy, transformer with RPE raises its accuracy even when the input sequence length increased to 256, while transformer with
50

---

Published as a conference paper at COLM 2025
Attention heazmep Sample
Attercion Heatmap Sample
61
120
122
124
92
322
266-
5.5 0 2
(a) Attention patterns of TFAPE
(b) Attention patterns of TFRPE
Figure Two-layer transformers with APE and RPE that are trained on sequences of length 128 show different first-layer attention patterns for length-256 sequences. The previous token head of TFA do not function well for positions 100, while TFRPE attends to previous tokens regardless of the positions.
acc. (f 128)
acc. < 256) 0.7925 + 0.0000
score (t 128)
score < 256)
TFAPE
0.8278 = 0.0000
0.5460 + 0.0007
0.1979 = 0.0001
TFRPE
0.8727 0.0000 0.8974 + 0.0000
0.9513 0.0000
0.9337 + 0.0001
Table 3: The evaluation results under the change in sequence length from 128 during training to 256 during evaluation. Transformer with APE degrades its accuracy and attention score as the sequence length grows, while transformer with RPE shows a stable accuracy and score.
APE dropped its accuracy to 79.25%. Although the difference may seem trivial, it is worth noting that the accuracy for f < 256 computed through the whole sequence including t < 128. This means that transformer with APE is more likely to fail to predict output tokens at positions 128 < t < 256. Looking at the attention scores, we can confirm that TFAPE is incapable of attending to previous tokens, even for positions f < 128. In contrast, TFRPE is trained to have a previous attention head that keeps its performance no matter how long the sequence is.
In fact, the attention pattern heatmaps in Fig. 3 in Appendix H tell us that the attention is no longer directed to the previous token after > 100 for TF APE while TFRTE continues to attend to the previous token. These experimental results match our theoretical result, where RPE is better suited for length generalization. can be seen from Fig. 3b that the first-layer attention head of TFRPE sometimes attends to the current token. This phenomenon is also explainable from the proof of Theorem 3, which shows that the attention matrix W not only learns the associations between wg( and 1 but also the pair of We (g(v) and ro.
H.2 Discussion of collision of context information
We discuss the result shown in Fig. 1b. The figure illustrates that is more likely to output global knowledge when the number of B₁ and B2 are the same, or when the number of B1 or B2 is almost maximum. Using Proposition 2, we can explain these phenomena by considering inequalities using the global knowledge log 7t,( A) and the strength of in- context knowledge T3. Using Proposition 2, we can explain these phenomena as follows.
If we have the relationships log (B, A)+ + 73 > log (B A) > log n. ( (B, () + /2, where B, represents B₁ and B2, the model outputs global knowledge B when AB₁ and AB2 occur equally frequently the prompt. On the other hand, if we have log Th (B2 A) + T3
log πτ, (B A) log 7t, (B1 A) + 73 and log TB A) > log Ty TogT,BB2 (B2 A), the prediction of the model changes to B2 to as the number of B1 increases.
51


Reviews:
[{"review_id": "WgVVhDJxLB", "reviewer": "Reviewer_1jjv", "summary": "This paper investigates in-context learning from the point of view of associative memory, where the core experiment analyzes how a model balances global knowledge vs knowledge given in context. The paper shows that using relative positional encoding (RPE) in two-layer transformers lets induction heads form more stably avoiding the recall decay seen with absolute positional encoding (APE). The authors motivate this theoretically and confirm empirically (on synthetic bigram data and analogy tasks) that RPE models outperform APE models especially at long context. These results clarify how positional encoding underpins reliable in-context learning.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "- **Timely Topic**: This paper investigates a timely topic of global pretrained statistics vs in-context information.\n- **Proof of RPE stability**: They provide a p roof showing that RPE forms a more stable induction head irrespective of token position\n- **Simple synthetic data**: The paper uses very simple synthetic data to validate their findings, which helps the interpretation of the results\n- **Explanation of undertrained token effects**: This paper addresses the recent curiosities of undertrained tokens in LLMs [1].\n- **Discussion of limitations**: They acknowledge the limitations: analysis constrained to two-layer transformer and induction heads.\n\n[1] https://arxiv.org/abs/2405.05417", "reasons_to_reject": "- **Heavily relying on theory of simple architectures**: Theorems proven on a two-layer transformer may not generalize to larger or more realistic models. While the small-scale experiments are helpful, it’s unclear how broadly the proofs apply.\n- **Lack of surprisingness**: It isn’t particularly surprising that RPE yields more stable induction heads—RPE explicitly encodes relative distances, which appear as a specific direction in the token representations. It isn’t that surprising that induction heads can naturally pick up this direction.\n- **Setting description could be improved**: The paper could more clearly specify exactly how the synthetic training data are constructed and sampled.\n- **Presentation of results**: It’s hard to gauge the real impact of the experiments without explicit accuracy-vs-context-length curves. Showing only two curves for first half and second half doesn’t show how the performance decays. Figure 1a in its current form is difficult to interpret, what the main findings are.\n- **Practical relevance**: It is unclear how relevant the study of positional encoding’s effect on ICL is when in practice, APE is almost never used compared to RPE.\n- **Lack of some critical references**: In general, the paper cites related works throughly, but significant work on global vs in-context correlations and in-context learning’s relation to the training distribution is missing: https://arxiv.org/abs/2312.03002, https://arxiv.org/abs/2401.12973, https://arxiv.org/abs/2412.01003, https://arxiv.org/abs/2306.15063, https://arxiv.org/abs/2404.07129", "questions_to_authors": "- **Bigram Model**: I am a little bit confused about this bigram model’s description. For this model to represent in context learning, shouldn’t the trigger token-output token pairs mostly be specified by the context? It seems like the trigger tokens are fixed for the whole data.\n- **Latter half**: I am quite confused about the mention of the latter half near line 165 and the comparison between sequence position <128 and >=128. Is there a sharp transition at half the sequence length? It seems like the transition is smooth. If there is a sharp transition at half and this is predicted by theory, I might update my score significantly as this would have real generalizable implications!\n- **Comparison with NOPE**: While I criticized the comparison of APE vs RPE since APE is seldom used, I think a comparison with no positional encoding (2) might be useful since then one is really comparing whether randomly initialized positional encoding is *worse* than a hard to represent but noise-free encoding.\n\n[2] https://arxiv.org/abs/2404.12224", "extra_scores": {}}, {"review_id": "tCSNEmYmo0", "reviewer": "Reviewer_SFDP", "summary": "The paper provides a theoretical analysis for understanding in-context learning in transformer LMs, specifically focusing on the associative memory and the induction heads. The authors first prove that a simplified 2-layer transformer, using Relative Position Encodings (RPE) can learn the previous token head without encountering the oversight issue that affects Absolute Position Encodings (APE). The authors also provide a construction for a 2 layer transformer that learns the bigram distribution of the training data through the feed-forward memories and the induction head behavior through the self attention layer. Through their construction, the authors show that the next token prediction is dictated by the bigram distribution as well as the frequencies of possible continuations that follow the current token in the sequence. The authors also provide toy experiments to provide evidence for the theoretical result, where models trained with APE fail to learn the desired associations with previous tokens as the sequence length increases but RPE models do. For global vs. in-context associative memories, the authors provide experiments on the Google Analogy Dataset modified to contain trigger tokens for the in-context associative memories, and find trained models show some evidence supporting their construction.\n\nOverall, the theoretical results of the paper are interesting and build upon the current theoretical understanding of transformer models. The experiments also provide some support for the theoretical findings, but do lack a more thorough exploration, which might be needed to show how much the theoretical constructions agree with what models learn in practice.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "- The result about transformers with RPE being robust to oversight as compared to APE is very interesting and also a potentially important result in the theoretical understanding of transformers.\n\n- While the constructions are provided for a simplified version of transformers with 2 layers, I appreciate the authors include feed-forward layer in their analysis as well, since those can be harder to construct given the non-linearity. The constructions very clearly show the interplay between the global and in-context associative memories, whereas the constructions in the prior work focused on one or the other but not both.", "reasons_to_reject": "- While through the experiments the authors show that APE models do not learn the desired associations at later stages of the sequence by calculating the memory recall metric, it would have been interesting to see whether this impacts the performance of models in correctly predicting the next token. In lines 303-310, the authors provide a discussion of how this can influence model prediction but, to my understanding, do not provide any experimental evidence to support it.\n\n- I would have liked stronger empirical evidence for the model's learning behavior consistent with the proposed construction of Definition 2.", "questions_to_authors": "- What form do you consider for the absolute position encodings $p_t$? Are these learned d-dimensional embeddings with Gaussian entries? If not, do they depend on the value of t? \n- In Remark 1, in this line “if the input sequence contains out-of-distribution tokens, the induction head mechanism does not activate”, how are out-of-distribution tokens defined?\n- The notation in Lemma 1 is a bit confusing. The $Q$ in $W_Q$ and the $Q$ which is the support of $\\pi_Q$ are not the same right?\n- In equation 8, what does the prime symbol over the logit symbols denote?\n- In the same equation, shouldn’t the difference be between the logarithm of the logits to be equal to the value in the right-hand side?\n- I would suggest the authors define induction heads at least informally before talking about them in the introduction.", "extra_scores": {}}, {"review_id": "6sVfNcsFh5", "reviewer": "Reviewer_dwGn", "summary": "This paper presents a theoretical analysis of associative memory in two-layer transformers, extending some of what is known about how induction heads work to more than one word. A new result shows that such a transformer has attention mechanisms that effectively mirror the transition probabilities in a bigram model.  An experiment on an analgoies task provides some suggestive evidence that such a transformer can be learned from data  (though only word-pair data, since the analysis is focused on bigrams).\n\nTo be honest: I have not followed all of the induction head literature closely in the last year or so, so I am not particularly well-positioned to analyze the novel contribution of this paper.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "* Introduces an explicit formulation of an associative memory transformer, and links this mechanism with bigram knowledge. \n* One new theoretical result on how associative memory implements bigram probabilities, together with one empirical study.", "reasons_to_reject": "* The paper relies very heavily on Bietti et al 2024.   While it does seem independent and a new contribution, a lot of the paper is hard to follow if the reader is not intimately familiar with that earlier paper.\n* Unclear whether / how the results generalize beyond bigrams, since both theoretical and empirical analyses depend on either text generated by a bigram model or word-pair tasks.  More discussion of how the simplifying assumptions in the theoretical results related to actual transformer implementations and natural language data would be helpful.", "questions_to_authors": "* Many papers cite only the arXiv version of works that have since been published (e.g. Friedman et al 2023, \"Learning transformer programs\").  I encourage the authors to check all their references for canonical published versions.", "extra_scores": {}}, {"review_id": "UoLncm60hm", "reviewer": "Reviewer_8sjk", "summary": "This paper aims to understand the ability of transformers to coordinate long-context information to answer in-context learning (ICL) queries and global knowledge acquired during the pre-training stage. It focuses on two-layer transformers, which have been shown to develop induction heads for ICL, from the perspective of associative memories. The model is trained on a bigram model with some trigger tokens for next-token prediction (NTP). The paper shows, theoretically and empirically, that (a) transformers with RPE can use long-context information, whereas transformers with APE don’t work well in this setting as the attention scores to latter parts of the sequence become smaller with increasing length, and (b) how transformers prioritize in-context knowledge over global knowledge by analyzing the weight matrices and resulting logits for specific prompts.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "The paper studies interesting questions of (a) effect of RPE vs APE in length generalization for ICL in transformers and (b) how global and in-context knowledge influences the output of transformer.  \n\nThe theoretical and empirical results presented in the paper are interesting and add to understanding ICL in transformers.", "reasons_to_reject": "The paper only considers a simple data model and one real dataset; it can be strengthened by including some (more) experiments on real data with larger models to support the key findings, and investigate how general they are. \n\nThe paper writing and organization should be improved. Currently, the title and abstract don’t clearly reflect the main message of the paper. It seems to have two different focus points - length generalization and global vs in-context knowledge, but they seem to be treated as two separate things and are not connected well together. Due to this, the paper doesn’t seem to have a unified focused narrative.", "questions_to_authors": "I like Table 1 which summaries and compares related work on associative memories. Can the authors also include Nichani et al. 2024 and Friedman et al. 2023 they discussed in Section 2 in the table?  The sentence in lines 59-60 is incomplete. \n\nThere’s a typo in the equation below line 97.", "extra_scores": {}}, {"review_id": "uq3zPkID9T", "reviewer": "Reviewer_Pgux", "summary": "This paper provides new theoretical results for induction heads on top of the associative memory viewpoint of (Bietti 2024). Specifically, it shows that for two-layer transformers, relative positional embedding (RPE) can help detect patterns in long sequences that previously have been shown to fail for absolute positional embedding (APE). Further, it shows how in-context knowledge and global knowledge (from pretraining) interact for next token prediction. \n\nOverall, this is an interesting set of results that improve our understanding of Transformer models. The writing is clear and convincing.", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1. New theoretical results for induction heads", "reasons_to_reject": "n/a", "questions_to_authors": "1. I don't fully appreciate the term \"oversight of patterns\". Oversight can mean either \"unintentional failure\" which is your meaning but also can mean \"action of supervising something\". Is there some other term available? \n2. Definition 1, Line 120: Is there a typo? \"A matrix is called...\" Do you mean W? \n3. Proposition 1: I don't understand the purpose of this for this paper?", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper provides a theoretical analysis of how a Transformer would be able to represent associative memory, and how it would balance the two goals of using in-context information vs. using in-weights knowledge. The paper also provides insights into the effect of RPE vs APE for length generalization.

The reviewers acknowledge that the topic considered here is timely; the contributions are interesting and novel.

The major concern that stood out to me was the potential gap to practical relevance (generalizing to multiple layers, beyond bigrams etc.,) and some lack of surprise (of one of the two main theoretical results) raised by one reviewer. However, I believe that the results in the paper are novel and insightful enough in light of what exists in literature --- hopefully future iterations on top of this work can extend these analyses to more interesting settings. 

I would argue it makes sense (and even desirable) for the theory to be in a simplified setting without unnecessary complications that make the proofs lengthy with no new insight; however the paper would be strengthened with more experiments verifying their insights on real-world data as suggested by multiple reviews (`SFDP `, `1jjv`, `8sjk`).  I hope the authors will incorporate the feedback about:
- stronger experiments and
- improving the clarity in the description of the setup, and recapping results from Bietti et al 2024.
- any missing citations
for future versions of the paper.

[Automatically added comment] At least one review was discounted during the decision process due to quality]
Decision: Accept