Towards Empirical Interpretation of Internal Circuits and Properties in Grokked Transformers on Modular Polynomials
Authors: Hiroki Furuta
Gouki Minegishi
Yusuke Iwasawa
Yutaka Matsuo
Venue: Accepted by TMLR ()

Abstract: Grokking has been actively explored to reveal the mystery of delayed generalization and identifying interpretable representations and algorithms inside the grokked models is a suggestive hint to understanding its mechanism. Grokking on modular addition has been known to implement Fourier representation and its calculation circuits with trigonometric identities in Transformers. Considering the periodicity in modular arithmetic, the natural question is to what extent these explanations and interpretations hold for the grokking on other modular operations beyond addition. For a closer look, we first hypothesize that (1) any modular operations can be characterized with distinctive Fourier representation or internal circuits, (2) grokked models obtain common features transferable among similar operations, and (3) mixing datasets with similar operations promotes grokking. Then, we extensively examine them by learning Transformers on complex modular arithmetic tasks, including polynomials. Our Fourier analysis and novel progress measure for modular arithmetic, Fourier Frequency Density and Fourier Coefficient Ratio, characterize distinctive internal representations of grokked models per modular operation; for instance, polynomials often result in the superposition of the Fourier components seen in elementary arithmetic, but clear patterns do not emerge in challenging non-factorizable polynomials. In contrast, our ablation study on the pre-grokked models reveals that the transferability among the models grokked with each operation can be only limited to specific combinations, such as from elementary arithmetic to linear expressions. Moreover, some multi-task mixtures may lead to co-grokking -- where grokking simultaneously happens for all the tasks -- and accelerate generalization, while others may not find optimal solutions. We empirically provide significant steps towards the interpretability of internal circuits learned through modular operations, where analytical solutions are not attainable.

Published in Transactions on Machine Learning Research (11/2024)
Towards Empirical Interpretation of Internal Circuits and Properties in Grokked Transformers on Modular Polynomials
Hiroki Furuta
Gouki Minegishi
Yusuke Iwasawa
Yutaka Matsuo
ffuruta, minegishi, awasawa, matsuo @hweblab.tt t. u-tokyo..cc.jp The University of Tokyo
Reviewed on OpenReview: https: openreview. net/ forun? id=Ma5f70uXJ0
Abstract
Grokking has been actively explored to reveal the mystery of delayed generalization and identifying interpretable representations and algorithms inside the grokked models is suggostive hint to understanding its mechanism. Grokking on modular addition has been known to implement Fourier representation and its calculation circuits with trigonometric identities in Transformers. Considering the periodicity in modular arithmetic, the natural question to what extent these explanations and interpretations hold for the grokking on other modular operations beyond addition. For closer look, we first hypothesize that (1) any modular operations can be characterized with distinctive Fourier representation or internal circuits, (2) grokked models obtain common features transferable among similar operations, and (3) mixing datasets with similar operations promotes grokking. Then, we extensively examine them by learning Transformers on complex modular arithmetic tasks, including polynomials. Our Fourier analysis and novel progress measure for modular arithmetic, Fourier Frequency Density and Fourier Coefficient Ratio, characterize distinctive internal representations of grokked models per modular operation; for instance, polynomials often result in the superposition of the Fourier components seen in elementary arithmetic, but clear patterns do not emerge in challenging non-factorizable polynomials. In contrast, our ablation study on the pre-grokked models reveals that the transferability among the models grokked with each operation can be only limited to specific combinations, such as from elementary arithmetic to lincar expressions. Moreover, some multi-task mixtures may lead to co-grokking where grokking simultancously happens for all the tasks and accelerate generalization, while others may not find optimal solutions. We empirically provide significant steps towards the interpretability of internal circuits learned through modular operations, where analytical solutions are not attainable.
1 Introduction
Grokking is a late generalization phenomenon when training Transformer (Vaswani et al., 2017) and other architectures with algorithmic data (Power et al., 2022) where training accuracy soon reaches 100% with low test accuracy (often 0%), and after several iterations, test accuracy gradually reaches 100%. Grokking has been actively explored to reveal the mystery of delayed generalization, and identifying interpretable circuits inside the grokked models should be a suggestive hint to understanding the grokking mechanism and dynamics. The interpretability analysis has mainly shed light modular addition, where grokking obtains the calculation with Fourier basis and trigonometric identities (Nanda et al., 2023; Zhong et al., 2023; Gromov, 2023; Rubin et al., 2023). Considering the periodicity in modular arithmetic, the natural question is to what extent these explanations and interpretations hold for the grokking on other modular operations beyond addition.
For closer look at the connections among the grokking phenomena in other modular operations, we first hypothesize that. (1) any modular operations be characterized with unique Fourier represenlation or

---

Published in Transactions on Machine Learning Research (11/2024)
algorithms (circuit formulation), (2) grokked models obtain common features transferable among similar operalions (transferabilityy) and (3) mixing funclionally similar operalions in dalasel promole grokking (multi-task training). Revealing these relations would help us understand and analyze the dynamics of grokking better. In this work, beyond the simplest and well-studiod operation, we observe the internal circuits learned through grokking in complex modular arithmetic via interpretable reverse engineering, and extensively verify our three hypotheses, while also investigating whether grokked models may exhibit transferability among the models grokked with other operations and scaling to the similarity and the number of tasks'.
First, analyzing modular subtraction, multiplication, and polynomials reveals that the operations that cause grokking have unique Fouricr representations (Section 5). For instance, subtraction poses a strong asymmelry on Transformer (Section 5.1), and multiplication requires cosine-biased components at all frequencies (Section 5.2). Grokking can easily occur in certain modular polynomials, such as the sum of powers and higher-degree expressions factorizable with basic symmetric and alternating expressions (Section 6). These polynomials have superposition of representations modular elementary arithmetic, while "non-grokked" operations do not have explicit patterns (Section 6.1). We also introduce the novel progress measure for modular arithmetic; Fourier Frequency Density (FFD) and Fourier Coefficient Ratio (FCR), which not only indicate the late generalization but also characterize distinctive internal representations of grokked models per modular operation (Section 6.3). We prove that our proposed FFD and FCR decrease accompanying the test accuracy improvement, and they reflect. features of internal circuits, such as the coexistence of addition and multiplication patterns in ab + b, or dependence of the factorizable polynomials on the parity of exponent n.
In contrast, the ablation study with pre-grokked models reveals that the transferability of grokked embeddings and models is limited to specific combinations, such as from elementary arithmetic to linear expressions (Section 7.1), and could be rarely observed in higher-degree expressions (Section 7.2). Besides, some mixtures of multiple operations lead to the co-occurrence of grokking and even accelerate generalization (Section 8.1). In contrast, others may interfere with cach other, not reaching optimal solutions (Section 8.2). These observations indicate that the mechanism of grokking might not always share the underlying dynamics with common machine learning. We provide significant insights in the empirical interpretation of internal circuits learned through modular operations, where analytical solutions are not attainable.
2 Related Work
Grokking Grokking has been actively studied to answer the questions: (1) when it happens, (2) why it happens, and (3) what representations are learned. In simple algorithmic tasks like modular addition, grokking would be observed with proper weight decay and the ratio of train-test splits (Power et al., 2022; Lyu et al., 2023). In addition to synthetic data (Liu et al., 2023b), grokking could occur in more general settings such as teacher-student (Levi et al., 2023), NLP (Murty et al., 2023), computer vision (Thilak et al., 2022), or molecular graph tasks (Liu et al., 2023a), which could be explained with the dynamic phase transitions during training (Rubin et al., 2023; Kumar et al., 2023) or mismatch between the train-test loss landscape against weight norm (Liu et al., 2023a). Recent findings have revealed that while grokking has initially been observed in neural networks (MLP and Transformer) with weight decay, it may also occur in Gaussian processes and linear regression models (Levi et al., 2023; Miller et al., 2023) or even the case without weight decay (Kumar et al., 2023). Our work focuses on complex modular arithmetic including subtraction, multiplication, polynomials, and a multi-task mixture, and then empirically analyzes the difference wet grokked and non-grokked modular operations.
Several works have argued that the late generalization dynamics has been driven by the sparsification of neural network emerging dominant sub-networks (Merrill et al., 2023; Tan & Huang, 2023) and the structured representations (Liu et al., 2022); the training process could be a phase transition divided into memorization, circuit formation, and cleanup phase (Nanda et al., 2023; Xu et al., 2023; Doshi et al., 2023; Davies el al., 2023; Zunkovič & llievski, 2022), and the formation of generalization circuits produces higher logits with small norm paramcters than memorization circuits (Varma et al., 2023). The sparse lottery tickets in neural
2

---

