$\mathbf{Li_2}$: A Framework on Dynamics of Feature Emergence and Delayed Generalization
Authors: Yuandong Tian
Venue: ICLR 2026 Poster (2026)
Keywords: grokking
gradient dynamics
generalization
memorization
modular addition
scaling laws
Primary area: interpretability and explainable AI
TL;DR: We study the gradient dynamics of grokking for 2-layer networks and proposes a mathematical framework to explain feature learning process.

Abstract: While the phenomenon of grokking, i.e., delayed generalization, has been studied extensively, it remains an open problem whether there is a mathematical framework that characterizes what kind of features will emerge, how and in which conditions it happens, and is still closely connected with the gradient dynamics of the training, for complex structured inputs. We propose a novel framework, named \ours{}, that captures three key stages for the grokking behavior of 2-layer nonlinear networks: (I) \underline{\textbf{L}}azy learning, (II) \underline{\textbf{i}}ndependent feature learning and (III) \underline{\textbf{i}}nteractive feature learning. At the lazy learning stage, top layer overfits to random hidden representation and the model appears to memorize. During lazy learning, the \emph{backpropagated gradient} $G_F$ from the top layer carries information about the target label, with a specific structure that enables each hidden node to learn their representation \emph{independently}. Interestingly, the independent dynamics follows exactly the \emph{gradient ascent} of an energy function $\mathcal{E}$, and its local maxima are precisely the emerging features. We study whether these local-optima induced features are generalizable, their representation power, and how they change on sample size, in group arithmetic tasks. When hidden nodes start to interact in the later stage of learning, we provably show how $G_F$ changes to focus on missing features that need to be learned. Our study sheds lights on roles played by key hyperparameters such as weight decay, learning rate and sample sizes in grokking, leads to provable scaling laws of feature emergence, memorization and generalization, and reveals the underlying cause why recent optimizers such as Muon can be effective, from the first principles of gradient dynamics. Our analysis can be extended to multi-layer architectures.

