Loss Landscape Degeneracy and Stagewise Development in Transformers
Authors: Jesse Hoogland
George Wang
Matthew Farrugia-Roberts
Liam Carroll
Susan Wei
Daniel Murfet
Venue: Accepted by TMLR ()

Abstract: Deep learning involves navigating a high-dimensional loss landscape over the neural network parameter space. Over the course of training, complex computational structures form and re-form inside the neural network, leading to shifts in input/output behavior. It is a priority for the science of deep learning to uncover principles governing the development of neural network structure and behavior. Drawing on the framework of singular learning theory, we propose that model development is deeply linked to degeneracy in the local geometry of the loss landscape. We investigate this link by monitoring loss landscape degeneracy throughout training, as quantified by the local learning coefficient, for a transformer language model and an in-context linear regression transformer. We show that training can be divided into distinct periods of change in loss landscape degeneracy, and that these changes in degeneracy coincide with significant changes in the internal computational structure and the input/output behavior of the transformers. This finding provides suggestive evidence that degeneracy and development are linked in transformers, underscoring the potential of a degeneracy-based perspective for understanding modern deep learning.

Published in Transactions on Machine Learning Research (07/2025)
Loss Landscape Degeneracy and Stagewise Development in Transformers
Jesse Hoogland
jesse@timaeus.co
Timaeus
George Wang
george@itimaeus.co
Timacus
Matthew Parrugia-Roberts
matthew@ffr.in..u in.
Department of Computer Science, University of Oxford
Liam Carroll
lemeykc@@gmail.com
Independent
Susan Wei
susan.vei@monnahh.cll
Department of Bconometrics and Business Statistics, Monash University
Daniel Murfet
daniel, muurfeefrrrett..m....
School Mathematics and Statistics, the University Melbourne
Reviewed on OpenReview: https: // openreview. net/ per/forun? id- 454
Abstract
Deep learning involves navigating high-dimenniionll loss landscape the neural network parameter space. Over the course of training, complex computational structures form and re-form inside the neural network, leading to shifts in suput/outpon behavior. It is priority for the science of deep learning to uncover principles governing the development of neural network structure and behavior. Drawing on the framework of singular learning theory, we propose that model development is deeply linked to degeneracy in the local geometry of the loss landscape. We investigate this link by monitoring loss landscape degeneracy throughout training, quantified by the local learning coefficient, for transformer language model and an in-context linear regression Iransformer. We show that training can be divided into distinct periods of change in loss landscape degeneracy, and that these changes in degeneracy coincide with significant changes in the internal computational structure and the input/output behavior of the transformers. This finding provides suggestive evidence that degeneracy and development are linked transformers, underscoring the potential of degeneracy-baad perspective for understanding modern deep learning.
1 Introduction
A striking phenomenon in modern deep learning is the sudden shift in a model's internal computational structure and associated changes in /ouputt behavior (e.g., Wei et al., 2022; Olsson et al., 2022; McGrath et al., 2022). As large models become more deeply integrated into real-world applications, understanding this phenomenon is a priority for the science of deep learning.
A key feature of the loss landscape of neural networks is degeneracy parameters for which some local per- turbations do not affect the loss. Motivated by the perspectives of singular learning theory (SLT; Watanabe,
Equal contribution

---

Published in Transactions on Machine Learning Research (07 (225))
Toa (oo))
Test (ow)))
3-
2-
Local teaiing ovedl Atwi
Local Atww)
135
150
100
100
50
10F
10
10P
10
10'
10
10P
Trining stop
Training step
Stage
LM1
LM2
LM3
LM4
LM5
Stage
LR1
LR2
LR3
LR4
LR5
End t
900
6.5k
8.5k
17k
50k
End t
1k
40k
126k
320k
500k
A/ -2.33
-1.22
-0.18
-0.40
0.34
AE
-..2
-2.21
-0.07
-0.05
-0.029
+26.4
+22.5
-1.57
+8.62
+1.77
A$
+21.4
+149
-2.3
44.1
+3.56
(a) Two-layer attention-only language transformer (LM).
(b) In-context linear regression transformer (LR).
Figure 1: Tracking loss landscape degeneracy reveals developmental stages. We train transformer
models on both (a) natural language data and (b) syuthetic in-context linear regression data. In addition tost loss (top row), track loss landscape degencracy quantified by the local lcarning coefficient (LLC) (middle row; Section 4). Critical points in the LLC curve mark boundaries between distinct developmental stages (bottom row; warm hues for increasing LLC, cold for decreasing LLC; Section 5). show in Sections 6 and 7 that most of these stages coincide with the formation of significant internal structures or changes in behavior. The language model first learns to predict using bigram statistics (LM1), then common n-grams (I.M2,, before lorming the induction circuit studied by Olsson et al. (2022) (LM3&TM4)). The regression model first learns the optimal context- indeeeddeed solution (LR1), then acquires robust in- context learning (LR2), then specializes to the pro-training distribution (LR3&LR4). These stage divisions and interpretations specific to the above training runs, but we show in Appendix B.4 that similar divisions arise with different training seeds.
2009) and nonlinear dynamics (Waddington, 1957; Thom, 1972), where degeneracy plays fundamental role in governing development, we believe that studying degeneracy in the local geomctry of the loss landscape is key to understanding the development. of structure and behavior in modern deep learning.
In this paper, we contribute an empirical investigation of the link between degeneracy and development for transformers in two learning settings. We track loss landscape degeneracy along with model structure and behavior throughout training, using the following methodology.
1. Transformer training (Section 3): We train two transformers, language model (LM) with around 3M parameters trained on a subset of the Pile (Gao et al., 2020; Xie et al., 2023), and an in- contert linear regression model (LR) with around 50k parameters trained on synthetic regression data following Garg et al. (2022).
Degeneracy tracking (Section 4): We quantify loss landscape degeneracy throughout training by estimating the local learning coefficient (LLC; Lau et al., 2025), a measure of degeneracy derived from SLT.
Degeneracy-based stage division (Section 5): Motivated by the singular learning process in Bayesian inference (Watanabe, 2009, $7.6; Chen et al., 2023), we use critical points in the LLC curve divide training into approximate developmental stages.
Developmental analysis (Sections 6, 7): We track shifts in each model's internal computational structure and onput/output behavior across training, quantified using various setting-specific metrics.
2

---

Published in Transactions on Machine Learning Research (07/2025)
Crucially, we discover that most of the developmental stages identified by changes in loss landscape degener- acy coincide with significant, interpretable shifts in the internal computational structure and iuput/outpul behavior of the transformers, showing that the stage division is meaningful. Our investigations are moti- vated by the hypothesis of a fundamental link between degeneracy and development in deep learning. This hypothesis is theoretically grounded in SLT but so far not. empirically validated except in loy models (Chen et al., 2023). We view the above discoveries as preliminary evidence for this hypothesis in larger models, and an indication of the potential of degeneracy as lens for understanding modern neural network development. Section 8 discusses this and other implications of our investigation.
2 Related work
Degeneracy and development in singular learning theory Our hypothesis that degeneracy and development are fundamentally linked is motivated by singular learning theory (SLT; Watanabe, 2009), a framework for studying singular statistical models (a class that includes neural networks, Hagiwara ct al., 1993; Watanabe, 2007; Wei et al., 2023). SLT proves that in singular models, Bayesian inference follows the singular learning process, in which degeneracy in the likelihood governs stagewise development in the posterior as the number of samples increases (Watanabe, 2009, $7.6; Lau et al., 2025; Chen et al., 2023). While there are many differences between Bayesian inference and modern neural network training, an analogy to the singular learning process informs our methooolggy for stage division.
Degeneracy and development in nonlinear dynamics Further motivation for our hypothesis comes from viewing neural network training as a stochastic dynamical system, in which the population loss is a governing potential encoding the data distribution. It is well-understood in nonlincar dynamics that degeneracy in the local geometry of potential can give rise to stagewise development of system structure (Waddington, 1957; Thom, 1972, cf. Franceschelli, 2010). This connection has been observed in biological systems at significant scale and in the presence of stochasticity (Freedman et al., 2023). We emphasize changes in degeneracy over a stage whereas in bifurcation theory the focus is more on the degeneracy at stage boundaries (Rand et al., 2021; MacArthur, 2022; Sáez et al., 2022).
Stagewise development in deep learning The idea that neural networks development occurs in stages goes back decades (Raijmakers al., 1996) and has received renewed attention in modern deep learning (e.g.. Wei et al., 2022; Olsson et al., 2022; McGrath et al., 2022; Odonnat et al., 2024; Chen et al., 2024; Edelman et al., 2024). In the case of deep linear networks, we understand theoretically that models learn progressively higher-rank approximations of their data distribution (see, e.g., Baldi & Hornik, 1989; Rogers & McClelland, 2004; Saxe et al., 2019) throughout training. Our findings suggest that studying degeneracy could help generalize this understanding to modern architectures that exhibit more complex internal computational structure, such as translormers.
Studying loss landscape geometry Given the central role played by the loss landscape in deep learning, it is unsurprising that there have been many attempts to study its geometry.
One approach is to visualize low-dimensional slices of the loss landscape (Erhan et al., 2010; Goodfellow et al., 2014; Lipton, 2016; Li et al., 2018; Tikeng Notsawo et al., 2024). Unfortunately, a random slice is with high probability a quadratic form associated to nonzero eigenvalues of the Hessian and is thus biased against geometric features that we know are important, such as degeneracy (Wei et al., 2023). Moreover, Antognini & Sohl-Dickstein (2018) have emphasized the difficulty of probing the loss landscape of neural networks with dimensionality reduction tools.
Other standard methods of quantifying the geometry of the loss landscape, such via the Hessian, are insensitive to important aspects of degeneracy. For example, the Hessian trace or maximum eigenvalues quantify the curvature of critical point but ignore degenerate dimensions, and the Hessian rank counts the number of degenerate dimensions but fails to distinguish between dimensions by the order of their degeneracy (c.g., quartic vs. zero). In contrast, the LLC a principled quantitative mcasure of loss landscape dogeneracy. Appendix B.5 includes experiments showing that Hessian statistics do not reveal the clear stage boundaries revealed by the LLC in our in-context linear regression setting.
3

---