Published in Transactions on Machine Learning Research (11/2024)
networks may also promote grokking (Minegishi et al., 2023). Moreover, our work highlights that in modular arithmetic such sparse representations are obtained interpretably through the discrete Fourier transform.
Mechanistic Interpretability While training neural networks is often accompanied by mysterious phe- nomena such as double descent (Nakkiran et al., 2019), many works along the mechanistic interpretability have attempted to systematically understand what happened during training and inference through extensive reverse engineering (Olah al., 2020; Olsson et al., 2022; Akyürek al., 2023; Elhage et al., 2022; Notsawo et al., 2023). Paying attention to the activation of neurons, those studies have tried to identify the functional modules or circuits inside neural networks (Elhage et al., 2021; Commy et al., 2023). Even for recent large language models, controlling activation patterns via activation patching can unveil the role of cach module (Vig et al., 2020; Meng et al., 2023; Zhang & Nanda, 2024). In grokking literature, several works have revealed what kind of algorithmic pattern was obtained inside the model when it worked on modular addition (Zhong et al., 2023; Nanda et al., 2023; Morwani et al., 2023) or group composition (Chughtai et al., 2023; Stander et al., 2023) through the Fourier transform of logits or investigating gradients. Gromov (2023) points out that the learned weights and algorithms in some arithmetic tasks are analytically solvable if ML. uses a quadratic activation. In contrast, we extend the analysis from addition to other modular operations, such as subtraction, multiplication, polynomials, and a multi-task mixture, which can bridge the gap between simple synthetic data from modular addition and complex structured data as seen in the real world.
3 Preliminaries
Grokking This paper focuses on grokking on the classification tasks from simple algorithmic data commonly invostigated in the literature (Power et al., 2022; Liu et al., 2022; Barak et al., 2022). We have train and test datasets (Strin, Stest) without overlap, and learn a neural network f(æ; 0) where input x is a feature vector of elements in the underlying algorithm space for synthetic data and 0 are weights of neural network. The small-size Transformers (e.g. one or two layers) or MLP are usually adopted as f. Specifically, prior works train the network using stochastic gradient decent over the cross-entropy loss L and weight decay:
0 <- argmin E. ((. [C[((;))) 0) + 101|
where y E 1 1} is a scalar class label (p is a number of classes) correspond to the inputs e {0, (O,.....-- P- 1) x {0, ( 1}, and X is a hyper-parameter controlling the regularization. The weight decay is one of the factors inducing the grokking phenomenon (Power et al., 2022; Liu et al., 2023a; Varma et al., 2023). We employ AdamW (Loshchilov & Hutter, 2019)) as an optimizer in practice. The fraction of training data from all the combinations is defined as:
Semail
S Strain
Suain
+ Seest
p²
It has been observed that a larger fraction tends to facilitate fast model grokking, whereas smaller fraction makes grokking more challenging and slow especially in complex settings such as modular polynomial tasks.
Transformers As discussed in Elhage et al. (2021), the functionality of a small-size Transformer can be written down with several distinctive matrices. We denote embedding weights as WF E RRdems xp, output weights at the last MLP block as Wont E Xddon, and unembedding weights as Wy E Rpxddmmb, The logit vector on inputs a, b can be approximately written with activatations from MLP block, MLP(a, b), Logits(a, b) z Wo W ouut MLP b) by ignoring residual connection (Nanda et al., 2023), and we investigate the ncuron-logit map Wt- = Wy E Roxdmip in the later analysis. See Appendix A for further details.
A Analysis in Modular Addition Nanda et al. (2023) first pointed out that Transformer uses particular Fourier components and trigonometric identities after grokking occurred in modular addition 2. The modular addition a basic mathematical operation as (a b) % p = e where a, b, e are integers. The model predicts c given a pair of a and b. As a slightly abused notation, a, b, may represent one-hot representation, and we
2Zhong et al. (2023) revealed the existence of algorithms independent of trigonometric identities if the model does not have attention. Since we employ a Transformer with attention, this paper assumes the algorithms based on trigonometric identities.
3

---

Published in Transactions on Machine Learning Research (11/2024)
will omit % p in later sections. In the case of modular addition, the way Transformer model represents the task has been well-studied (Zhong et al., 2023; Nanda et al., 2023), where the embedding matrix WE maps the imput one-hot vectors into cosine and sine functions for various frequencies W/ 26m k E {0, 1}, such cos(wra), sin(waa). It is also known that the addition is implemented inside the Transformer with trigonometric identities.
cos(ws(a + b)) = cos(waa) cos((u) sin(waa) sin (), sin(wi(( + b)) = sin(waa) cos(wxb) + cos(wia) sin (w,,)),
and then the neuron-logit map Wt reads off cos(wx + b- c)) by also using trigonometric identities,
cos(wn + c)) = cos(wee( + b)) cos( + sin(ws (a + b)) sin(w,.).
(1)
The logits of e are the weighted sum of cos(uz + bb- - c)) over k. Note that we only consider the first half of frequencies (i.e. & E (1...., [2])) because of the symmetry. We show the example Python code for Fourier analysis in Appendix B.
Experimental Setup In this paper, we expand the discussion above on modular addition to entire modular arithmetic: a b % p c where represents arbitrary operations (or polynomials) that take two integers and b as inputs, such as - (subtract), * b (multtpplication), 2a b, ab + b, +62 a + ab, (a 6)4 (polynomials) 3, Transformer takes three one-hot tokens as inputs, a, o, b. In addition to integer tokens, we prepare Rop special tokens representing the mathematical operations above. The models are trained to predict c as an output.
Inpur
Inout
Inpur
Emnbed
Embeed
Entted
f
Tenkner
Unanbed
Unented
Logit
Lagil
Logit
Our neural network composed of a singlc-layer causal Transformer architecture (Figure 1) with learnable embedding and unembedding (demb = 128). We use ReLU for the activation functions and remove positional embedding, layer normalization, and bias terms for all the layers. We initialize each weight from Gaussian distribution, where the mean is 0 and the standard deviation is Vdt.: former is trained via full batch gradient descent with AdamW (Loshchilov & Hutter, 2019) and weight decay X 1.0. We use 97 for all the experiments. For the dataset faction, we use = 0.3 unless otherwisc mentioned. We conduct the experiments with 3 random seeds and report the average*. Other hyper- pappee-erraeterss are described in Appendix C.
Figure 1: Grokking has been inves- tigated with training from scratch. shed light on the dynamics in- side Transformer, we introduce the notion of pre-grokked models, which pre-trained on a similar task un- til grokking and used to replace ran- domly initialized modules without any parameter updates (i.e. frozen). We pre-grokked embedding and Transformer in the later section.
1 (LeCun et al., 2012). This Trans-
4 Pre-Grokked Models and Fourier Metrics
In contrast to modular addition, the exact analysis of internal circuits across other modular arithmetic would be challenging, since not all the operations have analytical algorithms. To mitigate such interpretability issues, we introduce the notion of pre-grokked models, and propose a pair of novel progress measures for grokking in modular arithmetic; Fowrier Frequency Densily (FFD) and Fourier Coesficient Ralio (FCR), which are derived from our empirical observation sparsity and sinusoidal bias in embedding and neuron-logit map.
Pre-Grokked Models To dive into the internal dynamics, we leverage pre-grokked models, which are pre- trained on similar algorithmic lasks until grokking and used for another training to replace randomly initialized modules without any parameter updates (i.e. frozen). This allows us to consider learning representations and algorithms separately. We will use pre-grokked embedding (freezing Wemb) and Transformer (freezing all the weights except for W andbandWu) in later sections. We also use the same random seed between pre-gros toogrrkkinng and downstream grokking experiments to reduce identifiability issues (Singh et. al., 2024).
We omit the discussion on modular division, since it requires division cases while we also consider multi-task mixture. 4 We optimize the madels up 3e5 gradient steps and then judge whether grokking happens or not.
4

---

