Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficient
Authors: George Wang
Jesse Hoogland
Stan van Wingerden
Zach Furman
Daniel Murfet
Venue: ICLR 2025 Spotlight (2025)
Keywords: Developmental Interpretability
Mechanistic Interpretability
Singular Learning Theory
Learning Dynamics
Stagewise development
Model complexity
Primary area: interpretability and explainable AI
TL;DR: The paper studies how language model attention heads develop during training using new, refined variants of the Local Learning Coefficient (LLC)

Abstract: We introduce refined variants of the Local Learning Coefficient (LLC), a measure of model complexity grounded in singular learning theory, to study the development of internal structure in transformer language models during training. By applying these refined LLCs (rLLCs) to individual components of a two-layer attention-only transformer, we gain novel insights into the progressive differentiation and specialization of attention heads. Our methodology reveals how attention heads differentiate into distinct functional roles over the course of training, analyzes the types of data these heads specialize to process, and discovers a previously unidentified multigram circuit. These findings demonstrate that rLLCs provide a principled, quantitative toolkit for developmental interpretability, which aims to understand models through their evolution across the learning process. This work advances the field of developmental interpretability by providing a mathematically rigorous approach to understanding neural networks through the lens of their learning process. More broadly, this work takes a step towards establishing the correspondence between data distributional structure, geometric properties of the loss landscape, learning dynamics, and emergent computational structures in neural networks.