Published as a conference paper at ICLR 2026
PROVABLE SCALING LAWS OF FEATURE EMERGENCE FROM LEARNING DYNAMICS OF GROKKING
Yuandong Tian
Meta Superintelligence Labs (FAIR) & Independent Researcher yuandong. t tiant :: i com
ABSTRACT
While the phenomenon of grokking, i.e., delayed generalization, has been studied extensively, it remains an open problem whether there is a mathematical frame- work that characterizes what kind of features will emerge. how and in which con- ditions happens, and is still closely connected with the gradient dynamics of the training, for complex structured inputs. We propose novel framework, named Lig, that captures three key stages for the grokking behavior of 2-layer nonlinear networks: (!) Lazy learning, (II) independent feature learning and (III) interactive feature learning. At the lazy learning stage, top layer overfits to random hidden representation and the model appears memorize, and at the same time, the back- propagated gradient GF from the top layer carries information about the target label, with specific structure that enables each hidden node to learn their rep- resentation independently. Moreover, the independent dynamics follows exactly the gradient ascent of energy function E, and its local maxima are precisely the emerging features. We study whether these local-optima induced features are generalizable, their representation power, and how they change on sample size, in group arithmetic tasks. When hidden nodes start to interact in the later stage of learning, we provably show how GF changes to focus on missing features that need to be learned. Our study sheds lights on roles played by key hyperparameters such as weight decay, learning rate and sample sizes in grokking, leads to provable scaling laws of feature emergence, memorization and generalization, and reveals the underlying cause why recent optimizers such as Muon can be effective, from the first principles of gradient dynamics, Our analysis can be extended to multi- layer architectures. The code is available
INTRODUCTION
While modern deep models such as Transformers have achieved impressive empirical performance, it remains a mystery how such models acquire the knowledge during the training process. There have been ongoing arguments whether the models can truly generalize beyond what it is trained on, or just memorize the dataset and performs poorly in out-of-distribution (OOD) data (Wang et al., 2024b Chu al. 2025 Mirzadeh al. 2024).
Modeling the memorization and generalization behaviors have been a goal of many works. One such
behavior, know grokking (Power et al.) 2022 Doshi al. 2024 Nanda et al.) 2023 Wang al.)
2024a Varma et al. [2023] Liu et al. 2023 Thilak et al. 2022). shows that the model initially overfits to the training set, and then suddenly generalizes unseen test samples after continuous training.
Many explanation exists, e.g., effective theory (Liu et al. 2022 Clauw et al. 2024). efficiency of memorization and generalization circuits (Varma et al., 2023). Bayesian interpretation with weight decay as prior (Millidge 2022), etc. Most works focus on direct explanation of its empirical behaviors, or leveraging property of very wide networks (Barak et al. [2022] Mohamadi et al. [024]] Rubin et al. 22244), but few explores the details of the grokking learning procedure by studying the gradient dynamics on the weights.
In this work, we propose mathematical framework Lig that divides the grokking dynamics for 2-yeyer nonlinear networks into three major stages (Fig. D. Stage Lazy Learning: when training
httpp: http:////ccccccccoooo/o////o///e n/ss1 /real /real-dat aset/

---

Published as a conference paper at ICLR 2026
Memorization
Generalization
[Sharp p Optima)
Optina))
(l) Lazy Learning
[li] Independent
(00) Interactive
Initial Random Weights
(Ridge Reg resgrestionna on Rancimm Features
Feature Learning
Feature Learning
Noth energy (unction
Three stages of feature learning. In Stage I (Lazy learning), model first learns to overfit on random features provided by the bottom layer. For generalizable targets, the backpropagated gradients to the feature layer GF x YP TT becomes structured and signals task structures, triggering Stage (independent feature learning). Different hidden nodes independently learn to extract multiple task related features. In Stage III (interactive feature learning). these features interact and consolidate via backpropagated signal.
Figure 1:
begins, the top (output) layer learns first with random features from the hidden layer, and at the same time, the backpropagated gradient G to the hidden layer starts to carry useful signal about the target label (Proposition []. Stage II: Independent feature learning: GF drives the learning of hidden representations. In this stage, the backpropagated gradient of j-th neuron (node) only depends on its own activation, triggering independent feature learning for each node. Stage III: Interactive feature learning: When weights in the hidden layer get updated and are no longer independent, interactions across nodes adjust the learned feature to minimize the loss.
Each stage is studied in details. In Stage I, we show that the signal carried by the backpropagated gradient Gr either appears at the initial stage (Proposition I which depends the top layer ini- tialization, or the final stage (Lemma [ which depends on the weight decay 1). In Stage 11, we find that the dynamics that governs independent feature learning exactly follows the gradient ascent analysis
of an energy function & (Thm. [. which corresponds to nonlinear aanonical-correlan (CCA) between the input X and target Y. In group arithmetic tasks, we completely characterize the local maxima structure of & (Thm. 2). each representing one learned feature. We further show that the amount of training samples determine whether these features/local optima remain stable (Thm. D. whether they are generalizable or not, and thus determine the scaling laws of generaliza- tion and memorization. In Stage III, we provably show that there is a push for diversity between hidden nodes with similar activations (Thm. 6). when a subset of features are learned, backpropa- gated gradient FF changes to push the model to focus on missing features that need to be learned (Thm. 7. and optimizers like Muon (Jordan et al. 2024) further improves the diversity (Thm. 8)) Experiments on graph arithmetic tasks support our claims, e.g., the proved scaling laws about the generralizatiio/mmmmoiiaazzzz boundary (Thm. 4 fits well with the experiments (Fig. 5
Comparison with existing grokking frameworks. Our framework provides a theoretical founda- tion from first principles (i.e., gradient dynamics) that explains the empirical hypothesis Varma al. (023) that "generalization circuits Cgen is more efficient but learn slower than memorization cir- cuits Cmee" Specifically, we show that the data distribution determines the optimization landscape, which in turn governs which local optima the weights converge into, which lead to the behavior of memorization or generalization. We also show that the initial memorization, or lazy learning (Stage I), has to happen before feature learning (Stage II-III), since the former provides meaningful back- propagated gradient GF for the latter to start developing. In comparison, (Nanda et al., 2023) also provides a three stage framework of grokking, but mostly from empirical observations. Using our framework, we provide a systematic framework to explain when and why the grokking happens that are consistent with many empirical observations (Sec.
Our study goes beyond Neural Tangent Kernel (NTK) (Du et al. 2018 Jacot et al.] 2018 and mean field regime (Rubin et al.) 2024), in which the hidden weights do not update substantially from initializations. Our framework demonstrates concretely how features emerge from gradient dynamics and multiple key factors that affect the procedure.
Summary. Our framework Li2 splits the network training process into three stages (Lazy, independent, interactive), and provides insights into grokking dynamics and feature learning.
Explanation of Grokking behavior. At the beginning of grokking, the lazy learning cor- responds to memorization, in which the top layer finds a (temporary) solution to explain the target with the random features. Only after that, the backpropagated gradient GF be- comes meaningful, and causes the hidden layer to learn generalizable, emerging features. Grokking is accelerated with regularization.
Emerging features are the local maxima of an energy function E that governs the indepen- dent stage. These features are more efficient in label prediction than simple memorization.
2

---

Published as a conference paper at ICLR 2026
target?
larget P
target
Random
Vridge
V'udaated
upddatdd
Gy carries
Gr carniss
random
Go mindom
random
of
F upditted
D
Learned W
Complete W
[cl After Insependeent
(i) After Interactivv
(e] Initialization
(6) After Lazy learning
Festure Learning
Fruture learning
Figure 2: Left: The loss curves of training and test sets and their gap (generalization gap). Middle:
The, dynamics over different stages, stages, Right: Our theoretical results. For clarity, wtaagsssssssss........... is. plceedd as a coi
wiflle
concurren
Ja
sag
5.0
1.40
645-
1-000
3
260 250
400
110
25
Figure 3: Grokking dynamics on modular addition task with M 71, K = 2048, TL 2016 (40% training
out of 71 samples) with and without weight decay. Top: v = 0.0002 and grokking happens. Bottom: = 0 and no grokking happens. Weight decay leads to larger [Gr] around epoch 100 and induces grokking behavior. and 1, measured by cosine distance, shows two-stage behaviors: first there is huge update on the output weight V, then large update on the hidden weight W. Throughout the training, F f and P; FF remains diagonal with up 8% error, validating analysis (independent feature learning. Sec. J.. Experiments averaged over seeds.
The weights difference AW between consecutive weights at
Data governs the landscape of E. Sufficient training data maintain the shape of those generalizable local maxima, while insufficient data lead non-generalizable local maxima.
Scaling Laws of Feature Emergence, Generalization and Memorization can be derived by inspecting how the landscape changes with the data distribution.
2 PROBLEM FORMULATION
We consider a 2-layer network y =((W)) and la loss function on n samples:
min V.W
1 2
(Y 1) min V.W
1
1 I (Y -((W)
where Pt :- 1- /n is the zero-mean projection matrix along the sample dimension, Y gaxM is a label matrix (each row is a one-hot vector), X = [x1, X2,..., x,] E Raxd is the data matrix, V E IR RMM and W gdxk are the weight matrices of the last layer and hidden layer, respectively. a is the nonlinear activation function.
In the following, we show that grokking consequence of "leaked" backpropagated gradient Gp.
STAGE I: LAZY LEARNING (OVERFITTING)
Let F = (XW) be the activation of the hidden layer and P be the zero-mean version of it. Similarly define Y = P.Y. We first write down the backpropagated gradient Gp sent to the hidden layer:
as = P- - - -FVJT GF - OF
(2)
3.1 ANALYSIS OF THE BACKPROPAGATER GRADIENT GF
At the beginning of the training. both W and are initialized with independent vero-mean random variables. Therefore, the backpropagated gradient Gp is pure random noise. Over time, the hidden
3

---

Published as a conference paper at ICLR 2026
activation F is mostly unchanged, and only the output layer V learns. In this case, F can be treated as fixed during this stage of learning, and we can prove the following properties of GF (Sec. H): Proposition 1. If F is fixed and full column rank, entries of ((0) is initialized from normal distribution N NO, with 0 < a & L. then [[] F = D(evKM) and the backpropagated gradient Ge is dominated by the term YY F at initial time stamps:
Gr(() = V f - (aa) + (at) + O(())
(3)
and converges exponentially to the following fixed point when V = Vridge (f F + 111 T:
Gp(+oo)= = = +n11 "Y F(F F m1) m/1-
(4)
Gp at initial phase. The proposition suggests that for small top layer initialization (measured by o), ||Gr|l will first increase from ((x) to O(1) and then converge exponentially to O(n). Fig. shows that this is indeed the case for |Gr fl, regardless whether grokking happens or not.
Pp at later phase. The structure of F (+00) revealed by the following lemma:
Lemma 1 (Structure of backpropagated gradient Gp). Assume that (1) entries of follow standard
normal distribution N (O, 1), (2) |xx = const, (3) ||x; x - pll2 S. < for all + i' and (4) large
width K, then both FT F and FFT becomes a multiple of identity and Eqn. 4 becomes:
Gr(+00)= (Ke1 m)(nc2 +n)
"1
VVT F ((K-
(5)
where C1. C2 > O are constants related to nonlinearity. When 1j is small, we have GF & NYÝTE. F. Note that the input features and/or weights can be scaled and what changes is 01 and cz-
Interestingly, in both cases, we see that Gp contains a key term ÝÝTF. As we will see, it plays a critical role in feature learning. From Eqn. it is clear that if K - too, then Gp (+0o) + and there is no feature learning (i.e., NTK regime). Here we study the case when K is large (so that Eqn. i valid) but not too large so that feature learning happens.
3.2 INIO-INIT ACCELERATES THE FEATURE LEARNING PROCESS
Given the analysis in Sec. 3 of Stage I, it becomes tempting to think that if we initialize the top-layer weight V to zero, then at the initial stage the backpropagated gradient GP(() carries only the signal term Y T F according to Eqn. 3 and the feature learning process should be accelerated. We named this new approach as zero-init.
Experiments show that this is indeed the case. Fig. 15 16 and 17 show that for M = 41, 89, 127, zero-init accelerates the feature learning process compared to normal initialization.
In multi-layer setting, the boost brought up by zero-init is larger (Fig. 18). in particular when data scarce and the feature learning process becomes slow. The gap be large 10x (e.g., 100 epachs versus 1000 epochs) when training on MSE loss.
All the experiments in Appendix J.3
4 STAGE II: INDEPENDENT FEATURE LEARNING
4.1 THE ENERGY FUNCTION &
Now let us explore the feature learning process with the help of Gr. Let W [W1, W....., Wkl where w, E Rd is the weight vector of j-th node, and F F= [f, f2....,, fx] where f, (Xw,) R" is the activation of j-th node. For Gr V F. as the structure shown in both initial stage (Eqn. 3 and later stage (Fqn. (5), the j-th column g, of Gr is only dependent on j-th node w,, and thus decouple the dynamics into K independent ones, each corresponding single node:
w = D gj x Y T (Xw,)
(6)
where D, = diag(o' (Xw,) )) is the diagonal gating matrix of j-th node. Note that FTF ýT
since P is idempotent. A critical observation here is that Eqn. 6 actually corresponds to the gradi- ascent dynamics of the energy function E
4

---

Published as a conference paper at ICLR 2026
Theorem 1 (The energy function & for independent feature learning.. The dynamics (Eqn. 6 of independent feature learning is exactly the gradient ascent dynamics of the energy function w.r.t. W,, a nonlinear canonical-correlatinn correlatior analysis (CCA) between the input X and target Y:
E(w,) = Y oXX;))1
Therefore, the feature learned for each node is the one that maximizes the energy function E(yy). Since Eqn. 6 can be unbounded, in the following, we put an additional constraint that ||wyll2 due to weight decay regularization. Note that [Tian] (003) also arrives at an energy function when studying feature learning in the context of contrastive loss, the resulting function abstract and difficult to interpret its structure of its local maxima. Here the structure is much clearer, which we will explore below.
4.2 GROUP ARITIMETIC TASKS
To demonstrate a concrete example in solving the energy function E, we consider group arithmetic tasks, i.e., for group H, the task to predict h = h2 h22 given h1 h2 E H. One example the modu- lar addition task hn h22 h1 + h2 mod M. which has been extensively studied in grokking (Power et al. 2022
Gromov
2023 Huang et al. 2024 Tian 2025).
Definition of group. Here h2 is the group operation of two elements hn and h2 in the group H. The group operation satisfies few properties: (1) the group operation is associative, i.e., (h1 h2) h1 (hyh3) for all h1, h2, ha H. (2) there exists an identity element such that he = eh = h for all h E H. (3) for each h E H, there exists an inverse h-1 such that Note that the group operation not commutative in general, i.e., h1h2 & h2h1 for some h1, h2 E H. If the operation is commutative, then we call the group Abelian. Modular addition is an example of Abelian group.
= h- h
The task. We represent the group elements by one-hot vectors: each data sample x, E R2M is concatenation of two M-dimensional one-hot vectors [+) 9h2(() where h and h2[i] are the indices of the two one-hot vectors. The output is also one-hot vector y; @hrliha/i) where i n M Here the class number M | is the size of the group.
crash course of group representation theory. A mapping p(h) И + Cdxd is called a group representation if the group operation is compatible with matrix multiplication: p(h,) p(2)))(h)) = p(hgh2) for any h1, h2 E H. Let R E IR MxM be the regular representation of group element so that Chyh2 Ru, Ch₃ for all h1, h2 H, and P R MMM be the group inverse operator that Pej en. Note that p2 = and PT = P
A tale of two kind of "representations" We often refer the activation vectors in neural network as "neural representations", which means that we use high-dimensional vectors to represent certain semantic concepts (c.g., word2vec (Mikolov et al., (013)). For groups, the group element corre- sponds to "transformation" / "action" that acts on one object and turn it to another (i.c., rotation of a vector). Therefore, group representations often take the form of matrices (e.g., permutation / rotation matrix).
The decomposition of group representation. The representation theory of finite group (Fulton & Harris 2013) Steinberg 2009) says that the regular representation R, admits a decomposition into complex irreducible representations (or irreps):
((HH) ma
Rh = Q
+ K=
Ck(h)
Q"
(8)
where x(H) is the number of nontrivial irreps (i.c., not all / map to identity), (h(h) e Cde xd, is the k-th irrep block of Rr. Q is the unitary matrix (and Q* is its conjugate transpose) and my is the multiplicity of the k-th irrep. This means that in the decomposition of R2, there are m, copies of dy-dimensional irrep, and these copies are isomorphic to each other. So the k-th irrep subspace Hi has dimension d,d
For regular representation { ), one can prove that my = dg for all ls and thus H = M = -x di- For Abelian group, all complex irreps are 1d (i.e., Fourier bases). One may also choose to do the
5

---

in a real irrep of dimension d2 (with Published as a conference paper at ICLR 2026
Sufficient Data
Nonlinearity (+))
M
Elinear (w)
E(w)
Insufficient Data
Figure 4: Change of the landscape of the energy function & (Thm. [ Left: & with lincar activation reduces to simple eigen-decomposition and only have one global maxima. Middle: With nonlinearity, the energy
landscape now has multiple strict local maxima, each corresponds to feature (Thm. 2
More importantly, target prediction (Thm. Bp. Right: With sufficient
these features are more efficient than memorization training data, the landscape remains stable and we can recover these (generalizable) features (Thm. 4. with insufficient data, the landscape changes substantially and local maxima becomes memorization (Thm
decomposition in real domain. In this case, pair of 1d complex irreps will become 2d real irrep. For example, e0d and e -i0 becomes 2d matrix coos sin( sin( cos(
4.3 LOCAL MAXIMA OF THE ENERGY FUNCTION
Now we study the local maxima &. With the decomposition, we can completely characterize the local maxima of the energy & with group inputs, even that E(w) is nonconvex.
Theorem 2 (Local maxima of & for group input). For group arithmetics tasks with (xx) x² has multiple local maxima w
[n: PP]. Either E = M/8dj and E H&,), or in a puir of complex irrep of dimension dy (with E- = M/16ds and Hx 4 Hgg. These local maxima are not connected. No other local maxima exist.
Note that our proof can be extended more general nonlinearity (xx) + bxr with >> 0 since linear part will be cancelled out due to zero-mean operators. We can show that local maxima of & are flat, allowing moving around without changing E:
Corollary (Flatness of local maxima of & for group input). Local maxima of & for group arith- metics tasks with H = M > 2 are flat, i.e., at least one eigenvalue of its Hessian is zero.
We can apply the above theorem the popular modular addition task which is an Abelian group. The resulting representation is Fourier bases.
Corollary 2 (Modular addition). For modular addition with odd M, all local maxima are single frequency Uμ au [cos(kmm)) [ + by be((sin(kmw)) M-1 where w = 2 with Ex = M/16. For even M, uM/2 x [(-1)Tm[ mm=0 has 8 = -M/8. Different local maxima are disconnected.
Role played by the nonlinearity. Note that if (zx) is linear, then there is no local maxima, but only a global one, which is the maximal eigenvector of X T X. This corresponds to Linear Discriminative Analysis (LDA) (Balakrishnama Ganapathiraju, 1998) that finds directions that maximally separate the class-mean vectors. For group arithmetics tasks, for each target h h h2, each group element (h₃ and h2) appears once and only once, the class-mean vectors are identical and thus LDA fails to identify any meaningful directions. With nonlinearity, the learned w has clear meanings.
Meaning of the learned features. First, the learned representation offer more efficient construction of the target (see Thm. 3) than simple memorization of all M2 pairs. Second, learned representations naturally contain useful invariance. For example, some irreps of the cyclic group of Σ15 behave like its subgroup Z3 and Zs, by mapping its element h to div 3) and div(h, 5). If we regard be controlled by two hidden factors, then these features lead to focusing on one factor and invariant to others. More importantly, they emerge automatically without explicit supervision.
4.4 REPRESENTATION POWER OF LEARNED FEATURES
With Thm. 2 we know that each node of the hidden layers will learn various representations. The question is whether they are sufficient to reconstruct the target and how efficient they are.
Theorem (Target Reconstruction). Assume (I) optimized in complex domain C, (2) for each irrep k. there are m d& pairs of learned weights (uu") } form complete bases for Ha and (3) the top layer V also learns with n then Y = Y
Iu; + Pu whose associated rank- matrices

---

Published as a conference paper at ICLR 2026
1.0
0.7
1.0
0.6
0.6
0.5
train%
0.6
-0.6
0.4
0.4
0.2
0.2
20 60
80 1000
120
20 40
60
80
100
120
phase transition in modular addition tasks. When M grows, the train- ing data ratio p n/M² required achieve generalization decrcases. This coincides with Thm. which predicts p ~ M -1 log (dotted line). use learning rate 0.0005, weight decay 0.0002 and 2048. Results averaged over 20 seeds. Left: Simple cyclic group ΣM for prime M. Right: Σ for composite M. Other groups (product and non-Abelian) are shown in Appendix 0.1
Figure 5: Generalization//mmmoriiz
From the theorem, we know that K = 2 E m m I 2 [MM <H)))² + x(H) - 1] suffice. In particular, for Abelian group, x(H): = pure memorization solution that would require M² nodes, i.e., each node memorizes single pair (hi, E x H.
and K = 2M - 2. This is much more efficient than
Assumptions of the theorem. Training the model in an end-to-end manner automatically satisfies the assumption (3). If we initialize the weights randomly, then with high probability, the resulting u not collinear and the assumption (2) can be satisfied. For assumption (1), interestingly, no such theorem can be stated in the real domain R, because in R. the subspace of orthogonal matrices, which p representations belong to, is not covered by the subspace of symmetric matrices spanned by (uu' J. In contrast, in the complex domain C, the subspace of unitary matrices can be represented by Hermitian matrices. Does that mean that the real representation
not that useful? Not really. If we change w Iu; Pu] slightly to w [u: = Pu'] in which u' is small perturbation of u, then Thm. holdd for real solutions. This happens the stage III when end-to-end backpropagation refines the representation. Fig. 12 shows that using complex weights still works and may grok faster.
4.5 THE SCALING LAWS OF THE BOUNDARY OF MEMORIZATION AND GENERALIZATION
While Thm. 2 shows the nice structure of local maxima (and features learned), it requires training on all n M pairs of group elements. One may ask whether these representations can still be learned if training on a subset. The answer is yes, by checking the stability of the local maximum.
Theorem (Amount of samples to maintain local optima). If we select n 2 df M log(M) (M/)) data sample from H x H uniformly at random, then with probability least 1 ô, the empirical energy function & keeps local maxima for dę di-dimensional irreps (Thm. 2
The theorem above says that we do not need all 2 samples, but only O(M M) samples suf- fice to learn these features, which will generalize to unseen data according Thm. 3 Fig. 5 demonstrates that the empirical results closely match the theoretical prediction, and there is clear phase transition around the boundary (test accuracy + 1), where the training data ratio p n/ = O(M-1 log M).
Memorization. On the other hand, we can also construct cases when memorization is the only local maximum of E. This happens when only collect samples for one target /t but missing others, and diversity is in question.
Theorem (Memorization solution). Let ((x) (xr) and assume o' dxx) > for x > 0. For group arithmetic tasks, suppose we only collect sample (g. g for one target h with probability Pg. Then the global optimal of & a memorization solution, either (1) focused memorization w = v2 (eg=: (gr 9" = arg max Py if d is nondecreasing, or (2) a spreading memorization with w= Eg 8 e nn), if á strictly decreasing. Here 8g = - and X is determined by
Σ, g = 2. No other local optima exist.
We can verify that power activations (c.g., oxx) x²) lead to focused memorization, while more practical ones (c.g.. ReLU, SiLU, Tanh and Sigmoid) lead to spreading memorization. We leave for future work whether this property leads to better results in large scale settings.
7

---

Published as a conference paper at ICLR 2026
Noduraradiiion 2M1
Modura add Bonn W1-
Moduar 1M-1271
28%
3B5
36%
w3-304
6d-304
35% 0
04
62-
12%
0.2
385 0.2
34%
60
0.0
38902.3085 0.001
8.30
383
8:05
03
0 0.802333055 0.00
0,300
3.2 0.55
0.8023300055 0.003 6:305 502
0.08
0.1
Figure 6: Phase transition from generalizable (gsol) to non generalizable solutions (ngsol) in modular addition tasks (M 23, 71, 127) with K = 1024. Around this critical region, small learning rate more likely lead gsol. to the fact that small learning rate keeps the trajectory staying within the basin towards g gsol, while large learning rate converges to solutions with higher & (Fig. 7 Results averaged over 15 seeds.
2-323
334
2-26
2-325
X-325
4-4311
2_ 4.00
Amp
5-507
p
"MM1
c0ssa
2-120
c-207
2-3.0
1-000
Figure 7: In small data regime of modular addition with M = 127 and n 3225 (20% training out of 127? samples), Adam optimizer with small learning rate (00.011, left) and (0.002, middle)) lcads to generalizable so- lutions (Fourier bases) with low E, while with large learning rate (0.005, right), Adam found non-generalizable solutions (e.g.. memorization) with much higher E.
Boundary of generalization and memorization (semi-grokking (Varma et al. 2023 In between the two extreme cases, local maxima of both memorization and generalization may co-exist. this case, small learning rate keeps the optimization within the attractive basin and converges gsol, while large learning rate leads to ngsol which has better energy & (Fig.
Our analysis fits well with the observed empirical behaviors that there exists critical data size Varma et al. 2023) or more precisely a critical data ratio (Wang et al. [2024a] Abramov et al., 2025
Liu et 2022), above which the grokking suddenly leads generalization. Our framework Lig gives a theory- theory-backed explanation: different data distribution leads to different landscapes, which gives either generalization or memorization local maxima that the weights will fall into. also explains the ungrokking phenomenon: grokked model can move back to memorization when continues to train on a smaller dataset (Varma et al., 2023), and consistent with (Nguyen & Reddy) 2025) that shows that task diversity is important for generalization. The empirical hypothesis (e.g., "memorization/geee ietiaaatiiongenealizaaioooo circuits") now becomes a natural consequence of changed landscape after training.
5 STAGE III: INTERACTIVE FEATURE LEARNING
The starting point of Stage II is simplify the exact backpropagated gradient Gp = PTYÝTFF
(Eqn. 4 with B == (FT F + to Gp & n T F, by two approximations: (1) x I, and (2) P. 1/. The two approximations initialized. When training continues. W evolves from random initialization and the conditions may not hold anymore. In this section we put them back and study their behaviors.
valid due to Thm. I when the hidden weights W is randomly
8

---

Published as a conference paper at ICLR 2026
5.1 REPULSION OF SIMILAR FEATURES
We first study the effect of B, which leads to interplay of hidden nodes. Over the training, the activations of two nodes can be highly correlated and the following theorem shows that similar features leads to repulsion.
Theorem (Repulsion similar features). The j-th column of FB B is given by:
K
[FB), = =b,ff +
Σ
byffi
(9)
1=1
where sign(bşt = -sigu(f) Pa. Pa.-s Here == - P- F- + is projection matrix constructed from F-50 which F excluding the 1-th and j-th columns.
Remark.
Intuitively, if fy and f, are similar, then byt will be negative and the resulting and ! columns of FB will be pushed away from each other and vise versa.
5.2 NOP-DOWN MODULATION
Over the training process, is possible that some local optima are learned first while others learned later. When the representations are learned partially, the backpropagation offers a mechanism to focus on missing pieces, by changing the landscape of the energy function E.
Theorem (Top-down Modulation). For group arithmetic tasks with (x)) x?, if hidden layer learns only subset S of irreps, then the backpropagated gradient Gr ix (Øs & 1)) & aar) " F (see proof for the definition of 4s), which yields modified Es that only has local maxima on the missing irreps & S.
5.3 DIVERSITY ENHANCEMENT WITH MUON
In addition the mechanism above, certain optimizers (e.g.. Muon optimizer (Jordan et al. 2024 can also address such issue, by boosting the weight update direction that are underrepresented, enforcing diversity of nodes. While evidence (Tveit et al., (2025) and analysis exist (Shen et al.) 2025) to show that Muon has advantages over other optimizers, to our best knowledge, we are the first analyze it in the context of feature learning.
Recall that the Muon optimizer converts the gradient Gw = UGW DVT (its SVD decomposition) Gw Gww- UGW Vy and update the weight W accordingly (i.e., W x Gww). We first show that when Muon is applied to independent feature learning on each wy to make them coupled, it still gives the correct answers to the original optimization problems.
Lemma (Muon optimizes the same as gradient flow). Muon finds ascending direction to maximize the joint energy function Ejoime (W) = >; {(w,) and has stationary points iff the original gradient Gw vanishes.
Now we show that Muon optimizer can rebalance the gradient updates.
Theorem (Muon rebalances gradient updates). Consider the following dynamics Tian 2023
w = A(w)w, ||w||2 1
(10)
where A(w) E.A(W)<IGT. Assume that (1) {C,} form orthonormal bases, (2) for E: aiS we have >x(w) = iaii with μn I 1, and (3) {ai} is initialized from inverse-exponential
distribution with CDF(x) = exp(-x-0) with a > Then
Hf En ps. Then the expected #nodes to get all local maxima To 2 max 11 1/ming Pt, Et. 1/1).
Independent feature learning. Pr(w + Ci]
Pt P
Muon guiding. If we use Muon optimizer to optimize K nodes sequentially. then the ex- pected #nodes get all local maxima T. = 2- 2- For large Ta T..

---

Published as a conference paper at ICLR 2026
The intuition here is that once some weight vectors have "occupied" local maximum, say Sm their gradients point to the same direction (before projecting onto the unit sphere |/|l2 = 1), and the gradient correction of Muon will discount that component from gradients of currently optimized weight vectors, and keeping them away from Sm- In this way, Muon pressed novel gradient direc- tions and thus encourages exploration. Fig. shows that Muon is effective with limited number of hidden nodes K.
Note that Eqn. iO closely related &, under the assumption of homogeneouusrrversibbe activation. i.e., or (x) Ca'(x) (x)x with a constant C' (Zhao et al. 2024 Tian et al. 2020). In such setting, Eqn.
is related to the gradient dynamics with a PSD matrix A(w) = X D D(w) (w)YY"D(w)....
6 EXTENSION TO DEEPER ARCHITECTURES
The above analysis and the definition of the energy function & can be extended to deeper architec- tures. Consider multi-layer network with L hidden layers, Fj = ((F:-WWi with Fo X and Y = FFV. For notation brevity, let G3 = GF..- Let's see how the gradient backpropagated and how the learning fits our framework (Fig. D
Stage Stage I does not change since Ft is still a random representation. Then when V starts to learn and converges, the backpropagated gradient Gz now carries meaningful information: Gz xc YY Fy (Eqn. 5). which initiates Stage II.
Stage II. We assume homogencous activation (x) = (z) For the next layer L - 1, we have: (D,ÝÝ
= -i-, = -
(11)
since Wr, is randomly initialized, we have W.W,T I and thus Gr. DI Dr.Fr.-1
Doing this iteratively gives Gi (D4+ Di+1) Fi, where Dj = II Dm- Note that these D matrices are essentially reweighing/pruning samples randomly, since right now all {Wi} are random except for V. Now the lowest layer receives meaningful backpropagated gradient G₁ that is related to the target label, and it also exposes to input X. Therefore, the learning starts from there. Once layer learns decent representation, layer + 1 receives meaningful input Fi and starts to learn, etc. When layer is learning, layer l' I do not learn since their input Fi. remains random noise.
From this analysis, in which G, is definitely a much cleaner and stronger signal, compared to G₁ which undergoes many random reweighing and pruning of samples.
also see why residual connection helps. In this case, Gres. = Σle Gt.
Stage III. Once the activation Fi becomes meaningful, top-down modulation could happen (simi- lar to Thm. 7 among nearby layers so that low-level features can be useful support high-level representations. We leave the detailed analysis for future work.
7 CONCLUSION, LIMITATIONS, AND FUTURE WORK
We develop a principled mathematical framework Lig for the grokking dynamics in 2-layer net- works. Lig identifies three stages, lazy learning, independent feature learning, and interactive feature learning, each marked by distinct structure of backpropagated gradient Gp. Weight de- cay enables hidden nodes to independently extract lahel-related features via a well-defined energy function ε, before cross interactions consolidate them. The analysis clarifies how hyperparameters like weight decay, learning rate, and sample size shape grokking, and offers first principles insight into why modern optimizers (e.g., Muon) are effective. We also extend our framework to deeper networks.
While the derivation of energy is applicable to any input, analysis of its local maxima relies on restrictive assumption of group structure the input. Also our analysis does not include the transition time between consecutive learning stages. We leave them for future work.
Limitations.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Roman Abramov, Felix Steinbauer, and Gjergji Kasneci. Grokking in the wild: Data augmentation for real-world multi-hop reasoning with transformers. arXiv preprint arXiv:2504.20752 2025.
Jimmy Ba. Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Suresh Balakrishnama and Aravind Ganapathiraju. Linear discriminant analysis-a brief tutorial. Institule for Signal and information Processing, 18(1998)1-8,, 1998.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hid- den progress in deep learning: Sgd learns parities near the computational limit, Advances in Neural Information Processing Systems, 35:21750-21764, 2022.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shengbang Tong, Saining Xie, Dale Schuurmans, Quoc Le, Sergey Levine, and Yi Ma. Sft memorizes, rl generalizes: A comparative study of foundation model post-training. arXiv preprint arXiv:2501. 17161, 2025.
Kenzo Clauw, Sebastiano Stramaglia, and Daniele Marinazzo. Information-theerrtiic progress mea- sures reveal grokking is emergent phase transition. arXiv preprint arXiv:2408.00 08944, 2024.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pp. 1452. PMLR. 2022.
Clémentine CJ Dominé, Nicolas Anguita, Alexandra M Proca, Lukas Braun, Daniel Kunin, Pe- dro AM Mediano, and Andrew Saxe. From lazy to rich: Exact learning dynamics in deep linear networks. arXiv preprint arXiv:2409. 14623, 2024.
Darshil Doshi, Aritra Das, Tianyu He, and Andrey Gromov. To grok or not to grok: Disen- tangling generalization and memorization on corrupted algorithmic datasets. arXiv preprint 0rXiv:2310.13061 2023.
Darshil Doshi, Tianyu He, Aritra Das, and Andrey Gromov. Grokking modular polynomials, arXiv preprint arXiv:2406.0344 03495,, 2024.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054, 2018.
Philippe Flajolet, Daniele Gardy, and Loýs Thimonier. Birthday paradox, coupon collectors, caching algorithms and self-organizing search. Discrete Applied Mathematics, 39(3):207-229, 1992.
William Fulton and Joe Harris. Representation theory: firstt course, volume 129. Springer Science & Business Media, 2013.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02673 2023.
Yufei Huang, Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. Unified view of grokking, double descent and emergent abilities: A perspective from circuits competition. arXiv preprint 251.:2200..1115.. 2024.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks. 2024. URL https:
//kellerjordae github. io http/oosss/muueee s/muon/
Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh. Local signal adaptivity: Provable feature learning in neural networks beyond kernels. Advances in Neural Information Processing Systems, 34:24883-244 -2897,, 2021.
11

---

Published as a conference paper at ICLR 2026
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking the tran- sition from lazy to rich training dynamics, 2024. URL
06110
Aitor Lewkowyez and Guy Gur-Ari. On the training dynamics of deep networks with 1.2 regular- ization. Advances in Neural Information Processing Systems, 33:4790-4799, 2020.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. Advances in Neu- 2022.
ral Information Processing Systems,
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference on Learming Representations, 2023. URL https: //openzeview. httt:/oorumm11222DDoooo...
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking as competi- tion of sparse and dense subnetworks. arXiv preprint arXiv:2303. /183, 2023.
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeffrey Dean. Distributed rep- resentations of words and phrases and their compositionality. Advances in neural information processing systems, 26, 2013.
Beren Millidge. Grokking "grokking", 2022.
URL eetpp::/www.beeees io/
Iman Mirzadeh, Keivan Alizadeh, Hooman Shahrokhi, Oncel Tuzel, Samy Bengio, and Mehrdad Farajtabar. Gsm-symbolic: Understanding the limitations of mathematical reasoning in large language models. arXiv preprint 240.0522,, 2024.
Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica J Sutherland. Why do you grok? a theoretical analysis of grokking modular addition. arXiv preprint arXiv:2407.12332, 2024.
Andrea Montanari and Pierfrancesco Urbani. Dynamical decoupling of generalization and overfit- ting in large two-layer networks, 2025.
Depen Morwani, Benjamin L. Edelman, Costin-Andrei Oncescu, Rosie Zhao, and Sham Kakade. Feature emergence via margin maximization: case studies in algebraic tasks. arXiv preprint arXiv:2311.07568, 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learn- ing Representations, 2023.
net Forummidd 9XF SbD Pmdli
Alex Nguyen and Gautam Reddy. Differential learning kinetics govern the transition from mem- orization to generalization during in-context learning. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL httpt:////// /ttprr/v/eperreieew. net/ forum?
TNyi Nyyi?qUdjz
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02117,, 2022.
Lucas Prieto, Melih Barsbey, Pedro AM Mediano, and Tolga Birdal. Grokking at the edge of nu- merical stability. arXiv preprint arXiv:2403.03507 2024.
Lucas Pricto, Melih Barsbey, Pedro AM Mediano, and Tolga Birdal. Grokking at the edge of nu- merical stability. arXiv preprint arXiv:250J. 004697, 2025.
Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking as first order phase transition in two layer networks. ICLR, 2024.
Noa Rubin, Kirsten Fischer, Javed Lindner, David Dahmen, Inbar Seroussi, Zohar Ringel, Michael Krämer, and Moritz Helias. From kernels to features: A multi-scale adaptive theory of feature learning. arXiv preprint arXiv:2502.032210 2025.
12

---

Published as a conference paper at ICLR 2026
Wei Shen, Ruichuan Huang, Minhui Huang, Cong Shen, and Jiawei Zhang. On the convergence analysis of muon. arXiv preprint arXiv:2505.23777, 2025.
Maor Shutman, Oren Louidor, and Ran Tessler. Learning words in groups: fusion algebras, tensor ranks and grokking. arXiv preprint arXiv:2509.069) 2025.
Benjamin Steinberg. Representation theory of finite groups. Carleton University, 2009.
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. arXiv preprint arXiv:2206.0453 2022.
Yuandong Tian. Understanding the role of nonlinearity in training dynamics of contrastive learning. In The Eleventh International Conference Learning Representations, 2023. URL https:
X
Yuandong Tian. Composing global solutions to reasoning tasks via algebraic objects in neural nets. NeurIPS, 2025.
Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning with dual deep networks. arXiv preprint arXiv:2010.00578, 2020.
Joel A Tropp. User-friendly tail bounds for sums of random matrices. Foundations of computational mathematics, 12(4):389-434, 2012.
Amund Tveit, Bjørn Remseth, and Arve Skogvold. Muon optimizer accelerates grokking. arXiv preprint arviv:2504. 16041, 2025.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02390. 2023.
Boshi Wang. Xiang Yue, Yu Su, and Huan Sun. Grokked transformers implicit reasoners: A mechanistic journey to the edge of generalization. arXiv preprint 24Xx::2405..077 2024a.
Xinyi Wang, Antonis Antoniades, Yanai Elazar, Alfonso Amayuelas, Alon Albalak, Kexun Zhang, and William Yang Wang. Generalization vs memorization: Tracing language models' capabilities back to pretraining data. arXiv preprint 2rXiv:2407.14900 2024b.
Zhiwei Xu, Yutong Wang. Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in relu networks for xor cluster data, arXiv preprint arXiv:2310.02541, 2023.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. ICMI, 2024.
13

---

Published as a conference paper at ICLR 2026
A DISCLOSURE OF LLM USAGE
We have used SoTA LLMs extensively to brainstorm ideas to prove mathematical statements pre- sented in the paper. Specifically, we setup research directions, provide problem setup and intuitions, proposes statements for LLM to analyze and prove, points out key issues in the generated proofs, verify
adjust the statements accordingly and iterate. We also have done extensive experiments the resulting statements. Many proofs proposed by LLMs are incorrect in subtle ways and requires substantial editing and correction. We have carefully revised the proofs presented in the work, and take full accountability for their correctness.
B ETHICS STATEMENT
This work is about investigating various theoretical and empirical properties of neural networks. We do not rely on any sensitive or proprietary data, nor do we use any existing open source models that may produce harmful contents.
C REPRODUCIBILITY STATEMENT
datasets used in this work can generated synthetically. Models are pretrained from scratch with very small amount of compute. We will release code to support full Reproducibility.
D RELATED WORKS
Explanation of Grokking. Multiple explanations of grokking exist, e-g., competition of generaliza- tion and memorization circuits (Merrill et al.] 2023). a shift from lazy to rich regimes Kumar et al. (2024) etc. Dynamics of grokking is analyzed in specific circumstance, e.g., for clustering data (Xu et al. 2023), linear network (Dominé et al. 2024), etc. In comparison, our work studies the full dy- namics of feature emergence driven by backpropagation in group arithmetic tasks for deep nonlinear networks, and provide a systematic mathematical framework about what and how features emerge and scaling law about when the transition between memorization and generalization happens.
Usage of group structure. Recent work leverages group theory to study the structure of final grokked solutions (Tian) 2025) Morwani al., 2023 Shutman al.) (025). None of them tackle the dynamics of grokking in the presence of the underlying structure of the data as we do.
Scaling laws of memorization and generalization. Previous works have identified scaling laws for aetnorizationngeneralazaat
(Nguyen & Reddy 2025 Wang et al.) [024a] Abramov et al. 2025 Doshi et al. 2023) without systematic theoretical explanation. Our work models such transitions as whether generalizable local optima remain stable under data sampling, and provide theoretical framework from first principles.
Feature learning. Previous works treats the NTK as a holistic object and study how it moves away from lazy regime, e.g., it becomes more correlated with task-relevant directions (Kumar et al. 2024 Ba et al. 2022 Damian al.) 2022), becomes adapted to the data (Rubin al. 2025| Karp al. 2021), etc. In contrast, our work focuses on explicit learning dynamics of individual features, their interactions, and the transition from memorization generalization with more samples.
E WHEN DOES GROKKING HAPPEN?
We first want to clarify that grokking (i.e., delayed generalization) phenomenon/observaton that may suggest multiple underlying situations. When we observe "grokking happens", it is often the case that the training gets stuck lazy learning for a while before moving to feature learning. On the other hand, when we observe "grokking does not happen", then we may fall into one of the three very different situations: (1) training simply gets stuck in lazy training. (2) feature learning happens very early and the gap between training and test performance is always small, or (3) feature learning has converged to non con--eneraalzable solutions due to e.g., insufficient or noisy data, and thus test performance never improved. Note that (2) is very good outcome while (1) and (3) bad ones.
14

---

Published as a conference paper at ICLR 2026
Therefore, it would be tricky to directly relate the hyperparameters with grokking behaviors. Instead, using Lig framework, we focus on relating them with the backpropagated gradient GF and the underlying feature learning process, which is more fundamental. Only when the Gr become large and sends sufficient "amount" of the key term YÝT F down to the hidden layers, grokking could happen. Here we categorize these factors into several categories.
Learning rate. (Gromov 2023) reports that grokking happens without regularization, but with large initial learning rate (verified by the author). This corresponds to increasing the strength of Gr(() x TÝyTF F at the initial phase of learning so that the hidden layers receives enough correct gradient signal start the feature learning.
Stay longer before overfitting. (Prieto et al. 2025) uses stable softmax (linear form) rather than regular softmax (exponential form) in computing probability. This prevents the model from over- fitting to the label too quickly, and thus maintains a nonzero backpropagated gradient that can be useful for feature learning. (Kumar et al.) (024) also reports that grokking happens without regular- ization, using vanilla SGD optimizer. Our explanation is that it may take longer for SGD to converge to Vndge than Adam, and during that period, the hidden layer has already accumulated a sufficient amount of correct gradient signal.
Weight initialization. (Liu et al. 2023) reports that grokking happens with small initialization, regardless of the weight decay. This is straightforward from our framework, since G = O(a) + CY F + O(at) + Ott2) and if the weight initialization a is small, then Gx(t) is dominated by clear signal term E F, which leads to fast feature learning. If is large, then O(a) term is large and the initial phase of Ge contains too much noise, and we need to rely on the signal provided by the convergence phrase of G controlled by the weight decay 1/. This is consistent with the finding by (Liu et al.] 2023) that for large weight initialization, regularization is needed for grokking to happen, and small regularization leads to slow grokking transition.
Scaling factor 3 of the output. (Kumar et al. 2024 Chizat et al. 2019) reports that scaling the
output by a factor 3 > 1 will make the grokking faster. From Lig framework, this corresponds optimizing (V)
||8 - PVVVI + v/|V Following a similar derivation in Sec. H
H.3 we can show that at the initial phase, the backpropagated gradient Gx(t) = O(a) in Sec.
[SY F O( + 0(89a)) + Ot?). So if 3 > 1 is large and the initialization a <
then the signal term T F becomes more dominant than the case of B = 1, and the grokking happens faster. However, if we have too large 3 over a, then the term O(32a) may cause trouble in feature learning and the final we downscale the target label (Kumar et al. 2024).
becomes small (Eqn. 69) Equivalently, same theory applies
Weight decay 1). According to Eqn. since GF (+oc) x NYYT F. it is clear that the weight decay 17 becomes the learning rate of feature learning process. This coincides with findings in empirical works (Power et al., 2022: Clauw et af.] 2024) that low regularization leads to slow grokking transi- tion. This is also consistent with i ~ 1/1 laws to start grokking (Liu et al. 2023 reach maximal test performance (Lewkowyez & Gur-Ari
2020).
Data size n. Our sample analysis (Theorem. 4 shows the local maxima can be kept with sufficient number of samples (n Z M log M). Intuitively, more samples lead to better shaped local maxima
with less noise and thus the feature learning is faster, which is consistent with (Power et al. 2022
Nanda et al. 2023 Wang et al., 2024a] Gromov 2023 Abramov et al. 2025).
The number of hidden nodes K. Our analysis requires that we need a decent number of hidden nodes K to cover the diverse set of the local maxima of E. On the other hand, Lemma [ tells that very large K may reduce [+o(+ and makes grokking slower. This is consistent with the finding by (Chizat et al., 2019) and NTK regime (Du et al.) 2018 Jacot et al. 2018) that feature learning does not happen.
(onaliiiaaiiooMMotaaaarr shows that if data are scarce, generalization solutions have lower energy than memorization ones. trained long enough, weights will move solutions with larger E.
Early stopping. Model may overfit after generalization Montanari & Urbani 2025). Our framework
15

---

Published as a conference paper at ICLR 2026
F INDEPENDENT FEATURE LEARNING (SEC. 4)
Lemma 3. Let (2(2) == (=)(=)/ vil be the orthonormal Hermite system (7))). If (21, 22) are standard normals with correlation p. then
(l)n(()) 2m [2)] = p" 8mm
(n, 2 0).
Proof of Lemma 3 Use the generating functior exp(tz Then, for correlated normals (Z1, Z2) with correlation p.
= Exz0 (x for ~ ((0, 1).
E th o u7 exp(pbu)=
Σ p* (lu)*.
k20
Expanding the hand side by the generating functions and matching coefficients of th yields Bi (Z1)
[2)] = p 6nm.
To show why E e 12 u2 vnd
exp(ptu) is correct, decompose (Z1, Z2) into Gaussian
independent random variables (X2Y):
Z1 = X,
Z= = :x:=Xx+ + p2
Then we have
12, &
2 -
uZz 2
E. e -
u(px+ v 1-pa 1) -
E (tte 127 " p²
For G ~~ (), have B|eaG) ea²12 hence E 100
1 due Lemma 4
Applying this
twice,
E (t+pu)x 12 -epp
(t + pu)
t2
plu + p2u
= exp
2
27 E u "Ve
Y
12(11
p²)
u2
p2u2
exp
exp
2
2
2
Multiplying the two factors yields
p²u²
puu
exp
ptu
exp
exp(ptu),
2
as claimed.
Lemma (Moment identity). For X N' 1), [[c =
/2). Equivalently, 1feex-1112
1.
Proof. Complete the square:
(t2 22T / (11--))112 812 dx =cxp exp
E 1 = 2 J 1022 dx
2
Lemma (Structure of backpropagated gradient GF). Assume that (1) entries of W follow standard
normal distribution N (O, 1), (2) |x ll2 = const, (3) ||x; x, - pll2 S for all i + i' and (4) large width K, then both FT F and FFT becomes multiple of identity and Eqn. becomes:
Gr(+oo)= = (Ka, ()cn2 +
VVT F + O(K-1()
(5)
where C1. C2 > are constants related to nonlinearity. When 1j is small, we have GF x NY Note that the input features and/or weights can be scaled and what changes is eı and cz.
https: /tmp:s://e...ee wik wwt.:wiwiiiiiiiio/oooo.
wetrriteepplyooopooooooo al.
16

---

Published as a conference paper at ICLR 2026
Proof. In the following, will prove that (1) F F is multiple of identity and (2) FFT al + 811 Without loss of generality, we assume that entry of W follows standard normal distribution 1)..
FT F is multiple of identity. Since each column of F is Pf (XWy)) zero- mean n-dimensional random vector and columns are i.i.d. due to the independence of columns of W. With large width K, FT F becomes a multiple of identity.
FFT is diagonal plus all-constant matrix. Note that the i-th of F is [o(w] x,), (w, x.). o.....,wwwww ( x))], with large width K, the inner product between the i-th row and j-th row of F approximates (K(), where K(i,j) is defined follows:
(12) To estimate the entry K(i,,), we first do standardization by setting Z1 := w wx///, and Ze == w w'xx. x//s; where 8 ||x|12 and 8j ||5 Then (Z1, Z2) are standard normals with Corr(Z1, 22) Pij. and K[(i.) E [o (1)
() E (a(w x,)o(w'xx Xxy)]
Let (h(=)) == Hez(=)/v VT! be the orthonormal Hermite system on (,2(7), where y the standard Gaussian measure and Hez are the Hermite polynomials. For 8 2 0 define (z) (sz). By the (7()) assumption, fa = In= (s) oi with
1 ar (Js, (1)
B[o(sZ) (o(ss)Ho((Z)) VII
Thus
(s; = Σ (s:) (1)))),
(s, Σ
ai(s;) (z())).
120
120
By bilinearity and Lemma 3
k(i,j) =
Σ
ar(s))o
Σ
am (sj)@m (Z2)
Σ
ar(s:) (lss)am (8)) 22) m (22)
120
m20
1m220
(r(s))ar(3)) plj.
120
If = s and
||pis 2ll2 s E ffo + j. then
K(i, = Σ (s)
=:
120
Let == la)()) +oo (it is convergent due to the big factor in the denominator). Let b:= Eiz0 (x) p² and we have for all i + j:
[[((i,,) - bl12 Σ (s) lll p'112 Σ (a) cc
120
121
due to the fact that ||e|) 1/112 S 1g--1e for all / > and some & in between Pig and p. hence K(ii,) = 6)ōg + b Oe)) and thus FFT K(a 6)1 Kb11 +(Ke)]1TT, Note that by Parseval's identity, = =xx
Therefore, FFT = ((a (a-a+ K(a = K K(a
and P, K(a-b)+- Y. Since FT F is proportional to identity matrix, (FT F + v7 77 is also proportional to identity matrix and the conclusion follows.
F.1 THE ENERGY FUNCTION ε (SEC. 4.3
Theorem 1 (The energy function for independent feature learning). The dynamics (Eqn. o of independent feature learning is exactly the gradient ascent dynamics of the energy function & w.r.f.
Wj, nonlinear canonical-correlate analysis (CCA) between the input X and target Y:
E(w,) = Y Too oXxx
17

---

Published as a conference paper at ICLR 2026
Proof. Taking gradient of & w.r.t. w,. and we have ww x D which proves the theorem.
Theorem 2 (Local maxima of & for group input). For group arithmetics tasks with o(x) x?, & has multiple local maxima w" [u: Pul. Either it is in a real irrep of dimension d3 (with 3. M/8dg and Hz). or a pair of complex irrep of dimension dy (with E- M/16ds and E H& @ Hig. These local maxima are not connected. No other local maxima exist.
Proof. Following this setting, if ordered by target values, we can write down the data matrix X == [Xx, Xhgi... k (i.e., each Xn occupies M rows of X) in which each Xx [R] P] RMx2M Here R₂ is the regular representation (a special case of permutation representation) of Bhg group element h so that hh Ra, Bh₃ for all h1, h2 E H, and P is the group inverse operator so that Per_ e/- This is because each row of X that corresponds to the target h can be written as lenty 0,, P] [e,,, Ry e,, PI. Stacking the rows that lead to target h together, and order them by h1, we get X₁, = [R, PI.
Let W [u; Pv]. Let matrix Siy == ou; + vy), since RS is a permutation matrix, then (X,,w) = (Ht u + v) is row shuffling of Therefore, (X,,)) = diag H S) where diag(-) is the diagonal of a matrix. Note that in this target label ordering, we have Y IM @ 1m. So for each column h of Y, we have ya. = C₁ @ 1y. So
2h = vI oX w) = (X, و diag
(R) S)MM = un(R (Ra, S)
(13)
where (A, (A B) the Frobenius inner product. And the energy can be written as:
E(w) = 2 Σ (zh - 2)²
(14)
where: 3 == 17 Σn nn = M En, S) = 1 Σn Rr, S) = H7 (1M S). Therefore,
using 1M 1,, E(w) can be written
E(w)
Σ RR. S)?
(15)
where RA = RAP,' Now we study its property. We decompose {Ra} into complex irreducible representations:
+ C&(h) KSD r=1
Rr Q Q
Q*
(16)
where (h) is the k-th irreducible representation block of Rx, Q is the unitary matrix (and Q" is the conjugate transpose of Q) and my is the multiplicity of the k-th irreducible representation. Since Ra a zero-meaned representation, we remove the trivial representation (o(h) and thus Q*1 0.
Let S = Q SQ. Then
Σ (Ri, 5%% = E(e
+ 1= r=1
Ck(h)) Q², 0', = H C+()),
tr(C((h)8
h
1 r=1
ke r=1
(17)
Σ (kr
vec" (S
Σ vec(Cr(h)) vec(Ci vec(Swr.r)
(18)
Case 1. If k & k' are inequivalent irreducible representations of dimension de and de, then we can prove that En vec(C& (h)) vec(C), (h)) = 0. To see this, let Ag,x(()) = E (x) (h), then AKA (Z) is H -invariant linear mapping from dy de dimensional space. Thus by Schur's 's lemma, AK (Z) 0 for any Z. But since vec(Ax, (2)) (E, (h) @ (x)) vec(Z). we have En CN (h) ② (h) 0. Expanding each component, have En (e)(Cx(h)) vec
O. Case 2. If l = ke' are equivalent irreducible representations (and both have dimension de), then we M
prove that Σn reo(Cx(1)) vec(C](() dz
vec(Ta, vec" vec*(Ta,) Then with Schur's average
18

---

Published as a conference paper at ICLR 2026
lemma, have Axa(2) M tr(Z) A vectorization leads to (En Cx(h) & vech))vec())
du M tr(2) voc(Ta.). Notice that vec* (Ia,) vec(Z) = tr(Z) and we arrive at the conclusion.
dy Therefore, for the objective function we have:
E(w) = 2
E( (R S) = M Σ
Σ tr( tr(Sk.r)
(19)
2 hy0
d2
Special case of quadratic activation. If (x) =², then we have S = (uou)l + +1(vov)+ +uv and thus S = úv*, where ú Q*u and v = Q'v. Therefore, since Q*1 0. Sn, üs Vi.r and tr( tr(Skr)) = V
Vkrr- Therefore, with Sauchy-Schward inequality, we have
3 2 Σ RA, S)
M Σ 2
Σ
M S
Σ |
Σ Mkrrl
ug B.
VK.T
dx
2
dz
k
Re
(20)
Let ax
Er / bb Σ, / 1.
and ck
ak. + bk 2 0.
Then we have:
3 E
Σ Rr: 5)
M
axbR
M
c
2
subject to Σ C
(21)
di
8
dx
k
k70
which has one global maxima (i.e.. Cko = 1 for ku = arg ming d3) and multiple local maxima. The maximum is achieved if and only üko.r -vk.. for all and Σ, |ú Σ, 1koorl 1/2.
Local maxima. For each irreducible representation ko, Cka = 1 is local maxima. This is because for small perturbation e that moves the solution from Cj = I[k: = Ko) to de =
1 - i = ko
CN
if k Ko
with c 2 0 and Exetoo 6 = e, for & = E(fcx)) and E' = E((())) we have:
M
(c))
M
(Cks c)2
E'
(22)
8
8
dka
d2
100
M
2c
+
((2))
M Cko
M
Σ
c
=
(23)
8
dim
dko
8
deg
8
dp
All local maxima are flat, since can always move around within üg, and Varr, while the objective function remains the same.
Optimizing in Real domain. The above analysis uses complex irreducible representations. For real w, Sk.- will be a complex conjugate of S_k., for conjugate irreducible representations 1 and -k. This means that we can partition the in Eqn. 19 into real and complex parts:
E(w)
M 2
Σ
dR E
+ A
Σ
- = dx
(24)
k complex, take one
The above equation holds since R2 is real, and for any complex irreducible representation k, its conjugate representation -k is also included. Therefore, to optimize& & in the real domain R, we can just optimize only the real part plus the complex part taken one of the conjugate pair in the complex domain C.
Zero-meaned one hot representation. Note that if we use zero-meaned
hot representation
C₁ Pj. Cn. then RA, Ct2 Ch, and Pë, Ch-1 still hold, and X,, P, P [R]] P) still use X,, but enforee zero-meaned constraints
[R]] P|[P| J. This means that and v. which is already included since Q*I = 0.
Corollary (Flatness of local maxima for group input). Local maxima of E for group arith- metics tasks with H M 2 are flat, i.e., at least one eigenvalue of its Hessian is zero.
19

---

Published as a conference paper at ICLR 2026
Proof. For Abelian group H with [ M 2, all irreducible representations are 1-dimensional, and at least one it is complex. Since C is treated as 2D space in optimization. it has at least degree of freedom change without changing its function value (Eqn. 24 So the Hessian has at least zero eigenvalue. For non-Abelian group, there is at least one irreducible representation k with dimension greater than 1, which means it has at least I degrees of freedom to change Su.- without changing Σ,, tr tr((k,r)) and thus its function value (Eqn. 24 So the Hessian has at least 1 zero cigenvalue.
F.2 RECONSTRUCTION POWER OF LEARNED FEATURES (SEC.
4.4
Theorem (Target Reconstruction). Assume (l) is optimized in complex domain C, (2) for each irrep k. there are m de pairs of learned weights w = Iu; PP whose associated rank-1 matrices {uu') form complete bases for Hy and (3) the top layer V also learns with 17 0, then Y Y.
Proof. For each nontrivial irrep k, let 11, be the central idempotent projector onto the isotypic subspace H₂ Lmg @ Cde (for the regular rep, m. = dz). Let End(Hx) (Hx) be the space of all linear operators that map Hy to itself. Note that the dimensionality of Hr is D3 := mydk.
Let w ju. Pv] be the weights learned by optimizing the energy function with quadratic activa- tion (I) = x². From Thm. 2 we know that at local optima, u, vV and 1 u, = O. Therefore, the feature fA, IR is given by (o denotes the Hadamard product)
1 M
f = 22 ( o, + (R u,)
h
The third term 2 is a constant across all and was removed the zero-meaned projection. By our assumption we have node j and j' with both positive and negative signs. So 2 [fyh fy., 2 (R% uy) uj. If a linear representation of {fz can perfectly reconstruct the target Y, so does the original representation. So for now just let feature fy,A 2 (R, u,) uu 2diag( R uy, Let U, == u,uj, which is Hermitian in (nd)) then fJ,A = 2diag R Uy)-
Gram block diagonalization. For each irrep k. let Jg be the set of all node j that converges to the k-th irrep. For any Hermitian operator U supported in He (i.e. U = =lgUTT,), define the centered quadratic cross-feature
(h) 2diag (diag(RRA) CMM
and write cu, = Icu, (h)lnEH E C² as a concatenated vector.
For U, V e (Hx) define GUU V) = Σneh (cu (h), ev ch)). On Hr. RA = mms Ch), so the map U + cu (h) is linear and the bilinear form G is invariant under U + (I & Cx(g))U R Cx))))". By Schur's lemma, G(U, V) = as (U, V) = ag tr(UV*) for some scalar ag. Evaluating on rank-one U V (or by direct calculation) gives ag. 4, hence
Σ (cu ev ev (h)) tr(UV*).
For U, u,u; and U² u from H4 and He with k + E, we have
Σ (eu, cu),
41 T Σ
diag u,u²)= diag( R ú,ūū)
h
[
-11 T Σ
(R u2) RT u u 41
Σ Rh
(u, (
ū, o u
4/
This means that (f,, F2) = (cu, ) 0. And thus the Gram matrix G:= FT F is block diagonal with each block G2 corresponding to an irrep subspace k. Here G NN Note that since we sample D = m weights, then [Uybe becomes a complete set of bases (not necessarily orthogonal bases) and thus Gg is invertible.
20

---

Published as a conference paper at ICLR 2026
Right-hand side. For any U End(H:),
Tu(h)) Eeu co (k'),, = tr (1I.RR--1.)U)) = tr ((m, @ (a(HU))
and we have [f] MI Y| (k,(h')
Solve LS. Now we try to solve the LS problem GV FT Due to the block diagonal nature, this can be solved independently for each Gg. Consider GR = F Y. Here Fg = [fy]eeJ collects the subset column Ju from F.
Therefore, V =G F y and (y())) as the hi, entry of Vs, has 1y(h') " Daru, (h) 2 L ((Imy @ (Ui)).. Then we have ÿ V2V
(K) (NA),N
E0sh)) cu, (h) JEJ,
E DiG' jEJN
Ill (a1 Ca(h)))ii) diag(R), Uz).
By linearity in U and completeness of {U,} (the Hermitian bases span all operators in Hz), we have for any E End(HH)::
GG la tr())) diag(RN Uj) = diag(R) (R) A)
5
The last equality holds by noticing that (A,U)) = vec* (Ui) vec(A) and thus
4 EEilGh Eyl U)) = A. Take A = 2 Ck(h)) = II, E End(Hs), and we
have:
((()) (.h), = diag R H1 IT:
(h, E H).
To see why y Y, we have: (-))) (AA).
diag( R. (T:
Σ
(R)
(A)), diag R2
Σ
Th Rh:
Since Σx 1I I and Ry R
LR TI
Ra
Nlg
k=
where Hl = M 1 1 is the central idempotent projector onto the trivial irrep. Thus
(1 M 1M
h= h', h h h',
Σ
(k) (A),, diag(R,, Rw)) diag l1o)
1M,
k20
because diag(R), RN) = 1M iff h = and otherwise, while diag (R, lln) = 37 1M for all h. Hence Enso (A) P-Y . =
1
Remark. The above proof also works for real since can always take real decomposition of R₁ and all the above steps follow.
Property of the square term. With quadratic features the class-centered column for node j and block h decomposes F = [A, B), where for B each column j (and block h) is bj,, R, (u
I 112 las (the "square" part) and for each column (and block h) aj.h "cross" part we discussed above). The vector by is entrywise mean-zero, i.e. Σ. by(x) = 0 for all
2 2(RJ u,) o u, (the
h, hence it has correlation with any class-centered target column V(, 1: (b. V)
21

---

Published as a conference paper at ICLR 2026
Σg (x) = O. Moreover, under 1 u 1 ue 1u
0 one has En ac 0. So the normal
equation becomes
A
A B
Y
FV
V
B
A
B
B
B
which gives
A A
YTY
V
0
B
B
So even with the square term B in F, V will still have zero coefficient on them.
F.3 SCALING LAWS OF MEMORIZATION AND GENERALIZATION (SEC. 4.5
Theorem (Amount of samples to maintain local optima). If we select Z dg M log (M/())) data sample from H x H uniformly at random, then with probability at least 1 - 6, the empirical energy function & keeps local maxima for dy di-dimensional irreps (Thm.
Proof. Overview. We keep the setting and notation of the theorem in the prompt (group H, | = M, quadratic activation, S as defined there, 24 - -R,,, tr(R,, S), zero-mean removal already folded into R2)) We analyze random row subsampling and show that the empirical objective keeps the same maximma structure with 2 M log(MM/) retained rows.
Setup. There are M rows indexed by pairs (h1, h2) H H, with target h h2- For each h E !1 exactly M rows map to h; we index them by j E [M] after ordering by hı as in the proof, and write
M
Sh.j == (R. S)
so that
2A
Σ
Sh.j
= (Rx, S).
1=1
We subsample rows independently with keep-probability p E (0, 1]. Let Enis f0, 1) be the keep indicator for the row (h,
Pr(En, =) = ii.d. over (h,j).
The number of kept rows for target is
Σ End Bin(M.p).
E[ma] = pM, Var(fn) Mp(1 - p).
Estimator for 2/-- We use the linear/unbiased SHorvitz-Thompeer target- wise estimator
M End
E
8hJ-
S = 2h.
Define the diagonal sampling matrix
En,1
En.M
W = diag
so
Σn = tr( R. SWHTT) = ( WHT SS).
The empirical Gram operator. Set the normalized per-target weight Wh =
mh E[wa] = 1,
Var(wa) = pM pM
-
< 1
pM'
Decompose WHT into its mean and zero-mean parts:
WHT wh T Ahs
tr(An) =
E[A, mm) 0.
Therefore
2n = (Rp (W₁, A) = Wh 24 + Ehs
Eh = (RaAn, S).
E[EN S.m.l = (25)
22

---

Published as a conference paper at ICLR 2026
Using the decomposition
mk 2/
my vec(Sk.r) veo(Ck..))
tr(Ck,n Sk.r)
Σ
==0r=1
we obtain
Σ
2
2w + En)
Σ
Wh 2k
+
WhZsEn
Σ
E
(26)
signal
mixed
The signal term can be written as quadratic form over irrep blocks:
Σ J
vec(Sk.r
Σ
wh
vec(Ck,n) voo(Ck,,))) voc(Swe.)).
(27)
Recall that the full-data operator is
A
Eow, @ CK M
1
and vec(Ck,t)) voo(Ckr.h) just a column and row reshuffling of Ckr, @ Ck In the following we will study approximation errors of Ag, instead. Let
(2) R,k
1
1
Σ wh Cw @ Ck.b and A =
Σ wn Cw, @ C M
1
=
M
the second- and first-weighted empirical Gram operators, respectively, By construction, E[A
A and E
A + 1-p
pM Ak,k' (a tiny bias of order 1/(pM)).
k,k*
Error bounds for each (k:, k) block. We will control three deviations, uniformly over all (k, ().
|
log(M/5)
E1:
Av,
S Sav
(28)
Mp
op
log(M/5)
A
-Ar
E2:
I 12V
(29)
Mp
Mp
op
M (M/5)
M log( M/6)
E3
Σ Wh Z
S callz1ly
Σ
ch I C4
(30)
h
for numerical constants c; 4, with probability at least
5/3.
Tool: Matrix Bernstein (self-adjoint dilation form)
Tropp
2012).
Let (X,) be independent,
mean-zero random x d matrices with |X,|| < L and
|Σ E|x,X;1
x
I v. Then for all t > 0,
12/2
Pr Σ
X,
2 t
2d cxp
v + Lt/3
Proof of 288
Fix (k, W²) and define B, = Cw. @ Ca. (unitary, so | = 1). Consider
1w 1
X = M (wh, 1) Bh.
E[XA] = 0,
|X
<
M
M
We have
EX E((w - 1)
B B, = Var(wn)
PM3 1.
M2
M2
Summing over à gives variance proxy S M
Since d < M, with probability least
pdT
pAre
1 - 5/3, Matrix Bernstein yields
AN.N
Au,
loop
S V
log(M/0
Mp
which is 28
23

---

Published as a conference paper at ICLR 2026
Proof of 29 Write
A R,k
Ak.
1
Σ (wh wh) Bb.
M Σ ((wn. 1)2 + (wh 13) Bh.
M
For Σg we reuse the argument of 28 For Σ,, note that E[(w), 1)2] Var(wn) I 1/(p
and (wy, - 1)2 is sub-exponential with scale (11((pM) so matrix Bernstein again gives that with probability least - 6/3.
10)(M/6)
D lop S
Mp
Mp
Combining yields
29
Bounds for the mixed and noise terms in (30p Conditional on and {wa}, the {En} are inde- pendent, mean-zero, and
Eie wpl & ||
|xal (RR, S> & R.Aalle | S |Aalle ||
Hence by scalar Bernstein (and Mauchy-Schwapd for the mixed sum),
M Mlog(M/0 MM/))
M Mlog(M/5)
Σ wh 2
S |
Σ c
with probability at least 8/3, which is (30).
Combine the above three bounds, we know that with probability at least 1 1 8, 28 same time.
30
hold at the
Stability of local maxima. For the quadratic case (after mean removal), with the collinear and equal length and v required by local maxima, can written as positive semidefinite quadratic in the block masses (Eqn. 21::
E(e)
= M
of dz
Σ C = Ck 2 0.
8
kz0
The empirical energy has the form
E(c) M
ED+ + (terms independent of c),
where D diag(1/da) and E is the symmetric perturbation induced by replacing A with A.k and by the mixed/noise terms. By 28 30
log(M/6)
1
E op
(31)
Mp
Mp
with probability at least
Directional slope at vertex (no gap needed). mass to any other coordinate b o a: ca
Consider pure-irrep vertex e, and leak
E, cf = E, others 0. Population change:
M (1 (1--)²-1 8
22
+ O(=²).
da
ds
4d
M. Hence every leakage direction is strictly downhill at rate id, even if multiple de tic. Therefore, first-order approximation of A is
M A(cTEc)
M
M + O(E)) + dlEE))pp)) e).
Λ ΑE 8
4d
24

---

Published as a conference paper at ICLR 2026
Therefore A < 0 for all sufficiently small provided
M
M
1 BBop
Dilp
<
< da
4
4d
Combining with (31), a sufficient sampling condition is
(M/6)
1
M
Mp
2 de
log
Mp
Mp
Cd.
8
for a universal numerical constant C. Since the total number of kept rows is n = pdf this is exactly
" Z M d log 8
(up to universal constants). Under this condition, with probability at least - 8, every pure-irrep vertex remains strict local maximum of the empirical objective (energies shift by O(
When several irreps have the same dg (tied energies), which one is the global maximizer may swap, but the local-maxinn set is preserved.
F.4 MEMORIZATION
Setting. Fix group element h. The admissible training pairs are (g. 1)) for g E H with
probabilities Pg == P9.g-1h and unique maximum at g² i.e, Py > Pg for g # g Let
W [u; v] E M2M with budget ||u|| + ||v|β = 1. Define the pair-sums = + Ug 2 0. Then Eg 2 and the (single-target) objective reduces to g
F(s) = Egg subject to
8g 1 0,
Σ 2 22
where a E ((O x))) is strictly increasing on (0, oc). Maximizing the energy & is equivalent (up fixed positive factor) maximizing F.
Lemma 5 (KKT characterization via ø = d2/2). Assume o'(z) > 0 > 0. and define o(x)
of x x > O. Let s* be an optimal solution. Then there exists 2 such that for each g:
Pg (s²)) = 2X, i > 0,
(32)
Moreover, the budget is tight: then for every active coordinate >>0 0,
Eg (87) = (hence X > 0). If db is strictly monotone on (0, oo).
2/ = Φ-1 (2)
22
(33)
Pg
with X 2 0, Proof. Consider the Lagrangian L(s, A, = Eg (((g)) - X[E, -) - Mg
ly 2 0. Stationarity gives Py o' (8y) 2ds, 1&g 0. If > 0. then My = 0 and Pg o'(s)) = 2)8g. ie., Pg ((s,)
2 8, = O. complementary slackness allows 14 2 O and the stationarity reads Pg o (0) Me 0. Interpreting <(0+) lim o' o'(x)xx (possibly +oo), the inequality Pg x(0+) < 2A encodes the fact that activating 8g > would violate the KKT balance. Since
o' 0 and the objective is increasing in each 8g, the budget must be tight at optimum, hence
Σ, 82 = 2 and À > O. If @ is strictly monotone, 32 uniquely determines 8g as in 33
Lemma 6 (Memorization vs. spreading by o-monotonicity). Under the setup above and assuming ((x) = a' oxxx)x is continuous on o0):
(A) If é is nondecreasing (0, v2), then the unique maximizer is the memorization (peaked) solution
S 22,
Syyg"
realized by u 1 eg,
1 -
25

---

Published as a conference paper at ICLR 2026
(B) If ó is strictly decreasing (0, oo), then the unique maximizer spreads and is given by
Φ x-1(2))) 22
(for all y with 2\/Pg < <((0)),
Pg
and s 0 for any 9 with 2X/Ps 2 x(0+) (if (0+) xc). The multiplier À is
uniquely determined by the budget [g(s;)²= = In particular, if @0*) oo (e.g. ReLU sigmoid( sigmoid(2))), then
on [0, oo): x(z) 1/2; SiLU: x(x)=
sigmoid(()
all coordinates are strictly positive and
x
Pi Pj
O.
Proof. (A) Peaking when Φ is nondecreasing. Take any feasible with two positive coordinates
8t 2 > and pi > Pj. Define a squared-mass transfer preserving Σ ag. si(1) = V
sy(t) = v 55 t, and (t) = p:o(s, + (t)ss Then
V (t)
-
2 b1(p; (p)o(ss
0,
because (1) 2 8, (4) and á is nondecreasing. Hence v increases with I, so any two-support point can be strictly improved by pushing mass to the larger p. Iterating this collapse yields the single- support boundary Sn* 22, others 0. Uniqueness follows from strict inequality and the uniqueness of Pg-
(B) Spreuding when à is strictly decreasing. By Lemma 5 the optimal active coordinates satisfy Poo(s') (g) 2X. Since ó is strictly decreasing, o exists and is strictly decreasing, yielding sy = 6- (2X/p,) ) on the active set; complementary slackness gives the thresholding when <(0+) < oc. The budget Eg (s²)² = fixes X, and strict monotonicity implies the profile strictly ordered by Pg-
and assume o' (x) f x 0. For group arithmetic tasks, suppose we only collect sample (g.g- hh) for one target h with probability Pg. Then the global optimal of & a memorization solution, either (1) focused memorization
Theorem 5 (Memorization solution). Let (=x)==
V2 (eg-: -1) for g" arg max Pg if d is nondecreasing, (2) a spreading memorization with Eg 8g eg eg if do is strictly decreasing. Here 8g = 32>2(2\ and X is determined by
Σ 9 2. No other local optima exist.
Proof. The conclusion follows directly from Thm.
Some discussions. We know that
For power activations (z) x (q 2 have ((z)
q=-d2 nondecreasing: Thm. 6
gives memorization. In all these cases, the peaked solution is realized by even split
V2 fgr v v2 ) any profile s* can be realized with, e.g., Ug 1- $2/2.
ReLU on [0, oo): ((r) æ. o(x) 1/z strictly decreasing; Thm. [ yields 5" p.
SiLUU/Swshh/ Tt.LU//wwish//an//Sigoooo Tanh/Sigmoid:
strictly decreasing with ((0+) = oo; Thm. 6 B) gives
strictly ordered spread 8 @ (2\/pg))
G INTERACTIVE FEATURE LEARNING (SEC. 5.1)
G.1 FEATURE REPULSION (SEC. 5.1
Theorem (Repulsion of similar features). The j-th column of PB is given by:
K
[FB), = b, +
Σ bafe
(9)
where sign(b,i) = (fg Pr Pt. Here Pa--it == 1- F- FL F + -1
projection matrix constructed from F_ Pj which is F excluding the 1-th and j-th columns.
26

---

Published as a conference paper at ICLR 2026
Proof. Let Q : (F nI)- that preserves signs of the corresponding inverse entries), reorder columns so that the pair (j. E) becomes 21. Write the partition
Without loss of generality (by column permutation similarity
F fl 12 F.
F (1.2) E (=x(K-2)
Then the ridge Gram matrix G
FT
F n/K acquires the 2 x 2/ remainder block form
b
u
a: a:=ff fi +1
b = f f2 V f2
u:= f1
G
=
C
v
where
c c: f f2 +
F +nI.
H
Because > 0, H is positive definite and hence invertible. The inverse of block matrix is governed by the Schur complement. Define the 2 x Schur complement
u
a 8 8
S
=
H
Iu
v
=
7
where the entries are
a a -u, u,
B
==
A standard block inversion formula (e.g., via Schur complements) yields that the top-left x block of G-1 equals S-1 particular, the off diagonal entry of Q G-1 for indices (1, 2) the off diagonal entry of S Since
-3
a²
S-1
with a7 - 32
a7 32
-B
(because G > 0 implies S > 0), we obtain
B
912 = (S-1)12 = an B2"
It remains to identify o, 3,,, y in terms of ridge residuals with respect to F.. Note that
H F,T F, +n
= P. In par
by the definition Pr. == 1 F, -- Therefore
a= +
f, H F. 17 + f (I- = fT Pi
B = f. f F
fg = f F.H FTJ Fg = f, Ps.
7= 7 E.
Substituting these identities into the expression for Q/2 gives
f P
412
(17 f P.,, fi. (n fg Po.
(TT Pb)
The denominator is strictly positive (it is the determinant of the positive definite x 2 matrix S), hence
sign(y12) = - sign ( fo.
Undoing the preliminary permutation shows the same formula for the original indices (j. E), which proves the sign claim.
Finally, since Q is the inverse Gram with ridge, the j-th column of FQ is
K
(FQ).; = Σ Jmg im = E + Σ
Amj P.
m=l
Because qmj has sign opposite to the ridge-residual similarity FJ. Pr. features that are (resid- ually) similar to f, enter with negative coefficients and hence subtract from (FQ).; along those di- rections- "repellinn" similar features and promoting specialization. This completes the proof. I
27

---

Published as a conference paper at ICLR 2026
G.2 TOP-DOWN MODULATION (SEC. 5.2)
Theorem (Top-down Modulation). For group arithmetic lasks with o(x) = x², ihe hidden layer learns only a subset S of irreps, then the backpropagated gradient GF (#s @ ( 1 F (see proof for the definition of 2s), which yields modified Es that only has local maxima on the missing irreps $
Proof. Fix a nontrivial isotype (irrep) and we have
(()
(.h), = diag [R (x.)))))))))
Since Il, is central and idempotent, it commutes with Ra and 11 = II., hence
TT RT Rr R. TT.
Expand the central idempotent in the group algebra using unitary irreps {C&} and characters X
dx
Rg = de
1l =
Σ
(x(g))
Σ
Rg.
(34)
M
M
9EH
9EH
Therefore
dR
dR
1I
Σ (g
Rg
Σ xx(g)
RgN.
M
M
gEH
9EH
Taking the diagonal after the left shift by RT gives
de
diag( R (TI,
Σ (k(g) diag Rí Rygh)
M
DEH
Since RT, RgN Rh-ight we have
h-igh' e, 1Ms
diag Ryhh) =
0,
otherwise.
Only the unique term g hh'-1 survives, so
dt (hhh
de (x(h"""h)) 1m
diag RR (11))w))
1M
M
M
where we used (x(a) (x(a-)) for unitary irreps. Consequently,
dk M xa('1- 1
(k) (rows block hy
Summing subset S of isotypes yields
(R)
da (h) 1M-
(rows block h),
(rows for block A)
M KES
KES
Since summing over all h: + leads to & - (Thm.
3 for the residual ÿ - ý we have
y y mmws block i),,
Σ de (x)xx((( 1M.
M
which means that & - = =b 8 IMs where $ = =[V
E (I)
KOOORES
Since y = Pf 3 1M we have:
GF G x (V F= - 1 F (0s s (0s 8 1m) F
28

---

Published as a conference paper at ICLR 2026
Therefore, the energy function & now becomes
Es b [ @ 1 -1
where z [.,] = RR,, 5) CM defined in Eqn.
Computing each row in 4 and use the
property of projection matrix 11, (Eqn. 34). we have:
du
M V dx
[
Xu R,) S)
(11,, S)
M
hEH
In the Q space, we have (T2,, = E tr(Sk,r) and therefore
M
M
Es 2 k dk
Σ
(II Sy
Σ KA0.AES
Σ tr tr()k.r)
2
d2
which is exactly the same form as the decomposition (Eqn. 19) in Thm. (but a much cleaner derivation). Therefore, all the local maxima of &s are still in the same form Thm. 2 but we just remove those local maxima that are in isotype/irreps. A E S, and focus on missing ones.
G.3 DIVERSITY ENHANCEMENT WITH MUON (SEC.
Lemma (Muon optimizes the same as gradient flow). Muon finds ascending direction to maximize the joint energy finction Ejoine(W) = ", E(w;) and has stationary points iff the original gradient Gw vanishes.
Vw, E, EwgE,..., Vw El be the gradient matrix. Let G UDV be the singu- lar value decomposition. Then Muon direction is G = UV and thus the inner product between & and G
Proof. Let [V=[vwi
(G, (G,G)= tr G) =
tr(D) 2 0
(35)
So Muon always follows the gradient direction and improve the objective. Furthermore, (Ộ, G) 0 iff D = O, which means that G = 0. So the stationary points of the Muon dynamics and the original gradient dynamics are identical.
Lemma (Proposition of Fréchet log- log-Gumbel selection). Let I1. ,. In be i..d positive random variables with Fréchet(a) CDF
F(z) (exp(=x-1) x²a,,
0,
O,
and let Wi. Wn > 0 be fixed weights. Define
arg max Wj Zj.
15j5n
Then
w
Pr (i" =)
i= 1, ...... n.
Σ/= w
In particular, when a
Pr (i" =)
W/
Σラ wj
Proof. Set Yy = Wjzz. For / O,
Pr(max Y, 5(
exp
Σ
(w)/1)²)
Wy
==1
Differentiating gives the density of the maximum:
feax (t) d Pr (ax Y, 5t)
Σ
exp >
(wy (wz/t)a)
a
1=1
29

---

Tk- Published as a conference paper at ICLR 2026
The density that "i achieves the maximum at level t" is
(t)
II FI
w t-- exp - -
(wz/t)a
S#1 w.y
j=1
Hence the conditional probability that i is the argmax given max, Y = t is uff a--11
Pr(i" = max Y, =t) ΣE= a auy t- Σド= w
3
which is independent of t. Averaging over i yields the stated result.
Lemma 8 (The properties the dynamics in Equ. [O]. The dynamics always converges to St. for 1* arg max; (0), That is, the initial leader always win.
Proof. Note that due to orthogonality of {ci). the dynamics can be written as
0 Mya
M4 > 0,
with the constraint Σラ az < 1. Define
ry = pzag.
Interior. the interior, we have
ř = njā = py (2)0a))) =r.
For any pair i, k define the ratio Its derivative is
r1 ==
_
r
Ti
Ti
Pik
? rk = Pikk (r: rk).
Tk
Tk
Equivalently,
Ti
log
(1)
dt
T
Thus if (0) > r0), then d log(rs/rz) in r cannot be overtaken in the interior.
> 0 and pej(t) is strictly increasing. Hence a strict leader
zt
Boundary region C3 az 1). On the unit sphere, the projected dynamics is
àj =1j03 Aajj
maa 3
A=
In terms of ry,
= rr; - v),
2
k
k=1
k=
For the ratio Pi re/rk again obtain
d
Ti
Pu = (ra(Ta - -x))
log
ri
Tk.
(2)
dl
Fkk
Monotonicity of ratios. From (1)- r2())) > ry (0) then
Tt log
> Vt.
dt
rj
so (t) = (t) is strictly increasing for every j * E. Thus strict leader remains the unique leader for all time.
30

---

Published as a conference paper at ICLR 2026
Convergence to the vertex. Define weights
r3
Wi ====
Σ wj = 1.
2
Their dynamics is Taking ratios,
új = 200((rg - 1).
d
wi = 2(r; Tk).
log
di
In particular, we strictly increasing for every j * l. Therefore implying wx(() + 1 and wy(t) - 0. Hence
W (t)
+ 0 (++))
We (t)
(xtt) + c2
as t - 00.
Lemma (Muon projection). For the matrix A = 1Q, v] where Q is a column orthonormal matrix and v is vector with small magnitude, its Muon regulated version Â [A1, v) takes the following form:
v =
V + 1+ |ial
+ +0([v1113)
(36)
||vl||
where VIi = =QQTv v and v = - QQQ Tv.
Proof. Given A [Q, B] with QTQ= Ik. write = QC+ + Bi where C = BQQ Rixm and B1 = (1 - T)
Let = BI 0. For c > 0 define
[ A (A A)-1/o
A1o) = [A A...
We derive a first-order (in C) formula for the last block A/co.
The exact Gram matrix is
T C CTC
C
G = A =
Go + H,
: T),
H =
CT co.
Treat C as small, To first order in C we may drop the quadratic block:
0 CT
c + O(ICT1)..
И
Diagonalizing T'. Let T' T with A diag(A)..... Am), A3 0. Define the block orthogo- nal change of basis
P diag(Ik, U)
P:=PTGP :
diag(Iz, i==PTTPP -e
ÕT 0
where ő = CU. All first-order statements can be done in this basis and then mapped back by P.
First-order Taylor Expansion. Now let's do the Taylor expansion. Write
- To H Gb/2 (I [+ T. 1/2 H Go
1/2
Go
1/2
E
31

---

Published as a conference paper at ICLR 2026
Since Go diag(Ik: A),
ČA-1/2
E 1
is O(ICT))
0
For the scalar function f(x) = x 1/c,
(1+ E)-1/c I I--E -BB
+ O(1
c
Therefore
- (I E) -1/c G-1/22 G-1/0
G- EG G-122 1/2
O(IC11))
Compute the blocks using Go 1/2 diag( Lk. A -1/2):
0 ČA-1
To E
Go
1/2
A-1CT
Hence, to first order,
0-1/c=
Ik O
0
1
ČA-1
O(1[C11)).
(37)
A A-1/c
C
A
A
Back to the original space. Now
G-1/c- = PG-1/cpi
Using 37 and P diag(La, U),
I 0
0
CUA-
G-1/c=
+ 0(11)19).
0
U UA-1/EU
c UA- TT CT
Since A A T and A
1/eUT T/c
1k
0
1
CT
G-1/c
+ O(ICCI)).
0
T-1//-
T--CT
0
Now multiply
700 1Q, QC + BBB
G-1/e J
Taking the last columns (the 2nd block) and keeping first-order terms:
Alel = o( CT- + (QC + BB)TT + O([[II)
Bi 1/c + (cr--Ve-_CT-) + O(1CI)²).
c
Factor the Q-part columnwise via the spectral calculus 7f IfT = T then on each eigenvalue À the scalar factor
1 1-1
1 1- 1-11/
--c/
C 1-X
Thus, in matrix form,
1 _ - (T-T)-1.
and we have
100) B,T 1/c + By (1-1
(1 - + o(I[II).
(38)
where Bj
QQT B.
For polar case c = 2. the operator becomes (1 - T1/3) T)) -1 For = V, we have T
B_ B ||v1]3 1| and the conclusion follows.
32

---

Published as a conference paper at ICLR 2026
Lemma 10 (Bound of To).
To 2 max
min 1/pt Σ 1/1
(39)
Proof. To 2 min; 1/pr since the expected time to collect all the coupons is always larger than collecting the rarest coupon alone.
To prove To 2 E' 1/1, fix and consider the function
h(p) = log(1 - --),
p > 0.
direct computation shows
8²
k"(p) 4sinh (pt/2) 4sinh2ppt/2)
<0,
so h is concave. By Jensen's inequality and P; =1
log(1 e s Llog( -111).
Exponentiating gives the pointwise bound
II (1 - p-i S (1 - 11-6-
1=1
Therefore
000 (1 (1-(1---11 dt.
B[TT] 2 0
To evaluate the integral, set t//L dt - -Ldu/u and t: 0 + maps u: 1 0:
[. (1-(1-0011 dt [ 1- (1 4)th du = /
L-1 Σ w)'du du Σ i+1
L-1
Thus the conclusion holds. Equality holds if and only if pı case of equality in Jensen.
PL 1/L, since that is the
Theorem (Muon rebalances gradient updates). Consider the following dynamics TTan 2023
= A(w)w, ||w]|2 S
(10)
where A(w) == ENW(()GI Assume that (1) {c,} form orthonormal bases, (2) for w Σ, Si we have \e(w) = totl with hn S 1, and (3) {ai} initialized from inverse-exponential
distribution with CDF(2) (x))=xp(<>> with a > 1. Then
exp(-x=
Independent feature learning. Pr(w + Ci] = pt = U4/E, Then the expected #nodes to get all local maxima To 2 max (1/min) (//min pt, Etm11/1)
Muon guiding. If we use Muon optimizer to optimize K nodes sequentially. then the ex- pected #nodes to get all local maxima T. = 2To + - L.. For large a, Ta L.
Proof. From Lemma & know that the final mode ç; that the nodes converge into is the with
largest initial ay:
Pr(w - Ci] Pr(( arg max or (0)
(40)
By Lemma 7 we have Pr(w - C. = == 4.
Independent feature learning. In this case, getting all local modes (C,) } is identical to the coupon collector problem with L coupons. With the property of the distribution (Lemma 7). we know that the probability of getting (-th local maxima is pr / E,r4.
33

---

Published as a conference paper at ICLR 2026
Therefore, the expected number of trials to collect all local maxima is (Fajoolet et al. 1992):
4.00
To
1- I-II e --I----
dt
(41)
[=1
Note that To 2 max (1/ min pt, L. optimized, need K ~ To to collect all local maxima in K hidden nodes with high probability.
EkE
1/1)
(Lemma 39). Since each node is independently
Muon guiding. Consider the following setting that we optimize the hidden nodes "incrementally". When learning the weights of node we assume all the previous nodes (node 1 to node 1) have been learned, i.e., they have converged to one of the ground truth bases {(i}. but still keep the gradients of them (after deduplication) in the Muon update. Let Sy-1 S [L] {1...., L}} be the subset of local maxima that have been collected.
By Lemma know that
8%
B
8y.11
+ +(1gg, (13)
(42)
1
1
where gy.11 Pj-IP_-183 B and g,, gj formed by the previous nodes. Since
Bj.ll-
Here PJ-1 ISelvEs; the projection matrix
الدووو S ||g - Σ Maaoaaiill
Lail
af ofils S
(43)
We have |g T+ 1/2. Therefore, this means that the parallel components, i.e., the components nodes in the gradient was suppressed by at least 1/2,
that are duplicated with the previous compared the orthogonal components (i.e., the directions towards new local maxima). This is equivalent to dividing ii for all is that appear in Pj-1 by (at least) 2. By Lemma 7 for the node j. the probability converging to new local maximum other than Sy-1
Eigs,-1 Pt
Prew.S; 1 12 2-a ElES,-1 Pi + Eips,- P1
(44)
We do this sequentially starting from node j. then node j + 1, etc. Let m = |Sy-1l / be the number of discovered local maxima. Then the expected time that find local maxima is:
E
2 2 "T + 2
(45)
Prew. Pye
where mE[ = 1/ L/gs Egs Pt is the expected time for the original coupon collector problem to pick new local maximum, given S known ones. Adding the expected time together, we have
Tg T -T=% + (1-2-9))
(46)
Note that all the expected time are conditioned on the sequence of known local maxima. But since these values are independent the specific sequence, they are also the expected time overall.
H MORE DETAILED ANALYSIS ON STAGE I (LAZY LEARNING)
To analyze the Stage more thoroughly, we consider the gradient-flov dynamics of the output layer weights V.
Let i. E B be a fixed feature matrix and Y E R
We assume throughout that
(A1) F has full column rank K, and
(A2) col(Ý) S col(F), i.e. there exists V E KRKM such that ÿ FV*
34

---

Published as a conference paper at ICLR 2026
(A3) Small and independent random initialization entries of V(0), with mean zero and vari- ance a², where < 1, and thus |V(])) ((avKM) with high probability.
(A4) Zero-mean centering: 1 F and y 0.
Note that (A4) is optional. It simplifies some interpretations but is not needed for the main analysis. We train a linear readout V E IRKXM by minimizing
J(V) = 2 ( -V
2 O.
(47)
We define the (matrix) prediction error and the backpropagated gradient GF
B(I)
- RO)=YPFvO E RaxM,
Gz(1) == ((I)V(I) R
(48)
Note that in the main text, we use GF to denote the backpropagated gradient on the uncentered feature matrix F, i.e. GF = P-- Gg. where Pf := 1 11 /n is the zero-mean projection matrix along the sample dimension. As we will see below, the leading term of G3 is Y Y F and thus
GF P
(49)
because F = 0 and y = 0.
We consider continuous-tine gradient flow for V:
dV(()
V vvJ((V())
(50)
dt
The gradient of J with respect V is
= = FT
Y) nV AV - B, A: + FT=
(51)
We study the gradient flow dynamics
dV
Vv) = +B. di
(52)
Define the error matrix and the backpropagated gradient on F
E(() Ý FV(() E RR
Ge(t) = E(t)V(I) E RRaX
Our goal is to understand:
1. the small-time expansion of (t) and show that the leading term is ÿÿT F; and
2. the long-time decay behavior of G (tt), for both 77 0 and 17 >> 0.
H.1 THE DYNAMICS OF GF AT INITIAL. TIME STAMPS
H.1.1 SMALL-TIME EXPANSION AND LEADING TERM
Write the Taylor expansions 0 as
V(t) Vo tVi + 0(t),
E(t) Eg + +tE +(t)),
where Vo = V(() and Eg = y FV. From (2).
V1 = dV
= AVo + B (FT TTKKVo + FTV.
(53)
dt t=
Diffcrentiating E(t) = FV gives
dE
E dl 1=0
-
+ v PFTV. -F
(54)
35

---

Published as a conference paper at ICLR 2026
Now expand (t()):
Gp(() = (())V(L)) = (Fo +
+ +0(3) = FOV. +MEOVT + +
Using Eg = Y - FVVo and V from (53)
BoVi = (8 (V)) (-V ( F + n(k) +Y F)
-Y FV- VV, (F nIK) + F
Every term except Y F contains (at least factor of) Vo, hence is O(a) in Frobenius norm. Moreover, V also contains Vo:
E V = =(FFTF 1 -FFTVO
V | ((a) well.
We therefore obtain the small-time expansion
(p(() VVo +L +Ri(V) + O(f))
(55)
(
where R1 (V)) collects all order-t f terms that contain Vú and thus satisfy || (Vo)ll = O(a).
H.1.2 WHY T F IS THE LEADING TERM
We now compare the deterministic term Ý F to the Vo de-deennd terms using norm inequalities. Lemma 11 (Lower bound on | F) Let F have full column rank and be nonzero. Then
FYlF 2 Omia I!Y ||F > 0,
where miin (P) is the smallest singular value of F.
Proof. For any matrices A, B, ||AB tr(BBTAA ATA). PSD with eigenvalues bounded below by J (F)2,
A). Take == ÝÝT, B == F. Since BB is
[AB tr(BB ATA) A) 2 Jmin (F) = J (F).
Taking square roots gives the result.
Next, bound the Vo-deppedent part. For concreteness, consider the term FFT VV (other mixed terms are bounded similarly). Using ||A|| S ||
|IFF VV, [|x I |FFT Volle
Under the iid initialization with variance a V F = Olav hence
|F YV, || (aa).
The same argument applies all other Ve-deeendent order-! terms in Ri (Vo). Combining Lemma 11 with these upper bounds yields
Ri (V))
S C(
for some constant C independent of Thus, in the limit 0 (small random initialization), the term T F is the unique leading contribution at order I.
Proposition 2 (Small-time leading term of Gp). Under assumptions (AI) and small random initialization with scale <
Gp(t) YVo + +t)
in Frobenius norm. In particular. as +
Ge(1) -V. - and
(C+())l=~+[l\ Flle
for fixed small independently of whether nj 0 or 0.
36

---

Published as a conference paper at ICLR 2026
Remark the role of 7/. The weight decay parameter only appears in products involving Vo, and hence all dydeepenennt order-t contributions are also ((a) in norm. Therefore the leading deterministic term Y the same for both and > 0.
H.2
LONG-TIME DECAY OF Gp
We now analyze the behavior of G (t) as t + again for both 17 0 and 17 0.
H.2.1
GENERAL SOLUTION OF THE GRADIENT FLOW
From 522_
the gradient flow is a linear ODE with constant coefficients. The unique fixed point V
satisfies
AV = B + V ==
Define (V = V(t) V*
Then
d
AV(() = -VV())))) (V(() AAV((),
dt
and hence
V()) (V(0)))
(56)
Let Amin (A) denote the smallest eigenvalue of A. Since A = FT F and F has full column rank, (A) 2 Jmin ( for 1j = 0 and Amin (A) 2 min (F22 + v Standard bounds on matrix exponentials give
AV((I Amin Ammin (A)e |AV(O)]
(57)
The eπror satisfies
E(t) y Ý FV(t) - FV -VVV =: - FAV(t),
where E" == and 577
FV' is the residual at the minimizer. Using | (FA((I) S (IFI))
[E(()) F*Il $ [| Amnn |AV())]
(58)
H.2.2 CASE 1) 0
When 17 0, we have A FTF. By assumption (A2), y FV* is exactly realized by the model, so E = 0, i.e.
lim E(t) 0.
to0
Equations 57 and (58) imply exponential decay:
[[V(1) V* S mmmm (I)(0) - V*
|B(]| & [10111e-0mm0(f) (0 V/
omin
We can now bound Gp(t):
Gp(t) (tt)V(t)
so
||G < ||((]] [V S |E(O)l (IIV + [|((t) V*ll2). Using the exponential bounds above and the fact that (|V(t) - V* S [V( V*IlE, we obtain
(59)
(t) emmul (F)
for constant Co depending F, V* and V()) but not on t.
Thus in the realizable, unregularized case, the backpropagated gradient decays exponentially to zero. Proposition 3 (Exponential decay of G for = 0). Assume (A1)-(A2) and f = 0. Then lim (t) O,
and there exists Cy > 0 such that
[[G(())FF S - aminn(F))
for all 2 0.
more refined analysis using the SVD f UEWT shows that every singular direction of Gp(()) a finite linear combination of exponentials e (of+ and o so the slowest rate in the Frobenius norm is indeed - (
37

---

Published as a conference paper at ICLR 2026
H.2.3 CASE 17
When 7) > 0, the minimizer V* = A is the ridge solution. general it does nor exactly
interpolate Y, and the residual
E" = y FV
nonzero. Consequently the limiting backpropagated gradient
G : lim (t(t)= E 1-00
also nonzero in general.
To study the convergence, write
GG(t) G = -E( - (B(1)
+ -V())))
Using
57
58 and ||BB| = ||A1|=||B1|2 we obtain
(1) 1 & [[((() - E* = (V) - 1*
< (1 VT [V + [|] v V* e
Aman (A)e
Since [V(l2) is bounded (it converges to |2)))2), this shows exponential convergence of (t) to G Therefore, we have the following proposition:
Proposition 4 (Exponential convergence of Gp for 17 > 0). Assume (A1)-(A2) and > 0. Then
lim Gz(l) Gj V in general, t-000
and there exists C 0 such that
[GG - F s Oie Amin -mmiaapp
(A) o (f)) +
Finally, note that
G = E. = P,ýV = n + + -
(60)
where Py, = I F(FT
n = nFFFT + +nIh=1
by Woodbury matrix formula.
Summary.
For small t, the leading term in Gt(t) is +ÝÝTF, independent of All other terms (including those involving (0) and 17) are lower order in the initialization scale
For 17 = and realizable Ý col(F)) both the error E(t) and Gt)(t) decay exponentially to zero at rate at least amin ( (
For > O, Et) and (tt) converge exponentially to (E*, V*), and G converges exponentially to a nonzero limit G -*V*T
H.3 SCALING THE NETWORK OUTPUT BY 3
We now modify the objective by inserting a scalar factor 3 > 0 on the network output:
J(V) - -BV FERaX e RaxM V VERKXM
(61)
We study the gradient flow dynamics for V and, in particular, the initial behavior of the backpropa- gated gradient on F.
Gp(()= (t)) E(t) = (F)
As before, assume small random initialization V(0) V, whose entries iid, zero, with variance a², that [|V||| O(avKM) for < 1.
38

---

Published as a conference paper at ICLR 2026
Gradient flow with output scaling. The gradient of J with respect to V is
VvJ(VV = 3FF (BFV - 8) = F FV - BFTV.
(62)
The gradient flow is
dV -
(63)
We define
F(() - -BV((),
Gp(1) = F(O)(())
At t = 0,
Vo VVO),
Eo = E(0) & 3.V..
Small-time expansion of (t) We expand
V(t) = Vi +Vi + O(t),
E(t) = Eo + tEi +0(P),
dE -= From where V dV lt=0 and E₁
63
= -22 FV + BTT..
(64)
Differentiating E(I) Y -FV(() gives
E1 dE -3FV 8 -V- B.FFY. dt 1=
(65)
Now expand ()):
G1(() E(I)V(I)
[Eo + (V)) + + +0(2)
-VVV + (EEVVV + +.V. + +(P))
Zeroth order in f. Using Eo Ý - 3FVV.
(0) = VoVo = VVV - 3FV.V.
(66)
Thus at we have 0(3)) term and an O(31) term, both proportional to the random initial- ization.
First order in We next compute VVV and E, Vo and keep track of all powers of B.
From
64
VjT -3 F BÝ F. hence
V T
FT
F + (.
=V(-8VV FT B)-BBV(-PVVTTTT
F)
P BVV. PT + BVVVVO FT - PPFVVTT.
Vo. order /
From 65
V.V. = =IFFFTVVO -FFTTVV
Adding these,
V.V. E.V. = F
+ (F FF
(67)
Collecting terms. It is convenient to separate the Vn-independent and Vo depeenenn contributions. Define
(2(V)) = F FF
Ra(V)) = FVOV FF
39

---

Published as a conference paper at ICLR 2026
Then 66 can be written
V + E.Vo = =BY & + Ra(V).
Thus the small-time expansion of (p(
(F V. 3VVOV. +t BBYYYTF B R2(V)) + 1Vo(( + +(P)). (F
(68)
Orders in B and interaction with small init. Under the small-init assumption VVolle (lav KM) and bounded F, Y, we have
[22 = O(a),
[ O(1VV 12) -(a²).
Therefore, for fixed B,
B² =
1VVLl = -(09g))
In contrast, the Vo- Vo-iindepnndent term BY F has norm
|BYY Fll B
with strictly positive factor Y F|]FF whenever Ý * 0 and F has full column rank.
Hence, in the regime of small random initialization (a << 1), the leading O(t) contribution to ( is
L while the de-denend corrections are suppressed by factors 3Pa and a3a
Similarly the final fixed point of G F. (+00) can be computed by replacing Ý with 8/3 Eqn. 60
G G B2 G
(69)
Therefore, if 8 > 1 is large, then the final fixed point of Gp (++00 leads to a smaller backpropa- gated gradient GF and may delay grokking.
I DISCUSSION
Relationship between grokking and regularization. key insight of the Lig framework that a non-zero backpropagated gradient Gp initiates the feature learning process, setting the stage for optimizing the energy function & that determines which features learned and how. Intuitively, as long as the lazy learning phase produces a sufficiently large nonzero Gp correlated with the target label Y, feature learning will be triggered.
Several ways exist to ensure this nonzero gradient. Regularization mechanisms, such as weight decay, help guarantee the presence of nonzero Gr, and conveniently allow us demonstrate the feature learning process in mathematically rigorous way. However, we want to emphasize that regularization is not the only option. For example, GF might become large during lazy learning and later vanish due to perfect overfitting (see Eqn. Notably, (Prieto et al. (0224) replaces the standard softmax (exponential term) with a stable softmax (linear term), which from the perspective of the Lig framework, slows down overfitting to the labels. This maintains nonzero backpropagated gradient, thereby preserving the potential for subsequent feature learning.
In summary, our framework shows that weight decay can be sufficient condition for feature learning (and grokking), but not a necessary one.
Two different kinds of memorization. From the analysis, it is clear that memorization of grokking stems from overfitting on random features, which is different from landing on memorization solu- tions following the feature learning dynamics due limited/noisy data (Thm. 5). From this view,
40

---

Published as a conference paper at ICLR 2026
grokking does not switch from memorization to generalization, but switch from overfitting gen- eralization.
Flat versus sharp optima. Common knowledge often regards flat optima generalizable solu- tions, while sharp optima as memooriztiion/vvvriiingg From the point of view of Li2, sharp optima happen when the model overfits random features (Sec. and thus small changes of the weights will change the loss lot. On the other hand, we can prove that the local optima from & are flat (Corollary D. and thus small changes of the weights certain directions do not change &. the model is wver-parameterize. then multiple nodes may learn the same (or similar) set of features, which create freedom for the loss function be flat. If we learn memorizing features due lim- ited/noisy data (Thm. 5) then more nodes need to be involved to "explain" the target well, and the overall weights will appear to be less flat.
Small versus large learning rates. According to our analysis, in Stage I, we need large learning rate to quickly learn ridge solution V so that the backpropagated gradient G becomes meaningful to trigger Stage II. In Stage II, the best learning rate depends on the amount of data available. With a data, we can afford larger learning rate to find the features quickly. With limited data, we may need smaller learning rate to stay in the basins of generalizable features (Fig. 7). which may be contradictory to common belief.
J MORE EXPERIMENTS
J.1 MORE EXPERIMENTS ON PHASE TRANSITION
0.4
0.35
0.30
0.3
Test accurscy
ttanns
-0,6
0.6
0.20
0.15
0.2
0.2
0.1
80
100
120
140
100
120
Figure 8: Generrliiaation/memorraaaao phase transition in modular addition tasks (cont.). Left: Product group 74 & &7, Zg 07. a, 7g 872 Zo, F150711 Right: Non-Ahelian groups with max; dę = (maximal irreducible dimension 2). These non-Ahelian groups are generated from GAP programs (See Appendix Sec. J.2
7s &7 (7728777 072872, & 82, Z 2g @0 Z1r, Z & @ & Zg..
J.2 USE GROUPS ALGORITHMS PROGRAMMING (GAP) TO GET NON-ABELIAN GROUPS
GAP (https:// www. www.gaa---pa. system.smmmmmmmm org/) is programming language with library of thou- sands of functions to create and manipulate group. Using GAP, one can easily enumerate all non- abelian group of size M 127 and create their multiplication tables, which what have done here. From these non-Abelian groups, for each group size M, we pick one for our scaling law experiments (Fig. bottom right) with max; d2 2.
J.3 ZERO-INIT ACCELERATES THE FEATURE LEARNING PROCESS
41

---

Published as a conference paper at ICLR 2026
train%: 0.9, wd 0.0001
train%: 0.8, wd 0.0000001
train%: 0.6, wd-0.0001
0.6
0.6
0.6
muon
muon
muon
0.4
0.4
0.4
Test
0.3
0.3
0.3
0.2
0.1
0.1
0.0
0.0
0.0
60
80
100
120
60
80
100
120
60
80
180
120
train%: 0.9, wd=0,0002
train%: 0.8, wd=0,0002
train%: 0.6, wd=0.0002
0.6
0.6
0.6
adam
adam muon
adam
0.5
muon
muon
loss
0.4
0.4
0.4
0.3
0.3
0.3
0.1
0.1
0.1
0.0
60
80
100
120
60
80
100
120
60
80
100
120
train%: 0.9, wd=5e-05
train%: 0.8, wd=5e-05
train%: 0.6, wd=5e-05
0.6
0.6
0.6
adam muon
adam muon
adam
0.5
0.5
0.5
muon
Teat loss
0.4
0.4
0.4
0.2
0.2
0.1
0.1
0.1
100
100
100 Hidden size
Hidden size
Hidden size
Figure 9: Adam versus Muon optimizers in modular addition tasks with M 71, when the number of hidden nodes K relatively small compared to M. Muon optimizer achieves lower test loss compared to Adam.
42

---

Published as a conference paper at ICLR 2026
M
WOOL
WOOM
2layer
AWWKAY
VAMNAAA
MU
ANNMAIN
VOW
Lodoo 12500 15600 Eaecy
WWAW
DAWA
S4aype networe
NAMMAAI
MA000N
WWKAW
WMM
10000
wwww
MMN
MO00A
5000
30000
AWWN
Llayer nctMAr
2.0
beberbaw whm
M4*
WWW
mooool
W00
Ladk
Ialeand
Mande MOMNN aile. VANWIA
5000
1828s
15006 17500
Training modular addition tasks with and layer network with ReLU activations. Left: Training accuracy and losses. Right: Learned features the lowest layer. With more layers, the training takes longer and grokking (delayed generalization) becomes more prominant. However, features at the lowest layer remain (distorted version) of Fourier bases, which are consistent with the analysis in Sec. o
Figure 10:
10-
Trair
1.25
1.0
train Jecc
3.8
1.00
0.8
1.0
2.6
0.75
Acruney
0.6
Urain Joss
0.4
tast loss
0.4
task loss
0.6
0.4
02
0.25
0.2
3.0
0.00
0.0
0.0
M- 89
M4-127
train
-10
test_act
0.8
0.8
trann Joss
0.6
train Joss
0.6
test Joss
tedl
0.4
0.
0.0
20
60
80
100
20
100
Epoch
Epoch
Figure 11: Training modular addition tasks with real weights (M = 23, 41, 89, 127). Learning rate is 0.005.
weight decay is 5e - 5. Number of hidden nodes K = 256. Test sample is 20% of the full set of 2 Using Adam optimizer. Averaged over 5 seeds. This is baseline.
43

---

Published as a conference paper at ICLR 2026
M=41
10-
trai: acc
1.0
train acc
0.
test 302
0.8
Test.
0.8
04
5.6
tran,
teok loss
0.6
0.4
0.4
Lask loss
Urair Joss
0.6
02
0.2
0.2
0.0
0.0
0.0
0.0
100
Mi M4=127
trair
train
1.0
0.8
train loss
train Joss
0.8
Joss
teat
0.4
02
0.2
0.0
0.0
20
100
Epoch
Eporh
Figure 12: Training modular addition tasks with complex weights (M = 23, 41, 89, 127). Learning rate is 0.005, weight decay is 5e - 5. Number of hidden nodes K = 256. Test sample is 20% of the full set of M Using Adam optimizer. Averaged over seeds. Compared with the real case (Fig. II. models with complex weights seem to grok faster.
M=23
M=41
1.0
trair aCC
1.0
frain, acc
1.50
lest, scc
2.0
test acc
-15
Nrcusery
1.00
Loss
test
00.50
0.25
60
80
100
20
4.
60
80
100
Epoch
Epoch
1.0
trair.
1.0
train_300 30c
1.0
0.8
0.8
0.8
0.6
fram,
0.6
Accurery
0.6
train Joss
0.6
24
teat loss
8.4
Lend loss
0.4
82
0.2
00
0.0
0.0
20
100
weights (M 23, 41, 89, 127). Instead of using gradient descent to update the top layer V. in every gradient update we use ridge regression solution Viage with respect to the current (Proposition []. Learning rate is 0.005, weight decay 5e - 5. Number of hidden nodes K 256. Test sample is 20% of the full set of M Using Adam optimizer. Averaged over 5 seeds. The grokking still happens (for M 23 check Fig. for completeness). is slower for M 23 but actually faster for M 41, 89, 127, compared the baseline (Fig. IT).
Figure 13: Training modular addition tasks
44

---

Published as a conference paper at ICLR 2026
M=23
1.0
train acc
test acc
2.0
0.8
Accuracy
0.6
loss
0.4
test loss
1.0
0.2
0.5
0.0
0.0
100
200
300
400
500
Epoch
Figure 14:
Training modular addition tasks real weights 23 for 500 epochs, using Vridge as top layer weight. The grokking still happens but slower than the baseline (Fig. for M = 23.
45

---

Published as a conference paper at ICLR 2026
top_layer, faaee__nit_muutipplle M=41 %6train=00%
(Mm41, test size=0.5)
Accuracy (Mw41, test size=0.5)
1.25
muttppler-O ftes)
1.0
muttipler-1 Itest)
1.00
muttppeers
multiplirr] Brain)
Accuracy
0.75
0.50
muttiller-O Itest
mutiplier-1 ltesti
0.25
0.2
muttilier-O Itain)
mutipller-1 (tain]
0.00
0.0
1000
200
400
600
800
1000
Epoch
Epoch
top layer init multiplier: M=41
%6rainn40%%
test 0ize-0.6)
Accuracy (M=41, test size=0.6)
multipierro (eet)
1.0
1.0
mulipiir=3 (we)))
multipler-O (rain!
0.8
multiplor-3
Accuracy
multiplie -0 [test]
multisfier1 [tsstt
multifiir- [trai]]
0.0
0.0
Epoch
Epoch
top Lop_laye__mee h___jaye_ini__muiiiieee jnnit_iiiiiipiirrr:..... 0 vs M=41 %6train= 30%
Loss (M=41, size=0.7)
Accuracy (M=4], size=0.7)
1.25
mult oulerleer-- (test)
1.00
multslier-1 (test)
1.00
mui pliemo (main)
Accuracy
0.75
0.75
0.50
0.50
mu mutpllen-00 (test)
0.25
0.25
mutilllono
mutiplioe11 (test)
0.00
0.00
mutillerr11
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Epoch
Epoch
top layer init multiplier: 0 vs M=41
%5rain=25%
Loss (M-41. test size-0.75)
Accuracy (M-41, test test_ size-0,75)
1.25
muit isliere
mukelier=1 (as))
1.00
1.00
multilier-d (tain)
muttstler-11 (trii))
Accuracy
0.75
0.75
0.50
0.50
malipieedd (est)
0.25
0.25
mitplier11 (test)
mutipiieed (rwin)
mulipiee-1
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Epoch
Epoch
Figure 15: Zero-init accelerates the feature learning process for 41. Check Fig 17 for the hyperparam- eters setting. Blue lines are zero-init for (i.e., (0)) 0). red lines are regular initialization.
46

---

Published as a conference paper at ICLR 2026
M=89| %6triin50%%
(M-89, test, size=0.5)
Accuracy (M-89, size=0.5)
murtislier-0 test
1.0
1.00
muttsller-1
0.75
muttillerna
must ptlerel [tran))
Accuracy
0.50
muttipltr-O Itest!
mutiplier-1 ltesti
0.2
muttiliir-0 Itain)
muttiller-1 (trin]
0.00
0.0
600
800
1000
200
400
600
800
1000
Epoch
Epoch
top layer init multiplier: M=
%6rainn40%%
test size-0.6)
Accuracy test,
1.25
multipierro (eet)
1.0
mulipiir=] (we)))
multipler-O (rain!
0.8
multpplor-3
Accuracy
0.75
0.50
multiplier -0 [test]
multisfier1 [tsst]
0.25
0.00
0.0
multifiir- [tran]]
Epoch
Epoch
iiviipiittmuliiplieer. vs M=89| %6train= 30%
Loss (M=89, test size=0.7)
Accuracy (M=89, size=0.7)
1.0
0.8
multpfer=O best
Accuracy
mulipperro (tes))
muttpller11 (test)
0.6
muttpper-1 (best)
muttpler-O (trai)
muttpler- (tran)
0.5
0.0
0.0
400
800
1000
200
400
800
Epoch
Epoch
top layer init multiplier: 0 vs M=
20triin=0%%
Loss (M-89, test size-0.8)
Accuracy (M-89, test size-0.8
plie-= Bestt
1.0
mu mulpliar-1 Bestt
multipiis-00 (tain)
Accuracy
0.8
Loss
tillerl (trii)
0.6
malipiieed Best)
ma mutplier11 (test)
mutipieedd (owin)
mulipiee-1
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Epoch
Epoch
Figure 16: Zero-init accelerates the feature learning process for Check Fig 17 for the hyperparam- eters setting. Blue lines are zero-init for (i.e., (0) 0). red lines are regular initialization.
47

---

Published as a conference paper at ICLR 2026
top_ @op_layeriinit_mauee init_muulipplier vs 1 M=127| 55train=50%
size
Accuracy sizem0.5)
1.0
multislier- lest!
1.0
0.8
muttsller11 (test)
mutt pl [tran)
Accuracy
0.8
0.6
muttislerr1 itrain)
0.6
0.4
multiplie -0 [testt
multidfier11 (test)
0.2
0.2
mutiilier-0 Itrin)
mutipller-1 Itain]
0.0
0.0
400
800
1000
200
400
600
800
1000
Epoch
Epoch
top layer init multiplier: 0 vs 1 M=127| AMtrain=40%
Lass (M=127, test size=0.6
Accuracy (M-127, test size=
mult siliir [esl)]
1.0
1.00
multifiir- Bestt
0.8
munh pleerm (trik)
Accuracy
0.6
0.50
0.4
mubiplier-0 Itest
muttplierm1 [test]
0.25
multilie-- Itrin]
multisfier (rain)
0.00
0.0
400
600
800
1000
200
400
600
800
1000
Epoch
Epoch
top_ layer muiiiiiiiilliir: vs 1 Mm 127 99trrinn 30%
Loss (M-127 127, test size=0.7)
Accuracy (M-127, test size siem0.7)
1.5
multipiler-0 Ites
1.0
muttipterr1 (test)
muttppter-3
multipler-O (rain)
Accuracy
0.6
0.8
0.5
0.4
mutiifiar- Best
muttilier-1 Itstt
murtidlter-d Itrin)
mutipller11 Itai)
200
400
600
800
1000
200
400
600
800
1000
Epoch
Epoch
top_ hop_laye__iaae po__laye__nit__meee multipliiiii:: 0 vs 1
117
9ttrain=88%
Loss (M-127, test size=0.82)
Accuracy (M= 127, test size=0.82)
1.00
multplian11 (test)
1.5
mul m. plier-p (rww)
mul mollppier11 (trin))
Accuraci Cauurac
0.75
1.0
0.50
mutpiler-0
0.25
mutplier11 (teet)
mutipiieed Drwin)
0.0
mutipiirr1 (trin)
2500
5000
7500 10000 12500 15000 17500 20000 Epoch
2500
5000
7500 10000 12500 15000 17500 20000 Epoch
top layer jnit mullippler:: vs 1 M=127 %6train=15%
Loss (M-127, test size=0.85)
Accuracy (M-127. test sizem0.85)
2.0
1.5
0.8
multpfer=O
Accuracy
1.0
multpler-1
muttillrra crami
muttpllr-11 (trin)
mutiplierod
muttllian11 (tss))
muttlllr-0 (trai))
mutplier-1 (train)
0.0
2500
5000
7500
10000 12500 15000 17500 20000 Epoch
2500
5000
7500
10000 12500 15000 17500 20080 Epoch
Figure 17: Zero-init accelerates the feature learning process for M 127. Learning rate 0.005. Using Adam optimizer. 7) = 5e 5. K 1024. (xx)) = 3 and MSE loss. Trp.layerr.nntt.mll:pl..... means (0))= = 0, i.c., the top layer is initialized to he (blue lines). To btayye_init.mmutiipp..e = 1, then V is initialized regularly (red lines). It is clear that pero-initializationn leads to faster grokking (and feature learning) than regular initialization, in particular when data are scarce.
48

---

Published as a conference paper at ICLR 2026
top layer init multiplier: MM=1 %strain=80% other layers=1 layer_layers==
Loss (M-41, %0r, other_ ayers=1) (epochs 0-1999)
Accuracy (M-41, Strain-80% other, layers-1) (epochs 0-999
muttplier-0 ites)
1.0
muttitiir-1
muttisliar-0 itrin))
Accuracy
multsller-1 Itrain)
Loss
multisfie-O Itest
multillier11
0.2
muttplier-O Itrin]
0.0
multisfier [rrk])
500
1250
2000
200
600
Epoch
Epoch
top lop_laye__aae intiiinittttin multiplier 0 41 %train 80% laaher_layers22
(M=41, 80%, other, layers=2) (epachs 1999)
Accuracy (M=41, S4train=80% 80%, other Jayers=2) (epochs
mull plarró Best] multpliers1 Itest
6
multplirr] Itrai))
multipfier -1 [testt
muttilier-O [test]
4-
pller=1 [train]
2-
0
750
600
Epoch
Epech
top layer init multiplier: 0 M=
%srriin=0%% other layers=1 other_layers==1
% other_ layers=1)
Accuracy %train-80%, other, layers-
2.5
muiteller-0 Ites)
1.0
muttiliir-1 Ites)
muttitiir--
Acouracy
muttisller11
0.6
1.0
multisfierO0 Iestt
murtipliar11 Itsst
muttiller-- (trin)
multisfiar (rain)
0.0
250
500
750
1250
1500
1750
2000
200
400
600
ad0
Epoch
Epoch
top layer jili_initttttt multiplier: 0 vs 1 M- 71 %train - 80% other_ layers-2
7M=71, 9strain=80% other Jayers= (epochs 1999)
Accuracy (M=71, 9strain=80%, other Jayers= (epochs 0 0-999)
multpller-C (test)
muttplerm1 (est)
muttiplern (tain)
mudipperr]
Accuracy
mutiiller-0 (test)
muttilier-1 [tes]]
0.2
multilier-- (trin)
0.0
250
750
1250 Epoch
1750
2000
200
800
1000
Epoch
Figure 18: Zero-init accelerates the feature learning process multi-layer setting. other_layers means adding an additional hidden layer in between top layer V and bottom layer W with residual connection, etc. Learning rate is 0.005. Using Adam optimizer. = le - 4. K 2048. axx) = ReLU(z) and MSE loss. www.iiinniitmteeeeeee..em lines). If Top_aaer..inn www.iiiiitiitiiiommmmmmmmmmmiimt mpltir zero-iiiiialization leads to much faster feature learning than regular initialization multi-layer settings (often 10x).
means (0) = 0, i.e., the top layer is initialized be 0 (blue then initialized regularly (red lines). clear that
49


Reviews:
[{"review_id": "aCOq4joPpA", "reviewer": "Reviewer_LnxJ", "summary": "The paper proposes **Li²**, a mathematical framework to explain grokking (delayed generalization) in two-layer nonlinear networks via the structure of the back-propagated gradient matrix (G_F). Training is decomposed into three phases: **(I) Lazy learning**, where (G_F) is effectively random and the top layer overfits random hidden representations; **(II) Independent feature learning**, where each hidden unit learns independently because each column of (G_F) depends only on its own activation—weight decay injects label signal, the dynamics become exact gradient ascent on an energy (E), and the local maxima of (E) coincide with the emergent features; and **(III) Interactive feature learning**, where hidden units begin to interact and (G_F) reorients toward missing features that must be acquired for generalization. On group-arithmetic tasks, the authors analyze when these energy-induced features generalize, their representational power, and how they vary with sample size. The framework yields **provable scaling laws** for memorization and generalization as functions of weight decay, learning rate, and data size, and offers a first-principles explanation for the effectiveness of optimizers such as **Muon**. The analysis is argued to extend to deeper architectures.", "questions": "See weaknesses.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. The paper is overall solid and offers a detailed study of the dynamics of two-layer neural networks.\n2. It provides experiments that validate the theoretical results.", "weaknesses": "1. The writing feels very rushed: many symbols are undefined or unexplained, making the paper hard to follow and overly dense.\n2. The theoretical setup is quite restricted; for example, a projection function is deliberately designed so that the hidden layer receives gradients that are random noise.\n3. It is unclear whether the group (arithmetic) example pertains only to Stage II or also extends to Stage III.\n4. The theoretical analysis of Muon appears disconnected from the main body of the paper, and the setup and assumptions of Theorem 8 are entirely unclear.\n5. The relationships among the three stages are not well articulated; the analysis reads like heuristic case-by-case treatment rather than a genuinely unified three-stage dynamics analysis.\n6. The abstract mentions scaling laws, but it is not evident where in the paper this is actually developed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ahK71Ntm8y", "reviewer": "Reviewer_sYi6", "summary": "The paper analyzes grokking dynamics in the presence of regularization. They identify three phases in the learning process - lazy learning regime and independent and interactive feature learning regime, and prove various theorems which govern different aspects of the training dynamics.", "questions": "1) Line 105-106 : I suggest modifying \"grokking mostly happens ... regularization\" to \"grokking is accelerated ... regularization\"\n2) Why are the axes cut off till epoch 300 in Fig 2 last column?\n3) Since the discussion on Thm 3 involves optimization in the complex domain, were any follow-up experiments done with complex weights?\n4) The existence of maxima for the ascent functions is interesting (Thms 4-5). What can be said about the network's ability to find these maxima?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The theorems proved are mathematically rigorous, with detailed proofs in the appendices. \n2) The scaling law analysis in Section 5.4 will be useful since it gives a first-principles derivation of the scaling phenomenon.\n3) Extensions to modern optimizers and deeper networks were adequately discussed.", "weaknesses": "1) My main concern is that the key observation (that there is a lazy and rich learning regime) was already reported in [1], who study this in the context of polynomial regression. While the current paper offers mathematically rigorous proofs, the distinction with [1] needs to be more elaborately discussed in the main text.\n2) The mathematical framework presented displays the three stages, but fails to offer more insight on what drives these transitions or when do these transitions. For example, based on lines 471-475, it seems like it's the (inverse) learning rate which sets the scale for when these transitions occur, but it would be great if this could be explained in more detail.\n3) How does one relate the top-down modulation in Sec 6 to the continuous feature learning observed in [2]?\n4) Discussion on limitations not found in main text.\n\n\n[1] Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics, 2024. \n[2] Gromov, A. (2023). Grokking modular arithmetic. arXiv preprint arXiv:2301.02679.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6gZeGtrD5O", "reviewer": "Reviewer_NQrJ", "summary": "This work proposes a 3-stage framework to understand the emergence of feature learning in large-width two-layer neural networks, named **Li$_{2}$**. This consists of:\n\n1. Lazy regime: first-layer weights are effectively random. The second layer fits the data using random features. The back-prop term to the hidden layer, $GF$, carries little usable structure and, without weight decay, can vanish at the lazy fixed point (ridge solution).\n2. Independent: With weight decay \\(\\eta>0\\), \\(GF\\) acquires target structure. Under some assumptions and large width \\(K\\), it simplifies to\n$$\nGF = \\frac{\\eta}{(Kc_1+\\eta)(nc_2+\\eta)}\\tilde Y\\tilde Y^{\\top}F,\n$$\nso each neuron follows gradient **ascent** on the single-neuron energy\n$$\nE(w)=\\frac{1}{2}\\big\\|\\tilde Y^{\\top}\\sigma(Xw)\\big\\|_2^2,\n$$\ni.e., neurons learn useful directions independently.\n3. Interactive: As several features are learned, neuron–neuron interactions reshape $GF$: similar features repel, and the signal is steered toward missing directions. \n\nThis is discussed explicitly in a modular arithmetic task and quadratic activation $\\sigma(x)=x^{2}$. Using group-theoretical tools, it is shwon that local maxima of $E$ align with group irreducible representations which for this task coincide with Fourier modes, yielding closed-form descriptions of learned features and their attained energies $E^\\star$. \n\nFinally, an extension to the multi-layer case is also discussed.", "questions": "1. Different results in this work appear to rely on different assumptions about $n,d,K,M$, and it is not immediately clear whether these are mutually consistent. For instance, Lemma 1 assumes that $K$ is sufficiently large and that $x_i^\\top x_j=\\rho$, which is only possible when $n<d$ unless the $x_i$ are degenerate. By contrast, Theorem 2 and Corollary 1 take $d=2M$ and $n=M^2$, which is consistent with Lemma 1 only for $M<2$. Similarly, Theorem 4 requires $n\\gtrsim d_k^{,2} M \\log(\\delta/\\delta)$ via Matrix Bernstein. Overall, I found it confusing to determine whether the regimes considered for the different phases of Li$_2$ hold simultaneously.\n\n2. Feature learning for two-layer neural networks has been studied extensively in recent years, with several works arriving at a picture closely related to **$Li$_2$**. For example:\n\n- One line of work systematically analyzes one-pass SGD dynamics in teacher–student settings (a 2LNN learning another, not necessarily identical, 2LNN). These studies show that the dynamics decompose into plateau / saddle-to-saddle phases: first-layer weights move within a neighborhood of zero (the *mediocrity phase*), then individual neurons correlate with target directions independently before finally coalescing (the *specialization phase*); see Saad and Solla (1995), Goldt et al. (2019), and Arnaboldi et al. (2023). While these works differ technically (finite-width networks, one-pass SGD), the overall mechanism is closely related, with a key difference being the absence of an overfitting phase under one-pass SGD.\n\n- Another closely related line of work considers feature learning during the first few steps of GD with aggressive learning rates (Damian, 2022; Ba et al., 2022; Dandi et al., 2024). The energy in Eq. (8), often called the ``Correlation loss,’’ is a common approximation in this literature because it yields exact weak-recovery thresholds for the initial gradient steps. A notable observation is that, after a single aggressive step, the gradient can be asymptotically characterized—depending on sample complexity—by a rank-one matrix that correlates with the target, enabling the network to express nonlinear components with limited data. Is this the same mechanism as in Theorem 3?\n\n- More recently, Montanari and Urbani (2025) studied a related teacher–student setting for full-batch GD on a wide 2LNN learning a single neuron, and identified three timescales: a lazy timescale, a generalization timescale, and an overfitting timescale (motivating early stopping). (There is no specialization here because the target has a single neuron.) In Li$_2$, overfitting is associated with the lazy phase. How can this be reconciled with Montanari and Urbani (2025)? How should we understand the benefits of early stopping within the Li$_2$ framework?\n\n**References**\n\n- (Saad and Solla 1995) Dynamics of On-Line Gradient Descent Learning for Multilayer Neural Networks. NeurIPS 1995.\n- (Goldt et al. 2019) Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. NeurIPS 2019\n- (Arnaboldi et al. 2023) From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks. COLT 2023\n- (Damian 2022) Neural Networks can Learn Representations with Gradient Descent. COLT 2022.\n- (Ba et al. 2022) High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation. NeurIPS 2022.\n- (Dandi et al. 2024) How Two-Layer Neural Networks Learn, One (Giant) Step at a Time. JMLR 2024\n- (Montanari and Urbani 2025) Dynamical Decoupling of Generalization and Overfitting in Large Two-Layer Networks. arXiv 2025", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and provides a fairly generic picture of the mechanism for feature learning in two-layer neural networks, which is task independent and rely only on an investigation of the gradients.", "weaknesses": "The results for each section rely on different assumptions, which makes **Li$_{2}$** look more like different patches of results rather than a unified picture. Some of the observations appearing in this work have also appeared in other works in the feature learning literature, and more throughout comparison is lacking. I expand on these two points below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VfFiROjr46", "reviewer": "Reviewer_aGxx", "summary": "The paper proposes a three-stage framework to theoretically study grokking using two-layer neural network with MSE loss and quadratic activation functions. It find connections between stage II and the optimization of an energy function. It rigorously analyzes the role of learning rate, weight decay, sample size, and proposes a scaling law of generalization/memorization boundary. It also theoretically analyzes the benefits of Muon optimizer in the feature learning regime.", "questions": "1. In stage I, why the activations F are mostly unchanged? Any empirical observations to support such assumption? I guess it needs assumptions on the learning rate, weight decay, and initialization scale of V.\n2. In Lemma 1, why can we assume W always follow normal distribution at each step? Besides that, W is assumed to follow normal, but in the proof (line 714-716), you assume w_i follow N(0, I), which is stronger.\n3. In Lemma 1, why can we assume $<x_i, x_{i'}> = \\rho$ ? Does this assumption hold in any synthetic tasks that show grokking?\n4. In Theorem 5, what are focused and spreading memorization? What's the difference between memorization and overfitting?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The first work to rigorously analyze the role of learning rate, weight decay, and width on grokking dynamics.\n2. Interesting and novel analysis of the interactive feature learning regime.\n3. First theoretical analysis of Muon in the feature learning regime.\n4. The theoretical results of two types of memorization is intriguing.", "weaknesses": "1. MSE loss is uncommon in the training of deep learning models, though it's fine for the ease of theoretical analysis.\n2. The theory needs a nonzero weight decay to provably show the three phases, while in practice, weight decay is unnecessary for grokking to occur.\n3. Several assumptions need to be justified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)