The mechanistic basis of data dependence and abrupt learning in an in-context classification task
Authors: Gautam Reddy
Venue: ICLR 2024 oral (2024)
Keywords: in-context learning
mechanistic interpretability
language models
induction heads
Primary area: general machine learning (i.e., none of the above)
TL;DR: We characterize the loss landscape of an in-context classification task and identify the factors that lead to abrupt transitions during learning.

Abstract: Transformer models exhibit in-context learning: the ability to accurately predict the response to a novel query based on illustrative examples in the input sequence, which contrasts with traditional in-weights learning of query-output relationships. What aspects of the training data distribution and architecture favor in-context vs in-weights learning? Recent work has shown that specific distributional properties inherent in language, such as burstiness, large dictionaries and skewed rank-frequency distributions, control the trade-off or simultaneous appearance of these two forms of learning. We first show that these results are recapitulated in a minimal attention-only network trained on a simplified dataset. In-context learning (ICL) is driven by the abrupt emergence of an induction head, which subsequently competes with in-weights learning. By identifying progress measures that precede in-context learning and targeted experiments, we construct a two-parameter model of an induction head which emulates the full data distributional dependencies displayed by the attention-based network. A phenomenological model of induction head formation traces its abrupt emergence to the sequential learning of three nested logits enabled by an intrinsic curriculum. We propose that the sharp transitions in attention-based networks arise due to a specific chain of multi-layer operations necessary to achieve ICL, which is implemented by nested nonlinearities sequentially learned during training.

Published as a conference paper at ICLR 2024
THE MECHANISTIC BASIS OF DATA DEPENDENCE AND ABRUPT LEARNING IN AN IN-CONTEXT CLASSIFICA- TION TASK
Gautam Reddy
Department of Physics, Princeton University Physics & Informatics Laboratories, NTT Research Inc. Center for Brain Science, Harvard University
ABSTRACT
Transformer models exhibit in-context learning: the ability to accurately predict the response to novel query based on illustrative examples in the imput sequence. In-context learning contrasts with traditional in-weights learning of query-output relationships. What aspects of the training data distribution and architecture fa- vor in-context vs in-weights learning? Recent work has shown that specific dis- tributional properties inherent in language, such as burstiness, large dictionaries and skewed cank-frequency distributions, control the trade-off or simultaneous appearance of these two forms of learning. We first show that these results are recapitulated in a minimal attention-only network trained on a simplified dataset. In-context learning (ICL) driven by the abrupt emergence of induction head, which subsequently competes with in-weights learning. By identifying progress measures that precede in-context learning and targeted experiments, we construct a two-parameter model of an induction head which emulates the full data distribu- tional dependencies displayed by the frrention-based network. phenomenolog- ical model of induction head formation traces its abrupt emergence to the sequen- tial learning of three nested logits enabled by an intrinsic curriculum. We pro- pose that the sharp transitions in attention- chain of multi-layer operations necessary to achieve ICL, which is implemented by nested nonlinearities sequentially learned during training.
networks arise due to specific
INTRODUCTION
A striking feature of large language models is in-context learning (Brown et al. 2020 Dong et al. 2022] Garg et al. 2022, Dai al. 2022). In-context learning (ICL) is the ability to predict the response to query based illustrative examples presented in the context, without any additional weight updates. This form of learning contrasts with in-weights learning (IWL) query-respons. relationships encoded in the weights of the network. ICL emerges in transformer models (Vaswani eet al. (017)) trained on diverse set of tasks that contain a common structural element. TCL can be exploited perform zero-shot learning on novel tasks that share this structure. For example, transformer trained to solve numerous linear regression tasks learns to solve new linear regression task based on in-context examples (Garg al. 2022 Akyürek et 2022. Von Oswald al.) 2023 Ahn et al.) 2023). Specifically, given sequence of sample input-output pairs, the predictive error on target query is comparable an optimal Bayes predictor (Ahuja al., 2023]] Xie et al.) [02T]]
Li et al. 2023). This remarkable feature extends to other generative models such as hierarchical regression models that involve model selection (Bai et al.) 2023), random permutations of images
(Kirsch et al. 2022) and mixture models over sequential data (Wang al. 2023
Xie et al. 2021
Transformer models trained on language data exhibit another simple yet powerful form of in-context learning. Given sequence ... . ? for x, pairs unseen during training (for example, to- kens belonging to novel proper noun), these models learn the ability to predict y (Olsson et al.) 2022 In other words, the model leams empirical bigram statistics on-the-fly, thus displaying
greddy @princeton.edu

---

Published as a conference paper at ICLR 2024
(a)
(c)
(d)
logits (L labels)
31 41 Eg 21a
8g
2gg
Fig
Loss
3-layer MILP
labels
(b) In-weights (WY accuracy (largen not in context)
75
100
head 2
Terations (c500)
Text moc
In-contexi (IC) accuracyy (novel classes)
attention
Accuraty
head 1
0.5
In-context (IC2) eccuracy (swepped label)
0.00
100
200
Iterations (x500)
Figure 1: (a) Input sequences consist of N item-label pairs followed by a target. Items are drawn from K' classes assigned to L S K labels. At least one item belongs to the same class as the target. The network is tasked predict the label of the target. The number of classes (K), their rank-frequency distribution (a), within-class variability (s) and the number of items from a single class in an input sequence (B) parameterize the data distribution. (b) IWI. is measured using input sequences where the items' and target's classes randomly sampled. ICL is measured using items and targets from novel classes and by swapping the label of an existing class in the context. (c) Network architecture. (d) Loss and accuracy curves for six seeds (dark lines show averages over the seeds). Here, = 2, 512.
primitive form of zero-shot associative learning. Past work has shown that this computation involves an induction head (discussed in detail further below) and that minimal implementation requires
two-layer attention-onle network (Olsson et al. 2022 Bietti et al., 2024). Across networks of dif- ferent scales and task structures, the ability to perform ICL often increases abruptly during training (Olsson et al. 2022). The mechanistic basis of the abrupt transition remains unclear. Notably, this abrupt transition often preceded by the formation of induction heads intermediate layers of the network, suggesting that induction head formation may provide scaffold for the development of more complex in-context computations. Other work provides empirical evidence that ICL is the key driver behind the emergent abilities of large language models (Lu et al., 2023). Thus, elucidating the mechanisms that underpin ICL, and induction heads in particular, may provide crucial insights into the data distributional and architectural factors that lead to emergent zero-shot learning.
recent empirical study has highlighted key data distributional properties pertinent to language that promote ICL in hybrid in-context/in-weige a 12-layer transformer network trained to predict the class label a target item given a sequence of N item-label pairs in the context. The item classes are drawn from Omniglot (Lake al. 2019), a standard image-label dataset. By manipulating the distribution of classes shown during training, various data distributional properties that influence the ICL IWL trade-off were identified. This setup offers well-controlled paradigm for identifying the factors that enable attention-baase models learn in-context learning solutions without explicitly trained to do so.
classification task (Chan et al.) 2022). In this setup,
Our main contributions are as follows. We first show that the data dependencies highlighted in Chan jet al. (2022) are recapitulated in a task with simplified input statistics and a two-layer attention only network architecture. By identifying progress measures and designing careful experiments, we show that ICL driven by the abrupt formation of an induction head. We construct mini- mal two-parameter model of an induction head stacked with a deep classifier, which reproduces all data distributional dependencies and captures the dynamics learning. Finally, we develop a phe- nomenological model of an induction head's 's loss landscape. This analysis enables to trace the abrupt learning phenomenon to cliffs in the landscape created by nested nonlinearities in a multi- layer attention-based network.
2

---

Published as a conference paper at ICLR 2024
2
TASK AND NETWORK ARCHITECTURE
Task structure. The task structure is based on a common ICL formulation. The network is trained to predict the label of target By given alternating sequence of N items and N labels: Ti, f1, 222, f2...., IN, LN, Iq. (Figure la). We embed the items and labels in P + D dimensions. The first P dimensions encode positional information and the latter D dimensions encode content. Position is encoded by a one-hot hot P-dimensional vector (we use P = 65 throughout). The input se- quence occupies a random window of length 2NV between and P-1. 1. This choice of positional encoding biases the network to learn a translation-ine invariant computation.
The items are sampled from a gaussian mixture model with K classes. Each class k is defined by a D-dimensional vector l4l whose components are sampled i.i.d from normal distribution with mean zero and variance 1/D) The content of item T3, T;, is given by
14 + EN
T4
V1 + E²
where 7) is drawn from the same distribution as the p 's and E sets the within-class variability. The re-scaling with vI+ ensures that || = 1. Each class is assigned to one of L labels (L S K). The contents of the labels are drawn prior to training from the same distribution as the μa 's. Each label in an input sequence appears the same number of times as every other label in that sequence.
Importantly, at least one item the context belongs to the target's class. The network is trained to classify the target Tyy into one of the T. labels using cross-entropy loss. The network can thus achieve zero loss by either learning to classify targets from the K classes as in a standard in-weights classification task (IWL), or by learning more general in-context solution (ICL).
Parameterizing the data distribution. The data distribution is modulated by tuning various pa- rameters in addition to K and e. The burstiness B is the number of occurrences of items from particular class in an input sequence (N is a multiple of B). pø is the fraction of bursty sequences. Specifically, the burstiness is B for fraction pa of the training data. The classes (including the target) are sampled i.i.d for the remaining fraction 1 PB- The rank-frequency distribution over the classes is fkk) k a We use L 32, N 8, D 63, E 0.1, a unless otherwise specified.
Metries for tracking in-context and in-weights learning. To track IWL, we measure the prediction accuracy on input sequences. The target and item classes are sampled independently from the rank- frequency distribution used during training (Figure 1b). Since K >> N in our experiments, it is unlikely that the target's 's class appears in the context. The network therefore has to rely on IWL to correctly predict the target's 's class label.
The primary metric for tracking ICL is the prediction accuracy on input sequences. The target and items belong to novel classes (the y4 's are drawn anew). The novel classes are randomly assigned one of the existing L labels (Figure 1b). B copies of the target (within variability E) are included in the context. Since the classes are novel, the network has to rely on ICL for accurate prediction. We introduce a secondary metric for tracking ICL using input sequences where the items' labels different from those presented during training. We measure the accuracy of the network on predicting the swapped label. That is, the network has to rely on ICL rather than IWL.
Network architecture. The inputs are passed through two-layer attention-only network followed by a classifier. Each attention layer has one attention head with a causal mask. Given a sequence of inputs U1 u2 Un the outputs of the first (1;) and second (w:) layers are
02 = u4 + Vi
Σ
Ujs
wi =b; + 1/2
Pij
351
J51
where
(K,N;)T(QQuN))
(a) Pij
(3)
Ense "K
is the attention paid by query i on key j in the path layer. Qu, Ky, Vy are the query, key and value matrices, respectively. The classifier receives Wn, input.
The classifier is a lhree-layer MLP with ReLU activations and a softmax layer which predicts the probabilities of the I. labels. We use a deep classifier ensure perfect IWL is feasible. At least three
3

