Differential learning kinetics govern the transition from memorization to generalization during in-context learning
Authors: Alex Nguyen
Gautam Reddy
Venue: ICLR 2025 Spotlight (2025)
Keywords: in-context learning
mechanistic interpretability
small transformers
memorization
Primary area: interpretability and explainable AI

Abstract: Transformers exhibit in-context learning (ICL): the ability to use novel information presented in the context without additional weight updates. Recent work shows that ICL emerges when models are trained on a sufficiently diverse set of tasks and the transition from memorization to generalization is sharp with increasing task diversity. One interpretation is that a network's limited capacity to memorize favors generalization. Here, we examine the mechanistic underpinnings of this transition using a small transformer applied to a synthetic ICL task. Using theory and experiment, we show that the sub-circuits that memorize and generalize can be viewed as largely independent. The relative *rates* at which these sub-circuits learn explains the transition from memorization to generalization, rather than capacity constraints. We uncover a memorization scaling law, which determines the task diversity threshold at which the network generalizes. The theory quantitatively explains a variety of other ICL-related phenomena, including the long-tailed distribution of when ICL is acquired, the bimodal behavior of solutions close to the task diversity threshold, the influence of contextual and data distributional statistics on ICL, and the transient nature of ICL.

Published as a conference paper at ICLR 2025
DIFFERENTIAL LEARNING KINETICS GOVERN THE TRANSITION FROM MEMORIZATION TO GENERALIZA- TION DURING IN CONTEXT LEARNING
Alex Nguyen
Gautam Reddy
Princeton Neuroscience Institute Princeton University gaaguyyempprinccton.eee
Department of Physics
Princeton University
greddyy@riineeton.eet
ABSTRACT
Transformers exhibit in-context learning (ICL): the ability to novel informa- tion presented in the context without additional weight updates. Recent work shows that ICL emerges when models trained tasks and the transition from memorization to generalization sharp with in- creasing task diversity. One interpretation is that network's ' limited capacity memorize favors generalization. Here, we examine the mechanistic underpinnings of this transition using small transformer applied a synthetic ICL task. Using theory and experiment, we show that the sub-circuits that memorize and general- ize can be viewed as largely independent. The relative rates at which these sub- circuits learn explains the transition from memorization to generalization. rather than capacity constraints. We uncover a memorization scaling law, which deter- mines the task diversity threshold at which the network generalizes. The theory quantitatively explains variety of other ICI.-related phenomena, including the long-tailed distribution of when ICL is acquired, the bimodal behavior of solu- tions close to the task diversity threshold, the influence contextual and data distributional statistics ICL, and the transient nature of ICL.
sufficiently diverse set
INTRODUCTION
Large transformer models trained to predict the next token exhibit powerful generalization capabil- ities. One signature of such generalization capabilities is in-context learning (ICL): the ability to solve a task based on new information presented the context without additional weight updates (Brown al. (2020); Dai al.) (2022); Dong et al. (2022); Garg al. (2022); Xie et al. (2021): Olsson et al. (0222). Arguably, the ability to interpret novel inputs on-the-fly is core feature of any intelligent system. However, updating synaptic weights rapid behavioral timescales is chal- lenging, both for natural and artificial systems. The emergence of ICL in large language models (LLMss shows that finding network states that learn on-the-fly is indeed possible. Understanding how ICL emerges in LLMs promises insights into how such algorithms may be implemented in the brain and how the data distribution, training objective and network architecture interact enable ICL acquisition at scale.
Various methods have been used probe the ICL capabilities of LLMs (Brown et al. 2020 Dong et al. (2022); Pan (2023); Min et al.] (2022); Olsson et al. (2022). A common ICL paradigm is present exemplars as sequence
febellabel pairs, and measure the network's 's response
target item (Chan et al. (2022); Kirsch et al. (2022); Garg et al. (2022); Akyürek et al. (2022);
Von Oswald et al. (2023); Raventós al. (2023): y: Bai al. (2023). While LLMss display remarkable capabilities on such ICL. tasks, interpreting the underlying network mechanisms that give rise to
these capabilities remains challenging (but see WWnn et al. 2022 Recent work has approached this challenge by examining how small transformer models solve synthetic ICL tasks (Reddy (2023); Bietti et al. (004); Akyürek et al. (2022); Ahn et al. (003); y Von Oswald et al.) (2023); Edelman et al. (2024)). We highlight two notable aspects of ICL phenomenology relevant for current work: the influence task diversity on whether the network memorizes a finite dataset or acquires ICL (i.c.,
generalizes), and how ICL. is acquired (and lost) during training.

---

Published as a conference paper at ICLR 2025
First, data distributional properties (such as task diversity and their rank- funkffruuency distribution) in- fluence whether the network acquires ICL or encodes the response to queries seen during training within its weights (Kirsch et al. 2022); Chan et al. (2022); Raventós et al. (2023). Following pre- vious work, we refer to such memorization as in-weights learning (IWL.) Notably, the transition from memorization (IWL.) generalization (ICI.) is sharp with respect task diversity. A curious feature of this transition is that solutions close to the task diversity threshold are bimodal (Kirsch) al. 2022)). That is, across different random number seeds, solutions either show IWL or acquire ICL but intermediate solutions are unlikely.
Second, ICL is often implemented by multi-layer computations involving nonlinear attention heads and MLPs OOssoon et al.] (2022); Von Oswald et al. (2023); [eeddy (2023); Bietti et al. (2024). The rugged loss landscape induced by such multi-layer, nonlinear operations leads to long plateaus followed by a sharp drop in loss (Reddy (2023). Finally, ICL is seemingly transient, i.e., the network gradually loses the ICL capability if it trained for sufficiently long (Singh et al. (203)).
It is unclear what leads to the transition from IWL to ICL acquisition with increasing task di- versity. The loss is minimized when finite training dataset is perfectly memorized. Thus, naive hypothesis would suggest that the task diversity increases, the network's limited capac- ity to memorize favors the ICL solution (Figure D... Our goal is to test this hypothesis by deriv- ing a precise quantitative description of the tran- sition, and thus we set out to identify a minimal setting that captures the phenomenon.
(a) consaity
(b) Differential learning kinetics
data diverstty
ICL
IWL
ICL
IWL
IWL
Contributions and outline. We first identify one-layer transformer model trained on an context classification task that recapitulates the sharp transition from memorization ization. Despite its simplicity, the one-layer model displays surprisingly rich phenomenol- ogy, including abrupt ICL learning dynamics, transience and bimodal solutions close to the task diversity threshold.
IWL
ICL
Figure 1: (a) In the sapacity-constriine model, the network's limited capacity to memorize fa- vors ICL. acquisition with increasing task di- versity. (b) In the differential learning kinetics model, independent sub-circuits contribute to- wards IWL and ICL. IWL slower for greater task diversity. The network acquires ICL before the network can significantly memorize the train- ing set. IWL is significantly slowed down as ICL explains most of the loss, but does eventually memorize the training set. The network subse- quently loses the ICL capability due to regular- ization.
general-
Next, we derive an analytical framework that quantitatively characterizes ICL acquisition and its competition with IWL. We show that the transition from IWL to ICL for our network is governed by a dynamical competition be- tween memorization and generalization (Figure [[p). However, whether the network is capacity- constrained or rate-determined depends on the network architecture, and we derive quantita- tive measure to determine which of these constraints is at play. The theory predicts that the number of iterations before ICL is acquired is exponentially sensitive to the initial parameters, which in turn explains the bimodal behavior of solutions close to the task diversity threshold. The task di- versity threshold follows power law whose exponent has non-trivial relationship with another novel memorization scaling law. ICL transience naturally follows from the theory under standard La regularization. Finally, we validate our theory by empirically verifying these predictions using our transformer model.
TASK FORMULATION
We consider simplified version of an ICL task proposed by Chan et al. (2022). which allows for disentangling ICL and IWL performance (Figure 21). Before training, we generate a dataset D
that contains K item-label pairs, D {(11, (1), (122 fz),.... (xx, {.)]. Each item 2; is D-
2

---

Published as a conference paper at ICLR 2025
Dataset !
Model
1.0
(6)
0.9
0.9
4s)
IC), nemnecy
0.6
Areumcy
0.8
llayer
42
6.3
TWL
attondon heed
0.5
Input
Thining thenation:: (x.EP)
(c)
(0)
Training input sequence
0.0
ecersey
0.5
ICL eval sequence (novel lam-lsa
0.7
0.0
0.0
WL eval sequence (no exemplar)
0.5
10
10
10
Thiining teroationn (x:1PP)
Figure 2: (a) Data generation process: We create a dataset D consisting of K item-label (Iz, L.) pairs where each (O, D) and É, is randomly sampled from {-1,+1}. The network receives sequence of N + tokens. Each of the first N tokens is a concatenation x; @ l; of an (æ,, 4) pair sampled uniformly from D (details in main-text). The final N target token consists of only the item I; as its label component is zero-ed out. The network trained to correctly predict the label of the last item. Model architecture: The input normalized using LayerNorm, then fed to our network, which consists of a one-layer attention network followed by a 3-layer MLP. (b) ICL performance demonstrates sharp transition as function of data diversity K. Further, at the transition threshold K z 10*, we observe bimodality where the model either memorizes or generalizes. (c) IWI. accuracy vs K. IWI. and ICL. accuracies follow opposite trends. (d) ICL, accuracy curves show that ICL, performance plateaus at the beginning of training but undergoes rapid transition as ICI. is acquired. (e) ICL is transient, i.e., ICI. accuracy gradually decreases to chance levels when the parameters in the attention head are heavily regularized.
dimensional random vector with components drawn i.d from N(0,1/D) and is randomly assigned one of two labels, L {11,
The data is presented to the network as a sequence of N - 1 tokens, where each token Lj is the item concatenated with its label, t, 2y @ Ey. The first N tokens in the sequence are drawn uniformly from D. The N + 1th token (the target token) has an empty label vector. The target token is chosen uniformly randomly from the N items in the context, so that there is an exemplar always present in the context. Given an input sequence, the network is trained to predict the label of the target token using a binary cross-entropy loss. Since the total number of itemllabel pairs (K) is finite, the network can either memorize each item's label (IWL), can learn to use the exemplar within the context to predict the correct label (ICL).
To measure ICL, we construct test dataset D consisting of novel item-label pairs (sampled like D) and evaluate the network's accuracy on sequences sampled from Dvett using the previously described procedure. To measure IWL, evaluate the network sequences sampled from D, except that the target has no corresponding exemplar in the sequence. In this case, the context has useful information, and the network must rely the label's information encoded within its weights.
3 RESULTS
3.1 ONE-LAYER TRANSFORMER MODEL RECAPITULATES ICL PHENOMENOLOGY
We begin with a one-layer attention- based network followed by a multi-layer perceptron (MLP). Given tokens 4,, ........ ... 11, we first apply LayerNorm operation to obtain 1 = TayerNorm((.))
3

---

Published as a conference paper at ICLR 2025
(a)
(b)
(c)
1.0
1.0
0.9
0.9
0.9
IC1, Acrancy
0.8
Accuray
00.8
Avvirary
0.8
ICL
ICL
0.7
0.7
IWI.
0.7
IWL
0.0
0.6
0.6
0.5
0.5
10
10
500
1000
1500
2000
5000
10000
15000
20000
Training iterations (x 50)
Truning Iserataons (x 50)
Figure 3: Phenomenology of the minimal model. (a) ICL performance in the minimal model demon- strates sharp transition as function of data diversity K. (b) ICL acquisition is abrupt during training. (c) ICL is transient when exclusively regularized.
The attention operation then produces the output
N+1 TTK
Q
u N44 +
Σ TKTQUN Vt
where Q, K, V are (D x (D + 2) query, key and value matrices (L₃ 's are one-hot vectors for the one-layer model, and we use Lj =1 elsewhere). The MLP & takes input of dimension D 2 and produces logits s(u). If the target's true label is (1, 0), the loss is log or (d(u),, and the target's true label is 1), the loss is logoo- where is the logistic function. The MLP is a three-layer ReLU network with hidden dimension d. We use stochastic gradient descent (SGD) with batch size 128, learning rate 0.01, weight-decay 10-10, D 63, N 100 and d = 512, unless otherwise specified.
We examine how ICL accuracy scales with the number of item-label pairs, K. For K < K* where K ~ 10a, we find that the network shows IWL but fails to acquire ICL (Figure 2b.c). Conversely, for K > K*, the network consistently acquires ICL but does not show IWL. The transition from IWL to ICL is sharp. Solutions in the vicinity of K = K are bimodal (visible in Figure 2p and quantified later in Figure (o)). Examining the ICL accuracy during learning shows an abrupt transition from chance-level to perfect accuracy (Figure 2t). Previous work has also shown that ICI. accuracy gradually decays to zero if training continued for a sufficiently large number of iterations (Singh et al. (2023)). Transience is recapitulated when we significantly extend the number of training iterations (Figure 2b), provided the parameters within the attention head are regularized more heavily than those in the MLP. This skewed regularization scheme hints at a potential explanation for what causes transience (as noted in Singh et al. (2023)), which we will later explain quantitatively using theoretical framework.
3.2 DISENTANGLING ICL AND IWL IN MINIMAL MODEL
Figure 2 c-f) shows that, despite its simplicity, the one-layer transformer model captures the core features of the memorization to generalization transition and ICL training dynamies observed in more complex models. However, mechanistic analysis is still challenging due to the nonlinearities in the attention head, the MLP and how these two operations interact. To make progress, we further reduce our one-layer transformer model into disentangled model (which we refer to as the "mini- mal model" hereafter) by proposing two ansatz. We will show empirically that the minimal model also reproduces the phenomena in Figure c-f) This minimal model is amenable to theoretical analysis and leads to specific quantitative predictions. We then validate our ansatz by empirically testing these predictions using original transformer model (Section 4).
To motivate the ansatz, we observe that ICL in this task involves a simple match-to-sampee operation implemented by the attention head. The attention paid by the target token is determined by its dot- product similarity with the content (the first dimensions) of the tokens in the context. The value matrix reads the labels of the tokens weighted by the attention paid to those tokens and passes it
4

---

Published as a conference paper at ICLR 2025
to the MLP. Put together, these observations suggest that the relevant operations performed by the query-key product and the value matrix are captured by
[IIDDD 0lXD
0px1
V= (0Dp
ODxD
kJQ
01x1
where DXD D-dimensional identity matrix, 3, w are (learnable) scalars and the rest of the components are zeroes,
From equation 1 the MLP receives the sum of the target token t++1 and the output of the attention head as its imput. Note that all the information required by the MLP memorize the label of the target is present in + That is, IWL does not require the output of the attention head. Similarly, all the information required predict the target's 's label using ICL is contained in the output of the attention head. Based on these observations, we posit that the final logit used to predict the target's label is the sum of logits generated independently by an MLP (which takes x++1 as input) and the attention head (using the simplified K Q and V matrices in equation 2}. Specifically, given an input sequence t1, f2......................... ........,...... tN+1, we assume the MLP produces logit 2MLP o(EN+1) and the attention head produces the logit
N
BzfaN+ 2N+1
ZATT
G0 xx44 wljj
(3)
j=1
Put together, the predicted probability that the target's label is +1 given by a(ZMP + 2xTT)- The binary cross-entropy loss for given input sequence is then - log (Lo(EMLL 2xTT)), where le = +1 is the true label of the target.
In summary, assume independence ansatz, where the attention head and MLP perform ICL and IWL. respectively, and additively contribute towards the prediction of the target's label. Further, we assume that the majority of the ICL learning dynamics is captured by reducing the KT Q and V matrices to two "order parameters", 3 and w. That the strengths of the relevant attention operations determine ICL acquisition is our second ansatz.
The minimal model is parameterized by 8, w and the parameters of the MLP (a three-layer ReLU network of hidden dimension 512). The model trained and evaluated using the same procedures used on the transformer model. Figure aa-c) show that the three phenomena of interest, that is, the transition from memorization to generalization, abrupt ICI. learning and ICL. transience, are reproduced by the minimal model. We now use the minimal model to develop analytical theory. We outline the main results here and present more detailed derivations in the Appendix.
3.3 THE LOSS LANDSCAPE OF THE MINIMAL MODEL
We consider the asymptotic limit K > N >> 1 and the infinitediimensional limit D - ci (recall, K 103, NN = 10a, D = 63 in our experiments). From equation 3 ICL is acquired when w, 3 > 1. Our goal is to compute the time taken for the network to acquire ICL starting from w wo 3 Bo with |,0o, 1,11,,111 |.0l < 1.
In the limit D oo, the dot product 2. xx+1 1 is 1 if xy is copy of the target and 0 otherwise. It unlikely there is more than one copy of the target in the context when K > N. Let denote the index of this copy. From equation 3 we have
ZATT R w
( Lo
N N2
(4)
where n4 is the (binomally distributed) number of tokens with label +1 amongst the N tokens in the context. When N > 1, n.N 1/2 + n/2v/, where r7 (0,
Next, the MLP's contribution to the average loss appears only through the distribution of logits ob- tained by applying the MLP to each of the K items D. In particular, denote P+ as the distribution of logits obtained when the MLP is applied to the items in D with +1 label. We use the fact that the two labels are symmetric, and average over nt and p to show that the average binary cross-entropy loss L is (see Appendix)
(11
N
L
" VN
log Φ
+ w
(5)
c + N
+ - -1
1o-N(0,1).0000 IPP1
5

---

Published as a conference paper at ICLR 2025
The loss landscape throughout training is thus specified by Pt, and B. Retaining the fluctuations in is necessary to accurately describe ICL acquisition.
3.4 THE DYNAMICS OF ICL ACQUISITION
To examine the dynamics of ICL acquisition, we find expression for the loss when |w| < v N and e
< N (for arbitrary P+). Both these conditions are satisfied at initialization (wwbl, 130l < 1). a few steps of simplification leads to (Appendix)
From equation
5
L (og + e -- )),- C1
ee C2
2 2
(6)
N
where 0j = (a(-))))>> o+ and C2 = 1 (o(- ((o((()))+++ ///+//c Here, we used |w| < VN and
N to Taylor expand equation and retained terms to order 1/N (terms of order I/vN vanish in expectation). The distribution from which o drawn has been dropped for notational convenience.
Equation 6 allows us to make several important in- ferences. The first term in the r.h.s of equation 6 is the loss incurred by the MLP. It does not involve w. 1 and thus does not affect ICI. learning dynam- ics. Since the second term in the r.h.s of equation 6 is small at initialization, the rate at which the MLP memorizes is not affected by ICL learning. That is, IWL proceeds without any competition from ICL until ICL is acquired (which happens abruptly).
2.0
1.5
C
1.0
The scalar variables C1 and depend on pt and thus depend on the time since training began. Their evolution in general depends multiple factors, in- cluding MLP architecture, initialization scheme and the number of tasks K to be memorized. Impor- tantly, IWL influences ICL acquisition only through (f) and c2()). which in turn depend only how the MLP memorizes class labels. We proceed with analysis by retaining (t) and C2(t) yet-to- be-determined MI MI.--pecific dynamical "order pa- rameters", keeping in mind that their dependence t and K will play an important role in our analysis further below.
-2.
0.0
w
2.5
2
5.0
Figure 4: The approximate ICL loss land- scape C (fixing ZMIP 0) the minimal model as function of the key parameters 3, w exhibits a nearly flat region close to ini- tialization, but the dynamics always leads to ICL acquisition (w, ,3 >> 1).
Gradient descent dynamics over the loss in equa- g
tion 6 gives
dw Ci (e--c2w)
dt N
d3
(wee).
dt
Learning initially proceeds at slow rate c//N (since N > 1 and ea 1). Since ot average increases as the MLP memorizes, C1 decreases and slows down ICL acquisition. If the MLP (near) perfectly memorizes the K item-label pairs before ICL is acquired, then ICL is never acquired. In other words, the loss "explained away" due to MI.P memorization creates an effective competition between IWL and ICL acquisition despite the additive contributions of the MLP and the attention head to the logit, Since 0 < < 1, equation 7 shows that w eventually converges from its initial value a positive value w e /cz. increases monotonically when is positive until ICL is acquired. Thus, equation and equation 8 jiplyy that ICL will always be acquired, however slowly, the MLP unable to perfectly memorize the K item-label pairs (i.e., e, (00) > 0).
However, the choice of label statistics in the context matters. For example, consider the case when N is even and there are exactly N/2 tokens with +1 and -1 labels in the context. To compute the mean loss L' in this scenario, we set 1) = 0 in equation and Taylor expand w.r.t (c 1)/N to get
L' log(1 + ++
C1 w (c -1).
(9)

---

Published as a conference paper at ICLR 2025
The origin is saddle point. The parameters either flow to the ICL solution (w, B > 0) or to suboptimal solution (w, < 0) depending on the initial values Wo, Sp. That is, ICL acquisition is not guaranteed. Intuitively, when n is binomially distributed, the network learns that the target's label is more likely to be a +1 if there are more +1 labels than -1 labels in the context. This bias, however small, pushes w positive value and leads the network into the ICL. basin. When the numbers of +1 and -1 labels are forced to be equal, could flow into the basin that leads to ICL or flow to an alternative potentially suboptimal solution. We revisit the n. 2/2 case in Section 4.
3.5 EXPONENTIAL DEPENDENCE OF fICL ON INITIAL CONDITIONS
Equation 7 and equation 8 allow us to estimate the number of iterations it takes to acquire ICL (Appendix). Note that "time" here is a proxy for the number of iterations, which we can only determine up to a constant pre-factor. Exact integration of equations 7 and [ is infeasible, but an approximate expression is obtained when |wo] 1,5g] < We fix w/ = hereafter, though the more general case of we * 0 can be solved (Appendix). We find that the number of iterations it takes for ICL acquisition (denoted TK) satisfies
2 Po [g(TK), where (x(1) 2 (x f (a(()dt.
(10)
The subscript K introduced to highlight that e1 depends on K.
We first consider the case K is distributed symmetrically around O, in which case (1(t) = (a(-))))<+ Rs 1/2 and (t) = t.
that the MLP is unable to memorize D. The MLP logit o+
Solving for Too (which we call (ıC1. hereafter) using equation 10 we get
tICL
2 -8
(11)
The dynamics are qualitatively different when - -80 is large and e < 1. In this case, we obtain tICL Ne 29c (Appendix). We numerically verify the exponential dependence of fICL on the initial values of 30 (Supplementary Figure Ab).). A consequence of this exponential dependence is that normal-din dormal-disributed values of 30 will lead to a long-tailed distribution of fILL- In pictorial terms, due to the nearly flat loss landscape close to initialization (Figure 4 small variation in the initial parameters wo, 30 leads to large variation in when ICL is acquired.
3.6 MEMORIZATION SCALING LAWS AND THE TRANSITION FROM MEMORIZATION TO GENERALIZATION
Equation 10 shows that the behavior of an MLP-specific quantity, e1 (t) (via Ik), determines when ICL is acquired for different values of K. It is useful to introduce the quantity Ix (oo), which can be interpreted as the time taken for the MLP to memorize a dataset of size K. Equations and and together with the monotonicity of I (t) imply that ICL is acquired if
tICL < TK (co).
(12)
We delineate two distinct mechanisms depending whether (xc) is finite or not:
1. Capacity-constrained:: We call the network capacity-constrained if [x) diverges + co, i.c., the network never fully memorizes the dataset. Equation 12 then implies that the network generalizes when K Koe. where Koe is the smallest K at which the network is constttec.
2. Differential learning kinetics:
possible that Ix (oc) ( is finite. In this case, the network transitions from memorization to generalization at K = K" such that /IC.. Tx- (co). In other words, when K > K*, it takes longer for the network to memorize the dataset (even though it has the capacity to do so) than it takes for the network to generalize. We call this case the differential learning kinetics regime as the relative rates at which the network memorizes and generalizes determine when the transition occurs.
The divergence of (t) as + 3o may occur either because the network has limited capacity to memorize the K samples or because of the data distribution. For example, if the rank-frequency distribution of itbe-label pairs follows a Zipf's law p(f) f-a with exponent a S then the network's 's loss is dominated by rare item-label pairs that are not memorized. Previous work has shown that such skewed data distributions indeed favor ICL acquisition (Chan et al. (2022) )
7

---

Published as a conference paper at ICLR 2025
We examine the behavior of Ix for uniform distribu- tion over item-label pairs. To our knowledge, current deep learning theory does not inform how the distribu- tion of logits (or summary statistics such as c1) scales with and K for typical MI.P architectures. To make progress, we empirically measure e1 (t), c2(t) for dif- ferent K using an independent (Supplementary Figure A.2). We find that ea (t) for K between 500 and 50000 decays fast enough such that its integral T ((x) is finite (Supplementary Figure A. Our MLP is thus not capacity-constrained. Further, we uncover scaling law, Tx (00) K", where z 0.7 (Figure 5) From equation| [2, using the expression for fICL in equation [1 and the scaling law [x())) K" leads to an estimate for the task diversity threshold K for the transition from memorization generalization:
Slope = 0.70
[x((o))
10
of MLP experiments
10
10
K
Figure 5: Ix (oc) shows power-law scal- ing with with exponent 0.7.
K N 80/"
(13)
up to a constant prefactor. Note the exponential dependence on Sp. whose random initialization implies that there range of K for which the network will either show memorization or gen- eralization. The theory further predicts a non-trivial power-law relation between the task diversity threshold and the context length.
3.7 SLOW IWL EXPLAINS TRANSIENT ICL
We now explain why transience appears in our minimal model (Figure 3P) when the attention head is regularized more heavily compared to the MLP. For simplicity, we impose L2 regularization with parameter Aw only on w. We return to equation 5 for the loss, which applies throughout training. Since w, 3 > 1 after ICL is acquired, we simplify equation to obtain (Appendix)
Aw
L ~
(14)
Once ICL is acquired, memorization slows down dramatically due to the small factor e-"., Without L2 regularization w, continues to increase (at decreasing rate) and ICL is not transient. However, when w is regularized, w after ICL acquisition tracks Whr, where
(t)( A W (t)(t)))w), ca(t) =
(15)
The Lambert W function (xx) is monotonic in when is positive. C3 decreases the network memorizes dataset (Figure A. Thus, Wir decreases as c3 decreases. Whr decays to zero (and ICL fades away) when the dataset sufficiently memorized, i.e., when e3 < A Thus, the analysis suggests that extremely slow memorization coupled with regularization leads to ICL transience. We note however that in complex models the effects of a global regularization parameter different sub-circuits are hard to disentangle, which may explain the puzzling observations in (Sigh et al. 2023).
Equation 15 hints at relationship between the loss on ICL sequences (LICL) and the loss of IWL sequences (CTw1.) after ICL acquired. We use heuristic argument (Appendix) show that
LIWL R3 - Lig(Cicu), when CICL < 1,
1 LICL - log log(CCwL.) when Liw, 1.
(16)
These approximate relations between LICL and LiwL consequences of our two ansatz. If our ansatz are valid, the theory predicts that these relations should hold from the moment ICL acquired until it fades due to gradual IWI..
8

---

Published as a conference paper at ICLR 2025
(a)
Task diversity threshold vs context length
(c) Bimodality at the threshold
(d)
ICL IWL losses
2.5
10000
10
2.0
Crwx. -- Cloo)
Liws -
IWI. Ioes
1.5
&4 10t
Tiee Faclated
1.0
0.5
0.0
0
30
10
2
10
10P
10P
ICL Juee
(b)
Time to acquire ICL
ltenations
(+)
(e)
Contextual statistics
Fraciion di zuns anquite ICL
Lincar Fil
1.0
Shandan
itersionns X309)
Melista
Courr
75
0.0
80
0.2
25
0.0
100
150
200
250
0
Loc. (in ilssaaions x11))
K
Figure 6: (a) The critical task diversity threshold K* exhibits a power law relationship with respect to N. The experimentally determined critical exponent (lincar fit) closely matches predicted critical exponent 1/v where v z 0.7 (predicted fit). (b) The median time taken to acquire ICL (fICL) scales linearly function of N and the distribution of fILL is long-tailed. For clearer visualization of the full distribution, see Supplementary Figure A.4 (c) Close to = K", depending on the initialization seed (red and green dots), the network will either generalize (green) not (red). Dif- ferent seeds are predicted to show either small fICL. or do not acquire ICL, with few intermediates. A histogram of lICL for K 9666 confirms this prediction at the critical task diversity threshold K Note that the maximum number of iterations is = 4.5 x 10 (d) In runs that exhibit transience, we observe relationship between ICL loss and IWL loss after the model has acquired ICL, closely matching our predicted functional relationship.
(e) We measured the fraction of solutions that ac- quired ICL as we vary K, by training at least 20 models with different seeds for each K, on new dataset D' where every sequence is constrained to contain exactly N/2 items of each label. We ob- serve that the critical task diversity threshold K" greatly increased and that many more solutions fail to achieve ICL compared to models trained our standard dataset D.
4 EMPIRICAL VALIDATION
The theory makes a number of new quantitative predictions related to ICL acquisition. We em- pirically six nontrivial predictions that span various aspects of ICL phenomenology using the original transformer model in equation !
Power-law scaling of the task diversity threshold with context length. Equation 13 predicts highly non-trivial power law relationship between the task diversity threshold K and N. To test this prediction, we train our original transformer model (equation
varying N and K. At each N, we observe sharp transition from memorization to generalization as K increases (Supplementary Figure A.7 For each N, we determine K* by fitting a sigmoidal curve ICL performance as function of K. As RJ 0.7, Equation 13 predicts an exponent of 1/v 1.43, closely matching
measured exponent RJ 1.41 (Figure 64).
Linear scaling of the time taken to acquire ICL with context length. Equation 11 predicts that fICL (time taken to acquire ICL) scales linearly with N. To test this, we train our original transformer dataset D at every
model (equation I at varying N and take the limit K -> by resampling training iteration. We then determine fIC as the epoch at which ICL accuracy exceeds 95%. We train 100 seeds for each N to obtain the full distribution of lcL. Figure 6b confirms linear

---

Published as a conference paper at ICLR 2025
relationship between the median fICL function of N. We verified that the linear relationship also holds for two-layer transformer (Supplementary Figure A.5)
Bimodal solutions near the transition. For K near the transition, K K" equation 13 predicts that independent runs will either have f1CL. > 1x. (oc) or fICL. < 1x- (oo) depending on the net- work's initial parameters. In the former case, ICL is never acquired. The nonlinear form of Ix (t) (illustrated Figure G top) suggests that in the latter case, fICL. will be short, with few intermedi- ates in between. We verify this bimodal behavior of the solutions around the transition by generating runs with R 250 seeds near K 2 K" [Figure bottom).
Long-tailed distribution of the time taken to acquire ICL. Equation 11 predicts long-tailed distribution of EICL. We verify this prediction in histograms of tICL. for each N, using the same data as in Figure 6p (Supplementary Figure A.4).
Interdependence of the ICL loss and IWL loss after ICL acquisition. Equation 16 predicts non-trivial relationship between the ICL loss LEC, and the IWL loss LIwL after ICL is acquired. In function of ICL performance
runs in which we observed transience, we plot IWL performance and observe close match between equation 16 and data (Figure (+)).
Failure of ICL acquisition when every sequence contains N/2 tokens of each label. Equation 9 predicts that the network is more likely to fail to acquire ICL. when there are exactly N/2 tokens with +1 and 1 labels in the context. We test this prediction in dataset D with N = 100, where each training sequence (including the target token) viewed by the one-layer transformer has N/2 tokens of each label in its context. For each K, we trained the model using s 20 seeds. We measure the probability of acquiring ICL as the fraction of seeds whose final ICL performance exceeds 75%. We observe much lower probability of acquiring ICL for models trained D' compared to those trained on our standard 'D (6)) Moreover, the loss curves across seeds are diverse, with many seeds not acquiring ICL. and some saturating sub-optimal solutions (See Supplementary Figure A.8)
5 CONCLUSION
Here, we propose a theory based on the ansatz that the network contains sub-circuits that are inde- pendently involved in memorization and generalization. A trade-off arises simply because the rate at which one sub-circuit is optimized depends on how much loss is already explained by the other sub-circuit(s). Building on this theory, show that the transition from memorization general- ization our model is determined by the relative rates at which these sub-circuits memorize and generalize. However, the theory does not rule the possibility that capacity constraints play a role in other scenarios.
This ansatz, despite being cast in the context of simplified one-layer model, explains a surprising variety of ICL-related phenomena observed with much larger models. These include a tailed distribution in when ICL is acquired, an MLP memorization scaling law, the bimodality of solutions at the task diversity threshold, the transient nature of ICL, amongst other novel quantitative relations that our theory identifies. The two most striking predictions are (1) the non-trivial relationship between an MLP-specific memorization scaling law (Ix (oo) ~ K") and task diversity threshold scaling law w.r.t context length (K* N1/²), and (2) the long-tailed distribution of when ICL is acquired and its linear scaling with context length (IG). ~ N). Both these predictions have been validated in experiments.
Our results offer some hope that seemingly intractable phenomena observed in large models can be reproduced and analyzed using simpler, tractable models through careful experimental design. How- ever, further work is necessary to examine to what extent such insights provided by small models remain valid for larger models (that potentially contain many sub-circuits) and for more naturalistic tasks (where a clear distinction between memorization and generalization cannot be made) (Min et al. (2022); Wei et al. (2023): Pan 2023) Shi et al. (2024)).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Kenneth A. Norman, Albert Qin, Cole Gibson and anonymous reviewers for support and valuable feedback on the manuscript. AN is supported by NIH grant RFIMH125318 GR is partially supported by joint research agreement between Princeton University and NTT Research Inc.
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. 00297, 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 arXiv:2211.1566
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: 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:18 33:1877-1901, 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 in transformers. Advances in Neural Information Processing Systems, 15:18878-1888 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 as meta optimizers. arXiv preprint arXiv:2212. 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.002: 2022.
Benjamin Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint wwXiv:2402.11004, 2024.
Shivam Garg, Dimitris Tsipras, Percy 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.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprint arXiv:2212.0044: 2022.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202. 12837, 2022.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 11895, 2022.
Jane Pan. What in-context learning "learns" in-context: Disentangling task recognition and task learning. Master's thesis, Princeton University, 2023.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. arXiv preprint arXiv:2306. 15063, 2023.
11

---

Published as a conference paper at ICLR 2025
Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in in-context classification task. arXiv preprint arXiw:23/2.03002, 2023.
Zhenmei Shi, Junyi Wei, Zhuoyan Xu, and Yingyu Liang. Why larger language models do in-context
learning d 2024. URL htttttt///axiiiiit org./vv.rrraaaao wttpg//a//////////
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 wrXiv2231.08360 2023.
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. In International Conference on Machine Learning, pp. 35151-35174 PMLR, 2023.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint TrXiv:2211.00555 2022.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303. 03846, 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.
A APPENDIX
We present a detailed analysis of the minimal model outlined in the main text. Suppose the (D 1)- dimensional imput tokens are 11, i2,... .., ,tN+1, where t; = (2z, L,) for i S N and fN = (IN+1, In the minimal model, we consider two logits 2MLP = 0(xN+1) and
N ZATT
BrfaN ax
G Sa B wly.
(17)
J=1
is multi- layer perceptron (MLP), which throughout our paper a threc-layer ReLU network with hidden dimension 512. The final logit a used for classification is the sum of the contributions from 2MLP + ZXTT- The minimal model can be obtained from the
the MLP and the attention head, one-layer transformer model by (1) removing the LayerNorm operation, (2) the interaction strength ansatz, i.e., by assuming the query, key and value matrices in the attention head have the form
KTQ (BIDDD 01xD
0DX1 01x1
ODXD 0 W)
V
(18)
and (3) the independence ansatz, where the residual term ØN+1 is processed by the MLP to produce ZMLP. the output of the attention head is 2xTT and these two logits are summed to produce the final logit However, we stress that the minimal model serves phenomenological model and is not derived from the one-layer transformer model.
To reproduce the phenomenology shown in Figure we optimize B, and the parameters of the MLP é using the same procedure used to train the full model. In particular, we use stochastic gradient descent (SGD) with batch size 128, learning rate 0.01, weight-decay 10-10 D 63, 100 and MLP hidden dimension d = 512. To reproduce transience in Figure 3p with a fewer number of iterations, we increase the weight-decay parameter on w and 10-3
A. THEORETICAL ANALYSIS OF THE MINIMAL MODEL
To derive analytical expression for the loss landscape of the minimal model, assume (1) that only one copy (say x.) of the target x is present in the context, (2) that 1 = if = c and O otherwise, and (3) that the distributions of logits obtained when the ML. is applied to all the
12

---

Published as a conference paper at ICLR 2025
items in D with labels +1 and -1, (say, (Ø) and P (Ø) respectively) are identical except for the sign: P* (0) = P (-Ø).
These three assumptions are justified when K > N > 1 and D oc. Note that in the simulations presented in Figure 3 we use K > 10', N = 100 and = 63. Assumption (1) is justified when K > N it is unlikely that more than one copy of the target is sampled from the dataset DD [ID)= = K) in a single context of length N. Assumption (3) is justified when > as all the items in D statistically identical irrespective of the label. Assumption (2) justified when D + oo based on our sampling process, which will ensure that in this limit the norm of each item is one and different items are orthogonal. For finite D. our theory will break down when N is sufficiently large, though the precise scaling relationship between D and N for when the theory will break down remains to be examined.
Let c be the index of the target's copy in the context. Under the assumptions stated above,
e
1
ZATT =
w
e + NNN-
fot
ea + - Σ Ej
jot
8 lo
w
8 19+77 N - ( - N - -k-))
(19)
+ N
where n. is the number of tokens which have label +1 amongst the N tokens in the context. Sup- pose split the dataset D into two datasets D and D containing items with +1 and -1 labels respectively. The MLP (at any particular stage of training) when applied to the items from D produces two histograms of logits (*))
Recall that the binary cross-entropy loss for given input sequence is - log a (fa + =ATT)). The average loss is obtained by averaging over the cases when the target is drawn from D+ and D and by taking an expectation over n., (which follows a binomial distribution, B(n,)) )). Given that there n.+ items with label +1 in the sequence, the target's label is le =1 with probability n+/N and 1 - / respectively. The average binary cross-entropy loss is then
C ==
Σ
(logo(0t B(n+) (10go (0t (
(2m+ - N -N-1 ptapt
N
c + N
+ N
N
nh+ B(n+)) logo (-0- + N
(N 2n -0--1)
c +
3 + N
(20)
where the two sums consider the cases when the target's label is +1 and 1. Given the symmetry between positive and negative labels, the contribution to the loss when fe -1 (the second sum) is approximately equal to the loss when le +1 (the first sum). Formally, using P (Ø) P (-Ø), replacing n with N n4 in the second sum and using the property B(n+) B(N n4), we get
Σ m==0
N
- nt B(n+) (Doso (0= + @ N
e
+ ② + - (2n+ - N -11)
1
6ttPt
(21)
When N >> 1, n4/N is approximately cauussian-disributted n4/N ~ 1/2 + n/2/N, where n ~~ 10. This gives
C ((1+
e -
N N e c N
logo 0+
+
e N
--N(0.),
(22)
which is presented in the main text (equation 5
A.2 EXPANDING THE LOSS BEFORE ICL ACQUISITION
w, B are initialized at wo So where /wol. | < 1. Define 7 e - 1. Close to initialization, w, B satisfy 7/N & 1 and /wl/VN & 1. We compute the time it takes for the network to acquire ICI..
13

---

Published as a conference paper at ICLR 2025
ICL is acquired when w, B > 1. From equation 22 since 7 < N, we have
L= - ((1
log a
Φ+ + w
7 .N
+
"
(23)
VN
VN
where the distributions from which 77 and @+ are drawn has been dropped for notational convenience.
Since 7/N << and /w// wul/VN < 1, we have N + two Taylor expansions, we get
y VN
< Using (xx) 11(11 E) and
1 +ee 0l
-oll-m((+ (3+ {tx) (3+++t
-w
C ~
1
+
log
(24)
N
1 1- w
7 +
"
w²
7
7
+
log
N
VN
2
N
VN
7/++
(25)
< 1og((1 te =) + ((1+
7
w
7
77
log
1 -
+
+
+
VN
N
N
2
N
VN
(26)
[og(11+ (1og(L+0-07)
C1
(1+ 7)
c2
(27)
N
2
(1og(1 +e -0)11+)0)))
Ci
e w
c
(28)
where we define the two MLP-specific dynamical variables,
= (a(- ))
(29)
C2 1- 1ool-0"0 )2)
a+ /c1.
(30)
In the second to last step, we have expanded log(1 11 + x) for small æ, computed the expectation over 1) and ignored higher order terms as the 1 1/N term is the dominant contribution (the 1/VN
contribution vanishes in expectation). Equation 28 is presented the main text (equation 6
(x). er is one minus the probability of correct classification, averaged over the dataset. At initialization, the logits are centered around Øi+ x so that a R 1/2 and therefore C1, C2 ~ 1/2. ot increases on average during training, which means e₁ decreases and 22 increases during training. C1 = 0 and 02 = in the idealized scenario where the dataset is perfectly memorized. We train an MLP in independent experiments to memorize D with the same hyperparameters used for training the one-layer transformer. We verify empirically that our MLP can near meear-prfeecly memorize the dataset for dataset size K up to 10 (Figure/A.2 and Supplementary Figure A.3pu
Since a--x)
A.3 ICL ACQUISITION
We consider ICL to be acquired when w. 3 > To compute the time taken to acquire ICL, we consider gradient descent dynamics of w and B close to initialization. From equation 28 we get
dw
C
(e
-cw),,
(31)
dt
N
d3
Ci
(we²).
(32)
dt
Numerical simulations (Figure A.1 with c1, 02 fixed at 1/2 reveal qualitatively different dynamics depending on whether [,3p] < Bo 1. Note that _ - < N in both these cases, that our
approximation equation 28 is still valid. We first examine the relevant case |So| < 1. The analysis of the second case -0 is presented further below for completeness.
Equations 31 and 32 cannot be solved exactly and we resort to approximations. When |86l <
and [Wn] < 1, we have e 2s 1 and [w] < 1 close initialization. By definition, 02 1. Put
14

---

Published as a conference paper at ICLR 2025
together, the term on the right hand side of equation 31 can be ignored close to initialization.
Then, du /tt s (ci/N)e". Dividing this by equation 32 and integrating, we get the conservation
equation w²/2 /2 3 w fo. so that w = 2 - 8 + wg/2). Substituting this expression
for into equation
32
we have
dB
C1 2(3
3o + w
(33)
di
N
Integrating this equation from B = Bu to = 0o allows us to estimate the time fICL it takes to acquire ICL. Performing this integration, we get
NV2T(1/2, 2)e 3 2--8812 Ix (fiC..)
(34)
where I(...) is the upper incomplete gamma function and we have defined Ix (t) 2 2 e1 (t'))
The subscript K is introduced to highlight cj 's (and thus I 's) dependence on K. In the main text, we examine the case when wp so that (ICI. satisfies
2 ss (x(Cc))
(35)
When -30 >> 1 and Wol < 1, the dynamics be split into two parts: 1) w converges from its initial value to the nullcline, dw/d = 0, where w = e /c/c2. 2) w and 3 gradually increase along this nullcline until/ is large enough that the exponential dependence in equation [8 leads to abrupt ICL acquisition. The first regime is shorter and fICL is dominated by the duration of the latter regime. In the latter regime, since w e² /c2, we get
d3
C1
29
(36)
dt
No
Integrating this equation, we get
where [xx(t)= = L/ (t)) C2 (t)
Ne- 280 T (fCL),
de'
(37)
The e e and 200 scalings of MCL for |ol < and -Bp >> 1, respectively, are consistent with those obtained when equations 3T and 32 are numerically integrated beginning from wp 0 (Figure A.I. In these numerical simulations, we fix = 1/2, c2 1/2, which is equivalent to setting ZMLP 0 (i.e., the MLP does not contribute to the logit).
A.4 TRANSIENCE
We examine the influence of applying L2 regularization to the parameters of the attention head. For simplicity, we assume regularization (with Le regularization parameter Aw) is applied only to w. Our goal in this section is to show that such regularization parameter (however small) is necessary to induce transience, and to delineate the values of A for which transience is recapitulated in the minimal model. A sufficiently large regularization parameter will of course also affect ICL acquisition (Supplementary Figure A.6). Further analysis could delineate the range of values of X for which regularization will have significant effect ICL acquisition; this analysis is beyond the scope of this paper.
We write down the dynamics of w after ICL is acquired. Re-writing the expression for the loss in equation 22
C cx-((11
((+
e e + -
T/N e + N
logo
1+++ + w
(38)
VN
n
Note that the terms involving inside the logarithm are at most of order 1/VN. Since w, 8 > after ICI. is acquired, the first term (e 1 (e" + - 1) involving This leads to
will be much larger than the term
C R ((1+
17 VN
log (0t + +w)
(39)
15

---

Published as a conference paper at ICLR 2025
Since (n) 0, we have
L ~ log (ot + +w)) 0t
(40)
log 1 + e (wt
(41)
(--0)) c_
(42)
where we have used e e- < 1 and log(1 + x) for < 1. Define C3
+
o+
which is our third MLP-specific dynamical variable. 03 will decrease while the MLP memorizes the dataset. Equation 42 implies that, without any regularization, w will continue to increase (albeit at decreasing rate) and ICL will not be transient.
Introducing an L2 regularization term, we have
2 Cge " + uw
(43)
2
The gradient descent dynamics w is given by
dw
Cze An
(44)
dt
With regularization, once the network sufficiently memorizes the dataset (i.e., C3 sufficiently small), reaches a steady state, denoted Whs which satisfies uhe c3/ Re-expressing this relationship in terms of the Lambert W function, we have
Wa R W (cs/2w)
(45)
Since Cg = (e-
d+ we have c3 at initialization.
63 will continue to decrease (however
slowly) after ICL is acquired while the loss in equation 43 is minimized. X is typically chosen be small value. Due slow memorization (declining c3). wa slowly declines while tracking equation 45 ICL fades when Wy is of order one; since ((x) = 1 when æ z= 3, i.c., some number of order one, ICL fades when C3 R Au- Our MIP simulations show that 03 can decay further than 10-5 (Supplementary Figure A.2
This analysis suggests that ICL transient in our minimal model when the network sufficiently memorizes the dataset (c3 is sufficiently small). Importantly, whether ICL is transient or not depends both on the extent to which the network able memorize the dataset and how much regularization is applied to each of the network elements. In transformer networks, how regularization may affect each sub-cireuit is hard to disentangle. For example, it is possible that the output of the attention head passes through the MLP before producing the final logit. Increasing the regularization MLP parameters may decrease the rates at which the network memorizes and generalizes.
A.5 THE RELATIONSHIP BETWEEN ICL AND IWL LOSS AFTER ICL. ACQUISITION
Recall that IWL is measured on sequences where the item-label pairs are drawn from D but copy of the target item is (most likely) not in the sequence. This allows us to define an IWL loss LIWL. which corresponds the binary cross-entropy loss of the network on such sequences. Similarly, we define an ICL loss LICL, which is the binary cross-entropy loss on input sequences that contain N novel itee-label pairs and where the target is one of the N items in the sequence. Equation hints at relationship between LICI. and LiwI, after ICL is acquired. We now provide a heuristic argument to derive this relationship.
Immediately after ICL is acquired, the ICL loss is small, LICL < 1. Our arguments in the previous section show that ICL will be transient the MLP that can sufficiently memorize the dataset and appropriate regularization is applied. Once ICL fades and the network pear-perfectly memorizes the dataset, we expect LIWL < 1.
We first consider the when LICL < 1, that is, after ICL is acquired but before it fades. ICL acquisition corresponds to w. 3 >> 1. Recall from the independence ansatz that the final logit is ZMLP + ZATT- IWL sequences do not contain copy of the target in the sequence. Instead, since
16

---

Published as a conference paper at ICLR 2025
>>> 1, the target 2N+1 will pay attention to random item, say Ti, in the input sequence that maximizes the dot product 2 2N41-I, Item i's label & matches the target's label le with probability half. In the case that it does match, 2xTT we; we.- Since the MLP has not yet memorized D, the typical logit from the MLP is small compared to w, |EMMP| & w. The IWL. loss when l; = l. is 2ATT))
-log (Ce the alternative case, & = -l. the IWL loss is log a (wlflo) log (- log + e")) R w, where we used w > 1 in the final step. Since Ice with equal probability of 1/2, the expected IWL loss is LIWL. 2 w/2.
- llog (fc(wl:))
log (wr) which is negligible since >> 1. In
On the other hand, ICL sequences do indeed contain a copy of the target in the sequence. In this case, ZATT = wle Since ZMLPl & w. the ICL loss is LICL = log (ogoo(L((MM))) + zATT)) s log o (L.((( x log(1 + e==) R e-", where we used log(1 1 + x) x for & the final step. Putting the expressions for LIwI. and LICI. together, we get
LwI - log LieI.
(46)
after ICL. acquisition and before ICL fades away.
We now consider LIWL. < that is, after ICI. fades away and the network significantly memorizes the dataset. In this case, since |w], | < 1, we have |EXIT| < |aMMP]. The loss on ICL sequences dominated by the logit produced by the MLP novel items. Suppose that the typical magnitude of the logit produced by the MLP is Ş. The IWL loss is LIwI. - -log ( = log(1 + + E). LiwL. < 5 For novel items, the MLP will produce either + irrespective
implies > 1 and thus LIw of the true label of the target. The ICL loss is either negligible or log( + ef), depending on whether the sign of 2MPP matches the target's 's label or not. Since & > the ICL loss averaged over these two cases is LICL R E/2. We thus get
1
LeL ~ - log LIWL
(47)
after ICL fades away and the network significantly memorizes the dataset.
A.6 SUPPLEMENTARY FIGURES
(a)
(b)
0.0025
Uie io 2-2
Probudriity iiy dlrssity
0.0020
0.0015
102
0.0006
10
0.0000
500
1000
1500
tro1 (a.1))
Figure A.1: (a) tICL exhibits exponential dependence on the initial values of 3g. (b) Histogram of lic.. in the minimal model given normally distributed values of 3o. The distribution of 4ccr. is long-tailed due to its exponential dependence on 30-
17

---

Published as a conference paper at ICLR 2025
(a)
(b) 0.5
(c)
10
10
10
10
0.3
10
0.2
10
0.1
10
10
10-5
10
10
10
10
10
10
10
10
10 10 10
10
10
10
Ioenttion
Figure A.2: MLPs were trained to memorize K item-label pairs drawn as in our transformer exper- iments. We plot the time evolution of three MLP-specific "order parameters". Each plot shows runs from three random seeds. (a) c1 = (ar ( )) (b) C2 = (a(- /c (c)
ca(1)
The variability across seeds is negligible.
(a)
(b)
10-
10 -2
K
10
K 300
Ix(0)
/ 1622
10
2924
10
K
K 17008
10
90806
10
55504
10
10
10
10
10
10
10
10"
10"
10
Iterations
Iterations (t)
Figure A.3: (a) MLPs were trained to memorize item-label pairs drawn as our transformer
experiments. We plot C (e-*+) for different K. (b) Time evolution of Ix(t) = 2 S (t('))dt (upto constant prefactor) obtained by numerically integrating the curves in panel (a).
18

---

Published as a conference paper at ICLR 2025
2V =:00 50
N = 100
N = 130
(8)
(9)
(6)
0.3
0.20
0.15
Protabiiity
Probabaiity
Photubiitty
0.16
0.2
0.30
0.3
0.05
0.3
0.22
11.3
0.4
10.5
0.0
015
1.01
1.5
0.00
025
1.00
3.5
20
250
Sycz.
Iterations
xx1)))
x10")
1100
xc10
200
N 250
0.211
0.16
0.16
Pretetiltt
Probabiity
0.10
0.10
0.06
0.00
0.00
(in
x10P)
Figure A.4: For each N, trained the model across at least 130 seeds and show that histograms of tICL are long-tailed.
X 00:00
2 13
150
Linear
8.23
fic0. (in iterations x10")
X--
A15-
50
100
150
200
250
Net Proet x1P)
Nox Pestasa kLPP)
N
Figure A.5: For two-layer transformer, we replicate the finding that fic scales linearly with context length N. We trained the model across at least 29 seeds for each N and show that histograms of lIT are broadly long-tailed.
19

---

Published as a conference paper at ICLR 2025
N/A
+++
Transient itenatio
00
200K
10
3 10
x 10
x 10
Woight Doeay (x1009)
Figure A.6: To explore the impact of L2 regularization on ICL transience, we trained models while varying the weight decay Au applied the self-attention layer. For each run, compute the iteration at which we observe transience, which we define as the first iteration after the model has attained least 99% ICL accuracy and after the model's ICL accuracy has declined to 90%. The red plus signs indicate runs in which the model never attains the ICL solution (never attains at least 99% ICL accuracy). The yellow crosses indicate runs in which the model does attain the ICL solution but we do not observe transience, as the time to observe transience is too long. The blue dots are runs where we do observe transience, showing that increasing A. induces faster transience.
(6) 1.0
(6)
1.0
(6)
1.0
0.5
0.3
0.9
[CL MATIRnCy
[CI andiniyy
ICL ansaawey
0.5
0.5
0.8
1- 3EE 03
10
107
10
207
(4)
200
250
1.00
1.0
JCL uccInas
1CL accunee
1CL M03uter
0.5
Spevido
0.7
Rewddo7h
0.7
Beas
0.0
0.0
0.5
0.3
K
Figure A.7: ICL performance is measured for various N and K. For all N, we see a sharp transition from memorization generalization K increases. We fit sigmoid to the data determine the critical task diversity threshold K* for each N.
20

---

Published as a conference paper at ICLR 2025
1.0
0.9
ICL accurcy
0.8
0.6
0
2
Training iterations (x10*)
Figure A.8: ICL performance curves (for different random number seeds) for models trained our modified D', where every sequence contains exactly N/2 tokens with +1 and -1 labels for K 31336. We novel intermediate suboptimal solutions, and many of the runs never acquire ICL.
21


Reviews:
[{"review_id": "b9lHFNm65C", "reviewer": "Reviewer_6toB", "summary": "The paper explores the mechanism by which models acquire the ability to generalize and memorize. \nThe authors test the hypothesis that generalization arises because, in models with limited capacity, a preference for generalization over memorization emerges when task diversity is sufficiently high.\nThe authors find that generalization is not solely the result of capacity constraints but also due to different learning kinetics between sub-circuits responsible for generalization and memorization.\nAt sufficient task diversity, the sub-circuits responsible for generalization can be acquired faster than those for memorization.\nThey also find that a network can lose generalization abilities in favor of memorization with sufficient training under $L_{2}$ regularization.\nThe paper identifies a number of properties about the dynamics of generalization and memorization. \nAmong them are how MLP-memorization scales with task diversity, how the task diversity threshold scales with context length, and the sensitivity of generalization to initialization.", "questions": "Can the authors use this work to recommend methods to improve generalization? For example, can we choose an initialization method that might better support ICL?\n\nThe theory hints that ICL decays due to the use of regularization. Though using $L_2$ regularization makes the results more comparable with typical training, would it be possible to provide empirical results without regularization?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The way models acquire generalization during training is an interesting topic and the paper provides an insightful exploration of the mechanisms involved. \nThe authors convincingly support the premise of the paper with theory and empirical studies.", "weaknesses": "Because the model used for the empirical validation is stripped down to a one layer transformer model the authors note that they cannot be sure that the dynamics are the same for much larger models.\nThe authors could consider supplementing the one-layer analysis with tests on a deeper transformer model. \nEven if constrained by computational resources, experiments on a 2 or 3 layer transformer would help assess the applicability of their findings to larger networks.\n\nThe paper does not recommend methods to facilitate generalization. If possible, a discussion of how the ideas in the paper may be used to develop initialization and regularization techniques would make the submission even stronger.", "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": "IKBoZ9o0mk", "reviewer": "Reviewer_EVew", "summary": "This paper presents a small transformer model performing a synthetic task aimed at separating in-context and in-weights learning. A surrogate model which is amenable to analysis is also presented. The paper aims to combine the analysis of this ``minimal model'' and the empirical analysis of the original transformer to understand under what conditions in-context learning emerges. The training dynamics of the minimal model are characterised and used to analyse the loss landscape of the in-context pathway of the model. Further, a number of phenomena of transformers performing in-context learning are considered, such as the transient nature of this capability.", "questions": "- What is the $P^+$ value used to plot Figure 4 and how does it affect the loss landscape?\n- Figure 5: the power law is presented as being between $K$ and $I_K(t)$ but then Figure 6a presents it as being between $N$ and $K^*$. How do these two representations relate?", "limitations": "", "rating": 3, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "# Originality\nTo my knowledge this is a new combination of ideas from the literature. The tasks and network structure similar to the work of Chan et. al. (2022) is combined with the analysis of the loss landscape and loss dynamics in various high-dimensional limits (while it is not mentioned in the work, the analysis has D and K very larger - the input dimension and number of datapoints - which constitutes the well known thermodynamic limit). By working in this limit much of the dynamics of training become tractable and this is exploited here.\n\n# Quality\nThe tasks and model used here is well motivated and clearly relevant to in-context learning and useful for the aims of the work. The hypotheses or motivation of the work is clear.\n\n#  Clarity\nOverall use of language in this work is clear and it is well written. Figures are neat and legible.\n\n# Significance\nThis work touches on a number of established phenomena on in-context learning but also introduces new phenomena which would be useful for further research. The model itself could likely be built upon in further work.", "weaknesses": "# Quality\nThe work is vague and lacks the necessary degree of rigor. The fact that no full derivations are presented even in the appendix is already grounds for rejection in my opinion. I recommend the authors provide a more thorough presentation of the derivations in this work. \n\nThis somewhat casual approach to the mathematics means that some very important simplify assumptions are not described accurately and in some case the work changes assumptions or conditions to suit whatever point it is currently trying to make. The most clear example of this is the dimension of the input D. It is foundational to the analysis here that D is large enough that the key-query matrix will be diagonal. This appears to me to be a huge simplification, worthy of clear mention. But the work fails to even acknowledge this in Section 3.2 when describing the model, instead referencing the fact indirectly and terming it an ``independence ansatz''. This fact of needing infinite D is only brought up in Section 3.3. This lack of clear assumptions and formal definition results in the jump from Equation 3 to Equation 4 being unclear. A similar case can be made for the implicit simplicity of the setup of binary classification itself - which is necessary for the analysis which relies of a single variable specifying the probability of the  MLP being correct ($P^+$). Thus, the jump from Equation 4 to the equation after the sentence  \"since the two labels are symmetric the average loss $\\mathcal{L}$ can be written as\" (which again is not actually shown) is also unclear and missing multiple steps. Another example is the treatment of $N$ right before Equation 5, where it is never mentioned that N follows a distribution before this point, but then it is just defined as being approximated by a normal distribution and the loss equation updated again without derivation or explanation. This happens throughout the work and needs to be correct, for example before Equation 6 it is merely stated \"From Equation 5 a few steps of simplification leads to...\". I think I make my point but for clarity I will point out some quicker examples:\n- Lines 297 to 302: The conclusion of the loss from the MLP separating from the in-context pathway seems true by construction and if I am not mistaken is essentially the independence ansatz. I'm not certain why here it is being phrased as an emergent phenomenon.\n- Figure 4: The landscape depends on the probability of the MLP being correct but this is omitted entirely.\n- Lines 329 to 330: it says \"if the MLP is unable to perfectly memorize the K item-label pairs\" but on lines 225 and 226 it is said that \"all the information required by the MLP to memorize the label of the target is present in $t_{N+1}$\". These are contradictory statements. It seems the second sentence is an assumption which is used in the derivation (or else you cannot summarize the MLP contribution just with $P^+$ - it is difficult to tell without the derivation). This would mean that you cannot interpret the model in the case of the MLP not being able to learn.\n- Lines 336 to 338: The ultimate conclusion of this section is that the model behaviour depends on the initial values of $\\omega$ and $\\beta$ (less the MLP cannot learn the dataset in which case in-context learn will clean up the remaining loss as best it can). Given this conclusion I am confused what the point of this section is given that the condition of $\\omega, \\beta >> 1$ resulting in ICL is already noted from Equation 3. It's possible I am missing some nuance around this point but without derivations it is difficult to say.\n- Two Taylor expansions are used in this work with almost no mention of what sorts of terms or information is lost but this approximation.\n- Lines 355 and 356: It is not clear at all that $c_1(t),c_2(t) \\approx 1/2$ here. This is likely due to the information treatment and lack of explanation of $I_k(t)$.\n- Lines 369 to 372: The relationship between $c_1$ and time is not clearly discussed and $c_1$ in it's own right is not clear explained. Why is the mean of the sigmoid of the negative of the MLP output for the positive class useful? Why would it's integral over time diverge? Surely for a fixed dataset the MLP output stabilises? Is this integral then just a repeated sum of this stable mean value? Why does the divergence of $I_K(t)$ then imply ICL will occur?\n- For Section 3.7: How standard is it to regularize only the in-context pathway? Seems like a fairly trivial result that if only this pathway is regularized then it will disappear.\n- For Section 3.7: Since the dynamics are only considered once ICL is acquired, are the results of this section only valid if L2 regularization is applied once ICL is acquired? Since Equation 5 is not derived with L2 regularization I don't see how it can be just dropped in here and claimed that it ``applies throughout training'' and then the regularization term is just appended. I find it hard to believe that if the regularization term is present throughout training then ICL will emerge in the same manner as the  model which is used to option Equation 5.\n- Line 402: it is mentioned that the added regularization may be explicit or implicit. What justifies the conclusion that any sort of regularization will hold in a similar manner (or  is that not the implication of this statement?).  What sort of implicit regularization will only apply a weight normalization to t he in-context pathway?\n\nUnfortunately, given my confusion with the derivation due to the lack of rigor and derivations it is difficult to even assess Section 4. Many of the phenomena which are considered in this section are either due to fairly vague conclusions drawn on equations without clear meaning or seem to be true by construction (again having sufficiently large D that the key-query matrix is diagonal is a large simplification not given sufficient attention). Thus, the conclusions and insights of this section seem premature.\n\nIt is possible I am missing something in the research - indeed it is possible this research is very impressive. I also appreciate that theory is difficult and requires simplifying assumptions, especially to get to grips with something as complex as in-context learning in transformers. But I am certain that the presentation of this research is significantly lacking for the current version.", "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": "ai7ljRPTyX", "reviewer": "Reviewer_4B2U", "summary": "The paper studies the transition behavior of in-context learning (ICL) and in-weight learning (IWL) in single-layer, single-head, nonlinear Transformer, trained on a variant of the binary labelling task introduced by Chan (2022). The paper then reduces this setting to an effective, “minimal” model that treats ICL and IWL as separate terms, leading to a pair of coupled scalar ODEs stemming from gradient flow on a data-averaged loss function. The minimal model is able to qualitatively explain empirical phenomena in the literature. Importantly, for the setting considered, the provided explanations match quantitatively, and with high precision. \n\nIn particular, the model explains the following empirical phenomena\n\n1. Sharp acquisition of ICL w.r.t. task diversity (modelled by number of training samples $K$), with bimodal IWL-ICL distribution at the threshold K\\* (Kirsch, 2022\\)  \n2. Long-tailed distribution of ICL acquisition times (Reddy, 2023\\)  \n3. Transience of ICL (Singh, 2023\\)  \n4. (somewhat) The dependence on data statistics, in the paper modelled as label inbalance $n\\_+/N$\n\nIn addition, the model makes two predictions: 1) That acquisition time $t_{ICL}$ scales linearly with context length $N$, and that \"task diversity\" $K^\\star$ where ICL and IWL coexist depends on context length $N$ with a power law $ K^* \\sim N^\\{\\nu=0.72\\}$.\n\n### Recommendation\nThe explanatory power, originality, and simplicity of the model are striking and well outweigh my criticisms. Moreover, the paper is highly relevant for the current discussion on in-context learning. I recommend acceptance and a prominent highlight at the conference.", "questions": "I've expressed my questions in the Weaknesses section above.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The modelling approach is plausible and original, and the power despite its simplicity of the minimal model is striking.", "weaknesses": "(more important points first)\n\n- While the model qualitatively explains the literature as stated above, quantitative connections are missing. I understand that some quantitative predictions may be contingent on the one-layer-one-head architecture and task, but others may be transferable, such as the two predictions mentioned in the Summary. I would find it interesting to evaluate these on a naturalistic task. Two examples come to mind: The path-finding task introduces in (Brinkmann et al., 2024) and possibly also a lightweight NLP setting. \n- It is unclear to me how strong the dependence is in terms of data statistics. In particular, is there an intuition why the edge case of exactly $N/2$ positive $\\ell=\\+1$ tokens is so different? \n- It took me a while to understand the relation between the two losses at the end of Section 3.7. It would help the reader to find details, possibly in the Appendix. \n- Why does $\\\\beta=\\\\infty$ (or even just a large value) in Eq. (12) signify ICL dominance? I see this for $w$, but not for $\\\\beta$.   \n- The claim that $c\\_1(t)$ has a power-law up to $K=10^5$ is not supported by Fig. A.3. The claim needs to be weakened or additional simulations are required.   \n- The title of the paper, “differential learning kinetics”, that also is used to label the distinction in Fig. 1, is not picked up much in the main text, and in my opinion also not very descriptive overall. I feel the paper would benefit from connecting the theory to this intuition that the authors developed.   \n- In Fig. 1 and 2: At what time is $K^*$ evaluated?\n- No code is published, limiting transparency and reproduciblity.\n\n#### **Minor**\n\n- The paper states in line 376  \n> While previous work has shown that such skewed data distributions favor ICL acquisition (Chan et al. (2022)), our theory suggests that such data dependencies can be examined by simply measuring $I\\_K$.  \n\n  As I understand these findings are compatible, the wording here to me however makes it unclear whether they are or not.  \n- It is not clear how good of a proxy $K$ is for data diversity, generally, so this caveat should be mentioned.\n- Sometimes $\\rangle\\_{\\\\phi\\_+}$ and sometimes  $\\rangle\\_{\\\\phi\\_+ \\\\sim P\\_+}$ are used, that was slightly confusing to me. Also, $P_\\+$ and $P(n_+)$ could be differentiated more clearly notationally.\n- In line 413 it says\n> Equation 16 implies that $w$ after ICL acquisition tracks $w_{tr}$...\n\nI assume $w_{tr}$ stems from gradient descent on Eq. (15), but this and why it bears this particular name is not clear.\n- Data diversity is modelled by $K$, I would indicate this in Fig. 1 on the y-axis  \n- In line 361, the case $-\\beta\\_0 \\gg 1$ is discussed, and then a conclusion is being made on long-tailedness. I believe that the results show that long-tailedness should hold for both $-\\beta\\_0 \\\\gg 1$ and $-\\\\beta\\_0 \\\\ll 1$, but the wording here made it seem like it only applies to the former.   \n- Just before Section 3.3, the paper states   \n> …three phenomena of interest, the transition from memorization to generalization, abrupt ICL learning and ICL transience,...  \n  \n  of which the first and second seem identical, and the (Reddy 2023\\) long-tailedness property does not appear.\n- The Gaussian on the y-axis s Fig. 6c is not referenced.", "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": "u4Tx9itvXb", "reviewer": "Reviewer_SXqf", "summary": "The paper provides a mechanistic interpretability analysis of ICL and IWL, explaining observed phenomena from the literature and presenting testable predictions for other spacing laws. The authors hypothesize and experimentally demonstrate that the sub-circuits responsible for memorization and generalization can largely be considered independent.", "questions": "**Questions**\n\n- In previous works that observed the transient nature of the network, were weight regularizers applied?\n- Could there be an alternative explanation beyond the dynamical and capacity arguments?\n- For future work, it would be interesting to explore the dynamic effects under capacity limitations. Do you have any intuition about the phenomena that might arise, considering that both could occur in real-world models?\n- How do these insights from ICL and IWL models compare with findings at scale in previous literature? Do you expect some of your results to hold at scale?\n\nI hope this review is helpful for the further development of the paper. I encourage the author to continue this research and incorporate the feedback provided.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "First of all, congratulations on your ICLR submission! I’d like to commend the author for this excellent piece of research. Below is my review of the work:\n\n**Originality**\nThe paper is bringing a novel methodology to study ICL and IWL, and understand the trade of between the generalisation and memorisation.\n\n**Quality**\nThe paper is of general great quality\n- The paper effectively acknowledges its limitations.\n- The experimental results are well-detailed and thoughtfully motivated.\n- The findings demonstrate that the predictions made using the smaller model hold up in a slightly more general model used to study ICL and IWL.\n\n**Clarity**\nThe paper is very well-structured, making it easy to follow. The figures are clear, informative, and complement the content effectively.\n\n**Significance**\nThis work is highly relevant to the research community and plays a crucial role in improving the interpretability of neural networks. I would like to highlight its significance, particularly for the neuroscience community, where there is potential for fascinating cross-disciplinary ideas, as exemplified in this paper see . These paper explore bi level optimisation as a mechanism for effective learning and generalisation [1,2]. \n\n[1] Jensen, K.T., Hennequin, G. and Mattar, M.G., 2024. A recurrent network model of planning explains hippocampal replay and human behavior. Nature Neuroscience, pp.1-9.\n[2] Murray, J.M. and Escola, G.S., 2020. Remembrance of things practiced with fast and slow learning in cortical and subcortical pathways. Nature Communications, 11(1), p.6441.", "weaknesses": "**Originality**\nThe topic is highly relevant and aligns with current research trends, yet it still offers an intriguing and fresh perspective.\n\n**Quality**\nA more detailed discussion of the limitations and assumptions would enrich the paper. Additionally, delving deeper into the comparison with capacity theory could provide further insights (see question below).\n\n**Clarity**\nSuggestion: Including a small diagram of the model would improve readability, especially considering its compact size.\nCould the results be validated in a more step-by-step manner? This would clarify that the approximations are sound and the rationale behind them, potentially expanding on this in the appendix.\nProviding full derivations in the appendix would enhance completeness and transparency.\n\n**Significance**\nWhile the authors have noted the need for validation with larger models, doing so would strengthen the paper, despite the acknowledged material constraints.", "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 (Spotlight)