Published in Transactions on Machine Learning Research (11/2024)
1.00
6.79
575
0.75
2.75
0.75
0.75
1.50
0.23
6.00
101
107
0.00
0.20
1.00
1.00
2.30
0.75
1.75
0.73
0.73
0.75
0.50
0.50
0.50
100
30
p-
0.79
0.75
0.75
0.75
8.53
0.50
0.53
0.50
Opimization Stees
Optiizationn Steps
Opiimoation Stepe
Optimization Steps
Opiimization Stepe
Optnn on Steps
Scratch
Pre-Grokked Embedding
Pre: Grokked Transtomer
0.3
05
Figure 2: Test accuracy in modular elementary arithmetic (addition, subtraction, and multiplication) with pre-grokked models (embedding and Transformer). The x-axis is the logarithmic scale. Because of the task simplicity, grokking always occurs in elementary arithmetic, However, in certain combinations, pre-groke pre-grokked models hinder grokking even with 0.9 fraction. For pre-grokked embedding, addition and subtraction accelerate grokking each other (fig(0:2, 0:2]), while multiplicattin and those do not show synergy (+: fig[2, 0] and [0. 2), -: fig[2, 1] and [1, 2]). In contrast, for pre-grokked Transformer, subtraction is challenging in both directions, even transferring subtraction models into subtraction itself (fig(1, 4)). With small addition and multiplication accelerate each other (fig(0, 5) and [2, 3)).
Fourier Frequency Density (FFD) FFD quantitatively measures the sparsity of Fourier components in certain layer (embedding neuron-logit map),
H
| ll2 max |all|
| ll2
K. = 2 2 &
Σn
> 1
+
>
max 12l12
where uk E f = {U1, pk..... is a coefficient of cosine components (uu
-== W [t] cos (wyt); W[t] is a
t-th index of weight matrix W) and VA E v is a cocfficient of sine components (Di- with frequency Wk. We set 77 0.5. The low FFD indicates that few key frequencies are dominant in the Fourier domain, which can be often observed in modular addition.
W[t] sin (wxt))
Fourier Coefficient Ratio (FCR) FCR quantifies the sinusoidal bias of Fourier components in a certain weight matrix,
1
|| |
| ll2
FC
min
|
The low FCR means that Fourier representations of the weights have either cosine- or sine-biased components, which can be often observed in modular multiplication.
The decrease of either FFD or FCR (or both) correlates to the progress of grokking, and the responsible indicator varies for each modular operation; for instance, FFD a good measure for addition, and FCR is for multiplication. They are not only aligned with the late improvement in test accuracy but also can characterize each Fourier representation of modular operations a certain layer (Section 6.3).
5 Analysis in Elementary Arithmetic
We start with the analysis with internal circuits of pro-grokked models, which can reveal the characteristics of each arithmetic operation; if pre-grokked embedding encourages grokking in downstream tasks, the learned embedding should be similar, but if pre-grokked embedding prevents grokking, those tasks may require
5

---

Published in Transactions on Machine Learning Research (11/2024)
Fourier Components Embedding Matrix
Fourier Components of Neuron-Logit Map
coMerms of Logits 20 Fourier Basis
Nom Forter component
Nom Poutter Componett
ann
cos
cos
5in
14
5th
32
50
30
00
sn1s
03
50
00
5n
000.0
Coms:
Norm of Foutte Commonent
Nor of Componert
-6
003.5
2.1
14
40
27
41
sin 40
coe
os
58 20
06 23
27
32
58 36
06 41
Norm ofmorro Foufer Comporent
3.50
form Foufeer Comparent
Consi
002.5
225
1.00
27
2.00
2.25
cos 41
Prequency k
Frequency &
24
4m 28
Component
Figure 3: Frequency analysis in grokking with elementary arithmetic. Subtraction learns similar embedding addition with sparse Fourier components (fig(0, 0] and fig[1, 0J). However, it imposes an asymmetric neuron-logit map and norm of logits with cosine biases (fig[1, 1] and fig[1, 2). Multiplication obtains quite different embedding from others (fig(2, :)); it employs all the frequencies equally with cosine bias for both embedding and neuron-logit map.
p-
poool
(and p-
fred ym
(and p=1
1.00
0:75
0.75
0.75
0.35
0.50
0:25
0.25
5.25
0.25
0.5
6.00
0.02
Optmizziion Stess
0.00
Optmiiztioon
5.00
Optimization
JUl
0.00
Coumizzdon
0.00
Optniizztin
121
+
Figure& 4: Test accuracy in modular polynomials (univariate terms: a² + a² + a + 2b, the degree-1 with cross term: + a b). Grokking occurs even quadratic cubic expressions asymmetric with input and b.
different lypes of representations. Moreover, if pre-grokked Transformer accelerates generalization, it means the algorithms obtained internally would have similar properties, while the failure may hint at algorithmic differences.
Figure shows Lest accuracy in elementary arithmetic (addition, subtraction, and multiplication) with pre-grokked embedding and Transformer". Because of the task simplicity, grokking always occurs among those operations. However, in certain combinations, pro-grokked models hinder grokking even with r = 0.9 fraction. For pre-grokked embedding, modular addition and subtraction accelerate grokking (Figure 2(0:2, 0:2]), while modular multiplication and those two hurt the performances each other (+: Figure 2(2, 0] and [0, 2, Figure 2/2, 1] and [1, 2)). In contrast, for pre-gre rre-grokked Transformer, modular subtraction is challenging in both directions, even transferring subtraction models into subtraction itself (Figure 2 4)). Pre-grokked Transformer on addition or multiplication accelerates each other (Figure 2(0, 5] and [2, 3)). Those results might imply that (1) while there a similarity between the Iearned embeddings in addition and subtraction, their acquired algorithms significantly differ (Section 5.1), and that (2) multiplication requires representations independent of addition subtraction but the algorithm might be transferable (Section 5.2).
STo avoid the confusion, we will mention the sub-figures using phthonic coordinates like Figureft, i] column

---

Published in Transactions on Machine Learning Research (11/2024)
5.1 Modular Subtraction Imposes Strong Asymmetry
Considering the sign in trigonometric identities, Transformers should learn modular subtraction in the Fourier domain with trigonometric identities as the case of addition (Bquation 1):
cos(wi(a- b c)) = cos(wx(a - b)) cos(w)c) + sin(wx(a - b)) sin(wac),
and then would anticipate similar interpretable representations to addition. However, observe that the grokked models exhibit asymmetric properties for both embedding and Transformer. We transform the embedding into a Fourier domain along the input dimension and compute the L, norm along other dimensions. In Figure 3, subtraction learns similar embedding to addition with sparse Fourier components (Figure 3(0, and [1, 0j). On the other hand, it imposes an asymmetric neuron-logit map and norms of logits with cosine- biased components (Figure 3[1, 1] and [L, 2)), which may represent alternatings (a b a).
Such asymmetry is also observed in grokked Transformers. The pre-grokked Transformer subtraction could not be transferred to any downstream elementary arithmetic (Figure 2 =)), even subtraction itself (Figure 2[1, 4)), and pre-grokked models with addition or multiplication could not learn subtraction as well (Figure 2[:, 4]). This implics that while we could interpret subtraction as a part addition with negative numbers, it is possible that the embedding and algorithm inside Transformer are quite different.
Lastly, we examine the restricted loss and ablated loss in Appendix E, where the restricted loss is calculated only with the Fourier components of significant frequencies, and the ablated loss is calculated by removing certain frequency from the logits. The analysis emphasizes the subtle dependency other frequencies than significant ones.
5.2 Modular Multiplication Leverages All Frequencies
In contrast to modular addition and subtraction, we do not attempt to fully interpret the learned model in a closed form or form of pscudocode. However, following the empirical analysis in modular addition, we can observe that multiplication also leverages the periodicity in the Fourier domain.
Figure 3 reveals that multiplication obtains significantly different Fourier representation from addition or subtraction (Figure 3[2, :)); it employs all the frequencies equally with cosine bias for both embedding and neuron-logit map. Surprisingly, miltipicatiion-re-rrokeee (Figure 2[2, 3)) and addition--re--rrokee Transformer (Figure 2(0, 5]) causes grokking in multiplication. This implies that in contrast to the asymmetry of subtraction, addition, and multiplication leverage their symmetry in the operations. Since the embedding of multiplication quitc different from addition and subtraction, it seems to be reasonable to fail to grok with ieaiiiionnuubbaatttionn -raeriion-pp-ggrokeedd embeddings (Figure 2(0:2, 2] and [2, 0:2]). Morcover, we find that grokking in elementary arithmetic occurs even with frozen random embedding (see Appendix F; sampled from the same random seed) that does not have biased components nor sparsity, which also supports that some unique, non- non-transfed no-transffeable patterns learned in grokked models.
Transformer accelerates grokking in addition
6 Analysis in Polynomials
It has been known that grokking would be less likely to occur increasing the complexity of operators in general (Power et al., 2022), but the underlying reasons or conditions are still unclear. In addition to elementary operations, we examine the interpretable patterns of grokked models in modular polynomials. We first investigate the case simple polynomials (Section 6.1), quadratic, cubic, and quartic expressions (Section 6.2).
6.1 Polynomials Discover Superposition of Representations for Elementary Arithmetic
We here investigate the relatively simple polynomials that induce grokking (univariate terms: a² + 62, a² b, a3 + 2b, the degreo-1 with cross term: ab + + b). In Figure 4, grokking occurs even quadratic or cubic expressions asymmetric with input a and b, and suggests that the existence of symmetry or the cross term might be key for occurrence.
7

---

Published in Transactions on Machine Learning Research (11/2024)
Fourier Components of Embedding Matrix
Fourier Components Neuron-Logit Map
Nargne Logits in Fourier Basis
Norts of Componon
oM
of Componost
g12.5
oss
35.0
a.n
24
12.5
18.0
7.5
sin
500 52
9.5
0.0
sin 36
2.5
0.0
0.0
Comt
0055
4m
639
58 18
58:40
Comsonan. a.........
Nomm bo Companan
an1
500 14
sin
to-
sin 30
41
06
cG 14
2n 18
23
32
36
254 43
Const
Nomer Foulty Gomooyent
a4 Nam hanter Somoo
34
t00 23
cos 41
sin 45
Prequency
Frequency
Corat
to
smi8
Compoment Figure 5: Frequency analysis in grokking with modular polynomials (a² + 6², a² b, ab + b). Grokking discovers the superposition of frequency sparsity and bias seen in elementary arithmetic; a² inherits both biased sparsity in subtraction and significant cosine biases multiplication for embedding (fig(1,0).. Its neuron-logit map leverages addition-like sparsity (f1g(1.1)).
1.05 a! 2' (nei
1.30a eab- M (aod p- 90
+ 11PP---00 0m- (red you
(e & Geed yi
Ge- (rod 10)
1.00
1.00
0.75
0.00
0.25
0:23
0.23
5.23
0.25
0.5
0.00
-P7
5.03
427
0.00
0.00
F.s. p000)
+2 1k (aed p-
1.50
(a 1- Traal p-1
1.00
(red p-1
3.50
(6 Cood
1.00
6--8/ (1oody (0)
6.75
0.75
0.75
5.75
0.75
0.75
0.23
0.5
0.00
mittion Steps
0.03
Optimizstion Steps
Optimleation Steps
0.00
AuT
Steps
0.00
Optimization
Figure Test accuracy modular polynomials with quadratic, cubic, and quartic formulas. Compared to the elementary arithmetic (Figure 2), these requires longer optimization steps. Transformers suffer from late generalization in degree-n polynomials with cross-term (a² + + b² a² + ab + 5² + aa, a² + ab, a" ab² + b). If polynomials are factorizable with addition (a + b) subtraction (a 6), they are easy to grok (e.g. (a+ + a + b; fig[0[4]] although they also have a cross term (c.f. a² ab + b²). Even, cubic ((a + b) fig [[[2:4] or quartic ((a = b)*; fig[1][4:)) expressions, grokking occurs if they are factorizable.
Moreover, the grokked models exhibit partially-similar internal states the one in elementary arithmetic. Figure 3 provides frequency analysis with modular polynomials (a² + b?, a² - b, ab + b), where grokking discovers superposition of representations (frequency sparsity and bias) for elementary arithmetic. For instance, a² 2 52 finds cosine-biased embedding like multiplication and a sparse neuron-logit map like addition. a² inherits both biased sparsity in subtraction and significant cosine biases in mulliplication for embedding. Its nouron-logit map leverages addition like sparsity. ab + a a is similar to multiplication; leveraging biased all the froquencies while using sine components, because can be factorized as (a + 1)(b + 1) These trends are flipped between embedding and neuron-logit map. Norms of logits in 2D Fourier basis basically follow the trend in multiplication (Figure 5(;,2)), and especially a² b activates key frequency columns (Figure 5[1,2))
8

---

Published in Transactions on Machine Learning Research (11/2024)
Foprier Components Embedding Matrix
Fourier Componnnnts of Neuron-Logit Map
c.termss of Logits Fourier Basis
d' Norm Fourle Component
cod
Norm Fout Ff Comp
Sin
14
co5
45
40
00 14
006 23
s/n
005
Comis
Norm of Euutee Compunett
-0 Nem of routee Componant
(o- 11-
sin
41
45
033 34
an 28
055 25
20 27
32
20.36
cos 4)
10
Nem Eeattr component
Num baatee Component
34
by
605 32
5h 36
50.5
50
50
Comt
Jahh 14
20
di 36
Frequency
Frequency
Component
Figure 7: Frequency analysis in grokking with factorizable polynomials. Non- von-factorizable operation (a² ab fig(0, :]) cannot find sparse embedding representations. contrast, factorization with elementary arithmetic accelerates grokking in both quadratic ((a 6)², fig[1, :I) and cubic expression ((a + b)², fig[2, :I) with sparse Fourier features.
FFD
FFD
0.50
0.00
D1R22 D1022 Optimization STADS
20130
A0110
5000
Optniiztion stess
opamiaution Staps
1PP 10" ES 10
Opcim Staps
Accuracy
FFD/FRR
n
mm4,5
Figure 8: FFD and FCR progress measure of grokking. progress of grokking synchronizing with the lest accuracy improvement. The responsible indicator depends on each operation. See Appendix G for the details.
The decrease of either FFD FCR (or both) indicates the
6.2 High-Degree Factorization Allows Grokking
Increasing the complexity of operators, formulas in Figure 6. Apparently, Transformer fails to generalize in degree-n polynomials with cross term (a2+ab ab +6, ab +a, a, a² + ab, a² + at² b). However, if polynomials are factorizable as a product of addition or subtraction (in the form of b)"), they easily grok, even when they also have cross terms (c.g., (a b)2 + + b). Besides, the sum of powers also easy to be grokked. Even, cubic (Figure 6[1, 2:4) or quartic (Figure 6[1, 4:l) expressions, grokking occurs they are factorizable. Comparing a² + ab + b2 and (a + 6)2 a² ab 12 + and (a + b)2 + emphasizes the importance of factorizability for the emergence of grokking.
test modular polynomials with quadratic, cubic, and quartic
Figure analyzes the frequency components in factorizable polynomials. Non-factorizable operation (a² ab + (2) cannol. find the sparse embedding representation. In contrast, factorizable operations promote grokking in both quadratic ((a + 6)2) and cubic expression ((a + 6)3) obtaining sparsity embedding. The factorizable operations find more biased Fourier components than the non-factorizable ones the neuron-logit map. Morcover, factorizable polynomials exhibit clear logits patterns shown in elementary arithmetic (Figure 3), while non-factorivable ones only show significant. norms around constant component.

---

Published in Transactions on Machine Learning Research (11/2024)
20
p-1
1 pm
-. (mnd
3-3-- (mzé--11
0.75
035
0.59
0.90
550
0.50
030
0.2
0.0
5.b
0.33
025
0.00
20
20"
0.03
10
0.30
Optimization Steps
5.0
30"
Steps
12
0.00
301
12
Steps
SP
0.00
12t
12
Scratch
Pre-Grokked Embedding
Pre-Grokked Transformer
P-0.5
7-077
Tigure 9: Test accuracy in modular linear expression (2a 1 2a 1: 36) and degree-1 polynomials with cross term (ab + b). Pre-grokked erubedding in modular addition accelerates grokking in 2a + b, 2a + 36, and pre- Transformer modular multiplication accelerates grokking in ab = b, while the training from scratch could not generalize in 0.3.
6.3 FFD and FCR as Progress Measures
As shown in Figure 8, we measure FFD and FCR in embedding layer WE for various modular operations. See Appendix for the results in neuron-logit map Wr.
Elementary Arithmetic Addition (red) and subtraction (blue) decrease FFD and keep a high FCR, whereas multiplication maintains FFD as 1.0 and decrcascs FCR (groen). In all the cases, the saturation of accuracy and inflection point of either FFD or FCR almost match (verlical lines). Interestingly, ab + b (purple) exhibits decreasing both FFD and FCR, which reflects the feature of addition and multiplication simultancously.
Sum Powers In a" +b", FFD and FCR exhibit the same progress as multiplication, while the neuron-logit map has sparsity the same addition (Appendix G). We also observe different behaviors depending on the parity of exponent n; FFD decreases more when n is odd (blue) and FCR drops more when n is even (red).
Factorizable Polynomials (a + 6)a exhibits the same trend as addition: high sparsity and balanced components. In contrast, the neuron-logit map behaves similarly to multiplication (Appendix G). As in the sum of powers, the dynamics would be different depending on the parity of exponent FCR significantly drops when n is even. In the case factoriiiii a² + ab + b², FFD do not change during training, and the model cannot achieve late generalization.
7 Analysis in Transferability
Since all the modular arithmotic has periodicity, we could hypothosize that grokked models obtain common features among similar operations (transferability). Furthermore, pre-grokked models in a certain lask could promote grokking in other similar tasks because they already have useful basis. We first test the transferability of pre-grokked models from elementary arithmetic to linear expressions (Section 7.1), and then extensively investigate it with higher-order polynomials (Section 7.2).
7.1 Tro-krokked Models Accelerate Grokking in Linear Expression
We test whether frozen pre-grokked modules in elementary arithmetic (a + b, a b) are transferable grokking in modular linear expression (2n + b, 2a + 3b, ab + b). Those asymmetric expressions hard to grok from scratch, especially if the fraction is small (r 0.3) despite their simplicity. Figure 9 shows that pre-grokked embedding with addition accelerates grokking in 2a + b, 2a + 3b, and pro-grokked Transformer with multiplication does in ab + b. These support our hypothesis and imply that in complex operations, internal circuits struggle with finding interpretable patterns.
7.2 Pre-Grokked Models May not Help Higher-Order Polynomials
In Section 7.1, we demonstrate that pro-grokked models accelcrate grokking in lincar expressions. We here extensively test. pre-grokked models in higher-order polynomials (quadratic and cubic). Table 1 shows that pre-grokked models could not accelerate, and they even prevent grokking in higher-order polynomials, which implics that pre-grokked models may not always help grokking accelerations, except for lincar expressions. While the learned representation of polynomials seems be a superposition of that of elementary arithmetic (e.g. Section 6.1), their functionalities might differ significantly.
10

---

Published in Transactions on Machine Learning Research (11/2024)
Addition (o I 6)
Multiplication (a+ 6)
Subtraction (a 6)
Downstream Op.
PG-E
PG-T
PC-E
PG-T
PC-E
PG-T
From Seratch
2a +
0.4
e- 0.7
r 0.5
0.5
r-
2a 36
04
0.4
2a 36
T= 0.4
0.8
0=0.1
0.7
0.5
r=04
r= 0.1
0.7
0
+
0.9
(a 6)a
= 0.9
== 0.8
0.4
+ ab
0.9
0.7
F= 0.9
n 0.8
a²
0.6
0.9
0.4
0.6
0.7
Y
0.6
= 0.5
F= 0.7
F 0.4
(a+ 6)
0.6
(a Dj
0.4
0.6
0.5
+b
0=
at
Table 1: Summary grokked modular operators with pre-grokked models (both embbedding and Transformer). We provide the smallest train fraction where grokking happens. PG-E/T stands for pre-grokked embeedingg/Tannstrmeen Transsorrrrf The shaded ones are the results presented in Figure 9.
2.0
4+ (ssod 27)
(annd p 97)
a. Smod p 97)
0.8
0.0
0.2
0.2
0.2
10"
10:
10
1e
20
15a
10a
122
10²
10a
10
301
10"
127
Optimozation Staps
Optimization Steps
Optimization Steps
Fourler Components Embedding Matrix
Fourler Componants Neuron-Logit Map
JWaarnn Logits Fourier Basis
Norm Fodling Companene.
of Componote
15.0
car
457
12.5
10.0
7.5
ain
2.5
cou
35 Frequency h.
Frequency
Compoment
0.3
0.5
0.7
0.9
Figure 10: Test accuracy and frequency analysis in grokking with a mixture of elementary arithmetic. Co-grokking across different operations occurs, but it needs larger fraction than single task (r = 0.3 does not work).
These ablation studies reveal that the transferability of pre-grokked embeddings and models is limited to specific combinations, such as from elementary arithmetic to linear expressions, and could be rarely observed in higher-degree expressions. From the transferability of learned representation perspective, we should note that there is still an analysis gap between the grokking with synthetic data and common machine learning.
8 Analysis in Multi- Task Training
While previous works on grokking have only dealt with single task during training, the application of Transformers such as large language models (Brown et al., 2020) is usually trained on mixture of various tasks or datasets. Given the periodicity and similarity across entire modular arithmetic, we also hypothesize that mixing functionally similar operations in the dataset promotes grokking. To fill the gap between synthetic tasks and practice, we here investigate grokking on mixed datascts with addition, subtraction, and multiplication (Section 8.1). We also study multi-task, training mixing hard and easy polynomial operations (Section 8.2). We prepare 0.3 datasets and jointly train Transformers their mixture.
11

---

Published in Transactions on Machine Learning Research (11/2024)
eb- 8 (moo - 97)
a² ab+ (mod p p
2.00
+ (mod p 97)
a² + + (mod 96)
640-
0.30
Baa4-
0.00
2040
0.40
320
10
10
Opcinizatio
2
16
19'
Ostmmiation strps
23"
Optimiation Sid5s
20 20
optiicaionn SINES
SingleTeek
Must-Sask
F-43 -5 F-47
Singeeisk
Mulli-Takk
Figure 11: (Left) Test accuracy in grokking with a mixture of modular polynomials (fa+b ab bb and {a +b² a² + ab b2, (a + b)))). The multi-task training across similar operations promotes grokking. (Right) Test accuracy in grokking with a mixture of modular polyuomials ([(a + b)* a² + ab} and {(a 6)", ,0, + ab² 6)). The multi-task training across similar operations promotes the improvement of test accuracy.
8.1 Multi- Mixture Discovers Coexisting Solutions
Figure 10 reveals that co-grokking (i.e. grokking happens for all the tasks) occurs, but it requires a larger fraction of train datasct than a single task; for instance, r r= 0.3 could not cause grokking while it does in Figure 2. The Lest. accuracy of multiplication increases slower than the other Uwo, which implies the conflict. among different Fourier representations may affect the performance and generalization.
For the Fourier analysis of grokked models, training with a multi-task mixture seems to discover "Pareto- optimal" representations for all the operations in embedding and neuron-logit map (Figure 10(1, We can see the coexistence of component sparsity in embedding (addition), asymmetric cosine sparsity in neuron-logit map (subtraction), and cosine-biased components for the frequencies (multiplication). Furthermore, the norms of logits in 2D Fourier basis for addition and subtraction exhibit the same patterns. This means that addition and subtraction can be expressed on the same representation space originally, while they find quite different grokked models after the single-task training.
8.2 Proper Multi-Task Mixture also Accelerates Grokking in Polynomials
We also investigate the multi-task training with the mixture of polynomials; preparing the combination of easy and hard operations as {a + b, ab + b}, {a² + 12, a² + ab + b2, (a (a+ b) {(a + b)², as ab} and {(a + 6)3,0 a² + ab² + b]. As shown in Figure 11 (left), a proper mixture of polynomials, in terms of operation similarity, also accelerates grokking in multi-task settings. For instance, a2² and (a+b)² help generalization in a² + ab + b² This implics that the required representations among {a² + +²,,aa a² + ab + 6², (a (a+b)3 b)²} would be the same while original single-task a² + ab + 12 fails to grok due to the difficulty non-factorizable cross term. The test accuracy also improves in the cubic expression (Figure 11, right). However, hits a plateau before the perfect generalization.
The results imply that some multi-task mixtures may lead to co-grokking and accelerate generalization while others may not find optimal solutions. It would be an interesting future direction to further reveal the grokking dynamics and mechanism for multi-task training.
9 Conclusion
Our empirical analysis has shed light on significant differences in internal circuits and grokking dynamics across modular arithmetic. The learned representations are distinct from each other depending on the type of mathematical expressions. and despite the periodicity of modular arithmetic itself, the distinctive Fourier representations are only obtained in the operations that cause grokking. While grokking with complex synthetic data, we find that not all the insights related to the nature models. For instance, the ablation with frozen pre-grokked modules demonstrates that the transferability could only be limited to the specific combination of modular operations. The functional similarity between the mathematical expressions may not help. In addition, some multi-operation mixtures may lead to co-grokking and even promote generalization while others might not reach optimal solutions. We hope our extensive empirical analysis encourages the community to further bridge the gap between simple synthetic data and the data where analytical solutions are not attainable for better understanding of grokked internal circuits.
also happen in practical
12

---

Published in Transactions on Machine Learning Research (11/2024)
Limitation We have observed all modular arithmetic operations that can cause grokking have shown interpretable trends with the Fourier basis. However, except for a few cases, we may not. derive exact algorithms. It remains as future works to derive the approximate solutions covering the other modular operations beyond addition remain as future works. We also have cxamined a broader range of complex modular arithmetic than prior works and oblained some implications to bridge the analysis gaps between the synthetic and practical settings. However, our observations imply that the mechanism of grokking might not always share the underlying dynamics with common machine learning. Further investigations of internal circuits in practical models such as LLMs are important future directions.
Acknowledgments
We thank Tadashi Kozuno for helpful discussion on this work. HF was supported by JSPS KAKENHI Grant Number JP22J21582.
References
Ekin Akyürck, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations, 2023.
Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, eran malach, and Cyril Zhang, Hidden progress deep learning: SGD learns parities near the computational limit. In Advances in Newral Information Processing Systems, 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Tlya Sutskever, and Dario Amodei. Language models are few-shot learners. arXi preprint ar Xrv::20005.1 14165, 2020.
Voss,,
Bilal Chughtai, Lawrence Chan, and Neel Nanda. toy model of universality: Reverse engineering how networks learn group operations. arXic preprint orXiv:2302.0300 03025, 2023.
Arthur Commy, Augustine N. Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alons Towards automated circuit discovery for mechanistic interpretability. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
Xander Davies, Tauro Langosco, and David Krueger. Unilying grokking and double descent. arXio preprinl WrXiu:2303.061733 2023.
Darshil Doshi, Aritra Das, Tianyu He, and Andrey Gromov. To grok or not to grok: Disentangling generalization and memorization on corrupted algorithmic datasets. arXiv preprint arXiv:2310. 130061 2023.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya. Sanyal, Scan Welleck, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers on compositionality. ar preprinl arziv:2305. 18654, 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 HalfielddDodds, 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 Circuils Threud, 2021.
(0021//raawoor//iiii hinddeiiidxxxxiee html.
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravoc, Zac Hatfield-Dodee Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of superposition, 2022.
13

---

Published in Transactions on Machine Learning Research (11/2024)
Hiroki Furuta, Yutaka Matsuo, Aleksandra Faust, and Izzeddin Gur. Exposing limitations of language model agents in sequential-task compositions the web. arXie preprinl or orXiv:2311.1871. 2023.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.0267 2023.
Tanishq Kumar, Blake Bordclon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. arXiv preprint 0rXii:2310.06110, 2023.
Yann A. LeCun, Léon Bottou, Genevieve B. Orr, and Glaus-Roberr Müller. Efficient BackProp, pp. 9-48. Springer Berlin Heidelberg, 2012.
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arzio:2307.0338/ 2023.
Noam Levi, Alon Beck, and Yohai Bar-Sinai. Grokking linear estimators solvable model that groks without understanding. arXiv preprint arXio:2310. 1114:1230.11441, 2023.
Ziming Liu, Ouail Kitouni, Niklas Nolte, Fric J. Michand, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. arXie preprint or Xiv:2220 10343, 2022.
Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In
International Conference on Learning Representations, 2023a.
Ziming Liu, Ziqian Zhong, and Max Tegmark. Grokking compression: A nonlinear complexity perspective. arXiv preprint arXic:2310.0520 2023b.
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
Kaiseng Lyu, Jikai Jin, Zhiyuan Li, Simon Du, Jason D. Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provably induce grokking. arXiv preprint Xr11:2311.18817, 2023.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. arviv preprint arXiv:2202.0533 052662, 2023.
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking competition of sparse and dense subnetworks. arXie preprint arXiw:2303.11875 2023.
Jack Miller, Charles O'Neill, and Thang Bui. Grokking beyond neural networks: An empirical exploration with model complexity. arXiv preprint arXiv:2310. 2023.
Gouki Mincgishi, Yusuke Iwasawa, and Yutaka Matsuo. Bridging Lottery ticket and Grokking: Is weight norm sufficient explain delayed generalization?, 2023.
Depen Morwani, Benjamin L. Edelman, Costin-Andrci Oncescu, Rosie Zhao, and Sham Kakade. Feature
emergence via margin maximization: case studies in algebraic tasks. arXio preprinl arXiv:2311.07568, 2023.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher D. Manning. Grokking of hierarchical structure in vanilla transformers. arXio preprint arXw:22305.189 18741, 2023.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Tlya Sutskever. Deep double descent: Where bigger models and more data hurt. arXiv preprint arXi::1912.00220 2019.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In International Conference on Learning Representations, 2023.
Pascal Jr. Tikeng Notsawo, Hattie Zhou, Mohammad Pezeshki, Irina Rish, and Guillaume Dumas. Prodicting grokking long before it happens: A look into the loss landscape of models which grok. arXiv preprint arXiv:2306. 01003006.1253... 2023.
14

---

Published in Transactions on Machine Learning Research (11/2024)
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 10.23915/distill /iii/ssiiiil.000200011 https://dissilll.pub//// http://///////itt//////s/s//a/// cireuitsszzommiin
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 Hatficld-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitl, 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. https://trre /ttra//asooommccciriitttt. wttp:/iruuisiuuupp//e/ee pub/2022 ii-coonext-- annuaiinnuuctooo-eeaeass
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint 20X212201.02177, 2022.
Noa Rubin, Inbar Seroussi, and Zohar Ringel. Droplets of good representations: Grokking as a first order phase transition in two layer networks. arXiv preprint orXio:2310.03788) 2023.
Aaditya K. Singh, Ted Moskovitz, Felix Hill, Stephanie C. Y. Chan, and Andrew M. Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. arXiv preprint 0rziv:2404.07122 2024.
Dashiell Stander, Qinan Yu, Honglu Fan, and Stella Biderman. Grokking group multiplication with cosets. arXiv preprint SEft:22112065888 2023.
Zhiquan Tan and Weiran Huang. Understanding grokking through robustness viewpoint. arxiv preprint XNiv:3311.06591 2023.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. ar Xio preprinl wrXiiv2206.04811 2022.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXin:2309. 020000, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv:
2017.
Jesse Vig, Scbastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. Investigating gender bias in language models using causal mediation analysis. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, 2020.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunoi Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. arXiv preprint Niw:2206.0885. 2022.
Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi, and Wei Hu Benign overfitting and grokking in relu networks for xor cluster data. ar preprint orXiv:2310.025./ 2023.
Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metries and methods. Xiv preprint arXiv:2309.1000 16042, 2024.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. In Neural Information Processing Systems, 2023.
Bojan Žunkovič and Enej llievski. Grokking phase transitions in learning local rules with gradient descent. arXiv preprint arxiv:2210.11 15435, 2022.
15

---

Published in Transactions on Machine Learning Research (11/2024)
Appendix
A Mathematical Description of Transformer
In this section, we describe the structure of causal Transformer in our work, loosely following the notation of Elhage et al. (2021).
As defined in Section 3, we define embedding matrix as WE, query, key, and value matrices of j-th head in the attention layer W W K: Wyy. The input and output layer at the MLP block is denoted as Win, Wout, and the unembedding matrix is denoted as We. We use ReLU for the activation functions and remove positional embedding, layer normalization, and bias terms for all the layers.
We also denote the token (one-hot representation of integers) in position as 4, the initial residual stream
i-th token zi (0) causal attention scores from the last tokens (12, because the context length is 3) to all previous tokens al j-th head A', the attention output at j-th head W the residual stream after
the attention layer on the final token as x(1), the neuron activations in the MLP block "MLP", and the final residual stream the final token as (2), "Logits" represents the logits on the final token since only consider the loss from it.
We can formalize the logit calculation via the following equations.
Embedding:
T2 Wgt
Attention score: A softmax( (x) (0) W T )
Attention block: (1) = 22 + Ej ((
MLP activations: MLP ReLU(W)n)())
MLP block: x²2) W WoP + x())
Logits: (2x
(0) Note that these focus on the operations for the representation from the final token 22 and the above reflects the causal modeling. Following the discussion in Nanda et al. (2023), we ignore the residual connection and investigate the neuron logit map W1 = Wy Wout as a dominant part to decide the logits.
16

---

Published in Transactions on Machine Learning Research (11/2024)
B Example Python Code for Discrete Fourier Transform
In this section, we provide the example Python code to analyze the weights with discrete Fourier transform, 88 done in Section 5 and 6.
impuzt turch Tibrariee
impuzi
Iaport pardas at p4
Define ussful functions der to. expps (taneor, Flac=Palee):
if (tenxe))) terch. TanKer:
taseor
flattes
neapy
neapy
15
18
(tanzor)
17
artiadie
18
grid (g.ogrid[suplo (oayppppooyolllioio Coloco are.ahage))]
18
op.eeetty shage (nsti), rompes...oolt..pp... (orr. 1893)
20
cffset +1
21
range (6):
23
act (.r. Seoffees] grid(s) cieoffst] are
24
slape
columns-f "veloe (=
30
31
32
57
33
sodal Transforear
31
35
Cospats Pouries busis
36
Jouseer
basis basis
37
fourier, baris. (tors, (pppppppppp..ppt..())) (p))
(angel, p//2 +1):
tocrier_bazia spper sirther aras
(tarn (-2).aore
pl (arrr.aarrrngggp))))))))))))))))
43
stoci
44
45 Extract enbedcing veigits from Teansforeer
48 N.E aebel. abbe.. K.Hl. sts_op)
47 Retract the neiron-logts sap saighte fros Transforear
48 Koont blccss [0] .elp Moont
49 R. medel.ameebed M.D.S sin_opl.r
50 K.L V.5 4,00t
51 grbui
'sie', (coz")
Appy discreie Pourive trassiora vabvlding
to stcc.ats (t (orttr_eerss.tt(:::20))
58 Unt-Teigl] " aapaess laboleeb))))))
'cos's (Tdr 58 tabal f'aiats veiss (atii
Ponter
80
61
62 Jppy discreie Fouriur trassfora
neuron legit
63 fouriar_eabad_ens
(fourier, bariz e aore torch torch.stacz (tt.://seee.bed.aaaiiaaiie
84 cos_ais, cos_ais_rsbsd_ous
eoudser, (2::2]1
65 of snlntrar,ein. aebed aabed_ait)
66 [f_ons()Trig]] of (100))))pppp))))) (paaards x (rop__ahnte(s)))
67 labal nore Fonter cosponants
out Cuin's oo eeeeelll UTog Jlelsiaa), vo Prelav Idf, Urrii (3
17

