From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers
Authors: Ryotaro Kawata
Yujin Song
Alberto Bietti
Naoki Nishikawa
Taiji Suzuki
Samuel Vaiter
Denny Wu
Venue: NeurIPS 2025 spotlight ()
Keywords: Transformers
induction head
training dynamics
Primary area: theory

Abstract: Transformers can implement both generalizable algorithms (e.g., induction heads) and simple positional shortcuts (e.g., memorizing fixed output positions). In this work, we study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other. Focusing on a minimal trigger-output prediction task -- copying the token immediately following a special trigger upon its second occurrence -- we present a rigorous analysis of gradient-based training of a single-layer transformer.
In both the infinite and finite sample regimes, we prove a transition in the learned mechanism: if input sequences exhibit sufficient diversity, measured by a low “max-sum” ratio of trigger-to-trigger distances, the trained model implements an induction head and generalizes to unseen contexts; by contrast, when this ratio is large, the model resorts to a positional shortcut and fails to generalize out-of-distribution (OOD). 
We also reveal a trade-off between the pretraining context length and OOD generalization, and derive the optimal pretraining distribution that minimizes computational cost per sample.
Finally, we validate our theoretical predictions with controlled synthetic experiments, demonstrating that broadening context distributions robustly induces induction heads and enables OOD generalization. 
Our results shed light on the algorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.

From Shorteut to Induction Head: How Data
Diversity Shapes Algorithm Selection in Transformers
Ryotaro X..2 Yujin Song".l.2 Alberto Bietti Naoki Nishikawa 1.2
Taiji Suzukít.3, Samuel Vaiter Denny 303.6
'The University of Tokyo, PRIKEN AIP, Flatiron Institute, *CNRS, Université Côte d' Azur, Laboratoire J.A. Dieudonné, "New York University
ecc. u-c..u ac jp; roogg.reeearcc@@mll..cmm com; (abietti, dwu @flatironinstileeeeeeeeeo.innfee tetiittteee.... org;
y. song. reseerseerrrre whw.aail.cmmmaa.cccommmo saijjimistti..u-tooe i. ac. jp; samuel vwiterecnrs fr
Abstract
Transformers can implement both generalizable algorithms (e.g., induction heads) and simple positional shortcuts (e.g., memorizing fixed output positions). In this work, we study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other. Focusing output prediction task - copying the token immediately following special trig- ger upon its second occurrence based training single-layer transformer. In both the infinite and finite sample regimes, prove transition in the learned mechanism: if input sequences ex- hibit sufficient diversity, measured by a low "max-sum" ratio of trigger-to-trigger distances, the trained model implements seen contexts; by contrast, when this ratio is large, the model resorts to posi- tional shorteut and fails generalize out-of-distribution (OOD). We also reveal a trade-off between the pretraining context length and OOD generalization, and derive the optimal pretraining distribution that minimizes computational cost per sample. Finally, we validate our theoretical predictions with controlled synthetic experiments, demonstrating that broadening context distributions robustly induces induction heads and enables OOD generalization. Our results shed light on the al- gorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.
minimal trigger-
present rigorous analysis of gradient-
induction head and generalizes
Introduction
Large language models (LLMs) leverage circuits of attention heads [VSP+ 17] to perform (implicit) algorithmic reasoning. Certain attention heads implement discrete algorithms - notably induction heads [ENO+21, 21, 22]+ which for previously seen token patterns in the context predict subsequent tokens. Such heads enable in-contexi learning behaviors [BMR+20], allowing trans- former continue sequence such as [A, B...., ... A) text. By contrast, attention can also implement positional mechanisms that select tokens based solely on their location in the sequence IVTM* 19, AWKA24]. These mechanisms can yield contrasting generalization performance [CBKZ24], and we expect the pretraining data distribution to play central role determining which mechanisms a model learns rely on: depending on structural properties of the corpus, a transformer may either discover generalizable strategies content-based retrieval) or adopt position-baat based shortcuts.
B purely by leveraging patterns in the
Equual contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Pretraining Data
00D Data
Solution 1: positional shortcut
Output here
*
PiT+1X2
Pr
Pr
Match position T
Not Generalizann
Solution 2: induction head
2) Output next
...
C;
e; Generalizable E,
11 Natch trigger
t] - In the pretraining data, the size of irrelevant tokens hefore the occurrence of the first and second trigger l remains fixed per sequence, hence allowing two solutions: (i) positional shortcut that outputs the token position (T + +1)/2 1)/2 for input length Z'; and (ii) induction head using token embedding e. which finds the queried token and returns the ensuing token. Whereas on OOD sequences with varying & # l2, only (ii) remains a valid solution.
Figure 1:
Two mechanisms for the associative copying task [..., t, 0,
Motivation. We theoretically study how pretraining data influences the implemented circuit and -fitt--distributtin (OOD) generalization performance of the transformer. This perspective is mo- tivated from the empirical observation that pretrained models often leverage shortcut solutions that are brittle beyond the training distribution [MPL19, GJM+ LAG+22). For instance, a trans- former might utilize the aforementioned position-based attention head to memorize that certain output tends to occur at a particular position in the training text, instead of learning the underlying association (induction head); such positional shorteut is double-edge sword in algorithmic tasks: transformers can achieve near-perfect accuracy in distribution, but struggle on test sequences of un- seen lengths or structures. Since it is empirically known that the learned mechanism heavily depends on the structure of pretraining data [GTLV22, RLIGS22, RPCG23, WNB N25], we ask the following question.
How does the data structure decide whether a pretrained transformer implements a generalizable mechanism induction head) or a shortcut that fails OOD (e.g., positional memorization)?
1.1 Our Contributions
Trigger-output Copying. To investigate this question in a controlled setting, introduce min- imal trigger-output copying task inspired by [BCB [C]. In this synthetic task, each input sequence contains special trigger token that appears twice. The model must predict the token that imme- diately follows the first trigger when the trigger appears the second time. For example, given
... [trigger]][ [X] ... [trigger][]]
the correct prediction is X. Depending on the structure of the input sequence, this task admits multiple solutions. We focus on two mechanisms - see Figure 1.
Induction head. The model attends back to the location of the previous trigger and copies the token following it; this works for arbitrarily long gaps between trigger occurrences (up to context- length limit).
Positional shortcut. When the position of the first trigger is inferable from the second (c.g., under periodic structure), the model may copy the token using positional information alone. This shortcut is valid in-distribution but does not reflect the underlying association.
(OOD) generalization as performance on test sequence with altered structure, where the trigger appears positions not seen during pretraining (e.g., longer or aperiodic sequences). The induction head mechanism is robust to such shifts as it learns the correct association, whereas the positional shortcut typically fails OOD. Our goal is to identify a data-dependent transition between these two mechanisms that governs OOD generalization: intu- itively, increasing the diversity of pretraining sequences - by varying the distances between trigger occurrences dilutes positional signals and discourages the shorteut; conversely, as the number of
For this task, we define

---