---

Published as a conference paper at ICLR 2024
(b)
(c)
0.0
224 acburaay
0.75
IW aosuracy
0.75
Burstiness
0.60
0.4
W
0.25
0.25
Top
2
2
2
0.
0.25
0.50
0.75
0.00
0.0
0.5
1.0
1.5
Number classes R[
variation
Rank Rann-frequeency exponent
1.0
Burstnoss B
eecuracy
acsuraa
0.79
0.75
0.50
0.50-
02
0.25
Number classes K
20
0.0
0.00
0.00
Within-slass
0:25
0.30
0.75
0.00
0
Rank frequency exponent
0.5
155
Figure 2: In-weights (top row) and in-context accuracy (bottom row) against the number of classes (K), burstiness (B), within-class variability (e) and the exponent of the rank-frequency distribution
(a). Here K = 1024, = 0, B = 1, = 0.1 except when that parameter is varied.
layers necessary achieve perfect classification accuracy for the parameter ranges considered in this paper (since K >> L). The query/key dimension and the MLP hidden layer dimension are both 128. We repeat every experiment with six seeds (with random initializations and training/test. sets). For training, we use batch size of 128 and vanilla SGD with learning rate 0.01. Figure Id shows sample loss and accuracy curves, including the measures used to track IWL and ICL.
RESULTS
Recapitulating data distributional dependencies. In Figure 2, we quantify how IWL and ICL depend on the parameters of the data distribution. The upshot is that the highly simplified input statistics and network architecture considered here reproduce the core distributional dependencies observed in past work. The results are summarized below.
Increasing the burstiness B and the number of classes K promotes ICL while decreasing IWL (Figure 2a), highlighting the trade-off between ICL and IWL. Recall that the target and item classes randomly sampled when B 0. This implies that the network solution for the corresponding K. Similarly, within-class variation (E) promotes ICL and decreases IWL (Figure 2b). We find that the network always converges an IWL solution when the fraction of bursty sequences pn < 1. This is expected as the ICL solution is not a global minimum when PB
indeed learn a perfect IWL
A striking result is that Zipfiam rank-frequency distribution (a = 1) overcomes the trade-off be- tween IWL and ICL, and promotes both forms of learning. This is recapitulated our experiments (Figure 2c). Note, however, that while the network learns the IWL solution for the most common classes, it does not learn the less frequent classes even for
Moreover, we find that the network can support both ICL and IWL simultaneously. To show this, train the network on IC sequences, where the items are drawn from novel classes randomly assigned to one of the I. labels. The parameter Pc is the fraction of the training data containing IC sequences. The remaining fraction of the training data drawn as described previously. When 0 < PC 1 and 0 < pg < 1, the network can only achieve zero loss if it learns both the in-context and in-weights solutions. Figure A.I shows that the network is capable of learning both solutions simultaneously.
One potential explanation for the results in Figure 2 and Figure A.1 is that the network independently leams the in weights and in-context solutions at different rates until it achieves zero loss. The relative rates at which the network achieves ICL and IWL will then determine the fraction of loss explained by each mechanism after convergence. The rates of ICL and IWI. depend K, e, B. Specifically,
4

---

Published as a conference paper at ICLR 2024
(a)
(b)
Before
After
1.0
0.6
0.6
Query Index
accuracy
0.4
0.4
0.5
Query
10
10
0.2
0.2
15
0.0
0.0
200 Iterations (x500)
400
Prob.
10
Key
Key Index
Figure 3: (a) IC accuracy curve (pc 0.8, B 256) shows slow learning phase followed by the abrupt transition zero loss. (b) The layer and attention maps p(1) (top matrices) and (2) (bottom vectors) before and after the abrupt transition (marked in the IC curve in panel (a)). Pq.
increasing K, decreases the rate of IWL (as the classification task is harder) whereas increasing B increases the rate of ICL (as there are more demonstrations in the context). The Zipfian case of a further highlights the dynamic balance between ICL and IWL. Frequent occurrences of common classes allow the network to learn to classify them using IWL. On the other hand. the large number rare classes promotes learning of more general in-context solution. Once the in-context solution is learned, IWL freezes as the network incurs near-zero loss on all classes. When a > 1, the tail the rank-frequency distribution falls off rapidly and the rare classes do not contribute sufficiently to the loss to promote ICL. Conversely, when < 1, the network learns the in-context mechanism if K large enough such that IWL takes longer than ICL (see Figure 2a for a = 0 and varying K).
Attention maps and progress measures. now ex- amine the dynamics of ICL. We set PC 0, 1 as the IC sequences promote rapid convergence to the in- context solution and allow for more experiments. Figure 3a shows the IC accuracy, which displays a slow learning an abrupt transition to perfect aceu-
1.0
Progress messures
0.8
0.6
phase followed racy. To investigate network behavior at the transition, we examine the attention maps (for randomly chosen input sequence) before and after the transition (Figure 3b). Be- fore the transition, the attention map of the first layer p()) shows queries paying uniform attention the keys. For the second layer, we visualize the attention paid by the (2) the other tokens (as the other attention pat- target Pq.
0.4
context-lbbll
ilum-label assoc. (L.)
0.2
targes-ieemiabe assoc. (L2)
bargei-lbell
In-coniict acourocy
100
Aligned iterations (x500)
Figure 4: Progress measures for six seeds aligned based when the IC curacy crosses 50%. The color-progress measure pairings orange: (ILAI), green: (TILA2), blue: (CLA), red: (TLA2). black: IC accuracy. See text for more details.
terns do not influence classifier output), which also shows no clear pattern. After the transition. however, the atten- tion heads show clear structure: queries in the first layer pay attention keys that immediately precede them and the target pays attention to one particular key (here, the target's ''s correct label).
Another curious feature of the IC accuracy curves is the slow learning phase that precedes the abrupt transition (Figure 3a). This phase leads to a non-negligible increase IC accuracy despite the unstructured attention maps. What drives this slow learning? We hypoth- esize that the network learns to extract useful information from the context despite not learning the optimal ICL solution. Specifically, the total number of labels (L) is larger than the number of labels represented in the context (N). The network thus randomly pick of the N contextual labels to increase its accuracy from 1/L to 1/N. This picture suggests that the target pays attention to the N labels in the second layer.
5

---

Published as a conference paper at ICLR 2024
To test this hypothesis and quantify the patterns visualized in the attention maps, we define four progress measures. Item-label association (ILA1): the attention paid by a token to its previous one in the first layer. Target-item-labul association (TILA2): the attention paid by the target to the correct label in the second layer. Context-label accuracy (CLA): the probability that the network predicts a label present in the context. Target-labels association (TLA22): the total attention paid by the target to the N labels in the second layer. (ILAI) and (TILA2) quantify the changes that occur during the abrupt transition whereas (CLA) and (TLA2) quantify the changes expected during the slow learning phase. Each progress measure is obtained by averaging over 1000 test input sequences.
Figure shows alligedd progress measures (based on when IC accuracy reaches 50%). The dynamics of IC accuracy and the progress measures are remarkably reproducible across seeds. Figure 4 con- firms the hypothesis that the network learns to randomly pick contextual label in the slow learning phase (blue curve in Figure 4). Moreover, this is accompanied by the target paying attention to the labels (red curve in Figure 4). As visualized in Figure 3b, the item-label associations of the first layer and larget-item-label associations of the second layer appear precisely at the transition (green and orange curves in Figure 4).
Induction head formation drives the abrupt transition during ICL. The dynamics of the progress measures raises various hy hypotheses regarding the factors that lead to ICL. Specifically, we are in- terested in whether learning (CLA) or (TLA2) is necessary for the abrupt transition (tracked by (ILAI),(TILA2). We consider various hypotheses and design experiments to test them: H1. (CLA) + (TLA2) - (ILA1), (TILA2). H2. (TLA2) - (ILA1), (TILA2). H3. (CLA) (ILA1), (TILA2). It is also possible that none of these factors or a progress measure that we have no considered leads ICL.
We first note that progress measures (ILA1) and (TILA2) strongly suggest the formation of an induction head. Recall that an induction head enables zero-shot copying: given an input sequence ... I. ......... .......2.
position
read
an induction head allows for predicting even if x, l never appear together during train- ing. Clearly, this task in-context. An induction head implemented by two- layer attention-only network executes the following sequence of operations (visualized in Figure 5): (i) A token (say, pays attention to the token immediately preceding it (here, x) using positional information. (ii) The value matrix of the first layer now writes the content of x into f. Importantly, this is writ- ten "buffer" subspace orthogonal to the content of f. (iii) The target æ pays attention to é by matching its content to f's buffer, which now contains the content of the contextual x that preceded it. (iv) The value matrix of the second layer writes the content of & to the target x, which is then passed to the classifier. The classifier in turn uses this information to predict
write
content
buffer
mechanism that plausibly solves our
()
(ii)
(II)
(M)
Figure 5: An illustration of the four operations performed by an induc- tion head.
We construct a minimal layer induction head that emulates these core computations and also captures the four progress measures. We assume that the input embedding space can decomposed into two or- (c) thogonal D-dimensional subspaces. For token u, these orthogonal subspaces encode content "; and a buffer (b) (initially empty). Given a sequence U1, U2,..., un, the first and second layers of our minimal model compute
model of the two-
(6)
Σ
lij U
(1) (c)
u;
(c)
(4)
JSi
w
Σ
4i
m
0
(c)
(5)

---

Published as a conference paper at ICLR 2024
(a)
(b)
(c)
8
L=
1.00
1.00
Progress meseures
Progress mesoures
0.75
0.75
2
Loss
0.50
0.50
0.25
0.25
0.00
0.00
100
200
300
400
-50
Aligned Iterations (c500)
-25
0
25
50
-50
Aligned Iterations (x500)
-25
25
50
llerations (x500)
Figure 6: (a) Aligned progress measures (plotted as in Figure 4) for the minimal three- paramte model show similar dynamics the progress measures for the full network. For L 32, N 8. (b) As in panel (a) with B2 fixed to O. (c) Loss curves for six seeds when = N =8.
where
e (mn .01 (0.0)) +.22
(6)
4ij Eksi 881h
L.
4i
(2)
Σ Agi gork (8) of +.22..
(6)
The classifier receives the concatenated vector wA @ w Here, dū, is one only if and zero otherwise. 31 thus determines the attention paid by token to its previous token (progress measure (ILAI)). a determines the attention paid by the target's 's content to a token's 's buffer (progress measure (TILA2)). A.. is one only if i - j is odd and zero otherwise. 3g thus determines the attention paid by the target to the labels in the context (progress measure (TLA2)). Since the classifier receives the target's content and buffer, it has the capacity to capture progress measure (CLA). We optimize for a, 11,1,8 32 and the classifier's 's parameters using the same training procedure as the full network. Loss and accuracy curves are presented Figure A.2.
Progress measures from the minimal model exhibit strikingly similar dynamics (Figure 6a), includ- ing the abrupt transition in accuracy. Note that the slow learning phase in the IC accuracy curve truncated in Figure 6a compared to Figure 4. Nevertheless, the network does indeed gradually learn to predict the N contextual labels (blue curve Figure 6a). The abrupt transition appears sooner for the paremmerree model, which masks the slow learning phase.
Next, we repeat the experiment fixing 3, = In this case, the target cannot pay more attention to the N contextual labels relative to the items in the second layer. The dynamics of (ILAI), (TILA2) remain the same (Figure 6b), including the abrupt transition. This experiment rules out hypotheses H1 and H2, i.e., that the larget-labels association (TLA2) leads to (ILA1), (TILA2).
The two- wvo-parameter model (with 82 = 0 in equation 6 together with the deep classifier recapitulate all the data distributional dependencies exhibited by the full network (Figure A.3). Moreover, note that the two-parameter model contains only the two parameters that characterize an induction head. This reduction strongly suggests that induction head formation drives the abrupt transition in ICL in the full model.
To test hypothesis HI that (CLA) leads to (ILAI), (TILA2), ablate the slow learning phase. Re- call that during the slow learning phase, the network learns to randomly pick one of the N contextual labels. Since L N, this simple strategy increases accuracy from 1/L L to 1/N. The slow learning phase can be prevented by setting L = N and B 1. That is, the input sequence contains all the labels exactly once. This perturbation indeed affects robust ICL. Specifically, two of the six seeds acquire the IC solution. The other four of the six seeds exhibit distinct, slow dynamics and converge to sub-optimal minimum (Figure 6c).
The loss landscape of the induction head. We now examine the loss landscape of the induction head. Through this analysis, we aim to provide mechanistic insight into the abrupt transition and ex- plains the empirical results described above. We propose phenomenological model, which contains the key elements of the two- wwo-parameter induction head and the classifier. While this phenomenologi- cal approach helps identify core features of the learning dynamics, it ignores other elements. These other factors include the effects of stochasticity and the finite dimension (D) of the embedding. We assume B- 1; it is straightforward extend the model B > 1.
7

---

Published as a conference paper at ICLR 2024
(a)
(b)
(c)
3
6-
Loss
Parameter
PL
2
1030
2000
1000
2000
Ieertions
Iterations
Figure 7: (a) Loss curve for the phenomenological model obtained via gradient descent on the loss in equation 9 (b) The three parameters 31 (layer I), a (layer 2), & (layer 3) are learned sequentially. (c) The trajectory visualized on the loss landscape (green: initial point, red: final point).
Consider softmax classifier that receives input and classifies it into L labels. Given that the target's s correct label for particular input sequence is t, the classifier models the probability that the label is t as
e
Tt. ΣJ=1 6
where Yi is the D-dimensional regression vector for label j. The cross-entropy loss given target label t is L = - log Te. The input w is given by
1
Σ 1=1,k
EN
y ==
1 8 + N1
(8)
E + 2N
+2
+ 2N
where lj the D-dimensional embedding vector for the label at index j. the index of the target label t in the input sequence and Nj N for reasons discussed below. In equation g determines the attention paid by the target to the correct label in the second layer (recall that there are 2N + tokens in the input sequence including the target). Note that we have ignored the contributions to w from the N item vectors, which contain irrelevant information and add noise to w.
(c) From c y is the product of a and of .2g where q is the target's index. v .0g is the label at pays attention to the item before it in the first layer. The attention weight corresponding to this term *N (equation 6 where N1 is the number of other tokens that compete for the label's attention, namely, 2T 1. Since т varies from 1 to N across input sequences, we use an intermediate value, N1 N, for simplicity. A more elaborate model would consider an expectation over the N possibilities.
From equation 8 the exponents in equation 7 contain dot products of the form YE_e for arbitrary pairs i,j. If all labels are statistically identical and balanced, it is simpler to track the overlaps %.l, = for all and ?,l, = ç' for all j.
In summary, the loss after re- arrarangg terms is given by
C = log
[ + (N 1)e- (L- NN)=))
where
2 -1 E. z' =
el
e
2 = E + +2N
E,
y= @ + N
2N
where Ç ç- The loss contains three nested logits parameterized by 81, ox E. which correspond to the first attention layer, the second attention layer and the third softmax layer, respectively.
The learning curves generated by gradient descent on this landscape beginning from the initial point E, a, 3, 91 Indeed, aC/05 -(L (22( + 1)) at the origin. Intuitively, when L > N, the classifier gradually aligns the regression vectors with the labels (increasing {) when learning to randomly pick one of the labels in the context. This phase is slow the classifier cannot discriminate between the N contextual labels. The gradual rise in eventually drives the loss off a cliff and leads to rapid learning of and 31 (Figures 7b,c).
recapitulate the slow learning phase and the abrupt transition (for L > N. Figure 7a).
8

---

Published as a conference paper at ICLR 2024
As shown in Figure 6c, when L N, the slow learning phase is ablated and the learning dynamics show two distinct behaviors: ICL and slow convergence to a sub-optimal minimum. We reproduce these two distinct behaviors by setting L N in equation 2 and simulating gradient descent from two points near the origin (Figure A.4a). Examining the loss landscape shows that this divergence is due to a saddle point the origin (Figure A.4b). One path leads to the ICL solution whereas the other path gradually converges to a sub-optimal minimum. Moreover, the ICL solution takes much longer to acquire due to a shallower gradient at the origin (compare Figure 7a and A.4a). We tested this prediction in the full model (equation ) Consistent with theory, the full model robustly learns an ICL solution for L > N but not when L N (Figure A.5).
4 DISCUSSION
Summary. In summary, past work has found that particular features of the data distribution influ- ence the trade-off between ICL and IWL. The features that promote ICL are especially prominent in language. such as large number of rare tokens that are highly over-represented in specific con- texts. We reproduced these data distributional dependencies in a simplified model, highlighting the essential ingredients necessary to explain those observations. We present strong evidence that ICL. is implemented by an induction head. We build a minimal version of an induction head, which through careful experiments reveal the key factors that lead to its emergence. In particular, the learning of an independent sub-optimal strategy accompanied by a slow learning phase supports the induction head's abrupt formation. A phenomenological model of the loss landscape shows that this abrupt transition is likely due to the sequential learning of three nested logits.
Abrupt transitions during ICL. An abrupt transition in loss has been noted in wide variety of ICL tasks. However, a mechanistic understanding of ICL dynamics has been lacking. Known mech- anisms for ICI., such as an induction head, rely on series of specific operations performed by multiple attention-based layers. The attention operation involves a logit (or, in general, other non- linear operations), which creates sharp gradients. A chain of operations across attention layers will thus entail a series of nested logits, which together create "cliff's" in the loss landscape. Parallel work suggests that similar mechanisms underlie 'eureka' moments in other multi-step tasks involv- ing transformers/ Hoffmann et al. (2023). Note that this mechanism is distinct from the one that leads to sharp transitions in simple reinforcement learning tasks Reddy (2022).
Relationship with past work. Our work adds to existing evidence that induction heads play key role during ICL. Olsson et al. (2022): Bietti et al. (2024). We also recapitulate the data distributional dependencies delineated in Chan et al. (2022). Our results show that even simple networks such as ours are capable of simultancously learning ICL and IWL solutions (see Figure A.1 for example). However, ICL is not transient in our simulations. This contrasts with recent work (Sngh et al. (2023) who use a much larger transformer network (12 layers and 8 heads) and finite training data. It is possible that larger networks slowly memorize the training data, leading to gradual degradation of ICL.
Implications for LLMs. We show that an intrinsic curriculum may be necessary overcome the shallow gradients at the top of the cliff and guide networks towards the ICL. solution. This observation consistent with empirical results in Garg et al. (2022). who use curricula to robustly train transformers to solve complex ICL tasks. An intriguing possibility is that learning of simpler ICL operations enables the learning of more complex ICL strategies in LLMs. An hierarchy of increasingly complex sub-tasks may lead to cascading effect and potentially explain the sudden emergence of zero-shot learning abilities in LLMs. Testing this hypothesis will require careful mechanistic analysis of minimal networks that solve complex ICL tasks. More generally, while automatic curriculum learning has been used to train foundational models for RI. Team et al. (2023). the role of curricula for accelerating ICL in LLMs remains relatively unexplored.
Limitations. While our formulation provides a minimal model that exhibits ICL, it is possible that larger models use different mechanisms than the ones that we have identified here. Methods for mechanistic interpretability Waang et al. (202)) may help probe these mechanisms in LLMs. We have not used heuristics such weight tying Inan et al. (2016); Press & Wolf (2016), which are used to accelerate training of LLMs. Such heuristics may make the slow learning phase unnecessary by aligning the classifier's regression vectors with the labels (increasing è) from the outset.

---

Published as a conference paper at ICLR 2024
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306. 2it.:2306000099,, 00297, 2023.
Kabir Ahuja, Madhur Panwar, and Navin Goyal. In-context learning through the bayesian prism. arXiv preprint arXiv:2306. 230.:233000 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arw.v:2211.156000 2022.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877 2020.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learn- ing transformers. Advances in Neural Information Processing Systems, 35:18878-18891, 2022.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in- context? language models secretly perform gradient descent meta optimizers. arXiv preprint arXiv:2212. 10559, 2022.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. survey for in-context learning. arXiv preprint arXiv:2301.0022 2022.
Shivam Garg. Dimitris Tsipras, Perey S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
David T Hoffmann, Simon Schrodi, Nadine Behrmann, Volker Fischer, and Thomas Brox. Eureka- moments in transformers: Multi-step tasks reveal softmax induced optimization problems. arXiv preprint arXiv:2310. 12956, 2023.
Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying word vectors and word classifiers: A loss framework for language modeling. arXiv preprint arXiv: /6/ 1.01462, 2016.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General General-purpose in-context learning by meta-learning transformers. arXiv preprint 2tth:2222204458,, 2022.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua Tenenbaum. The omniglot challenge: 3-year progress report. Current Opinion in Behavioral Sciences, 29:97-104, 2019.
Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers al- gorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067, 2023.
Sheng Lu, Irina Bigoulaeva, Rachneet Sachdeva, Harish Tayyar Madabushi, and Iryna Gurevych. Are emergent abilities in large language models just in-context learning? arXiv preprint arXiv:2309.01809 2023.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint 2209:2200.11889 2022.
10

---

Published as a conference paper at ICLR 2024
Ofir Press and Lior Wolf. Using the output embedding to improve language models. arXiv preprint arXiv:1608.000 05859, 2016.
Gautam Reddy, A reinforcementt-assed mechanism for discontinuous learning. Proceedings of the National Academy of Sciences, 119(49)::2215352211 2022.
Aaditya K Singh, Stephanie CY Chan, Ted Moskovitz, Erin Grant, Andrew M Saxe, and Fe- lix Hill. The transient nature of emergent in-context learning in transformers. arXiv preprint arXiv:2311.08360, 2023.
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmice Michael Chang, Natalie Clay, Adrian Collister, et al. Human-timescale adaptation in an open-ended task space. arXiv preprint arXiv:2301.06608, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learming, pp. 35151-35174 PMLR, 2023.
Kevin Wang, Alexandre Variengien, Arthur Commy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593 2022.
Xinyi Wang. Wanrong Zhu, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. arXiv preprint arXiv:2301.11916, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2021.

---

Published as a conference paper at ICLR 2024
A
APPENDIX
0.25. = -0.25
0.25, Pc. = 0.75
1.0
1.0
Test ace
Test ace
IC ncc
IC acc
Accuracy
IC2 acc
Accuracy
IC2 acc
0.5
IW
0.5
IW
0.0
0.0
100
200
300
400
0
100
200
300
400
Iterations (x500)
Iterations (x500)
0.75,
0.25
0.75, Pc 0.75
1.0
1.0
Test acc
Test acc
IC acc
IC acc
Accuracy
IC2
Accuracy
IC2
0.5
0.5
0.0
0.0
100
200
300
400
0
100
200
300
400
Iterations (x500)
Iterations (x500)
Figure A.1: Accuracy curves for the full model when 0 < pn < 1 and 0 < PC < 1. In all cases, the network learns both the ICL and the IWL solutions. Here 256, B = 1, a = 0, O.
1.00
IC
0.75
IC2 acc
Accuracy
Loss
IW acc
0.50
0.25
0.00
100
150
200
150
200
Iterations (x500)
Iterations (x500)
Figure A.2: Loss and accuracy curves for the minimal model. Here K 512, D = 64, B
12

---

Published as a conference paper at ICLR 2024
0.75
0.50
Tip
0.25
Zop 25
Top 500
0.0
0.6
0.00
0.0
0.5
1.0
2.0
Within-clas varintion
Class
scaling
0.75
06
0.50
0.0
0.0
0.2
0.8
0.0
0.00
Nombeer classes
Within class vasiation.
Class Drequency scaling exponent
Figure A.3: Data distributional dependencies recapitulated by the minimal model. Plotted in Here 512, D 64, B a = 0, e = 0.1 (except when that parameter is varied) Figure 2.
(a)
(b)
2.0
2.0
5
1.5
Loss
1.0 4
1.0
0
0.5
0.5
-5
0.0
0.0
10000
20000
-5
0
Iterations
Figure A.4: (a) When L N, the loss curves starting from two initial values recapitulate the two distinct behaviors noted in Figure 6c. (b) The loss landscape has saddle origin such that small fluctuations lead the path either to the ICL solution (top right quadrant) or a sub-optimal minimum (bottom left quadrant).
13

---

Published as a conference paper at ICLR 2024
N=
1.0
--2 2
= 4
IC accuracy
8
0.5
0.0
100
200
300
400
Iterations (x500)
N
1.0
8
IC accuracy
16
32
0.5
0.0
0
100
200
300
400
Iterations (x500)
N=8
1.0F
8
16
IC accuracy
32
0.5
0.0
100
200
300
400
Iterations (x500)
Figure A.5: IC accuracy curves for different N and L (six seeds for each pair of values of L and
N shown). Consistent with the theory and the minimal network, the full network (equation 2 robustly learns the in-context solution L > N but not when L = N. Here K = 256, B = 1, PC 0.8, PB = 1,a 0,
14


Reviews:
[{"review_id": "zbDUCdPi5c", "reviewer": "Reviewer_sq8T", "summary": "The authors build on recent work on induction heads as the mechanism of in context learning (ICL) and give a characterization of how they are learned.  They start with a minimal transformer architecture, show that it capture previous findings on data dependence of ICL vs in weights learning (IWL) and proceed to carry our further model reductions that focus in on crucial abstractions that characterize the learning dynamics in terms of abstract underlying variables rather than particulars of particular connection weights.  They test several causal hypotheses, zeroing in on factors that jointly influence how learning occurs.", "questions": "Suggestion 1: Confirm that they setting L=N and B=1 disrupts learning of the induction head in the full model as it does in the 2/3 parameter reduction and in the corresponding phenomenological model.\n\nSuggestion 2: Redirect some of the presentation to an appendix to provide at least 1/3-1/2 page for discussion of implications for LLMs.\n\nSuggestion 3:  Provide an intuitive understanding of the shape of the loss landscape in both the 2 and 3 parameter version of the phenomenological model.\n\nSuggestion 4: Reduce the cognitive load on the reader: Specifically, increase the conceptual content of referential expressions throughout.  Don't use arbitrary letters/roman numerals for quantities and hypotheses; avoid former/latter and 'respectively' type constructions where possible.  Also, place figures as close as possible to the point in the text where they are described.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The ability to use information in context to respond appropriately to later queries (called 'in context learning' or ICL), is central to the capabilities of AI systems like ChatGPT.  ICL was enabled by the attention mechanism in transformer-based neural networks.  ICL is exemplified by the simple item-label association task (introduced by others) that the authors have selected for the focus of their analysis.  By shedding light on how this task is solved (building on an earlier paper taking initial steps in this direction) the current paper deepens our understanding of this core property of today's performant AI systems.\n\nThe authors have created new variants of the task that further simplify it, and have introduced minimally-sufficient transformer architecture containing two attention layers, which together implement what they call an 'induction head', arguably the core emergent computational structure enabling ICL.  They have provided a insightful analysis of the (ultimately simple, but nevertheless important) computations performed by the network that allow the effective use of information in context in their task.  They gone on to attempt to understand how this attention head computation emerges as the network learns to solve the item-label association task.  They identify progress measures in both the network's input-output performance and of its attention head computations and establish clear alignments of several of these measures.  I consider these measures and their alignments enlightening contributions and consider them to be strengths of the paper.\n\nThey go on to further support their analysis by developing a three-parameter reduction of the induction head, and show that the learning dynamics of this reduction is sufficient to reproduce many of the features of the learning dynamics of their complete neural network; they then use the reduction to test hypotheses about the relationships between the progress measures, showing that a further reduction that eliminates one of the progress measures makes learning success initialization dependent.  Finally, they make an even further reduction in the form of a 2- or 3- parameter 'phenomenological model' whose loss landscape can be fully characterized.  The parameters now directly reflect the efficacies of the two attention layers making up the induction head and of their mapping to the correct label, and allow the loss landscape of each of the variants to be visualized.  This phenomenological model provides an abstract characterization of the emergent learning dynamics of the 2- and 3 parameter reduction models that allows a full explanation of why these models learn reliably under the condition that the number of possible labels is greater than the number of item-label pairs in context, and fails to learn reliably when the number of possible labels is equal to the number of item-label pairs in context.  All the paper provides us with important clues toward understanding the computations performed by transformers and of the processes that give rise to their learning dynamics.  Along the way the paper provides an approach to analysis of neural network learning dynamics that others could adopt to understand the learning dynamics as they arise in other setting, another valuable contribution to the effort to understand the complex computations performed by neural networks.", "weaknesses": "Although I consider the analysis presented a tour de force, possessing all the strengths describe above, it is not perfectly clear that the analysis of the 2-3 parameter reduction would carry over to the full 2-attention-head network of Figure 1c.  A hunch I have is that the L=N case might not be quite as susceptible to failure in the full network because the full network might have a more complex loss landscape with a lower likelihood of being initialized in a place in that landscape that doesn't allow a complete solution.  An important and simple step toward addressing this would be to repeat the L = N simulation in the full network.  If the full network fails to learn in that case, it would confirm the applicability of the analysis to the full network.  Success would not fully invalidate the analysis, but would leave something left to explain.\n\nMore generally, I believe more consideration of what will happen in a larger model will be useful for the field.  Clearly things will not work just in the way they do in these reductions when the task is learned in a larger transformer.  While fuller characterization of that will be a task for future work, noting this issue as a limitation of the present effort and pointing considering how these results inform us about what is happening in LLMs will be valuable.\n\nThere are two less important weaknesses I'd like to see addressed. \n\nFirst, I don't feel I have an intuitive understanding of why the loss landscape of the 3 parameter model does not have a saddle point at the point were all three parameters are equal to 0.  Perhaps an understanding of this is latent in the equations and I could work it out with a bit of effort, but to help me (and possibly others) understand, it would be useful if the authors could work out such an intuitive understanding.  Such an understanding could help address reasons why the behaviors of the 2- and 3-parameter reduced models might or might not be applicable to the full model.\n\nSecond, paper is harder to read than it should be.  The main deficiency of the paper was its failure to take cognizance of the difficulty of extended chains of arbitrary associative bindings requiring long-distance leaps across context.  It is just such binding that lie at the heart of the mechanisms the authors are investigating, but they are hard for human readers when arbitrary as they often were in this paper. \n\nAs examples, we are treated to terms like the former vs the latter as referring expressions, arbitrary labels (a-d) for key phenomena, random ordering of the assignments of these labels to lines in graphs, arbitrary labels for hypotheses (I-V), and the unhelpful placement of figures (esp fig 4) on pages remote from the place in the paper where they are discussed.  Although ultimately the conclusions are stated in (what I find myself to be) conceptual terms, there should be engagement with this conceptual structure in the referential expressions used. I know space is limited, but I'm sure it is possible to do a better job. As examples, H3 could be abbreviated sCLA -> ILA+TILA (slow-learned context-label attention -> Item-label attention and Target-item-label association). Just let a,b,c,d and I-V go.  H4 and H5 should each be expressed directly, or at the very least the order of defining the symbols x and 0/ should correspond to their order of appearance in these hypotheses. \n\nI am also not sure that the difficulty of the L=N case in the", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oxmXlHjg2U", "reviewer": "Reviewer_Y7Wm", "summary": "The paper attempts to provide a deeper understanding of in-context learning of LLMs -- which is also related to the broader discussion about their \"emergent learning capabilities\". \nTo provide such an understanding, the authors wisely choose to abstract away many details of LLMs, and of the tasks that LLMs usually perform, and to focus instead on a simple experimental setting in which they can easily control or monitor whether the task is performed through in-weights learning (IWL) versus in-context learning (ICL). Additionally, the data is generated through a parsimonious gaussian mixture model in which they can also control some important aspects, such as the burstiness with which certain classes appear in the sequence, or the rank-frequency relation. \nThen, and based on the insights from the previous experiments, they design a very small (in terms of number of parameters) phenomenological model of an \"induction head\" that reproduces quite well the observed behaviors of the more complex attention-based networks used in the experiments.", "questions": "As written in the Strengths section, I am very positive about this paper and so I do not have many technical suggestions or questions for the authors. \n\nI would like to see at the end a clear discussion about the limitations of this simple/abstract model. Which aspects of an LLM's behavior may still be important but not captured by the proposed simple task and model that the paper proposes? \n\nI would also like to see a more clear discussion of how your observations/conclusions agree (or disagree?) with earlier results in the literature about ICL and the emergent abilities of LLMs.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "I have read several papers recently that try to provide some insights about the emergent capabilities of LLMs -- through abstract modeling and experiments with tasks such as linear regression (learned through ICL). This paper is the best I have read so far in that direction. \nThe experiments are wisely designed, allowing us to understand the complex tradeoff between IWL and ICL -- as well the effect of some key data distributional parameters such as the rank-frequency exponent. \nThe simple model proposed in the second part of the paper is also intriguing, explaining how the abstract model of an induction head can explain mechanistically the ICL capabilities of an attention-based network.", "weaknesses": "The paper can be improved in terms of writing/presentation. For example, you can explain early on in the paper what \"induction head\" means for readers that are less familiar with this area. \n\nThere are also several other parts of the paper in which the writing can be improved -- mostly by writing simpler/shorter/more clear sentences.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VokzE19go0", "reviewer": "Reviewer_3R5R", "summary": "Expanding on the seminal contributions of Chan et al and Olsson et al, this paper investigates the hugely important topic of the emergence of in-context learning in supervised learning via transformers. Specifically, it starts with experiments about the concurrence of in-context (IC) and in-weight (IW) learning, in the spirit of Chan et al, but systematically deconstructs their findings in such a way as to expose several potential hypotheses for the mechanism responsible for the ICL transition. The paper's main contribution is then to replicate and explain empirical phenomena via a three-parameter nested logits ansatz (based off a two-layer attention network and a linear classifier), with the appearance of an induction head controlled by a main parameter \\xi seen as the difference between overlaps in on-diagonal and off-diagonal dot products feeding the final softmax.", "questions": "What are the authors' intuition as to why does \\xi undergo a phase transition driving ICL ? And how to, ideally, accelerate it using insights derived here ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is excellently written and fairly easy to follow despite the depth of insights proven. The scientific investigation is very well conducted : of note is that it alternates particularly well between empirical elements, formulating subsequent hypotheses (section \"Induction head formation drives the abrupt transition during ICL\"), disproving some of those, and finally introducing a theory that accounts for those findings, replicating empirical stylized facts, whilst much simplifying the problem. In particular, the phenomenological model Equation 10 (and its illustration Figure 7) is a standout novel contribution, and clearly worthy of publication, in our view.", "weaknesses": "In a sense, the paper is tantalizing, as it invites further work, for instance on the interplay of overlap difference \\xi and data Zipfianity parameter \\alpha.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2ZCcyMAYL9", "reviewer": "Reviewer_ciFG", "summary": "This paper explores how and why properties of the data distribution control a transformer's propensity to use in-context vs. in-weights learning strategies.  The authors design a synthetic, parameterized family of tasks that expose dependencies of learning strategies on the data distribution that have been identified in prior work.  They then analyze the behavior of models trained on these tasks to understand how the dynamics of learning give rise to different strategies over time.  One of the key findings is that models learn a shortcut strategy that enables better-than-chance (but suboptimal) performance -- choosing an answer among one the labels presented in-context, without regard for the inputs presented in context.  In a stripped-down model that maintains the qualitative behavior of the full transformer parameterization, the authors show how this shortcut strategy facilitates (but is not necessary for) learning of a true ICL strategy via an analysis of the loss landscape.", "questions": "-- In my opinion, the connection between the phenomenology observed in the full transformer model and the insights the authors derive from the three-parameter model could be made stronger with more experiments.  Have the authors considered techniques like activation/path patching and knockouts (see e.g. https://openreview.net/forum?id=NpsVSN6o4ul for examples of this approach) to test whether the explanations they come up with actually have explanatory power in the original model?\n\n-- Why do the authors switch from the p_C = 0 case to the p_C > 0 case halfway through the paper?  I find this confusing as it makes the relevance of the mechanistic analysis to the p_C = 0 regime unclear.\n\n-- The authors note that when p_B < 1, the network learns an IWL solution, and then fix p_B = 1 thereafter.  Later, they note that p_C > 0 always leads to an ICL solution (presumably with p_B fixed at 1).  To me, the authors have failed to consider the most realistic and interesting regime, where 0 < p_B < 1, and 0 < p_C < 1, where neither an ICL solution alone nor an IWL solution alone is optimal.  Is there a reason the authors choose not to consider this regime?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This paper tackles a timely and interesting topic, and contains several insights and useful contributions.\n\n-- The synthetic task family introduced is a clean and intuitive way of exposing dependencies of learning strategies on pretraining data distributions.  Showing that key phenomena identified in (Chan et al., 2022) can be replicated in this setting is a useful contribution.\n\n-- The characterization of the initial slow learning phase as driven by an increase in context-label accuracy is interesting\n\n-- The idea that a strategy that results in good context-label accuracy can facilitate (despite not being necessary for) learning of a true ICL strategy is very interesting, and some evidence is provided for this idea\n\n-- The evidence provided that the emergence of induction heads is strongly linked with the development of the ICL strategy, while not entirely novel, is nice to see", "weaknesses": "I have the following concerns about this paper.  Many of them involve claims that I feel are made too strongly in the paper relative to the level of evidence provided.\n\n-- The paper illustrates a set of phenomena in figure 2, and promises a mechanistic understanding of these phenomena.  But the mechanistic analysis provided later in the paper does not speak to most of the phenomenology -- for instance, the dependence of the ICL/IWL tradeoff on B, epsilon, K, and alpha.  In fact, the mechanistic analysis focuses on the p_C > 0 case, which is different from the p_c = 0 regime that gives rise to all the tradeoffs observed in Figure 2.  Thus, the connection between pages 1-4 of the paper and the rest is not entirely clear.\n\n-- -- The following sentence, while intuitively reasonable, is written as a key strong claim and as far as I can tell is not really justified with evidence: \"Therefore, the relative rates at which the network acquires ICL and IWL control the fraction of loss explained by each mechanism after convergence.\"\n\n-- The paper makes strong causal claims based only on correlational evidence.  For instance, \"Induction head formation drives the abrupt transition during ICL.\"  As far as I can tell no evidence is given for this claim, other than the (very suggestive, I agree!) fact that they coincide in time.\n\n-- The paper makes strong claims about the three-parameter model proving or ruling out certain hypotheses.  An example is the sentence \"This rules out hypothesis V as only the factors corresponding to the progress measures (a) through (d) have been included in the minimal model.\"  In my opinion, such claims are much too strong.  The three-parameter model is ultimately a different model from the original transformer architecture being used!  While the analysis of its behavior is suggestive of the learning strategies used by the original architecture, it is not conclusive.  The strength of the claims should be adjusted accordingly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)