---

Published in Transactions on Machine Learning Research (11/2024)
C
Experimental Details
We summarize the prpee-parameters for the experiments (dimension in Transformers, optimizers, etc.) in Table 2. We provide the code in supplementary material.
Name
Value
Mod
97
Epochs
1e6
Optimizer
AdamW (Loshchilov Hutter, 2019)
Learning Rate
0.001
AdamW Betas
(0.9, 0.98)
Weight Decay
1.0
Batch Size
(Full batch)
Max Optimization Steps
3e5
Number of Seeds
3
Embedding Dimension demb
128
MLP Dimension dmip
512
Number of Heads
4
Head Dimension
32
Number of Layers Activation
ReLU
Layer Normalization
False
Bias Term Weight Matrix
False
Weight Initialization
\00 Vil.at
Vocabulary Size p' Context Length
+ Rop (including operation tokens)
Table 2: Hyper-par appe-paaameters for the grokking experiments. We follow the previous works (Power et al., 2022; Nanda et al., 2023; Zhong et al., 2023).
D Terminology for Mathematical Expressions
As reference, summarize the terminology for mathematical expressions in Table 3.
Term
Expressions
Modular Arithmetic
(a) % ==c
Addition
++b
Subtraction
Multiplication
a
Elementary Arithmetic
all the above (+, -,*)
Polynomials
a² + a ab. (a b).... (including all the below)
Linear Expression (degree-1)
2a b, 2a + 3b, ab + b...
Cross Term
ab, at?
Quadratic Expression (degree-2)
(a - 6)3 1),,a² + ab, a²
Cubic Expression (degree-3)
(a + b)3 ....
Quartic Expression (degree-4)
(a b)+ b)+...
Factorizable Polynomials
(a + b)".
(a + 6)nh + (aa bb (n 2,3..... <))
Polynomials with Cross Term Non-Paactoiiab Polynomials)
a² + ab 12, at? +b...
Sum of Powers
aa 6 =n=
Table 3: Terminology for mathematical expressions in this paper.
18