trigger tokens in the data grows, the effective signal for induction weakens. We make these intuitions precise our theoretical analysis.
Main Findings. We provide a quantitative account of how pretraining data diversity shapes the mechanism learned by pretrained transformer in the trigger-output copying task introduced above. Specifically, we rigorously analyze the in-distribution and out-of-distribttion performance of a shal- low (single-layer) transformer trained on this synthetic task. By studying "carly-phase" simplifi- cation of gradient descent in both the infinite-data (population loss) and finite-data (empirical loss) regimes, show that the pretraining distribution directly selects the model's algorithm: when pre- training data are sufficienlly "diverse" as measured by a max-sum ratio trigger distances the transformer learns induction head; when diversity is low, the model adopts a positional shortcut that fails to generalize OOD. Using this diversity measure that governs the phase transition, we dis- cuss various tradeoffs and how choose pretraining distribution that induces the desired induction mechanism with minimal computational cost. Finally, we empirically probe the learned circuits by visualizing attention scores, and present evidence that similar mechanism transition arises under standard gradient-based training beyond our theoretical setting.
1.2 Related Works
The induction head mechanism in transformers was first presented in the mechanistic interpretability literature [ENO+21, 21, OEN+ 22), and followup theory investigates when such circuits emerge under simplified training dynamics and tasks [BCB + 23, ETE 24, NDL24, Red24, CSWY24). Empirical studies on algorithmic tasks (copying, arithmetic, sorting) demonstrate that transformers often rely on spurious "shortcut" solutions that fail to generalize, often due to poor use of positional informa- tion [ 23, 23, GJB+ the OOD brittleness of shortcut solution is also doc- umented in [LAG++22]. complementary thread links the structure of pretraining data in-context behaviors: the function classes a transformer implements in context and the sensitivity of perfor- mance to data statistics such as corpus coverage and frequency [GTLV22, RLIGS22, MLH or task diversity [RPCG23, LLZV+25), 25]. Our analysis aligns with this view by making explicit how diversity in trigger distances steers the learned mechanism. Methodologically, we borrow the "early-phase" simplification of training dynamics and study the loss improvement after the first few gradient descent step [BES + 22, DLS22, ORST23, BCB+ 233].
2 Problem Setting
Notations. For positive integer N, we denote [N] : 11. For integers N1 S N2, we define [N1 : N₂] == (Ni, N1 + 1,..., N2). The Softmax function for an N-dimensional vector v E R is defined as Softmax(V); == E For a vector v, we write v O2(((N)
||l||2 = O(((N), and e = O (f(N)) if max, |os1 = Of((N). Similar notation is used for a matrix A, where ||l and || denote its l2 lg spectral and max norms, respectively.
2.1 Data Generating Process
We study the trigger-oulput setting to investigate how transformers acquire the induction head mech- anism. Let N E N denote the vocabulary size and L E N the maximum input sequence length. We designate special tokens as trigger tokens. We define our data model as follows:
Definition 1 (Data Distribution). Let E,, lg N such that T == l1 + fg +3 < L --1- 1. Let Nug S N
denote the number of trigger tokens. A sequence 21-T+1 [N]T+] is sampled as follows:
1. Sample trigger token E [Nur]] and an output token E [Narg + 1: N] uniformly at random, where Norg o(N1/3).
Construct the sequence:
Z1.T+1 =
t trigger-outuul pair
2e, 43,
t, trigger-output pair
é1 irrelevant tokens
Es irretevant tokens
where irrelevant token 2; (i [1
(] +3: ( + le + 2]) is drawn i.i.d. from [Nug + 1 Nj.
We refer to such sequence trigger-output model with subtext lengths & and lz.
3

---

In our data model, the task is to identify the output token 27+1 = from the sequence 31:T (21....;22) 2t, t, o. 2t, -f,-3... 3............................... 24, +x+2: t). This can be achieved by implementing the induction head mechonism [ENO+21, 2EN+ which copies the token that follows the first occurrence of the trigger token and outputs it upon encountering the second occurrence of the same trigger.
Due to structure of the input sequence, transformer may also rely on positional shortcuts achieve low loss; in particular, when the lengths of irrelevant tokens are identical within each sequence, i.c., 4 = f2 L, a transformer can achieve 100% training accuracy simply by inferring the correct position to attend to (T' + 1)/2 = L + 2 from the position of the second trigger I = 2 - 3. Such positional solution does not make use of the semantic information and generally fails when & # l2-
To study the transition between the mechanisms, assume the pretraining data consists of mixture of sequences with different lengths determined by {= E1 = l2--
Definition 2. Consider language model po 2222- 27) that is pretrained on M sequences
(1) M
(iTT generated as follows:
Sample ((0) from distribution De-
Generate (0) 117644 according to Definition with & ((0), i.e.
(0) 11:0+1 = (21....; Egtts 1, O, Eg00 +2: 1, 0).
OOD Generalization. Note that the pretraining distribution (defined by De) may not cover all possible sequences. We say that Pg generalizes out-of-diitrrbbution (OOD) if it implements the cor- rect copying mechanism across all possible f's. that is, for any é1, l2 such that éi + &a + 3 S L-1 (possibly & f f2), and for any test sequence 21:7+1 generated from the trigger-output unigram model with subtext lengths & and &2 (Definition 1), we have
arg max Po [ 21.2....0 27) 2T+1- NE[N]
2.2 Gradient-based Training of Single-layer Transformer
Architecture and Embedding. We consider a single-layer transformer block fTF defined as ( Jii WKQ, Wv) = W. XI: Softmax X WKQT);) E RRV.
(2.1)
where WKQ E DDD Wv NN and X1 = (æ)
xx)) E Rx denotes the input embeddings of 21:1, with embedding dimension D. We define the embedding as follows:
Definition 3. Let D = I 2N. Let Pt E RR_ denote the one-hot vector with I at the t-th position (representing the positional embedding), and let e₂ E R.N denote the one-hot vector with a at the z-th position (representing the token identity).
We then construct the input embedding æ;
Pi
Eze
E 122N
(2.2)
ezz-1
The prediction probability is given by
Pw = k 2 ... ......)T)=))) = [Softmax(fTr (Xa:: WKQ, Wv)lle.
Remark 1. We make the following remarks on the design of our architecture and embedding.
The architecture (with the FFN is absorbed into the value matrix Wv, and tied key and query projections) is commonly used in theoretical analyses and mechanistic studies [L.R223, BCB+23, NDL24]; the simplification allows us focus the inductive bias by simple attention mecha- nisms, while retaining sufficient expressiveness to implement algorithmic behaviors.
Two-layer architecture is typically needed to implement the induction head mechanism, where the first layer often learns to detect the trigger and identify of the following token via attention to the previous token [SHT24]. To reflect this inductive step in our simplified single-layer setting, we ex- plicitly encode the identity of the previous token 2:-1 in the third component of the embedding x.. This choice also echoes recent empirical developments that incorporute information of previous tokens directly into the current state, such as Mamba [GD23), RWKV [PAA+ 23], and convolution augmentations [LZHO25, AlI25].

---

Algorithm 1: Gradient-based training of single-layer transformer
Input Learning rate TJQQ; ' Initialize (k)(O) = ((L42))x(L+2N) Wv(() ON((L22N)
Gradient descent on Wv
(1( W (0) TΙ Ww My EiS C(X 1.70i WKQ(O), Wv(0)
Gradient descent WKQ
WkQ(1) - @ γкQ V VKK Miq Ei=M++1 c(x (1 LTO WkQ(O), Wy(1))
Output: Prediction fFF(.)
Gradient-based Learning Algorithm. We use gradient descent (Algorithm 1) on the cross- entropy loss to pretrain our shallow transformer (2.1),
C(X LTTO) WKQ, Wv) = CrossEntropy( (2:,,, (oftmax(frr(X (X WKQ. Wv))).
In Algorithm 1, we apply single gradient descent step with large learning rate on the value and key-query matrices. This is motivated by recent studies [ORST23, BCB+23, WS24] showing that the first gradient step can induce associative memory tied to specific components of the input em- bedding. In particular, the gradient can often be expressed as a linear combination of outer products we where either dient structure is sufficient to construct simple forms of associative memory within the model. We remark that similar single-step update is commonly used in the analysis of feature learning in shallow neural networks [BES* f 22, DLS22, BEG 22] and transformers [OSSW24, NSO WNB + 25].
corresponds to embedding vectors such ex2 Ezt- 1- Pr. Such gra-
3 Main Result: Data-driven Transition Between Mechanisms
3.1 Positional Shorteut vs. Induction Head
In this section, we illustrate how the diversity of pretraining distribution influences which algorithm the trained transformer implements - either the positional shortcut or the induction head. The following quantity plays central role in characterization.
Definition 4. For each f, let qe denote the probability mass assigned under De. and S the support Dr. We define the max-sum ratio as
R(D:) = maxtes E-1y
Etest \1'
Interprefation of max-sum ratio. The max-sum ratio can be seen as diversity measure of De- The following example provides an intuitive illustration: Example 1. Ler De Unif((lo, tu 1...., Eu K -1)). Then the max-sum ratio is given by
R(Lo, K)
(3.1)
Ek- (Lo + kk
which monotonically decreases with K; hence greater diversity of DD gives smaller max-sum ratio.
Note that the max-sum ratio does not merely capture the width of the distribution: in Example 1, in- creasing lo while keeping fixed decreases the proportion of 6 [6] . (6++ Ko++-1)-= 1))- thus reducing the max-sum ratio. Hence, even with a narrow range, shifting the distribution rightward placing more probability on larger ? naturally yields a smaller max-sum ratio. This because the max-sum ratio weights each probability mass qe by f-1
Learning under Population Loss.
The next theorem shows the existence of a threshold in the max-sum ratio that determines whether OOD generalization achieved, in the infinite-data limit. Theorem (Infinite Sample Setting). Suppose we run Algorithm on the expected loss N²
([A(Xi:T; WkQ, Wv)) with learning rates JN S 1, TyvK Z No log N. Then, there exist E1 (Nugg) (Nurg) ((N such that:
5

---

R(Di) <
é1, then the pretrained transformer generalizes OOD, defined in Definition 2.
If R(D)) > <2, then there exist OOD test sequences such that the pretrained transformer fails.
Remark
Note that the training data only contain sequences with & Eg, and thus a positional shortcut (as illustrated in Figure 1) can still achieve 100% training accuracy. However, since the OOD test data include sequences with & o l2, such shortcuts inevitably fail. Our main theorems show that the pretrained transformer avoids such shortcuts when the max-sum ratio is below certain threshold, ie., when the data distribution sufficiently diverse.
We also provide light (N charucterization of the max-sum ratio threshold, indicating that increasing the number of possible triggers makes OOD generalization more difficult. The under- lying mechanism is discussed in the ensuing subsection.
Learning under Empirical Loss. Our next result establishes (via gradient concentration) similar transition behavior in the fmite-sample setting.
Theorem 6 (Finite Sample Setting). Suppose ing rate scaling as Theorem 5, and with sample sizes MkQ Z poly log N No N
run Algorithm 1 with the same learn-
(E+v and My 2 poly N. No Cees v4.
Nin (Eis qet)) (E
Then, with probability at least 0.99
there exist ef (Nuug) (N ((Nigg) such that the assertion of Theorem holds by substituting (<),e) 62) for 62).
3.2
Mechanism of Algorithm Selection
Now take closer look how the positional shoricul and the induction head implemented in the attention. We begin with the case where the support of 'D; is a singleton and Nirg = 1. After single gradient step, the parameter matrix WKQ can be shown to implement form of associative memory over the relevant embedding vectors.
Lemma (Informal). Let De = gradient step Algorithm W KQ takes the form
(f), and assume the trigger consists of a single token w. After one
(p4+2 + +pe+3)
e 0).
WKQ & T(t)-1
O
Pr(o
ew
where T(O) 28 + 3 denotes the position of the second occurrence of the trigger token.
To further simplify the exposition, we ignore the cross terms between p and e and assume that WKQ takes the following form:
(pe+2
D&+3)
WKQ -())-1
Pr(s)
T
oT + T(C)
0 0
0]
(3.2)
positional shortcut
induction head
Now consider an OOD test sequence as in Figure 1, whose total length matches the training sequence but whose first and second subtext lengths differ: E la 26, E + lz. In this case, the two terms in (3.2) contribute to the attention score
Softmax
WKQ @Tic)) with
Trest E1 + lg + =
follows (see Figure 2), noting that TT...
(PT(e) Ew
1st term (positional shortcut). Regardless l,, it attends to the positions E+2 = (Tiss + 11/2 and & + 3 (Tiest + 3)/2. In particular, for the former, even though E * Eg, the transformer incorrectly associates the second trigger position Trest with (Tees + 1)/2 as e = Eg.
FFor the latter position (Teas + 3)/2, the model also attends embedding. This follows from detailed computation of W which we omit here.
same token via the previous-token

---

?]
7]
2
[*
*
t
?]
6
8
12345677
15
323456789
SE
3000
pedlicedl
1200
878
EF8
indurttoo
beod
555
Pj
e,
4-1
P
24-1
(a) Attention heatmap for l ~ Unif([3}).
(b) Attention heatmap for & ~ Unif ([3, 4, 6)).
Figure 2: Attention heatmaps of W.QQ when the pretraining sequence diversity is small (left) and large (right). In the left figure, there is strong positional shortcut that links position 9 to position (the correct position in pretraining data), whereas in the right figure, the trigger positions more dispersed, weakening this shortcut. Instead, a signal corresponding to induction head detecting tokens after trigger becomes dominant.
2nd term (induction head). It attends to tokens whose third embedding block equals ews i.e., tokens whose previous token is the trigger w. In other words, it scans for the trigger TTes and then attends to its next token this is precisely the desired induction head behavior.
Thus, the learned attention matrix implements a mixture of positional shortcut and induction head, and the relative strength of these components determines which algorithm is ultimately selected. Two factors affect this balance: the diversity of irrelevant token length & and the trigger size Ner.
Length distribution Dj. Equation (3.2) describes the case where é is deterministic. When is distributed according De. WKQ becomes superposition over &
(pe42 + p3-3) 0
WKQ(1) & Eqe
PTO OT 0T + B
0 e 0].
0
Here, the first term spreads its mass across multiple positions and is consequently weakened, whereas the second term does not depend on l and retains its strength. As result, the magnitude of the former is at most max; q: T())-1, while that the latter is Ee qe T(E) Since T(E) x the ratio between the strengths of positional memory and the induction head is nothing but the max-sum ratio R(D)). This explains why the max-sum ratio governs algorithm selection.
Trigger size Nirg In (3.2), when the trigger size Nirg 2 2, the second term is replaced by
Σ T()) 1
0
e o].
trx
wE[Nengl
while the first term remains unchanged. Hence, the induction-heal signal is split across trigger types and its strength decreases proportionally to N This explains O ((N:rg) threshold in Theorem 5. The above intuition visualized an experiment reported in Figure
Example 2. Figure 2, we set = and Nirg 2, train the model with D(E) = and D()) Unif([ 8)), and visualize the resulting KQQ The trigger-token set (1, 2). The training setting is the same that in Section 4.1.
7

---

(Left): when D(E) {3}, WKQ has strong component that maps position 9 to position 5. Although also contains an induction head component that maps between trigger tokens, is comparatively weak compared to the positional signal.
(Right): when D(E) (ni 8), W.Q exhibits a superposition of signals mapping position k to (I + 1)/2, which results in each individual signal being weakened. In contrast, the induction head signal does not diminish.
3.3 Tradeoff between Context Length and OOD Generalization
As discussed in Section 3.1, the max-sum ratio captures not only the overall "width" of the distri- bution but also decreases mass shifts toward larger f. This effect becomes especially pronounced near the ((Nerg) arg) threshold identified in Theorems 5 and 6:
Example 3. Consider the max-sum ratio for the uniform distribution (3.1). If to = 1, then R(lo, K) ((log K)-1). To atlain max-sum rafio of order ((Nug) the OOD general- ization threshold in Theorems 5 and 6 the support width must satisfy K exp(Nur)- By contrast, l = (Nurg). then it suffices to take K- (Nurg) to obtain a max-sum ratio of ((N
Therefore, merely "widening" the distribution may not be efficient reduce the max-sum ratio; biasing pretraining toward longer contexts is substantially more effective. This, in turn, suggests that reliably learning the induction-head mechanism (and hence achieving OOD generalization) may incur greater computational cost due to longer training sequences.
We now consider the "optimal" shape the pretraining sequence (under the constraint in Defini- tion 2) that learns the induction-head mechanism with minimal compute. Since the forward cost scales quadratically with context length, we seek short contexts while maintaining a favorable Nur. consider the optimization problem
max-sum ratio. Formally, for U
minimize
Σ/ gel2 get
subject to
-et-T S &Nurg
P:
1 a1 =1
11, QU 2
This objective is the sample-average forward-pass cost in pretraining: the constraints enforce the OOD threshold from Theorem 7 and the normalization of (qt) This problem is a linear program whose optimizer is characterized below.
Proposition The optimal solution of problem assigns limearly increasing probability mass to the first Nug context lengths and zero to the remaining ones:
(g1. 92. ........ ) - (1, Nurg 0, O,...,0,,
where the normalization constant = Nurg (Ntrg + 1)/2.
other words, to minimize average forward-pass cost per sample while meeting the OOD general- ization constraint, the pretraining distribution should be linear in the context length, making g uniform over S Nerg- We note that if one optimizes a different objective (e.g., incorporating sample complexity), the optimal pretraining distribution may change.
Numerical Experiments
4.1
Experiments for Theoretical Setting
To observe the transition from positional shortcut to induction head, we first consider the architecture defined in (2.1) and conduct experiments under the data model described in Definition
4.1.1 Experimental Setup
Dataset. We generate training and test data according to the trigger-output setting in Definition 1.
In the pretraining data, the lengths of irrelevant tokens f,, l are always equal. We choose two integers Emin and E (fmin S Emax), and length E is sampled from Unif([Pmins (max])

---

OOD Accuracy (Nirg=4)
OOD Accuracy (Ntr==8)
40-
aln oubhort tongghh Liwn
14h
145
(aa subtest lingth l max
max subtesd lrighh maa
(a) number of triggers Nug =
(b) number of triggers Nug =
Figure 3: Out-of-distribution accuracy map over varying biin (vertical) and max (horizontal); moving right indicates greater diversity the pretraining distribution.
In the OOD test data, we shift the position of the first trigger to produce non non-periodic sequences. Specifically, we first sample & ~ Unif ([emi + maxx), and then sample & Unif((1..... ,2- 1} (f)), defining & 2 l so that & + E2.
Model architecture, embedding, and training. We implement one-layer transformer architec- ture as defined in (2.1) with embeddings defined in (2.2). Training follows Algorithm 1, and the learning rates for W and W set to 103 and 10 respectively. Both matrices trained with the empirical cross entropy loss computed on 8192 training examples.
4.1.2 Empirical Observations
OOD Accuracy. We conduct experiments for combinations of bmin E [3, 15] and fmax E [3, 15) such that mmin < lmax, and evaluate all models on 1024 OOD test samples. The test accura- cies (with different trigger size Nurg) are presented in Figure
OOD accuracy tends to increase as fax increases (with mii fixed). This suggests that a greater diversity in the training data biases the model towards the induction head.
Comparing the left and right figures, OOD generalization is achieved shifts rightward, suggesting an increased difficulty of induction head learning with larger Nug. as predicted by Theorem 6.
that the trigger size increases, the region where
Error Visualization. Our theory predicts two characteristic error modes:
Pseudo trigger position. non-periodic OOD evaluation data with & + 82 and E * lz. let (fi + f2) The positional shorteut maps the second-trigger position f + la the pseudo output position + 2. Accordingly, we measure the fraction of instances where the model outputs zح-2 and report this frequency as the pseudo accuracy rate.
position. positonnn Since the leftmost trigger in the pretraining data typically provides the strongest positional signal, the model may output m-a independent of the second trigger position. This error mode is especially likely when Nurg is small. We record its frequency the lefimost rate.
Figure in Appendix illustrates the existence of these positional shortcuts. We observe that the error rate due to the pseudo-trigge mechanism is higher near the diagonal, and both errors decline as l fmax increases.
4.2 Experiments for Practical Settings
Next we examine whether a similar transition from positional shortcut to induction head occurs in more standard gradient-based pretraining beyond our theoretical simplification. We consider three- layer transformer architecture with separated key-query matrices, MI.ss, and residual connections, where all parameters are learned jointly using the AdamW optimizer [KB 14, LH19]. The dataset

---

is generated in the same way in Section 4.1: we set N = 32 and Nugg 1, and varied fmin .,.... 20, max 4, 8, ..,............................. 40. More experimental details can be found Appendix E.
000 Arcuracy (hull thiiiig!
Peeuds Accuracy Sate Thaiiin!
Leitmot Rate (hut thaining)
30
(a) OOD accuracy
(b) pseudo accuracy error rate
(c) leftmost error rate
Figure 4: Accuracy map over mmin (vertical), kmax (horizontal) for a 3-layer transformer trained with AdamW.
Empirical Observations. Figure 4 shows the OOD accuracy, pseudo accuracy rate, and leftmost rate, following the same setup as in Section 4.1. Note that as the diversity of the pretraining dis- tribution increases, the OOD generalization accuracy improves and the errors due to the positional shortcut decrease this is consistent with our theoretical prediction in Section 3. We also observe that the transition point is less sharp compared to our theoretical setting.
Conclusion
In this work, using simplified trigger-output task, we developed theoretical analysis showing that gradient-based training implicitly selects between two distinct mechanisms with different out- of-distribution generalization properties- an induction head or a positional shortcut. We introduced the max-sum ratio as key quantity governing this selection. Our results demonstrate that the statistical structure of pretraining data critically shapes the algorithms internalized by transformers, offering quantitative insights into steering learning via data design.
We conclude with several directions for future work. First, beyond absolute positional embed- dings, important to characterize which positional shortcuts can arise under relative position embeddings and related variants. Second, while our analysis centers single-layer architecture, a two-layer model naturally delegates retrieval the first layer (recovering the token correspond- ing to 6z, ,): analyzing the coupled dynamics that emerge from this decomposition is an intriguing next step. Finally, developing methods to analyze and quantify richer classes of algorithmic biases beyond the induction-shorten dichotomy would deepen our understanding of how pretraining distributions induce specific computational circuits.
Acknowledgements
RK was partially supported by JST CREST (JPMJCR2115). NN was partially supported by JST ACT-X (JPMJAX24CK) and JST BOOST (JPMIBS2418). TS was partially supported by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2015). This research is supported by the Na- tional Research Foundation, Singapore, Infocomm Media Development Authority under its Trust Tech Funding Initiative, and the Ministry of Digital Development and Information under the AI Visiting Professorship Programme (award number AIVP-2024-004). Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore, Infocomm Media Development Au- thority, and the Ministry of Digital Development and Information.
References
[AI125] Zeyuan Allen-Zhu. Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers. SSRN Electronic Journal, May 2025.
[AWKA24] Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learn- ing: Architectures and algorithms. arXiv preprint arXiv:2401. 12973, 2024.
10

---

[BCB* 23] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. In Advances Neural Information Processing Systems, 2023.
[BE 22] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-211 257644 2022.
[BES 22] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptoties of feature learning: How one gradient step improves the representation. In Advances in Neural Information Processing Systems 35, 2022.
[BMR=20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33, 2020.
[CBK724] Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning in solvable model of dot-product attention. Advances in Neural Information Processing Systems, 37:36342-36389, 2024.
[CSWY24] 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 (NeurIPS), 2024.
[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pages 5413- 5452. PMLR, 2022.
[ENO 21] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021.
[ETE*24] Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. In Advances in Neural Information Processing Systems, 2024.
[GB10] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feed- forward neural networks. In Yee Whye Teh and Mike Titterington, editors, Proceed- ings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 249-256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR.
[GD23] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0075: 2023.
[GJB+25] Noah Golowich, Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. The role of sparsity for length generalization in transformers. arXiv preprint arXiv:2502.16792. 2025.
[GJM+20] 20] Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):665-6)) 673, 2020.
[GTLV22) Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can trans- formers learn in-context? a case study of simple function classes. Advances in neural information processing systems, 35:30583-30598 2022.
[3DD+ Samy Jelassi, Stéphane d' Ascoli, Carles and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. / 15400, 2023.
Yuhuai Wu, Yuanzhi Li,

---

[KB14] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.9980 2014.
[LAG 22] Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749, 2022.
[LH19] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Interna- tional Conference on Learning Representations, 2019.
[LLR23] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic struc- ture: Towards a mechanistic understanding. In International Conference on Machine Learning, 2023.
[LLZV*25] Yue M Lu, Mary Letey, Jacob A Peratone-Veth. Anindita Maiti, and Cengiz Pehle- van. Asymptotic theory of in-context learning by linear attention. Proceedings of the National Academy of Sciences, 122(28)e2202259912, 2025.
[LZHO25] Mingchen Li, Xuechen Zhang, Yixiao Huang, and Samet Oymak. On the power of nnnvllttion-uugmentee transformer. In Proceedings of the AAAI Conference Arti- ficial Intelligence, 2025.
[Man79] O.L. Mangasarian. Uniqueness of solution in linear programming. Linear Algebra and its Applications, 25:151 1979.
[MLH* 22] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Ha- jishirzi, and Luke Zetllemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202. 12837, 2022.
[MPL19] R Thomas McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong rea- sons: Diagnosing syntactic heuristics natural language inference. arXiv preprint arXiv: 2019.
[NDL24] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal struc- ture with gradient descent. In International Conference on Machine Learning, 2024.
[NSO*22 25] Naoki Nishikawa, Yujin Song, Kazusato Oko, Denny Wu, and Taiji Suzuki. Nonlinear transformers can perform inference-timg feature learning. In Forty-second Interna- tional Conference on Machine Learning, 2025.
[OEN 22] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly. Dawn Drain, Deep Ganguli, Zac Haddield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
[ORST23] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thram- poulidis. On the role of attention in prompt-tuning. In International Conference on Machine Learning, 2023.
[OSSW24] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Pretrained transformer efficiently learns low-dimensional target functions in-context. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, Tomczak, and C. Zhang, editors, Ad- vances in Neural Information Processing Systems, volume 37, pages 77316-77365. Curran Associates, Inc., 2024.
[PAA 22] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305. 13048, 2023.
[Red24] Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in- context classification task. In International Conference on Learning Representations (ICLR), 2024.
12

---

[RLIGS22] Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint urXi:2220..07206,, 2022.
[RPCG23] Allan Raventós. Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diver- sity and the emergence of non-bayesian in-context learning for regression. Advances in neural information processing systems, 36: 14228-14246 2023.
[SHT24] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail to solve the induction heads task. arXiv preprint arXiv:2408. 14332, 2024.
[VSP*17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett, editors, Advances in Neural Information Prucessing Systems, volume 30. Curran Associates, Inc., 2017.
19] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyz- ing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint MXXix:190500915 2019.
[WNB+25 25] Zixuan Wang, Eshaan Nichani, Alberto Bietti, Alex Damian, Daniel Hsu, Jason D Lee, and Denny Wu. Learning compositional functions with transformers from easy-to-hard data. arXiv preprint wrXiv:2505.2368: 2025.
[WS24] Shuo Wang and Issei Sato. Understanding knowledge hijack mechanism in in-context learning through associative memory, 2024.
[ZBB+22] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04300 04301, 2022.
[3BI++23] Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization. arXiv preprint 1eni:2311..0022, 2023.
13

---

A Preliminaries
Throughout the paper, O(-) notation is taken with respect to the vocabulary size N and we assume the following scaling:
Assumption 1 (Scaling between problem parameters). Number of triggers Norg satisfies Norg o(N1/3). Context length T' satisfies T o(N/Ner) trg) almost surely.
We also assume the following: Assumption 2.
We assume 2 4 almost surely at pretraining and L poly
I(A) denotes the indicator function of the event A, that is, I(AA 1 if A holds, and 0 otherwise. For a I b, lasb denotes the vector whose a, a+++a++ 1.... ......,be b-th entries are one and others are zero. For matrix A R=*x, we use the notation A(;, [J to denote the submatrix consisting of all rows and columns from index / to J. For an arbitrary matrix A, (x(A) is the i-th largest eigenvalue of A.
B Infinite Sample Analysis
In this section, analyze Algorithm with infinite sample size. Recall that defined
JTF 145 WKQ, Wv) Wy X1x Softmax
X la W WKQ R
and next token prediction loss
C ((
WKQ, Wv) Cross Softmax(frr(XX (X 1ero:
WKQ. Wv))).
For
simplicity, [ (X(1:);
we
denote
the
population
loss
L(WRQ, Wv)
Ey
WKQ. Wvl]]-
B.1 Population Gradient of Wv
Note that Softmax( (X F.T W kQ [1/T,..., 1/T and Softmax(frr(X WKQ. Wv))
[I/,...., 1/N]T is satisfied at initialization, for any X1a- From [BCB+23, Lemma the popu- lation loss can be calculated as
VW, L(WKQ. KQ. Wv)
1
exET
E E(x
Σ
exET
Σ
(3T = 33)+1==)))
T
N=N
+1
1
1
Σ ExET
Σ E[xx]T
EgE
Σ B|x ( =
B| 19T
T
Nug
T
N=Nuxx
We conduct block-wise calculation for the population gradient: let
Wi W
W (2)
W( (3)
and let
Wi = W (1)
W (2)
W (3)
be Wy after one GD step. i.e., Wi -Πy Ww [WKQQ Wv), where w E RND, w ( NxN In this section we show the following, using the rescaling A = Niiv for notation simplicity.
Lemma 9, If we use stepsize Nivv for Wv. then it holds that
-vniv aviv Nux N N-Nur
(I: [Nugl), (I: & [Nugl),
(ex, W
P.) =
(B.1)
14

---

21iv E I
(1 E (Nergg)
1/v N -
1 1-23/T
(j [NurR & [NurR]
(e,, W. ex)
21/v
N- N-Nur
E $1
(1 & [Nug), hs E (Nurg)
(B.2)
'/
Nerg +
INNCN 1) Non
Nov
(j (Nerg))
1/v
Narg 21T (N22
( + k; j, & (Nerg))
(N (N-Nurx)
and
1/v B[T BIN-
(1 k E [Niræl)
1iv N Nun
2E(TT 4
(j E [Nirgl, k & [Ntrg])
(ej, W. (3) ex)
11V N N Ner
EIT
(i E [Ntrel k E [Ntra])
TT
Niv + El(T
11 KNIN 1) (N+2))
(N- Nire)
(j k & [Nuug]
Nug E 1NN
(j + k, j. k & [Nurg):
(NV -Nrr)
where we defined 04 (EE [ (1-111111) 1:T].Pt): specifically, at 1r S TH S Er is satisfied.
Proof. For the first block, we have the following evaluation of the gradient of population loss L:
way
L(WKQ, KQ
Wv)
N
T
N
T
1
1
ET
E
[P]]
ExEr
E
pp
2T
=
K]]
N
T
N Nor k=Nurt1
T
k=1
t=1
t=1
N
T
N
T
Er
T-1
pt
exET
7
Σ
N
N
Nirg k=Nugt1
k=1
t=1
This immediately yields (B.1).
Now for the second block, we have
L(WKL) KQ. Wv)
T
N
T
Ex
T
Σ
Elea.]
TT
NerR 1
Σ
ENET
T
ele, 2T = K]T
where the first term is evaluated as
T
N E3
Nar
1 1Nart11N1 )
Σ ETET
Ele.
ET
1
2e +
T-
T
N
Nurg
T
N W-Nrg
k=1
w=1
and similarly for the second term
Σ
egET
Elez, |2T+1 =kj7
N Nur Ne
N
Nus
T-3
NN + +ex) e )
Σ
ek
ET
2e +
N - Nug
Nur
T
N - NNrr
KNN
w=1
Putting everything together yields (B.2).
The third block can be computed in the same fashion
L(WKQ. Wv)
(
N
T
N
T
e
ET
E e
exET
Ble 127+1 K]]
N
T
N
NerR
T
N=Nutt
15

---

where we have the evaluations
1
eT
Elez. 11T
N
Ntrs
T --2
ET
Ew
+
1 Nux + I:N
(: eg 0)
N
Norg
T
N Nirg
k=
1=1
and
1
enLT
e 2T+1 Kij
Nux No
T
N
NicK
1
-
NN + tet) )
Ck
Cw +
N - Nug N
Nug
T
N Nurg
w=1
B.2
Population Gradient of W KQ
B.2.1
Preparations
denote transformer's predicted probability of token & given an input sequence after one-step GD Wy
p(k|2) = Softmax(frr (f WKQ. Wy)).
We can approximate p(k|z) by considering sufficiently small rjy: The following corollary is ob- tained by Lemma 9.
Corollary 10. If set riv & 1/N, then |//N - p(k]=) = 0 (E (ETTTT11//' uniformly holds for any k and
Proof. From Lemma 9, it holds that [(STF (X 1:7; WKQ. Wx))al E O(B + 1/N2) ((E/)-)//N (:. Assumption 1) for all k. If 0s O(E(T-)//N for all E [N] where RN- then it holds that
((fr((XIr)
1 N exp(O(((T-~))N))
exp( ( Softmax(w), S
From Taylor's s theorem we obtain the assertion. Following [BCB+23 Lemma 4], starting from WKQ = 0,
V Wix (WxQ,, W
E
(Ø(k))) 1|2T-1
V (ex, X x X Ii WKQNT))
where
T e W
Vwk
(.Sotmaa((x].. WKQ
T
c
T
Wia; V V X WKQEF),
t=1
Σ e Wyx, (z. #1TT)TT
for uT.T 7 -1=1 Ct. Hence the population gradient simplifies to, assuming πñv S 1/N,
WWweC(Wre Wi)
16

---

N
1
1
BT
Ex
{ Wix: (æx - 21:T)011
T
N
Br
Σ
p k)Ex lee Wi (an $1.7 ) | = ki]
T
Nir++1
+1
t=1
N
T
ET
Ex
[ex
Witx:
(at $1.r) )ao
N
T
k=
1=1
1
ET
Ex lex Wi'z (xx, 2(.))0/22T11 = k]]
N
Nux No
t=1
(D(klz) Nv E WVx (xr, B1T)3T
Ex
(+1)
Note that each entry of e Wi (æ., +1 is of O from Lemma 9, and lp(k|=)
+
E from Corollary 10. Therefore, using [B[ab] BB we conclude that
is O (nv ///N) E O~ (ong
Moreover, have
1
ET
Ex
e
Wi (ær - x11))01]
N
T
k=
1=1
N
1
Σ
Er
Ex [e] Wixx (xe $1.1 [
N
Nug kmm Nux++
T
1=1
Nes
T
+ N NErg
Er
Ex lej Wirr (ax; 21.T)=77]
Nirg
N
T
Er
Ex Wi'x (x:: Zi.T =11))
N(N Nug)
T
t=1
Here, again using the fact that each entry of e/_ Wi (n, - ).t)) is of O(jiR((T-)), the second and third terms can be bounded by O0 (iv Norg/N). In conclusion, obtain
V Wike EWKQ WV)
N
T
Er
Ex [e Wi (xt - F )
N Nug km Nuxt1
T
f=1
N
T
1
Er
Ex [et Wi (x. - # 2 -k]
N Nurg k= A=Nuxt1
T
1m1
+ O. (ivVET '1 Nirgg//)
B.2.2
Detailed Calculations
Let
A (k, T) T Σ
Ex [ej Wia, (z, - -1.7)271
Ex [er Wi (x, - 2..T Tx |2741 k))
T t=1
17

---

for Nug + 1 S k
N. Then, it holds that VWko WV) N-Norg Lk= Nug BT| (k: [)] - 7/ Ti So (Nug/N).
For the first and second (block) columns, we have the following lemma:
Lemma 11. Let W KQ niQ V VWkQ L(WKo, W Then, it holds that
W KQ l 1: L +
Mixk
P442 ((())
PetsBe [T(E)=
E p(e) T(e)-)
proo e
IJy TKQ Nug
0 e (T(E")-1
w=l
-T(
2e [1()
(T()) e
+ ( (N N-
Here ONNoT -) denotes matrix whose entries are all of O(Ner: N and p()) is the probability of drawing l pretraining data.
Proof. [Proof of Lemma 11] Suppose Wy is obtained by Lemma 9. Note that
A (k, T[:;; L]
Σ Ex cx Wy - (xt B1:T)] Ex e Wv (x) T1rr)le441 k) k)
New
Σ t=1
Ex [xx 1. ØT w] Ex [(xe= #1.T)a, |2T+1 NT Ør = w/
Nerg
m=
Wy TeepT and similarly
T),, + L N]
Nux
T
Ex
[
pr w] Ex [(x= xx 12141 k. æ.r = w]
Nug w==1
T
- Wi
T e
Now let us consider the difference
Ex (xx: - xx [eT = w] Ex [(x: 21.))= 12r+1 = k or w)
T
x a T lo w| Ex a |2T+1 k xr =vrl)
T
Ex [
127 pr
w) Ex [ 2744 k x w|
Ex 1 WT w] - Ex 20 |2T41 ær w)
Σ
E a ØT = m0 -Ex a
12T41 Ir, 2 m/
E
18

---

Then it suffices to calculate the vector d(t,t f k, w) M(t,t', k. w
= (B|zz poT = w| - B 2T41 = k, ær = w)) Wie
for each t, and w. Recall
Ners
N -Nrg
du Nerg
N Narg
[T '1 1 Nerc
N Nick
BTT + O N )) R -N ET 10 (Nerg trg N 2- O(Nerg N 222
W c rñv
ET
E O N 2))
N N B "[INN N11
ET
+ O(Nug N () =-)
E
T 10 1N O -N -N- 22)
[[
O(Nix N 1
For preparation, we define some vectors: let
+ N +
1
a<k:
0, 0.0....... N Nug
N Nug
N
Nug
N Nirg
N
Nug
Nerg
zeros
k-th
be an N-dimensional vector for Norg
+ 1 < k < N and
β = 0.0..... 0
E ER..
N Nugg
N Nerg
Nus
7eros
First, if then (t, t', k, w) zero unless + ort= + 3, as 2141 is independent of 2, (i E [T), + f + 2) and only 2442 and 2443 include the information of 2242. For each case, we have
OLXL a(k)p
Pk)20(()T diag(a(k) (xk))
Olxn (k)e
M(E + k,
ONXL
Nxn
and
OLXL
Pk+ OLXx
+ 3,
ONXL
ONNN
Ba(k) diag(o(())) (k))
akk)
(kk)
then we obtain
P1+2E E T-1
Nug
Nur N2
dll+ 2,6 2, 1/v
+ 1 [E
e. E
and
PF++ - [[ 0 ET
Nur
Nerg N2
d(l + 3, k,k,
+ +0 (B
+
N
For the case * t', deal with the following three cases:
19

---

(i (((,t)) (t, - (f+ then we have
M 2, k, M( 3, 2, k, w)
O1xl
P
Pe+20((k] (iag(a(k))
akk)p
a(k) + (a(k)
aak)pe42
diag(a(k))
akk)e + T
and
d(e+ + 2, + 3, w) + d( 3, + k, w)
-P442 P:+3B[T
N Vug
Nug ). N
=Ijv
C2
+ 00 O
-T
C. E[T
N
(ii Fort x 2,6 2
have
0 = a(k)p?
(ka (k)
O a
(t)
0
Ba(())
where 7(1) = p: if + ort and (1) = 33 otherwise. To summarize,
P. e., E
Nur N
S(t,e+2,k,u) ( rÿv
+ + 20 BB
N
fft= and
Pr PTT
Nurg N
Nug
( 1/V
0
+ Jÿv 000 (E[
N2
0
otherwise.
(iii) For l + 2, + we have
0
o
Pia(k)
M
0
O
(()a(k)
aak))(!)T
Ba((k) aKk)
and we obtain
P
Nerg
Nerg N2
d(t, + rjv
O
N
ift 1 T and
pr E
Nus N
Nug N2
d(1, k w)) - + 1jv
0
7jv 000 (E
0 AKk,
otherwise.
Now we are ready to calculate
N Nus Ek=Nax
ΣN
N] as
A A(k.T);,1:Lt
N Nug km Nerg+1
Nurg
N
P422E
Pe+s - E
n/v
11
+
Ntrg w=l
N
Nug
T= N=Nret +1
T
T2
ET
-e
P++2 111 De+3E(17 E
1 1.mh (642,843) E
2
e ETE(TT1
2e, 12
Ipf e
7
72
eM e. 11
0
20

---

+ TVv DivE[T JJO(ee 2 + 1iv O(Nur 2
Nirs
(T Pe+3
niV
(p) ew]
MIR w=1
T
Cu R. T-
27iv 7
11:
2e
ppt Td
+ BiVBTT (Nerg N ! 1iv O0 (Norg N 3),
which concludes the proof together with B 2 N-1 from Assumption 1.
We also bound the last column: Lemma 12. It holds that
Σ
(k: T) L +11
N Nerg k=Norg +1 (T(E) (Nrgg* NN
Proof. Note that
(k, [ +
Ex
6k
Wi';
(xt 31 )1r))r
Ex e Wix; (at Zir)e 1x11 1
1 N
T
Ex e Wax (z, 01.T)22-11 =1]
N Nirg 1
T
Ex c Wi': (a, - R.r) = 2T =
Then the assertion immediately follows from the fact that each entry of
Ex [ Wi ac (x, (1-r) 1] Ex [e] Wi m: (æ. T1.r (1:r)22r-1
2,+1 =k]
upper bounded by " up to constant, from Lemma
B.3 Max-sum Ratio and Algorithm Selection
Now we are ready establish analysis on transformer's algorithm selection based on max sum ratio. From Lemmas 11 and 12, it holds that
W
Niw
(T 2T ) pet3)
1.T
i Nurg
27
2
[p.
e
o
T-l r-lew ew
w=l
+ * N-1)
where fj E
Assume that a test sequence = [21. 21 is made from subtext lengths (41, Eg) (hence T 4 + Ez 3). Now let yA P(e = and g" P(2e + 3 T** respectively (probability is defined by pretraining distribution). If the trigger Øт satisfies IT w", it holds that
NAs
((T*)-1 + 2(7*) (p: +2 p
I 1ir-
WK =
-2
2
N
trg
(T - c..
===1
21

---

(T
2T
(ptt2 + Pe+s)
1 1:f
1
E
0
-7
27
w
Nirg
T ew- Cw
+ O (i/-ug _N
where 6* = (T 3 if such f is not integer, then do not define f (in such case q" holds and we don't need to define such a quantity).
Hence, for any we can calculate the attention logit as
St = xt WiQ
=1 9" Nerg
((T))) ((T))- +)]
Nerg
+ (PP)) = w)) - TT)) 21((z = =w))
+ 1) Nirg
(T(L- - 2) + 27(1 2) 2 9) + (T(1 3)- + 27(1 - 3) 3)---3
+ B[T [][[t[][[tt11_] = w²) - 2 I[( S T) 4ET 2 [[( = w²))
+ Oz (ūj.err N 1-
(B.3)
We begin with showing that if max-sume ratio is not sufficiently large, we can construct an OOD test sequence such that transformer mistakenly use the positional shorteut:
Lemma 13. If holds that
max, qet Eeyet-I
e(Nurg)
where
there exists an OOD test sequence such that the pretrained transformer via Algorithm / fails to generalize.
Proof. Assume that
2"
U. U. U,..., U. w", u. U...., u, w., v).
and 7 = f7 + + 20 + where lr = arg max; (())- Furthermore, we assume
(1 1,8, L- 1,8 +
(B.4)
Since we have Assumption 2. there exists ej 2 such that (B.4) holds. We show the following sub-lemma:
Lemma 14. There exists e1 (Nug) = ((Nigg) such that if
max; gel
2 61 (Nug);
Σ3 Eggee-1
then
1
811+1 S 81
42-
Proof. Here we show 28 S 5 - other properties can be deduced in the same vain.
Note that, from (B.3),
82--2 5 7 9" Nara
-((rr)-'1(w = w²)l
N w=1
ñ NerR
((20 + +3) 31 2(207 + 33 +-3--210 2)g; + ((26) + 1) 2(26) 1) - -)
22

---

+ 1[ + Occ (7.Nug _
$ 119" (T*) 1 6îj4"
(T"))
+ (f.Veeg N
Nug
Norg
Nug
On the other hand, it holds that
Sg* +-
1" Narx
2 1 Nirg
+2( - +2
21 = ==))]
ñ Nerg
[112f + -33-1+ 2(2 + +3) ((28' +1- 2(2 + +1)-))ar-1 9---1 6
+ 000 (T.Verr -N-1)
2 7/9" (T**)-1 --- 6îj4" (T**
E[T "EETTTS + O0 (ñ.Nrr N N
7
Nurg
Note that T-2 v
1 T- holds from Assumption Therefore,
10
1
13 5
BTT 7ńja" (T*)-1 1 0 (ñ.Ver -N
5 D/ (T-)-1-
52 -
10 Nug
Nug
Together with Assumption 1, if the max-sum ratio is ((N we obtain 2817+2 S 86 as desired.
Since now we have Lemma 14, when
N2
Nirg
îл 2 Clog N
2 Cllog N.
(B.5)
Na
E[T-112
for sufficiently large C we obtain st/ +2 I exp(-Clog N) N-C where G 1, 2 + and T This immediately implies that XI-T- Softmax(X, XT WKQ = -, agTz + +00 (N-C") for sufficiently large C' where Ex4r Og N-- Therefore, we get
Sortmaa X Ir- WKQUT.) = (1 - N- eu e
+ N N-C C'
From the structure of wV (Lemma 9), observe that the predicted logit ( 1 W WKQZT.) has a peak the token u, meaning that OOD gener-
alization fails.
Remark 3. Here we worked on the ratio between qr- T(E)- = max; + 3) 40 and
B[T(()-4) 1010)-1=>0 E, -, (26+3)-1 We can immediately show (using Assumption 2) 3 maxx(()) 4 S
ma ++) S maxe(e)- and 3 Ee(e)= ge S Ee(2l + 3) S 2 Ee) qe. then
we do not distinguish these two definitions of max-sum ratio.
Similarly we can show the following upper bound:
Lemma 15. it holds that
max, ge--1
Eeget- e(Nuug
where e(Ntrg) ((NTT)), the pretrained transformer via Algorithum can generalize O0D.
Proof. In the same vain the proof of the lower bound, it suffices show $ 2 2s, for any t + E + 2, going the other way around Lemma 13.
First we have
Nerg
9"
82 2 1 N
+2
trg
w
23

---

ñ [((20] + +3 2(26 + -
+ Nurg
Dge; + ((28) + +)= + 2222 -
"
n0 (7/Nrrg .N-1)
f) - - E 2n max qeT) Oz (ñNiig NN
2 Nug
Nurg
For all 1 fj + 2, observe
Ntre
9
17
81 S 17
13(T") .2] 2]
qc-2 3 - 3-- 4x-33 + 000 (T.Nerr N N
Nerg w=1
Ntrg
S max ((T(()) +6 " max (()) + +0x00ii N=
Norg
Therefore, we obtain 22 2 8t as desired, if max-sum ratio is (
B.4 Proof of Theorem 5
Theorem 5 is directly obtained by combining Lemmas 13 and 15: only remains to adjust the stepsize.
From Conorally 10 we need ijv = nv/NN $ 1/N, and from (B.S) we need ijv 2 log N
Therefore, suffices to set
N
T S and лил 2
log N.
N3
trg
C Finite Sample Analysis
Now we turn to make analysis for finite sample size setting.
Proof Sketch. We explain how to evaluate the concentration of the empirical gradient Vw, ((fTF). Concentration for V Wwka ((frF) can obtained similarly.
Let {a
(0) T(O.) 1 i-1
My
be My i.i.d. sample sequences, and {{ (He.a) } ) ars" (f Sel i, be M i.i.d. sub-
samples conditioned on and 2T+1 k. Note that Σe, M My. The empirical gradient (wrf(frr) is expressed as
Vw Vw [(frr) = BAND, Bx-Uni([K]] Acce] (similar terms omitted)
(3.,)) (de, 1..... * with the sample size M We focus bounding the first term in this sketch. The gap between the empirical and population gradients is bounded as
where A the empirical average of Ei= lxz,
Vv, ((fTF) - Vw, C(fTF)ll
E. Aet - E[A,,
Ad.x-E(AA.. e, [I]]|, + LE (Bx Ex) EE)
(1( A1||, +
E. BL (E(AcH
The first term is bounded by the matrix Hoeffding's inequality using || 22. I1 1 and M ~ N My for each pair (e, k). Note that Ee / emerges in this bound because | Ej Ac 1i]l22 ~ qe 1/2 and E. ~ Ee qe- The second and third terms can also be bounded by Ae,ell < and using the standard Hoeffding's 's inequality.
24

---

C.1 Value Matrix
We first establish an upper-bound for the difference between empirical and population gradient with respect Wy.
Lemma 16. Let A: 11:2
1I,
e e
Then we have
A,A,
E.
0
A, A
where (x(E,) S N f 1 E,) 0f 2 3 and Then, we have
EEll2 N. Similarly, let Bi ei,,z
B B
0
3 E
B
Bu.k
where |2/ll2 1.
Proof. We provide the proof for A.. Proof for B, be derived in the same vain.
A,A, 1,:2 x 3
3 2 (a)
where d1 (2((1))) 3N, (ZE 0 for i 2 2.
A A. 2,1 II Noæ 3 =x/
where (1 = 3N, X ( (b) = 0 for 1222 2. Using
rank(A.A) rank(A) A₂) $
and
A O
A. A 1)
0
A
m A.).
we obtain the conclusion.
Lemma 17. Let {a)))) [=1 i=1
(1) To
Mv
{{
(tr)) 1== ir=1
MT
be i.i.d. sample sequences,
be conditionally
i..d sub- sab-samples with fixed T. and {{ (in 1 (((ttw)) Sel 17 be conditionally i.i.d. sub-samples with
fixed T and 2T+1 = k. Note that ΣT MV = ETM' = My. Then, with probability at least 1 - -(c),
My
Amax
Vwy
L(X (1)
1:7 WKQ. Wv) - Vw, LWWk, Wiv)
Mv
N Nog (NL(N + L)c-1)
Σ v
Mv
Proof. The empirical gradient is, noting that = 28 3 5,
T
Vw, ((FTF)
ETT
B(ac]]
N
T
1=1
Σ
1
B[x.I[ 1 k1]
T
Nerg++11
Mf
1
Σ
(tr) T
My
TN
T=l
25

---

N
M
T
Σ
(TT,))
ex
Mv
My
My
T
T=5
Neg +1
ir,
Let
I - E
A (i))
(ir)
Σ
11:
(ir)
(tr)
IT
and
B(o,,e)
Σ
(TT,))
E
Σ
(tr.e)
1 ki
t=1
:=1
Using Lemma 16, we can bound the operator norms as
1
A (
(i)
0
E SVN
T
O
A (ir)
A (ir)
2
and
B B(r,.))
O
21 1.
T
O
B B(r,))
B ()T
By combining the matrix Hoeffding's 's inequality and the union bound, we have
My
A (t)
log(L(( L)E
MTTN
NM
1==
and
M
Σ 1T.==1
B(ire)
Vog (NL(N I)e-1).
T(T
MTT
with probability at least
Using the union bound Mt Mvgr(e + My 1/2 11/2 Vlog(Le-1)), the matrix Hoeffding's 's inequality, standard Hoeffding's s inequality, and (T(E)
O(e).
1 1
(tr)
1l2 S 1, we obtain
L-1 M
T
T
(TT)T)
(T)
Mv T= 1r=
TN
TN
1
My MT
MY
M
TN
1
A (ir)
+
Σ
L-1
MY
My
(r(
E
TN
1
(ir)
IT)
7=
tm
M
My
MT
MT
1
TN
1
(67)
Σ
MT Mv
(T())
E
TN
1
1 1
(iT) T
IT
T
t=1
Probb(1)
SN-1/2
log (LIN+))--1) NAC Prob(T)
VMV
EL
Vlog (L(N+ + L)e
VNMy
Let B(ira) = ΣT eyat (TT.a) Using MT ~ MF(N - N1rg)-111 + 1///2/11// Vlog(Lc-1)), the ma- trix Hoeffding's inequality, standard Hoeffding's inequality, and 1 S 1. we obtain
1/27)112N1122 (log(NLe-))) (//)=l/RN1//ggg(LLL I), MT MI 1/ 9T())
Σ T=S
M
N
M
1
M
Σ
T
(tr,))
exat
My
M
M
T.
T
k= No
+1
1r,t=1
t=1
26

---

N
r -
ET
Σ
(fTa)) eno
extoy
N
Nerg
Nugt1
+1
MY
N
My
Σ
((T
Mv
My
M
k= G=Nuxx
ir,==1
Mr
N
M
B (tr.a)
y = k
Mv
MY
N
Nerg
T
k= N1
MY
(r())
N Σ
1
E(T
B (6T..e)
Kk: T
My
N Nerg
Rmm NeNergt!
&
MY
T ((T,
MY
M T
T=5
br,2=1
(T())
N-1
N Mlg (NL(N+D)e FIybIT
T-1
M
N
M
E
B B(i
1 k 7 T]]
My
MT
N
New
T=U
k= Nerg +1
Toee N Sv
S1
NMVOR(()
Mr
qr(e)
E[T ((x..a) k k,
Mv
Nirg
k=
N log (NL(N + D)e-1)
Σ Te
Mv
N5 Eies Te
Note that if Mv Z poly log N (Tr N-1, where Wi and W. are Wy after one GD step with infinite and finite sample size, respectively. The following corollary obtained by combining Lemmas and 17.
then we obtain [W Willa S
No Exes
N Excs o 2 Etes qit then holds that
Corollary 18. If we set Viv S and My 2 poly log N. ///N=((k|=]]=aa 1|ET" for any k and z, where p(k]2)) the transformer output regarding token it after pretraining Wv.
Nox
C.2 Key-Query matrix
Now let us consider the KQ-matrix with finite samples
L-1
Mrk Σc,, - Σ
Mko M
WKQ = NKQ Σ
Din
MRQ Mkq ir
MkQ Mro Mrd in.
T=0
Teb km Nu +1
+ nKQniV
E 100 (Nirg/N)
where
(G-
Cip =
Σ
1
Σ
(ir) (tr)
W. exa (tr)
N Nug
T
x x
Amm Nestt +1
N
Σ
a
(ir) (ir)
x
x
W
WTeralir (ir)
N Nurg
T
N
27

---

and
Dina= =
(
--
Σ
Wy Wverr exat (r.)))
T
ax
2t
1
Σ
( (iT,e) Zy
W exay (TT
T
zt
type
where
0k-1
wVex 1[
1
+
112NN.
ON
0k-1
1
ON-
Lemma 19. Let BLNTT z (α; Wy e ) Then, we have
BLTT ((1(2) = k] + 1 B)) '1 onjv E(T x
Therefore,
B B T. B T.
O
B Bric 3 Σ
where >x(E) OV for = 2 and {x = forr > 3.
Proof. The proof is straightforward. Note that 2; 1142N= = (1) for all b.
Lemma 20. Let Wiq be WKQ after one gradient descent step with finite sample size (Al- Mv be
gorithm 1) and Wiq be the counterpart for infinite sample size. Let {{
{(a)))
My
{{
(0((6T))
T
i..d sample sequences,
be conditionally ii.d sub-samples with fixed T, and
4=1
M
{1 (1T.K)
{
conditionally i.i.d. sub-samples with fixed T and 27-1 = k. With probabil-
ity at least O(),
Nlog (LN(L NNe-1)
Amax Wike Wix) S IKQ "KQ'iVB(T"") Σv V
MxQ
+ T(QQIVETTT1 10 NN).
Proof. We have
br
S ET
and
Dina DT tr
S EIT
DT kr
D
by Lemma 19. Then obtain the error bound as
WKQ=
W
T
28

---

Mik
Mko MrQ
Σ
Olz Biiin (T))
Σ
B B|C 12]
(r(e)
MRQ MiQ in
Mkq
MrA KQ
E
Σ (Der. B[D er [T, [T,
MKQ
Mik
T.
T
km Nux+1
KQ kink
MkQ
MKQ
Σ
E E[D Prr [T,
MKQ km Nux+1
Mak
N Nur
MMR QT(()
E.. E(DiT,al 17]
+ + (Nirg/N)
MKQ
2
Nlog (L+
ST Σ ve
B (Ntrg/N)
MKQ
in the same way as bounding Wv. We used M = BKTTMMQQ M Hoeffding's s inequalities, and ||ć,, 12, |i lla S ñv E by Lemma 19.
KQ = N N KK (matrix-)
Based on these finite sample analyses, Theorem 6 can be obtained similarly to Theorem 5.
Proof. [Proof of Theorem 6] Note that, the approximation in Section B.2. still applies, if the finite-
sample error with respect W XQ is falling into (oO(Nrr/ From Lemma 20, it suffices to
set Mxq. 2
poly log N
N² .
(E =
Together with the requirement My 2 poly log N
Eiss v in Corollary 18 we obtain the assertion.
NJ
Eies Stesvet) yet
D Proof of Proposition
We first show that
21, ............ Nirg 0, O,...,.)
(Nrer + 112
(gi
22
(Z
7
satisfies the KKT condition of the LP
minimize
El 9062
mas U' get
subject
1:8-1 S NNrg
P
Σ1 11 =1 91...., QU 20
and show its uniqueness. The KKT condition of IP is
U
@ (+e Σ
Nun No )e Mi + 0 (e [U]),
(D.1)
8=1
Aelget - No1
qe (e') = 0 (e (U]), Σ
(D.2)
8
me(-qe) = 0 (E (U]),
(D.3)
916-1 < No11 Σ
qrl 8'1 (eE (U]),
(D.4)
8=1
91 +
qu =
(D.5)
29

---

A 2 0 0 (e (t]).
(D.6)
We construct (X (eet μ {uz), 2) such that (q. X, μ, v) satisfies these conditions: by struction, (D.4) and (D.5) are already satisfied. Here, note that
{ 7
(L S Nurg),
gel
10
(e> Nurg)-
Thus, from (D.2) and (D.3) we have d = 0 (f Nurg) and 0 (e S Nurg)
Now remains to satisfy (D.1), not braking (D.6). For E [Nirg]. (D.1) reduces to the following linear equations:
N
A1
1 Fv-1
urg
tng
Iuu
Nirg
N Ning
XNux
N + Nug
Noting that the sum of the all entries in the vector obtained by evaluating left-hand side zero, we obtain
1 N
Nurg (Nug +1)
Nirg
2
We can also observe \ 2 ... ANng since the right-hand side is decreasing w.r.t the vector index.
N
No
Since
[1, 1,
belongs the right kernel of I I
can
Nirg
N
trg
shift À by this vector to ensure ANug = 1 (then X 2 0), meaning
1
1- 1-
N33 Nerg + ì
N
2
trg
2
erg
+
where Ā = =Nerg E Norg
L==1 de. We now consider (D.1) with e> Nirg Since
= v ++ l2 1-1
2 (Nirg + 1)2 Norg(Nergg +
1
1 =
N2 + Nerg
2
2Nug +
Nug 2 1,
2
2 trg
2
Nug
and then KKT condition.
determine μ satisfying (D.6). Therefore, obtained (q" A, pe, 1/) satisfies the
From [Mnn79], show the uniqueness suffices to show that for any p E RRU there exists such that even if we replace the objective function to Eieuo (f2 (f=+ep))tt q" is optimal. We can eas- ily see this by reconsidering KKT wondition-while the only effect by changing the objective is the nonnegativeness À and μa (D.6), these parameters are continuous with respect to the perturbation, and we can still ensure nonnegativeness since for we already obtained positive parameters.
E Detailed Experimental Settings and Results
E.1 Detailed Settings for Section 4.2
We introduce the detailed settings for the full-traning experiment.
Architecture.
Embedding. We use embeddings obtained by concatenating the positional embedding and the token embedding, i.e., Pi where p: and e., are one-hot vectors with ones at /-th and z,-th e.,
entries, respectively. The previous-token embedding in (2.2) is omitted.
30

---

Transformer blocks. Each layer consists of single-head attention module with separate Key- Query matrices, GeLU-based MLP, and residual connections:
3t - 3t + MLP Wy 1 Softmax X W Wox:)).
where
(a)
ML. + b2-
Three such layers are stacked, followed by a linear projection of size (N. D) that maps the final embeddings (dimension D) to the vocabulary of size N. We initialized Wx: Wq. Wv, WMLP, and WMLP MLP.2 using Xavier initialization [GB10) while biases bj and b are initialized from the zero vector. The size of WMLP,1 and WMLP.2 are (4D, D) and (D, 4D), respectively, where D = N + L is the embedding dimension. The transformed embedding at the last layer is fed into the trainable linear output layer W of size (N, D), initialized using Xavier, before softmax.
Training. Training was performed using A with both the learning rate and weight decay set to 10-2, using 32,768 training samples. We prepared 1,024 in-distribution samples drawn from the same distribution as the training data and stopped training once the accuracy exceeded 90% on these samples.
Pseudo Accuracy Rate
Leftmost Rate
16B
TLCK
SLEE
1A0%
min puttott
192%
200
IEN
10.% 10.N
160
max subtest length Limax
max subtent length Lmax
(a) pseudo accuracy rate (6))222))2+2
(b) leftmost rate (3r 42)
Figure 5: Map of types of errors due the positional shorteut. Note that both errors probabilistically coincide with the correct answer, such cases are not excluded.
31

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The paper is primarily theoretical analysis and its problem setting and con-
clusion are accurately aligned with the statement in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitation and future work are summarized in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specificattin, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [es]]
32

---

Justification:
The problem setting and assumptions are summarized in Section 2, and the all complete proofs are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
Experimental details are provided in Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data also be provided via
often one good way accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
33

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The experiments are conducted only on toy simulations.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https: //tp::/niiicccclccccec
for more details.
While we encourage the release code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //niss... "tippsccppppiii cuc/public/guti http//uddde//ee/e Sud
onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Experimental details are provided in Section 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [NA]
Justification: Our experiments conducted the qualitative tendency of each setting and thus is not aimed to report statistical significe.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
34

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The experimental details are summarized in Section 4. Since all experiments are conducted on small size synthetic data, it does not require special computational sources,
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https /eti::/nneuupp/e cc/ppblicc/thicccuuueeeeee
Answer: [Yes]
Justification:
The authors have checked that the research conforms with the NeurIPS Code Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper is primarily aimed to reveal a specific characteristic of a trans- former model with simple structure, and immediate societal impact is expected.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
35

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper is primarily thcoretical.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions to access the model or implementing safery filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
36

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswithcode com/datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
37

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Execution of the theoretical analyses and numerical experiments in this pa- per does not involve LLMs in important, original, or non-standard components. We just exploited them for auxiliary use.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https /ttp:///nn.p.p/eeeee C 220pcceL///2//LLLLL. for what should or should be described.
38


Reviews:
[{"review_id": "N8KB8wNO8j", "reviewer": "Reviewer_aZH3", "summary": "This paper investigates how the diversity of training data affects the algorithm a transformer learns in a synthetic in-context learning task. The task involves sequences with two trigger tokens; given a sequence that ends at the second trigger, the model must look back and predict the token that followed the first trigger. Here, diversity in the training set can be considered as variations in the distance and positions of the two triggers across sequences. The test set then evaluates the model’s ability to generalize to longer sequences, i.e., out-of-distribution (OOD) generalization to contexts where the triggers appear later and/or farther apart.\n\nIn a simplified theoretical setting with a one-layer transformer and fixed input embeddings, the authors show that a single step of gradient descent can lead the model to adopt one of two strategies: either the so-called induction head mechanism, which copies the token following the first trigger, or a position-based strategy that memorizes fixed post-trigger positions from the training set. Which strategy is learned depends on a *max-sum* ratio defined over the training distribution—a quantity that captures how spread out the trigger positions are.\n\nThe authors also present empirical results using a two-layer transformer trained with a standard setup to support their theory: when the training distribution over the trigger positions is narrow, the model fails to generalize. However, increasing the support size of this distribution leads to a sharp transition in OOD generalization", "questions": "1. Can you clarify the intuition behind the max-sum ratio? It is true that it captures how pointed or spread the distribution is, but it also depends on the absolute values of the positions in the support. For example, shifting the entire distribution while preserving its shape could change the max-sum value and potentially affect generalization (based on the theory), despite the diversity remaining the same. \n\n2. Regarding experiments in Fig 1:\n\n    a) In some setting with large $l_0$, the model still generalizes even with a narrow distribution (e.g., $l_0=7, k=1$). So, also related to the previous question, the model might generalize even if the data is not very diverse?\n\n    b) Since the theory is based on a single gradient step, does the length of training affect the results and the phase transition?\n\n   c) The theory suggests two mechanisms for prediction: induction head vs. position-based memorization. In the figure, you're just measuring loss to evaluate generalization. Is it possible to also check the mechanism? For example, for induction head, measuring the prediction accuracy (instead of loss) verifies if the correct copying mechanism is implemented. Is there a way to directly measure how much the model relies on position information—especially in the left side of Figure 1, where the model performs poorly?\n\n3. Minor typos: 1) $M_{QK}, M_V$ are not defined in text, 2) (line 126) $W_{KQ}\\in\\mathbb{R}^{D\\times D}$, 3) In Figure 1, using both $i$ and $l_0$ is confusing.", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "The paper introduces a neat and minimal synthetic setup for theoretically studying the impact of training data distribution. The main takeaways—that data diversity helps the model choose the right prediction mechanism, and that there is a trade-off between generalization and statistical complexity depending on the training distribution—are intuitive and nicely supported by the theory. The proof sketches are also helpful in breaking down the main argument.\n\nThe theoretical transformer model is simplified, but most of the simplifications—such as merging key-query weights and removing the FFN—are standard in similar theoretical works and are reasonable given the paper’s focus.\n\nA key limitation is that the diversity measure used in the analysis, the max-sum ratio, seems very specific to this task and not easily generalizable. It also depends on absolute position values (see Question 1), which makes its interpretation less straightforward. Please, also see the questions.", "was_revised": true, "final_justification": "My main concern, the lack of empirical verification for the proposed mechanism, was resolved by the new experiments discussed in the rebuttal. Separately, the expanded discussion of the max-sum ratio was helpful. The study’s design and theoretical message are solid, so I lean toward acceptance. I keep the rating at borderline, however, because elements such as the max-sum ratio diversity metric seem specific to the toy setup and may not generalize to other setups, even though the broader message appears extensible.", "rebuttal": "We thank the reviewer for the helpful feedback. We will correct the typos in the manuscript.  We address the comments and questions below.\n\n### Intuition behind the max-sum ratio\nAs you pointed out, the max-sum ratio is not shift-invariant, since the inverse of the context length — measured by $\\ell^{-1}$ — varies under a global shift of the distribution.  For instance, if we consider a uniform distribution over an interval of width $K$, namely $[\\ell_0, \\ell_0 + 1, \\ldots, \\ell_0 + K - 1]$, and we shift this distribution to the right (i.e., toward longer contexts) while keeping the total width $K$ fixed, the max-sum ratio decreases. As suggested by Theorem 7, a smaller max-sum ratio is favorable for OOD generalization. Therefore, shifting the distribution to the right reduces the required width $K$ for achieving OOD generalization.\n  The intuition is that the task distribution, when reweighted by the inverse of the context length (i.e., $q_\\ell$ weighted by $\\ell^{-1}$), should be approximately uniform in order for the max-sum ratio to attain its minimum.  In other words, $q_\\ell$ associated with smaller $\\ell$—which corresponds to a task where positional information can be memorized easily—should take relatively smaller values.\n\n  A more precise interpretation of the max-sum ratio, then, is that it primarily encourages diversity in the task distribution, while assigning greater weight to longer contexts where positional signals are harder to obtain. In light of your feedback, we plan to incorporate this precise implication into the paper.\n\n### Regarding experiments in Fig 1\nWe address each of your points in the following:\n\na) As mentioned above, the wider the rightward shift of the distribution over the period length $\\ell$ during pretraining, the narrower the required support of the pretraining distribution becomes for implementing an induction head. Thus, this point is consistent with the prediction made by the max-sum ratio.\n\nb) Our learning setup — based on one gradient descent step with a large learning rate — is intended to model the behavior of parameters escaping from saddle points during the initial phase of gradient-based algorithms. When using other algorithms such as one-pass SGD, we expect that the length of pretraining influences the implemented algorithm. This is analogous to the compatibility observed in some feature learning settings between the sample complexity of one-step GD and the computational complexity of one-pass SGD [1,2].\n\n[1] Abbe, Boix-Adsera, and Misiakiewicz: SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics.\n\n[2] Barak et al: Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit.\n\nc) One method to internally verify the mechanism by which each algorithm is implemented in the transformer is to examine the heatmap of the key-query matrix that governs the attention scores. Theoretically, this is formalized in Equation (3.1) (more precisely, Lemmas 10 and 11 in Appendix B), which show that the key-query matrix exhibits peaks corresponding to each algorithm (positional memorization and induction head). Furthermore, using the same data model and algorithm as in our theory, we conducted experiments demonstrating that both peaks indeed appear in the key-query matrix, and that increasing the width of the pretraining task leads to the selective attenuation of the positional memory peak.  We will include these experimental results in the revision. \n\nWe would be happy to answer any follow-up questions in the discussion period."}, {"review_id": "r5APIsfMIw", "reviewer": "Reviewer_pSCb", "summary": "This paper investigates how the diversity of pretraining data influences the mechanism selected by transformers, particularly whether they adopt robust induction heads or brittle positional memorization. The authors formalize a minimal synthetic task and provide a rigorous theoretical and empirical analysis. They define a max-sum ratio to quantify pretraining diversity, prove a phase transition in mechanism selection, and derive optimal sampling strategies for data. Controlled experiments validate the theory.", "questions": "- Can the authors clarify how the current analysis might extend to capture the full training dynamics?\n- How do the authors envision the extension of their results to more realistic pretraining settings, such as natural language corpora? Can the current theory predict induction head behavior in standard language modeling benchmarks?\n- To what extent do the authors believe the layerwise training assumption affects their main theoretical conclusions?\n- To what extent does the simplification of the quadratic sturcture in attention affect the overall theoretical conclusions?", "limitations": "See Weaknesses.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths.**\n- The paper focus on a fundamental theoretical and practical question: how pretraining data influences mechanism selection in transformers.\n- It introduces a principled criterion, the max-sum ratio,that characterizes the phase transition between positional shortcuts and induction heads.\n- Gradient-based analysis provides clear insights into how different training regimes shape model behavior.\n\n\n**Weaknesses.**\n- The theoretical analysis is incomplete; it does not capture the full training dynamics across time.\n- The synthetic trigger-output task is overly simplified and may not reflect structural properties of natural language. It remains unclear how the proposed theory and experiments relate to induction heads in real-world language modeling.\n- The assumption of layerwise training, though common in recent works [1–3], may significantly affect the validity of the theoretical conclusions.\n- The theoretical model (changing $W_K$ and $W_Q$ to $W_{KQ}$) simplifies the attention mechanism. A recent work [3] suggests that the quadratic form $W_KW_Q$ is crucial for the emergence of induction head. (Note that while the two forms are equivalent in terms of expressivity, their training dynamics are not equivalent.)\n\n\n[1] Chen et al, 2024. Unveiling induction heads: Provable training dynamics and feature learning in transformers.\n[2] Nichani et al, 2024. How transformers learn causal structure with gradient descent.\n[3] Wang et al, 2024. How Transformers Get Rich: Approximation and Dynamics Analysis.", "was_revised": true, "final_justification": "Although other concerns have been addressed in the rebuttal, two critical issues remain unresolved: (1) the theoretical analysis is incomplete; (2) the relevance to realistic settings remains unclear. I will therefore maintain my current score.", "rebuttal": "We thank the reviewer for the helpful feedback. We address the comments and questions below.\n\n## Extension to the full training dynamics and simultaneous training on the all layers\nWe remark that our learning setup — based on one gradient descent step with a large learning rate — is intended to model the behavior of parameters escaping from saddle points during the initial phase of gradient-based algorithms. This setup has been extensively studied in the neural network theory literature, and in many feature learning settings it is known that alternative gradient-based training algorithms such as one-pass SGD yield similar statistical efficiency [1,2].\n\nMoreover, the experiments described in Section 4 were conducted using a multi-step Adam optimizer, where all layers were trained jointly. This suggests that the transition from memorization to out-of-distribution generalization is not merely an artifact of the one-step gradient descent setting assumed in our theory. We will include a more thorough empirical analysis of how the phenomena we identified unfold in multi-step optimization, including an in-depth examination of inner mechanisms such as attention patterns.\n\n[1] Abbe, Boix-Adsera, and Misiakiewicz: SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics.\n\n[2] Barak et al: Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit.\n\n## Insights on the realistic NLP settings\n\nThanks for the question. Our goal is to provide a first step towards understanding how to design the most efficient data distributions for learning a range of algorithmic tasks in LLMs. While our focus is on a simple synthetic task involving copying, you could imagine scaling this up to much more involved algorithmic tasks with synthetic data, and possibly turning them into natural language reasoning tasks with the help of LLMs, something which has been recently successful (e.g., [3]). Indeed, while the copying operation is just a tiny step towards this grander goal, copying is plausibly an important building block in many more complex algorithms, and failing to have a reliable copy mechanism would easily cause the larger tasks to also fail. We hope that our precise characterization of how data affects training dynamics, and how diversity helps the model transition from relying on spurious (positional) heuristics to generalizable solutions, will help to study data efficiency of more complex tasks consisting of multiple such components.\n\n[3] Gunasekar et al.  Textbooks Are All You Need.\n\n## Merged key and query matrices\n\nIn our theoretical analysis, we apply a simplification that ties $W^K$ and $W^Q$ in our model, yet we still observe the emergence of induction heads. This suggests that a factorized form is not strictly necessary for induction heads to appear (although they may appear at different speeds in the two models due to different dynamics and implicit biases, as studied in the referenced paper “How Transformers Get Rich”). On the other hand, the experiments described in Section 4 use a setting where keys and queries are separated.  Taken together, these findings indicate that whether keys and queries are tied or separated is not essential for reproducing our theoretical insight.\n\nWe would be happy to answer any follow-up questions in the discussion period."}, {"review_id": "e1zai8JM5n", "reviewer": "Reviewer_tEKW", "summary": "The in-context learning behavior of transformer-backed large language models is known to be implemented through induction heads. This paper aims to study how the pre-training data distribution and composition influence shaping of induction heads and other attention mechanisms. The authors first examine how gradient-based training works on shallow, single-layer transformers in the infinite data regime and show that pre-training data distribution does indeed \"affect the models' learned algorithms.\" Then, the authors theoretically derive the threshold in data diversity that instigates the model to transition from shortcut learning to the development of induction heads. Lastly, the authors empirically show that the models trained on a wider variety of context length distributions show stronger OOD generalization performance than those trained on a narrow scope of data to further support their theoretical findings.", "questions": "Please refer to the weaknesses section.", "limitations": "Yes", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "> **Strengths**\n\n- Understanding how transformers develop induction heads and other attention mechanisms that eventually lead to the emergent behaviors in LLMs is an important research topic that can advance our understanding of AI.\n\n- The authors' findings regarding the relationship between data diversity and OOD generalization are technically sound and well-aligned with common beliefs in machine learning.\n\n- Albeit minimal, the authors do make an attempt at providing empirical evidence for their findings.\n\n> **Weaknesses**\n\n- I do not quite grasp the authors initial statement and assumption on the transformers' inability to perform OOD length generalization: \"despite their capacity for in-context learning, models often fail to generalize out-of-distribution.\" To the best of my understanding, in-context learning allows LLMs, transformer-backed architectures, to learn new OOD tasks through in-context description and examples (as suggested in a seminal paper \"language models are few-shot learners.\") Although LLMs do have context limit, but isn't that due to their limited context length? If so, I am not sure whether I can agree with the above statement especially in the era of LLMs with a million+ tokens worth of context length and multi-agent LLM systems that can handle multi-turn long-horizon tasks.\n\n- Although I concur with the authors on that developing theoretical frameworks to understand the inner-workings of transformers is important, I remain quite doubtful of the impact of the findings and insights provided in this paper. There are several papers that report the importance of pre-training data on LLMs' performance and investigate how to optimize the pre-training data recipe to improve their downstream performance on various tasks [1, 2, 3]. I have only cited a few papers here, but the related works sections of these cited papers discuss more works. If so, what additional insight can a theoretical analysis on a single-layer transformer contribute to the LLML community?\n\n[1] DataMan: Data Manager for Pre-training Large Language Models \n\n[2] Deciphering the Impact of Pretraining Data on Large Language Models through Machine Unlearning\n\n[3] Rewriting Pre-Training Data Boosts LLM Performance in Math and Code\n\n[4] Investigating the Pre-Training Dynamics of In-Context Learning: Task Recognition vs. Task Learning -> not exactly on data distribution but studies how pre-training affects in-context learning. \n\n- The numerical experiment section is quite hastily written, and many of the necessary details for interpreting the results are missing. In particular, the data generation process was difficult for me to follow. I think this section could offer more details, given that NeurIPS now allows up to 9 nine content pages, and the current version only has 8 (and a few more lines).\n   - According to Definition 1, $N$ denotes the vocabulary size and $L$ denotes the maximum sequence length. What vocabulary corpus did the authors use? What are the values of $N$ and $L$ for this dataset?\n   - How is the token 1 designed as the trigger token? Do we treat it as a special token like <EOS> with an XML-style tag?", "was_revised": true, "final_justification": "I have read the authors' clarification regarding my questions. My concerns are partially resolved, and thus I am raising my score to borderline reject. Given my limited expertise in this area, I kindly ask AC to weigh the other reviewer's judgment more when making the final decision.", "rebuttal": "We thank the reviewer for the helpful feedback. We address the comments and questions below.\n\n## OOD length generalization\n> I do not quite grasp the authors initial statement and assumption on the transformers' inability to…\n\nThank you for raising this point. While LLMs tend to work well even with very long contexts, they sometimes rely on heuristic or spurious features for solving certain tasks, which can manifest itself in a failure to generalize to certain OOD scenarios. These failures to generalize have been observed, e.g., in linguistic tasks (e.g. [1]), or in simple algorithmic tasks requiring length generalization (e.g. [2,3]). Generally, these tasks can be much more complicated than simple few-shot learning, and often involve the composition of small algorithmic primitives that need to be robust to changes in their inputs.\n\n[1] McCoy, Pavlick and Linzen.  Right for the Wrong Reasons: Diagnosing Syntactic Heuristics in Natural Language Inference.\n\n[2] Anil et al. Exploring Length Generalization in Large Language Models.\n\n[3] Zhou et al. What Algorithms can Transformers Learn? A Study in Length Generalization.\n\nOur setup describes possibly the simplest algorithmic scenario in a transformer that is prone to such failures, namely copying a token based on heuristic positional information instead of learning a general copy mechanism that works regardless of position. The simplicity of our task allows us to obtain a precise characterization of the phenomena at play, such as how the choice of data distribution affects gradient dynamics, and how much diversity is needed in order to properly avoid the heuristic “shortcut” solutions and obtain reliable models. Our study paves the way towards understanding more complex scenarios which may involve compositions of such small algorithmic components.\n\n## Value of theoretical works on transformers\n> Although I concur with the authors on that developing theoretical frameworks…\n\nThis is a good question. While it is clear at this point that the community has realized the importance of pre-training data and data efficiency, many important questions remain open in this area, and the potential to further improve data efficiency could unlock significant improvements in model performance, as well as cost savings in pre-training. This is even more the case when you start considering synthetic / simulated data, instead of filtering existing available data, as in some of the referenced works. An important long-term outcome that we’re aiming towards is a way to find the smallest possible dataset that would allow us to train a model that perfectly solves a range of complex algorithmic reasoning tasks. Even just formalizing this question is very difficult. Our work takes a first step by considering possibly the simplest algorithmic task (copying with induction heads), and trying to understand how the data distribution affects the learning dynamics, and to characterize how diversity leads to a transition from using spurious positional heuristics to generalizable copy mechanisms. While the task in itself is quite basic, you could imagine extensions of our work that compose many such small algorithmic components in a larger transformer in a reliable way to achieve more complex algorithmic tasks. We hope this clarifies our motivation, and the significance of this initial step towards a much grander goal.\n\n## On the numerical experiment\n\nThank you for the feedback.  We will elaborate on the experimental setting in the revision.\n\n- **Data Model.** The data used in the experiment of Figure 1 is synthetic data following the data-generating process in Definition 1. In particular, \nWe set the vocabulary size to $N = 36$ and maximum sequence length to $L = 36$.\nThe vocabulary itself consists of natural numbers $\\{1, 2, \\dots, N=36\\}$; each sample sequence is of length up to $L = 36$, such as $[10, 8, 1, 4, 5, 8, 1, \\dots]$, and the probability distribution over data sequences is defined in Definition 1.\nWe do not treat token 1 as a special token; rather, we embed it in the same manner as any other token, as described below (This is because it is required to construct the induction head without explicitly signaling that token 1 carries a different role from the other $N-1$ tokens.)\n- **Embeddings.** Each token $ z_i \\in [N]$ (where $i \\in [L]$) is represented by the concatenation of a zero vector of length $L$ and a one-hot vector of length $N$ whose $z_i$-th component is one. For positional embeddings, we use the concatenation of a one-hot vector of length $L$ (with one at the $i$-th position) and a zero vector of length $N$. See also Appendix E.1 for further details.\n- **Architecture and Training.** Based on the above, we train a two-layer architecture as described in Section 4. We perform optimization using the Adam optimizer with a learning rate of 0.0001, with early stopping every 250 steps.  We plan to integrate the explanations of architectural and training details, which are currently split between Section 4 and Appendix E.1.\n\nWe would be happy to answer any follow-up questions in the discussion period."}, {"review_id": "seMU51Qe7s", "reviewer": "Reviewer_DDFp", "summary": "This paper studies how the diversity of pretraining data influences whether a Transformer learns a generalizable algorithm (such as an induction head) or relies on a brittle, position-specific shortcut. Using a simplified trigger-output copying task, the authors analyze a single-layer Transformer trained with one step of gradient descent. A key contribution is the identification of a sharp transition in the model’s behavior, controlled by a metric they introduce—the “max-sum ratio”—which captures the diversity of the training distribution. When this ratio exceeds a certain threshold, the model shifts from positional memorization to learning an induction-like mechanism that generalizes to unseen data. The theoretical findings are further supported by experiments on a two-layer model with relative positional encoding, which exhibit the same phase transition and corresponding improvement in out-of-distribution performance.", "questions": "I really enjoyed the theoretical framing in this paper—especially how it connects data diversity to the emergence of associative recall. That said, I’m curious about how these insights might carry over to more realistic or complex settings. For example, challenges like long-context retrieval (e.g., in “Needle in a Haystack” or multi-key recall) remain open. Do you think the max-sum ratio you introduce could be used to inform how we design synthetic pretraining datasets to better support such abilities?\n\nAlso, I wonder how your theory might generalize to hybrid architectures like Transformer-Mamba. These models often show different recall failure modes, possibly due to the limited state size in the recurrent component. Do you have any thoughts on how your framework could help us even mitigate these limitations?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strength\n\nThis paper gives a sharp and intuitive take on how data diversity influences what kinds of algorithms Transformers end up learning. I found the idea of \"positional memory\" especially compelling—it neatly explains why models often struggle to generalize when the input structure shifts. The theoretical results are solid, especially the phase transition between learning a positional shortcut vs. a more general induction-based solution. I also appreciated the experimental setup: validating the theory on a stronger model with relative position encodings really drives home that the insights aren't just artifacts of toy settings.\n\n\n# Weakness\n\nI do not have major concerns about this work. My feedback consists of some minor points.\n\n1. The theoretical tools used are standard, and the gradient-based perspective has been widely adopted in the study of in-context learning. While the absolute embedding is arguably restrictive, it serves well for building intuition. From this perspective, the theoretical results are not surprising, though I appreciate the authors' effort in formalizing them.\n\n2. The provided theoretical analysis cannot subsume multiplicative positional encodings like RoPE. Although the authors argue that their findings should generalize to relative positional encoding schemes (given their setup where the relative distance is also fixed), and the experimental results appear to support this claim, this extension lacks a fundamental theoretical justification. Understanding how RoPE-equipped models might form such positional memories is a highly non-trivial research question worthy of further investigation. It would be valuable to build a dedicated theoretical framework for this case.", "was_revised": true, "final_justification": "I maintain my positive score", "rebuttal": "We thank the reviewer for the helpful feedback. We address the comments and questions below.  \n  \n### Implications of the max-sum ratio for realistic settings  \n\nWhile the max-sum ratio is derived under our simplified theoretical setup, we believe it provides a high-level guideline for designing datasets that promote better generalization.  \n\nPrimarily, the max-sum ratio encourages diversity in the task distribution, as discussed in the main text.  Further intuition behind the max-sum ratio is that the task distribution, when reweighted by the inverse of the context length (i.e., $q_\\ell$ weighted by $\\ell^{-1}$), should be approximately uniform in order for the max-sum ratio to attain its minimum.  In other words, $q_\\ell$ associated with smaller $\\ell$—which corresponds to a task where positional information can be memorized easily—should take relatively smaller values.\n Accordingly, the max-sum ratio suggests the following simple high-level principles for dataset design: (i) selecting a more diverse range of tasks, and (ii) placing greater emphasis on tasks in which positional signals are harder to obtain.  \n  \n### Architectural and Embedding Considerations    \nThank you for the suggestion.  Indeed, it is challenging to straightforwardly extend our theoretical framework to multiplicative embeddings like RoPE or to SSMs that depart from the transformer architecture with absolute embedding.   \nThat said, the high-level insights regarding dataset design discussed in the previous section—such as promoting some notion of diversity in the pretraining data distirbution and emphasizing tasks where positional signal is harder to obtain—are not specific to any particular architecture or embedding scheme. As such, we believe it is worthwhile to explore whether these dataset design principles remain effective in these more advanced settings.\n\nA step toward the RoPE setting is to incorporate an additional layer. While the analysis may become more involved, the use of a residual stream potentially allows us to extend our one-layer architecture to a multi-layer architecture equipped with relative positional embeddings [1,2].\nAlthough further extensions would be required to fully capture the multiplicative structure of RoPE, we view this as a meaningful future step in that direction.  \n\n[1] Bietti et al. Birth of a Transformer: A Memory Viewpoint.\n[2] Wang and Sato.  Rethinking Associative Memory Mechanism in Induction Head.\n\nWe would be happy to answer any follow-up questions in the discussion period."}, {"review_id": "Fz77W4awOC", "reviewer": "Reviewer_eKuv", "summary": "This paper provides a theoretical analysis of how pretraining data diversity influences whether a transformer learns a generalizable induction head or a positional shortcut. Focusing on a minimal trigger-copying task, the authors show a phase transition governed by a max-sum ratio, which quantifies data diversity. With high-diversity data (low ratio), the model learns the induction head, while low-diversity data leads to positional memorization that fails on out-of-distribution inputs. The analysis is first carried out under infinite sample size and then extended to the finite-sample regime. The latter yields sample complexity bounds and identifies an optimal data distribution that balances generalization and statistical efficiency. Synthetic experiments on a two-layer transformer validate the theoretical prediction that increasing data diversity robustly induces OOD generalization.", "questions": "1. **Impact of Multiple Trigger Tokens:** \"Your analysis assumes a single, fixed trigger token (`e_1`). What happens if the trigger is instead drawn from a set of `K` possible tokens? This extension intuitively seems to increase the learning complexity of the induction head, requiring `W_KQ` to learn a more general identity-matching function (a sub-identity matrix) rather than a single association. How would this affect two of your core results:\n\t- **The Transition Point:** Would the `R(D_l)` threshold required to learn the induction head become stricter?\n\t- **The Sample Complexity:** How would the sample complexity bound in Theorem 7 change? Would it now need to scale with the number of possible triggers `K`?\"\n\n2.  **Formalizing the Max-Sum Ratio's Role:** \"The core of your paper's theoretical results is that the max-sum ratio, `R(D_l)`, determines whether the model learns a positional shortcut or an induction head. This is intuitively explained as a competition between the two mechanisms.\n\n\tWould it be possible to formalize this intuition by directly mapping the terms of the ratio to the attention score?\n\t\n\tSpecifically, using the attention score `s_t` from the proof of Proposition 5, can you show that:\n\t- The effective magnitude of the positional shortcut component of `s_t` is proportional to the numerator of the ratio, `max_l q_l l^{-1}`?\n\t- The magnitude of the induction head component of `s_t` is proportional to the denominator, `\\sum_l q_l l^{-1}`?\n\t\n\tThis would provide a clear, formal bridge between the `R(D_l)` criterion and the final learned attention scores.\n\t\t\n3. The experiments show a sharp drop in OOD loss as K increases, but do the empirical transition points align quantitatively with your predicted $\\epsilon_1$,$ \\epsilon_2$?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n\n1.  **Theoretical Framework:** This paper addresses an interesting and timely question with clarity and rigor. The authors introduce the \"max-sum ratio\", a novel and mathematically well-defined quantity that enables rigorous study of how data diversity influences algorithm selection in transformer models when solving in-context associative recall tasks. This provides a concrete, predictive framework for analyzing training dynamics in transformers.\n2. **Analysis:** The analysis of the phase transition in both infinite and finite-sample settings is a significant theoretical contribution. \n3.  **Clear Mechanistic Story:** The manuscript is well written, and I enjoy reading it. It provides a clear, guided explanation of how positional signals \"diffuse\" with diverse data while content-based signals \"concentrate\", leading to the emergence of one mechanism over the other. This is supported by detailed proofs and validated by experiments.\n\n**Weaknesses:**\n\n1.  **Simplifying Assumptions:** The analysis relies on key simplifications.  The input embedding `x_t` is hard-coded with previous-token information (`e_{z_{t-1}}`), meaning the model is given the feature required for induction rather than learning to construct it. This sidesteps the full complexity of how induction heads form in standard two-layer models, but it is also fairly standard in related works. \n\n2.  **Focus on a Single, Special Trigger:** The analysis is restricted to a single, fixed trigger token. This simplifies the induction mechanism to a single association rather than the more general identity-matching function (`W_KQ` learning a sub-identity matrix) required for real-world induction heads that operate on arbitrary token patterns.\n\n3.  **Single-Step Gradient Analysis:** The results are derived from a single gradient step from a zero initialization. This approximation nicely captures the investigated mechanisms but may not fully represent the complex dynamics of a full training run with an adaptive optimizer.", "was_revised": true, "final_justification": "I recommend the acceptance of this paper. The discussion simply served to confirm my already positive evaluation.\n\n- Resolved Issues: The points I raised in my review were not critical flaws but rather clarification on the extensibility and formal underpinnings of the work. The authors' responses were satisfactory, providing intuition for the multi-trigger case while confirming the role of the max-sum ratio and the model's alignment with experiments.\n\n- Unresolved Issues: There are no unresolved issues.\n\n- Final Recommendation: I find the main question studied in the paper interesting, and the theoretical framework is solid. The discussion did not fundamentally change my assessment but rather reinforced it.", "rebuttal": "We thank the reviewer for the helpful feedback. We address the comments and questions below.  \n\n## Impact of multiple trigger tokens  \nExtending the analysis to the case of multiple triggers is an interesting and natural direction, and we are actively working on this extension.  \n The key idea behind our current analysis is that the relative strength of two peaks in the attention key-query matrix—namely, the positional memory and the induction head components corresponding to terms (a) and (b) in Equation (3.1)—determines which of the two algorithms is implemented by the transformer. Our results suggest that when the pretraining dataset (i.e., the positions of the triggers) is diverse, the peak associated with the positional memory becomes dispersed across multiple trigger positions, thereby making the induction head dominant.  \n  \nIn contrast, when there are multiple triggers—drawn from a set of $K$ tokens, as you suggested—the induction head signal, which depends on the semantics of the triggers, also becomes dispersed. This yields two implications:\n- **Transition Point**. Since multiple triggers make the induction head more difficult to implement, we expect that a higher level of diversity in the pretraining data is required to induce a phase transition from positional memorization to the induction head regime.   \n- **Sample Complexity**. As the induction head signal is attenuated due to the semantic dispersion across multiple triggers, the number of pretraining samples required to prevent the signal from being buried in noise—in other words, the sample complexity required for implementing the induction head—is expected to increase and dependent on $K$.\n\nWe believe these implications have the potential to make the transition point more complex but to provide a rich characterization of the competition between positional and semantic signals.  \n\n## Formalizing the Max-Sum Ratio's Role  \n\nAs you pointed out, the effective magnitude of the positional shortcut and that of the induction head component in $s_t (= x_{t}^\\top W_{KQ} x_{T})$ are proportional to the numerator and the denominator, respectively. In short, the max-sum ratio governs the relative strength of the positional shortcut and the induction head in the attention logits. In light of your comment, we will discuss this intuition in the main text.  \n\n## On the experiment  \n\nIn practice, the precise location of the transition point can depend on architectural and embedding-specific details, and thus cannot be fully predicted in a quantitative manner based solely on the max-sum ratio.  \n\nBeing said, some qualitative insights can still be obtained.   For instance, if we consider a uniform distribution over an interval of width $K$, namely $[\\ell_0, \\ell_0 + 1, \\ldots, \\ell_0 + K - 1]$, and we shift this distribution to the right (i.e., toward longer contexts) while keeping the total width $K$ fixed, the max-sum ratio decreases. As suggested by Theorem 7, a smaller max-sum ratio is favorable for OOD generalization. Therefore, shifting the distribution to the right reduces the required width $K$ for achieving OOD generalization.　This is consistent with the experimental result.  \n\n## Simplified setup  \nWe provide the following remark concerning our simplified setup:  \n- **Simplified Embedding.** Our simplified embedding scheme is intended to model the behavior in deeper architectures, where information from previous tokens flows into the residual stream [1].  While the analysis may become more involved, a simplified multi-layer architecture equipped with residual connections and relative positional embeddings [2] would reproduce our theoretical insights.  \n- **Single-step analysis.** We remark that our learning setup — based on one gradient descent step with a large learning rate — is intended to model the behavior of parameters escaping from saddle points during the initial phase of gradient-based algorithms. This setup has been extensively studied in the neural network theory literature, and in many feature learning settings it is known that alternative gradient-based training algorithms such as one-pass SGD yield similar statistical efficiency [3,4].  \n\n[1] Bietti et al. Birth of a Transformer: A Memory Viewpoint.\n[2] Wang and Sato.  Rethinking Associative Memory Mechanism in Induction Head.\n[3] Abbe, Boix-Adsera, and Misiakiewicz: SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics.\n[4] Barak et al: Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit.\n\nWe would be happy to answer any follow-up questions in the discussion period."}]

Author rebuttal:

Meta-review: *Main Contribution*  
This paper investigates the influence of pre-training data on the test-time computation of language models. Its primary contribution is to establish a phase transition between data memorization and a form of test-time generalization (commonly referred to as induction heads). This transition is governed by the “max–sum” ratio, which quantifies input diversity. Importantly, the analysis is restricted to training with a single iteration of gradient descent.  

*Strength*  
The work introduces a highly interesting theoretical question and proposes a quantitative metric for characterizing a specific phase transition in test-time computation.  

*Weakness*  
The main limitation lies in analyzing only a single gradient descent step, rather than studying convergence. This restriction stems from the complexity of the problem. Nonetheless, analyzing a single step can be seen as an important first step toward a deeper theoretical understanding of LLMs.  

*Rebuttal*  
Most reviewers were convinced by the authors’ response. However, one reviewer maintained that restricting the analysis to a single gradient descent step is a significant limitation.  

*Decision*  
In my view, the theoretical results presented in this paper are both insightful and novel, and they open a promising new direction for theoretical research on understanding LLMs.  

*AC Comment*  
There are prior studies examining the interplay between training data distribution and in-context algorithm selection. In particular, several recent works show that ill-conditioned training data can encourage the learning of preconditioned gradient descent for in-context linear regression. The paper may benefit from reviewing these related works.
Decision: Accept (spotlight)