Published in Transactions on Machine Learning Research (07 (07/2025)
3 Training transformers in two settings
We study transformers trained in two learning settings, namely language modeling and in-contert linear regression. These settings have been the subject of recent work on the emergence of in-context learning (TCL),, a compelling example of a sudden shift in a model's internal computational structure in modern deep learning (Olsson et al., 2022).
In this section, we describe both settings and introduce their loss functions and data distributions. Common to both settings is transformer model denoted fw with parameters w, which takes as input sequence of tokens, also called contert. We describe specific architecture details and training hyperparameters in Appendices F.1 and F.2.
Language modeling Elhage et al. (2021) and Olsson et al. (2022) observed that two-layer attention- only transformers (transformers without MLP layers) form interesting internal computational structures supporting ICT, including induction heads. In order to compare with their behavioral and structural analysis we adopt the same architecture. In Appendix E we also study one-layer attention-only transformers. We note that, while we don't study language models with MLP layers (following prior work), we do MLP layers for in-context linear regression.
We consider the standard task of next-token prediction for token sequences taken from subset of the Pile (Gao et al., 2020; Xie et al., 2023). We denote the input context by Sx = (l...., (x) where K is the context length. We denote by Se* the prefix context (t1...., tx) of context Sk. Our data is a collection of length-K contexts, {5k) ) Thus Sy denotes a prefix of the zth context, Sk.
Given the context Skkk: the transformer model fw outputs a vector of logits (u such that softmax(Sft (fflmax((f(($z)) is probability distribution over all tokens (we denote by (fftmex([f the prob- ability of token t). The per-token empirical loss for language modeling is then the average cross-entropy between this distribution and the true next token at each index E [1....,,--11
1 11
(w,((w)=
log (s
The empirical loss is then (n(w) ক- EK (w), with the test loss {(w) defined analogously on a held-out set of examples. The corresponding population loss E(w) is defined by laking the expectation with respect to the true distribution of contexts (see also Appendix A.6).
In-context linear regression Following Garg et al. (2022), a number of recent works have explored ICL in the setting learning simple function classes, such as linear functions. This setting is of interest because we understand theorotically optimal (in-context) lincar regression, and because simple transformers are capable of good ICL performance in practice (see, e.g., Garg el al., 2022; Raventós et al., 2023).
We consider standard synthetic in-context linear regression problem. A task is vector t E RR,, and an example is a pair (x, E RD x R. We sample context by sampling one task t ~ NN(0, ID) and then sampling K i.i.d. inputs æ1.-.., TK ~ -N(0, and outputs y1,..., VK ~ Nit "N(TTx, 0²) This results in the context Sx = (I1, Vi...., 1- 11 with label yк- We denote by Sek the prefix context (z1, Vii, Tx) of context Sk, its label is yk- Appendix F,2.2 describes how we encode the xa and y as tokens. Our data is a set of contexts [(t,, vic 1% sampled i.id. as described above.
Running context Sigu through the transformer yields a prediction o (w leading to the per-token empirical loss for in-context linear regression for k E {l....,K? KJ.
(n.(ww = Σ (ol v1.
(2)
The associated empirical loss is (w)) = ΣK ln (w). The corresponding test loss ((w) and population
loss (ww) are defined analogously in the language modeling setting.
4

---

Published in Transactions on Machine Learning Research (07 (07/2025)
E
V(c) α e² Locall learning coefficient
f(x,y) x² v²
E(a,y) x² v²
(z, v) n'y² 1(2,0)
((a,)) (s 1) + x²)²
1 X(0,0)= 8 A(1.3) =
X) =1
x(0,0) () =
Figure The local learning coefficient (LLC) measures loss landscape degeneracy. The LLC can be defined in terms of the rale al which the parameter space volume (within a given neighborhood and with given maximum loss) shrinks the loss threshold is reduced to that of the local minimum. We show four population loss landscapes for a two-dimensional parameter space with decreasing LLC (incrcasing degeneracy). In these examples, the local multiplicity is 1. See Appendix A.2 for detailed description of each example, as well as several additional examples.
4 Quantifying degeneracy with the local learning coefficient
We track the evolution of degeneracy in the local geometry of the loss landscape throughout, training by estimating the local learning coefficient (LLC; Watanabe, 2009; Lau et al., 2025) at model checkpoints. In this section, we review the LLC and the estimation procedure of Lau al. (2025).
The local learning coefficient (LLC) Given local minimum w" of population loss é (a negative log likelihood), the LLC of w², denoted X(w"), is a positive rational nunber that measures the amount of degeneracy in l near w' (Watanabe, 2009; Lau et al., 2025), i.e., how many ways can be varied near w" such that f(w) remains equal to f(w"). Formally, the LLC is defined as the volume-scnling rale near w". This illustrated in Figure 2, further described in Appendix A.1, and treated in full detail in Lau et al. (2025). Informally, the LLC is a measure of minimum "flatness." It improves over conventional (second- order) Hessian based measures of flatness because the LLC sensitive to more significant, higher-order contributions to volume-scaling.
Estimating the LLC Tau et al. (2025) introduced an estimator for the LLC based on stochastic-gradient Langevin dynamics (SGLD; Welling & Teh, 2011), which we use in our experiments. Let w² be a local minimum of the population loss E. The LLC estimate )(w*) is
A(w") = nB R " , [w]]]] ) - en (w²)
(3)
where R&
w/w denotes the expectation with respect to the localized Gibbs posterior
p(uo w", .,,7) a exp (3)3n((w) - ? 2 w w11
with inverse temperature 8 (controlling the contribution of the empirical loss landscape) and localization strength (controlling proximity to w"). The basic idea behind this estimator is the following: the more degenerate the loss landscape, the easier it is for a sampler exploring the Gibbs posterior to find points of low loss, and, in turn, the lower '(w*). Appendix A.3 discusses technical SGLD details, Appendix A.4 docunents the hyperparameters used in our experiments, and Appendix A.5 outlines our hyperparameter tuning procedure.
Assumptions of LLC estimation Strictly speaking, the LLC defined only for loss functions arising a negative log likelihood, whereas our loss function includes terms from overlapping context prefixes.
5

---

Published in Transactions on Machine Learning Research (07 2025)
(((w) > (2(22)
Fn(W))
2(w4) {(w5)
Fx (W2)
Wi
W
W
WI preferred Perit. W2 preferred n
Figure 3: In the singular learning process, the Bayesian posterior shift between neighbor- hoods with different degeneracy. Watanabe's free energy formula (4) highlights a tradeoff between loss in (the linear term coefficient) and degeneracy À (the LLC, the logarithmic term coefficient). Consider two local minima wi, wh with neighborhoods WT, W As the number of samples n increases, if wh has lower loss and higher LLC than wi. W3 will suddenly achieve lower free energy than Wf at some critical sample size Merit causing the Bayesian posterior to shift from concentrating in WT to W2.
It is possible to define a negative log likelihood- loss for transformer training we show empirically in Appendix A.6 that this does not have a significant effect on LLC estimates, and we proceed with overlapping contexts for efficiency.
Moreover, the LLC is defined only for local minima of such loss functions, whereas we note equation (3) is defined for arbitrary w" and we apply the estimator throughout training. This approach has precedent in prior work on LLC estimation: Lau et al. (2025) showed that when applied to trained parameters, the estimator accurately recovers the learning coefficient associated with a nearby minimum, and Chen et al. (2023) found that the estimator produces reliable results for parameters throughout training. In our case, we obtain stable estimates throughout training given sufficiently strong localization y. See Appendix A.7 for more details.
5 Degeneracy-based stage division
We use critical points (that is, platcaus, where the first derivative vanishes) in the LLC curve to define slage boundaries that divide training into developmental slages. This approach is motivated by the singular learning process in Bayesian inference, which we review below.
Bayesian local free energy Let W be neighborhood of a local minimum w" of the population loss (a negative log likelihood). Given n samples we can define the local free energy of the neighborhood (Lau et al., 2025),
(W")) = = /
ww- (xp(-nen(ww)))
dw,
where (w)) is a prior positive on the neighborhood W*. The lower the local free energy of a neighborhood W*, the higher the Bayesian posterior mass of W*. In fact, by log-sum-exp approximation, the Bayesian posterior is approximately concentrated the neighborhood with the lowest local free energy (cf., Chen et al., 2023).
The singular learning process Watanabe's free energy formula gives, under certain technical conditions, an asymptotic expansion in n of the local free energy (Watanabe, 2018, Theorem 11; Lau et al., 2025):
Fa(W") = (w"" + A(w") log + (log (ologgoog)) logn)
(4)
Here, (n(w") ) is the empirical loss, X(w*) ) is the LLC, and the lower-order terms include constant contribution from the prior mass of W*
The first two terms in equation (4) create a tradeoff between accuracy (e,) and degeneracy (A). Morcover, increases, the linear term becomes increasingly important relative to the logarithmic term, changing the

---

Published in Transactions on Machine Learning Research (07/2025)
nature of the tradeoff. At certain critical the neighborhood with the lowest local free energy may rapidly change to a neighborhood with decreased loss and increased I.I.,, as illustrated in Figure 3.
A sequence of such posterior transitions between increasingly degenerate neighborhoods is prime example of the singular learning process (Watanabe, 2009, $7.6). We note that this is not the only possible dynamic- lower-order terms may also play a role in the evolving competition.
LLC plateaus separate developmental stages While the general connection between the singular learning process in Bayesian inference and stagewise development in deep learning remains to be understood, Chen et al. (2023) showed that, in small autoencoders, both Bayesian inference and stochastic gradient descent undergo rapid transitions between encoding schemes, and these transitions are reflected as sudden changes in the estimated LLC.
This perspective suggests that changes in the loss landscape degeneracy, as measured by the LLC, reflect qualitative changes in the model. In larger models, we expect that these qualitative changes may be more gradual, while still being delineated by brief moments in which the posterior is stably concentrated around a given local minimum. This motivates our approach of identifying plateaus in the estimated LLC pauses before and after given increase or decrease in degeneracy as sluge boundaries which divide training into approximate developmental stages. The resolution of these stage boundaries depends on the density of checkpoints used for LLC estimation and the precision of those estimates.
Results In our experiments, we identify plateaus in the estimated LLC curve by first lightly smoothing the LLC curve with a Gaussian process to facilitate stable numerical differentiation with respect to log training time. We identify plateaus as approximate zeros of this derivative, namely local minima of the absolute derivative that fall below small threshold (see Appendíx B.1). Figure 1 and Appendices B.2 and B.3 show the results. Appendix B.4 shows that similar stage divisions arise for independent training runs.
6 Results for language modeling
Platcaus in LLC estimates (Figure 1a) reveal five developmental stages for our language model. In order to validate that this stage division is meaningful, we search for concomitant changes in the model's input/oulpun behavior and its internal computational structure. In this section, we report a range of setting-specifi metrics that reveal the following significant, interpretable changes coinciding with each stage: in LM1 the model learns to predict according to bigram statistics; in LM2 the model learns to predict frequent n-grams and use the positional embedding; in LM3 and LM4 the model respectively forms "previous-token heads" and "induction heads" as part of the same induction circuit studied by Olsson et al. (2022). Note that we did not discover significant changes in LM5, and we do not claim that these are the only interesting developmental changes occurring throughout training. There may be other interesting developmental changes that are not captured by our metrics, or are not significant enough to not show up in the LLC curve.
6.1 Stage LM1 (0-900 steps)
Learning bigram statistics Figure 4(a) shows that the bigram score the average cross entropy between model logits and empirical bigram frequencies (see Appendix C.1.1) is minimized around the 1.11LM, boundary, with a value only 0.3 nats above the irreducible entropy of the empirical bigram distribution. This suggests that during LM1 the model learns to predict using bigram statistics (the optimal next-token prediction given only the current token).
6.2 Stage LM2 (900-6.5k steps)
Using positional information During LM2 the positional embedding becomes structurally important. Figure 4(b) shows that here the test loss for the model with the positional embedding zero-ablated diverges from the test loss of the unablated model (sce Appendix C.2.1). Specifically, we mcan setting learned positional embeddings to zero during evaluation. Conditional on our architecture this establishes whether the model effectively uses positional information. A similar method could be used in a model without learned
7

---

Published in Transactions on Machine Learning Research (07/2025)
Model Rigram entrpy
(d)
0.3
Prex tokem hrad 1:2
Bignin scoe
Preo uilco scoRe
Prox token head
0.2
0.0
107
10'
Pos ablated
Induction head 2:7
Tea loax Rowe
Teai
Induction head 28
0,050
Pedft
0,000
107
10t
10t
Prgrum xore
A-grams
ICL scre
-0.1
perf
-0.2
Training
Trining
Figure 4: Language model stages coincide with significant structural and behavioral changes. (a) The model learns bigram statistics in LM1, (b) then the positional embedding becomes useful from LM2, (c) enabling the learning of common n-grams. Induction circuit formation begins with (d) previous-token heads in LM3, followed by (c) induction heads LM4, leading to (f) a drop ICL score indicating the acquisition of in-context learning. Note: in (d,e), Ith denotes attention head h in layer I; dark lines indicate heads comprising the induction circuit.
positional embeddings. There is also an uptick in previous-token attention among some first-layer attention heads shown in green in Figure 4(d).
Learning common n-grams We define an n-gram score as the ratio of fimal-position token loss on (1) a baseline set samples from a validation set truncated to n tokens, and (2) a fixed set common n-grams (see Appendix C.1.2). To compute the "common n-grams" score after extracting the top 1000 n-grams, we compute the loss on contexts like
[<bos_token] <token_1>> 1>, <token_22 2>, <token n>]
using the loss on <token_n> and normalize against the average loss on the n-th token of similar-length contexts drawn from the pretraining distribution, then divide the n-th token loss of truncated pretraining contexts by the n-gram loss to get the n-gram score.
Figure 4(c) shows large improvement in n-gram score for n n= 3,4 4 during LM2. This suggests that during LM2 the model memorizes and learns to predict common n-gramss for n > 2 (note this requires using the positional encoding and may also involve previous-token heads).
Foundations induction circuit In this stage, the heads that eventually become previous-token and induction heads in future stages begin to compose (that is, read from and write to shared residual stream subspace; Figure C.4 and Appendix C.2.2). This suggests that the foundations for the induction circuit are laid in advance of any measurable change in model outputs attention patterns.
6.3 Stages LM3 & LM4 (6.5k-8.5k 8.5k-17k steps)
Formation of induction circuit (as studied in Olsson et al., 2022) Figure 4(d) shows the previous-token matching score (Appendix C.2.3) rises over LM3 and L.M4 for the two first-layer heads that eventually par-
8

---

Published in Transactions on Machine Learning Research (07 (075)
ticipate in the induction circuit (as distinguished by their composition scores, Appendix C.2.2). Figure 4(c) shows that during LM4 there is an increase in the presta-malching score (Appendix C.2.4) for the two second- layer induction heads that complete the induction circuit. Figure 4(f) shows a corresponding drop in the ICL score (Appendix C.1.3) as the model begins to perform in-context learning.
The LLC decreases during LM3, suggesting an increase degeneracy (a decrease in model complexity). This may be related interaction between heads. It would be interesting to study this stage further via mechanistic interpretability.
7 Results for in-context linear regression
(a)
Por-e Norn
Plateaus the LLC estimates (Figure 1b) reveal five developmental stages for our in-context linear regression model. We validate that this stage di- vision is meaningful by identifying significant, con- comitant changes in the model's structure and be- havior: in LR1 the model learns to predict with- out looking at the context; in LR2 the model ac- quires a robust in-context learning ability; and in LR3 and LR4 the model becomes more fragile out-of-distribution inputs. We did not discover sig- nificant changes in LR5, do claim this is exhaustive list of developments.
30
(b)
Xa Simmm
x, gl0.
g
10
(cf
7.1 Stage LR1 (0-1k steps)
Learning to predict without context Fig- ure 5(a) shows that the mean square prediction for all tokens [ decrcases during LR1, reaching a minimum of 0.1 (smaller than the target noise o² = 0.125) slightly after the end of LR1. Similar how the language model learned bigram statis- ties in LMI, this suggests the model first learns the optimal context-independent prediction dx = æt where the mean of the task distribution (zero in this case).
Figure 5: In-context linear regression model stages coincide with significant structural and behavioral changes. (a) During LR1, the model learns to make context-independent predic- tions, 2g + On = 0. (b) During LR2, ICL per- formance improves, then during LR3 the model be- comes worse at ICL OOD inputs rą ~ (O, gTp) for g 3. (c) During LR3 and LR4, layer normal- ization weights "collapse," " possibly contributing to the LLC decrease.
7.2 Stage LR2 (1k-40k steps)
Acquiring in-context lcarning Figure 5(b) shows that during LR2 there a drop in ICL score (Ap- pendix D.1.2), indicating that the model acquires in-context learning.
Embedding and attention collapse Appendix D.2 documents additional changes. Near the end of LR2, token and positional embeddings begin to "collapso," effectively losing singular values and aligning with the same activation subspace (Appendices D.2.1 and D.2.2). the same time, several attention heads form concentrated, input-independent attention patterns (Appendix D.2.3).
7.3 Stages LR3 & LR4 (40k-126k & 126k-320k steps)
Reduced robustness to input magnitude While performance continues improve on typical se- quences, Figure 5(b) shows that during LR3 and LR4, the model's in-context learning ability deteriorates for outlier sequences with higher-than-average [mx]]

---

Published in Transactions on Machine Learning Research (07/2025)
Layer-normalization collapse Figure 5(c) shows the individual weights in the final layer normalization module. A large fraction of these weights go to zero in LR3 and TR4. This occurs in tandem with similar collapse in the weights of the unmbeddding transforms (Appendix D.2.4). This results in the model learning to read its prediction 0K from a handful of privileged dimensions of the residual stream. Since this means that the network puts become insensitive to changes in many of the parameters, we conjecture that this explains part of the striking decrease in estimated LLC over these stages (Appendix D.2.4).
This collapse is most pronounced and affects the largest proportion of weights in the unembedding module, but in LR4 it spreads to earlier layer normalization modules, particularly the layer normalization module before the first attention block (Appendix D.2.5).
8 Discussion
In this paper, have examined the development of transformer models in two distinct learning settings, We quantified the changes in loss landscape degeneracy throughout transformer training by estimating the local learning coefficient (LLC). Motivated by the singular learning process in Bayesian inference, we divided these training runs into developmental stages at critical points of the LLC curve. We found that these developmental stages roughly coincided with significant changes in the internal computational structure and the input/output behavior of each model. In this section, we discuss several implications of these findings.
Towards degeneracy-baseed understanding of deep learning That significant structural and be- havioral changes show up in the LLC curve is evidence that the development of our transformers is closely linked to loss landscape degeneracy. This finding underscores the potential of loss landscape degencracy as a crucial lens through which to study the development of deep learning models.
While we studied two distinct learning settings (including language modeling with a nontrivial transformer architecture), it remains necessary verify the connection between degeneracy and development across a more diverse range of emergent model structures and behaviors. Moreover, future work should investigate this connection in more depth, seeking to establish causal connection between changes in degeneracy and changes in structure and behavior.
Towards developmental interpretability We showed that degeneracy can reveal meaningful changes in transformers. We emphasize that our analysis is not exhaustive we expect only certain "macroscopic" changes, such as the emergence of in-context learning, will have significant enough effect on loss landscape degeneracy to appear separated by plateaus in the LLC curve. Recent work has extended these ideas by measuring the LLC with respect to network sub-modules and with different data distributions, providing a more refined picture of model development (Wang et al., 2025). We expect this research direction will lead to insights into the development of more complex models.
Loss landscape degeneracy offers potting-agnoste "unsupervised" alternative to setting-specific progress measures such those derived by Barak et al. (2022) developed using mechanistic insights from similar models by Nanda et al. (2023). Both approaches landscape degeneracy is able to detect changes without requiring a mechanistic underslanding in advance. Of course, once a change is detected through its effect on degeneracy, it remains to interpret the change.
reveal developments invisible in the loss, but loss
Cases studics in transformer development We do not claim that the structural and behavioral do- velopments we observed in each setting are universal phenomena. Transformers trained with different archi- tectures, data distributions, algorithms, or hyperparameters may develop differently. Rather, our detailed analysis contributes two "case studies" to the growing empirical literature on the emergence of structure in transformers.
On this note, our observations extend those of Olsson et al. (2022) and Elhage et al. (2021). We show that before the induction circuit forms, our 2-layer language model learns simpler interpretable strategics (basod on bigram statistics and common n-grams). This shows that a single training run follows a progression akin to that found by Olsson et al. (2022) for dully-developed models of increasing depth (they showed that
10

---

Published in Transactions on Machine Learning Research (07/2025)
"O-layer" models learn bigram statistics and 1-layer models learn takipptriggams"" A similar progression was observed by Edelman et al. (2024) in Markovian sequence modeling task.
Moreover, in both settings, we saw that before in-context learning emerges, the model learns to predict tokens using the optimal prediction given only the current token (bigram statistics for language modeling, zero for in-context linear regression with this distribution of tasks).
Development and model complexity While we have described the LLC as measure of loss landscape degeneracy, it also be understood measure of model complexity (ef. Appendix A.2). Tt is natural for changes in model's internal structure to show up change in complexity. For example, Chen et al. (2024) showed that the emergence of syntactic attention structure coincides with a spike in two model complexity measures, namely the model's Fisher information and the intrinsic dimension (Facco et al., 2017) of the modol's embeddings.
Notably, we observe stages which the LLC decreases, corresponding to a simplification of the computational structure of the model. Such model simplification has empirical precedent, for instance with Chen et al. (2024) and the recent literature on grokking (Power et al., 2022; Nanda et al., 2023; Tikeng Notsawo et al., 2024). In our case, the mechanistic nature of the simplification is not fully clear, with the collapse of various weights and attention patterns arising as candidates in the in-context linear regression setting.
This phenomenon is currently not accounted for by theories of ncural network development. In the theory of saddle-to-saddle dynamics, deep linear networks learn progressively more complex approximations of the data (Saxe et al., 2019). Likewise, the example transitions the singular learning process outlined Section 5 and Figure 3 describe LLC increases. Though we note that decreasing the LLC while holding the loss constant would be another way to decrease the free energy according to equation (4), providing full theoretical account of these stages is an open problem.
Author contributions
All authors communicated regularly about all aspects of the methodology and writing. The following is non-exhaustive list of some particular areas of individual contribution.
JH led the coding, experimentation, and analysis for the in-context linear regression setting; as- sisted with analysis of the language modeling setting; contributed substantially to the writing of the manuscript; and designed the conceptual illustrations.
GW led the coding, experimentation, and analysis for the language modeling setting; and contributed substantially to the writing of the manuscript.
MFR led the writing of the current main text, based on earlier versions written with others; and assisted with coding and tuning hyperparameters for the in-context linear regression setting.
LC assisted with coding, experimentation, and analysis for the in-context linear regression setting.
SW advised on technical details including assumptions LLC estimation; and assisted with writing.
DM had the original idea for the degeneracy-based stage division methodology; contributed sub- stantially to earlier versions of the main text; and coordinated the project.
Acknowledgments
We thank Edmund Lau for advice on LLC estimation and Mansheej Paul for advice on training transformers in the in-context linear regression setting. We are grateful to Evan Hubinger and Simon Pepin Lehalleur for helpful conversations. For their valuable feedback on the manuscript, we thank Andres Campero, Zach Furman, Simon Pepin Lehalleur, and Nandi Schoots.
We thank Google's TPU Research Cloud program for supporting some of our experiments with Cloud TPUs. JH and GW completed part of this research through the AT Futures Fellowship. MFR completed part of
11

---

Cited on Published in Transactions on Machine Learning Research (07/2025)
this research independent rescarcher sponsored by private individuals, part at Timaeus, and part the University of Oxford. T.C's work was supported by Lightspeed Grants.
References
Joseph Antognini and Jascha Sohl-Dickstein. PCA of high dimensional random walks with comparison to neural network training. In Advances Neural Information Processing Systems 31, pp. 10307 10316. Curran Associates, 2018. Access via NeurIPS. Cited on pages 3 and 29.
Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53-58, 1989. DOI 10. page 3.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. In Advances in Neural Information Curran Associates, 2022. Access via NeurIPS. Cited on page 10.
Processing Systems 85, pp.
Angelica Chen, Ravid Shwartz-Ziv, Kyunghyun Cho, Matthew I. Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in MLMs. In The Twelfth International Conference on Learning Represenlations. OpenReview, 2024. Access via OpenReview. Cited on pages 3 and 11.
Zhongtian Chen, Edmund Lau, Jake Mendel, Susan Wei, and Daniel Murfet. Dynamical versus Bayesian phase transitions in toy model of superposition, 2023. Preprint arXiv:2310.066 06.000001111 [cs.LG]. Cited on pages 2, 3, 6, 7, and 24.
Ezra Edelman, Nikolaos Tsilivis, Benjamin L Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning Markov chains. In Advances in Neural Information Processing Systems 37, pp. 64273 64311. Curran Associates, 2024. Access via NeurIPS. Cited on pages 3 and 11.
Ronen Eldan and Yuanzhi Li. TinyStories: How small can language models be and still speak coherent English?, 2023. Preprint arXiv:2305.07755 CL). Cited on page 46.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. Access via Transformer Circuits Thread. Cited on pages 4, 10, 32, and 33.
Dumitru Erhan, Yoshua Bengio, Aaron Courville, Antoree Manzagol, Pascal Vincent, and Samy Bengio. Why does unsupervised pre-training help deep learning? Journal of Machine Learning Research, 11(19):625-60 660, 2010. Accoss via JMLR. Cited on page 3.
Elena Facco, Maria d'Errico, Alex Rodriguez, and Alessandro Laio. Estimating the intrinsic dimen- sion of datasets by a minimal neighborhood information. Scientific Reports, 7(1):12140, 2017. DOI 00.1088/41188-011-187777 Cited on page 11.
D H Fowlor. Structural Stability and Morphogensis: An Outline of a General Theory of Models. W. A. Benjamin, Inc., 1975. English translation of Thom (1972), with updates from the author. Cited on page 14.
Sara Franceschelli. Morphogenesis, structural stability and epigenetic landscape. In Paul Bourgine and Annick Lesne (eds.), Morphogenesis: Origins of Patterns and Shapes, pp. 283-293. Springer, 2010. DOI 10.1007/9978-3-662-1377775 16. Cited page 3.
12

---

Cited on Published in Transactions on Machine Learning Research (07/2025)
Simon L Freedman, Bingxian Xu, Sidhartha Goyal, and Madhav Mani. A dynamical systems treat- ment. of transcriptomic trajectories in hematopoiesis. Development, 150(11)devv0112)) 2023. DOL 10.1242/dev.20122 201280. Cited on page 3.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pilc: an 800GB dataset of diverse text for language modeling. Technical report, EleutherAI, 2020. Preprint arXiv:2101.0002/ [es.CL]. Cited on pages 2, 4, 46, and 47.
Shivam Garg, Dimitris Tsipras, Percy $ Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. In Advances in Neural Information Processing Systems 35, pp. 50583-30599 Curran Associates, 2022. Access via NeurIPS. Cited on pages 2 and 4.
Hamidreza Ghader and Christof Monz. What does attention in neural machine translation pay attention to? In Proceedings of the Bighth International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 30-39. Asian Federation of Natural Language Processing, 2017. Access via ACL Anthology. Cited on page 37.
lan J Goodfellow, Oriol Vinyals, and Andrew M Saxe. Qualitatively characterizing neural network optimiza- tion problems, 2014. Conference paper presented at ICLR 2015. Preprint arXiv:1412.6544 [cs.NE]. Cited
on page 3.
Katsuyuki Hagiwara, Naohiro Toda, and Shiro Usui On the problem of applying AIC to determine the structure of a layered feodforward noural network. In Proceedings Of 1993 International Joint Conference On Neural Networks, volume 3, pp. 2263 2266. IEEE, 1993. DOI D..1109/JCNNN1993..
page 3.
Andrej Karpathy. NanoGPT, 2022. Access via github.com. com /kappathy. Cited on page 47.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014. Conference paper presented at ICLR 2015. Preprint arXiv:1412 6698 [cs.LG]. Cited on page 48.
Edmund Lau, Zach Furman, George Wang, Daniel Murfet, and Susan Wei. The local learning coefficient: A singularity-aware complexity measure. In The 28th International Conference on Artificial Intelligence and Statistics. OpenReview, 2025. Access via OpenReview. Cited on pages 2, 3, 5, 6, 17, 18, 19, 22, and 24.
Hao Li, Zheng Xu,, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Advances in Neural Information Processing Systems 31, pp. 1389-6399) Curran Associates, 2018. Access via NeurIPS. Cited on page 3.
Zachary C Lipton. Stuck in what? Adventures in weight space, 2016. Workshop paper presented at ICLR 2016. Preprint arXiv:1602.07320 [cs.LG]. Cited on page 3.
Ben D MacArthur. The geometry of cell fate. Cell Systems, 13(1):1-3, 2022. DOI 20.1011/ Cited page 3.
Thomas McGrath, Andrei Kapishnikov, Nenad Tomašov, Adam Pearee, Martin Wattenberg, Demis Hassabis, Been Kim, Ulrich Paquet, and Vladimir Kramnik. Acquisition of chess knowledge in AlphaZero. Pro-
ceedings of the National Academy of Sciences, (12(44)))2206625111 2022. DO1 1ww.1////ppa.222066000 Cited on pages and 3.
Neel Nanda and Joseph Bloom. TransformerLens, 2022. Access via githb..omm/nechandd-i Cited page 46.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistie interpretability. In The Eleventh International Conference on Learning Repre- sentations. OpenReview, 2023. Access via OpenReview. Cited pages 10 and 11.
Ambroise Odonnat, Wassim Bouaziz, and Vivien Cabannes. Clustering head: visual case study of the training dynamics in transformers, 2024. Preprint arXiv:2410.24050 [es.LG).. Cited page 3.
13

---

Published in Transactions on Machine Learning Research (07 (225))
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Halfield-Dodds. Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodci, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Trunsformer Circuils Thrend, 2022. Access via Transformer Circuits Thread. Cited on pages 1, 2, 3, 4, T, 8, 10, 28, 29, and 32.
Mary Phuong and Marcus Hutter. Formal algorithms for transformers. Technical report, DeepMind, 2022. Preprint arXiv:2207. 009038 [cs.LG]. Cited on pages 40 and 47.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets, 2022. Preprint arXiv:2201.0117 [cs.LG]. Cited on page 11.
Ofir Press, Noah A Smith, and Mike Lewis. Shortformer: Better language modeling using shorter inputs. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp- 5493-5505. Association for Computational Linguistics, 2021. DOI 10.18653 /v/ 110/3//122221///// /22/2...c--ooongee..e.
Cited on page 46.
Maartje E. J Raijmakers, Sylvester van Koten, and Peter C M Molenaar. On the validity of simu- lating stagewise development by means of PDP networks: Application of catastrophe analysis and an experimental test of rule-like network performance. Cognitive Science, 20(1):101-136, 1996. DOI 4. Cited on page 3.
David A Rand, Archishman Raju, Meritxell Sáez, Francis Corson, and Eric D Siggia. Geometry of gene regulatory dynamics. Proceedings of the National Academy of Sciences, 118(35):22109729118 2021. DOL 00.1073/pnas22109720
Cited on page 3.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non- Ban-Bayesian in-context learning for regression. In Advances in Neural Information Processing Systems Curran Associates, 2023. Access via NeurIPS. Cited on pages 4, 34, and 35.
36, pp. 11228-14246
Timothy T Rogers and James L McClelland. Semantic Cognilion: A Parallel Distributed Processing Ap- proach. MIT Press, 2004. DOT 10.7551 /mim//msss//111100000 100.00000000001.....011111000000 Cited on page 3.
Meritxell Sáez, Robert Blassberg, Flena Aagilho-Aguilar, Eric D Siggia, David A Rand, and James Briscoe. Statistically derived geometrical landscapes capture principles of decision- making dynamics during cell fate transitions. Cell Systems, 13(1):12 28, 2022. DOI /..016/j.cels.200 000.013. Cited on page 3.
Andrew M Saxe, James L McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23):115377-1144) 2019. DOI 10003//nas.180022200.
Cited pages 3 and 11.
Maxwell Shinn. Phantom oscillations in principal component analysis. Proceedings of the National Academy of Sciences, 120(48):o231 1420120, 2023. DOT 10.1073/ htt.:/mas.2200 2011231111200 1420120. Cited on page 29.
René Thom. Stabilité Structurelle et Morphogénèse: Essai d' d'une Théorie Générale des Modèles [Structural Stability and Morphogensis: An Outline of a General Theory of Models]. W. A. Benjamin, Inc., 1972. In French. Translated into English by Fowler (1975). Cited on pages 2, 3, and 12.
Pascal Tikeng Notsawo, Jr., Hattic Zhou, Mohammad Pezeshki, Trina Rish, and Guillaume Dumas. Pre- dicting grokking long before it happens: A look into the loss landscape of models which grok. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models. OpenReview, 2024. Access via OpenReview. Cited on pages and 11.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackbonNLP: Analyzing and Interpreting Neural Networks for NLP. pp. 63-76. Association for Computational Linguistics, 2019. DOI 10. 10.18633/11/W19-10 Cited on pages 37 and 38.
14

---

Published in Transactions on Machine Learning Research (07/2025)
C H Waddington. The Strategy of the Genes: A Discussion of Some Aspects of Theoretical Biology. Allen & Unwin, 1957. Cited on pages 2, 3, and 24.
George Wang, Jesse Hoogland, Stan van Wingerden, Zach Furman, and Daniel Murfet. Differentiation and specialization of attention heads via the refined local learning coefficient. In The Thirteenth International Conference Learning Representations. OpenReview, 2025. Access via OpenReview. Cited page 10.
Sumio Watanabe. Almost all learning machines are singular. In IEEE Symposium on Foundations of Computational Intelligence, pp. 383-388, IEEE, 2007. DOI M.1100/FOCC.2207.337777
Cited on page
Sumio Watanabe. Algebraic Geometry and Statistical Learning Theory. Cambridge University Press, 2009.
DOI 10.1017/C309780118000 Cited on pages 1, 2, 3, 5, 7, and 17.
Sumio Watanabe. Mathematical Theory of Bayesion Statistics. 10.1 /91211//7811113333010 Cited on page 6.
CRC Press, 2018.
DOI
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilitics of large language models. Transactions Machine Learning Research, 2022. Access via OpenReview. Cited pages 1 and 3.
Susan Wei, Daniel Murfet, Mingming Gong, Hui Li, Jesse Gell-Redman, and Thomas Quella. Deep learning is singular, and that's good. IEEE Transactions on Neural Networks and Learning Systems, 34(12):
2023. DOI 10. 10.1109/TNNI.S. 202. Cited on page
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning, pp. 681 688. ACM, 2011. Access via ICML. Cited on page 5.
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy S Liang Data selection for language models via importance resampling. In Advances in Neural Information Processing Systems 36, pp. 34201-34227. Curran Associates, 2023. Access via NeurIPS. Cited on pages 2, 4, 46, and 47.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. PyHessian: neural networks through the lens of the Hessian. In 2020 IEEE International Conference on Big Data (Big Data), pp. 581 590, 2020. DOI 10.1109 htt///iiDDaaa50022222000200 200.000000000222207088.78111 90.7999971. Cited on page 27.
15

---

Published in Transactions on Machine Learning Research (07 (025))
Appendix
Appendix A reviews the learning coofficient, providing some simple toy examples contrasting the learning coelficient with Hessian-based measures. This section also discusses SGLD-based LLC estimation including experiment hyperparameters (Appendix A.4), and offers detailed example of the calibrations involved in applying LLC estimation to regression transformers to serve as reference (Appendix A.5). Appendix B provides further detail on our procedure for LLC-based stage identification, including stages identified in additional training runs and a brief comparison with Hessian statistics. Appendices C and D examine the developmental stages of language models and in-context linear regression in more detail and explain the various metrics we use track behavioral and structural development. Appendix E describes some additional experiments one-layer language model. Appendix F covers transformer training experimental details, such model architectures, training procedures, and hyperparameters.
To facilitate reproduction of our analyses, we have made our codebase available. A repository containing additional figures and code can be accessed at the URL https:/ http://ggttthbboooo www/timam/tttssstteeeeemmmmmaaee htt/:/tmaausssseee researchhre
A The local learning coefficient (LLC)
17
A.1 Formal Definition of the LLC
A.2 Interpretations and examples of the L.C.
17
A.3 Estimating LLCs with SGLD
18
A.4 LLC estimation experiment details
18
A.5 A guide SGLD-based LLC estimation
19
A.6 LLC estimates for a han-log-likelihood based loss
................................
23
A.7 LLC estimates away from local minima
B LLC-based stage boundary identification
25
B.1 Procedure for stage boundary identification
25
B.2 Stage boundary identification details for language model
................................
25
B.3 Stage boundary identification details for in-context lincar regression
26
B.4 Stage identification for additional training runs
26
................................
B.5 Comparison Hessian statistics
27
C Developmental analysis of language models
28
C.1 Behavioral development
................................
28
C.2 Structural development
29
D Developmental analysis of regression transformers
34
D.1 Behavioral development
................................
34
D.2 Structural development
36
E One-layer language model experiments
44
F Transformer training experiment details
46
F.1 Language models
46
F.2 In-context linear regression transformers
47
16

---

Published in Transactions on Machine Learning Research (07 (275)
A The local learning coefficient (LLC)
A.1 Formal Definition of the LLC
In the setling of Section 4, let B be elosed ball around w" such that w" is a global minimum B, by which mean point with (equal) lowest loss. If there are multiple such global minima, the volume asymptotics are determined by the geometry of one that is most degenerate in the precise sense of SLT, formalised in Lau et al. (2025), roughly corresponding to having the lowest LLC, We call this minimum the maximally degenerate global minimum on B. Consider the volume of the set of nearby low-loss parameters,
v(e) = f [(((w) & {(w") + +e) dw.
As e + 0, V(e) is asymptotically equivalent to
oe A(w") log 1og(1/() log()/))((w>))
where \(w") is the LLC, m(w") is another geometric quantity called the local mulliplicity, and c 0 is a constant.
A.2 Interpretations and examples of the LLC
In Section 4, introduced the LLC quantification of geometric degeneracy. In this section, discuss additional perspectives on the LLC as count of the "effective" dimensionality of parameter, and we give additional examples of the LLC. We refer the reader to Watanabe (2009) and Lau et al. (2025) for more discussion.
The LLC has some similarity to an effective parameter count. If the population loss looks like a quadratic form near w" then X(=*)= = half the number of parameters, which we can think of d contributions of 2 from every independent quadratic direction. If there only d - 1 independent quadratic directions, and one coordinate wi such that small variations in wy near w do not change the model relative to the truth d-1 (this dimension is "unused") then X(w") -
2
The situation bocomes more intricate when certain dimensions are degenerate but not completely unused, varying to quartic or higher order near the parameter (rather than being quadratic or flat). While every unused coordinate reduces the L.I.C by 3, changing the dependency on a coordinate from quadratic (wf) quartic (w?) (increasing its degeneracy while still "using" it) reduces the contribution to the LLC from y to 4'
As source of intuition, we provide several examples of exact LLCs:
((wi, W2, wa) awi + bwg + cuž with a, b 0. This function is nondegenerate, and X(O, 0, 0) + - 3. This is independent of a, b. That is, the LLC X does not measure curvature. For this reason, it is better to avoid an intuition that centers on "basin broadness" since this tends to suggest that lowering a, b, c should affect the LLC.
((w,, w2, wa) = wf +w m
+ 1 in R3 is degenerate, but its level sets are still submanifolds and A00,0,0) + In this case the variable Wy unused, and so does not contribute to the LLC.
f(wi, W2, wa) = wp + we + wg is degenerate and its level sets are, for our purposes, not submanifolds. The singular function germ (f, 0) is an object of algebraic geometry, and the appropriate mathemat- ical object is not manifold or variety but scheme. The quartic terms contribute to the LLC, 1. The higher the power of variable, the greater the degeneracy
that A(O, 0) 1 + 1 and the lower the LLC.
Figure 2 offers several additional examples, from left. right:
17

---

Published in Transactions on Machine Learning Research (07/2025)
A quadratic potential & (1 w2) wî + whz for which the LLC is maximal in two dimensions, (0, = 11
A quartic potential (y(,,, ([) = wi + w2, for which the LLC is (0, 1/2.
An even more degenerate potential (z(w), w2) wiws, for which Aa(0, - 1/4. We note that Hessian-derived metrics cannot distinguish between this degeneracy and the preceding quartic de- generacy.
A qualitatively distinct potential (sw1,, w2) = (wi - ()? + w3)* from Lau et al. (2025) with the same LLC at the origin, A4 0) = 1/4.
While nondegenerate functions can be locally written quadratic forms by the Morse Lemma (and are thus qualitatively similar to the approximation obtained from their Hessians), there is no simple equivalent for degenerate functions, such as the population losses of deep neural networks.
A.3 Estimating LLCs with SGLD
We follow Lau et al. (2025) in using SGLD to estimate the expectation value of the loss in the estimator of the LLC. For a given choice of weights w" we sample C' independent chains with TSGLD steps per chain. Each
chain c is sequence of weights {w, Tr=1 of an observable O by
(c) ToeLD, From these samples, estimate the expectation E Bu|w-- .aluir,,,0 [O(w)]
TSGLD Σ
Oww(1),
CTsGLD c=1
with an optional burn-in period. Dropping the chain index c, each sample in a chain is generated according to:
W--1 Wy + Awr.
(6)
W1 w",
(7)
where the step Aw, comes from SGLD update
Aw,
Vn + (w+ (w: ww') + +N((0,).
(8)
In each step I we sample mini-batch of size m and the associated empirical loss, denoted elr) is used to compute the gradient in the SGLD update. We note that LLC estimator defined in (3) uses the expectation ((w)) which in the current notation means we should take O(w) to be (n(w). For computational efficiency follow Lau et al. (2025) in recycling the mini-batch losses l. (w()) computed during the SGLD process. That is, we take O = e(T) rather than O = En-
Time and Space Complexity. The computational cost. per LLC estimate proportional a slandard training step, denoted S. We expect to require a constant number CTSGGLD of samples (on the order of 102 104) to yield robust estimates, independent of model size. Using logarithmically spaced checkpoints, the total computational complexity for generating an LLC curve over the entire training process scales O(SCTSGLD log N), where N is the total number of training steps. The space complexity incurs a modest linear overhead compared to standard SGD, requiring storage for one additional copy of the weights to enable localization.
A.4
LLC estimation experiment details
A.4.1
LLC estimation details for language models
For language models, we use SGLD to sample 20 independent chains with 200 steps per chain and sample per step. For the one-layer model, we used 0.003, 7 300, and for the two-layer model used
18

---

Published in Transactions on Machine Learning Research (07 (07/2025)
c= 7 100. Estimating the LLC across all checkpoints took around 200 GPU hours for the two-layer model on single A100 and around 125 GPU hours for the one-layer model. For additional runs of the two-layer model, we ran fewer chains, bringing the time down to about 2 TPU hours per training run.
We sampled a separate set of million lines (lines 10m-11m) from the DSIR filtered Pile, denoted as Dagid- The first 100,000 lines from this SGLDD set. (lines 10m-10.1m) were used as validation set. The sampling of batches for SGLD mirrored the approach taken during the primary training phase. Each SGLD estimation pass seeded analogously, so, at different chockpoints, the SGLD chains encounter the same selection of batches and injected Gaussian noise.
Table 1: Hyperparameters for estimating the LLC for language models.
Hyporparameter
Category
Description/Notes
1-Layer 2-Layer
C
Sampler
# of chains
20
TSGLD
Sampler
# of SGLD steps / chain
200
c
SGLD
Step size
0.003
0.001
7
SGLD
Localization strength
300
100
nB
SGLD
Inverse temperature
21.7
m
SGLD
The size of each SGLD batch
100
μ Data
Dataset size for gradient minibatches
13m
A.4.2 LLC estimation details for in-context linear regression
For in-context linear regression models, we generate a fixed datasel of 22 samples. Using SGLD, we sample 10 independent chains with 5,000 steps per chain, of which the first 1,000 are discarded as a burn-in, after which we draw observations once per step, at temperature nB 66.7, 0.0003, and y 13.3, over batches of size 1024. LLC estimation takes up to 72 CPU-hours per training
Table 2: LLC estimation hyperparameters. A summary of the hyperparameters involved in estimating the LLC and the default values we use.
Hyperparameter
Category
Description/Notes
Default Values
C
Sampler
# of chains # of SGLD steps / chain
10
TSGLD
Sampler
5000
SGLD
Step size
0.0003
Y
SGLD
Localization strength
13.3
n//
SGLD
Inverse temperature
66.7
SGLD
The size of each SGLD batch Dataset size for gradient minibatches
1024
Data
220
A.5
A guide to SGLD-based LLC estimation
This section walks through some of the hyperparameter choices and sweeps involved in calibrating LLC estimates. We provide it as a reference for others sceking to adjust LLC estimation to novel settings.
A.5.1 Varying the temperature
In Lau et al. (2025), the inverse temperature / set a fixed "optimal" value 8 1/log n, where is the number of training samples. In practice, we find that it can be advantageous to sample a higher temperature.
19

---

Published in Transactions on Machine Learning Research (07 (07/2025)
(a) )(we) Dependence on Dataset Size
#= 1024
200
2048
# 4096
150
== 8192
((w))
16384
100
μ 32768
μ 65536
50
μ 131072
μ= 262144
524288
10
10a
10
Figure A.1: Past some threshold, the choice of validation set size from which SGLD batches are sampled has little effect on learning coefficient estimates. Estimation hyperparameters are 8, TSGLD 000, m 1024, 0.0003, 9 0.01, 3 0.01. Loss is evaluated gradient minibatches at representative selection of checkpoints. LLCs quickly converge to constant value the size increases.
12 2024
12 2048
ja 4096
ju 8192
100
Step
507
8303
27272
122222
0193333
12:00
200
200
200
3000
Figure A.2: The size of SGLD minibatches has negligible effect LLC estimates (at least among the batch sizes considered). Top: Loss is evaluated the same minibatch as the SGLD gradients. Bottom: Loss is evaluated newly sampled minibatch from the SGLD gradients (of the same size). Estimation hyperparameters are C = T.GLD = 2, 000000 220
20

---

Published in Transactions on Machine Learning Research (07/202) 2022)
Token
2111
A
200
ALP
150
A141
2
50
A
Dift. seed
Same seed
10a
10a
104
105
Step
Figure A.3: Consistently seeding SGLD estimates at each checkpoint smooths out the resulting LLCC-ver- time curve. Except towards the end training (this plotted over a log time axis), the difference is barely noticeable. Variable seeds yield a noisier set of estimates.
Since 3 always shows up in a product with n (in (8) for the SGLD step and in (3) for the LLC), we can view the inverse lemperature as a multiplier that adjusts the effective size of your dataset. In Bayesian setting, 3 2 would mean updating twice on each of the samples in your dataset.
The problem with the default choice of 3 is that size to prevent the update from becoming ill-conditioned, and this eventually causes the gradient term to suppress the noise term. This, in turn, leads to requiring larger batches to suppress the gradient noise and requiring longer chains to sufficiently explore the local posterior (Appendix A.5.3).
increase
have to decrease the SGLD step
Instead n33 = n/ log n, we perform LLC estimation at n,3 = m/ log m, where m is the SGL.D batch size.
A.5.2 Seeding the random noise
To smooth out the Å, curves, we reset the random seed before LLC estimation run at each checkpoint. This means the sequence of injected Gaussian noise is the same for LLC estimation runs at different checkpoints. Additionally, if the batch size held constant, the batch schedule will also be constant across different estimation runs. Figure A.3 shows that this does not alfect the overall shape of the learning coefficient curves; it simply smooths it out.
A.5.3 Calibrating B, and 7
As a rule of thumb, should be large enough that the Å estimate converges within the TSGLD steps each chain but not too large that you run into issues with numerical stability and divergent estimates. Subject to this constraint, 7 should be as small as possible encourage exploration without enabling the chains to "escape" nearby better optima, and 8 should be as large as possible (but no greater than 1/logn).
To determine the optimal SGLD hyperparameters, we perform grid sweep reparamctrization of the SGLD steps in terms of 7,5.
Aw; (Ve() + 7("" wt) + +N(,ee),
where 3 = e,nn2, - e7/4.
The results of this hyperparameter sweep are illustrated in Figure A.4 for final checkpoints. Separately (not pictured), we check the resulting hyperparameters for subset of carlier checkpoints. This noeded since, for example, a well-behaved of hyperparameters at the end of training may lead failures like divergent estimates (Figure A.5) earlier in training when the geometry is more complex and thus the chains less stable.
21

---

Published in Transactions on Machine Learning Research (07 (07/2025)
ElG.ol) So Varion e.y and Por 27072)
2-800L--00
Figure A.4: Results of grid sweep over SGLD hyperparameters for model at t = 27k.
A.5.4 LLC traces
As a useful diagnostic when calibrating the LLC estimates, we propose an online variant for learning coef- ficient estimation. When overlaid on top of individual-chain LLC traces, this helps reveal common failure modes like divergent estimates, non-converged estimates, and escapes (Figure A.5). These traces display the running estimate of Å as a function of the nunber of steps taken in a chain (with the estimate averaged across independent chains).
Define (wWw), the LLC at wi after t time-steps for a single SGLD chain as follows (Lau al., 2025):
1 T
(u(w) (n(cv)))
(wo) = n2
Σ
T
Moving terms around, we get,
{-(w)) logn n
Σ
En(w,)) (x(ww)
(9)
r'=1
nn
T 1-1
La(w() - -w(ww) + ln (n(w)))
+ 1 (wr))))) - la ta(w)))
(10)
r
x==1
T ) + nB
(n(w))) + --- -1) en (wo)
(11)
1
( (wo) + nz (Ln (wr) (n(ww))).
(12)
where
(n(wn) = 0.
22

---

Published in Transactions on Machine Learning Research (07/202) 2025)
23
00
33
530
.250
2500
Drewy
Diar.
(a) Numerical instability
(b) Non-convergence
-50
Sech
6.2
-00
1000
255
1200
Drow.
(c) Negative estimates
Figure A.5: Failure modes of SGLD estimation. Top left: the gradient term is too large, leading to issues with numerical instability and exploding À estimates. Top right: e is too small, leading to å not converging within each chain. Bottom: the localization term is too small, which allows the chain to escape to better optima.
This can be easily extended to an online estimate over chains by averaging the update n33 (L,(w,) ln (wo)) over multiple chains.
A.6 LLC estimates for a non-log- likelihoodl----ooeeo basedd loss
the main body, we apply the LLC to empirical loss functions that do not arise as the log likelihood of independent random variables, due the repeated use of dependent sub-sequences. Here we explain that it is possible define proper negative log likelihood over independent observations for the in-context linear regression setting: similar observations can be made in the language modeling setting.
Let (T(k) be probability distribution over the context length k. Tdeally, the transformer would be trained to make predictions given a context of length & where k is sampled from 11. With the given distribution over contexts this leads to a negative log likelihood of the form
I(w)) [peT(p(w))
(13)
where Pk is the probability of sampling & from II and
= (w) - (SS,, yk Uxlt, k))(t) (w (f(()) ve] dS& dyk dt
(14)
using the notation of Section St (x1, Vi....,, Ft- 1k-1, .x) is a context of length k. is straight.or- ward to check that this negative log likelihood L agrees with the population loss associated the empirical loss defined in Section 3. However the empirical quantities L, (w) and (x(w) defined for set of samples of size n are not the same.
Since we use the empirical loss e, in our calculation of the estimated LLC, whereas the foundational theory of SLT is written in terms the empirical negative log likelihood Ln, it natural to wonder how much of a difference this makes in practice. Figure A.6 depicts LLC traces (Appendix A.5) for highlighted number of checkpoints using either likelihood- based estimate (with variable sequence length) or loss-bascd estimate (with fixed sequence length). The relative orderings of complexities does not change, and even the values
23

---

Published in Transactions on Machine Learning Research (07/2025)
of the LLC estimates do not make much of difference, except at the final checkpoint, which has higher value for the sub=sequencc--based estimate.
Subsequence
Step
270
507
[- 1124
3247
-
200
9375
150
121212
250
500
1000
1250
1920
1750
2000
230 500
750
1900
1290
1550
1350
2000
Figure A.6: Loss-based (left) and likelihood- (right) LLC estimation yicld identically ordered LLC estimates. With the exception of final checkpoint is LLC estimate (which is larger for the loss-based estimale), the values are close to identical. These plots display LLC traces, which show the LLC estimate as a function of SGLD steps. This is a useful tool for calibrating LLC estimation (Appendix A.5).
A.7 LLC estimates away from local minima
Our methodology for detecting stages is to apply LLC estimation to compute \(w*) at neural network parameters w* = W, across training. In the typical case these parameters will not be local minima of the population loss, violating the theoretical conditions under which the LLC is defined.
It not surprising that the estimator appears to work if w' is approximately a local minima. Lau et al. (2025) validated their estimator at both parameters constructed to be local minima of the population loss and also at parameters found through training with stochastic gradient descent (possibly not local minima of the empirical loss, let alone the population loss). They showed that in both cases the estimator recovers the true learning coefficient associated with the global minimum of the population loss.
On the other hand, if w² is far from any local minima, it is priori quite surprising that the SGLD-based estimation procedure works at all, as in this situation one might expect the chains to explore directions in which the loss decreases. Nevertheless, Chen et al. (2023) found that, empirically, LLC estimation away from local minima appears give sensible results in practice. In our case, with sufficient localization we see stable estimates throughout training.
Theoretically accounting for this phenomenon is interesting open problem. Perhaps there is notion of slubly evolving equilibrium in the setting of neural network training, echoing some of the ideas of Waddington (1957), such that the LLC estimation procedure is effectively giving us the LLC of different potential to the population loss potential for which the current parameter actually is at a critical point. We leave addressing this question to future work.
24

---

Published in Transactions on Machine Learning Research (07 (07/2023) (025)
B LLC-based stage boundary identification
B.1 Procedure for stage boundary identification
To identify stage boundaries, look for plateaus in the LLC: checkpoints at which the slope of X(w) over vanishes. To mitigate noise in the LLC estimates, first fit Gaussian process with some smoothing to the LLC-over-time curve. Then we numerically calculate the slope of this Gaussian process with respect to log'. The logarithm corrects for the fact that the learning coefficient, like the loss, changes less as training progresses. We identify stage boundaries by looking for checkpoints at which this estimated slope equals zero. The results of this procedure are depicted in Figure B.1 for language and Figure B.2 for in-context linear regression.
At local minima maxima of the estimated LLC curve identifying a plateau from this estimated slope is straightforward, since the derivative crosses the x-axis. However at saddle point, the slope may not exactly reach zero, so we have to specify a "tolerance" for the absolute value of the derivative, below which we treat the boundary as an effective plateau.
In this case, we additionally require that the plateau be at a local minimum of the absolute first derivative. Otherwise, we may identify several adjacent points as all constituting stage boundary.
To summarize, identifying stage boundaries is sensitive to the following choices: the intervals between check- points, the amount of smoothing, whether lo differentiate with respect to or log l, and the choice of tolerance. However, once a given choice of these hyperparameters is fixed, stages can be automatically identified, without further human judgment.
B.2 Stage boundary identification details for language model
Figure B.1 displays the test loss and LLC curves from Figure 1a in addition to the weight norm time and associated slopes. Stage boundaries coincide with where the slope of the LLC crosses zero, that is, where there is a plateau in the LLC.
6x 10P
180
3x
5x 10"
i(w))
i(w.))
160
[Wr]
140
10a
10P
10
10
10P
10P
0.25
0.04
20
0.00
64(wg)6llgt
0/w///tt
-0.25
0.02
-0.50
-0.75
0.00
10a
102
10
10F
102
10
102
10F
10"
Training stop
Training stop
Trinnnng
LMI
L.M2
1.M3
LM4
LM5
Fit
Figure B.1: A more detailed version of Figure 1a for two-layer language models. Top: Loss, LLC, and weight norm, along with an overlaid Gaussian process fit these curves (red dotted lines). Bottom: Associated slopes, both numerically estimated finite differences (transparent blue) and of the Gaussian process (red dotted lined). Note that stage LM5 may be subdivided into further stages (Appendix B.1). However, the noise in LLC estimates late in training is high, so do not draw any conclusions from this.
25

---

Published in Transactions on Machine Learning Research (07 (07/2025)
B.3 Stage boundary identification details for in-context linear regression
Figure B.2 displays the test loss and LLC curves from Figure 1b in addition to the weight norm over time, and numerically estimated slopes associated to these three metrics. As in the case of language models, we identify stage boundaries by looking for plateaus in the LLC. Unlike the language models, here the boundaries LR1 LR2 and LR2 LR3 are clearly visible in the loss.
4x
200
i(we)
3
i(w-))
10
2x10P
10
10
109
0.0010
bitwcl/dlogt
(Am))/ooggt
6jw//0t
0.0005
-100
0,0000
10a
10P
10
10
10a
102
10
10
10c
102
10
10a
Training stop
Training
Tmining
LR1
LR22 LR3 1.R4 1R3 Fit
Figure B.2: A more detailed version of Figure for in-context linear regression. Top: Loss, LLC, and weight norm, along with an overlaid Gaussian process fit. these curves (red dotted lines). Bollom: Associated slopes, both numerically estimated finite differences (transparent blue) and of the Gaussian process (red dotted lined). Top middle: Error bars displaying the standard deviation over the 10 SGLD chains are displayed in the background. Note that large error bars across chains are to be expected. Between different SGLD estimations, the variance is much lower. For example, averaged over training, the standard deviation over different seeds is only 4.2.
B.4 Stage identification for additional training
Tent los Itwe)
kow (tw)
10c
10*
101
10
10
10a
Lacal loming caett Atee)
Local leaning coett Nome
150
150
125
100
50
10
10'
10f
10
10'
10a
10
Trining step
Thining stop
(a) Two-layer attention-only language transform-
(b) In-context linear regression transformers.
Figure B.3: Figure la and Figure 1b for multiple seeds. In both settings, LLC reveals consistent set of stages across five seeds. Late-training behavior shows more variance across seeds (see Appendix B.4).
26

---

Published in Transactions on Machine Learning Research (07 (025))
Figure B.3a shows loss and LLC curves for five seeds (differing in model initialization and batch schedule). In each seed, LLC estimation reveals stage I.MILMM4 In three of the five seeds, stage LM5 is subdivided into two additional stages.
Figure B.3b shows loss and LLC curves for five unique seeds (differing in model initialization and batch schedule). each seed, LLC estimation reveals stages LR1-L. There remarkably little variance across different seeds.
B.5 Comparison to Hessian statistics
Figure B.4 shows a quantification the curvature-based notion of flatness captured by the Hessian (in contrast the degeneracy- notion of flatness captured by the LLC) for our in-context linear regression transformer. To estimate the trace and maximum eigenvalues shown in this figure, library (Yao et al., 2020) over batch of 1024 samples.
the PyHessian
Crucially, we observe that these Hessian-dent -erivee metrics (henceforth, "curvature") and the LLC are not consistently correlated. During the first part of LR2, the LLC and the curvature are jointly increasing. Starting at around = 20k, while the LLC is still increasing, the curvature starts decreasing. In the first part of LR3, both metrics decrease in landem, but as around 120k, the curvalure turns around and starts increasing.
The Hessian fails to detect three of the four stage boundaries identified by our LLC-based methodology. Since these Hessian-based metrics are dominated by the largest eigenvalues the directions of maximum curvature- they fail to observe the finer-graines measures of degeneracy that dominate the LLC. Moreover, we observe that LLC estimation is more scalable (empirically, it seems to be roughly linear in parameter count) than estimating the full Hessian (which quadratic).
10a
Hessian Statistics
Eigawalue 1
Eigenvabee
Eigenvatoe d
Hessian
10
10
Figure B.4: Hessian-based statistics reveal only stage boundary in the development of our in-context linear regression transformer.
27

---

Published in Transactions on Machine Learning Research (07 2025)
C Developmental analysis of language models
In this section, we present further evidence on behavioral (Appendix C.1) and structural (Appendix C.2) development of the language model over the course of training.
C.1 Behavioral development
C.1.1 Bigram score
We empirically estimate the conditional bigram distribution by counting instances of bigrams over the training data. From this, we obtain the conditional distribution @((')(, the likelihood that a token 1' follows t. The bigram score D at index of input context S is the cross entropy between the model's predictions p(t)+1lfa) at that position and the empirical bigram distribution,
Avocals
B =
Σ a(c (ogl4s)oogp(((0)) logg(((ff) ((11)s)
(15)
where the t range over the possible second tokens from the tokenizer vocabulary. From this we obtain the average bigram score
n B
B kg
(16)
where we take fixed random sequences of kę and St for isisn= 5, 000, which is displayed over training in Figure 4(a). This is compared against the best-achievable bigram score, which is the bigram distribution entropy itself, averaged over the validation set.
C.1.2 n-gram scores
In stage LM2 we consider n-grams, which are sequences of n consecutive tokens, meaning 2-grams and bigrams are the same. Specifically, we consider common n-grams, which defined heuristically by comparing our 5,000 vocab size tokenizer with the full GPT-2 tokenizer. We use the GPT-2 tokenizer as our heuristic because its vocabulary is constructed iteratively by merging the most frequent pairs of tokens.
We first tokenize the tokens in the full GPT-2 vocabulary get list of 50,257 n-grams for various The first 5,000 such n-grams are all 1-grams, after which 2-gramss begin appearing, then 3-grams, 4-grams, and so on (where 2-grams and 3-grams may still continue to appear later in the vocabulary). We then define the set of common r--rams as the first 1,000 n-grams that appear in this list for a fixed n, n 2 2.
If we track the performance on n-grams and see it improve, we may ask whether this is simply a function of the model learning to use more context general, rather than specifically improving on the set of n-grams being tracked. We measure performance against this baseline by defining an n-gram score. For a fixed n, obtain the average loss en gram of the model predicting the final tokens of our set of 1,000 n-grams and
en also obtain the average loss test of the model on validation set at position n of each validation sequence. The n-gram score is then defined to be / gram
C.1.3 In-context learning score
The in-contezt learning score is a behavioral measure of the relative performance of model later in sequence versus earlier in the sequence. We define ICLK, to be the loss token lo minus the loss token ki, so more negative score indicates better relative performance later in the sequence. A more negative ICL score does not, however, mean that a model is achieving better overall loss on later tokens; it is only about the relative improvement. For the language model, we follow a similar construction as Olsson et al. (2022), where we take h/ to be the 500th token and k; be the 50th token. This is then averagod over 100k-row validation dataset. The performance of the language model over the course of training can be seen at the bottom of Figure 4(f).
28

---

Published in Transactions on Machine Learning Research (07/ 2225))
C.1.4 Visualizing behavioral changes
In Figure C.1, we visualize changes in the model's imput/output behavior by comparing model predictions before and after developmental stages and highlighting tokens with the grealest differences.
(a) LMI (0 90-
(c) LM3 LM4 6,,00 17 17,000)
<lendoftextb> should like, before progeeding further, lo tell you how I feel about the State which we have described. I might compare myself to person who, on beholding beautiful animals either created by the painter's art, or, better still, alive but al rest, is seized with desire seeing them in motion or engaged some strugele or conftici to which their forms appeur suited;
<endoftexti>-Mr and Mrs. Dursicy, of number four, Privet Drive, were proud to say that they were perfectly normal, thank you very much. They were the last people you'd expect be involved anything strange mystcrious. because they just didn't hold with such nonsense. Mr. Dursley was the director of firm called Grunnings. which made drills, He was big. beefy man with hardly any neck, although he did have very large mustache. Mrs. Dursley was thin and blonde and had nearly twice the usual amount of neck, which came in very useful as she spent so much of her time craning over garden fences, spying on the neighbors. The Dursleys had a small son called Dudley and in their opinion there was no finer boy anywhere.
(b) LM2 (900 - 6,500)
dendoftexth> the midst of unexpected circumstances with Lingx and Python, the honorable Supreme Court in Boston delivered a ruling ampiiasizigg a crazy database framewrr last week.
Figure C.1: Samples are shown with tokens highlighted to indicate changes in logits during a given range. Red is improved performance (higher logit output for the true next token) and blue is worse. Sample (a): improvement in bigrams (LM1) such as "te/ll, ab/out, des/ire, mot nooninn, eng/aged, strugg/le, etc." Sample (b): improvement in common n-grams (LM2) such as "L/in/ux, P/y/th/on h/on/or/able S/up/reme, dat/ab/ase f/ram/ew/onn LM4), visible in the improved predictions in the word "D/urs/ley' after the first time it appears in the context, as initially observed by (Olsson et al., 2022).
Sample (c): development of in-context learning via induction circuits (LM3,
C.2 Structural development
C.2.1 Positional embedding
In Figure C.2, we mcasure the effect of the positional embedding on model performance by comparing the model's performance at particular context positions on a validation set over the course of training against performance on the same validation set but with the positional embedding zero-ablated The full context length is 1024, and we measure test loss at positions 1, 2, 3, 5, 10, 20, 30, 50, 100, 200, 300, 500, and 1000. In the transition from stage LM1 to LM2, the model begins using the Icarnable positional embedding to improve performance. The difference between Lest loss with and without the positional ablation is negligible at all measured positions until the LM1 LM2 boundary.
Structurally, might predict that the positional embeddings should organize themselves in particular way: in order to understand relative positions, adjacent positions should be embedded close to each other, and far-away positions should be embedded far apart.
In Figure C.3, we examine the development of the positional embedding itself over time from two angles. The first is to take the embeddings of each position in the context and to run PCA on those embeddings. The result is that as training progresses, the positional embedding PCAs gradually resolve into Lissajous curves, suggesting that the positional embeddings might look like a random walk (Antognini & Sohl-Dickstein, 2018; Shinn, 2023). However, if we look to the explained variance, 94.2% at training step 6400. This is much higher than we would expect for Brownian motion, where we expect to see about 61% explained variance in PCI (Antognini & Sohl-Dickstein 2018).
that it grows very large for PC1, reaching
The second perspective we use is to look at how the magnitudes of positional embeddings over the context length develop. In this case, we observe that the magnitudos seem to have a fairly regular structure. In conjunction with the PCAs and explained variance, we might infer that the positional embeddings look approximately like (possibly eurved) line in dmadel 256 dimensional space. A positional embedding
29

---

Published in Transactions on Machine Learning Research (07/202) 2025)
with posiiional ablation
Truining step
Training step
Figure C.2: The model learns to start using the positional encoding in LM2, when the performance starts to worsen when ablating the positional encoding. In both plots, earlier token positions are colored more purple, while later token positions are more yellow, and the overall mean loss is colored in red. Both sets of per-token losses are shown both graphs for case of comparison. Left: original test loss is emphasizod. Right: test. loss with the positional embedding ablated emphasized.
organized in this way would make it easier for attention head to attend to multiple recent tokens, which is necessary if a single head is to learn n-grams.
30

---

Published in Transactions on Machine Learning Research (07/202) 2025)
Training susp
Training stop 400
Training step 800
Trining siep 1600
Tmining siep 3200
Tmining step 64000
50
2.5
2.5
80
00
00
00-
2-
-25
-02
PC2
PC2
PC2
PC2
PC2
PC
0.5
1.
0.1
04
0.90
00-
-025
-0.2
0.0
02
-0.
0.0
0.2
0.0
0.5
022
0.5
0.2
-02
-025
0.0
0.2
-0.25
0.00
0.25
-2.5
0.0
PC1
PC
PC1
PC1
PC1
PC
Bygirdd
0.008
vaiaace
04
0.
10.
0.5
0.6
02
04
04
04
02
02
0.90
1.2
0.8
Manitaic
Manitaia
Manitix
Manitiia
3.25
0.
1.00
6.70
0:75
500
1000
500
1000
500
1000
500
1000
500
1000
500
1000
Figure C.3: Columns progress through training time at training steps 0, 400, 800, 1600, 3200, and 6400. The first. three rows are plots of the first three principle components of PCA on the positional embedding weights, while the fourth row shows the explained variance for each of the principal components. The fifth plots the magnitude of the embedding of cach position in the context length of 1024.
31

---

Published in Transactions on Machine Learning Research (07 (225))
C.2.2 Composition scores
Let W W W be the query, key, and value weights attention head h respectively. There are three types of composition between attention heads transformer models Elhage et al. (2021):
Q- Q-Composition: the query matrix W of an attention head reads in a subspace affected by previous head
K- Composition: the key matrix W of an attention head reads in subspace affected by previous head
V-Composition: the value matrix W of an attention head reads in a subspace affected by previous head
If W is the output matrix an attention head, then W = W TW and W = wowe The composition scores are
MIM ||
(17)
Where M Whot M Who QK: and M Wo for Q-, K- and V-Composition respectively. See Figure C.4 for K-composition scores over time between attention heads the induction circuits.
QK
C.2.3 Previous-token matching score
The previous-loken malching score is a structural measure of induction head altention. It is the attention score given to [A] by an attention head at [B] in the sequence.. ... [A][B] (i.e., how much the head attends to the immediately preceding token).
We compule this score using synthetic data generating process, generating 10k fixed random sequences with length between 16 and 64. The first token is special "beginning of string" token, and the remaining tokens are uniformly randomly sampled from other tokens in the vocabulary.
For each sample in this synthetic dataset, we measure the attention score that an attention head gives to the previous token when at the last token in the sequence. These scores are averaged across the dataset to produce the previous-token matching score for that attention head at a given checkpoint. The progression of previous-token matching scores over time can be seen in Figure 4(d).
C.2.4 Prefix matching score
The prefiz matching score from Olsson et al. (2022) is defined similarly to the previous-token matching score. Given a sequence [A][]]... [A], the prefix matching score of a particular attention head is how much the attention head attends back to the first instance of [A] when at the second instance of [A].
We compute this score using synthetic data-generating process. We first generate 10k fixed random sequences of length 128. The first token is always special "beginning of string' token and the [A] and [B] tokens are selected and placed randomly. One [A] token is placed in the first half of the sequence, the other is placed in the sccond half, and the [B] token is placed directly after the first [A] token. The remaining tokens are randomly sampled from the lokenizer vocabulary, excluding the [A], [B], and beginning of string tokens.
For cach sample in this synthetic datasct, we mcasure the attention score that each attention head assigns to the earlier instance of [A] from the latter instance of [A]. These scores averaged across the dataset. to produce the prefix matching score for that attention head at given checkpoint. The progression of prefix matching scores over time can be seen in Figure 4(c).
32

---

Published in Transactions on Machine Learning Research (07/2025)
2:1
<> 2:7
1:1 <>
0.20
0.15
0.10
1:2<> 2:1
<
2
K compesiton
0.10
0.05
1:5 2:1
1:5 <3 2:7
1:5 <>>28
0.20
X composson
0:00
10
10
10
10'
10
10a
10a
Training
Training stop.
Teiining slep.
Figure C.4: The K-composition scores (Elhage et al., 2021) between first and second layer attention heads. The hth attention head in layer is indexed by ( h. The attention heads that eventually become previous token heads are h = 2, 5 in layer (subplot rows and 3); and the attention heads that eventually become induction heads are = 8 layer (subplot columns 2 and 3). The attention heads included for comparison. The induction heads 7 and 2 8 begin K- co-oomosingg with first layer heads near the start of stage LM2. They continue to compose with the previous token heads in stages LM3 and LM4 (highlighted in green) while their K-composition scores drop with other attention heads in layer in later stages.
and 2 are
33

---

Published in Transactions on Machine Learning Research (07 (07/2025)
D Developmental analysis of regression transformers
In this section, we present further evidence on the behavioral (Appendix D.1) and structural (Appendix D.2) development of the transformer the setting in-context linear regression.
D.1 Behavioral development
D.1.1 Task prior score
In addition to training models on data distribution in which tasks are generated on-the-fly, we examine the setting of Raventós et al. (2023), in which a finite set of M tasks is generated ahead of time, and training samples involve randomly selected tasks from this set.
Figure D.1 depicts (a) the mean square distance between the model's predictions and the zero prediction addition to (b) the mean square distance between the model's predictions and the "task prior" prediction, using the component-wise averaged over the set of tasks encountered during training. For all models, the minimum distance to the task prior prediction is lower than the minimum distance to the zero prediction. Hence, we call stage LR1 "learning the task prior" rather than simply learning the zero prediction.
(a) MSE from 0 Prediction
(b) MSE from Prediction
10
10
M = 2
M = 2
MSE
MSE
2
10a
10a
M 26
M 2
10-1
10
10
105
101
10
10
Step
Step
Figure D.1: Learning the task prior is universal across models trained on very different data distributions. Each line represents a model trained on & data distribution with a different number of M distinct tasks ("task diversity" in Raventós et al., 2023). In addition to taking finite M, the models depicted here differ from the other models considered in this paper in that the former were trained with maximum learning rate of 0.01, and the models (inadvertently) lack an output matrix after the multi-head attention layer.
D.1.2 ICL
We consider two variants of the ICL score: ICL1.D, and ICLD.K.
If the noise term o² equals zero and both tasks t and inputs are normalized (i.e., t E SD-1), then D-1 observations of input/output pairs are enough to precisely identify t. Therefore, ICL1DD measures how successful the model at initially locating the task. The fact that the tasks and inputs are not normalized changes this only slightly: the task will still sit near SD-1 within shell of vanishing thickness as D - co.
Once localized, ICLD.K measures how successfully the model refines its internal estimate of t with additional examples, which it use to reduce the error due to noise.
In terms of implementation, it's not necessary for the model to internally make distinction between locating and refining its estimate of the lask. For example, ridge regression makes no distinction. Still, we find it useful for reasoning about the progression of the model. In particular, we note that early stage LR2, while the model begins develop ICL for carly tokens, becomes worse at ICL over tokens late the context.. Later, at around 23k steps, ICLD.K stabilizes, while ICL1.D continues improving over the entire training
34

---

Published in Transactions on Machine Learning Research (07 (07/2025)
(a) ICL Scores
(b) ICL Scores (Zoomed)
0
0.0
-0.5
11:4
-1
14:8
1..
-1.
-2.
10a
10
10
10
Training step
Training step
Figure D.2: ICL scores for the in-context linear regression model. Right: ICL scores between imputs 1 and 4 and imputs 4 and 8 over time. We see that ICL emerges during the first half of LR2. Left: Highlighted ICL score curves from the end of LR1 to halfway through LR2. Note that when the model first starts improving on early tokens, it temporarily becomes worse at predicting later tokens. Note also that the model ceases to become better at later tokens as the second half LR2, whereas ICL on early tokens continues to improve throughout training.
D.1.3 OOD generalization
To further investigate behavior in stages LR2 and LR3, we probe the model on data sampled from different distributions than encountered during training." We evaluate behavior two families of perturbations: "OOD inputs" Ik: sampled according to different scale
N(00,gI))
(18)
for some gain parameter g. and "OOD tasks"
t t~~(00.
(19)
Note that these inpuls and tasks are not out-of-listributtio in the sense of coming from a distribution with a different support than the training distribution. However, the samples drawn from these "extreme" distri- butions are exponentially suppressed by the original training distribution. Figure D.3 plots the normalized MSE for these distributions over training time.
Between 1k and 4k the model's outputs rapidly diminish in scale for out-of-distribun samples, both for > 1 and g < 1, especially for out-of-distribution inputs. While the model is moving away from predicting with the task prior for in-distribution samples, it moves closer to predicting with the task prior for-in-distribution samples.
Between and = 23k, the model recovers on moderately out-of-distributiion inputs g < 10 with performance remaining close constant beyond this range. Past this stage, performance improves constantly for out-of-distribation tasks.
For out-of-distribution inputs, performance eventually worsens for some ranges of Between / 23k and t = 80k the model further approaches the task prior prediction for extreme dut-of-distribution inputs > 101.5 Subsequently, between = 75k and t = 130k the model moves away from the task prior prediction for extreme inputs, and performance deteriorates for inputs with g > 100.5. As LR5, performance roughly constant.
Cf. Raventós (2023) evaluating models trained of discrete tasks "true" distribution consisting of novel tasks..
35

---

Published in Transactions on Machine Learning Research (07 (07/2025)
(a) Normalized MSE Extreme Inputs
(b) Normalized MSE on Extreme Tasks
10'
10
MSE g2
6x 109
MSE g
6 109
2.5
100
4 109
3x 109
3 109
2.0
2x
x 109
1.5
mulippierr logg
10a
10a
10
10
10a
10
10a
10a
Step
Step
Average Prediction Norm Extreme Inputs
Average Prediction Norm Extreme Tasks
0.5
0.0
Norm
Norm
10P
-0.5
10-1
-1.0
10-1
10a
10a
10
10
10a
10a
104
10a
Step
Step
Figure D.3: Performance extreme inputs time may reveal additional substages in LR2 and in LR3. Left: The model first becomes better, then worsens at ICL on inputs sampled from ((0,gT)) for large Right: The model continues to improve on ICL at tasks sampled from gl, Top: Normalized loss (divided by g²) over time for OOD inputs and tasks. Bottom: Average [g] over time for OOD inputs and tasks.
D.2
Structural development
D.2.1 Embedding
The embedding matrix WF is a linear transformation from RD+ - Reembed Plotting the + 1 singular values of this matrix, we notice that the embedding partially loses one of its components starting at the end of LR2 (Figure D.4a).
The input "tokens" Ia span D-dimensional subspace of the (D + 1) 1)-dimensional "token space." The target tokens yk span orthogonal 1-dimensional subspace. The collapse of one of the embedding matrix's singular values means that the model learns to redundantly encode the inputs and targets in the same D- dimensional subspace of the space of residual stream activations. The almost order magnitude separation in the magnitudes of the square singular value means that the (D 1)th component. the token embedding explains only 2.9% of the variance in activations of the residual stream immediately after the embedding, whereas the dominant components explain roughly 24% each.
Contributions to degeneracy Given linear transformation Ti PR. + D2 followed by another linear transformation Tz R= + Rs reducing the rank of Ti from r to r' < renders Ps(r-r)) components of the second transformation irrelevant. This would mean a decrense in the learning coefficient of D3(r-r')/2 (a decrease in the effective dimensionality of d leads to a decrease in the LLC of d/22). In the actual model, we don't see an exact decrease the rank, and a layer normalization sits between the linear transformation of the embedding and the linear transformations of each transformer block and unembedding, It is unclear what the precise relation between structure and degeneracy is in this case (Appendix D.2.6). Still, suggestively, the onset of embedding collapse coincides with decrease in the rate of increase of X(w)).
Note that this quartic some higher power would also lower the LLC,, by fractional amount.
only passible way for the LLC to decrease. Changing local loss landscape from quadratic
36

---

Published in Transactions on Machine Learning Research (07 (075))25)
(a) Singular Values Token Embedding
(b) Singular Values Positional Encoding
(c) Cossim between Wembed postn
16
of/Tr
03
10
10
0.6
10
10
10
10a 10
10
10 10"
10'
101
Training
Training step
Training step
Figure D.4: Left: The embedding partially "collapses" during the second half LR2. At the start of stage LR2, the minimum singular values explains only 3% of the variance in residual stream activations due the sample. By the end of training, it explains half that. Middle: The positional encoding goes through similar shift during LR3 (that begins carlier during LR2). Right: The cosine similarity between the 5 rows of Wmdbed and the projection of those rows onlo the subspace spanned by Wnemmbed shows that the model learns to write to the same write tokens and positional information to the same subspace.
D.2.2 Positional encoding
The positional encoding goes through a similar collapse to the unembedding starting during the second part of LR2 and continuing into LR3 (Figure D.4b). Additionally, throughout these stages, the subspace spanned by the embedding becomes more aligned with the subspace spanned by the positional encoding (Figure D.4c).
Contributions to degeneracy For the same reason as with the token embedding, a decrease in the dimensionality of the subspace occupied by activations reduces the effective number of dimensions and thus the learning coefficient. This occurs both as the positional encoding's effective dimensionality decreases (vanishing singular values, Figure D.4b) and the token embedding subspace and positional embedding subspace align (increasing cosine similarity, Figure D.4b).
D.2.3 Attention collapse
Over the course of training, we observe that some attention heads learn to attend solely (soft attention becomes hard attention) and consistently to certain positions (the attention pattern becomes content- independent). We call this phenomenon allention collapse in parallel with the other observed forms of collapse. Not only does this potentially contribute to a decrease in the LLC, but it also makes the attention heads identifiable: we find a self-attention head, previons-attention heads, crrrrous---attettoon heads, and previous-g-attention heads.
r-attention vs. g-attention For convenience we separate each attention head in two: one part for the z-tokens, and the other for the y-tokens.
Attention entropy score To quantify attention hardness, we use the allention entropy score (Ghader & (b.A)
Monz, 2017; Vig & Belinkov, 2019). Given the attention pattern ag. for how much token k in head h in
block b attends back to token W, its attention entropy score H.Ch (b.k) is the Shannon entropy over preceding indices 3'' < k
((0,))
(b.h) (b.A) O log, O'k.k'
(20)
From this, we compute the normalized entropy (b.k) which divides the attention entropy by the maximum entropy for the given context length,
(b.b)
A (b,N)
(21)
log, (k)
37

---

Published in Transactions on Machine Learning Research (07 (07/2023) /025)
b 1, h Xk
b 1, h 2, Xk
1, 3, Xk
b 1, h Xk
2 h)
10+
10a
10a
10
Training stop
Training step
Training step
Training step
b = 1, h = 1
b 1, h h yk
b 1, h 3, yk
b = 1, h 4, Yk
E
10a
10+
10a
10a
10a
10
10a
10
Training step
Training
Training step
Training step t
b = 2, h = 1, Xk
b 2, h = 2, Xk
b = 2, h 3, Xk
b = 2, h 4, Xk
1200 hl
0
0
10a
10
10a
10a
10a
10
10a
10
Training step
Training step
Training step
Training step t
b = 2, h = Yk
b 2, h = 2,
b = 2, h = 3, Yk
b = 2, = 4, Yk
2/ h)
102
10
10a
10a
10a
10a
10a
10
Training step t
Training step
Training step
Training step t
X1
X2
X3
X4
X5
X6
X7
X8
Y1
Y2
V3
Y4
ys
Yo Y7
Ye
Figure D.5: Attention hardening measured by the normalized attention entropy score (Appendix D.2.3). Block 1 heads 1y/3y and block 2 head ly harden over training. In combination with the fact that these attention heads become less variable (Figure D.6), this may contribute to decrease in the LLC (discussed in Appendix D.2.3) The x-components of the attention heads remain much softer over the entire training run.
This accounts for the entropy being calculated over different numbers of tokens and is displayed in Figure D.5. Notably, the identified stages line up closely to stages of these attention entropy curves.
Constant attention Accomplishing constant attention requires the presence of biases in the query and key transformations, or there no bias (as the case for the models we investigated), requires attending to the positional embedding. With the Shortformer-style positional encoding used for the language models (Appendix F.1.1), this is straightforward: the positional information injected directly into the key and weight matrices. With the in-context linear regression models, where the positional embedding is added to the residual stream activations, this is less straightforward: achieving constant attention requires scparating residual stream aclivations into orthogonal positional- and inpul-dependent subspaces, then reading from the former with the query and key weight matrices.
Attention variability score To quantify how constant the attention pattern is, we use measure attention
variability (Vig & Belinkov, 2019),
38

---

Published in Transactions on Machine Learning Research (07 (07/2023) /025)
b 1, h 1, Xk
1, h 2, Xk
b 1, 3, Xk
b 1, h 4, Xk
0.25
0.25
0.00
0.00
0.00
0.00
10a
10
10a
104
10
Training step
Training step
Training step
Training step
b 1, h = 1, yk
b 1, h 2, Yk
b = 1, h 3, Yk
b = 1, h = 4, Yk
Vib.
0.25
0.25
0.25
0.00
0.00
0.00
0.00
10a
10+
10a
10+
10a
10+
10a
10+
Training
Training step
Training step
Training step t
b 2, h = 1, Xk
b 2, h = 2, Xk
b = 2, h 3,
b = 2, h 4, Xk
V
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
10a
10+
10a
10+
10a
10
10a
10
Training step
Training step
Training step
Training step t
b = 2, h = 1, yk
= 2, h 2, Yk
b = 2, h = 3, Yk
b = 2, h 4, Yk
0.5
h)
0.25
0.25
0.00
0.00
0.0
0.00
10a
10
10a
10
10a
10
102
10
Training step
Training step
Training step
Training step t
X1
X2
X3
X4
X5
X6
X7
X8
Y1
Y2
Y3
Y4
ys
Yo
Y7
Ye
Figure D.6: Attention variability over time. The heads that develop hard attention in Figure D.5 (block heads 1y, 3y, and 4y) also become less variable over time.
(b.h) (s))(s(()) (b.h) Ok,
Em ΣN'R
(b.b)
Okkck'
(22)
2n Evec Ok.kk'
(b.1)
where the division by ensures the variability lies in the range [0, 1]. This is displayed in Figure D.6. These reveal that though attention hardness and variability are independent axes differentiation, empirically, wo observe that hard attention is correlated with low variability.
Self-attention score Self-attention is measured by the average amount token k attends to itself, 0
(b.k)
Previous-token attention score Previous-toker token attention is measured the same in the language model setting (Appendix C.2) with one difference: we compute the previous-token score not over synthetic dataset but over validation batch.
z-attention score
(b.h) = >K== The total amount attended to inputs that is 0
=== 0g
(b.h)
y-attention score
Defined analogously (b.b)
ok
ΣK
=b==1
0k224+11
(6.1)
39

---

Published in Transactions on Machine Learning Research (07 (025)
1, - 1y
1,0 3y
1, 4y
b - 2. 1yy
1.0
1.0
0.5
0.4
0.6
0.4
0.2
0.6
0.2
105
100
10:
10
10
101
10
107
10f
107
Texining eup
Thaining cup
Traming dupt
Tounnng dtep
Figure D.7: Collection of attention heads identified by their consistent and recognizable attention patterns. Left to right: provious-zs head, previous-token head, previous-ys head, previous-ys head
Several attention heads are easy identify by virtue of being both concen- trated and consistent. These are depicted in Figure D.7 and include: (B1H3y) previous-token heads (also present in the language model case), (B1Hly) previous-x, and (B1H4x, B2Hlly previous-y heads. Other training runs also include self-attention heads.
Classifying attention heads
Contributions to degeneracy Suppose an attention head block has the following constant attention pattern (after the softmax) (bbh)) = E, di()). That is, for each token i, that attention head attends solely single earlier token 1(i) S i and others. Restricting to single-head attention (the argument, generalizes
straightforwardly), the final contribution of this attention head to the residual stream is the following (Phuong & Hutter, 2022):
= Wo (V. -A)
(23)
where A E IR x RC is the attention pattern, V E Rdout x R& is the value matrix, and Wo Rd x R is the matrix of residual stream activations, and V E Rdout R& is the value matrix. The result of this operation is subsequently multiplied by the output matrix and then added back into the residual stream. Plugging in the hard and constant attention pattern, writing out the matrix multiplication, and filling in the definition of A we get
Oi Σ Vaggg
(24)
For each column in A, the hard attention picks out a single element of V at column ((j) for each row k. Now suppose that there is token l' that receives no attention from any position j. That is, there exists no j such that 1' = I()). Then, there is a column in which does not contribute to the result of V A, and, in turn, column l' in Wo, which does not contribute to the output of the head. As discussed for the embedding and layer norm, this decrease in effective dimensionality leads to docrease in the learning coefficient.
Note that this argument does not hold for all hard and constant attention patterns. It holds solely for attention patterns that consistently ignore some earlier token across all positions, such as the previous-2 and previous-y heads, but not the self-attention and previous-token heads. As discussed in Appendix D.2.6, it remains unclear what exactly the threshold for "ignoring" a token should be before contributes to degeneracy and whether any of the heads we examine actually meet this threshold.
D.2.4 Unembedding collapse
The unembedding block consists of a layer normalization layer LN(x) followed by a lincar transformation Wyz + by and finally a projection my to extract the y-component. Given the 64-dimensional vector of activations in the residual stream right before the unembedding (for a specific token), the full unembedding operation is:
E[x] [
2
Ty
We
07+
+ bu
+c
40

---

Published in Transactions on Machine Learning Research (07/202) 2025)
Weights
1.0
0.0
10
10a
10
107
10
10"
1. Noves
Litoer Blue
0,025
1xt Blas
0.0
0.0
-0.1
-0.5
-0.025
10
104
10P
10P
104
10
107
10#
10"
Training step C
Training stop
Triining step
Figure D.8: Unembedding weights over time for the RT transformer undergo "collapse" that begins lowards the end of LR2. When these weights reach zero in LR3 and LR4, il may contribute to the observed decrease in the LLC. Top: Weights over time. The outlier in the positive direction is the weight for the yr token output. Bottom: Biases over time. Left: Unembedding layer normalization weights over time. Middle: Unembedding linear weights over time (restricted to y-subspace). Right: Effective unembedding weights over time (obtained by element-wise multiplication of preceding columns, and focusing on the bias for only the yoteen.
where @ denotes element-wise multiplication of two vectors and 8 are the layer normalization weights and biases respectively.
Effective unembedding weights and biases
Moving terms around, we can represent this
B[2] V +
(WW)) On)
+ ((W6)1433) + (bu)
where we order the outputs that the y-token corresponds to the 0th row. Because we are reading out a single y component, we can express the unembedding transformation in terms of "effective" unembedding weights and biases
We (Wu)p.) @++ by ((Wo) (bu)
Unembedding weights over time In Figure D.8, we plot (7, (Wo))o.) (bu)) and (We, by) as a function training steps, along with the mean weight over time. These are 64- and 1-dimensional vectors, so we can display the entire set components. During stage LR3 the majority of weights 3 and Wy "collapse" zero. Additionally, the layer normalization biases temporarily experience large increase in variance before returning to small values. Despite this, the mean of the lincar weights, layer normalization biases, and effective weights remains remarkably constant and close to zero throughout the entire process.
Contributions to degeneracy Suppose that D of the layer normalization weights have vanished, say 7= = 0 for i D. Then the corresponding columns of Wy only contribute to the uncmbedding via their product (Wi) with the first D rows of 3. This creates a typical form of degeneracy studied in SLT and found, for example, in deep linear networks, where we can change the weights to (Wu) A A113p for any invertible x matrix A without changing the function computed by the network. in addition the B, vanish for I i < D then the entries of (Wo)j.nn) are completely unconstrained, creating further degeneracy.
41

---

Published in Transactions on Machine Learning Research (07/2025)
Blick: Altention Weight
Blace I Ph-MLP LN Warght
Block Pre Tre-Atetionn LN Weight
Rlock PM-MI PLLP Weight
Unershddiny 1X Waysh
1.0
10P
127
137
107
30
307
10
Block Poo-Aticttos LN Biss
Slack Pre LN Bias
Block Are-Anariee LN Bio
Block 2 PreMIPP 1NN Bis
Unceisoking LN Bia
3.-
08
Figure D.9: Layer norm weights over time. Top: After LR3, the layer normalization collapse expands from the unembodding to carlier layers, most notably in the first pre-attention layer norm. This occurs without explicit regularization and may contribute to the concurrent decrease in LLC. Bollow: During layer normalization collapse, the variance of layer normalization biases increases drastically while the mean of the biases remains relatively constant. Inset: Plotting the fraction of weights biases whose magnitude is less than 0.1 over time reveals that the collapse is more measured for intermediate layer norms: weights shrink to small values but not extremely close to zero as in the unembedding and first attention layer.
D.2.5 Layer normalization collapse
The "collapse" in layer normalization weights not unique to the umembedding. As depicted in Figure D.9, this behavior occurs in all layer norms except for the second MLP. The biases also remain centered close to zero even as the variance in biases grows much larger. Unlike in the unembedding, these layers begin to change earlier (starting halfway through LR2).
What is most striking about the layer normalization collapse is that it occurs without any explicit regulariza- tion (neither weight decay nor dropout). As such, it demonstrates a clear example of implicit regularization. i.e., inductive biases in the optimizer or model that favor simpler solutions.
Contributions to degeneracy In the previous scction, we describe how layer norm collapse in the un- embedding is linked to an increase in degeneracy because it ensures that parameters in the subsequent linear layer become irrelevant. The same is true for layer norm which precedes the attention and MLP blocks.
D.2.6 Degeneracy and development
In the previous subsections, we provide a set of theoretical arguments for how embedding collapse (Ap- pendix D.2.1), layer normalization collapse (Appendix D.2.5), and attention collapse (Appendix D.2.3) can lead to an increase in degeneracy, even while leaving the implemented function unchanged.
The free energy formula tells us that, for two different solutions (sets of weights) with the same loss, the Bayesian posterior will asymptotically prefer the model that has the lower learning coefficient (i.e., higher degeneracy). This suggests that these different forms of collapse may be driven by a bias towards higher degeneracy, as captured the free energy formula. However, in idealized Bayesian inference, we do not expect the postcrior to concentratod around the neighborhood of an equal-loss- -ighssut---hhh---eeeeeeee minimum to begin with. That this kind of transition arises in practice might arise from one of the various differences between Bayesian inference and madient--aseed training.
local
Actually establishing a causal link between incrcasing degeneracy and structure development is beyond the scope of this paper. For one, the theoretical arguments hinge on the collapse being complele, that is, the components that go to zero must become exactly zero in the limit, where we take the number of samples to compute the loss to infinity. In practice, we expect there to be some threshold below which we can treat weights as effectively zero. Second, even if these explanations are correct, we do not know that they account for all of the empirically observed decrease in the LLC during these stages. There may be other drivers
42

---

Published in Transactions on Machine Learning Research (07 (025))
missed. Finally, establishing causal link requires theoretical progress in relating the Bayesian learning process to the SGD learning process. The arguments are suggestive, but currently only source of intuition for how structure and degeneracy can be related, and a starting point for future research.
43

---

Published in Transactions on Machine Learning Research (07 (07/2025)
E
One-layer language model experiments
We also trained and ran some experiments on onc-layer language model (sce Appendix F.1.1 for details). We aggregate results for the one-layer language model here, mirroring the experiments for the two-layer language model where possible. The early development the one-layer model has many parallels with the two-layer model. At a single stage boundary, just as it occurs in the two-layer model, the one-layer model minimizes its bigram score (see Appendix C.1.1), begins utilizing the positional embedding to noticcably improve performance (see Appendix C.2.1); and starts making sudden improvements to the same n-gram scores (sce Appendix C.1.2). Remarkably this occurs the same checkpoint in the 2-layer model (at 900 training steps).
One key difference, however, is that this occurs at the second stage boundary as discerned by the plateaus of the LLC cstimation. We did not closely invostigate why the LLC cstimation appears to drop between steps 400 and 900 this model. a result though, we do observe an interesting qualitative similarity the drop in LLC in stage LM3 of the two-layer model, that this drop precedes noticeable bump in the loss function.
Local leaming Atuu)
Test lose (w))
140
102
Training
Traising
Modic
Por alland Tust
Bandiie
Traningrop
Traning App
Traniig sxpd
Figure E.1: We train a one-layer transformer model in the language setting to compare with the two- layer model. The development of certain behavioral and structural metrics over time closely mirrors the development of the same metrics in the early stages of the two-layer language model. Top: test loss and LLC estimations over time for the one-layer attention-only transformer, compare with Figure la. Bollom: bigram score, test loss with positional embedding ablated, and n-gram scores for the one layer attention-only transformer, compare with Figure 4(a,b.c).
44

---

Published in Transactions on Machine Learning Research (07/202) 2025)
10
150
x
140
2x 10
Rwd
130
Iwd
4x
120
110
x 10
Training step
Training step
Training stop
0.2
0.05
40
0.04
dilne)llogg
0.0
0A/wn1/0logt
20-
d(w,j/0t
-0.2
0.02
-04
-20
0.01
0.66
-40
10a
10'
10
10a
10
10
10P
10a
10a
Step. 2
Step, 2
Step. t
A1
B1
A2
C1
Fit
Figure E.2: more dotailed version of Figure E.1 for the one-layer language model. Top: Loss, LLC, and weight norm, along with an overlaid Gaussian process to these curves (red dotted lines). Bollom: Associated slopes, both numerically estimated finite differences (transparent blue) and of the Gaussian process (red dotted lined).
45

---

Published in Transactions on Machine Learning Research (07 (07/2025)
F Transformer training experiment details
F.1
Language models
F.1.1 Architecture
The language model architectures we consider are one- and two-layer attention-only transformers. They have a context length of 1024, residual stream dimension of dmodde 256, H attention heads per layer, and include layer normalization layers. We also used learnable Shortformer positional embedding (Press et al., 2021). The resulting models have total of d 3, 091, 336 parameters for L = and d 3, 355, 016 parameters for L = 2. We used an implementation provided by TransformerLens (Nanda & Bloom, 2022).
Component
1 L-Layer 2-Layer
Token Embedding Weights
1, 280, 000
Positional Embedding Weights
262, 144
Logits
Layer Layer Norm Weights
256
Layer Layer Norm Bias
256
Unembedding
Layer Attention Query Weights
636
Unembed
Layer 1 Attention Key Weights
65, 536
LN
Layer 1 Attention Value Weights
65, 536
Layer Attention Output Weights
65, 536
X L
Layer Attention Query Bias
256
Layer
Attention Key Bias
256
Attention
Layer 1 Attention Value Bias
256
Projection
Layer 1 Attention Output Bias
256
Layer 2 Layer Norm Weights
N/A A
256
h0
Layer 2 Layer Norm Bias
N/A
256
Layer 2 Attention Query Weights
N/A
65, 536
LN 1
Layer 2 Attention Key Weights
N/A
65, 536
Layer 2 Attention Value Weights
N/A
65, 536
Positional embed
Layer 2 Attention Output Weights
N/A
65, 536
Layer 2 Attention Query Bias
N/A 256
Layer 2 Attention Key Bias
N/A 256
Token embed
Layer 2 Attention Value Bias
N/A 256
Layer 2 Attention Output Bias
N/A 256
Tokens
Final Layer Norm Weights
256
Final Layer Norm Bias
256
Unembedding Weights
1, 280, 000
Unembedding Bias
5, 000
Figure F.1: Attention-only transformers with Shortformer position-infused atlention and pre-layer norm. The lny-layer model has a total of 3,091,336 trainable parameters, while the two-layer model has 3,355,016.
F.1.2 Tokenization
For tokenization, we used truncated variant of the GPT-2 tokenizer that cut the original vocabulary of 50,000 tokens down to 5,000 (Eldan & Li, 2023) to reduce the size of the model. We think this may contribute to the prominence of the the plateau at the end of LMI: the frequency of bigram statistics depends on your choice of tokens, and a larger tokenizer leads to bigrams that are individually much less frequent.
F.1.3 Training
The models are trained on single epoch over 50, 000 steps on ~5 billion tokens using a resampled subset of the Pile (Gao al., 2020; Xie et al., 2023) using a batch size of 100. A snapshot saved every 10 steps
46

---

Published in Transactions on Machine Learning Research (07 (07/2025)
for total of 5000 checkpoints, though majority of analysis used checkpoints every 100 steps. The training time was around 6 GPU hours per model on an A100. Additional seeds were trained on v4 TPUs at around 1.5 TPU hours per model.
Training was conducted on the first 10 million lines of the DSIR-filtered Pile (Xie et al., 2023; Gao et al., 2020) but did not exhaust all million lines. The model was subject to weight decay regularization, without the application of dropout. We did not employ a learning rate scheduler throughout the training process.
Table 3: Summary of hyperparameters and their values for transformer language model training experiments.
Hyperparameter
Category
Description/Noter
Value
n
Data
# of training samples
5, 3,000, 000
T
Data
# of training steps
50, 000
Neest
Data
# of test samples
512
Tokenizer Type
Data
Type of Tokenizer Vocabulary size
Truncated GPT-2 Tokenizer
Data
5,000
Data
Context size
1,024
L
Model
# of layers in the model
Model
# of heads per layer
8
dmlp
Model
MLP hidden layer size
N/A
dembed
Model
Embedding size
256
dhead
Model
Head size
32
seed
Model
Model initialization
1
m Optimizer Type 1
Training
Batch Size
100
Optimizer
Type of optimizer
AdanW
Optimizer
Learning rate
0.001
Awd
Optimizer
Weight Decay
0.05
$1.2
Optimizer
Betas
(0.9, 0.999)
F.2
In-context linear regression transformers
F.2.1
Architecture
In the following L refers to the number of layers (blocks) in the transformer, H is the number of heads in each layer, D the dimension of inputs x E RDD and K is the number of (xr, y) pairs provided to the Transformer in-context.
The architecture is a pre-layer-norm decoder-only transformer modeled after NanoGPT (Karpathy, 2022; see also Phuong & Hutter, 2022) with learnable positional embedding. For the models discussed in the main body, we consider L Hl 4 transformers (with d = 51, 717 parameters), ie., two transformer blocks with four attention heads each.
F.2.2 Tokenization
To run contexts Sx through the above model requires an initial encoding "tokenization step" and final "projection step." The context is encoded as sequence of "tokens" T3 as follows:
91
yk 0
0
T
T1
:
Through the main text, we write (5x) for (Tx). Note that this tokenization includes the final ys token even though this receives no training signal. For this reason, we omit this token from the attention entropy and variability plots (Figures D.5 and D.6).
47

---

Published in Transactions on Machine Learning Research (07/2025)
Component
# of Parameters
Logits
Token Embedding Weight
320
Unembedding
Positional Embedding Weight
1, 024
Linembad
Tayer
Layer Norm Weight 2
64
LN
Layer 1 Layer Norm Bias 1
64
L
Layer 1 Attention Weights
12, 288
Layer 1 Autention Output Weights
4, 096
Block
Layer 1
Layer Norm Weight
64
MLP
Layer 1
Layer Norm Bias 2
64
Projection
Layer 1 Feed-Forward MLP Weight
4, 096
Layer 1 Foed-Forward MLP Bias
64
Layer 1 Feed-Forward Output Weight
4, 096
Layer 1 Feed-Forward Output Bias
64
Layer 2 Layer Norm Weight
64
Attention
Layer 2 Layer Norm Bias
64
Layer 2 Attention Weights
12, 288
Projeciion
Layer 2 Attention Output Weights
4, 096
PA
Layer 2 Layer Norm Weight 2
64
Layer 2 Layer Norm Bias 2
64
LN
Layer 2 Feed-Forward MLP Weight
4, 096
Layer 2 Feed-Forward MLP Bias
64
Layer 2 Feod-Forward Output Weight
4, 096
Embedding
Layer 2 Feed-Forward Output Bias Unembedding Layer Norm Weight Unembedding Layer Norm Bias 1 Unembedding Weight 2
64
Positional embec
Token embed
64
64
Tokens
320
Unembedding Bias 2
Figure F.2: Transformer parameters in the in-context linear regression setting. two transformer blocks for a total of 51, 717 trainable parameters.
The model has
The transformer outputs a series of tokens of the same shape as Tg. To read out the ÿk predictions, we read out the first component of every other token, ie.,
Ty R (D+1)x2K + R
(25)
(0)
(1 :
Uk
( + (fI...,,''' 9k)
(26)
F.2.3 Training
We train from a single seed for cach choice of architecture and optimizer hyperparameters using minibatch stochastic gradient descent. We train without explicit regularization and use the Adam optimizer (Kingma & Ba, 2014). The training runs take 1 to 5 TPU-hous on TPUs provided by Google Research. Models are trained from the same initialization and on the data vectors within each batch (but for different sets of tasks and task orderings).
Models are trained on a single epoch: each of the = 500, 000 batches consists of a new set of sequences with batch size 256. For the LLC estimates, we save 190 checkpoints: 100 are lincarly spaced over the training run, and the remaining 90 are logarithmically spaced. We perform LLC estimation and other analyses on these checkpoints.
48

---

Published in Transactions on Machine Learning Research (07/2025)
Table 4: Summary of hyperparameters and their default values for in-context linear regression transformer
model training experiments.
Hyperparameter
Category
Description/Notes
Default Values
Data
# of training samples
128,000,000
B
Data
Batch size during training
256
T
Data
# of training steps
500k
Ntest
Data
# of eval samples
2048
Data
Dimensions of linear regression task (Task size)
4
K
Data
Maximum in-context examples
8
2²
Data
Variance of noise in data generation
0.125
L
Model
of layers in the model # of attention heads per layer
2
H
Model
4
dmlp
Model
Size of the hidden layer in MLP
64
dembed
Model
Embedding size Training run seeds Type of optimizer Maximum learning rate
64
seed Optimizer Type
Misc
{0, 1, 2, 3, 4}
Optimizer
Adam
7)
Optimizer
0.003
Awd
Optimizer
Weight Decay
0
$1.2
Optimizer
Belas
(0.9, 0.999)
Scheduler Type
Scheduler
Type of learning rate scheduler
OneCycleLR
Stratogy
Schoduler
Strategy for anncaling the learning rate
Lincar
% start
Scheduler
Percentage of the cycle when learning rate is increasing
0.5
49


Reviews:
[{"review_id": "HeZTw5GPL6", "reviewer": "Reviewer_EQrd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* I think this paper is quite well structured. The concepts necessary to understand the method are laid out in a clear and logical order. It does feel like a lot of information to be presented just in the setup, but this seems necessary as the method relies on and builds upon several existing concepts.\n* I think the Figures, particular Figure 2, are thoughtfully created and do a really good job of conveying the ideas.\n* I appreciate that the paper presents experiments with both real and synthetic data.\n\nWeaknesses:\n\n* My primary concern with this work is the necessarily handwavy aspect of the experimental evidence. The paper's writing often seems to suggest an extremely strong connection between LLC and the network's internal structure, yet the experiments involve looking at curves and showing with shading/color that they roughly coincide. I think it would be more meaningful, for example, to try to design a metric that captures the degree to which LLC succeeds in identifying something meaningful. Say, some unsupervised learning of whether or not a model has yet achieved a particular capability, such as 2-grams, 3-grams, etc.\n* I think the introduction of the concepts of LLC and Bayesian local free energy could be made perhaps a bit more high level, with a more detailed description relegated to the appendix. Alternatively, maybe some sort of indication that the details are not absolutely essential for the rest of the paper. It was sometimes a little frustrating to take the time to carefully read and understand the notation and ideas being presented, only to realize in the next section that I could have simply grasped the general idea that was being conveyed and that would have been enough for this paper's purposes.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies how training can be stratified into distinct segments where the model is updating its internal computational structure in different ways. They propose to identify these different segments by analyzing an existing metric of the loss landscape called the Local Learning Coefficient.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think it would be helpful to shorten/simplify the discussion of the existing concepts like LLC and Bayesian local free energy and move the more detailed description to the Appendix. That space could then be used to present more experiments, which would also make a more compelling case for the claimed experimental findings, particularly in settings that used real world data.\n\nIdeally, these findings and the connections they imply could be made a bit more precise by designing some sort of explicit measure of how well the curves \"align\". But I recognize that this could be difficult and whatever measure is chosen would necessarily be a bit arbitrary anyways. Maybe just something to think about.\n\nOverall, I think this is a good paper and certainly worth publication in TMLR.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "SAcm0WePnK", "reviewer": "Reviewer_Afb3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Overall, the work presents a set of interesting results, and I appreciate that the authors examined two different task settings.\n- The paper provides good context on loss landscape degeneracy and LLC for an unfamiliar audience.\n\nWeaknesses:\n- My main concern is that the relationship between LLC changes and model changes is characterized incoherently throughout the paper. The title and abstract suggest \"degeneracy drives development\" and that they're\"deeply linked\", however later in the paper it was described as \"roughly coincides\". I appreciate that the authors made clear that the interpretation is observational, non-exhaustive, and non-causal, but I do feel like a more in-depth discussion of the nature of this link would be needed to justify the stronger framing. For example, the authors mention being motivated by bayesian inference/singular learning process, it would be nice if these concepts can be more clearly unpacked within the context of the two particular case studies.\n- It would be nice to include some analysis on checkpoints where the actual loss curve shows stage-like transition, which serves as a baseline to compare how LLC changes capture additional meaningful model changes beyond any model changes that correspond well to the actual loss.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work examines the relationship between loss landscape degeneracy, as measured by local learning coefficient (LLC), and changes in transformers' internal circuits and behavior. Specifically, in two tasks (language modeling and in-context linear regression), critical points/plateau moments in the LLC score over training were identified and found to correspond to changes in model learning e.g. acquiring heuristic prediction (e.g., n-gram), ICL performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- I find the discussion on the actual case studies most interesting but these sections were a bit thin compared to the background on LLC. I think it would be nice to condense the background on LLC and move some details about the actual experiments (e.g., train/eval details, the different metrics in Figs 4 and 5) up from the appendix to help interpret the results (or at least, provide more pointers to the relevant appendix sections).\n- In Fig 5B, why does the model achieve better ICL for g=3 (OOD?) compared to g=1 (ID)? Does this ICL metric correspond to accuracy at all?\n- It would be nice to discuss if there are more principled ways of discovering changes in model behavior/internal computation that may be reflected in the LLC changes. For the cases studied, some metrics were directly inspired by prior work, what motivated the analysis on the other metrics? Were there metrics that failed to link to any LLC changes?", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PrquwMuDsE", "reviewer": "Reviewer_eimZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- Linking loss-landscape degeneracy to model formation stages is a novel application of singular learning theory to deep transformers.\n- Two contrasting settings (natural language and synthetic regression) both exhibit stagewise shifts, supporting generality across tasks. Multiple seeds, detailed appendices on LLC estimation, and comparisons to Hessian-based measures bolster confidence in the findings.\n\n**Weaknesses**\n\n- Estimating the LLC via SGLD is computationally intensive and demonstrated only on small (≤ 3 M-parameter) models. Only two small, attention-only architectures are studied. It is uncertain whether similar degeneracy-driven stages occur in deeper, hybrid (MLP+attention) transformers or other architectures (e.g., GPT-style models). It therefore remains unclear how this approach scales to production-scale transformers.\n\n- While stage boundaries are empirically aligned with structural changes, the paper lacks a causal or mechanistic theory explaining why degeneracy shifts trigger specific computational structures.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates how training dynamics in transformers reflect transitions between developmental stages. Using the Local Learning Coefficient (LLC), a metric from singular learning theory that measures local loss landscape degeneracy, the authors identify stage boundaries in two small transformer models. These stages correspond to qualitative shifts in attention patterns and task performance, such as the emergence of induction heads or robust function learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- To make a more convincing case of the applicability of their theory, the authors should either include degeneracy tracking results on a larger transformer, provide a complexity analysis for computing LLC (or an approximate version) as model size scales, or make a convincing argument of why scaling up experiments is not necessary.\n\n\n- It would be helpful to provide preliminary analysis or conjectures on the link between LLC drops/spikes and specific mechanistic events (e.g., neuron specialization, attention head formation).", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "9Zt81sedRB", "reviewer": "Reviewer_3d9h", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths*\n\n- The paper is well written \n- Enough context is given on the mathematical tools, such that one does not need to be an expert on SLT to understand the main message\n- I appreciate the intuitive illustrations (Figure 2/3) to convey the idea\n- The experiments validate the theory well\n\n*Weakness*\n\nI list below a mix of remarks and questions that are not weaknesses per se, but could be clarified by the authors.\n\n1) Several (recent) works have studied the phase learning in transformers [1-3]. Although the mathematical/methodological tools of study are different, I believe it would be relevant to discuss them in the current work.\n2) Could the authors provide more details on the benefits of LLC to study the stage learning compared to the techniques used in [1-3] like visual study, loss landscape study, use of optimization tools, etc.?\n3) It would be interesting to see how the method proposed by the authors scales to bigger models. I do not request such experiments, but out of curiosity, I wonder if there would be some computational issues with the SLT theory and LLC estimation. I would appreciate it if the authors could elaborate on that.\n\n*Questions*\n1) Is there any connection between SLT theory (and the LLC) and the sharpness study of the loss landscape? From Figure 2, one might think that a flat loss landscape could lead to lower LLC (e.g., first plot vs second, but several landscapes can indeed be flat but with different LLC, like the third and fourth plots).\n\nOverall, I find the paper interesting, and the idea of applying SLT to the problem is interesting. While the benefits compared to other techniques of studies are not clear to me, I believe the paper is valuable and that the analysis could be used on other problems.\n\n[1] Edelman et al. The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains. NeurIPS 2024\n\n[2] Odonnat et al. Clustering Head: A Visual Case Study of the Training Dynamics in Transformers. arxiv 2024\n\n[3] Chen et al. Sudden drops in the loss: syntax acquisition, Phase transitions, and simplicity bias in mlms. ICLR 2024", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors study the phases of transformer learning during training using Singular Learning Theory (SLT). They posit that the learning phases observed are linked to degeneracy in the loss landscape geometry. To support this claim, they observe the degeneracy along training by computing the local learning coefficient of small transformers in language modeling and context linear regression. The authors identify distinct phases in the learning where changes in degeneracy match significant changes in the internal mechanism at play in transformers\nand input/output behavior.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would appreciate it if the authors could discuss relevant works on phase learning and training dynamics in transformers [1-3].\n\n[1] Edelman et al. The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains. NeurIPS 2024\n\n[2] Odonnat et al. Clustering Head: A Visual Case Study of the Training Dynamics in Transformers. arxiv 2024\n\n[3] Chen et al. Sudden drops in the loss: syntax acquisition, Phase transitions, and simplicity bias in mlms. ICLR 2024", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)