---

Published in Transactions on Machine Learning Research (11/2024)
E
Analysis of Restricted Loss in Modular Subtraction
In Figure 12, we test the restricted loss and ablated loss, the metrics proposed by Nanda et al. (2023), where the restricted loss is calculated only with the Fourier components of key frequencies, and the ablated loss is calculated by removing a certain frequency from the logits. The results show that modular subtraction has several dependent frequencies, which cause worse restricted loss if ablated, while they are not key frequencies (we set the threshold to AC le - 9). Those dependent frequencies are not observed in modular addition. Moreover, the restricted loss for modular subtraction significantly gets worse than the original loss, which also emphasizes the subtle dependency
other frequency components.
Moreover, we extensively evaluate the relationships between loss and Fourier components. We here decompose the logits as follows:
Logits = (Logits from key frequencies) + (Logits from non-key frequencies) + (Logits from residuals),
where logits from residuals are estimated by subtracting logits of all the frequencies from the raw logits.
The results are presented in Table 4. In modular addition, we find that key frequencies contribute to the prediction and non-key frequencies only have a negligible effect on the loss (e.g. train loss v.s. ablation (d), restricted loss v.s. ablation (c)). The residuals actually hinder prediction accuracy (e.g., train loss v.s. ablation (c)). In modular subtraction, any ablations drop the performance and all the components contribute the predictions, which implics that the grokked models in modular subtraction have informative representations to some degree over all the frequencies, even residuals in the logits.
Ablating Requunciiss (Addition)
Ablating Frequencies ction)
4/
Other Freq Figure 12: Loss of Transformer when ablating each frequency (k 18) and everything except for the key frequencies (restricted loss). In modular subtraction, we find several dependent frequencies (orange), which cause worse restricted loss if ablated while they are not key frequencies.
Key Freq
Dependent Freq
Restricted Loss
Original Loss
Logits
Loss (4)
Key Freq.
Non-key Freq.
Residuals
Add (+)
Sub (-)
Train Loss Restricted Loss
1.008e-7
1.336e-7
4.985e-8
7.141c-1
Ablation
4.576
7.741
(b)
5.385
2.179e+1
(c)
4.989c-8
5,582c-1
(d)
1.015c-7
5.348c-6
5.383
2.188e+1
Table 4: Loss of Transformer when ablating the components of key frequencies, non-key frequencies, and residuals, from the logits.
19