Published as a conference paper at ICLR 2025
DIFFERENTIATION AND SPECIALIZATION OF ATTEN- TION HEADS VIA THE REFINED LOCAL LEARNING CO- EFFICIENT
George Wang Timaeus
Jesse Hoogland Timaeus
Stan van Wingerden
Zach Furman Timaeus
Timaeus
Daniel Murfet
School Mathematics and Statistics The University of Melbourne
ABSTRACT
We introduce refined variants of the Local Learning Coefficient (LLC), measure model complexity grounded singular learning theory, to study the development of internal structure transformer language models during training. By applying these refined LLCs (rLLCs) to individual components of a two-layer attention- only transformer, gain novel insights into the progressive differentiation and specialization of attention heads. Our methodology reveals how attention heads differentiate into distinct functional roles over the course of training, analyzes the types of data these heads specialize to process, and discovers a previously unidentified multigram circuit. These findings demonstrate that rLLCs provide a principled, quantitative toolkit for developmental interpretability, which aims to understand models through their evolution across the learning process. More broadly, this work takes step towards establishing the correspondence between data distributional structure, geometric properties of the loss landscape, learning dynamics, and emergent computational structures in neural networks.
INTRODUCTION
Structure in the data distribution has long been recognized as central to the development of internal structure in artificial and biological neural networks (Rumelhart et al. 1986 Olshausen & Field 1996. Rogers & McClelland] 2004). Recent observations have renewed interest this topic: language
models progress through distinct stages of development during training, acquiring increasingly sophisticated linguistic and reasoning abilities in ways that seem reflect the structure of the data
distribution (Olsson et al. 2022 Chen et al. 2024 Belrose et al. 2024 Tigges et al. 2024 Edelman
al. 2024 Hoogland et 2024).
deeper understanding of how structure in the data determines internal structure in trained models requires tools that provide information about which components of model are being shaped in response to what structure in the data distribution. Our foundation for the study of such questions begins with the local learning coefficient (LLC: Lau et 2023) from singular learning theory (SLT; Watanabe 2009), which is a measure of model complexity. In this paper, we introduce the refined local learning coefficient (rLLC), which measures the complexity of a component of the model (for example attention head or whole layer) with respect to a given data distribution, which may differ from the training distribution. in transformer trained on the Pile (Gao et al. 2020 Xie et al. 2023 with respect to the distribution which conditions on a token sequence representing code.
example, we can measure the rL..C for particular attention
We focus mainly the rLLCs of individual attention heads and demonstrate the utility of these metrics studying the progressive differentiation and specialization heads. The diversity of attention heads at the end of training has been established in recent years through mechanistic interpretabilityy which has provided numerous examples of attention heads that appear to have specialized functions, including previous-token heads (Voita al. 2019; Clark et al. 2019) and

---

Published as a conference paper at ICLR 2025
50
HIGHER COMPLEXITY (KLLC)
Refined Local Learning Coefficient (rLLC)
LAYER ATTENTION HEADS
MULTIGRAM HEADS
=layer starting
N-GRAN & SKIP GRAN HEADS
40
[A][]] [C]: [A].. [C]
MEMORIZATION
210 Histrral numbers pries
consinccioons
30
112 WiRispoe
50k
HEAKs
20
Ex: (... "...
1:3 Solt cortuoontt
10
INDUCTION HEADS
ALGORITHMIC/
[A] [A] [0]
ABSTRACT
10
104
langaage 116 PunctustinsSymbbel
Training step t
LOWER COMPLEXITY
(ALLC)
Figure 1: The weight-refined local learning coefficient (wrLLC) measures the complexity of model components (such as attention heads) over training. At the end of training, heads with lower wrLLC can be described by simple algorithms (e.g., induction heads, bracket-matching), whereas heads with higher WrLLC memorize n-grams and skip 7-grams ("multigrams"), Shown on the left are the wrLLC curves over training for layer 1 heads, automatically clustered by K-means (clusters are indicated by dominant color, within which individual heads are distinguished by shading). The clusters match the head types (middle-right, classified in Appendix B). Final rL..CC correlates with the number of memorized multigrams for each multigram head (far-right).
induction heads (Olsson et al. 2022) among other kinds (Wang al. 2023 Gould al. 2024). In this paper, we are interested in understanding how this diversity emerges over the course of training, in way that may eventually replace the need for detailed head-by- wead-by-head head mechanistic analysis.
Our methodology yields several novel insights. the same two-layer attention-only transformers studied in Hoogland al. 2024
show that:
Weight-refined LLCs reveal how attention heads differentiate across training: although initially the rLLC curves for heads across training are similar, they progressively diverge with distinctive patterns for different types of heads (Figure I and Section 4.1 4.1
Data-refined LLCs reveal how attention heads specialize across training: a leading hypothesis for the driver of specialization is structure in the data, which imprints order into neural network. We show that the data rLLC reflects this specialization, for example by showing that one induction head appears partly specialized to code (Section
These findings are validated using independent, established techniques, such as clustering algorithms and ablations. We then demonstrate how use these refined L.CCs in combination, which results in our third key contribution:
The identification of a novel multigram circuit: refined LLCs reveal evidence of internal structure related to multigram prediction, which we corroborate with other interpretability techniques (Section
The developmental perspective, where we pay attention to the transformer across the entire training process, is crucial to this methodology. The internal structure of mature organisms constrained by the fact that they have to develop from embryos (Barresi & Gilbert 2023). Similarly, the fact neural network has to develop from initialization in way that decreases the
that the structure loss each stage should be source of valuable insight into the trained model, but to date the field of interpretability has not focused on this source of information. We view our results as helping to establish the validity of field of developmental interpretability. which emphasizes the interplay between distributional structure in the data, geometric structure of the loss landscape, structure in learning dynamics, and resulting computational structure in the model.
2

---

Published as a conference paper at ICLR 2025
2 SETUP
Following Elhage et al. 2021 Olsson et al. (2022). and Hoogllnd et al. (2024), we study two-layer attention-only (without MLP layers) transformers (architecture and training details in Appendix F trained on next-token prediction on a subset of the Pile (Gao et al. 2020 Xie et al. 2023). Throughout, we refer to the attention head in layer with head index h by the notation Lh (starting at O). We also contextualize the development of different components of this model by using the same macroscopic stages LMI through LM5 from (Hooglad al. (2024) as the backdrop for figures (see top of Figure These stages correspond to (LM1) learning bigrams, (LM2) learning n-grams and skip n-grams ("multigrams"", (LM3) developing previous tokknn heads, and (L.M) developing induction heads before (LM5) converging.
We denote an input context, a sequence tokens tu. by SK = (t1...., (.....,,k) where K is the context length. We denote by Sek the sub-sequence (41,...,ts) of Sx. Our data Dn is a collection of length-K contexts, {S i from a data distribution indexed by the superseript i.
The empirical loss with respect to Dn is
K-
(u(w; 9)
- n Σ
Σ
log
(I)
where for a probability distribution P over tokens we denote by P[t] the probability of token t, and fu the function from contexts to probability distributions of next tokens computed by the transformer. The corresponding population loss Eww; q) is defined by taking the expectation with respect to the true distribution of contexts. When q is the pretraining distribution, suppress g and write e, (w)).
METHODOLOOGY
The (global) learning coefficient À is the central quantity in singular learning theory (Watanabe] 2009). In this section we review the local learning coefficient (LLC) from Lau et al.] (2023) before defining the refined variants that new to this paper. The LLC at neural network parameter w* denoted A(w"), ). is a positive scalar measuring the degeneracy of the geometry of the population loss near w". The geometry is more degenerate (lower LLLC) if there more ways in which be varied near w such that ((w) remains equal to f(w*).
3.1 LLC IN PRACTICE
In the setting of Section 2 where we have compact parameter space W, model with parameter w of the conditional distribution of outputs given inputs (in our case transformer neural network with weights w parametrizes predictions of next-tokens y given contexts x), and samples D, from a true distribution with associated empirical loss L,, we define the estimated local learning coefficient at a neural network parameter w" to be
\(w") = nB3 E
w|w , [[,(w)] | - en (u=))
where E w/wt is the expectation with respect the Gibbs posterior (Bissiii al. 2016
p(w; w", B,rr exp (1) ? w w0112]
||w
2
The hyperparameters are the sample size n, the inverse temperature / which controls the contribution of the loss, and the localization strength 7 which controls proximity w². For a full explanation of these hyperparameters the reader is referred to Watanabe (2013); Lau et al. (2023); Furrmnn & Lau (2024); Hoogglnd et al. (2024). Further, the expectation is approximated by using cechaasti--gradicoe Langevin dynamics (SGLD: Welling & Teh 2011) which introduces additional hyperparameters such as the step size; see Appendix F2 for the settings used this paper.
Intuitively, the quantity in represents the typical deviation in empirical loss E,, (w) e, (w*) under perturbations away from w" that are likely according to a local tempered posterior distribution. For more theoretical insight into this intuition see Appendix A
3

---

Published as a conference paper at ICLR 2025
3.2 REFINED LLC
The LLC X(w") depends on the parameter space and the true distribution. If we view some directions in parameter space at w" fixed and view the model as function of the remaining directions we obtain the weight-refined LLC. If we instead allow the true distribution to vary we obtain the data-refined LLC. We now explain both in more detail.
Weight- and data-refined LLC (wdrLLC). Let o' be a data distribution, not necessarily the training distribution 9, and let e' and e, be the corresponding population and empirical loss. Given product decomposition W = U x V corresponding to choosing set of weights V belonging to partieular component of the model (with U denoting the rest of the weights), with associated decomposition of the parameter w² = (u",v"), v²), we let B be neighhorhood of v² small enough that ((u", ,)) 2 E' ('"", v*) for all v E B and define
vol(e, w", VV.') - (G) dv.
(4)
The weight- and rata-refined LLC
A(w" V.d') = lim loga (ooliee, w", V. (/o)/vol((r, /r, w", v.o')
(5)
c-0
The associated estimator X(w"; V, q') is defined by modifying (2 as follows: the expectation E replaced by the expectation with respect Gibbs posterior defined over V by
"//w",
p(u;"" ,.', x exp -n3B,, (u", v) v1
In practice, the estimator is implemented by projecting the SGLD update steps used produce approximate posterior samples onto V and computing both SGLD updates and average posterior loss using samples from o' For further background on the theoretical definition in (5) see Appendix A
When q' = 9. we suppress & and refer to this as the weight-refined LLC (wrLLCC A(w"; V), and when = W, we suppress V and refer to this as the data-refined LLC (drLLC) :(w":g'. When
both of and V = W, we recover the original LLC \(w²).
3.3 LIMITATIONS
There are numerous limitations to LLC estimation in its present form, including:
The justification of the LLC estimator x(w*) presumes that w" is a local minima of the population loss but there is no clear way to ascertain this in practice, and we typically perform LLC estimates during training where this is unlikely.
Accurate estimates can be achieved in some cases where we know the true LLC (Furman & Lau 2024), but in general, the ground truth LLC is unknown. As such, we cannot guarantee the accuracy of estimated LLC values in transformer models, but we do have reason to believe that the ordinality is correct, e.g. that if the wrLLC estimates of two attention heads are in a particular order, then this also true of the underlying true wrLLCs.
We expect SGLD -bweeddddddddd LLC estimation to mature as a technique. In the meantime, a series of papers (Lau al.) 2023 Chen al. 2023] Furman & Lau 2024 Hoogland al. (0224) have demonstrated that despite these limitations, the estimated LLC does in practice seem to offer useful signal for studying neural network development. In the appendix, we compare our analysis using rLLCs against Hessian-based methods (Appendix D) and ablation- based methods (Appendix Some analysis is given for other seeds (Appendix
EMPIRICAL RESULTS
4.1
DIFFERENTIATION VIA WEIGHT LLCS
The weight-refine LLC for an attention head a measure of the amount of information needed to specify configuration of the weights in the head which achieves certain relative improvement in
4

---

Published as a conference paper at ICLR 2025
UMc N/U0GSANG
LiA: CONVERGNEE
LMC
LMC
LMagt: BRAKTION
UN:
5SCAMS
55AM05
LAYER
LAVER
60
Weight-refind LLC
Weight-istteed LLC
30
40
20
20
10
10a
10a
10a
10a
Step
Step
Lenx Hos Previous-Tolen
Mu tigram
Current-laten
FiRA TYPEE
Induction
Multigram
Figure 2: The weight-refined local learning coefficient (wrLLC) reveals how different types of attention heads differentiate during training. The wrLLC curve for each head is shown colored by its functional type (Appendix [B. Remarkably, the partition of the heads by type coincides with the clustering of their wrLLC curves, viewed as time series and clustered by Euclidean K K-means (Appendix [F5)). This suggests that heads which compute differently, develop differently, as revealed by the wrLLC. Throughout this paper, developmental stages LM1-LM5 are colored in the background according to the classification of Hoogland et al. (2024).
the loss (see Section Thus we should expect complexity be the principal axis along which the wrLLC differentiates attention heads. However, a priori it is not obvious how the complexity as measured by the wrLLC should relate to other properties of the attention heads, such as the classification of heads by their functional behavior or the number of multigrams that they memorize. Our first key contribution, contained in FigureTann Figure 2 is to show that there in fact very natural relation between the wrLLC and these other axes of diffcrentiation.
As explained in detail in Appendix B current-token heads) if they strongly and systematically attend to the previous token (resp. current token). Induction heads are identified as in [lhage et al. (202T). All other heads are referred as multigram heads, ablating them tends to highly impact prediction of multigrams. These categories give the rype of an attention head.
we classify attention heads as previous-token heads (resp.
Figure shows that when the attention heads in both layers colored by their type, it is immediately visible that heads of the same type tend to cluster together not only in terms of their wrL.LC at the end of training but also in terms of the overall shape of the WrLLC curve across training. More precisely, curves within a cluster share notable features such as scale, shape, and critical point locations. This visual impression is backed up by the results of various clustering algorithms (Appendix [:35). As further test that the clusters are semantically meaningful, we examine the subdivisions that occur as the number of clusters is increased in Appendix [F5..3]
In the case of the layer 1 heads, it is further the case that heads with higher wrLLC tend to memorize more multigrams (see Figure [ and Figure IT]. and heads with lower wrL.LC tend to described by simple algorithms like induction or bracket matching (see Figure I Thus, the differentiation of attention heads by their WrLLC lines up with an intuitive sense of the description length of the head's computational behavior. For more in-depth and quantitative treatment, see Appendix B.
In summary, the WrLLC reveals differentiation of the attention heads across training that we can independently verify is semantically meaningful.
4.2 SPECIALIZATION VIA DATA-REFINED LLCS
In the previous section saw that the weight-refined LLC reveals the differentiation of heads into functional types, which are useful for prediction on different kinds of patterns (e.g. induction patterns vs. multigrams). We now demonstrate how further refining the LLC measurements by changing the data distribution (such as to one more heavily featuring certain kinds of patterns) provides additional information on model components specializing particular patterns in the data.
5

---

Published as a conference paper at ICLR 2025
40
THE PiLe
Relstive Difeeretoo (N)
50
Anmun - Avna
Apik
1:6
-50
10a
10a
10a
10+
10a
10
Step
Step
Step
LIVER HEAR TYPES
induction Heads
Mutigram Heads
Figure The data-refined local learning coefficient (drLLC) reveals how attention heads specialize to different types of data. The data-refined LLC for GitHub (middle, CodePProo 2023 indicates that on code samples. perturbations to the weights in the multigram heads in layer have significantly less impact the loss than perturbations the induction heads. Informally, the drLLC suggests these heads are differentially more important for predicting code than natural language. This distinction is especially pronounced for (1:6..
we think of the weight-refined LLC A(w"; V) as a measure of the information in V about all patterns in the pre-training distribution, then the simultancous weight- and data-refinement A(w"; V, g') measures the information about the subset of those patterns that occur a subdistribution q'
For example, when head V is specialized multigrams that uncommon in code predict that A(w"; V, (Gast)) ) < A(w"; V) when g' = 9uaHub is a distribution code (CodeParrot) 2023). In the opposite direction, since induction patterns are frequent in code (e.g. repeated syntax, repeated variable names), we expect that X(w"; V. gGatHh)) > A(w"; V) when V is an induction head.
Both of these predictions are borne in Figure 3 Further, we observe that the wrLLC values of the two induction heads are nearly identical at the end of training. but they are pulled apart by the data-refinement to code, with 1166 having a significantly higher value. This leads to the prediction that (1:6) is specialized further, within the set of induction patterns, to patterns common in code.
We verify this prediction in Figure by examining some examples of contexts on which prediction is most negatively affected by ablating 1.6 and 177 and see that, indeed, the former set of examples the latter, which is consistent with the
tend to have a more syntactic or structural flavor relative data-refined LLC. for code of (1:6) having higher value. In Appendix B we cover more examples validating this behavioral difference, including Figure TO which shows that changes in rLLC curves accurately reveal that in-context learning develops at different times for each head.
Just as we can decompose a neural network into its architectural components (e-g. attention heads) we can imagine decomposing the pre-training distribution into its structural components (e.g. dif- ferent data sources). The development process sets up a rich interplay between these two types of components, and the simultancous weight- and data-refines LLCs quantitatively illustrate which components of a model are being shaped in response to what structure in the data. For example, (1.6 may have been significantly influenced by code.
4.3 A NEW MULTIGRAM CIRCUIT
A multigram of length m is common sequence of tokens f1, f2..................... .......,t,. tm in the data distribution, where the tokens may appear contciniigy in context (Shen et al. (2066). In this paper, multigrams are typically of length 3 or 4 and often do involve consecutive tokens. It well-known that transformers can implement the prediction of bigrams (m 2) using the embedding and unembedding layers, and the prediction of skip-trigrams (a subset of m multigrams) using individual attention heads (Elhage et al., (02I). However, the prediction of more complex multigrams may require coordination between attention heads in different layers. In this section, we explain how used refined LLCs to investigate this coordination and provide evidence for a new circuit involved in multigram prediction.
As noted in Hoogland al. (2024) and revisited in Figure 13
two-layer attention-ond -nlly transformers pass through consecutive developmental stages where they behave like zero- and one-layer transform- It is around stage LM3 that the behavior of the two-layer transformer starts diverge from that

---

Published as a conference paper at ICLR 2025
Rise the Witsh-king New Member imc
HANNAHHAREES BOWILER, JAMES WL.
Cocies with nocther Ploced
sumaments
kitchiing
Jaruaty 23, Backkown NSW, Austral
Duckan redly
Ner Dirnet Action & High Treescn Act
Apri 24, 1978 Tarworh, NSW Austrula,
Heutrnght
soded for bach paries tn
cas Conanata i- Soxa Gemenals Discussion
August
Lascablle
Asian Dynostis Disoussions Ago
(51) was born December 26, 1o in Blace
sprees Mouse thats Tunadii onr
Creer Ground
Tectics Tipo
NoE3: The
KASCHBupkember 1974 Tom
Surry at trate Gymssaam 49C
204 Plavel Now Frantrip Kod
rathes pince Seaa novd
undezz cancaled rechesdudd Allooggt
Thurs the mone co Mlss ORC
Donull Excaika les Silto Monan
oms Promiere cast Ovaddo Huso Prove
toppeil caliegee fort soree
they
derie
Thursday
Franterbb Derter Erzskals brother
of and now the plecas duarpe Mayy
Enesing por apric Tor be wene pare
(a) Induction head 116
(b) Induction head 177
[betw MasDoogga 6th Avee
The
the Throe
assimloting U.S. culture
Rubell Wilamm
the group congkiing ng of gold (Pu), (Age N
sanning Kiraa thee cond eacord leter FK
iin suronni
(D11,
palton Loing the parti New Yore" (ptoure Son ne
"hough he was paved convicted.
hour furt
ha.
-hhrrsnng service prooms tatee
sox Avarun, Now NCTICE VEETING
Tauth db", "calcutalo hal",
Moas lnduenia Devsloprant Auticity (Te (u
"pornibouf" (toogh not fo igatodif
(c) Multigram head (Dyck) 0:7
(d) Multigram head (Dyck) 1:5
Figure 4: Induction heads and multigram heads develop subspecializations. Compared to 177. the induction head 116) is more involved in predicting induction patterns (Appendix B.2) that feature punctuation and special characters. Multigram heads 077) and (1:5) learn skip re-grmm involved in bracket-matching (" ("Dyck patterns" Appendix [B.3). Blue indicates the token to be predicted. Orange indicates the strength of the attention pattern at the current token. Samples are selected by filtering
for tokens where ablating the given head leads to the largest increase in loss.
of a one-layer transformer. Thus, if it exists, coordination between heads for complex multigram prediction is likely to emerge in LM3.
The development of such coordination might require "displacing" learned structure for the prediction of simpler multigrams. To test this hypothesis, we investigated how the information in attention heads about simple multigrams changes over training by using data-refined LLCs with one-layer transformer as the generating process for the data distribution (see Appendix F.3). These drLLCs are shown in Figure 3and compared with the wrLLCs for the full pre-training distribution. In line with our expectations, we see that the two sets of LLCs are similar early in training and start to diverge around LM3, which we interpret as relative decrease across all attention heads of the information about simple multigrams that can be predicted by the one-layer model.
If examine the cluster of heads with the largest relative decrease, we see some examples (e.g. the previous-token, current-token and induction heads, classified as such according to their behavior at the end of training) for which the decrease is easily explained: they are involved in predicting simple multigrams at t = 8.5k steps but acquire other roles by the end of training (see Figure 9 and Appendix B.6). However it is a priori surprising find the layer 0 multigram heads in this cluster, since ablation experiments indicate that they are primarily involved in multigram prediction throughout training.
This suggests that it is in these layer 0 multigram heads that simpler multigrams are displaced by the development of coordination for more complex multigram prediction. We use mean ablations to test this theory. In Figure we see that the head (1:5) seems to involve Dyck patterns (closing parentheses and brackets, see Appendix B.3). Curiously, ablating (0:7) at the end of training heavily impacts the same kinds of patterns. However, this is not the case earlier in training: at training step / = 8.5k, 07:) is responsible for a different set of multigrams (Appendix B.6). Combined with the observation from Figure [ this leads to the hypothesis that this layer 0 multigram head may be going through some transition that that causes it to switch to primarily passing information forward to layer multigram heads, including (1:5).
7

---

Published as a conference paper at ICLR 2025
60
THE PILE
LATER OUTPUTS
40
Rolatine Diffornnce
Layer Muftgram
Datereined
20
Apie
-40
Apple
Laver Multiram others
0
10a
10
10a
10*
10a
10*
Step
Step
Step t
Hoa Tives
Layer O Mullgram Heads
Layer Mullggrem Heads
Figure Using a 1-layer (LI) model as the data distribution data-refined LLCs helps locate skip-trigram-relatee structure. During stage LM3, the drLLC begins decreasing for layer 0 multigram heads while increasing for layer heads (middle). In stage LM5, when layer drLLCs also start decreasing, the decline significantly less pronounced than in the layer multigram heads and the rest the heads (right).
10
UNNESTED (...)
("...")
Loss D)]
8
Loss on
6
4-
2
10a
10a
10
10a
10a
10*
Step t
Step
TYPE Loss:
Unablated
Path
1:5
Head 1:5
Figure 6: The formation of the multigram eircuit enables nested bracket-matching, Head 1::5 is a multigram head that specializes to matching brackets. On a synthetic dataset of sentences with parentheses, mean-abe abbaaing this head causes loss increase sharply (blue line to green line). With nested brackets (right), this may require coordination with the layer O multigram heads; mean-ablating the multigram heads and patching their activations just into the input of (1:5) (Appendix (.22) causes an increase in loss (orange), precisely when predict the shift in computational role occurs. With unnested parentheses, 1:5) does not need the layer O multigram heads; same procedure leads to decrease in performance during LM4, but by the end of LM5, the effect is small.
We validate this hypothesis with path patching (Wang et al. 2023) in Figure 6 where we note change in interaction between layer multigram heads and 1:5 the expected time. In Appendix B.6 we present examples of multigrams migrating from layer to layer between & 8.5k and = 50k.
To further corroborate the hypothesis that layer 0 and layer I multigram heads are coordinating, check their K K-composition scores, defined by Elhage et al. (20TT). Informally. these measure how much an attention head in later layer reads from the write subspace of the residual stream of an attention head in earlier layer (for more formal definition, see Appendix [.4). In Figure 7 we that the layer 0 multigram heads all undergo the same turnaround in K-composition scores with the layer multigram heads, including unusually strong coupling of composition scores from given layer head to each of the layer multigram heads. We refer to this pattern of coordination between layer and layer multigram heads as the multigram circuit.
In summary, these results suggest that layer 0 and layer 1 multigram heads independently memorize simple multigrams until around the end of LM3. The layer 0 heads then begin forgetting some simple multigrams as they transition to a supporting role within the multigram circuit. At the same time, the layer multigram heads specialize
complex multigrams, such nested Dyck patterns.
8

---

Published as a conference paper at ICLR 2025
Layer 0 Head
Layer 0 Head
Layer Head
Layer Head
0.15
k-composition
0.08
0.08
0.08
0.06
0.10
0.06
0.06
0.04
0.04
0.04
0.05
0.02
10a
10a
10*
10a
10
10
10a
10
10"
102
10
10
Layer 0 Head 4
Layer 0 Head
Layer 0 Head
Layer Head
0.15
K-composttinn
0.08
0.08
0.10
0.06
0.06
0.1
0.05
0.02
0.02
10a
10a
10*
10a
10a
10
10a
10a
10*
102
10a
10
Step,
Step,
Step,
Step,
Induction
1:6
Induction Head
Multigram Heads
Figure 7: K-composition scores between layer multigram heads and layer multigram heads increase in tandem during stage LM4. After decrease in LM2, K-composition between multigram heads begins increasing during LM4, while composition between multigram and induction heads begins decreasing. Subfigures correspond to heads in layer O. Individual lines show K-compositions over training between that head and layer head. Border color indicates head type: orange for multigram heads, blue for previous-token heads, and green indicates for current-token head.
5 RELATED WORK
Data distributional structure. It is clear that structure in the data distribution plays significant role the kinds of structures learned in neural networks and how they are learned (Rumelhart et al., 1986, Olshausen & Field 1996] Rogers & McClelland, (004)). For instance, properties of the data distribution have been linked to the emergence of in-context learning by Chan et al. (2022b), and Belrose et al. (2024) note that networks learn lower-order moments before higher-order ones.
In the experiments and accompanying theory of Rogers McClelland (2004) p. 103, 169). "waves of differentiation"
triggered by coherent covariation in the data distribution. They claim that the "timing different waves of variation, and the particular groupings of internal representations that result, are governed by high-order patterns of property covariation in the training environment" (Rogers & McClelland, 2004 p.3). Some of these ideas were given a theoretical basis in [axx et al. (2019) which contains mathematical model of Rogers & McClelland (2004).
20044
Specialization. Krizhevsky al. (2012) observed that in the first layer of AlexNet, the two branches specialized to different kinds features. Voss al. (202I) later hypothesized that this phenomenon of branch specialization is driven ultimately by structure in the data. Some of the earliest work in mechanistic interpretability (Cammarata et al., (220) demonstrated interesting specialization in convolutional neural networks. A more automated search for local specialization was pursued in Filan et al. 2021); Hod et al. (2022):Casspr et al. (2022).
6 DISCUSSION
In this paper have introduced the refined LLCs (rLLCs) principled tool for understanding internal structure in neural networks and shown that this tool can used to study the differentiation and specialization of attention heads over training. This builds the theoretical foundations of
(attanaee 2009), the introduction of the ordinary LLC (Lau et al.) (023) and recent results
showing that changes in the LLC over training reflect developmental stages Hoogland et al. (2024).
This section puts these contributions into the broader context the science of deep learning and interpretability. From structural point of view, the problem of interpretability for neural networks is

---

Published as a conference paper at ICLR 2025
to understand internal structure and how it determines the map from inputs to outputs. We take the point of view that this problem cannot be solved in a deep way without first addressing the question: what is the true conception of internal structure in neural networks?
There is a long tradition in mathematics (Langlands) 1970). computer science (Howard et al. 1980) and physics (Maldacena 1999] Greene & Plesser 1996) of understanding the nature of a mathematical object or phenomena by putting it in correspondence or duality with other phenomena. is therefore interesting note the literature (reviewed in Section (5) arguing that data distributional structure is an important factor in shaping internal structure in neural networks, and that this structure is further linked structure in the learning process. To this we may add the singular learning theory perspective, which relates geometric structure of the population loss landscape to the structure of the Chen et al. 2023).
(singular) learning process (Watanabe 2009
The synthesis of these perspectives suggests novel approach to the problem of understanding the fundamental nature of internal structure in neural networks, which to place the problem in the broader context of studying the correspondence between four categories of structure:
Data distributional structure: the inherent patterns and regularities in the data which exist independently of any model (Cristianini & Shawe- Taylor 2004). In this paper: induction patterns, Dyck patterns and multigrams (Appendix B.I).
Geometric structure: the analytic and algebraic geometry of the level sets of the population
loss (Watanabe, 2009 Ammrii 2016). In this paper: the learning coefficient (or real log canonical threshold, as it is known in geometry).
Learning process structure: developmental stages, critical periods, and the sequence in which different capabilities or internal structures emerge (Rogers & McClelland, 2004). In this paper: the overall developmental stages of (Hoogland et al. 2024) and the staggered development of individual attention heads.
Computational structure in the model: the functional organization within the neural network itself and computational motifs that emerge during training. In this paper: attention heads, the induction and multigram circuits.
Here by structure we loosely mean the "arrangement of and relation between parts or elements of something complex" (McKean 2005). Since there can be no structure without differentiation of the whole into parts, the foundation of these correspondences is a relation between the "parts elements" in each of the four categories. From this perspective, the contribution of the present paper is to begin establishing such correspondence for two-layer attention-only transformers by using refined LLCs to quantitatively track which components of the model are being shaped in response to what structure in the data distribution. More precisely:
In Section 4.1
we related computational structure (the behavioral type of attention heads) to learning process structure and geometrie structure measured by the LLC, by showing that the behavioral type can be recognized by clustering wrL.I. curves (Figure and Figure 2
In Section
we related data distributional structure (the difference between the frequency of certain kinds of induction patterns in code versus natural language) to the differences in geometry between particular induction heads (Figure B)
4.2
In Section 4.3 we related data distributional structure (nested brackets in Dyck patterns) to a new multigram circuit whose emergence (a structure in the learning process) seems linked to geometric changes in the layer O multigram heads (Figure 5)
Finally, we highlight that many of these results depend on a developmental perspective. We could not cluster attention heads by their weight-refined L.I.ss without seeing their evolution during training, nor could we clearly see the connection between induction patterns and performance on code samples without observing changes in the data-refined LLC curves. Even the ablation and composition score analyses the multigram circuit depend on comparing results across training.
The techniques pioneered in this paper for understanding internal structure in two-layer attention-only transformers can be applied to models at a larger scale or with different architecture. We refer to this approach, which combines the refined LLCs from singular learning theory with an emphasis on studying networks over the course of development, as developmental interpretability. Using this set of ideas, we hope to open new paths towards systematic understanding of advanced Al systems.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
Detailed descriptions of our experimental setup, including model architecture, training procedures, and hyperparameters, are provided in Appendix
Our LLC estimation procedure is documented in Appendix F.2 which lists the SGLD hyperparameters used for estimating the Local Learning Coefficient and references detailed resources for implementing LLC estimation. Methodological details for implementing the various refinements are provided in
Section and Appendix
Details for the procedure used to classify heads are provided in Appendix B.1 We also provide
link an anonymized repository that contains extended results of this analysis along with additional figures. More implementation details for ablation- based analyses are provided in Appendix E.2 Implementation details for the comparison with Hessian-ban metrics (Appendix D can found in the associated subsections.
CONTRIBUTION STATEMENT
The following is non-exhaustive list of some particular areas of individual contribution.
GW led the project, developed the methodology for LLC hyperparameter calibration, identi- fied the multigram circuit, conducted experiments using other interpretability techniques, contributed to the classification of attention heads, and made substantial contributions to manuscript writing.
JH contributed to engineering, led the classification of attention heads, performed clustering analyses, designed figures, and made substantial contributions to manuscript writing.
SvW led the engineering and conducted the LLC estimation experiments.
ZF conducted the experiments comparing against Hessian-based methods, including design methodology and associated writing and figures.
DM had the original idea for refined local learning coefficients, played role in identifying the multigram circuit, and contributed substantially to writing the main text.
ACKNOWLEDGEMENTS
We thank Matthew Farrugia-Roberts for his feedback on earlier drafts of this manuscript.
REFERENCES
Alessandro Achille and Stefano Soatto. Emergence of Invariance and Disentanglement in Deep Representations. Journal of Machine Learning Research, 19650):1-34, 2018.
Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical Learning Periods in Deep Networks. In International Conference on Learning Representations, 2019.
Madhu S. Advani, Andrew M. Saxe, and Haim Sompolinsky. High-dimensional Dynamics of Generalization Error in Neural Networks. Neural Networks, 132:428 446, 2020.
Shun-ichi Amari. Information Geometry and its Applications. Springer, 2016.
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A Convergence Analysis of Gradient De- scent for Deep Linear Neural Networks. In International Conference on Learning Representations, 2019.
Haim Avron and Sivan Toledo. Randomized algorithms for estimating the trace of an implicit symmetric positive semi-definite matrix. Journal of the ACM (JACM), 58(2):1-34, 2011.
Michael Barresi and Scott Gilbert. Developmental Biology. Oxford University Press, thirteenth edition, 2023.
11

---

Published as a conference paper at ICLR 2025
Nora Belrose, Quintin Pope, Lucia Quirke, Alex Troy Mallen, and Xiaoli Fern. Neural Networks Learn Statistics of Increasing Complexity. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 3382-3409. PMLR, 21-27 Jul 2024.
Leonard Bereska and Stratis Gavves. Mechanistic Interpretability for Al Safety A Review. Transac- tions on Machine Learning Research, 2024. ISSN 1835-8856. URL /tllattap//pperevieees net piet/Eozee Eorum? ?imummmiiiia 1e- Pu6 Survey Certification, Expert Certification.
Donald J Berndt and James Clifford. Using dynamic time warping to find patterns in time series. In Proceedings of the 3rd international conference on knowledge discovery and data mining, pp. 359-370, 1994.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Transformers to Recognize Formal Languages. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (eds.), Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7096-71 16, Online, November 2020. Association for Computational Linguistics. doi: 10.18653/v wet.000emmllmmmmiii
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pp. 2397-2430. PMLR, 2023.
Pier Giovanni Bissiri, Chris C Holmes, and Stephen G Walker. A general framework for updating belief distributions. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 78 10330,0-1130,, 2016.
Tadeusz Caliński and Jerzy Harabasz. A dendrite method for cluster analysis. Communications in Statistics-theory and Methods, 3(1):1-27, 1974.
Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: Circuits. Distill, 2020. doi: 0123391/distilll0002 http://diisilllpu///00oeeeeeee
Stephen Casper. Shlomi Hod, Daniel Filan, Cody Wild, Andrew Critch, and Stuart Russell. Graphical clusterability and local specialization in deep neural networks. In ICLR 2022 Workshop on PAIR*2Struct: Privacy, Accountability, Interpretability, Robustness, Reasoning on Structured Data, 2022.
Lawrence Chan, Adrià Garriga-Alonoo, Nicholas Goldwosky-Dill Ryan Greenblatt, Jenny Nitishin- skaya, Ansh Radhakrishnan, Buck Shlegeris, and Nate Thomas. Causal scrubbing, a method for rigorously testing interpretability hypotheses. Al Alignment Forum, 2022a.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances Neural Information Processing Systems, 35: (188818889911 2022b.
Angelica Chen, Ravid Shwartz-Ziv, Kyunghyun Cho, Matthew L.. 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 Representations, 2024.
Zhongtian Chen, Edmund Lau, Jake Mendel, Susan Wei, and Daniel Murfet. Dynamical versus Bayesian phase transitions in toy model of superposition. Preprint arXiv:2310.06301 [cs.LG], 2023.
Noam Chomsky and Marcel P Schützenbergeer. The algebraic theory of context-free free languages, In Studies in Logic and the Foundations of Mathematics, volume 26, pp. 118-161. Elsevier, 1959.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? An analysis of BERT's s Attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276-286 286. Association for Computational Linguisties, 2019.
12

---

Published as a conference paper at ICLR 2025
CodeParrot. codeparroo/github-code Datasets Hugging Face, October 2023. URL https
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Fran- cisco Guzmán, Edouard Grave, Myle Out, Luke Zettlemoyer. and Veselin Stoyanov. Unsupervised cross-lingual representation learning at scale. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 8440-8451, Association for Computational Linguisties, 2020.
Nello Cristianini and John Shawe-Taylor. Kernel Methods for Pattern Analysis, volume 173. Cam- bridge University Press, 2004.
David L Davies and Donald W Bouldin. A cluster separation measure. IEEE transactions on pattern analysis and machine intelligence, (2):224-227, 1979.
Nicole Dehé. Particle verbs in english. 2002.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp Minima Can Generalize For Deep Nets. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1019-1028. PMLR, 06-11 Aug 2017.
Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How can self-attention networks recognize Dyck-n languages? In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4301-4306. Association for Computational Linguisties, 2020.
Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. Preprint arXiv:2402: 1004 [cs.LG]. 2024.
Bradley Efron and Trevor Hastie. Computer age statistical inference, student edition: algorithms, evidence, and data science, volume 6. Cambridge University Press, 2021.
Armin Eftekhari. Training linear neural networks: Non-local convergence and complexity results. In International Conference on Machine Learning, pp. 2836-2847, PMLR, 2020.
Ronen Eldan and Yuanzhi Li. TinyStories: How small can language models and still speak coherent English? Preprint arXiv:2305.07755 [cs.CL], 2023.
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 Hatield-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.
Daniel Filan, Stephen Casper, Shlomi Hod, Cody Wild, Andrew Critch, and Stuart Russell. Cluster- ability in neural networks. Preprint arXiv:2103.03338 [cs.NE], 2021.
Zach Furman and Edmund Lau. Estimating the Local Learning Coefficient at Scale. Preprint arXiv:2402.03698 [cs.LG], 2024.
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 Pile: an 800GB dataset of diverse text for language modeling. Preprint arXiv:2 20121001000227 [cs.CL) 2020.
Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. Preprint arXiv:1909.12000 2019.
Nicholas Goldowsky- -Dil,, Chris MacLeod, Lucas Sato, and Aryaman Arora. Localizing Model Behavior with Path Patching. Preprint arXiv:2304.05969 [cs], May 2023.
Rhys Gould, Euan Ong, George Ogden, and Arthur Conmy. Successor heads: Recurring, interpretable attention heads in the wild. The Twelfth International Conference on Learning Representations, 2024.
13

---

Published as a conference paper at ICLR 2025
Brian R Greene and M Ronen Plesser. Mirror symmetry in algebraic geometry. Bulletin of the American Mathematical Sociery, 33(1):1-44, 1996.
Peter D Grünwald. The Minimum Description Length Principle. MIT press, 2007.
Michael Hahn. Theoretical Limitations of Self-Attention in Neural Sequence Models. Transactions of the Association for Compulational Linguistics, 8:156-171, 01 2020.
Chris Hamblin, Talia Konkle, and George Alvarez. Pruning for Feature- Preiurr-rreeevving Circuits in CNNs, April 2023.
Michael Hanna, Ollie Liu, and Alexandre Variengien. How does GPT-2 compute greater-than?: Inter- preting mathematical abilities a pre-trained language model. Advances in Neural Information Processing Systems, 36:76033-76061, December 2023.
John Hewitt and Christopher D. Manning. A structural probe for finding syntax in word representa- tions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4129 4138. Association for Computational Linguistics, 2019.
Sepp Hochreiter and Jürrenn Schmidhuber. Flat minima. Neural computation, 91)::144,, 1997.
Shlomi Hod, Daniel Filan, Stephen Casper, Andrew Critch, and Stuart Russell. Quantifying local specialization in deep neural networks. Preprint arXiv:21 10.08058 [cs.LG], 2022.
Hoogland, G. Wang, Farrugia-Roberts, L. Carroll, Wei, and Murfet. The developmental landscape of in-context learning. arXiv preprint arXiv:2402.023664 2024.
William A Howard et al. The formulac-as-typer notion of construction. To HB Curry: essays on combinatory logic, lambda calculus and formalism, 14:479-490, 1980.
Michael F Hutchinson. A stochastic estimator of the trace the influence matrix for laplacian smoothing splines. Communications in Statistics- Simulation and Computation, 18(3):1059-1076, 1989.
Mark H Johnson and Michelle DH de Haan. Developmental Cognitive Neuroscience: An introduction. John Wiley & Sons, 2015.
Eric R Kandel, James H Schwartz, Thomas M Jessell, Steven Siegelbaum, and A James Hudspeth. Principles of Neural Science. McGraw-Hill New York, 5th edition, 2013.
Stefan J Kiebel and Karl Friston. Free energy and dendritic self-organization. Frontiers in systems neuroscience, 5:80, 2011.
Michael Kleinman, Alessandro Achille, and Stefano Soatto. Critical learning periods for multisensory integration in deep networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 24296-24305, 2023.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet Classification with Deep Convolutional Neural Networks. Advances in neural information processing systems, 25, 2012.
Robert P Langlands. Problems in the theory of automorphic forms. Lectures in modern analysis and applications III, pp. 18-61, 1970.
E. Lau, Wei, and D. Murfet. Quantifying degeneracy in singular models via the learning coefficient. Preprint arXiv:2308. 12108, 2023.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. Preprint arXiv:2012.00 22212.0990 2,09999, 2020.
14

---

Published as a conference paper at ICLR 2025
Tom Lieberum, Matthew Rahtz, János Kramár, Neel Nanda, Geoffrey Irving, Rohin Shah, and Vladimir Mikulik. Does Circuit Analysis Interpretability Scale? Evidence from Multiple Choice Capabilities in Chinchilla, July 2023.
J Macqueen. Some methods for classification and analysis of multivariate observations. In Pro- ceedings of 5-th Berkeley Symposium on Mathematical Statistics and ProbabiliyyUniivessity of California Press, 1967.
Juan Maldacena. The large N limit of superconformal field theories and supergravity. International journal of theoretical physics, 38(4): 11-11133, 1999.
James Martens, New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21 21(116):1-76, 2020.
Massive Text Embedding Benchmark. mteb/raw arxiv Datasets at Hugging Face, July 2024. URL
http:///i..faaac..ccaaaaaaeeeeee wttp://////...o..totto https/////////////m rxiv
Erin McKean (ed.). New Oxford American Dictionary. Oxford University Press, New York, 2 edition, May 2005. ISBN 978-0195170777 70771. Second Edition. Includes nearly 3,000 new words, senses, and phrases. Large format, 83/2" by 11". Accompanied by a CD-ROM with full text for Palm OS devices. Used as the built-in dictionary for Mac OS X, Amazon Kindle, and later published electronically at Oxford/AmericcnDeetionayy.oooom in 2006 and as part of Oxford Dictionaries Online in 2010.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. Preprint (Tti::6609.78880 [cs.CL], 2016.
Richard Meyes, Constantin Waubert de Puiseau, Andres Posada-Moreno, and Tobias Meisen. Under the Hood of Neural Networks: Characterizing Learned Representations by Functional Neuron Populations and Network Ablations. Preprint arxiv:2004.01255 [cs, q-bio], May 2020.
Ari S. Morcos, David G.T. Barrett, Neil C. Rabinowitz, and Matthew Botvinick. On the importance of single directions for generalization. In International Conference on Learning Representations, 2018. URL //tp::: http://perreeeeeeeeeeeeeeeeeeeee
wwt. fket/forrumiddee Luu
Neel Nanda and Joseph Bloom. Transformerl.ens, 2022. URL
com/
neel wwwaanaanoorrnnnaooeeeeeeeeeeeee
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
Bruno A Olshausen and David J Field. Emergence of simple-cell receptive field properties by learning sparse code for natural images. Nature, (6583))) 1996.
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 Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context Learning and Induction Heads. Transformer Circuits Thread. 2022.
Ioannis Paparrizos. Fast, Scalable, and Accurate Algorithms for Time-Series Analysis. PhD thesis, Columbia University, USA, 2018
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization beyond overfitting small algorithmic datasets. Preprint arXiv:2211.0059) [cs.LG], 2022.
Paul C Quinn and Mark H Johnson. The emergence of perceptual category representations in young infants: A connectionist analysis. Journal of experimental child psychology, 66(2):236-663 1997.
15

---

Published as a conference paper at ICLR 2025
T. Rauker, A. Ho, S. Casper, and D. Hadfield-Menell. Toward Transparent AI: A Survey Interpreting the Inner Structures of Deep Neural Networks. In 2023 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML), pp. 464 483, Los Alamitos, CA, USA, 2023. IEEE Computer Society.
Timothy T Rogers and James L McClelland. Semantic Cognition: A Parallel Distributed Processing Approach. MIT Press, 2004.
Peter J Rousseeuw. Silhouettes: graphical aid the interpretation and validation of cluster analysis. Journal of computational and applied mathematics, 20:53-65. 65, 1987.
David E Rumelhart, James L McClelland, PDP Research Group, et al. Parallel distributed processing, volume 1: Explorations in the microstructure of cognition: Foundations. The MIT press, 1986.
Jill Sakai. How synaptic pruning shapes neural wiring during development and, possibly, in disease. Proceedings of the National Academy of Sciences, 119629):1096-11699, 2020.
Hiroaki Sakoe and Seibi Chiba. Dynamic programming algorithm optimization for spoken word recognition. IEEE transactions on acoustics, speech, and signal processing, 26(1):43 49, 1978.
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); 11537-11546, 2019.
Dou Shen, Jian-Tao Sun. Qiang Yang, and Zheng Chen. Text classification improved through multigram models. In Proceedings of the 15th ACM international conference on Information and knowledge management, pp. 672-681, 2006.
Ravid Shwartz-Z4 Ziv and Naftali Tishby. Opening the black box of deep neural networks via information.
Preprint 1rXiv:1703.000 703.00810, 2017.
Romain Tavenard, Johann Faouzi, Gilles Vandewiele, Felix Divo, Guillaume Androz, Chester Holtz, Marie Payne, Roman Yurchak, Mare Rußwurm, Kushal Kolar, and Eli Woods. Tsleam, machine learning toolkit for time series data. Journal of Machine Learning Research, 21(118):1-6, 2020.
URL http:////ii...oo/aaass//////////
091 .html
Curt Tigges, Michael Hanna, Qinan Yu, and Stella Biderman. LLM circuit analyses consistent across training and scale. Preprint arXiv:2407. 10827, 2024.
Shashanka Ubaru and Yousef Saad. Fast methods for estimating the numerical rank of large matrices.
In International Conference on Machine Learning, pp. 468-4767777 PMLR, 2016.
Stan van Wingerden, Jesse Hoogland, George Wang, and William Zhou. Devinterp. https: 2024.
//qithub.
Roberto Viviani and Manfred Spitzer. Developmental pruning of synapses and category learning. In ESANN, pp. 287-294, 2003.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing muulti-hedd self-attention: Specialized heads do the heavy lifting. the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5797-5808. Association for Computational Linguistics, 2019.
Chelsea Voss, Gabriel Goh, Nick Cammarata, Michael Petrov, Ludwig Schubert, and Chris Olah. Branch specialization. Distill, 2021. doi: 10.239 //diistll.00024....
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023.
Joe H Ward Jr. Hierarchical grouping to optimize objective function. Journal of the American statistical association, 58(30) (:2301):336244 1963.
Sumio Watanabe. Algebraic Geometry and Statistical Learning Theory. Cambridge University Press,
USA, 2009.
16

---

Published as a conference paper at ICLR 2025
Sumio Watanabe. A Widely Applicable Bayesian Information Criterion. Journal of Machine Learning Research, 14:867-892 2013.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pp. 11080-11090, PMLR, 2021.
M. Welling and Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamies. In Proceedings of the 28th International Conference on Machine Learning, 2011.
Kaiyue Wen, Yuchen Li, Bingbin Liu, and Andrej Risteski. Transformers uninterpretable with myopic methods: a case study with bounded dyck grammars. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 38723-38788 --7766...6... Curran Associates, Inc., 2023.
Guillaume Wenzek, Maric-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Edouard Grave. CCNet: Extracting high quality monolingual datasets from web crawl data. In Nicoletta Calzolari, Frédéric Béchet, Philippe Blache, Khalid Choukri, Christopher Cieri, Thierry Declerek, Sara Goggi, Hitoshi Isahara, Bente Maegaard, Joseph Mariani, Hélène Mazo, Asuncion Moreno, Jan Odijk, and Stelios Piperidis (eds.). Proceedings of the Twelfth Language Resources and Evaluation Conference, pp. 4003-4012, Marseille, France, May 2020. European Language Resources Association. ISBN 979-10-95546-34-4.
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. Preprint arXiv:2302.03169 [cs.CL], 2023.
Bolei Zhou, Yiyou Sun. David Bau, and Antonio Torralba. Revisiting the Importance of Individual Units in CNNs via Ablation. Preprint arXiv: I 02690 [cs]. June 2018.
17

---

Published as a conference paper at ICLR 2025
APPENDIX
Appendix provides theoretical background on the local learning coefficient.
Appendix B
provides further details the head classification. We describe the methodology followed to manually classified each head, and offer more explanation and examples of each head's classification and specializations.
Appendix C discusses the significance of critical points, increases, and decreases in the (r)L.IC in relation stagewise development artificial and biological neural networks. We provide additional results for the full-weights data-refined LLC for variety of common datasets.
Appendix D compares the rLLC the Hessian trace, Fisher Information Matrix (FIM) trace, and Hessian rank. We show that the LLC consistently outperforms these other techniques.
Appendix E compares the rLLC against ablation- based metrics that are common mechanistic interpretability (zero ablations, mean ablations, and resample ablations). We discuss the strengths and weaknesses of each of these methods in relation to the rL..C.
Appendix provides more experimental details on the architecture, training setup, LLC hyper- parameters, expected cost scaling the methodology, using models as the generating process for data-refined LLCs, composition scores, and automated clustering analysis.
Appendix exxminne the consisteney of our findings across different random initializations, support- ing the generality of our conclusions. This analysis further supports the robustness of our observations across various model components.
Appendix H examinne consistency across models scales, rerunning the analysis of the induction circuit for Pythia-70m (Biderman al. 2023).
Additional figures & data can be found at the following anonymized repository: 2024
Refined
Github (dala)-refined LLC
(model)-refined
60-
60-
Mallgran es
Heads
Pros Teano3
Laywr
40
40
Teanno
Cun Take
10a
10a
10a
10'
40
40
40
Layer Heads
30
Layer
10-
10-
10
107
10'
10a
107
10*
107
10a
10"
Figure 8: Detailed reference for hea- rLLCs. This figure combines data from figures Figure 2 Figure and Figure 5 for detailed cross-referencing.
18

---

Published as a conference paper at ICLR 2025
A LLC IN THEORY
What not necessarily clear from (2) is that this an estimator for a theoretical quantity A(w=) which is an invariant of the geometry of the population loss. To explain we recall one form of the definition of the learning coefficient from Watanabe (2009). For a triple (p, 9, sp) consisting of parameter space W with model plyla, w), truth g(y(z) and prior yp on W consider the volume
vol(c) [
(ow)dd
K(w)<e
where K(w)= = D (a(yl=)llp(( (x)d. Under some conditions rem 7.1) the learning coefficient given by
(Watanabe 2009 Theo-
A lim loga c-0
vol vol()=)/vol((e]
(7)
The local learning coefficient X(w") at a local minima w" of ((w) is defined by restricting the volume integral neighborhood B of w" where K(w) 2 K (ww))
vol(e, (0)
p(w)
(8)
K(w)))
and then defining (Lau al. 2023 Furman Lau
2024
X(w*) lim log2 vol( w w²) //volee,, w*)
c-+
This is the asymptotic number of bits necessary to specify parameter near w" which half again closer to the truth. This has some relation to intuitive notions of "flatness" (Hochreiter & Schmidhuber 1997) or description length (Grünwald, 2007).
In practice we do not have access to the function K since it depends the true distribution. Nonetheless there are several methods available to estimate these quantities empirically (Watanabe) 2013; 2009), using the negative log-likelihood of set of samples D,. In practice we substitute the population loss for the KI. divergence and use (2) to approximate the LLC, rather than
try directly approximate (9) Nonetheless this formula offers valuable ceermaaton-theorefiee interpretation of the LLC that we employ in this paper.
B CLASSIFICATION OF ATTENTION HEADS
In this section, we classify attention heads by their behavior. At high level, the attention heads can be thought of being part of one of four groups, illustrated in Table Unless specified otherwise. all description of attention heads refers to behavior at the end of training.
B. METHODOLOGY: TOKENS IN CONTEXT
Identifying tokens in context. To identify which patterns in data each attention head associated to (independently of rLLCs), mean-ablate that head (Appendix (B). then filter subset of 100k samples from the training dataset for the 1K most affected "tokens in context" (i.e., a sample index combined with position index), measured by the increase in the per-token loss. This results in pairs of token in context and local attention pattern for the ablated attention head that generated that sample.
Classifying tokens in context. We attempt to classify each (token in context, attention pattern) pair as belonging to one of the following "patterns:"
Induction patterns (Appendix B.2
Dyck patterns (Appendix B.3
Skip n-grams (Appendix B.4
n-grams (Appendix B.5
19

---

Published as a conference paper at ICLR 2025
Head
Classification
Comments
Section
0
Multigram Previous-toke
R-grams esp. prepositional phrases
B.5
1
Esp. Proper nouns
Multigram
n-grams; Post-parenthetical tokens n-grams
B3.
0
Multigram
B
4
Previous-token
Esp. dates and multiple spaces
0 5
Current-token
B.6
Multigram Multigram (Dyck)
Tt- R-grams; Post-quotation tokens
0
(Nested) bracket- matching; Double spaces / start of sentences
B.3
Multigram
R-grams, esp. numbers & prices (thousands place), hyphenated phrases ("first-ever", "year-ago"), "[#] [timespan] ago"
B.5
Multigram
n-grams, esp. proper nouns; double spaces, com- parison completion (more...th -> an)
B.5
Multigram Multigram (Dyck)
n-grams, esp. dates & start of sentence
B.5
Correlative conjunctions ("neither" "..oorr""""""""" Abbreviations ("N.A.S.A."), Superlatives ("most" ... "ever"), Questions ("Why"..."" "..".
B.3
"...' r-grams, esp. news-related (date ranges, "As- sociated press", "copyright", "spokeswoman," "); phrasal verbs ("prevent..fromm" "keep...safe") post-quotation tokens (" said").
Multigram
B.5
1
5
Multigram (Dyck)
[...], (...), (... .......)) [..), "...",
B.3
6
Induction (Code)
[A] [B] ... [A]
[B]
Induction
[A] [B] [A] -> [B]
Table 1: Attention Head Taxonomy
This classification is automated and serial: if token in context cannot be classified induction pattern, we subsequently check can be described as Dyck pattern, then a skip n n-gram, then an n-gram. The criterion for inclusion varies for each pattern (and is described in the associated subsection) but typically consists of checking whether the token receiving maximum attention ("max- attention token"), current token, and next token match particular template. If pattern matches, we say that pattern "explains" the token in context.
Classifying heads by tokens in context. We classify each head by whichever pattern explains the most of its tokens in context, in relative sense, rather than an absolute sense (so 30% Dyck patterns, 10% skip n-grams, 20% n-grams, 40% unexplained is enough to classify head as Dyck head). This has obvious limitations: there could be another missing pattern that explains the remaining unexplained patterns. We therefore supplement this classification with manual inspection of both explained and unexplained tokens in context (see Figure 9.
We begin by identifying previous-token and induction heads (0:1), (0:4, (1:6), and (177), which we then subdivide based on their previous-toke and induction scores. The remaining heads are classified as multigram heads that memorize Dyck patterns, skip n-grams, and contiguous n-grams. Notably, by the end training, all heads have almost all of their tokens in context explained by these classifications (in some cases only after additional manual inspection). The results of this classification process are displayed Table T anonymized repository:
For full analyses and datasets tokens in context, we refer readers to our http://wwwommoooooooo//////////p
Counting multigrams. To count the "# Multigrams" in Figure and Figure 11 unique number of Dyck patterns, skip n-grams, and n-grams associated to each head. These unique counts defined in their respective sections.
add up the
20

---

Published as a conference paper at ICLR 2025
Pm
P-M
(a) 8.5k steps
100%
Prosent
Wolioee
(b) 50k steps (Final Checkpoint)
Figure 9: Token-in-context attributions for each head at 8,5k steps (top row) and at 50k steps (bottom row). Each column corresponds to single head: the bars (left axis) show the percentage of tokens in context that are explained by either induction patterns, Dyck patterns, skip n-grams, or n-grams (as described in Appendix B.T]. The heads are ordered by the number of unique multigrams for each head.
21

---

Published as a conference paper at ICLR 2025
B.2 INDUCTION PATTERNS
An induction pattern is a sequence of tokens that repeats within given context, where the model learns to predict the continuation the repeated sequence. The simplest form of an induction pattern is an in-context bigram, [A] [B]
[A] [B], where [A] and [B] are arbitrary placeholder tokens. In this pattern, after seeing the sequence [A] [B] once, the model learns to
predict [B] when it encounters [A] again later in the context. In our analysis, we extend the
definition of induction patterns to include arbitrary in-context n-grams the form ([1]... (NJ)
([1) [N-1]] + [N].
B.2.1 CLASSIFICATION
Olsson al. (2022) showed that two-layer attention-only transformers develop induction circuit which completes such patterns, predicting the second [N] from the second [N-1]. This circuit involves two components:
A previous-token head in layer O, which attends to the second [N-1].
2. An induction head in layer 1, which attends to the first [N].
We classify token in context as part of an induction pattern if it fits this form: if (1) the next token is the second [B]) and (2) the max-attention token is the second [A] (previous-token head) or the first. [B] (induction head).
B.2.2 RESULTS
Hoogland al. (2024) identified heads 0:1 and 00:4 previous-toker heads and 1:6 and 1:7 as
induction heads, using the previous-token score and induction score introduced in Olsson et al. (022).
These heads can also be identified such by their tokens in context, illustrated by the selection in Figure 4 and the automated classification in Figure
The rLLC reveals staggered development of induction heads. The Github drLLC reveals only that the heads specialize to different data, but also that they form different times: Figure 10 shows that the inflection point and peak in the rL.LC several thousand steps later for head (16 than head (1:). To validate this, we consider the in-context learning (ICL) score from Olsson al. 2022 which takes the average loss at the 500th token minus the average loss at the 50th token.
Bonus:
1:7 Ablated
ICL Score
0.0
myshe
1:6 WILLC
20
-0.05
-0.1
Full (Unablated)
1:6 + 1:7 WILLC
ICL Score
00.0
30
0.0
1:6 Ablated
ICL Score
1:7 WILLC
-0.15
20
20
-0.1
-0.20
MMJ
10
M
102
10a
10c
102
10a
10a
Step
Step
ICL Score
Pile Github
Figure 10: Induction head development correlates with critical points in the Github drLLC.
The two induction heads not form at the same time; 1:7 precedes (1:6) by several thousand steps, as is visible in the inflection points of both the Github drL.I. for these heads and the per-head ICL. score (obtained by mean-ablating the opposite induction head, left column). When restricting to both heads' weights at the same time, the critical point in the Github drLLC coincides with the end of the ICL score drop (right column). This suggests (critical points in) rLLCs can be used to study the development of model components.
22

---

Published as a conference paper at ICLR 2025
We plot the ICL score upon mean ablating either of the two heads: see that the drop in the ICL score upon mean ablating head 117 coincides with the peak in the 1:6 Github drLLC, and vereve Likewise, the peak in the rLLC restricted to both heads' weights at the same time coincides with the standard ICI. score.
B.3 DYCK PATTERNS
In predicting the next token in natural language, some of the most explicit occurrences of hierarchy come in the form of nested brackets and punctuation. This is formalized in the context-free gram- mars Dyck-k ("deek"), which involve sequences with matching brackets kinds (Chomsky & Schützenberger 1959).
take a wide definition of bracket, including:
Delimiters:
Traditional brackets/rraces/paretttes (...), [...],
Quotation marks: " '...', as well as mistokenized
Markdown formatting symbols:
Split constructions:
Correlative conjunctions: "both...and", "not only...but also", "neither...nor, "ei-
ther...or"..
Correlative Comparatives: "more...th,, "less...th", "better..th",
Correlative Superlatives: "most..ever", "least...ever",
Questions: "Why...?", "What..??", "How...?..
There are similar constructions that we did not analyze but that would make natural candidates for follow-up analysis, such additional correlative comparatives ("as ... as") conditional statements ('if... then"), cleft sentences ("it is who"), comparative correlatives ("the [more you practice], the [better you get]"",, result clauses ("[it was] so [cold] that [the lake froze]"),
The number of unique Dyck patterns for each head is the size of the set of unique (opening token, closing token) pairs for each token in context that is classified as a Dyck pattern.
B.3.1 CLASSIFICATION
We classify token in context as part of Dyck pattern if (1) the next token contains a closing bracket corresponding to an earlier opening bracket, and (2) the subsequence spanned between those two brackets has valid nesting.
Often, we find that the max-attention token is the corresponding opening bracket, especially for layer 1 Dyck heads. However, we do not require this to be true to classify a token in context as a Dyck pattern.
B.3.2 RESULTS
Figure shows that automated classification identifies three heads as Dyck heads: 1:5) (98% ex- plained), 11:3 (43.2%), and (0:7) (23.4%). Manual inspection of their tokens in context confirms these diagnoses and reveals further subspecializations.
1:5 Delimiter-matching head. Head 11:5 specialized to traditional brackess/braccsppaeenttess quotation marks, and other symbolic delimiters.
1:3 Split-construction--ateeeee matchinn head. Head 1:3 is specialized to the natural language Dyck patterns, and one variant of quotation marks.
0:7 Dyck-support head. Head 0:7 is specialized similar tokens in context 1:5. However, only about a quarter of (0:7's tokens in context are recognized by our automatic procedure as Dyck patterns. Manual inspection reveals that many remaining unexplained tokens in context are either
23

---

of (for with mistokenized scare quotes rather than direct quota- Published as a conference paper at ICLR 2025
misclassified as non-Dyck (e.g., because identify whether quotation marks are opening closing by checking whether the previous/subsequent character is not a letter, which is sometimes too restrictive) Dyck-like (c.g., all-caps text surrounded by multiple spaces, where the spaces serving as delimiters). Most of this head's other remaining tokens in context seem to be involved in skip n-grams (analyzed in the next subsection).
One additional way in which these heads are different their max-attention tokens: the layer 1 Dyck heads both primarily attend the opening token corresponding the closing token being predicted. Head 077 casts more diffuse attention pattern.
B.3.3 DISCUSSION
The ability to correctly close brackets underlies all nonregular context-free languages, in the formal sense that by the of Dyck-2 through intersection with regular language and homomorphisms (Chomsky & Schützen berger 1959). For this reason the ability of transformers recognise Dyck languages has been studied at some length (Hahn 2020; Bhattamishra al.) 2020) Ebrahimi et al. 2020). In Weiss et al. (2021) an algorithm in RASP is given which compiles to a transformer which recognises Dyck-k languages for any k. We did not examine how this relates the heads investigated here. It is unclear whether transformers actually learn similar algorithms to these in practice (Wen et al. (002)..
theorem, any context-free language arises from a variant
As far know this paper is the first time that circuit recognizing a nested Dyck language has been found "in the wild", that is, in a transformer trained on natural language. This seems interesting in connection with the ability of transformers learn the hierarchical and recursive structure in natural language. is however well-known that, in general. syntactic structure in natural language is represented within transformers (Hewitt & Manning 2019
B.4 SKIP N-GRAMS
The simplest form of skip n-gram is skip trigram of the form [A] ... [B] -> [C], where the distance between tokens [A] and [B] is variable. We consider more general skip n-grams of
the form [1] ... ([2] ... (N-11) -> [N], as well as skip n-grams involving multiple steps [1] [2]. [N-1] -> [N].
The Dyck patterns considered in the previous setting are special case of skip n-grams.
B.4.1 CLASSIFICATION
We match tokens in context against preset list of skip n-grams, including:
Post-quotations: e.g., "..." said, "..." says.
wrote,
and + tions),
Post-parentheticals: e.g., (@...) October,
Post-correlative comparative: Finishing a comparative of the form "more...th" or predicting what comes after "than."
Abbreviations (Acronyms/Inittaaismmm)) "N.A.S.A.,, "...", "D.C.".ete
Phrasal verbs / verb-particle constructions (with object insertion) like "prevent... from". "keep safe", "let... go"". These were compiled by manually inspecting tokens in context.
The unique number of skip n-grams for each head is the size of the set of unique ([1), [2], [N] tuples for each token in context that is classified as skip n-gram. There are likely quite few other skip n-grams that our search neglected.
Though phrasal verbs fall under the category of (natural language) "split constructions", we choose to analyze these separately from the split constructions listed as Dyck patterns for two reasons. First, Dyck patterns typically have obligatory closing elements (like matching brackets question marks), the particle in phrasal verbs is often optional or can be omitted without rendering the sentence ungrammatical (though correlative comparatives and superlatives sometimes an exception to this
24

---

Published as a conference paper at ICLR 2025
rule). Second, the Dyck patterns are primarily syntactic while phrasal verbs are more lexical in nature, meaning their behavior and interpretation are more closely tied to specific lexical items and idiomatic meanings rather than general syntactic rules (Dehé 2002).
B.4.2 RESULTS
One of the remaining heads appear to be primarily involved in skip n-grams: 1:4 (43.6%). Heads 1:1 (37.5%) and (1:3) (32.6%) are close to being classified as skip-n-gram heads but are more involved in n-grams and Dyck patterns, respectively. few other heads are marginally involved: 0:2) (3.2%), 0:6 (5.7%).
0:2, 0:6,, 1:11, 1:4 Post-Dyek patterns. Heads 0:2, 0:6, and 114 are not themselves Dyck heads. Instead, they involved in predicting tokens that follow closing parentheticals and two variants of closing quotation marks, respectively. Head (111 is responsible for finishing correlative comparisons such "more...th" -> "an". This is not itself a pure Dyck pattern, since the preceding token has already started to close the opening bracket. Head (1:4 specializes in post quotation-marks, such as says...
133
Abbreviation head. In addition to its role in predicting split-construction Dyck patterns, head is also specialized to predicting skip n-grams periods in acronyms and initialisms.
O33
1:4 Verb-particle phrases & other set phrases. Head 1:4 seems specialize in verb-particll phrases and other set phrases, including examples such as prevent from, keep safe, let die, le. go, meet ments), remove from, accused ... of, asked ... whether, and turn ... into. Additionally, this head appears to handle other types of set phrases or common patterns, such as !/... followed by / (for URLs), email followed by @, and at followed by @. This head also involved in recognizing common verb-object pairs such as solve... problems, and violated... laws.
B.5
n-GRAMS
An n-gram is a commonly occurring contiguous sequence of n tokens, e.g., bigram is [A] -> [B], trigram [A] [B] -> [C], and so on.
B.5.1 CLASSIFICATION
For all remaining tokens in context, we classify the sample as an n-gram if the subsequence starting at the token the receives maximum attention up to and including the next token occurs in at least two different tokens context. We enforce no restrictions on the attention threshold or the length of the n-gram.
B.5.2 RESULTS
All multigram heads memorize some number of regramm, but there is clearly an ordering, where Dyck heads memorize the fewest, the skip n t-ggra--ppeciazized heads slightly more, and the remainder (0:0. 0:2, 0:3, 0:6), 1:00 1:4) are almost entirely focused n-grams.
Start-of-sentence heads. Both 1: and 1:2 are responsible for predicting multiple contiguous spaces/newlines and tokens that follow these spaces, contexts where periods are followed by double space.
1.1. 1.2
0:0, 1:4) Prepositional phrases. In addition to learning miscellaneous n-grams, head 0:00 spe- cializes prepositional phrases. Many of its tokens context are prepositions ("of", "by", "with", "to". "from", "with"" that appear in set phrases, like "Chamber of Commerce", "plenty of room", "by email", "went to college", "prevented from doing." 1:4 is also involved in some these. "at". "on <blank> basis", "as time goes".
1:0 Numbers, hyphens, periods of time. Head 11:0 specializes to predicting numbers and prices, especially following the thousands place, hyphenated phrases such "first-ever" and "year-ago",
25

---

Published as a conference paper at ICLR 2025
0:1
Head Type
Final WLLCC Ayalt== 50k)
50
Multigram
Lo Previous-token
L0 Current-token
40
4:0
0:
Multigram
1.4
P.P:
J.2
Induction
D:6
3:3
415
D:7
J1667
( = 0.78. 0.881
0
20
40
60
80
100
120
140
160
# Multigrams
Figure 11: The n-grams memorized by each multigram head correlates with the final weight-refined LLC. By counting the number of unique n-grams associated to each head after the head-atribution procedure outlined in Appendix B.I we find a correlation with the wrLLC
as well as a period of time ending in "ago" ("years ago". "month ago". "day ago"). It seems to be involved more generally in predicting n-grams that common in historical or legal context, e.g.. "defamation", "disengagement", "argues".
1:4 News head. Head [1:4 is especially relevant for predicting n-grams that show up in news articles. This includes words like "Associated Press", "spokeswoman" well as date ranges and post-quotation tokens. It's also involved several shorter skipgrams involving citations ("- <Name> (@ <citekey>) <date>"), attributions ("copyright <Source> image").
B.6 OTHER OBSERVATIONS
0:5 Current-token head. We classify the current-token head 0:5 separately from the above tokens- in-context analysis. This head has two main distinguishing features: it attends almost entirely to the current token, and its composition with the layer multigram heads increases dramatically towards the end of LM4 before decreasing equally dramatically early in LM5.
do not yet fully understand this head's role. From looking at its tokens context, this head appears to be involved in predicting several regrams, and potentially also in predicting Dyck patterns (though its attention concentrated on the current-token and not the opening bracket/quotation mark/etc.).
0:0 Space head. Head 0:0 appears to initially become a "space head." At 8.5k steps, its attention pattern tends distribute itself across all the single-space tokens over the preceding context. This seems to largely revert by the end of training, where it becomes a standard multigram head. This behavior may be linked why the (0:0-LLCC diverges from the other multigram heads during stages LM1-LM3.
Migration of the ago n-gram. Early in training, around 8.5k steps, head 0:0 seems to be responsible for predicting time spans ending in "ago," c.g. "years ago" and "one year ago." By the end of training, head 11:0 takes over this role, even though did not appear have any involvement at the earlier timestep. Head 0:00 retains only the role of predicting ago when it completes a word (as in "Santiago" or "archipelago").
Migration away from multigram heads. Heads 0:1. 0:4, 0:7 ultimately become the two previous- token heads and (layer 0) Dyck head respectively. However, these heads appear to start out as simple n-gram heads. For example, at 8.5k steps 0.7 (Figure 9) initially learns trigrams like "Fitzgerald", "transient". "transactions". and "transformation" (by running the same procedure to detect tokens in context). At the end of training, these trigrams do not feature among any of 0:7's maximally affected tokens in context. Head (07) is also initially involved in predicting the completion of comparative skip-trigrams ("more...thi th" "an"), role that gets overtaken by (:11)
26

---

Published as a conference paper at ICLR 2025
C INTERPRETING CHANGES IN THE (R)LLC
C.I INTERPRETING CRITICAL POINTS IN THE (DATA-REFINEL) LLC
Hoogland al. (2024) studied the (non-refined) LLC as tool for analyzing stagewise development in neural networks. The authors show that critical points in the LLC correspond to boundaries between distinct developmental stages (see Figure
In Figure 12 we show that the full-weights data-refined LLC respects the stage boundaries discovered with the non-refined LLC, for variety of datasets:
Common training evaluation corpora: The Pile (Gao et 2020 Xie et al. 2023
TinyStories (Eldan & Li, 2023), and Wikitext (Merity et al. 2016).
2. Scientific domains for which we create datasets by filtering Arxiv abstracts by category (Massive Text Embedding Benchmark, (2024): math (math.*), physics (astro-ph.* cond-mats
gr-qc, hep-*, math-ph m.lin., nucl-*, physics. *, quant-ph) computer science (cs. +), and economics (econ.*).
3. Human languages for which datasets sampled from the CC-100 (restricted to languages with a Latin alphabet because the 5k vocabulary size has poor support for other languages).
(Conneau al. 2020: Wenzek al., 2020)
4. Programming languages for which datasets are subsampled from Github (CodeParrol 2023
With the exception of the programming languages datasets (discussed in Section the drLLCs resemble vertically shifted copies of the original LLC (evaluated on subset of the Pile). In the scientific domains, the decrease during LM3 is especially pronounced and extends further into LM4 than in the original trajectory. We
surprised to that AMah Àcs
APhysics
AEzon-
There are some slight changes in the locations of the boundaries: the decrease in (r)LLC during LM3 flattens out for many of the datasets (TinyStoores, natural & programming languages), and it shifts slightly earlier (5.5k-8k steps rather than 6.5k-8.5k 5.5k steps) for the other datasets. We believe that these changes are more likely be due changes in the hyperparameters used for LLC estimation (a slight increase in n,3 from 23 to 30) than related to these particular datasets, see Appendix F2
C.2 INTERPRETING INCREASES IN THE (MODEL- [MODEL-)REFINED RREFNND LLC
Many works have observed that the training process of small neural networks can have step-like appearance, with the loss decreasing between plateaus at the same time as a complexity measure
Corpora
Languages
Programiing Languages
102
10a
10*
1.7
10r
10a
1.07
10a
10a
3.7
16'
10"
150
15.0
150
150
1.0
10.0
1.0
0-
0-
102
10F
10a
Stop
Whltot
50e
Figure 12: Data-refined LLC retain coarse developmental stages. Except for the Github drLLC. the data-refined LLC on its own (without additional weight refinement) yields curves that are close to the original unrefined LLC.
27

---

Published as a conference paper at ICLR 2025
1.5
K-L1
2
1.0
Github
0.5
10
10
10a
10a
10a
10*
150
Github
100
M=L1
100
Am
50
10a
104
10a
104
Step t
Step
Figure 13: Two-layer attention-only transformers undergo stagewise development, learning (LMI) bigrams, (LM2) multigrams, and (LM3 & LM4) the induction circuit, before (LM5) converging (Hoogland et al. 2024). This development may be "hidden" from the training loss (top right, blue line), but discovered by looking for plateaus in the (data-refined) LLC (bottom left, blue line). Alternatively, certain datasets highlight particular developmental stages, e.g., Github (left, orange, CodeParrot 2023) accentuates the development of induction. Likewise, evaluating the model against the outputs of a reference zero-layer (right, green) or single-layer (right, red) model highlights the bigram and multigram stages, respectively. Shown in the bottom right the data-refined LLCs using zero and one-layer trained transformers as the data generating process (Appendix F.3
(e.g. rank) increases between plateaus (Arora et al. 2019 Li ct al. 2020 Eftekhari 2020; Advani ce al. 2020: Saxe et al. 2019). Indeed, by viewing traiming as starting at point of low complexity, this monotonic increase in complexity in simple systems has been put forward as an explanation for the generalization performance of neural networks (Gissin al. 2019). Similar explanations for generalization performance of infants have appeared in child psychology (Quinn & Johnson 1997).
In Chen al. 2023); Furman & Lau (2024) this phenomenon was put into the context of the singular learning process. In this vein, Hoogland et al. (2024) and [2 show that, in simple language models, developmental stages appear typically (though not always) involve increases in model complexity (as measured by the LLC), with critical points corresponding to some structure or behavior finishing development.
C.3 INTERPRETING DECREASES THE REFINED LLC
The picture of development of neural networks (biological artificial) in which complexity monotoni- cally increases is incorrect, and thus cannot be a complete explanation for generalization performance. In neuroscience it conflicts with the phenomena of critical periods and synaptic pruning (Kandel et al.) 2013; Viviani Spitzer 2003), and in neural networks the experiments in Hoogland et al. (2024) show that is common for developmental stages involve decreasing complexity (as measured by the LLC and mechanistic analysis). This is entirely consistent with the singular learning process.
Memorization and forgetting. Intuitively, decreases the LLC represent an increase in the rigidity of the network's 'ss computation (at constant loss) with corresponding decrease in the LLC, which we can interpret as either increasing geometric degeneracy or decreasing model complexity.
Some of the dyads used to describe similar phenomena in the literature include:
Memorization and forgetting / reorganization (Achille et al., 2019) in vision networks.
Memorization compression (Chen al. 2024. Appendix G) language model training.
Memorization and foreutt-oomation// cleanup (Nanda et al. 2023) the setting of grokking
Power et al. 2022).
According to the natural information theoretic interpretation of the LLC (Watanabe (000g $7.1, Furman & Lau 2024 see also Appendix A) when this quantity is increasing more bits required
28

---

Published as a conference paper at ICLR 2025
to specify neighbourhood of a low loss parameter: the amount of information in the weights is increasing. This information could include "memorization" of the inputs. The reverse is true when the LLC decreases, this is consistent with "forgetting".
However we prefer to avoid these terms since their precise meaning is unclear. more principled perspective follows from noticing that quantities like the LLC provide upper bounds on both the mutual information between inputs and activations within the network, and the total variation of those activations (Achille & Soatto 2018 Proposition 5.2). is consistent with these results that decreasing LLC could be related to a process whereby the representations "discard" extraneous information about the imput, and become disentangled: it would be interesting to extend the results of (Achille & Soatto 2018) using singular learning theory.
Critical periods. In the development of biological neural networks there are periods where tem- porary "sensory deficits" (interpreted as a shift in data distribution) can lead to permanent skill impairment; these are called critical periods (Kandel et al., 2013). It was argued in Achille al. (2019) that this phenomena is linked to the existence of "memorization phase" in which the infor- mation in the weights increases followed by a "reorganization phase" or "forgetting phase" which the amount of information is reduced without negatively affecting performance; interventions in the learning process during the memorization phase disproportionately affect the final performance on the test set (making these analogous to the critical periods studied in neuroscience).
In Achille et al. (2019) these two phases were related to increasing and decreasing trace of the Fisher Information Matrix (FIM), respectively. In our small attention-only language model attention heads undergo periods of rapid specialization followed by periods of refinement, linked respectively to increasing and decreasing weight rLLC; this closely mirrors the concept of critical periods in developmental neuroscience explored in Achille et al. (2019); Kleinman et al. (2023) for vision networks.
we show increases and decreases in the rLLC are correlated with changes in the Hessian trace (a quantity closely related to the FIM trace) during stage LM1, LM2, and LM5, but that they are less reliably correlated for the intermediate stages. In Appendix D.2 we show similar behavior in the FIM trace though this noisier than the Hessian trace due methodological limitations.
In Appendix D.I
In Chen et al. (022)) it was observed that an increase and decrease complexity (measured by e-g. the intrinsic dimension) coincided with an apparent phase transition in the model; this was related to the model the training process explained in (Shwartz-Ziv & Tishby (2017).
Pruning in neuroscience. In developmental neuroscience it is understood that after an initial phase of proliferation during childhood, synapses (the points where signals are exchanged between neurons) are then pruned, in process that accelerates during adolescence and stabilises in adulthood (Johnson] & de Haan] 2015 $4.3). The apparent plasticity of the young brain is often attributed to the initial overproduction of synapses and synaptic pruning is closely related to critical periods.
There are several aspects of pruning that are particularly relevant to the present paper:
Synapses are kept or pruned based on activity (Johnson & de Haan 2015) $4.3, for a survey see (Sakai 2020). This data dependence of pruning has led some claim that "to learn is to eliminate". Functional structure in the brain is constructed through growth and elimination, with both being required.
Some have argued that the differentiation of the brain into separate processing streams and specialized regions is related synaptic pruning (Johnson & de Haan 2015] $12.4). A famous example is the emergence of ocular dominance columns (and thus the separation of input from the eyes to facilitate binocular vision) and more general separation of modalities.
Different regions of the brain have different timings for the reduction of synaptic density. For example the primary sensory areas of cortex show faster growth and decay curves than the prefrontal cortex (Johnson & de Haan 2015 $4.3).
It has been suggested that pruning is the biological basis of online Bayesian model selection Kiebell & Friston 2011). This is particularly interesting in connection with Watanabe's
29

---

Published as a conference paper at ICLR 2025
discovery (Watanabe 2009 $7.6, also Chen et al. 2023) that in singular models, internal model selection happens automatically as a consequence of free energy minimization.
D COMPARISON AGAINST HESSIAN-BASED ANALYSIS
In addition to the LLC, we explore several other geometric measures based on the Hessian of the loss, and evaluate their ability capture structural development and differentiation over training time. The Hessian, which represents the second-order derivatives of the loss function with respect to model parameters, has been widely used in machine learning analyze model complexity, optimization dynamics, and generalization properties (LeCun et al. 1989; Dinh et al. 2017).
is worth noting that the LLC has stronger theoretical foundations as a model complexity measure compared to Ressian-basee methods. Unlike Hessian-basee methods, the LLC captures higher-orde order information about the loss landscape, and not susceptible nigh-frequency noise the empirical loss.
In this section, we investigate three Hessian-based measures:
Hessian trace (Appendix D.I)
Fisher information matrix (FIM) trace (Appendix D.2)
Hessian rank (Appendix D.3)
The Hessian trace shows the greatest empirical success, identifying similar developmental stages as the LLC, albeit with some differences in interpretation. The FIM trace largely agrees with the Hessian trace but introduces additional noise due to estimation methodology. The Hessian rank, while theoretically appealing, proves challenging to estimate reliably and does not clearly reflect the structural development observed through other methods.
Overall, these fessian-basee measures offer complementary perspectives to the LLC analysis, both independently supporting the findings of the LLC, while also highlighting the advantages and disadvantages of the LLC in comparison. The following subsections detail our methodology, results, and discussions for each of these measures.
D.1 HESSIAN TRACE
We measure the trace of the Hessian of the loss function over the course of training. We measure this for the entire model on the original training distributio,, as well as for subsets of the parameters and for different data distributions, similar what was done for the LLC in Section
The Hessian trace results (Figure (4) show both similarities and differences compared to rLLC findings. It detects some stage transitions and, when weight-refined, differentiates attention head clusters. While the Hessian trace provides valuable insights into model development, its interpretation can be less straightforward than rLLCs some cases.
D.1.1 METHODOLOGG
For large models, the Hessian trace can be too expensive to compute explicitly; we use the Hutchinson trace estimator (Hutchinson) 1989 Avron & Toledo 2011) to efficiently estimate the Hessian trace using Hessian-vector products. This has computational cost comparable that the L.C (a constant multiple higher than that of single SGD step).
The Hutchinson trace estimator relies on the fact that
Tr(H) E(v
(10)
for any matrix H and with entries drawn IID from the Rademacher distribution (Avron & Toledo 2011). When H is the Hessian matrix, we may efficiently compute the product H using standard automatic differentiation techniques, even in cases where explicit computation of H itself is pro- hibitively expensive. The two hyperparameters of this method are the number of dataset samples to in calculating the Hessian, and the number of samples used compute the expectation in Eg 10
30

---

Published as a conference paper at ICLR 2025
Full Model
Layer
107
101
104
102
1.0
10
10
10P
400
2002
300
Teffural
150z
1k-
200
100
100k
50k-
500
Step
Hend
Hed?
Figure 14: Estimated Hessian traces for Pile-13m (top row) and for Github (bottom row). Compare with rLLCS in Figure
For all the results in the paper, we used value of 100 for the former and 50 for the latter. However, we have found little dependence of the results on the values of these hyperparameters, beyond the fact that lower values lead to more noisy estimates. This is in contrast to the LLC, where we have found that hyperparameter tuning is important for the final results.
We also compute the Hessian trace for subsets of parameters, and for different datasets, in a similar manner to the LLC. Weight refinement is performed by considering the loss function of only the selected parameters rather than all parameters, and taking the Hessian of that function. Data refinement is performed by using different dataset for computing the loss. We measure the same parameter subsets and datasets as for the LLC (Section))
D.1.2 RESULTS
The Hessian trace results are plotted in Figure 14 which we compare the rLLC results in Figure We immediately see both similarities and differences between these curves and the rLLCs, both supporting the validity of the information provided by the rLLCs while also providing tentative evidence that the rLLCs do not merely recapitulate the information contained the Hessian trace.
The unrefined, full-model Hessian trace does appear detect the transition between LMI and LM2 and between LM3 and LM4, but fails to distinguish between LM2 and LM3, Additionally, the Hessian and the LLC disagree on what happens after LM4 from the LLC's perspective, the complexity of the model mostly stops increasing, whereas the Hessian trace continues to rise.
Zooming into specific components by weight-refining the Hessian shows number of connections to the structural changes that were investigated earlier. The first notable similarity of the Hessian trace its ability to (partially) distinguish the differentiation that different clusters of attention heads go through.
In layer O, we see that the previous token heads, 0:1 and 0:4, extend their middle-of-training plateau region up until their composition scores with the induction heads are maximized. We also the layer 0 multigram heads sharing a peak in value around the end of LM2 or LM3, which approximately matches the timing of critical point in their K-composition scores (cf. Figures 5p. Finally, 0:5) is visually separated from the rest of the heads just by its magnitude, and the end of its plateau region approximately matches the time when it has its spike in K-composition scores (cf. Figure 7).
31

---

Published as a conference paper at ICLR 2025
In layer 1, the clustering seems more ambiguous. Most of the multigram heads seem to have similar shape, but it also seems plausible to include one of the induction heads, (166, with the rest of the layer 1 multigram heads. The other induction head, 1:7, is differentiated from other heads by a later peak.
Overall, we do see non-trival developmental signals from the weight-refinet Hessian traces. However, there are also significant drawbacks, such the apparent instability. As result, it is clear where the critical points are for wrLLCs, but this is often not the case for the Hessian trace. For example, heads like (011, (0:4) and (0:5) appear have noisy plateaus as opposed obvious peaks. One might be tempted to adopt a heuristic that the end of a plateau is what matters, but this heuristic becomes problematic for heads like 166. A priori, there is no way distinguish the second peak of (16) between a real, distinct critical point, the end of plateau starting from its first peak, and simply an artifact due to the instability of the trace.
We may also examine the data-refined Hessian trace, specifically on the GitHub dataset. This appears substantially harder to interpret than the Hessian trace on the original Pile dataset: the data is much more noisy, does not appear to respond to most developmental stage boundaries, and does not appear to add information beyond that the original weight-refined Hessian traces. However, does remain possible that this is methodological issue some earlier runs with less compute yielded less noisy data. Therefore, we believe the data-refined Hessian trace yields inconclusive results.
D.1.3 DISCUSSION
Given that the Hessian trace appears to find developmental information related to that found by the LLC, is worth comparing the two. Compared to the LLC, the Hessian trace has both advantages and disadvantages. The Hessian trace is more popular in the literature, and is easier to tune hyperpa- rameters for. On the other hand, the Hessian trace has less theoretical support compared to measures like the Hessian rank or LLC, and (unlike the LLC) the Hessian trace cannot measure the behavior of the loss function beyond second order, by definition. Empirically, for the language model we studied, the LLC appears to present a clearer picture of the model's development than the Hessian trace.
As a separate point, is worth emphasizing that the Hessian trace is measured in a substantially different manner to the LLC: the Hessian trace uses second order information whereas LLC estimation uses first order information at many parameters near the original parameter. This decreases the likelihood of correlated mistakes between the two methods, and makes any agreement between them non-trivial.
single parameter,
D.2
FIM TRACE
The Fisher information matrix (FIM) can be seen as particular kind of Hessian matrix (Martens 2020), making the FIM trace closely related to the Hessian trace discussed in Appendix [D.]] Com- pared to the Hessian, the FIM has the advantage of always being positive semi-definite, as well as having deeper anformatiion-therred trace qualitatively similar to that of the Hessian trace, but with more noise due to the estimation methodology.
roots (Amari (016). We find that the results of the FIM
D.2.1 METHODOLOGY
The Fisher information matrix (FIM) is given by
( i
log p()p))))))))) ) (Bug log log p()lw))) ] (Dw,
( a
Iyx = Ezwp(s)w)
where p(a w:) gives the model's 's probabilities over data samples o given a parameter w (Efron & Hastie 2021), Implicitly, the use of p(z/w) requires model to be probabilistic one; for models which are not immediately probabilistic, this step may require some interpretation.
For language models, we choose to interpret the data samples strings of tokens of condext-window length, and given such a string x (for a fixed parameter w), the language model outputs the probability of this string, ((z/w). Note that this differs from more "literal" probabilistic interpretation of the model, where the model is treated as a supervised model p(y) xx w), and yields next-token predictions over next-tokens y given input previous-tokens s.
32

---

Published as a conference paper at ICLR 2025
We prefer the former for two reasons: the p(z|w) representation directly encodes the fact that language model really about natural language sequences (rather than next-tokens), and inference is more natural and direct from p(r) w) (whereas inference from p(y [z, w) requires multiple evaluations to string next tokens together).
In order to estimate the FIM trace, first note that the FIM parameter w' coincides with the Hessian of the population loss at w" if the true distribution the training data is given by p(x (/zl/"), the model w" (Martens, 2020 Via this relationship, we may reuse the methodology for Hessian trace estimation.
Then, in order to estimate the FIM trace, we need only replace the training data with samples from p(zlw") (i.e., generated by the model itself), and repeat the Hessian trace methodology from Appendix D.I. For hyperparameters, we used 30 for the dataset sample count, and 5 for the Hutchinson sample count. All other methodology is identical to Appendix D.I
Note as the parameter w" changes over training, new samples need drawn from p(a (=") at each training checkpoint where we wish to estimate FIM trace- this introduces additional noise and computation cost into the estimation process.
It is possible to estimate the FIM trace in weight-refined manner similar to the LLC and Hessian trace (on the other hand, data refinement would need be done differently), but we only estimate the FIM trace for the entire model.
D.2.2 RESULTS
50000
40000
FIM trace
30000
20000
10000
102
103
104
Step
Figure 15: The trace of the Fisher information matrix (FIM) plotted over the course of training, with stages colored. Note the similarity to the Hessian trace from Figure 14
The FIM trace plotted over the course of training in Fig The shape of the FIM trace appears to largely agree with the shape of the Hessian trace observed in Appendix D.I
The extra noise is due the fact that must sample from p(=l(=") each training step. This variance could possibly be reduced by e.g. keeping samples the same over all training steps, and just adjusting importance weights for the samples.
D.2.3 DISCUSSION
Given that the FIM trace appears to largely match the Hessian trace, while being significantly noisier for a given amount of compute, we chose not to pursue it further.
D.3 HESSIAN RANK
We also measure the (approximate) Hessian rank. From the perspective of singular learning theory, the Hessian rank is more principled than the Hessian trace, because it lower bounds the LLC
'Thi can be concluded by applying the Thom splitting lemma to the population loss function K (w), which changes variables that ((w) may be written of quadratic function K2(w)) wn Huy and
33

---

Published as a conference paper at ICLR 2025
However, empirically, we find that compared the Hessian trace, our estimation of the Hessian rank is more computationally intensive, depends more on hyperparameter choice, and does not reflect structural development in the model as well. We consider these results inconclusive, as it is unclear our methodology is measuring Hessian rank accurately; hetter estimation methodology could potentially avoid these problems.
D.3.1 METHODOLOGY
The approximate rank of square matrix A can be defined the number of eigenvalues of A above some threshold т. Note the choice this threshold can be a nontrivial hyperparameter practice, and this was significant difficulty for
Efficiently estimating the approximate rank of large matrices requires more sophisticated techniques and approximation than the typical algorithms used for smaller matrices. We use the technique described in Ubaru & Saad (2016), which we summarize here.
The rank estimation algorithm exploits the following facts:
Spectral mapping theorem: if p is a polynomial, A is a square matrix, and À is an eigenvalue of A, then p()) is an eigenvalue p(4).
A step function is not a polynomial, but it may be approximated by one; in partieular by Chebyshev polynomials, which are easy to compute efficiently.
The trace of a matrix may be computed efficiently via matrix- vector products (see Appendix D.Ib, and if p is a Chebyshev polynomial, the matrix-vector product p(A)e is easy to compute efficiently.
Combining all of the previous facts: if A is square matrix, and p a polynomial approximat- ing unit step function about a threshold r, then Tr(p(A)) efficiently yields approximately the number of eigenvalues above r (the approximate rank with threshold r).
The details of this algorithm may be found in Ubaru & Saad (2016). As Hessian- vecctor products may be computed efficiently via autodifferentiation, we may use this algorithm to efficiently compute the approximate rank of the Hessian.
Because this method is equivalent to the Hessian trace estimator applied to p(H)) where is the polynomial from above and H is the Hessian, the computational cost of this method determined by the cost of computing the product p( HH).. Each one of these matrix- products requires some constant number of Hessian-vector products, determined by the degree of p, making the Hessian rank require a constant multiple more compute than the Hessian trace.
As this method relies on the trace estimation algorithm, it shares the two hyperparameters from that algorithm (the dataset sample count and Hutchinson sample count). It also has several additional hyperparameters: the degree of the polynomial p. the range of values for which we require p to accurately approximate a step function, and the desired eigenvalue threshold T.
higher degree for p gives a more accurate approximation to the step function, allowing p to rise faster and giving more tineigrined rank approximation, but requires more compute. The larger the range we require the polynomial p to be good step function approximation, the slower p will rise at the threshold value; however, the valid approximation range must at minimum include the entire eigenvalue spectrum of the Hessian, or risk nonsensical results?) The threshold r determines which cigenvalues are considered "zero" for the purposes of the approximate rank.
The degree can be chosen based on required accuracy, or computational resources available. The range can be either set manually, or chosen automatically based the min/max Hessian eigenvalues. The value for the threshold т is harder to choose; there is often no clear eigenvalue gap or other set it arbitrarily.
indicator for where to set т,
arbitrary function K₂ K2(w>) (w2). Then Remark 7.2.3 from (Watanabe 2009) tells us that the learning coefficient of K (w) equal to that of Ki (w)) plus that of K2(w2), and because the learning coefficient of Ki K1(w)) is where r is the rank of H. then the leaming coefficient of K(w) must be least
"This is because the value of Chebyshev polynomial outside its approximation range typically explodes to positive negative infinity, so if any eigenvalue is outside the approximation range, the resulting trace also explodes.
34

---

Published as a conference paper at ICLR 2025
There is another complication for the latter two hyperparameters: Hessian rank over the course of training, not for a single training step. But the Hessian eigenvalue spectrum changes significantly over training; is not clear if, or how, the approximation range and threshold should change over the course of training.
trying to estimate the
We try two strategies set these hyperparameters, which we call the fixed method and the adaptive method.
Fixed method. Keep the approximation range and threshold the same over training. This means we are applying a consistent function to the eigenvalues over training, which is desirable. On the other hand, need set the approximation range large enough to deal with the largest eigenvalue seen over all of training (+= 5000), even if most of the time the cigenvalues are much smaller, So we probably are not getting good resolution estimate of the rank, especially early in training where the Hessian eigenvalues are smaller.
Adaptive method. Adjust the approximation range and threshold over training. We adapt the approximation range to the maximum and minimum eigenvalue at each training step, and set the threshold to a constant fraction of this range. This means we get better resolution when the eigenvalue spectrum is smaller, which is desirable. On the other hand, the function of the eigenvalues we measure will change at each training step, and this may confound the results.
For both strategies, we set the remaining hyperparameters to the same values: 3 for the dataset sample count, for the Hutchinson sample count, and 100 for the degree of p. For the fixed method, we set the approximation range to (-1000, 6000), and the threshold to 500. For the adaptive method, we set the approximation range to (-1.2 x.). and the threshold т to 0.07A, where [^] is the largest eigenvalue of the Hessian in absolute value (estimated by power iteration).
While possible to measure weight-refined or data-refined Hessian rank, we did not measure this, given the difficulty faced with estimating the Hessian rank for the full model on the original dataset.
D.3.2 RESULTS
The Hessian rank over the course of training is plotted in Fig T6 for both hyperparameter selection methods discussed in the prior section. The fixed method appears to somewhat resemble the Hessian trace curve, but it barely changes early in training, and even when it does change later in training, it changes a relatively small amount in proportion to the overall value. The curve from the adaptive method has similarly low variance in comparison its value, but in addition, does not appear interpretable.
D.3.3 DISCUSSION
We believe these results are preliminary and inconclusive. Our assessment that this methodology is not producing good results stems from two main observations: (1) the failure to detect stage boundaries that are consistently identified by other methods (e.g., LLC, ED, Hessian trace, and behavioral metrics), and (2) the high sensitivity to estimation methodology, evidenced by the discrepancy between fixed and adaptive methods and the dependence on numerous hyperparameters.
8350
Agpret 1605 oxdeptiee meihoo
22050
22000
8340
21950
8330
21850
8320
21800
8310
21750
21700
107
10a
10t
12
12a
10f
Figure 16: The approximate rank of the Hessian plotted over the course of training, estimated using the fixed method (left) and the adaptive method (right), with stages colored. Note y-axis scale in both charts: while both curves appear to show large changes in relative terms, they have little variation in absolute value.
35

---

Published as a conference paper at ICLR 2025
These issues cast doubt on whether our estimated Hessian rank accurately captures the structural development in the model.
It remains unclear whether the problem lies in the theoretical underpinnings of the Hessian rank approach (e.g., its potential inability to capture higher-order degeneracy), in the practical implementa- tion (due to flawed estimation methodology), or both. Given the apparent limitations of the current estimation methodology, we are unable to draw further conclusions at this time.
However, it does seem possible to improve the Hessian rank estimation methodology to deal with the issues we have highlighted. We suggest a few ideas for future work:
To determine if the Hessian rank could even work well in principle, measure it in some toy model small enough to compute the Hessian cigenspectrum exactly (but still large enough to exhibit nontrivial stagewise development).
Attempt something in between the fixed and adaptive methods: perhaps fixing the approxi- mation range and threshold for each period of model training where the Hessian spectrum does not change significantly, or for each developmental stage.
Fix noise in the adaptive method caused by noise in the max eigenvalue, by e.g. filtering the value of the eigenvalue over time.
Find principled way set the threshold based some desired loss after dropping all eigenvalues below т. In turn, the desired loss could be set to a desired "effective compute" value via scaling laws.
E COMPARISON AGAINST ABLATION ANALYSIS
F.1 BACKGROUND
Ablation techniques are widely used in mechanistic interpretability test hypotheses about the importance of specific components or activations in neural networks (Chan et al. 2022a Wang et al., 2023 Rauker al. 2023 [Bereska & Gavves 2024). By selectively removing or modifying parts of a model and observing the impact on performance, researchers aim to identify critical elements of the model's s computation.
In this analysis, we compare three common ablation methods: zero ablation (Meyes et al. 2020 Hamblin et al.) 2023 Nanda et al. 2023 Morcos al., 2018 Zhou et al., (018). mean ab ablations
(Bereska & Gavves, 2024), and resampling ablations (Chan et al. 2022a Hanna et al., 2023
Goldowsky-Dill et al. 2023 Lieberum al. 2023). Our goal to understand the relative strengths
and weaknesses of each approach in comparison to the weight-refined LLC.
E.2 METHODOLOGG
We compared the following ablation techniques:
Zero ablation: Setting the targeted activations to zero.
Mean ablation: Replacing the targeted activations with their mean value across the dataset.
Resampling ablation: Replacing the targeted activations with those from randomly selected different input (Chan et al. 2022a).. In practice, we rolled the activations forward over the batch index
(Chan et al. 2022a).
The ablation score is the difference between the loss before and after the ablation (computed over the same dataset over which we average the activations for the mean ablations).
Besides treating ablations as a baseline to compare rLLCss against in this section, we use ablations at several points throughout the rest of the paper (Figure 6 and Appendix B.I to complement the other analyses. This requires introducing one additional technique:
Path patching (Wang et al.) 2023: Goldowsky-Dii Dill et al. 2023 : Path patching involves
first running two forward passes: one uncorrupted inputs/activations, and one involving
36

---

Published as a conference paper at ICLR 2025
comrupted inputs/ablations. Then, runs final forward pass, patching in the uncorrupted & corrupted activations so as to isolate the role specific computational path.
For example, for the path ablation in Figure 6 mean-ablate the layer multigram heads, then save the outputs of head (15), then we run another (clean) forward pass and patch in (155ss activations before the layer readout layer. This lets determine the particular influence of the layer 0 multigram heads on head (15,, without having to worry about layer O heads' other roles.
E.3 RESULTS
Figure 17 displays the results of the different kinds of ablation over training time.
Ablations differentiate heads. As with the rLLC, we can use these ablation scores to distinguish several types of heads:
The previous-token heads 0:1 and 0:4 can be identified by an increase in the ablation scores during stage LM4.
The current-token head 0:5 is also distinguishable by its increase across the ablation scores starting towards the end of LM2 until it reaches a peak during LM4, after which decreases. This is especially pronounced in the resampling ablation scores.
The induction head 177 is clearly distinguished by the increase the ablation scores in LM4. The other induction head (1:6) is less clearly distinct, though its ablation scores do have a different shape from the multigram heads.
The multigram head 0:0 has similar rLLC curves to the other layer 0 multigram heads (though is relatively larger, see Figure 2 The ablation scores suggest that 00:0 is a distinct type of head throughout much of training, with substantially higher values and qualitatively distinct shape. It is not until LM5 that this heads ablation scores settle to a value comparable to the other layer 0 multigram heads. This complements the analysis in Appendix th.t suggests 0:0 starts out as a "space" head that attends to whitespace tokens.
The shape of ablation scores is mostly conserved across the different types of ablations, with one exception: 1:3, which reaches a local maximum at the LM2-LM3 boundary for the zero ablation score that is not visible in the other ablation scores. There are other more minor differences for 00:5 (described above).
Intriguingly, the ordering of the final resampling ablation scores is very close to the ordering of the final rL..Cs. These are both also similar to the ordering of the final Hessian traces (Figure
E.4 DISCUSSION
Weights vs. activations. The most pronounced difference between the wrLLC and the ablations considered here is that the WrLLC operates in weight space, while ablation methods work in activation space. In practice, the two approaches seem be complementary. Indeed, the LLC can be interpreted as a kind of ablation: it measures the expected change in behavior (as measured by the loss) under "typical" perturbations to weights, where "typical" means the perturbation is drawn from the local posterior.
Discrepancies between the wrLLC and ablations. Both approaches identify key developmen- tal stages and specialized heads, such as previous-token and induction heads. However, notable discrepancies exist:
The ablation scores better at identifying discrepancy in 0:0. On the other hand, the wrLLC is better at identifying that this head ultimately matures into multigram head (which we confirm separately).
Some heads (c.g.. induction head (166 and the current-token head 0:5) are more distinguish- able in the wrI.L. than in ablations.
37

---

Published as a conference paper at ICLR 2025
Ablation methods generally computationally more efficient. However, they lack the theoretical grounding of the wrLLC (which reflected in the existence of many different possible choices of ablation).
Overall, the resampling ablation score most clearly differentiates the different kinds of heads, which we take as positive evidence in favor of the method (Chan et al., (022a).. However, the practical distinctions are small enough that default to using mean ablations.
ADDITIONAL. EXPERIMENTAL DETAILS
TRAINING MODEL DETAILS
Model architecture: We considered the same model architecture in Hoogland al. (2024)
Olsson et al. (2022) : two-layer attention-only transformer, with the following specifications:
Context length: 1024 tokens
Residual stream dimension: 256
Number of attention heads per layer: 8
Layer normalization: Included
Positional embedding: Learnable Shortformer-style
The resulting models contained approximately million parameters. We implemented these models using the TransformerLens library (Nanda & Bloom 2022).
For tokenization, we employed a modified version of the GPT-2 tokenizer, reducing the vocabulary size from 50,000 to 5,000 tokens. This adjustment allowed decrease the overall model size.
Layer 0 Heads
Layer 1 Heads
Zero Ablation Score
0.20
0.15
0.2
0.05
0.1
0.0
0.00
10'
10
16
Mean Ablation Scote
0.15
0.3
0.10
0.1
0.00
10F
Ressampling Ablation Score
0.3
0.4
0.2
0.1
0.0
0.0
Stap
Stap
IN4
11637
Figure 17: Comparison of attention-head ablations over training time. This figure shows the results of three different ablation techniques (zero ablation (top), mean ablation (middle), and resampling ablation (bottom)) applied to attention heads from layer 0 (left) and layer (right). The y-axis shows the ablation score (difference in loss before and after ablation).
38

---

Published as a conference paper at ICLR 2025
Training: The model and training considered in the main body is the same in Hoog land et al. (2024). This was trained on a subset of the DSIR-filtered Pile (Gao et al. 2020 Xie et al. 2023) for total of around 50, 000 steps, with batch size of 100.
F.2 LLC ESTIMATION DETAILS
For estimattnn the Local Learning Coefficient (LL.), we employed Stochastic Gradient Langevin Dynamics (SGLD) with the following hyperparameters:
The SGLD step size le-3
The inverse temperature 3 = 30/n
The localization strength 200
Number of independent chains: 4
Burn-in steps: 0
Draws per chain: 200
The procedure is the same as described in Hoogland et al. (2024) using code from van Wingerden al. (2024). The cost per SGLD-update is comparable SGD with linear increase in memory requirements.
We estimate that the computational cost of producing rLLC curves like those in Figure for all heads scales like CHT log(N), where C' is roughly constant O(100) number of SGLD samples, H is the number of heads, N the total number of training steps, and T is the cost of a single training pass.
F.3 DATA-REFINED LLCs USING OTHER MODELS AS THE GENERATING PROCESS
In Section 4.3 we make use data-refined LLC where a particular trained one-layer attention-only transformer used to provide the data distribution. For computational reasons, make use of modified form
n
K-1 Σ
(ur; M)
Σ
DEL (SI
(11)
n
of the empirical loss where M denotes the one-layer (LI) transformer.
Lo and L1 models: The LI models used for these data-refined LLCs were trained analo- gously to the model considered in the main text and described in Appendix F.I
F.4 COMPOSITION SCORES
Let Wh. Wh. W be the query, key, and value weights of attention head h respectively. There three types of composition between attention heads transformer models in Elhage et al. (2021):
C-Coppssiiion the query matrix W an attention head reads 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 subspace affected by previous head
If Wh the output matrix of an attention head, then Whh Whh Wh and Wh Whwh. The composition scores are
[|MWAL ot |
(12)
Where M Wha T M Who Qk. and M why for Q- K-, and V-Composition respectively.
39

---

Published as a conference paper at ICLR 2025
F.5 CLUSTERING RLLCS
To better understand the patterns and relationships among the attention heads based on their rLLC trajectories, we applied various clustering algorithms the rLLC data. This analysis aims to quantitatively group attention heads with similar developmental patterns.
F.5.1 CLUSTERING ALGORITHMS
We employed four different clustering algorithms to ensure a comprchensive analysis, using imple- mentations provided by sklearn (Pedregosa et al. 2011) and tslearn (Tavenard al. 2020
1. Euclidean K-means: centroid-based algorithm that partitions n observations into &t clus- ters, minimizing the within-cluster sum of squares based on Euclidean distance (Macqueen 1967).
2. DTW K-means: A variation of K-means that uses Dynamic Time Warping (DTW) the distance measure, which is particularly suitable for time series data as it allows for non-linear alignment of sequences (Berndt & Clifford 1994 Sakoe & Chiba 1978).
3. Hierarchical Agglomerative Clustering (HAC): A bottom-up approach that starts with each observation as a separate cluster and merges them iteratively based on a chosen linkage criterion (in this case the minimum variance criterion of Ward Jr 1963).
4. Shape-Basee Based K-means: A custom implementation of K-means that uses Shape-Basee Based Distance (SBD) as the distance measure. SBD designed to capture similarities in the shape of time series, regardless of differences in scale offset (Paparrizos) 2018).
In particular, apply the clustering algorithms the per-head Pile drLLCs and per-head Github dr- LLCs concatenated together. That is, the space in which we perform our clustering is 27 dimeennsiinnl,, where T' is the number of checkpoints. We found it unnecessary apply standardization techniques to the trajectories before clustering as this led to marginal changes in the learned clustering.
F.5.2 EVALUATION METRICS
To assess the quality of the clustering results, we used the following metrics:
1. Silhouette Score: Measures how similar an object is to its own cluster compared to other clusters. The ranges from -1 to I, where higher score indicates better-defined clusters (Rousseeuw 1987).
2. Calinski-Harabass Harabasz (CH) Index: Also known as the Variance Ratio Criterion, this index is the ratio of the sum of between-cluster dispersion and of within-cluster dispersion. A higher score indicates better-defined clusters (Caliñski & Harabasz, 1974).
3.
Davies-Bouldiy (DB) Index: Calculates the average similarity between each cluster and its most similar cluster. A lower score indicates better clustering (Davies & Bouldin 1979).
F.5.3 RESULTS
Clustering Layer 0 heads. Figure 18 shows that as the number of clusters e increases, the first distinction that emerges (for e = 2) within layer is between the current-token head, previous-token heads, and multigram head (0:0) and the remaining multigram heads. Subsequently (c = 3). the first
cluster splits apart: either the current-token head (0:5) or 00:0 breaks off from the rest, depending on the exact choice of clustering algorithm. The next distinction emerge is between 0:7 and the bulk of the layer attention heads.
As analyzed in Appendix B heads 0.0 and 0:7 are, in fact, special kinds of multigram heads. Head (:7) specializes to Dyck pattern (see Appendix B.3) and head 0:0 initially specializes to long -ra-ggrmms of spaces (see Appendix B.6, the exact function remains unknown).
40

---

Published as a conference paper at ICLR 2025
Number clasterscm c
Number cicislescc
Number disless
Kumber clasbers c.m
Clates Clate
-olec 11:4-33
Saaer 32.35
Outer C5.33
1.54
Ositon
D10000 C4
Clattr
Clatery Caale B.
Mest!
Stype
Figure 18: Clustering results for Layer 0 rLLC trajectories using various algorithms (rows) and number of clusters (columns). The x-axis represents training steps, and the y-axis shows the rLLC values. Each line represents an attention head, colored by its assigned cluster. As the number of clusters increases, we observe a clear separation between the current-token head, previous-token heads, and multigram heads. These clusters are fit on concatenation of Pile drLLCs and Github drLLCs, but only the Pile drLLC component is displayed.
Clustering Layer heads. Figure 19 shows that the clustering extremely consistent within layer 1. The two induction heads are clearly distinguished from the bulk of the layer heads. Upon increasing the number of clusters, the bulk of layer 1 multigram heads splits into pairs (see column 3 for 4 in Figure 19). Additionally, given enough clusters, the two induction heads become separated from one another.
As in the case of the layer O heads, these finer distinctions are meaningful: heads 1::5 and 1:3 are specialized to Dyck patterns (Appendix [B.3), and the remaining heads are involved more prosaic (skip) n-grams. This clustering is not perfect: in particular, clustering suggests a distinction between pairs 1:11 and 1:4/100, which is not obvious in the analysis by tokens context (Appendix B
41

---

Published as a conference paper at ICLR 2025
Number cm
Number clalers
Number disless
Kumber clabeess
Clater
Cudree
Claters
42
CARTIN
Clatery
Culler
Curtre
Stpprr
Segs
Stype
Figure 19: Clustering results for Layer 1 rLLC trajectories using various algorithms (rows) and number of clusters (columns). The x-axis represents training steps, and the y-axis shows the rLLC values. Each line represents attention head, colored by its assigned cluster. The results consistently show clear distinction between induction heads and other heads across all methods. These clusters fit on concatenation of Pile drLLCs and Github drLLCs, but only the Pile drLLC component is displayed.
We note that the clusters are largely unaffected when first standardizing the trajectories (by subtracting the mean and dividing by the standard deviation computed across steps and possibly heads).
we show that these clusters appear to be universal across training runs: clusters to this seed generalize to different training runs.
Generalization to other seeds. In Appendix G
F.5.4 DISCUSSION
Capturing shape similarities: We chose use the Euclidean K-means algorithm for our ultimate classification, with = 3, 4 for layer 0 and c 2, 4 for layer 1. We manually chose the number of clusters based on visual inspection.
Generally, the different clustering algorithms were in tight agreement. However, the Shape Shape-Based K-means algorithm performed best when the clusters were fit just the Pile drLLC data and not the Github drL.LC data. This technique also transfers better to other seeds (Appendix G and is closer to the results of clustering by eye. However, shape-based clustering performs worse when we increase the cluster count to attempt to resolve finer distinctions within the layer I multigram heads (Figure [ In practice, we recommend taking majority vote approach, with shape-base clustering getting additional weighting.
Limitations of evaluation metrics: It's important to note that the standard evaluation metrics (Silhouette Score, Calinski-Harabasz Index, and Davies-Bouldin Index) did not appear work well
42

---

Published as a conference paper at ICLR 2025
in our context. This is likely due to the low-sample setting of our analysis, with only 8 trajectories per layer, and due to the fact that selideen-distancee based diagnostic metrics may not be appropriate when the shape appears to be the more salient feature for clustering. These metrics are designed for larger datasets and may not provide reliable guidance in our case.
G COMPARISON WITH OTHER SEEDS
We trained three additional seeds identically to the original setup described in Appendix F and Hoogland al. (2024). Applying weight- and data-refined LLCs to these training runs yields the figures in Figure 23
The different training runs are very similar (Figure 23 but there some differences:
1. Some models develop only single induction head and previous-token head. Two of the runs have two heads each. Two of the runs have single head each. Other than this, each run has single current-token head, and the remainder of the heads seem to be multigram heads,
2. In some runs, there is no LM3 stage (in which the full-model LLC is decreasing). However, it is always the case that the multigram-head rL.LCCs begin decreasing after I.M2.
Remarkably, despite these differences, clusters that fit to the original seed generalize per-head rLLCCs from different training runs. Figure 20] Figure 21 and Figure 22 show the results of applying the shape-based clusters fit seed (Figure [8 and Figure [9| bottom rows) different training runs. There are exceptions: e.g., for seed 2, layer O with c = 3 clusters, one of the previous-token heads is identified as a current-token head, and for seed 3 the gurrent-token head does not get distinguished from the other heads by any crossing.
Number clusters
Number clusters
Number dusters
Number
Share gotes
Chaisa
20
Cupter
Oudders
Outer
10
Supaa
Steps
Stppe
Number c
Number of clateexes
Number dusies.c
Number of claters
18
Chillera
10P
30
10
12
127
Sepprr
Stepr
Stype
Figure 20: Clustering validation on seed 2 for layer 0 (top row) and layer (bottom row). Applying shape-based clusterings fit to the original seed Pile drLLC and Github drLLC generalize to held out seeds (Figure 18 and Figure 19 bottom rows). Here, the current-token head is 0:6), the previous-token head is 0.11 and the induction head is 1:0..
43

---

Published as a conference paper at ICLR 2025
Number c
Number clalers-
Number disiesc
Kumber of cluteers c
Chatess
2=2
dusters
Chote's
Cuslers
Outer
Figure 21: Clustering validation seed 3 for layer 0 (top row) and layer (bottom row). Applying shape-based clusterings fit to the original seed Pile drLLC and Github drLLC generalize to held out seeds (Figure 18 and Figure 19 bottom rows). Here, the current-token head is 0:2, the previous-token head 004,, and the induction head is 1:6).
54
20
Chictess
Guilrr
Outer C3
18
Supar
Steps
Stape
Number cm
Number cl
Number dustess
Number of cuaters
Gader L-
10 107
30a
10
12
127
Supe
Stapa
Stppa
Figure 22: Clustering validation on seed 4 for layer 0 (top row) and layer (bottom row). Applying shape-based clusterings fit the original seed Pile drLLC and Github drLLC generalize to held out seeds (Figure 18 and Figure 19 bottom rows). Here, the current-tokes token head is 0:6), the grevious-token heads are 0:0 and and the induction heads are and
H COMPARISON WITH LARGER MODELS
In Figure 24 we show that the refined LLC generalizes to larger model settings. In particular, for
Pythia-70ne 70m (Biderman et al. 2023). the Github wdrLLC (cf. Figure 3 right) distinguishes the previous-token heads and induction heads (that together make up the induction eircuit) from other types of heads. To verify the identity of these heads, we use the previous-token and matchin scores following Olsson al. (2022).
44

---

Published as a conference paper at ICLR 2025
Ful Model
Layer 0
Layar
30-
100
40-
30-
160
20-
10
Step
Stege
Stoge
Full Model
Layar
Layar
10
60
20
11*
40-
160
Steg
Full Model
Layer
Layer
160
170
20-
30
16*
60
50
Figure 23: Additional seeds show similar her-head rLLCs the original training run. Every seed has current-toker head. Every seed has at least one previous-token head and one induction head (the top two seeds only have single one each). All of remaining heads appear to multigram heads.
45

---

Published as a conference paper at ICLR 2025
150
Frevius Taltn Scote
1.0
Laper Reantve Difference (8)
150
107
All heods Profit Score
100
0.8
100
30.1
Heads
0.6
211
Restive
10-7
0.4
-50
hoass
-50
10-4
3:6
10-1
197
10'
10
107
10"
10
107
109
10a
107
307
30*
Step
Stop
Stap
Stop
Figure 24: The induction circuit in Pythia 70M revealed through data-refined LLCs. The panels show different metrics tracking the development of the induction circuit: (Left) Relative difference between Pile and GitHub LLCs for Layer 2 heads with the outlier identified as the model's sole previous-token head. (Center-left) Previous-token scores across all heads confirm this identification. (Center-right) Relative LLC differences for Layer heads reveal two clusters the end of training: candidate induction heads and non-induction heads. (Right) Prefix scores confirms the cluster of layer heads with higher relative rLLCs as the induction heads, with head 3:2 being an exception.
46


Reviews:
[{"review_id": "9w8k0nZS8i", "reviewer": "Reviewer_RBF6", "summary": "The authors introduce a novel approach to understanding the internal specialization of attention heads in Transformer models through Refined Local Learning Coefficients (rLLCs). By measuring the complexity of attention heads and tracking their developmental process during training, the authors reveal how different heads gradually specialize in performing specific tasks.", "questions": "1. Could insights from rLLC tracking be used to optimize training processes, such as by dynamically adjusting learning rates or selectively freezing layers based on head specialization patterns?\n\n2. How well does the rLLC methodology generalize to larger models, such as GPT-like architectures?\n\n3. How do rLLCs compare with the existing interpretability metrics? The authors could refer to these related works: [Michaud, Eric, et al. \"The quantization model of neural scaling.\" Advances in Neural Information Processing Systems 36 (2024).] ; [Xiao, Xiongye, et al. \"Exploring neuron interactions and emergence in llms: From the multifractal analysis perspective.\" arXiv preprint arXiv:2402.09099 (2024).]", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The rLLC provides an innovative, quantitative tool for tracking specialization in Transformer models, which could have applications beyond this study.\n\n2. The paper includes a comprehensive analysis of attention head specialization, backed by clustering and ablation studies that demonstrate the functional roles of different heads.\n\n3. The discovery of \"multi-phrase circuits\" and the evidence of cross-layer head collaboration provide valuable insights into how Transformer models handle complex linguistic patterns.", "weaknesses": "1. While the method is effective on smaller models, the paper does not discuss the computational implications for larger models.\n\n2. The study focuses on understanding Transformer models without discussing how these insights could guide architecture optimization or training strategies in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Dv3cz3zQ0y", "reviewer": "Reviewer_qe7W", "summary": "This paper introduces refined Local Learning Coefficients (rLLCs) as a method to analyze the internal structural differentiation within transformer models. \nIn particular, they introduce weight refined LLCs and data-refined LLCs.\nBy applying rLLCs to individual attention heads of a two-layer attention-only transformer model on different types of data, the authors aim to track how these components specialise and differentiate into distinct functional roles across training. The work provides insights into attention head specialisation based on the types of data processed, discovers specific structural patterns (e.g., multigram circuits), and highlights the correspondence between data distributional structure, loss landscape geometry, and model learning dynamics.", "questions": "- it was not clear to me how the K-means clustering was obtained to generate the colours for the different lines in e.g. Plot 1. What is the K-means over? \n- Have you considered how the rLLCs would scale with deeper, more complex transformer models? Are there any preliminary results or hypotheses about the performance and practicality of this method on architectures with more layers or MLP components?\n- Beyond interpretability, have you explored any practical use cases where rLLCs could directly influence model design or training strategies (e.g. model pruning)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper presents a novel metric to track structure during training based on data structure and tracking the evolution and specialisation of attention heads, overcoming some of the limitations of previous metrics. The paper presents some novel insights into how heads specialise and the effect of data.\n- The presented metric (rLLC) overcomes some of the existing limitations of the LLC method, including the assumption that the $w^*$ is a local minima (which during training, it is very unlikely to be)\n- rLLC seems to in practice be able to differentiate attention heads by their specialisation, it is interesting to observe the higher complexity of heads performing memorisation and n-gram heads vs the induction ones (Figure 1)\n- The authors show how with a different data (Github) the head specialisation changes and more importance is given to the induction heads (Figure 3)\n- The experiments are well-documented, with comprehensive comparisons against other interpretability measures like Hessian and Fisher Information Matrix traces, adding robustness to the findings.", "weaknesses": "- the experiments limit themselves to a two layer transformer, it would have been interesting to see how the proposed rLLC metric would behave in a bigger transformer. Is it still trackable? Does it lead to a meaningful interpretation of results?\n- what about different architectures? Line 465 in the related section mentions AlexNet, it would have been interesting to carry out a similar analysis on layers of AlexNet", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qe9dlE69c2", "reviewer": "Reviewer_YkMj", "summary": "The authors introduce refined variants of the Local Learning Coefficient (LLC), a measure of model complexity grounded in singular learning theory, to study the development of internal structure in transformer language models during training. By applying\nthese refined LLCs (rLLCs) to individual components of a two-layer attention- only transformer, they were able gain novel insights into the progressive differentiation and specialization of attention heads. The methodology used in this paper reveals how attention heads differentiate into distinct functional roles over the course of training, analyzes the types of data these heads specialize to process, and discovers a previously unidentified multigram circuit. The authors conclude that their findings demonstrate that rLLCs provide a principled, quantitative toolkit for developmental interpretability,", "questions": "See comments above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The work is well grounded in the existing literature and theories. A good number of appropriate citations back this up that the authors build the work in this paper on. \n\nThe mathematical descriptions are brief and concise but sufficiently detailed and understandable. Appropriate level of detail.\n\nSection 3.3 Limitations and the comparisons to related work within was much appreciated. \n\nThe degree and amount of supplementary information and detail in the appendices was excellent and greatly supports the main paper.", "weaknesses": "It would be beneficial to even if briefly define how the authors are explicitly using the terms 'component' and 'arbitrary data distribution' earlier on since it is the differentiating elements to other LLC metrics. The reader is left wondering and speculating as they continue to read in the way it is presented in the original paper. \n\nThe immediate next statement says \"We focus mainly on the rLLCs of individual attention heads ...\", implying diversity in how their new rLLC metric can be defined, tweaked, and/or applied.  But the authors have not even defined what it is. It is introducing a degree of variability prematurely that comes across as confusing. Before implying what they are 'focusing' on among several possibilities of (how?) their rLLC metric is defined, how about saying what it is first clearly and explicitly. \n\nSimilarly, the use of the term 'development' at the bottom of page 2 would be strengthened by giving the reader at this stage a brief one or two sentence understanding of the explicit meaning of the term as adopted in the paper. Italicizing it alone is not enough. The reader understands its important, but please help them in building understanding along the way.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j9qJ1o7Fy5", "reviewer": "Reviewer_w7vo", "summary": "This paper proposes two types of refined local learning coefficients (LLC), weight-refined LLC and data-refined LLC, and applies them to a two-layer attention only transformer to study the progressive differentiation and specialization of the attention heads. Both LLCs (defined in eq 5) are based on the estimated LLC (defined in eq 2), where suppressing q or V leads to weight-refined or data-refined LLC. The limitations of the proposed LLCs are given in Sec 3.3. Experimental results show that the refined LLCs reveal how attention heads differentiate (featured in Fig 2) and specialize (featured in Fig 3). In addition, combining the refined LLCs reveal internal structure related to multigram prediction that enables nested bracket-matching.\n\nI am not an expert in this area, and therefore read the paper multiple times. I enjoyed reading it. My rating is 6, but I'm willing to raise the score if weakness #1 is addressed because it is preventing me from understanding the general applicability/feasibility of the proposed method.\n\n[update during discussion] The authors have added experiments on a larger model, I therefore raised my rating.", "questions": "1. If the analysis is applied to transformers learned in different domains (e.g., decision transformer in RL), will we have similar observations?\n2. Tiny issues in equations:\n    1. In eq (1), if $f_w$ is a function from contexts to probability distributions, do you still need $\\texttt{softmax}$?\n    2. Is $\\lambda$ in equation (5) missing a hat?\n    3. Is the threshold $t$ in eq (4) and (5) the time steps? If not, consider changing it as $t$ is used in the figures as time steps.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is well-written, even a non-expert can easily follow the content if they are willing to spend some time figuring out the math.\n- Although on a toy network, the results presented by this paper are strong and interesting.", "weaknesses": "- In Sec 6, the authors mentioned that *\"The techniques pioneered in this paper for understanding internal structure in two-layer attention-only transformers can of course be applied to models at a larger scale or with different architecture.\"*. It would be more convincing if they present some results in the paper. [I'm willing to raise my rating if this can be addressed]\n- [This is a suggestion] When I read the paper for the first time, I didn't immediately understand \"what this paper unlocks\" or \"why this work is important\" after finishing the introduction section. I suggest improving the writing of this section. For example, moving some text from Sec 6 to the introduction may be a good option.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)