---

Published in Transactions on Machine Learning Research (11/2024)
F Grokking with Frozen Random Embedding
We here show that even if the sparsity and non-trivial biases are not realizable in embedding, grokking could occur in Figure 13. In this experiment, we initialize embedding weights from Gaussian distribution, where the mean is 0 and the standard deviation is vdout (LeCun et al., 2012), and then freeze them not allowing any parameter updates during training. Even with the restricted capacity, the models exhibit grokking in elementary arithmetic.
Fourier Components Embedding Matrix
Fourier Components Neuron-Logit Map
(mod p 97)
1.0
Norn of Faurier Componens
Component
1.0
80
0.8
a
0.8
15
30
Accuracy
0.6
0.6
0.4
10
0.2
0.0
0
10
20
30
40
50
0
10
20
30
40
50
10a
10a
10t
Optimization Steps
(mod pm 97)
1.0
Norm of Faurier Compenent
Fourler Component
80
1.0
70
0.8
60
0.8
50
Accurasy
0.6
0.6
80
30
0.4
20
0.2
0
10
30
40
50
100
30
40
50
10a
10
10t
Optimization Steps
a (mod p=97] 97]
1.0
Norm of Foutter Component
1.2
Norm Fourter Component
12
10
Accuracy
de
8
0.2
0.0
0
10
20 30
40
50
0
10
20
30
40
50
10-
10a
10*
Frequency
Frequency
Optimization Steps
T=
0.3
r= 0=5
T= 0.7
r= 0.9
Figure 13: Test accuracy with elementary arithmetic from random embedding. Grokking can occur even using frozen random embedding, while unmbbddding obtains similar Fourier representation as discussed Section
20

---

Published in Transactions on Machine Learning Research (11/2024)
G FFD and FCR in Neuron-Logit Map
Figure 14 presents our progress measures: FFD and FCR in neuron-logit map WT. For clementary arithmetic operators, the dynamics seem to be the same as seen in embedding (Figure 8). This might be due to the similarity of embedding and neuron-logit map (Figure 3). For sum of powers (a" + b*) and the factorizable b)") behaves differently from embedding (Figure 8). The sum of powers decreases FFD while keeping FCR relatively higher. The factorizable polynomials maintain both FFD and FCR relatively higher. This might. be due to the representation asymmelry between embedding and neuron-logit map in polynomials (Figure 7).
/0
+b
FCR
FFD a"++""
(c
(o=A)*
1.00
100
1.00
1,50
1.00
1.00
0.50
0.50
3.50
0.50
0.50
0.50
0.25
0.25
0.25
25
10000
30000
30000
15000
20000
30000
10
30
10:
107
10
10
10
10
10
10
10'
10
Optimization
Optimization
Test Loss
FFD/FCR
m=4,5
6,6,7
Figure 14: FFD and FCR in neuron-logit map for each operation (a + b, a b, a + b. + 6", (a (a+ b)").
Detailed Analysis on FFD and FCR
In this section, we provide experiments to mcasure the FCR and FFD with different train datasct fractions r or threshold y,. As shown in Figure 15 (above), we change the dataset ratio, the needed optimization steps to be grokked are also changed, and the inflection point of the corresponding progress measure is changed too. For instance, when change 0.3 with 0.5 in + and - b, the decrcase of FFD is also accclerated along with the increase of test accuracy. We also observe that if the grokking does not happen, the progress measure does not exhibit the change, such as the case of FFD in ab - b with 0.3.
As for different threshold 7, Figure 16 demonstrates that lower delays the decrease of the metrics in contrast to the progress of grokking; for instance, the decrease the metrics starts after test accuracy reaches 1.0 (in the case + b, b, ab + b). On the other hand, larger v) results in too fast convergence. Based on those observations, we have set 77 0.5 in this paper.
FFD
FFD
FFD(.))
FFD
1.0
04
0.3
6:0
00
Optimization Steps
Optimization Steps FCRlo 6)
Optimization Steps FCR(ol)
Optimization Steps
FCRto 6)
FCRlab 6)
61
04
622
132
Optimization Staps
Optimization Steps
Optimizatin Steps
Optiiization Steps
FFD/FCR
Test Accuracy
0.6
Figure 15: FFD and FCR in embedding for each operation (a + b, a b,a b, ab b) with different r.
21

---

Published in Transactions on Machine Learning Research (11/2024)
FFDXe 6)
FFDle 6)
FFD(e))
FFD +1)
04-
Optimization Steps
Optimization Steps
Optimization Steps
Optimization Steps
Test Accuracy
Figure FFD in embedding for each operation (a b, b, a ab 6) with different threshold n
Summary of Grokked Modular Operators
Table 5 summarizes if cach modular operator can cause grokking in y= 0.3 or not. We provide the best test accuracy if they do not grok.
Elementary Arithmetic
Lineur Expeession
Praction
a
24 +.b
a-t-b-9 20
2a =
6 b- Za
2a
a 0- 26
2a
a 0- 24
2.556
3.355
3.735
0.4
9.0%
0.9
Cross Them (Degren-) 1)
Univariate Terms
a
0.3
6.156
5.0%
9.0%
0.736
7
0.8
Cross Term (Degree-n)
Sum of Powers
Praction
c² ob.bbbb
ab 6 +
a² ab
o 42 +-6
6"--8
6'4.8
4'
d'
2'48
d'
y 0.3
34%
4,8%
4.9%
4.00%
47%
8.256
9.4%
0.5
56%
10%
11%
10%
0.7
74%
1736
14%
13%
1250
16%
15%
0756
18%
Pactorizable
Fraction
ts + 8:2
(n++A) +-+
[o - 8/3
(o + A2
(o AY
(n +-*
(a-M) b)*
(n Ay*
(n (-*
(n b1
0.4
12%
91%6
Y-
9150
9256
0.7
0.9
Table 5: Summary of grokked modular operators tested in this paper (p = 97). We provide the best test accuracy if the operator does not cause grokking.
22

---

Published in Transactions on Machine Learning Research (11/2024)
J Analysis on Grokking with Different Modulo
J.1 Periodic Patterns in Fourier Components are Common Characteristics
We here provide the ablation study with different modulo p 59, 113 (from Figure 17 to Figure 21). These results show that, basically, our findings and observed trends in the main text (done with = 97) can be seen in different modulo p as well. The periodic patterns in Fourier components can be common characteristics
among different p in most cases.
o.6 (szed pm 59
laaod pm
eeb Irned p E9)
of. 5 (zzd p 593
(saod ym E1)
ob- (mnd pm
1.00
1.00
1.00
5.75
3.50
0.50
0.50
3.50
0.25
0.25
0.25
0.25
0.25
0.25
102
1.00
0.75
0.75
3.50
0.50
0.5
8.00
20
122
10"
10 0.00
10
107
307
10
0.00
20
10
101
107 Stops
10
00
10
101
20
00
20
101
1.2
207
10
00
10
up
2pt
107
Opdmizazion
0ptimization
Optimizatinn Stops
Optima ation Stops
Opāmization Steps
0.3
r-
6.5
0.7
r-
Figure 17: Test accuracy with different modulo p (59, 113} (a a b,
b²
b, ab + b).
Fourier Components Embedding
Fourier Components Neuron-Logit currrr Logits
17.5
fouggr Componer
12.5
of foudg
30.0
ofrmol
5.0
23
sin 23
Conss
003
40
20 34
Canat-
horm fogter Comporeng
Norm of Fourdd Componetes
004
600 23
27
25
Cons
033
40
2
dee Compontes
Burm Founa Companare
1.00
14
1..5
ch 18
0.25
0.00
10
13
20
25
Conm
Frequency
Frequency
003
Component
Figure 18: Frequency analysis in grokking with modulo, 59 (a
b.
b.
23

---

Published in Transactions on Machine Learning Research (11/2024)
Fourler Components of Embedding Marrix
Fourier Components of Neuron- ogt Map
Harms of Log logits in 20 Fourier Basis
Compotent
000
-in
cou
ro
5n
40
00 23
560 27
500
Const
Norg of Cemponent
0o5
الليسسس
cou
sin
500
kn
00225
20
2n336
temm beiier Compatent
ced hom of fogriee Component
sin
cos
ain
cou
b4
m
40 Frequancy
20
Frequancy
sin
000
the
Com
Figure 19: Frequency analysis in grokking with modulo 113 (a
Pourier Components of Embedding Matrix
Fourier Components of Neuron Logit Map
Morms Logits in 2D Fourier Basis
Norm Tee der Componeet
2.
Norm Componeent
[Componens oonene
C0G 14
18
23
0.0
10
15
20
23
30
10
15
20
com
000
23
Const-
Comporant
Norm of Founle Componett
15.0
L.Coppunion cot
30.0
K00 14
18
20
20
com
055
cos 27
Constr
Merm Teuntr component
Norm Funner Comporent
Nori of Fourter Component
000
E0G 14
8.6
18
10
15
20
25
30
10
15
Com
0555
Frequency
Frequency
ommpoter
Figure 20: Frequency analysis in grokking with modulo 59 (a² 6.
b
ab
+ b).
24

---

Published in Transactions on Machine Learning Research (11/2024)
Fourier Components of Embedding Matrix
Fourier Components of Neuror-Legit Map
Narzps Logits in Fourier Besis
Nom Fouter Componeee
37.55
35.0
12.5
c05 23
12.00
9th 27
055 52
082
kn77
Norm Fourter Componatt
Norn of bouter Component
20
15
006 53
C0 21
27
200
Const
Congoren
horm of or Compurane
9.0.6
0.4
55
cos 50
Frequency 8
Frequency
Component
Figure 21: Frequency analysis in grokking with modulo p 113 (a² + 62
b, 6).
J.2 Grokking Can be Function of Modulo p
In addition mathematical operation and dataset fraction, grokking can a function of modulo p. Figure shows that p = 97 causes grokking with a² + ab, while p 59 and p 113 do not. Surprisingly, p 59 has fewer combinations than 97, but p 59 does not generalize to the test set even with = 0.9. The results suggest that we might need to care about the choice of p for grokking analysis.
a ab (mod p 59,91 113)
1.00
Accuracy
0.75
0.50
0.25
0.00
10"
10
102
10a
10*
Optimization Steps
5 59 (r= 0.9)
p= 97 (r 0.9)
p= (r (r= 0.9)
Figure 22: Test accuracy in grokking with 8 + ab (r 0.9). 97 only causes grokking among (59, 97, 113).
25

---

Published in Transactions on Machine Learning Research (11/2024)
K Dataset Distribution Does not Have Significant Effects
One possible hypothesis why some modular polynomials are hard to generalize is that some polynomials bias the label distribution in the dataset. To examine this hypothesis, we calculate several statistics on label distribution in the dataset. We first randomly split train and test dataset (r= 0.3), and get categorical label distributions. We compute the KL divergence between train label distribution drrain and test label distribution deest. train label entropy, and test label entropy, averaging them with 100 random seeds.
Figure 23 shows KT. divergence between train and test datasets (top), train dataset entropy (middle), and test dataset entropy (bottom). While those values slightly differ across the operations, there are significant difference between generalizable (e.g. a² + bi; a² + b²) and non-generalizable (e.g. a² + ab, a² + ab + b²) polynomials despite their similarity. The results do not imply that dataset distribution has significant impacts on grokking.
(DKL (darain dress) + DEa (drest (reann)))
0.00
drenin log darain
0.04
0.02
0.00
dhea log dhest
0.04
0.02
0.00
1.11000 qo
9 qo
qp
or.
96 02
d0 nlq m)
m)
6l m)
60 + o)
( o)
qo + m
2f a
+ P
ofe o)
alq o)
.e o)
9
9 0t-000 +
59
c0+ 9o
Figure 23: KT, divergence between train and test datasets (top), train dataset entropy (middle), and test dataset.
entropy (bottom).
L Extended Limitation
Our work extends the grokking analysis from simple modular addition to complex modular polynomials. However, those tasks are still synthetic and far from LLMs (Brown et al., 2020), the most popular application of Transformers. Connecting grokking phenomena or mechanistic interpretability analysis into the emergent capability (Wei et al., 2022), or limitation in compositional generalization (Dziri et al., 2023; Furuta et al., 2023) and arithmetic (Lee et al., 2023) would be interesting future directions.
26

---

Published in Transactions on Machine Learning Research (11/2024)
M Initialization and Pre-Grokked Models
To clarify that our experimental obscrvations do not trivially come from random initialization, we provide a detailed version of Figure in Figure 24, which has a visualization of standard deviation among three random seeds (shadow area). These results show that the overlap is small and thus our observations are consistent and do not come from the initialization or variance.
0. 5- a40 (med 991
0.7
2=0.
o-B-vaad (and 90)
c-5 -aa- (mod 91
a (mod 92)
Acoursck
Acoura
0.50
0.50
0:50
0.25
0.25
0:25
0.25
0.00
0.00
0.00
0.00
10"
102
102
10t
107
10"
201
10'
10"
10"
107
10"
10t
10
- +
(and p 90]
C-6
6+40 (mad 9-99)
c-5 - 8 (mod
97)
4-5
-- (mod
97
0.75
0.75
0.75
Socursc
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.00
10'
10
16
10
10
107
10
10
10
10
0.00
10a
10
10"
10
10"
10
10
10
10t
10
- a. (and p- 07)
1.00
5 0.0
(nnd
p
1.00
6 (mod
97)
3
0.0
p 97)
1.00
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
10
10a
12
10
10
107
36
10"
12
10
0.00
10
10"
10"
12r
13
10
10"
10"
(mod
871
1.00
p
1.00
(wed p
97)
(mod
1.00
1.0
20.75
Acturasy
0.75
0.75
Acuraay
0.50
0.50
0.25
0.25
107
10
10
107
10
107
10
10
10
10
10
10a
102
107
107
107
Scratch
Pre-Grokked Embedding
Pre-Grokked Transformer
0.3
0.5
0.7
0.9
Figure 24: Test accuracy in modular elementary arithmetic (addition and multiplication) with pre-grokked models (embedding and Transformer) trained with subtraction. In addition to the plot in Figure 2, we visualize the standard deviation (shadow area).
N Identifiability and Pre-Grokked Models
Through the experiments with pre-grokked models, we observe that some pre-grokked models prevent. grokking in certain combinations. However, Singh et al. (2024) point out that freezing a subset of models may suffer from identifiability issues, which may not prompt the downstream performance. To mitigate the potential issues, our experiments use the same random seed between pre-grokking and downstream grokking experiments. Moreover, we observe that the combinations that cannot accelerate grokking have quite different. Fourier components (for instance, modular addition and multiplication seen in Figure 3) in the weight matrix, while the rotating operations to the weight matrix do not change the Fourier basis. We think that our claims are not just due to the ident ifiabiliy issues.
27


Reviews:
[{"review_id": "r9krQJrxRr", "reviewer": "Reviewer_Sx7C", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "(numbered for clarity)\n\n# Strengths\n\n1. The paper contains a large array of experiments across a range of modular arithmetic tasks with clear results figures and summaries (I especially enjoyed table 1 and Appendix H to be very instructive/a nice summary)\n2. The insights drawn are interesting (though not fully supported – see weaknesses) and thought stimulating. Specifically, considering the usefulness of shared representations across tasks may provide an interesting angle on analogs of grokking in larger models, which are likely closer to the \"multi-task training\" setting.\n3. Appendices are thoughtful and helpful:\n\n      a. Appendix F is cool. It seems in line with Zhong et al. (2023)'s findings that not all models learn circular embeddings. Any intuitions as to what's going on there?\n\n      b. I like Appendix J – it's a nice confound to check.\n\n# Weaknesses\n\n1. (big) In Section 5, the analysis methodology relies on the statement \"if pre-grokked embedding encourages grokking in downstream tasks, the learned embedding should be similar, but if not, those tasks should require different types of representations.\" While the former is true, the latter isn't necessarily. Just because pre-grokked embeddings don't accelerate learning doesn't mean those tasks require different types of representations. The methodology for restoring pre-grokked embeddings involves fixing weights of a submodule (e.g., $W_{emb}$). However, this methodology suffers from the identifiability issue – namely, a rotation could be applied to all matrices. Prior work [2] has demonstrated the large effect this can have on dynamics (see Appendix D.1.1 of [2]), so I would be wary of negative claims (if grokking is discouraged, different representations are required) that rely on weight fixing (the positive claims are fine). Many of the main claims of the paper require on this negative \"pre-grokked didn't help\" argument.\n2. (big) While the authors conduct many experiments, most conclusions revolve around grokking timing (with comparisons between timing for \"from scratch\" and \"pre-grokked\" runs being the main methodology) are only done on a single initialization. How do timings vary across initializations? Are the effects consistent, or can they be explained by this variance?\n3. I would be careful with statements like \"However, if polynomials are factorizable with addition (subtraction) or are the sum of powers, they easily grok, although they also have cross terms\" (Section 6.2) – the results only show that the quick grokking runs are for factorizable polynomials, not the other way around? To test the provided statement, one would want to explicitly control for \"factorizability\". E.g., we saw that $(a^2 + b^2)$ in Figure 4 can be grokked. Does this imply that $(a+b)(a^2 + b^2)$ can always be grokked?\n4. In Section 4, \"The decrease of either FFS or FCR (or both) indicates the progress of grokking, and the responsible indicator depends on each modular operation\" – could this statement be justified/is it obvious?\n5. In Figure 2, \"Addition and multiplication accelerate each other (fig[0, 5] and [2, 3]).\" I'm not sure if fig[0,5] actually shows this? At least at high values of $r$ from scratch seems to learn way quicker/there generally seems to be a lot of variance so I wouldn't necessarily say addition accelerates multiplication. Also, generally, it would be important to see if these results transfer across initializations for the claims being made/how much noise there is in the timing of grokking.\n\n[1] Shalizi (2024). \"Attention\", \"Transformers\", in Neural Network \"Large Language Models\". http://bactra.org/notebooks/nn-attention-and-transformers.html#identification\n\n[2] Singh et al., (2024). What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formation. https://arxiv.org/abs/2404.07129", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present an extensive analysis of grokking in various modular arithmetic functions of the form `f(a,b) = c (mod p)`. They go beyond prior work by exploring more complicated functions f (subtraction, multiplication, polynomials). Their central claims revolve around how representations from pretraining one task may enhance grokking on other tasks. They apply this methodology and find some positive transfer, indicating that common (periodic) representations are useful. They also introduce two novel progress measures (explored in Figure 8).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "(numbered for clarity)\n\n1. (critical) Overall, I would suggest the authors strengthen the claims made in Section 5 and 6 with a more rigorous experimental methodology. Namely, they should describe the effect of initialization on timing (given how central timing is to their claims). They should also consider the identifiability/rotation issue when restoring-and-freezing pre-trained weights (another key aspect of the methodology). These are both good methodology ideas, they just need to be made rigorous.\n\n2. (critical) Please define in the main text what you mean by pre-grokked Embedding and pre-grokked Transformer. I think the former is intuitive. The latter I'm still not clear on which matrices are being set and frozen.\n\n3. (important) In Section 3 \"Analysis in Modular Addition\", the authors portray the \"clock\" algorithm (to use the term from Zhong et al. (2023)) as the way modular addition is implemented. This portrayal lacks the nuance that Zhong et al. (2023) introduced – specifically, that other algorithms may emerge (e.g., the \"pizza\"), and more importantly, not all models learn circular embeddings (which seem to be a central claim of this paper). I would strongly suggest revising this introduction to incorporate the nuance from prior work.\n\n4. (important) In Section 4, I found the definition of FFS and FCR a bit hard to follow – namely, I would suggest the authors explicitly mathematically define the calculation of $\\mu_k, \\nu_k$ for a given weight matrix (e.g., $W_{emb}$).\n\nMinor nits (MN):\n\n1. The paper emphasizes performing over 1000 experiments. While impressive, experiments alone do not increase the significance of the work/I would suggest softening the use of this statement.\n2. Page 2 first full paragraph: When introducing FFS and FCR, add abbreviations in parentheses before using (e.g., \"Fourier Frequency Sparsity (FFS)\").\n3. Section 3, I think the formula for $r$ assumes binary tasks (two inputs), which hasn't been specified yet – I assume that's why $|S_{train}| + |S_{test}| = p^2$\n4. Page 7, I would be wary of the boldness of using the word \"causes\" in  \"addition-pre-grokked Transformer (Figure 2[0, 5]) causes grokking in multiplication.\" This isn't true at all values of $r$, results are shown for a single seed, and one could only claim that addition-pre-grokked *speeds up* grokking.\n5. Figure 6 x-axis is a substantially different scale than Figure 2. This isn't an issue, I would just mention this in the caption since the x-tick-labels are quite small/hard to read.\n6. What are the colors in Figure 16?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "mjhD7tepyD", "reviewer": "Reviewer_YGon", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "On the positive side, the paper can be seen as a fairly comprehensive study on the internal representations of simple transformers trained on modular arithmetic tasks. This significantly extends beyond the scope of the early work of Nanda et al and Zhong et al. \n\nThe paper covers a wide range of experiments and settings. Most details are provided. \n\nThe technical sections of the paper are mostly well-written. \n\nTo my knowledge, (1) the main observations on the properties of the Fourier spectrum of learned embeddings (symmetry vs asymmetry, transferability etc), and (2) the notion of 2 measures of progress are both novel.\n\n\n \n\nIn terms of weakness, I find the early section (including abstract) hard to parse sometimes. I provide detailed comments below, and I do believe the paper, as it currently is, requires revision before publication.\n\nIt is somewhat underwhelming that the paper provides no attempt to fully reverse engineer a learned model on modular arithmetic (beyond addition) in an end to end fashion, similar to the prior work. While the experiments cover a wide range of settings, I find the paper lacks technical depth.\n\nI find the two notions of progress measure a bit arbitrary. It is unclear whether they are causally responsible for grokking. Looking at the second picture of Figure 8, for example, the green line shows that the FCR measure continues to drop even when the model has fully learned to generalize.  Similarly, in the blue line, the measure almost plateaus throughout, though the accuracy goes up with a phase transition. \n\nFinally, the paper studies p=97 only.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies grokking behavior of training transformers to solve modular operations, including addition, subtraction and multiplication. \n\nFirst, the paper finds that all the modular operations exhibit grokking in training. Second, they provide analysis on the Fourier representations within the models weights, extending the early work of Nanda et al on modular addition. Third, the paper proposes two new measures of hidden progress (based on Fourier coefficients) that seem to correlate with the occurrence of grokking. Finally, the paper also studies transferability of pre-trained embeddings and multi-task training.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The second part of the abstract, along with the introduction, reads vague at times. Multiple terminologies are undefined. It’s unclear what the main message is. For example:\n\n>  superposition of the patterns from elementary arithmetic\n\nWhat does this really mean?\n\n> challenging cases\n\nWhat defines easy vs hard in this setting? \n\n> the ablation with frozen pre-grokked modules\n\nIt should say: our ablation study on the pre-grokked models reveals that …. Also, it’s unclear what transferability means here. \n\n> co-grokking\n\nI am pretty sure this is not a standard terminology. If it’s a new phenomenon discovered by this paper, it’s worth expanding and highlighting what it is.\n\n> entire modular operations.\n\nI’d replace this to be “other modular operations (beyond addition)”.\n\n> through over a thousand experiments,\n\nI don’t know how this is counted. It’s not necessary to make such a claim here. It perfectly suffices to just say “through extensive experiments…”\n\n> while also investigating whether grokked models may exhibit transferability and scaling to the similarity and the number of tasks\n\nThis sentence sounds vague. Do you simply mean “whether grokked models may exhibit transferability to other tasks”? (Note that at this point of the paper, transferability hasn’t been defined.)\n\n> End of introduction: We provide significant insights in the empirical interpretation of internal circuits learned through modular polynomials, where analytical solutions are not attainable.\n\nHow is it connected to the co-grokking stuff you were talking about?\n\n> In contrast, we provide a detailed analysis of entire modular arithmetic, while extending the range of operations from addition to subtraction, multiplication, polynomials, and a multi-task mixture,\n\nYou can simply say, we extend the analysis from addition to ….\n\n> weight decay is one of the key factors inducing the grokking phenomenon\n\nCite Varma et al here.\n\n> Nanda et al. (2023) have first pointed out\n\nRemove “have”\n\nSec 3: also specify the initialization scheme here. The Omnigrok paper, for example, suggests that the scale of initialization matters here.\n\n> Sec 5.2:  we may not describe possible acquired algorithms for modular multiplication in a closed form, since trigonometric identities do not have multiplication formulas.\n\nThis requires some qualification. In principle, a full reverse engineering of the learned model, in the style of Nanda et al, is still conceivable. I would just say “we do not attempt to fully interp the learned model (in a closed form or pseudocode). However, we observe …”\n\n> Section F:  In this experiment, we initialize embedding weights from Gaussian distribution\n\nSpecify the scale here.\n\nFinally, I encourage the author(s)  the experiments for different values of p.", "broader_impact_concerns": "The paper does not raise ethical concerns, in my view.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "nFjhVtHVvn", "reviewer": "Reviewer_XqUP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Generally, I think the paper provides good empirical insights into grokking that build on the existing literature. I believe it fulfils both criteria of the TMLR review process. Thus, I recommend that the paper be accepted. \n\nGeneral strengths:\n1. I really like the empirical analysis completed in the paper. The introduction of new progress measures is always a welcome addition to the literature.\n2. I think the idea of ablating certain frequencies from the logits is quite insightful and could be used in further work.\n3. I like the observation that certain combinations of operations can either accelerate or inhibit grokking.\n4. I really like the discussion around degree-$n$ polynomials with the cross term. I do, however, have a question about this listed under *general critiques and questions*.\n5. It's interesting that co-grokking across different operations needs a larger fraction than training on a single task. I would not have necessarily expected this before reading the paper.\n\nSome general critiques and questions:\n1. I am unsure how to interpret the right-hand column of figures 3, 5, and 7. What are these providing to the paper? Perhaps they could be moved to the appendix.\n2. How does the choice of $\\eta$ influence the progress measures? Were there any ablations done with different values other than $0.5$?\n3. Concerning grokking in degree-$n$ polynomials, did the authors consider increasing the size of the network? Or with the given amount of data, will the models always fail to generalise for some polynomials?\n4. There seems to be a lack of theory as to why certain trends occur in the paper. For example, why are the identified progress measures connected to generalisation? This is not a reason to reject the paper, but it would have been nice to see!\n\nMinor gripes:\n1. I think the language could be improved. For example, in the sentence, \"Specifically, they train the network using stochastic gradient descent over the cross-entropy loss L and weight decay,\" it is unclear who \"they\" refers to.\n2. The use of the word \"inducing\" seems to be misapplied in the sentence, \"weight decay is one of the key factors inducing the grokking phenomenon.\" Grokking can occur without weight decay, so \"inducing\" may not be the best choice of word here. I can provide references for this if the authors need.\n3. There isn't any stated motivation for the choice of prime or the training fraction.\n4. What is meant by \"grokking discovers superposition of representations (frequency sparsity and bias) for elementary arithmetic\"? Is this relating to superposition in the [Toy Models of Superposition](https://transformer-circuits.pub/2022/toy_model/index.html) sense of the word?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors lay out three hypotheses:\n1. Any modular operations can be characterised with distinctive Fourier representations or internal circuits;\n2. Grokked models obtain common features transferable among similar operations;\n3. Mixing datasets with similar operations promotes grokking.\n\nThey then explore whether these hypotheses hold. In doing so, they find two progress measures for grokking: Fourier frequency sparsity and Fourier coefficient ratio. They also find that transferability is limited to specific combinations of operations, degree-$n$ polynomials sometimes do not grok, and that grokking relies on subtle frequencies within the logit domain.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have a few requested changes:\n1. It might be worth renaming FFS to be Fourier Frequency Density as it seems if the number increases, the spectrum is denser?\n2. I would be interested in a more rigorous statistical analysis of whether or not the progress measures actually do align with the point of grokking in the loss/accuracy. Perhaps one could vary the training ratio or another variable controlling the grokking gap in order to provide data for such a test?", "broader_impact_concerns": "I have no reason to be concerned regarding the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper expands the literature on grokking in small transformer models by considering broader set of operations (including more complex ones like polynomials), and evaluating transfer across tasks by reusing pretrained model components, as well as through multitask training. The authors identify various interesting findings in the different representation structures learned by different tasks (in the fourier basis) which relate to which tasks transfer well. Overall. the reviewers and I agree that the paper offers some interesting findings and supports its claims reasonably well. The main limitation, as with other related works, is that the focus on simple modular arithmetic tasks and small scale dataset and models necessarily limits our ability to draw more general conclusions.
Decision: Accept (poster)