A Stochastic--Geometric Theory of Scaling Laws in Grokking

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

A Stochastic--Geometric Theory of Scaling Laws in Grokking

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Preprint

A S TOCHASTIC –G EOMETRIC T HEORY OF S CALING L AWS IN G ROKKING

arXiv:2606.30388v1 [stat.ML] 29 Jun 2026

Róisı́n Luo∗ Research Ireland – Centre for Research Training in AI J.E. Cairnes School of Business & Economics University of Galway Jonas Ngnawé Université Laval Mila – Québec AI Institute

Christian Gagné Université Laval Canada–CIFAR AI Chair Mila – Québec AI Institute

Ihsan Ullah School of Computer Science & Data Science Institute Visual Intelligence Lab University of Galway

Karyn Morrissey J.E. Cairnes School of Business & Economics University of Galway

A BSTRACT Delayed generalization (i.e. grokking) refers to the phenomenon in which a neural network fits its training data early in training but only begins to generalize after a prolonged delay, often through an abrupt transition. Despite extensive empirical study, its underlying mechanism remains poorly understood. In this work, we first theoretically characterize a shell–core topological configuration of the reachable solution space induced by Adam’s optimization dynamics with weight-shrinkage regularization, supported by empirical evidence. This optimization-induced topological configuration gives rise to grokking. In model’s parameter space, random initialization solutions concentrate on a thin outer spherical shell, enclosing another spherical shell of memorization solutions, which in turn contains a core corresponding to the generalization solutions. Leveraging stopping-time theory, we then analyze the geometry of this topological configuration and the solution transition time at which optimization trajectories escape the memorization manifold and first reach the boundary of the generalization manifold. Our theoretical analysis derives grokking scaling laws for the learning rate, batch size, and ℓ2 regularization coefficient, which are further validated through experiments and shown to recover results from prior literature.

I NTRODUCTION

Neural networks trained on noise-free, highly structured learning tasks have been observed to exhibit an epiphany phenomenon known as delayed generalization, or grokking (Power et al., 2022). In these settings, models rapidly achieve near-zero training loss, often early in training, yet fail to generalize for an extended period before abruptly transitioning to strong test performance. This behavior has been reported across a range of tasks with exact underlying structure (Hwang & Park, 2026), such as modular arithmetic (Power et al., 2022; Liu et al., 2022; Zhong et al., 2023), algorithmic reasoning (Nanda et al., 2023), and group-theoretic learning (Stander et al., 2024; Notsawo et al., 2026). Moreover, prior work has shown that the onset of grokking is highly sensitive to hyperparameter choices (Zhong et al., 2023), including the fraction of training data, learning rate, batch size, and regularization coefficient (Power et al., 2022; Liu et al., 2022). Despite extensive empirical investigation, however, its underlying mechanism remains poorly understood. ∗

Correspondence to: [email protected]

Preprint

200

Late stage

100

Equilibrium

Early stage

100

200

100

(a) Weight-Shrinkage Grokking

(b) Trajectories w/ S5

200

100

200

Figure 1: Optimization-Induced Shell–Core Topology and Empirical Evidence. Left (a): grokking dynamics under weight-shrinkage regularization induce a shell–core topology in parameter space, where initialization points √ θ0 ∼ N (0, σ 2 Ip ) concentrate on a hyperspherical shell Θ with √ radius ρΘ ≈ σ p and thickness σ/ 2, enclosing the memorization shell M \ G and the generalization core G; the radii ρG , ρM , and ρΘ are determined by the task and optimization configuration. Middle–right (b–c): empirical evidence on S5 and Z127 using isometric manifold learning (e.g., MDS) to visualize ten optimization parameter trajectories, where θ0 , θτM , and θτG denote initialization, first memorization, and first generalization states. The projections show that Adam-reachable M and G are approximately symmetric embedded in ambient space Rp .

Empirical findings consistently suggest that training dynamics in the grokking regime exhibit threestage behavior (Power et al., 2022; Nanda et al., 2023; Kumar et al., 2024). In the initial stage, the model rapidly interpolates the training set and reaches the memorization regime. In the second stage, the trajectory remains near the memorization solutions for an extended period before escaping and eventually reaching the generalization regime. The length of this extended period scales with the task, learning rate, batch size, and regularization coefficient. In the third stage, drift and diffusion equilibrate in the optimization-induced dynamics, and the trajectory settles within the generalization solution region centered at its minimizer. Under the regime with ℓ2 weight-shrinkage regularization, these observations imply the existence of two distinct classes of solutions: memorization solutions, which only interpolate the training data and does not generalize, and generalization solutions, which additionally achieve low loss with respect to the underlying data distribution. Their topological configuration is thus induced by optimization dynamics, and characterized into three stages: the early stage, the late stage, and the equilibrium, as shown in Figure 1a. The early-stage dynamics are driven by deterministic drift from large gradients, which dominates the optimization trajectory and rapidly drives it toward the memorization manifold. At this stage, the gradient fluctuations are anisotropic. The late-stage dynamics are driven more slowly by decaying but nonzero gradients, together with isotropic gradient fluctuations, which guide the optimization trajectory across the trivial memorization solutions. Once the trajectory reaches the generalization set, the drift and diffusion in the optimization dynamics equilibrate and confine the trajectory within the generalization manifold, referred to as the equilibrium. Figure 1b–1c provide empirical evidence consistent with this optimization-induced topological configuration by visualizing optimization trajectories via isometric manifold learning. This optimizerinduced topological configuration, as illustrated in Figure 1a, gives rise to grokking and forms a shell–core configuration: initialization solutions concentrate on a thin hyperspherical shell, which encloses a shell of memorization solutions containing a compact core of generalization solutions. Detailed empirical results of Adam–induced shell–core radii are provided in Appendix A.2. We theoretically analyze the geometry of this topological configuration and derive scaling laws for the manifold radii and the solution transition time from memorization to generalization through a stopping-time analysis (Dynkin, 1965) of stochastic differential equations (SDEs) (Øksendal, 2003). The remainder of the paper is organized as follows. Section 3 presents initialization concentration, and Section 4.1 establishes the joint-state SDE for Adam (Kingma & Ba, 2015) designed to analyze grokking dynamics. Building on this formulation, Section 4.2 derives the scaling law for the 2

Preprint

memorization-manifold radius. Under the late-stage reduction observations of grokking, Section 4.3 derives a closed-form preconditioned radius SDE for the evolution of manifold radii. Sections 4.4– 4.5 derive the scaling laws for the generalization-manifold radius and the solution transition time from memorization to generalization. Finally, Section 5 validates the theory through experiments on modular arithmetic and group-theoretic learning tasks, and by comparison with the grokking literature. The contributions are summarized below: 1. Optimization-Induced Shell–Core Topology. We introduce an optimizer-induced shell– core prior for grokking, where initialization concentrates on a thin hyperspherical shell enclosing memorization solutions, which in turn contain a compact core of generalization solutions. 2. Stochastic–Geometric Characterization. Leveraging the topological prior, we formulate grokking as a stochastic transition problem and characterize the geometry by the first stopping time at which stochastic gradient flows escape the memorization manifold and reach the generalization manifold. 3. Empirical, Literature, and Symbolic-Algebra Validation. We validate our theory through experiments on group-theoretic learning and modular arithmetic tasks, and by comparison with the literature, recovering known scaling laws for the learning rate, batch size, and regularization coefficient. We additionally machine-check all closed-form and asymptotic results appearing in the proofs by a symbolic algebra system (SAS) with Julia.

N OTATIONS

Data and Network. Let Dtr := {(xi , yi )}ni=1 denote a finite training set of size n, where each pair (xi , yi ) ∈ X × Y is drawn i.i.d. from the empirical distribution PbX Y , with underlying true distribution PX Y . Let fθ : X → Y be a model parameterized by θ ∈ Rp . Throughout, we use Xt to denote the time-indexed value of a variable X at time t; for example, θt denotes the parameter at time t. Instance, Batch, Empirical, and True Loss. Let ℓf (x, y; θ) denote the instance loss evaluated at a point (x, y). Let Lf (ξ; θ) denote the batch loss evaluated on a mini-batch ξ = {(xξi , yξi )}bi=1 cf (θ) := Eξ∼D [Lf (ξ; θ)] denote the empirical loss evaluated over the training set of size b. Let L tr Dtr . Let Lf (θ) denote the true loss with respect to the underlying data distribution PX Y . Initialization, Memorization, and Generalization Manifolds. We study the solution manifolds reachable under the dynamics induced by Adam (Kingma & Ba, 2015). The ambient parameter space is Rp endowed with the Euclidean metric. Let Θ denote the set of initialization solutions cf (θ) < ϵ ∧ pAdam (θ | θ0 ∈ under some initialization distribution PΘ . Let M := {θ ∈ Rp | L Θ) > 0} denote the set of memorization solutions under Adam’s dynamics with a loss tolerance ϵ. Similarly, let G := {θ ∈ Rp | Lf (θ) < ϵ ∧ pAdam (θ | θ0 ∈ Θ) > 0} denote the set of generalization solutions. Throughout, ∂S and S ◦ denote the boundary and interior of a set S, 2 respectively; ∂M is the outer2 boundary of M \ G, 2and ∂G is its inner boundary. Let ρM := 2 Eθ0 ∈Θ,θ∈∂M ∥θ∥2 | θ0 and ρG := Eθ0 ∈Θ,θ∈∂G ∥θ∥2 | θ0 denote the mean squared radii of the memorization and generalization manifolds, respectively.

I NITIALIZATION C ONCENTRATION

We first demonstrate the topological concentration of network-parameter initialization under normal (k) (1) and uniform distributions. Consider that θ0 = (θ0 , · · · , θ0 ) ∈ Rp consists of k sub-vectors such P Pk (j) k that θ0 ∈ Rpj where j=1 pj = p. Define effective dimension p̃ such that p̃1 := k1 j=1 p1j . For (j)

coordinate-wise initialization [θ0 ]i ∼ N (0, σj2 ), the initialization radius square ! pj k k X X X 2 (j) 2 2 ρΘ (θ0 ) = ∥θ0 ∥2 = [θ0 ]i ∼ σj2 χ2pj , j=1

i=1

j=1

(1)

Preprint

follows a sum of scaled chi-square distributions, where each σj2 χ2pj has pj degrees of freedom. Hence, the radius of Θ admits the closed-form expression k

σ̃ 2 :=

1X pj σj2 , p j=1

Eθ0 ρ2Θ (θ0 ) = σ̃ 2 p = O(k),

k X k Varθ0 ρ2Θ (θ0 ) = 2 pj σj4 = O( ), p̃ j=1

√ under a coordinate scaling σj = O(1/ pj ) (LeCun et al., 1998; Glorot & Bengio, 2010). This suggests that initialization solutions concentrate on a thin annulus of radius O(k) with thickness p O( k/p̃) in parameter space. Similar concentration also holds for coordinate-wise uniform initial(j) ization distribution [θ0 ]i ∼ U(−εj , εj ), k

ε̃2 :=

1X pj ε2j , p j=1

pε̃2 Eθ0 ρ2Θ (θ0 ) = = O(k), 3

k k 4 X Varθ0 ρ2Θ (θ0 ) = pj ε4j = O( ), 45 j=1 p̃

√ under a coordinate scaling εj = O(1/ pj ). Proofs for the normal and uniform concentration results are provided in Appendices A.3 and A.4, respectively.

M ANIFOLD R ADIUS AND S OLUTION T RANSITION A NALYSIS

We characterize Adam’s dynamics through its continuous-time limit, formulated as an SDE on the joint state of the parameters and first- and second-moment estimates. Based on this joint-state SDE, we derive the scaling law for the memorization radius ρM , defined as the parameter norm at the firsthitting time of the empirical-loss boundary ∂M , via perturbation method (Bender & Orszag, 1999). Further leveraging two reduction properties of the grokking dynamics, we reduce the intractable joint-state SDE for Adam (Section 4.1) to a tractable radial SDE. This reduction enables an analytical characterization of the generalization radius ρG , defined as the parameter norm at the first-hitting time of the generalization boundary ∂G, and the solution transition time E[τM →G ], defined as the mean time required to transition from memorization boundary ∂M to generalization boundary ∂G. A high-level sketch for theoretical analysis framework is illustrated as in Figure 6 of the appendix. 4.1

A DAM ’ S J OINT-S TATE C ONTINUOUS -T IME SDE L IMIT

We consider Adam dynamics in the small learning rate and large batch size regime η → 0, b → ∞, with exponential moving average (EMA) coefficients (β1 , β2 ) ∈ (0, 1)2 and an ℓ2 regularizer R(θ) = 21 ∥θ∥22 with coefficient λ > 0. Let gt := ∇[Lf∗ (ξt ; θt ) + λR(θt )] denote the regularized mini-sample gradient flow, with per-sample mean ḡt := Es [∇ℓf (s; θt )] and covariance Σt := Covs [∇ℓf (s; θt )]. For sufficiently large batch size b, the mini-batch gradient admits the dis tributional limit gt ∼ N ḡt , Σt /b . We refer to this setting as coupled weight decay, since the ℓ2 regularizer contributes to the gradient used to update Adam’s first- and second-moment estimates. In contrast, the decoupled weight decay scheme applies weight decay as a separate parameter-shrinkage step and is commonly referred to as AdamW (Loshchilov & Hutter, 2019). Adam’s Parameter Update Rules. Let θ ∈ Rp be model’s parameters. Let m, v ∈ Rp be the firstand second-moment estimates, respectively. Let m̂, v̂ be the m, v with bias corrections. Adam’s discrete update rules at iteration k are given as mk+1 = β1 mk + (1 − β1 ) gk ,

vk+1 = β2 vk + (1 − β2 ) (gk ⊙ gk ),

m̂k+1 = mk+1 / (1 − β1k+1 ),

v̂k+1 = vk+1 / (1 − β2k+1 ), p θk+1 = θk − η m̂k+1 ⊘ v̂k+1 + ε 1p ,

with initial values θ0 ∈ Θ, m0 = 0, v0 = 0, where ⊙ is element-wise product and ⊘ is element-wise division. m̂ and v̂ are Adam’s bias corrections, which compensate the warm-up underestimation of (mk , vk ) caused by the zero initialization m0 = v0 = 0. We drop Adam’s numerical-stability constant ε throughout the analysis as it does not affect the analysis. Defining the Adam joint state as St := (θt , mt , vt ) ∈ R3p and using the continuous-time interpolation t = ηk, we state the resulting joint-state SDE for Adam in Lemma 1. The proof sketch is illustrated in Figure 7. The 4

Preprint

proof is provided in Appendix A.5. The technical correctness is verified through an induced radius SDE with Itô’s lemma in Appendix A.6. Lemma 1 provides a stochastic-theoretical framework for studying grokking dynamics induced by Adam optimizer. ⊤ ⊤ Lemma 1 (Adam’s Joint-State Continuous-Time SDE Limit). Let St := (θt⊤ , m⊤ ∈ R3p t , vt ) denote the joint state of Adam, where θt , mt , and vt are the parameters, first-moment estimates, and second-moment estimates, respectively. Taking the continuous-time interpolation t = ηk, the discrete Adam updates admit the following Itô SDE limit on the joint state r η dSt = µ(St ) dt + σ(St ) dWt , (2) b where dWt ∈ Rp denotes the infinitesimal increment of a Wiener process adapted to the filtration generated by the mini-batch sampling process {ξt }, with drift factor µ(St ) ∈ R3p and diffusion factor σ(St ) ∈ R3p×p √     √ 0 − B(t) mt ⊘ vt 1 − e−α2 t 1/2     − α1 mt − ḡt µ(St ) = , (3) , σ(St ) = α1 Σt , B(t) := 1 − e−α1 t − α2 vt − g ⊙ g t α2 Dt where αi := (1 − βi )/η and B(t) is the bias-correction factor in continuous time and g ⊙ g t := ḡt ⊙ ḡt +

1 diag(Σt ) ∈ Rp , b

1/2

Dt := 2 diag (ḡt ) Σt

∈ Rp×p .

(4) −1/2

In particular, to decompose this SDE into a tractable form, we define π(θt ) := diag(g ⊙ g t ) ∈ Rp×p as Adam’s preconditioner. We write θt (θ0 ), θt (S0 ) to denote the Adam’s θ-evolution starting from initial states θ0 ∈ Θ, S0 ∈ ΘS , respectively, where ΘS is the initialization distribution of S. Remark 1. While the continuous-time SDE limit of ordinary stochastic gradient descent (SGD) is now well established (Li et al., 2017; Mandt et al., 2017; Luo et al., 2025), a rigorous continuoustime treatment of Adam that explicitly accounts for stochasticity remains, to the best of our knowledge, comparatively less developed. For example, the joint (m, v) deterministic limit of Adam appears in (Da Silva & Gazeau, 2020), and approximations for adaptive methods such as RMSprop and Adam were established in (Malladi et al., 2022; Compagnoni et al., 2025). Our joint-state SDE limit only partially overlaps with this line of work and yields an SDE limit that accounts for stochasticity without relying on a first-order approximation. 4.2

M EMORIZATION R ADIUS

⊤ ⊤ Starting from the initial state S0 = (θ0⊤ , 0⊤ p , 0p ) , Adam’s joint-state SDE (2) generates a stochastic trajectory {St } adapted to the filtration generated by mini-batch sampling. The first time at which θt reaches the memorization boundary ∂M defines the stopping time (Karatzas & Shreve, 1991) and mean squared memorization radius as τM (S0 ) := inf{t ≥ 0 : θt (S0 ) ∈ ∂M }, and ρ2M := ES0 ∈ΘS ∥θτM (S0 ) ∥22 , (5)

respectively. This defines a Dirichlet exit problem on the state St , with absorbing boundary ∂M and the infinitesimal generator (i.e., a differential operator) (Dynkin, 1965) induced by Adam’s SDE in Lemma 1 as: 1 η tr ΣS (S) ∇2S [•] , (6) LS [•] := µ(S)⊤ ∇S [•] + 2 b which contains an operator modulated by a coefficient η/b with ΣS (S) := σ(S)σ(S)⊤ . The resulting Dirichlet partial differential equation (PDE) admits no closed-form solution, we apply a low-order regular perturbation expansion (Bender & Orszag, 1999) in ε = η/b to obtain Theorem 1. The proof is provided in Appendix A.12. Experimental verification is provided for S5 in Figure 2 and for Z127 in Appendix A.15. Theorem 1 (Scaling Law of Memorization Radius (Perturbation Solution)). The memorization radius approximately admits the scaling law with respect to η/b as η 2 η η 3 (0) 2 (2) 2 ρM = (ρM ) + cM + cM + O , (7) b b b 5

Preprint

×104

×104 mean over seeds ±1 theoretical

3.4 3.2

2.6

2.5

4.0

2.0

3.8

2.4 2.2

1.5

3.6

2.0

1.0

3.4 0.2

0.4

0.6

0.8

1.0

×10 2

mean over seeds ±1 theoretical

2 M

2.8

×104

3.0

4.2

3.0

1.8

3.5

mean over seeds ±1 theoretical

4.4

100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

×10 3

Figure 2: Scaling Laws of Manifold Radius ρ2M on S5 . We show the scaling law of ρ2M with respect to the learning rate η, batch size b, and ℓ2 regularization coefficient λ on the S5 task. For each hyperparameter configuration, we train for ten runs. The results show that larger η/b modifies the dynamics with a stronger diffusion variance, whereas λ does not affect the diffusion variance. We also overlay the theoretical fits. As predicted by theory, the scaling law of ρ2M with respect to the learning rate η exhibits a U-shaped curve. An additional experimental results for Z127 are provided in Appendix A.15. (0)

(2)

where (ρM )2 , cM , cM are task-determined constants independent of η and b, as defined in Appendix A.12. The quadratic regularizer λ2 ∥θt ∥22 induces deterministic contraction. The leading con(0) stant (ρM )2 depends on the deterministic gradient flow ∇Lf (θt ), when the task gradient ∇Lf∗ (θt ) 2 (0) (0) 1 d ρ = −2λ dt. is small, λ enters (ρM )2 through d∥θt ∥22 = 2θt⊤ dθt ≈ −2λ∥θt ∥22 dt ⇒ (0) M 2 (ρM )

(0)

Thus, (ρM )2 ∝ exp(−O(λ)), and hence ρ2M ∝ exp(−O(λ)). Remark 2. In the radius-shrinkage regime, diffusion first shifts the hitting point inward, giving (2) cM < 0, while boundary curvature produces a stabilizing second-order correction, giving cM > 0. 2 Our theory predicts that ρM is U-shaped in η/b with two regimes: a sub-linear decrease ρ2M ≈ (0) (2) (0) (2) (ρM )2 +(η/b)cM for η/b ≪ |cM |/(2cM ), and a sub-quadratic increase ρ2M ≈ (ρM )2 +(η/b)2 cM (2) (2) for η/b ≳ |cM |/(2cM ), with minimum at (η/b)⋆ = |cM |/(2cM ). This U-shaped effect is observed in Figure 2, where the scaling law with respect to η exhibits a U-shaped curve under large b. 4.3

L ATE -S TAGE R ADIUS DYNAMICS

To characterize the manifold geometry, we derive the radius SDE induced by Adam’s SDE (2) via Itô’s lemma under the late-stage reduction observations. Empirically, as shown in Figure 3, once the trajectory {St } reaches the memorization manifold M (i.e., memorization regime), the latestage grokking dynamics exhibit two reduction properties. First, refer to equation (3), the first- and second-moment estimates mt and vt rapidly converge to ḡt 1 − e−α1 t and g ⊙ g t , respectively, at an exponential rate exp(−O(t)). We refer to this property as slow-manifold reduction, formalized in Observation 1. Second, refer to equation (4), the gradient fluctuations, characterized by Σt , become small and isotropic at an exponential rate exp(−O(t)). We refer to this property as small-isotropic gradient covariance, formalized in Observation 2. Observation 1 (Slow-Manifold Reduction on (mt , vt )). If (mt , vt ) varies slowly as t → ∞, the first- and second-moment estimates mt , vt rapidly converge to its mean-field limit (8) mt → ḡt 1 − e−α1 t → ḡt and vt → g ⊙ g t (1 − e−α2 t ) → g ⊙ g t , respectively, at an exponential rate exp(−O(t)). The proof for this mean-field limit is provided in Lemma 7 of Appendix A.7. Observation 2 (Small-Isotropic Gradient Covariance). In Adam’s late-stage dynamics, the gradient covariance Σt rapidly becomes small and approximately isotropic as t → ∞. With this isotropy and by equation (4), the Adam’s preconditioner is therefore approximated by π(θ) ≈ s(θ)−1 Ip , as t → ∞, where the state-dependent scalar s(θ) > 0 is defined by 1 1 −1/2 := tr π(θ) , π(θ) := diag(g ⊙ g(θ)) , (9) s(θ) p 6

Preprint

100 mt gk 2 vt gk gk 2

10 3

M G

10 4

Var[ ( t)]

10 2

10 2 10 3

10 5 0

5000

10000 t

15000

20000

(a) Reduction on (mt , vt )

10 4

Var[ ( t)]

10 1

M G

10 1 [ ( t)]

10 1 Residual

[ ( t)]

100

M G

10 3 10 5 10 7

5000

10000 t

15000

20000

(b) Gradient noise magnitude

5000

10000 t

15000

20000

Figure 3: Observations in Grokking Dynamics. The markers τM and τG denote the first-hitting times of ∂M and ∂G, respectively, and the y-axis is shown on a logarithmic scale. Left (a) shows the exponential reduction of moment estimates; (mt , vt ) rapidly converges to ḡt (1−e−α1 t ), g ⊙ g t (1− −α2 t e ) , with exponential decay exp(−O(t)). Middle–right (b–c) shows the exponential reduction of the gradient noise; the mean and variance of the singular values of the gradient covariance Σt decay as exp(−O(t)), indicating small and isotropic gradient fluctuations in late-stage dynamics. where p is the parameter dimension. In particular, let G(θ) := √1b B(t)π(θ)Σ(θ)1/2 be the preconditioned diffusion factor, which admits G(θ)G(θ)⊤ ≈ Ip . See Appendix A.10. Under Observations 1–2, Lemma 2 gives the approximate closed-form dynamics of rt2 = ∥θt ∥22 . The proof is provided in Appendix A.11, where we also experimentally verify that the lemma accurately characterizes the late-stage dynamics of manifold radii in Figure 10. Lemma 2 (Reduced Late-Stage Radius SDE). Let rt2 := ∥θt ∥22 . In Adam’s late-stage regime underlying grokking, assume that the diffusion covariance becomes isotropic G(θt )G(θt )⊤ ≈ Ip as t → ∞, then the reduced late-stage squared-radius SDE admits 2λ 2 2 ⊤ √ (r) drt2 ≈ − rt − θt ḡt − λθt + η p + RSM (t) + Rπ (t) dt + 2 η rt dWt , (10) s(θt ) s(θt ) where slow-manifold residual RSM (t) and preconditioner residual Rπ (t) are q √ ⊤ −α1 t RSM (t) := −2B(t)θt E [mt ⊘ vt ] − ḡt 1 − e ⊘ g ⊙ gt , (11) Rπ (t) := −2θt⊤ π(θt ) − s(θt )−1 Ip ḡt , (12) R t (r) respectively, and a one-dimensional Wiener process Wt := 0 er (θu )⊤ G(θu )dWu with er (θ) := (r) θ ∥θ∥2 by Lévy’s characterization: ⟨Wt ⟩t = t. Empirically, in the memorization regime, the residual

sum RSM (t) + Rπ (t) becomes negligible. Hence, we set RSM (t) + Rπ (t) = 0 in the subsequent analysis, see Appendix A.11. 4.4

G ENERALIZATION R ADIUS

Let SτM denote the state at which the optimization trajectory first reaches ∂M . We treat SτM as the initial state for the subsequent first-hitting problem (Karatzas & Shreve, 1991) on ∂G. Define the stopping time for first hitting ∂G and the corresponding mean squared generalization radius as τG (S0 ) := inf{t ≥ τM : θt (S0 ) ∈ ∂G}, ρ2G := ES0 ∈ΘS ∥θτG (S0 ) ∥22 . (13) Once the trajectory enters ∂G, the drift and diffusion in equation (2) equilibrate, and Adam’s SDE approaches its equilibrium. The center of this equilibrium starting from S0 is the regularized local minimizer λ θ⋆ (S0 ) := arg min Lf∗ (θ(S0 )) + ∥θ(S0 )∥22 , (14) 2 θ(S0 ) where Lf∗ (θ(S0 )) is task loss. This yields an equilibrium-distribution problem for the reduced latestage dynamics around θ⋆ (S0 ), initialized from the post-memorization state SτM (S0 ). Linearizing the reduced dynamics around θ⋆ (S0 ) solves Theorem 2. The proof is provided in Appendix A.13. Experimental verification is provided for S5 in Figure 4 and for Z127 in Appendix A.16. 7

Preprint

×103 5.1

5.6 5.4

mean over seeds ±1 theoretical

5.0 4.5

5.0 2 G

2 G

×103 mean over seeds ±1 theoretical

2 G

5.8

×103 mean over seeds ±1 theoretical

4.9

4.0

4.8

3.5

5.2 5.0 0.2

0.4

0.6

0.8

1.0

×10 2

100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

×10 3

Figure 4: Scaling Laws of Manifold Radius ρ2G on S5 . We show the scaling laws of ρ2G with respect to the learning rate η, batch size b, and ℓ2 regularization coefficient λ on the S5 task. For each hyperparameter configuration, we train for ten runs. The results show that larger η induces stronger diffusion variance, whereas λ does not affect the diffusion variance. We also overlay the theoretical fits. An additional experimental results for Z127 are provided in Appendix A.16. Theorem 2 (Scaling Law of Generalization Radius). The mean squared generalization manifold radius admits the asymptotic expansion 2 η (0) 2 ρ2G ≈ ρG + cG + O λη 2 , (15) λ (0) with (ρG )2 := ES0 ∈ΘS ∥θ⋆ (S0 )∥22 a landscape constant and cG > 0 a task-determined constant independent of η and λ, as defined in Appendix A.13; a weak batch size b-dependence enters through cG = O( √1b ). 4.5

S OLUTION T RANSITION T IME

Starting from S0 ∈ ΘS , the transition time on the trivial annulus M \ G is given as τM →G (S0 ) := τG (S0 ) − τM (S0 ),

(16)

where τM (S0 ) and τG (S0 ) are the stopping times that the trajectory first hits ∂M and ∂G, respectively, defined in Sections 4.2–4.4. Taking θτM (S0 ) ∈ ∂M as the initial condition (Theorem 1), we formulate the transition from ∂M to ∂G as a mean first-passage problem (Karatzas & Shreve, 1991) for the radial process rt := ∥θt ∥2 on the annulus M \ G. Its equilibrium distribution characterizes the equilibrium around the generalization region, while the mean first-passage time from the memorization radius to the generalization radius gives the solution transition time. Solving this first-passage problem yields Theorem 3. The proof is provided in Appendix A.14. Experimental verification is provided for S5 in Figure 5 and for Z127 in Appendix A.17. Theorem 3 (Scaling Law of Solution Transition Time). The mean number of optimizer iterations between the first hit at ∂M and the first hit at ∂G admits the asymptotic expansion (0)

τM →G := ES0 ∈ΘS τM →G (S0 ) ≈

(2)

η ρ cτ cτ s̄ log M + + O , + 2 (0) ηλ bλ λ λ3 ρ

(17)

where s̄ is determined by the harmonic mean of the effective deterministic gradient flow, while (2) cτ and cτ collect the trajectory corrections, as defined in Appendix A.14. All three are taskdetermined constants independent of η and λ; a weak batch-size b-dependence enters through the (2) scalar preconditioning scale s̄ = O( √1b ), and hence cτ = O( √1b ) and cτ = O( 1b ).

VALIDATION

We validate the three scaling laws in Section 5.1 by directly measuring ρ2M , ρ2G , and τM →G on the group-theoretic learning task S5 and modular arithmetic tasks over Z127 , and further compare them with published results from the grokking literature in Section 5.2. 8

Preprint

×104

×104 mean over seeds ±1 theoretical

1.2

×104 mean over seeds ±1 theoretical

1.2

1.0

1.1

0.8

1.0

mean over seeds ±1 theoretical

1.2 1.1

0.6

0.9

0.4

0.8

M G

1.0 0.9 0.8 0.7 0.6 0.2

0.4

0.6

0.8

1.0

0.7

×10 2

100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

×10 3

Figure 5: Scaling laws of solution transition time on S5 . We show that the solution transition time τM →G from the memorization manifold M to the generalization manifold G scales with the learning rate η, batch size b, and ℓ2 regularization coefficient λ. For each hyperparameter configuration, we train for ten runs. We also overlay the theoretical fits. An additional experimental results for Z127 are provided in Figure 13. 5.1

E XPERIMENTAL VALIDATION

Experimental Setting. We use a two-layer MLP: each input token is embedded into 256 dimensions, the two embeddings are concatenated, passed through a width-512 ReLU hidden layer, and projected to |Y| output logits, where |Y| = 120 for S5 and 127 for Z127 . The embedding layers are initialized from a standard normal distribution, while the linear layers use Kaiming initialization. Training is performed in 32-bit floating-point precision. We find that 16-bit precision can introduce numerical instability in grokking, suggesting that grokking dynamics are sensitive to numerical precision errors. We use base hyperparameters η = 10−3 , (β1 , β2 ) = (0.9, 0.999), ℓ2 regularization coefficient λ = 10−4 , and batch size 100. These base values are varied as needed to study scaling laws. Learning Tasks. We adopt two structured learning tasks: a group-theoretic learning task on Sn and a modular arithmetic learning task on Zp . For Sn , each input is an ordered pair of permutations (f, g) ∈ Sn2 , and the model is trained to predict their group product f ◦ g under permutation composition. For Zp , each input is a pair (a, b) ∈ Z2p , and the model is trained to predict the modular sum a + b mod p. In both tasks, the dataset consists of all possible input–output pairs, with a randomly sampled subset used for training and the remaining pairs used for evaluation. More details are provided in Appendix A.1. Results. Figures 2 and 4 validate the predicted scaling laws of the manifold radii ρ2M and ρ2G with respect to the learning rate η, batch size b, and regularization coefficient λ, as stated in Theorems 1 and 2. Figure 5 validates the predicted scaling law of the solution transition time τM →G with respect to the same hyperparameters, as stated in Theorem 3. Additional results on Z127 are provided in Appendices A.15, A.16, and A.17. Overall, the experiments support the theoretical predictions; in particular, ρ2M exhibits the predicted U-shaped dependence on η/b. Discussions and Limitations. The scaling laws originate from how η, b, and λ modulate Adam’s p SDE: η and b module the stochastic diffusion by scaling the gradient covariance with a factor η/b, whereas λ modulates the deterministic gradient drift flow by regularization. Our results inherit the conditions required by Adam’s SDE limit, including an adapted Wiener-process approximation. These conditions require a sufficiently small learning rate and a sufficiently large batch size; thus, our results do not directly apply to the large-η/b regime. 5.2

L ITERATURE VALIDATION

Theorem 1 characterizes the scaling of the memorization radius ρM with respect to η/b, which, to the best of our knowledge, has not been derived in prior work. Theorem 2 derives the scaling law of the generalization radius ρG , showing that it decomposes into the landscape-dependent term √ (0) 2 ρG = Eθ0 ∈Θ ∥θ⋆ (θ0 )∥22 and an O(η/( bλ)) fluctuation correction; this is consistent with the weight-norm “Goldilocks zone” reported in prior studies (Liu et al., 2022; Varma et al., 2023). 9

Preprint

Theorem 3 recovers the inverse dependence of the grokking delay on the regularization coefficient, including the scaling τ ∝ 1/(ηλ) observed or implied in prior empirical and mechanistic studies (Liu et al., 2022). Together, these results show that the proposed stopping-time analysis recovers known scaling behavior while refining it into separate contributions from learning rate, batch size, and regularization coefficient.

C ONCLUSIONS

Optimization dynamics induce a topological configuration of solution manifolds underlying grokking. In this work, we characterize this configuration through a shell–core prior of the reachable solution space, in which training first reaches memorization solutions and then transitions, after a prolonged delay, toward generalization solutions. We derive scaling laws for the manifold radii and the solution transition time with respect to the learning rate, batch size, and regularization coefficient through stopping-time and equilibrium-distribution analyses of Adam’s joint-state SDE. These results connect the geometry of solution manifolds with the stochastic dynamics of Adam and provide a theoretical framework for understanding grokking through optimization dynamics.

ACKNOWLEDGMENTS This research was financially supported in part by Taighde Éireann – Research Ireland under Grant No. 18/CRT/6223, and by the J.E. Cairnes School of Business & Economics, University of Galway, Ireland. It was also supported in part by computational resources and services provided by Calcul Québec (calculquebec.ca) and the Digital Research Alliance of Canada (alliancecan.ca). Partial computational support was also provided by Taighde Éireann – Research Ireland under Grant No. SFI/12/RC/2289 P2 and the Insight Research Ireland Centre for Data Analytics. The authors gratefully acknowledge Prof. Karyn Morrissey of the J.E. Cairnes School of Business & Economics, University of Galway, Ireland, for the support. The authors also acknowledge helpful discussions with Yann Pequignot of Université Laval and Mila – Québec AI Institute. For the purpose of open access, the author has applied a CC BY public copyright licence to any Author Accepted Manuscript version arising from this submission.

R EFERENCES Carl M. Bender and Steven A. Orszag. Advanced Mathematical Methods for Scientists and Engineers I: Asymptotic Methods and Perturbation Theory. Springer New York, New York, NY, 1 edition, 1999. ISBN 978-0-387-98931-0. doi: 10.1007/978-1-4757-3069-2. Originally published by McGraw Hill, 1978. Enea Monzio Compagnoni, Tianlin Liu, Rustem Islamov, Frank Norbert Proske, Antonio Orvieto, and Aurelien Lucchi. Adaptive methods through the lens of SDEs: Theoretical insights on the role of noise. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=ww3CLRhF1v. André Belotto Da Silva and Maxime Gazeau. A general system of differential equations to model first-order adaptive algorithms. J. Mach. Learn. Res., 21(1), January 2020. ISSN 1532-4435. E. B. Dynkin. Markov Processes: Volume II. Grundlehren der mathematischen Wissenschaften. Springer Berlin, Heidelberg, Berlin, Heidelberg, 1st edition, 1965. ISBN 978-3-662-23320-7. doi: 10.1007/978-3-662-25360-1. Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington (eds.), Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pp. 249–256, Chia Laguna Resort, Sardinia, Italy, 13–15 May 2010. PMLR. URL https://proceedings.mlr.press/v9/glorot10a.html. Hyeonbin Hwang and Yeachan Park. Intrinsic task symmetry drives generalization in algorithmic tasks. In Proceedings of the Forty-Third International Conference on Machine Learning, 2026. URL https://arxiv.org/abs/2603.01968. 10

Preprint

Ioannis Karatzas and Steven E. Shreve. Brownian Motion and Stochastic Calculus. Graduate Texts in Mathematics. Springer New York, New York, NY, 2nd edition, 1991. ISBN 978-0-387-976556. doi: 10.1007/978-1-4612-0949-2. Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: //arxiv.org/abs/1412.6980. Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=vt5mnLVIVo. Yann LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient backprop. In Neural Networks: Tricks of the Trade, This Book is an Outgrowth of a 1996 NIPS Workshop, pp. 9–50, Berlin, Heidelberg, 1998. Springer-Verlag. ISBN 3540653112. Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 2101–2110. PMLR, 06–11 Aug 2017. URL https://proceedings.mlr.press/v70/li17f.html. Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: an effective theory of representation learning. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS ’22, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 9781713871088. Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=Bkg6RiCqY7. Róisı́n Luo, James McDermott, Christian Gagné, Qiang Sun, and Colm O’Riordan. Optimizationinduced dynamics of Lipschitz continuity in neural networks, 2025. URL https://arxiv.org/abs/ 2506.18588. Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS ’22, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 9781713871088. Stephan Mandt, Matthew D. Hoffman, and David M. Blei. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18(134):1–35, 2017. URL http://jmlr.org/papers/v18/17-214.html. Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=9XFSbDPmdW. Pascal Jr Tikeng Notsawo, Guillaume Dumas, and Guillaume Rabusseau. Grokking finitedimensional algebra. In Proceedings of the Forty-Third International Conference on Machine Learning, 2026. URL https://arxiv.org/abs/2602.19533. Bernt Øksendal. Stochastic Differential Equations: An Introduction with Applications. Universitext. Springer Berlin, Heidelberg, Berlin, Heidelberg, 6 edition, 2003. ISBN 978-3-540-04758-2. doi: 10.1007/978-3-642-14394-6. Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets, 2022. URL https://arxiv.org/abs/2201. 02177. Dashiell Stander, Qinan Yu, Honglu Fan, and Stella Biderman. Grokking group multiplication with cosets. In Proceedings of the 41st International Conference on Machine Learning, ICML’24. JMLR.org, 2024. 11

Preprint

Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency, 2023. URL https://arxiv.org/abs/2309.02390. Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=S5wmbQc1We.

Preprint

A PPENDIX

This section provides proofs, experimental verification of the technical derivations, and additional experimental results. A high-level sketch of theoretical analysis framework is illustrated as in Figure 6. dSt

Adam’s Joint-State SDE (Lemma 5) q η/b σ(St )dWt , St =

µ(St )dt +

(θt , mt , vt )

Induced Radius SDE (Lemma 6) ⊤

drt = −2B(t)θt (mt ⊘

√

Memorization-Regime Identities (Lemma 10) Late-stage identities

vt ) dt

√ −1/2 tr(πt ) ≈ b tr (diag Σ(θt )) , ⊤ tr G(θt )G(θt ) ≈p

Mean-Field Limit of (mt , vt ) (Lemma 7) −α1 t

mt → ḡt (1 − e

), vt → g ⊙ g t

Scaling Law of Memorization Radius (Theorem 4) 2

Joint-state first-exit analysis for ρM

Preconditioned-Decomposition of Adam’s θ-SDE (Lemma 8) dθt ≈ −a(t)π(θt )ḡ(θt )dt + Rθ,SM (t)dt √ + η G(θt )dWt

Preconditioned-Decomposition of Adam’s Radius SDE (Lemma 9) Deriving the preconditioned-decomposition of radius SDE for analyzing manifold radius precisely. h i √ ⊤ 2 ⊤ ⊤ drt ≈ −2a(t)θt π(θt )ḡt + RSM (t) + η tr G(θt )G(θt ) dt + 2 η θt G(θt )dWt

Reduced Memorization-Regime Radius SDE (Lemma 11) Use memorization-regime identities to reduce the preconditioned radius dynamics 2λ 2 2 √ (r) 2 ⊤ drt ≈ − r − θ ḡt − λθt + η p + RSM (t) + Rπ (t) dt + 2 η rt dWt s(θt ) t s(θt ) t

Scaling Law of Generalization Radius (Theorem 5)

Scaling Law of Solution Transition Time (Theorem 6) 2

Stationary late-stage radial analysis for ρG

First-passage analysis from ∂M to ∂G

Figure 6: High-Level Sketch of Theoretical Analysis Framework. Table of Appendix Contents: A.1 Experimental Settings: Learning Tasks . . . . . . . . . . . . . . . . . . . . . . . .

A.2 Additional Results: Adam–Induced Shell–Core Radius & Stopping-Time Concentration . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

A.3 Proof: Concentration of Normal Initialization . . . . . . . . . . . . . . . . . . . .

A.4 Proof: Concentration of Uniform Initialization . . . . . . . . . . . . . . . . . . . .

A.5 Proof: Adam’s Closed-Form Continuous-Time SDE Limit . . . . . . . . . . . . .

A.6 Proof: Induced Radius SDE . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

A.7 Proof: Mean-Field Limit of First- and Second-Moment Estimates . . . . . . . . . .

A.8 Proof: Preconditioned-Decomposition of Adam’s θ-SDE . . . . . . . . . . . . . .

A.9 Proof: Preconditioned-Decomposition of Adam’s Radius SDE . . . . . . . . . . .

A.10 Proof: Memorization-Regime Preconditioner and Effective Diffusion Identities . .

A.11 Proof: Reduced Memorization-Regime Radius SDE . . . . . . . . . . . . . . . . .

A.12 Proof: Scaling Law of Memorization Radius . . . . . . . . . . . . . . . . . . . . .

A.13 Proof: Scaling Law of Generalization Radius . . . . . . . . . . . . . . . . . . . .

A.14 Proof: Scaling Law of Solution Transition Time . . . . . . . . . . . . . . . . . . .

A.15 Additional Results: Scaling Laws of Manifold Radius ρ2M on Z127 . . . . . . . . .

A.16 Additional Results: Scaling Laws of Manifold Radius ρ2G on Z127

. . . . . . . . .

Preprint

A.17 Additional Results: Scaling Laws of Solution Transition Time on Z127 . . . . . . . A.1

E XPERIMENTAL S ETTINGS : L EARNING TASKS

Group-theoretic learning task on Sn . Let Sn denote the symmetric group on n elements. Each group element f ∈ Sn is a permutation sending its index set [n] := {1, 2, . . . , n} to a fixed enumeration of all n! permutations. Let f [i] denote the indexed permutation i 7→ j from an index i ∈ [n] to a target j ∈ [n]. For example, given two permutations f, g ∈ S5 : ! ! 1 2 3 4 5 1 2 3 4 5 g= ↓ ↓ ↓ ↓ ↓ , f= ↓ ↓ ↓ ↓ ↓ 3 1 5 2 4 2 5 1 4 3 the permutations are given by f [1] : 1 7→ 2, f [2] : 2 7→ 5, etc. The product of f and g is defined by permutation composition (f ◦ g)[i] = g f [i] given by ! 1 2 3 4 5 f ◦g = ↓ ↓ ↓ ↓ ↓ . 1 4 3 2 5 The learning task is to predict the group product f ◦ g given the pair (f, g). The dataset consists of all ordered pairs (f, g) ∈ Sn2 , with a randomly sampled subset used for training and the remaining pairs used for evaluation. Modular arithmetic learning task on Zp . Let Zp denote the cyclic group of integers modulo p. Each input consists of a pair (a, b) ∈ Z2p , and the task is to predict their modular sum a+b mod p. As in the group-theoretic task, training is performed on a randomly selected subset of all p2 input–output pairs, while evaluation is conducted on the full set.

Preprint

A.2

A DDITIONAL R ESULTS : A DAM –I NDUCED S HELL –C ORE R ADIUS & S TOPPING -T IME C ONCENTRATION

Table 1: Adam–Induced Shell–Core Radius & Stopping-Time Concentration on Learning Task S5 . The experiment shows the concentration of Adam–induced shell–core radii and stopping times on the learning task S5 , across 10 seeds, with learning rate η = 10−3 , (β1 , β2 ) = (0.9, 0.999), ℓ2 regularization coefficient 10−4 , and batch size 100. Seed (S0 ) ρ0 (S0 ) ρM (S0 ) τM (S0 ) ρG (S0 )

τG (S0 ) τM →G (S0 )

1 2 3 4 5 6 7 8 9 10

249.27 247.09 247.87 248.13 248.85 248.44 247.98 248.03 247.95 248.88

186.07 184.59 185.19 185.16 185.64 186.61 185.05 185.19 186.23 185.59

1006 1009 1007 1009 1008 939 1009 1009 937 1004

71.39 69.22 68.84 69.67 70.72 69.34 69.50 68.33 67.80 70.15

14000 14850 12150 13950 12500 12100 15450 13850 13050 12750

12994 13841 11143 12941 11492 11161 14441 12841 12113 11746

E[·] Std[·]

248.25 0.60

185.53 0.59

993.7 27.9

69.50 13465.0 1.02 1083.3

12471.3 1071.5

Table 2: Adam–Induced Shell–Core Radius & Stopping-Time Concentration on Learning Task Z127 . The experiment shows the concentration of Adam–induced shell–core radii and stopping times on the learning task S5 , across 10 seeds, with learning rate η = 10−3 , (β1 , β2 ) = (0.9, 0.999), ℓ2 regularization coefficient 10−4 , and batch size 100. Seed (S0 ) ρ0 (S0 ) ρM (S0 ) τM (S0 ) ρG (S0 ) τG (S0 ) τM →G (S0 )

A.3

1 2 3 4 5 6 7 8 9 10

256.25 254.24 254.76 255.47 256.10 255.76 255.09 255.45 255.46 256.06

186.57 183.95 184.10 184.79 185.09 184.74 184.36 184.88 185.97 186.53

1041 1115 1121 1118 1119 1121 1116 1114 1047 1042

71.47 72.40 71.35 71.32 70.61 71.06 72.32 72.21 72.66 73.44

8000 8600 8300 8300 8400 8450 8400 8500 7650 7800

6959 7485 7179 7182 7281 7329 7284 7386 6603 6758

E[·] Std[·]

255.46 0.60

185.10 0.90

1095.4 34.2

71.88 0.82

8240.0 299.8

7144.6 269.4

P ROOF : C ONCENTRATION OF N ORMAL I NITIALIZATION

Lemma 3 (Gaussian Initialization Concentration). Let (1) (k) (j) θ0 = θ0 , . . . , θ 0 ∈ Rp , θ0 ∈ Rpj ,

k X

pj = p,

(18)

j=1 (j)

where the subvectors are independent and [θ0 ]i ∼ N (0, σj2 ). Define the effective coordinate variance σ̃ 2 and the effective dimension parameter p̃ by k k 1X 1 1X 1 σ̃ 2 := pj σj2 , := . (19) p j=1 p̃ k j=1 pj Then the squared initialization radius satisfies E ∥θ0 ∥22 = σ̃ 2 p,

k X Var ∥θ0 ∥22 = 2 pj σj4 . j=1

(20)

Preprint

In particular, for initialization schemes used for numerical stability, each block-wise coordinate −1/2 scale is often chosen so that σj = O(pj ). Under this scaling, k k σ̃ 2 = O , E ∥θ0 ∥22 = O(k), Var ∥θ0 ∥22 = O . (21) p p̃ Proof. We consider a Gaussian initialization with block-wise coordinate scales. Write (1) (k) (j) θ0 = θ0 , . . . , θ 0 , θ0 = σj Z (j) , Z (j) ∼ N (0, Ipj ),

(22)

where the subvectors Z (j) are independent. Then ∥θ0 ∥22 =

k X

(j)

∥θ0 ∥22 =

k X

j=1

σj2 ∥Z (j) ∥22 ,

(23)

j=1

and it suffices to study the distribution of each ∥Z (j) ∥22 . Let Rj := ∥Z (j) ∥22 .

(24)

The random variable Rj follows a chi-square distribution with pj degrees of freedom, whose density is pj 1 fRj (r) = pj p r 2 −1 e−r/2 , r ≥ 0. (25) j 22Γ 2 Computing E ∥θ0 ∥22 .

The expected squared Euclidean norm of the initialization is therefore k X E ∥θ0 ∥22 = σj2 E[Rj ]

(26)

j=1

k X

σj2

Z ∞ rfRj (r) dr

j=1

k X

σj2

j=1 2

pj 2

Z ∞ pj 2

(27)

r 2 e−r/2 dr.

(28)

Applying the change of variables u = r/2 yields Z ∞ p p pj j j r 2 e−r/2 dr = 2 2 +1 Γ +1 . 2 0

(29)

Substituting back, we obtain k p X Γ 2j + 1 E ∥θ0 ∥22 = 2σj2 p Γ 2j j=1

k X

pj σj2

(30)

(31)

j=1

= σ̃ 2 p.

(32)

Computing Var ∥θ0 ∥22 . We next compute the fluctuations of the squared initialization radius. Pk Since the subvectors are independent and ∥θ0 ∥22 = j=1 σj2 Rj , it follows that k X Var ∥θ0 ∥22 = σj4 Var[Rj ]. j=1

(33)

Preprint

To compute Var[Rj ], we first compute the second moment: Z ∞ 2 r2 fRj (r) dr E[Rj ] = 0 Z ∞ p j 1 = pj p r 2 +1 e−r/2 dr. j 0 22Γ 2 Applying again the change of variables u = r/2 gives Z ∞ p p pj j j r 2 +1 e−r/2 dr = 2 2 +2 Γ +2 . 2 0 Thus, p Γ 2j + 2 2 E[Rj ] = 4 p Γ 2j p p j j +1 =4 2 2 = pj (pj + 2). Therefore, 2 Var[Rj ] = E[Rj2 ] − E[Rj ]

(34) (35)

(36)

(37) (38) (39) (40)

= pj (pj + 2) − p2j

(41)

= 2pj .

(42)

Consequently, k X Var ∥θ0 ∥22 = σj4 Var[Rj ]

(43)

j=1

k X

pj σj4 .

(44)

j=1

Coordinate Scaling. Finally, in common initialization schemes used for numerical stability, each block-wise per-coordinate variance is scaled inversely with its block dimension. Equivalently, one may write 1 1 2 σj = O √ , σj = O . (45) pj pj Substituting this scaling into the definition of the effective coordinate variance gives k 1X k σ̃ 2 = pj σj2 = O . (46) p j=1 p Therefore, k X E ∥θ0 ∥22 = pj σj2 = σ̃ 2 p = O(k),

(47)

j=1

  k k X X 1 Var ∥θ0 ∥22 = 2 pj σj4 = O . p j=1 j=1 j

(48)

k 1X 1 1 = , p̃ k j=1 pj

(49)

k 2 Var ∥θ0 ∥2 = O . p̃

(50)

Using the definition

we obtain

Preprint

A.4

P ROOF : C ONCENTRATION OF U NIFORM I NITIALIZATION

Lemma 4 (Uniform Initialization Concentration). Let (1) (k) θ0 = θ0 , . . . , θ 0 ∈ Rp ,

k X

(j)

θ0 ∈ Rpj ,

pj = p,

(51)

j=1 (j)

where the subvectors are independent and the coordinates of each subvector satisfy [θ0 ]i ∼ U(−εj , εj ). Define the effective coordinate variance ε̃2 and the effective dimension parameter p̃ by k

ε̃2 :=

k 1 1X 1 := . p̃ k j=1 pj

1X pj ε2j , p j=1

(52)

Then the squared initialization radius satisfies k

ε̃2 p , 3

Var(∥θ0 ∥22 ) =

4 X pj ε4j . 45 j=1

(53)

 k X d ε̃ p 4 , pj ε4j  , ∥θ0 ∥22 ≈ N  3 45 j=1

p → ∞.

(54)

E[∥θ0 ∥22 ] =

Moreover, by the central limit theorem, 

−1/2

In particular, for initialization schemes with block-wise coordinate half-widths εj = O(pj squared initialization radius satisfies k k 2 2 2 , E[∥θ0 ∥2 ] = O(k), Var(∥θ0 ∥2 ) = O . ε̃ = O p p̃ −1/2

For example, a common numerically stable block-wise choice is εj = pj k 4k 2 d , . ∥θ0 ∥2 ≈ N 3 45p̃

), the

(55)

, in which case (56)

Thus uniform initialization concentrates on a thin hyperspherical shell with squared radius k/3 and squared-radius variance 4k/(45p̃). (j)

(j)

Proof. For each coordinate Xi := [θ0 ]i ∼ U(−εj , εj ), the density of Xi is fX (j) (x) = 2ε1j for i x ∈ [−εj , εj ]. Then the second moment is computed as Z εj Z εj ε2j 1 (j) 2 2 E[(Xi ) ] = x fX (j) (x) dx = x2 dx = . (57) i 2εj −εj 3 −εj The fourth moment is computed as Z εj Z εj ε4j 1 (j) 4 4 E[(Xi ) ] = x fX (j) (x) dx = x4 dx = . i 2εj −εj 5 −εj (j)

Defining Si

(j)

(58) (j)

:= (Xi )2 , we have E[Si ] = E[(Xi )2 ] and E[(Si )2 ] = E[(Xi )4 ], so that (j)

(j)

Var((Xi )2 ) = Var(Si ) = E[(Si )2 ] − E[Si ]2 (j) = E[(Xi )4 ] − !2 ε4j ε2j

= =

−

4ε4j . 45

2 (j) E[(Xi )2 ]

(59) (60) (61) (62)

Preprint

Now set R2 =

pj k X X

(j)

(Xi )2 =

(63)

ε̃2 p 1X pj ε2j = , 3 j=1 3

(64)

j=1 i=1 (j)

Since the random variables Si

(j)

Si .

j=1 i=1

are independent, we obtain

E[R2 ] =

k X

(j)

pj E[Si ] =

j=1 2

Var(R ) =

k X

4 X pj ε4j . 45 j=1

(j) pj Var(Si ) =

j=1

Moreover, by the central limit theorem,     k k k 2 X X X d ε̃ p 4 (j) (j) R2 ≈ N  , pj ε4j  , pj E[Si ], pj Var(Si ) = N  3 45 j=1 j=1 j=1

(65)

p → ∞.

(66)

For numerically stable uniform initialization, the coordinate half-width of each block is often chosen to scale as 1 εj = O √ . (67) pj Substituting this scaling into the preceding identities gives k k 1X pj ε2j = O , ε̃2 = p j=1 p

(68)

and therefore ε̃2 p = O(k), E[∥θ0 ∥22 ] = 3

  k k X X 4 1 . Var(∥θ0 ∥22 ) = pj ε4j = O 45 j=1 p j j=1

(69)

k 1X 1 1 = , p̃ k j=1 pj

(70)

Using the definition

we obtain Var(∥θ0 ∥22 ) = O

k . p̃

(71)

For example, taking the common block-wise choice 1 εj = √ , pj

(72)

we obtain 2 d

R ≈N

k 4k , 3 45p̃

(73)

Preprint

A.5

P ROOF : A DAM ’ S C LOSED -F ORM C ONTINUOUS -T IME SDE L IMIT Modeling Mini-Batch Gradient as Wiener Process by Central-Limit Theorem gk

1 1/2 ∇Lf (θk ) − √ Σ(θk ) ζk , b

ζk

∼

N (0, Ip )

∈

Adam’s Discrete Update Rule mk+1 = β1 mk + (1 − β1 )gk , mk+1 m b k+1 = , 1 − β1k+1 p θk+1 = θk − η m b k+1 ⊘ v bk+1 .

vk+1 = β2 vk + (1 − β2 )(gk ⊙ gk ), vk+1 v bk+1 = , 1 − β2k+1

Continuous-Time Interpolation (Discrete–Continuous Scaling) t

ηk,

βi

1 − ηαi ,

Continuous-Time Limit of First-Moment Estimate

Continuous-Time Limit of Second-Moment Estimate

dmt = −α1 (mt − ḡt )dt q 1/2 + η/b α1 Σt dWt .

dvt = −α2 (vt − g ⊙ g t )dt q + η/b α2 Dt dWt .

→

Expansion of ElementWise Squared Gradient g ⊙ g t := ḡt ⊙ ḡt +

1 diag Σt , b 1/2

Dt := 2 diag(ḡt )Σt

Continuous-Time Limit of Parameter Dynamics √ dθt = −B(t) (mt ⊘ vt ) dt, p where B(t) =

1 − e−α2 t . 1 − e−α1 t

Joint-State SDE q dSt = µ(St )dt + η/b σ(St )dWt , St = (θt , mt , vt ),    √  0 −B(t)mt ⊘ vt 1/2     −α1 (mt − ḡt ) µ(St ) = . , σ(St ) = α1 Σt −α2 (vt − g ⊙ g t ) α2 Dt

Sanity Check via Lemma 6:Induced Radius SDE (Derived via Itô’s Lemma) √ 2 ⊤ drt = −2B(t)θt (mt ⊘ vt ) dt.

Sanity Check via Numerical Experiment in Figure 8 numerically compare the radius dynamics predicted by the induced radius SDE.

Additional Technical Correctness Verification Technical correctness of Adam’s joint-state SDE is also checked by the later decomposed preconditioned radius SDE in Lemma 9 and an approximated form in Lemma 11. Further numerical verification experiments are in Figures 9 and 10.

Figure 7: Proof Sketch for Adam’s Closed-Form Continuous-Time SDE Limit. This diagram illustrates the proof sketch and the corresponding correctness checks for Adam’s continuous-time SDE limit. The mini-batch gradient is modeled as a Wiener process by Central-Limit Theorem. Combining this stochastic gradient model with Adam’s discrete update rules and the continuoustime interpolation yields the continuous-time limits of the first- and second-moment estimates, as well as the continuous-time parameter dynamics. Together with the expansion of the element-wise squared gradient, these components give the joint-state SDE for St := (mt , vt , θt ). Lemma 5 (Adam’s Joint-State Continuous-Time SDE Limit). For a network f parameterized with θ ∈ Rp , let ℓf (s; θ) be the loss for sample s, let Lf (ξ; θ) be the mini-batch loss for batch ξ, and 20

Preprint

let Lf (θ) be the population loss over true dataset distribution. Let θk , θt denote the parameter at iteration k and time t, respectively. Consider Adam with small learning rate η, large batch size b, and

Discrete-Time Update Rules. coefficients (β1 , β2 ). Let

⊤ ⊤ 3p Sk := (θk⊤ , m⊤ k , vk ) ∈ R ,

m0 = 0p ,

v0 = 0p ,

(74)

denote the discrete-time representation of the parameter, first-moment, and second-moment estimates at iteration k. For a batch ξk at step k, the discrete Adam update rules with the bias-correction step are: gk = ∇Lf (ξk ; θk ) mk+1 = β1 mk + (1 − β1 )gk vk+1 = β2 vk + (1 − β2 )(gk ⊙ gk ) mk+1 m̂k+1 = 1 − β1k+1 vk+1 v̂k+1 = 1 − β2k+1 θk+1 = θk − η h(m̂k+1 , v̂k+1 )

(mini-batch gradient) (first-moment estimate) (second-moment estimate)

(75) (76) (77)

(bias-corrected first-moment estimate)

(78)

(bias-corrected second-moment estimate)

(79)

(bias-corrected parameter update),

(80)

where h(m̂k+1 , v̂k+1 ) is the per-coordinate Adam update direction, defined as   m̂(1) √  (1)   v̂ + ϵ  √   . .. h(m̂, v̂) = m̂ ⊘ ( v̂ + ϵ1p ) =  (81) ,   (p)  m̂  √ (p) v̂ + ϵ ⊘ denotes element-wise division, ⊙ denotes element-wise product, and ε is a small positive constant for numerical stability. Distributional Limit of Mini-Batch Gradient. For the parameter θk , the mini-batch gradient gk := ∇Lf (ξk ; θk ) is an empirical average of b i.i.d. per-sample gradient gk (s) := ∇ℓf (s; θk ). Hence, by the central limit theorem, as batch size b → ∞ b 1X 1 d gk = gk (si ) −→ N ḡk , Σk , ḡk := Es [gk (s)], Σk := Covs [gk (s)], (82) b i=1 b where ḡk and Σk are the per-sample gradient expectation and covariance at step k, respectively. Continuous-Time SDE Limit. Under the continuous-time interpolation t = ηk, with the conditions of sufficiently small learning rate η → 0 and large batch size b → ∞, the discrete Adam updates admit the Itô SDE limit r η dSt = µ(St )dt + σ(St )dWt , (83) b where Wt ∈ Rp is a Wiener process adapted to the filtration generated by the mini-batch sampling process {ξt }. The drift and diffusion factors are √     0 − B(t) mt ⊘ vt , µ(St ) =  − α1 mt − ḡt  , σ(St ) = α1 Σ1/2 (84) t − α2 vt − g ⊙ g t α2 Dt with

√ B(t) :=

1 − e−α2 t , 1 − e−α1 t

αi :=

1 − βi , η

i = 1, 2,

(85)

and g ⊙ g t := ḡt ⊙ ḡt +

1 diag(Σt ), b 21

1/2

Dt := 2 diag(ḡt )Σt .

(86)

Preprint

Proof. We derive the continuous-time limit of the Adam optimizer starting from its discrete update rules. For a network f parameterized with θ ∈ Rp , let ℓf (s; θ) be the loss for sample s, let Lf (ξ; θ) be the mini-batch loss for batch ξ, and let Lf (θ) be the population loss over true dataset distribution. Let θk , θt denote the parameter at iteration k and time t, respectively. Mini-Batch Gradient as Wiener Process. Let gk (s) := ∇ℓf (s; θk ) be the per-sample gradient for sample s with the mean and covariance ḡk = Es [gk (s)], Σk = Covs gk (s) , (87) respectively. For a batch ξk := {si }bi=1 of size b, the gradient at time k is given by: b 1X gk = gk (si ), b i=1

(88)

by central limit theorem, so that gk distributionally converge to 1 d gk −→ N ḡk , Σk , (89) b as batch size is sufficiently large b → ∞. We therefore can model mini-batch gradient gk as a Wiener process 1 1/2 1/2 1/2 gk = ḡk − √ Σk ζk , (Σk )(Σk )⊤ = Σk , (90) b where ζk ∼ N (0, Ip ) is Wiener process adapted to the filtration generated from mini-batch sampling sequence {ξk }. Because gradient descent uses negative gradient, we take a negative sign on ζk for simplifying the later analysis algebraically. Discrete Adam with Bias Correction. Let η > 0 be sufficiently small learning rate, β1 , β2 ∈ (0, 1) be exponential decay rates, and ϵ > 0 be numerical stability constant. Let θt ∈ Rp , mt , vt be parameter, first-moment, and second-moment estimates, with the initialization m0 = 0 and v0 = 0. Let k = 0, 1, 2, . . . be discrete iteration time points. The discrete Adam update rules with the bias-correction step written explicitly are: λ (coupled ℓ2 -regularizer) (91) ḡk := ∇[Lf (θk )] = ∇[Lf∗ (θk ) + ∥θk ∥22 ] 2 1 1/2 gk = ḡk − √ Σk ζk (mini-batch gradient) (92) b mk+1 = β1 mk + (1 − β1 )gk (first-moment estimate) (93) vk+1 = β2 vk + (1 − β2 )(gk ⊙ gk ) (second-moment estimate) (94) mk+1 (bias-corrected first-moment estimate) (95) m̂k+1 = 1 − β1k+1 vk+1 v̂k+1 = (bias-corrected second-moment estimate) (96) 1 − β2k+1 θk+1 = θk − η h(m̂k+1 , v̂k+1 ) (bias-corrected parameter update), (97) where λ is the regularization coefficient, and Lf (θk ) is the population coupled loss and Lf∗ (θk ) is the population task loss, then the per-coordinate Adam update direction is   m̂(1) √  (1)   v̂ + ϵ  √   . .. h(m̂, v̂) = m̂ ⊘ ( v̂ + ϵ1p ) =  (98) ,   (p)  m̂  √ (p) v̂ + ϵ ⊘ denotes element-wise division, and ε is a small positive constant for numerical stability. Remark 3. Throughout this paper, we use the convention that the ℓ2 regularizer is coupled through the loss function. Thus, the gradients induced by the regularizer enter the computation of the optimizer’s first- and second-moment estimates. When the ℓ2 penalty is instead decoupled and applied as a standalone update term, θk+1 = θk − η h(m̂k+1 , v̂k+1 ) − η λ θk , (99) where the resulting optimizer is referred to as AdamW. 22

Preprint

Continuous–Time Interpolation.

Taking continuous-time interpolation in discrete iteration k by: t = kη,

(100)

dt = lim ∆t = lim η = η.

(101)

therefore: ∆t→0

Wiener Increment in Gradient. s with the mean and covariance

∆t→0

Let gk (s) := ∇ℓf (s; θk ) be the per-sample gradient for sample Σk = Covs gk (s) ,

ḡk = Es [gk (s)],

(102)

respectively. For a batch ξk := {si }bi=1 of size b, the gradient at time k is given by: b

gk =

1X gk (si ), b i=1

(103)

by central limit theorem, so that gk distributionally converge to 1 d gk −→ N ḡk , Σk , b

(104)

as batch size is sufficiently large b → ∞. We therefore can model mini-batch gradient gk as a Wiener process 1 1/2 gk = ḡk − √ Σk ζk , b

1/2

(Σk )(Σk )⊤ = Σk ,

(105)

where ζk ∼ N (0, Ip ) is Wiener process adapted to the filtration generated from mini-batch sampling sequence {ξk }. Because gradient descent uses negative gradient, we take a negative sign on ζk for simplifying the later analysis algebraically. For the adapted noise vector ζk ∼ N (0, Ip ) ∈ Rp at iteration k (time t = kη) corresponds to the Wiener increment over [t, t + η]: 1 √ 1 1 √ ζk = √ ζk η = √ ζk ∆t = √ ∆Wk , η η η where ∆t = η and the increment follows the property of Brownian motion: √ ∆Wk = ζk ∆t ∼ N (0, ∆t Ip ).

(106)

(107)

Remark 4 (Sanity-Check Reference: Continuous SDE for SGD). Therefore, for SGD, the continuous-time SDE limit is, with t = kη and ∆t = η: dθt = lim θt+∆t − θt (108) ∆t→0 " # 1 1/2 = lim −η ḡk − √ Σk ζk (109) ∆t→0 b " # r η 1/2 √ = lim −ηḡk + Σ ηζk (110) ∆t→0 b k " # r η 1/2 = lim −ηḡk + Σ ∆Wk (111) ∆t→0 b k r η 1/2 = −ḡt dt + Σ dWt . (112) b t 23

Preprint

Discrete Increment of First-Moment Estimate. From equation (93): mk+1 = β1 mk + (1 − β1 )gk mk+1 − mk = β1 mk + (1 − β1 )gk − mk = (β1 − 1)mk + (1 − β1 )gk = −(1 − β1 )mk + (1 − β1 )gk = (1 − β1 )(gk − mk ) = −(1 − β1 )(mk − gk )

(113) (114) (115) (116) (117) (118)

Substituting equation (105): 1 1/2 gk = ḡk − √ Σk ζk b

(119)

into equation (118) yields: 1 1/2 mk+1 − mk = −(1 − β1 ) mk − ḡk + √ Σk ζk b 1 1/2 = −(1 − β1 )[mk − ḡk ] − (1 − β1 ) √ Σk ζk b 1 − β1 1/2 = −(1 − β1 )[mk − ḡk ] − √ Σk ζk . b Continuous-Time Limit of First-Moment Estimate. dm = lim ∆m

Use ∆t = η and ∆Wk =

(120) (121) (122) √

(123)

∆t→0

= lim (mk+1 − mk ) ∆t→0 "

(124)

= lim −(1 − β1 )[mk − ḡk ] − ∆t→0

1 − β1 1/2 √ Σ k ζk b

1 1 − β1 1/2 1 √ = lim −(1 − β1 )[mk − ḡk ] ∆t − √ Σk √ ∆tζk ∆t→0 η η b =−

∆ζk ,

1 − β1 1 − β1 1/2 [mt − ḡk ]dt − √ Σ dWt . η bη k

(125) # (126) (127)

Expansion of Element-Wise Squared Gradient. To derive the continuous-time limit of the second-moment estimate dv. We need to expand gk ⊙ gk of equation (96) into a tractable form. Consider: 1 1/2 gk = ḡk − √ Σk ζk , (128) b we expand the element-wise quadratic product gk ⊙ gk by: 1 1/2 1 1/2 gk ⊙ gk = ḡk − √ Σk ζk ⊙ ḡk − √ Σk ζk (129) b b 1 1/2 1 1/2 1/2 = ḡk ⊙ ḡk − 2 ḡk ⊙ √ Σk ζk + Σk ζ k ⊙ Σk ζ k (130) b b h i 1 2 1/2 1/2 ⊙2 = [ḡk ] − √ ḡk ⊙ Σk ζk + (Σk ζk )⊙2 . (131) b b Using the fundamental identity, for two vectors u and v: u ⊙ v = diag(u)v, then the Hadamard cross product can be written as: h i 1/2 1/2 ḡk ⊙ Σk ζk = diag(ḡk )Σk ζk .

(132) (133)

We write gk ⊙ gk as: h i 1 2 1/2 1/2 gk ⊙ gk = ḡk⊙2 − √ diag(ḡk ) Σk ζk + (Σk ζk )⊙2 . b b 24

(134)

Preprint

Statistics of Element-Wise Squared Gradient. gradient gk ⊙ gk admits a mean Eξk [gk ⊙ gk | θk ] = E

Fix θk over all batches, the element-wise squared

1 1/2 diag(ḡk ) Σk ζk +

2 ḡk⊙2 − √

1/2 Σk ζ k

⊙2

b b h i 1 ⊙2 2 1/2 1/2 ⊙2 = ḡk − √ diag(ḡk ) E Σk ζk + E Σk ζk b b 1 = ḡk⊙2 + diag(Σk ) , b

(135)

and a covariance 1 1/2 ⊙2 2 1/2 Σ k ζk Covξk (gk ⊙ gk | θk ) = Cov − √ diag(ḡk ) Σk ζk + b b ⊙2 4 1 1/2 = diag(ḡk ) Σk diag(ḡk ) + 2 Cov Σk ζ k b b 4 1 = diag(ḡk ) Σk diag(ḡk ) + 2 [2 (Σk ⊙ Σk )] b b 4 2 = diag(ḡk ) Σk diag(ḡk ) + 2 (Σk ⊙ Σk ) , b b

(136)

where the last identity follows componentwise from ⊙2 2 2 2 2 1/2 1/2 1/2 1/2 1/2 Cov Σk ζ k = E Σk ζ k Σk ζk − E Σk ζ k E Σ k ζk i

= (Σk )ii (Σk )jj + 2(Σk )2ij − (Σk )ii (Σk )jj = 2(Σk )2ij = 2(Σk ⊙ Σk )ij . Continuous-Time Limit of Second-Moment Estimate.

(137) From equation (96):

vk+1 − vk = β2 vk + (1 − β2 )(gk ⊙ gk ) − vk = −(1 − β2 )vk + (1 − β2 )(gk ⊙ gk ) = −(1 − β2 )[vk − gk ⊙ gk ],

(138) (139) (140)

substituting equation (134) yields: i h 1 1/2 2 1/2 vk+1 − vk = − (1 − β2 ) vk − ḡk⊙2 + √ diag(ḡk )Σk ζk − (Σk ζk )⊙2 b b

(141)

Expanding yields: h i 2(1 − β ) 1 − β2 1/2 ⊙2 √ 2 diag(ḡk )Σ1/2 vk+1 − vk = − (1 − β2 ) vk − ḡk⊙2 − (Σk ζk ) . (142) k ζk + b b Consider the diffusion terms: 2(1 − β2 ) 2(1 − β2 ) 1/2 1/2 ∆Wk √ √ diag(ḡk )Σk ζk = diag(ḡk )Σk √ η b b 2(1 − β2 ) 1/2 √ = diag(ḡk )Σk ∆Wk , bη 25

(143) (144)

Preprint

and: 2 h i2 1/2 1/2 ∆Wk Σ k ζ k = Σk √ η i i 2  p (j) X ∆W 1/2 =  (Σk )ij √k  η j=1 p

(145)

(146)

1 XX 1/2 1/2 (j) (h) (Σk )ij (Σk )ih ∆Wk ∆Wk η j=1

(147)

1 η j=1

(148)

h=1 p p XX

1/2

(Σk )ij (Σk )ih δjh ∆t

h=1

1X 1/2 [(Σk )ij ]2 ∆t = η j=1 =

(149)

∆t ∆t (Σk )ii = diag Σk i , η η

(150)

so (Σ1/2 ζk )⊙2 = (∆t/η) diag(Σk ). Take limits ∆t → dt, ∆Wk → dWt , then the continuous-time limit of ∆v = vk+1 − vk is: i 1 − β2 h dvt = − vt − ḡk⊙2 dt η 2(1 − β2 ) 1 − β2 1/2 − √ diag(ḡt ) Σt dWt + diag Σt dt. ηb bη Continuous-Time Limit of Parameter Dynamics. crete update uses the bias-corrected moments m̂k = mk /(1 − β1k ),

(151)

From equation (97) (with ε dropped), the disv̂k = vk /(1 − β2k ),

(152)

then θk+1 = θk − η m̂k ⊘

p v̂k = θk − η

√ 1 − β2k mk ⊘ vk . 1 − β1k

(153)

Under the joint scaling η → 0 with (1 − βi )/η = αi fixed and t = ηk, 1 − βik = 1 − (1 − αi η)t/η −−−→ 1 − e−αi t , η→0 p √ k 1 − β2 1 − e−α2 t −−−→ =: B(t), k η→0 1 − e−α1 t 1 − β1 with B(t) defined in equation (85). The continuous-time θ-dynamics is therefore √ dθt = − B(t) mt ⊘ vt dt,

(154)

(155)

which is time-inhomogeneous through B(t). Joint-State SDE.

To simplify discussion, by equations (134), (135), (136), h i 1 2 1/2 1/2 gt ⊙ gt = ḡt ⊙ ḡt − √ diag(ḡt ) Σt ζt + (Σt ζt )⊙2 , b b 1 Eξt [gt ⊙ gt | θt ] = ḡt ⊙ ḡt + diag(Σt ) , b 4 2 Covξt (gt ⊙ gt | θt ) = diag(ḡt ) Σt diag(ḡt ) + 2 (Σt ⊙ Σt ) , b b 26

(156) (157) (158)

Preprint

we set g ⊙ g t := E[gt ⊙ gt | θk ] = ḡt ⊙ ḡt + 1b diag (Σt ), (159) 1/2 1 1 Dt := 2 diag ḡt Σt , (160) Dt Dt⊤ = Cov[gk ⊙ gk | θk ] − O( 2 ), b b where Dt is the diffusion factor of the leading cross-term in gk ⊙gk ; the sub-leading term has Wick’s variance at a scale by O(1/b2 ) and is dropped if b → ∞. Set αi := (1 − βi )/η, and combining equations (155), (127), (151), √ dθt = − B(t) mt ⊘ vt dt, 1 − β1 1/2 1 − β1 [mt − ḡt ] dt − √ dmt = − Σ dWt η bη t r η 1/2 α1 Σt dWt , = −α1 (mt − ḡt )dt − b i 1 − β2 h dvt = − vt − ḡt ⊙ ḡt dt η 2(1 − β2 ) 1 − β2 1/2 diag Σt dt − √ diag(ḡt ) Σt dWt + ηb bη r i h η 1 = −α2 vt − ḡt ⊙ ḡt − diag(Σt ) dt − α2 Dt dWt b b r η α2 Dt dWt , = −α2 (vt − g ⊙ g t ) dt − b so that:   √ ! ! r 0 −B(t) mt ⊘ vt θt η  p×p 1/2 −α1 (mt − ḡt ) d mt = dt + α1 Σt  dWt . b vt −α2 (vt − g ⊙ g t ) α2 Dt

(161) (162)

(163)

(164)

(165)

Write the joint-state as St :=

θt mt vt

! ∈ R3p ,

(166)

then Adam’s optimization is characterized by r dSt = µ(St ) dt + with Wt a p-dimensional Wiener process and √ ! −B(t) mt ⊘ vt −α1 (mt − ḡt ) µ(St ) = ∈ R3p , −α2 (vt − g ⊙ g t )

η σ(St ) dWt , b

 0p×p  ∈ R3p×p . σ(St ) = α1 Σ1/2 t α2 Dt

(167)



(168)

The technical correctness verification is provided in later sections and experiments, such as Appendix A.6 through an numerical experiment.

Remark 5 (SGD’s SDE as Special Case). The joint-state SDE in equation (167) contains the SGD’s SDE as a limiting case. Specifically, take the no-momentum limit β1 =√0, so that mt = gt , and ignore the second-moment state vt equivalently, take β2 = 0 and replace vt by 1p . Then Adam’s θ-update reduces to the SGD update θk+1 = θk − ηgk , (169) and its continuous-time interpolation gives the standard SGD SDE p 1/2 dθt = −ḡt dt + η/b Σt dWt . (170)

Preprint

A.6

P ROOF : I NDUCED R ADIUS SDE t 22 (measured) rt2 (SDE)

60000

t 22 (measured) rt2 (SDE)

60000

50000

40000

rt2

40000

30000

20000

20000 10000

10000 0

2000

4000

6000

8000

10000

2000

(a) Problem S5

4000

6000

8000

10000

(b) Problem Z127

Figure 8: Sanity Check with Radius SDE. We use the induced radius SDE (Lemma 6), derived via Itô’s lemma, as a sanity check for Adam’s continuous-time SDE limit. We use the induced radius SDE as a sanity check on the joint-state SDE of Lemma 1. Applying Itô’s 3p×3p lemma (Øksendal, 2003) to the quadratic form rt2 = S⊤ t E St with the θ-block projector E ∈ R defined in equation (176) below must reproduce the Itô decomposition stated in Lemma 6; empirical verification on Adam runs is reported in Figure 8. ⊤ ⊤ ∈ R3p evolve according to Adam’s Lemma 6 (Induced Radius SDE). Let St = (θt⊤ , m⊤ t , vt ) joint-state SDE r η dSt = µ(St )dt + σ(St )dWt , (171) b where √ ! − B(t) mt ⊘ vt − α1 (mt − ḡt ) , − α2 (vt − g ⊙ g t )

µ(St ) =



 0 , σ(St ) = α1 Σ1/2 t α2 Dt

(172)

Wt ∈ Rp is a Wiener process, and B(t) is the Adam bias-correction factor. Define rt2 := ∥θt ∥22 ,

(173)

then rt2 admits the dynamics drt2 = −2B(t)θt⊤ (mt ⊘

√

vt ) dt.

(174)

This is not an ordinary differential equation (ODE) as the mt and vt are stochastic through its dependence on the joint state St . ⊤ ⊤ 3p Proof. Recall from Lemma 5 the joint-state SDE on St = (θt⊤ , m⊤ t , vt ) ∈ R , p dSt = µ(St ) dt + η/b σ(St ) dWt ,

(175)

with µ(St ) and σ(St ) as in equations (168). Define a projector E ∈ R3p×3p E :=

Ip 0p 0p

0p 0p 0p

! ,

(176)

where each block is p × p. Then E is symmetric and idempotent, E⊤ = E and E 2 = E, and for S = (θ⊤ , m⊤ , v ⊤ )⊤ the squared parameter norm is the quadratic form r2 = ∥θ∥22 = S⊤ E S.

(177)

Apply Itô’s lemma to f (S) := S⊤ E S under equation (175). The gradient and Hessian of f are ∇2S f (S) = 2 E,

∇S f (S) = 2 E S, 28

(178)

Preprint

so p η ⊤ ⊤ drt2 = 2 S⊤ t E µ(St ) dt + 2 η/b St E σ(St ) dWt + b tr σ(St ) E σ(St ) dt.

(179)

Each contraction with E keeps only the θ-block: ⊤ ⊤ S⊤ t E µ(St ) = θt µθ (St ) = − B(t) θt mt ⊘

S⊤ t E σ(St ) ⊤ tr σ(St ) E σ(St )

= =

θ⊤ t σθ (St ) = 0, tr σθ (St )⊤ σθ (St )

√ vt ,

(180) (181)

= 0.

(182)

Equations (181)–(182) use σθ ≡ 0. Substituting equations (180)–(182) into equation (179) yields √ vt dt, (183) drt2 = − 2 B(t) θ⊤ t mt ⊘ which is equation (174). A.7

P ROOF : M EAN -F IELD L IMIT OF F IRST- AND S ECOND -M OMENT E STIMATES

Lemma 7 (Mean-Field Limit of First- and Second-Moment Estimates). Let mt and vt be Adam’s first- and second-moment states, let gt be the mini-batch gradient at time t, let ḡt := E[gt | θt ] be the mean of gt , and let g ⊙ g t := E[gt ⊙ gt | θt ] be the mean of gt ⊙ gt . Assume that ḡt and g ⊙ g t vary slowly in the long-term dynamics as t → ∞. Then, under initialization m0 = 0, v0 = 0, the first- and second-moment estimates admit the mean-field limit mt → ḡt 1 − e−α1 t , vt → g ⊙ g t 1 − e−α2 t , (184) so that mt ⊘

√

q vt → ḡt 1 − e−α1 t ⊘ g ⊙ g t (1 − e−α2 t ) q −α1 t ≈ ḡt 1 − e ⊘ g ⊙ gt .

(185)

Proof. We start from the stochastic moment dynamics of the joint-state Adam SDE, r η 1/2 dmt = −α1 (mt − ḡt )dt − α1 Σt dWt , (186) b r η dvt = −α2 (vt − g ⊙ g t ) dt − α2 Dt dWt . (187) b The mean-field approximation replaces the stochastic moment dynamics by their conditional mean dynamics. Since the stochastic terms are martingale increments, E[dWt | θt , mt , vt ] = 0.

(188)

Taking conditional expectation in (186) and (187) gives E[dmt | θt , mt , vt ] = −α1 (mt − ḡt )dt,

(189)

E[dvt | θt , mt , vt ] = −α2 (vt − g ⊙ g t ) dt.

(190)

Thus, under the mean-field dynamics, dmt = −α1 (mt − ḡt )dt,

(191)

dvt = −α2 (vt − g ⊙ g t ) dt.

(192)

Because ḡt varies slowly on the mt relaxation time scale, we freeze ḡt when solving equation (191). Hence dmt = −α1 mt + α1 ḡt . (193) dt Multiplying by eα1 t gives eα1 t

dmt + α1 eα1 t mt = α1 eα1 t ḡt dt d α1 t e mt = α1 eα1 t ḡt . dt 29

(194)

Preprint

Integrating from 0 to t, eα1 t mt − m0 =

Z t

α1 eα1 u ḡt du 0 = ḡt eα1 t − 1 .

(195)

Using m0 = 0, we obtain mt = ḡt 1 − e−α1 t .

(196)

Similarly, because g ⊙ g t varies slowly on the vt relaxation time scale, we freeze g ⊙ g t when solving (192). Hence dvt = −α2 vt + α2 g ⊙ g t . dt

(197)

Multiplying by eα2 t gives eα2 t

dvt + α2 eα2 t vt = α2 eα2 t g ⊙ g t dt d α2 t e vt = α2 eα2 t g ⊙ g t . dt

(198)

Integrating from 0 to t, eα2 t vt − v0 =

Z t

α2 eα2 u g ⊙ g t du = g ⊙ g t eα 2 t − 1 .

(199)

vt = g ⊙ g t + (v0 − g ⊙ g t ) e−α2 t .

(200)

vt = g ⊙ g t 1 − e−α2 t .

(201)

Therefore, Using v0 = 0, this becomes Combining equation (196) and (201), we get q √ mt ⊘ vt = ḡt 1 − e−α1 t ⊘ g ⊙ g t (1 − e−α2 t ). (202) √ For large t on the vt relaxation time scale, 1 − e−α2 t decays faster than 1 − e−α1 t , so that q √ mt ⊘ vt ≈ ḡt 1 − e−α1 t ⊘ g ⊙ g t . (203)

A.8

P ROOF : P RECONDITIONED -D ECOMPOSITION OF A DAM ’ S θ-SDE

We decompose Adam’s exact θ-dynamics into a preconditioned drift and diffusion with a residual. This decomposition helps us to formulate the late-stage evolution of training dynamics with Adam in grokking. The technical correctness is verified via numerical experiments in Appendix A.9. ⊤ ⊤ Lemma 8 (Preconditioned-Decomposition of Adam’s θ-SDE). Let St = (θt⊤ , m⊤ t , vt ) evolve under Adam’s joint-state SDE. Define √ √ 1 − e−α2 t −α t 2 a(t) := 1 − e , B(t) := , (204) 1 − e−α1 t 1 −1/2 1/2 , G(θt ) = √ B(t)π(θt )Σt . (205) π(θt ) := diag(g ⊙ g t ) b where π(θt ) is referred to as the preconditioner of Adam’s SDE and G(θt ) is referred to as the preconditioned diffusion factor. By dropping higher correction orders, Adam’s parameter dynamics admit the preconditioned decomposition √ dθt ≈ −a(t)π(θt )ḡt dt + Rθ,SM (t)dt + η G(θt )dWt , (206) 30

Preprint

where q √ Rθ,SM (t) := −B(t) E [mt ⊘ vt ] − ḡt 1 − e−α1 t ⊘ g ⊙ g t ,

(207)

is referred to as the slow-manifold residual, which represents the residual term to the preconditioned term. This result immediately recovers standard SGD’s SDE with B(t) = 1, mt = gt , vt = 1p , π(θt ) = Ip . Proof. We decompose the θ-component of Adam’s joint-state SDE √ dθt = −B(t) (mt ⊘ vt ) dt, where

√

1 − e−α2 t , 1 − e−α1 t The SDE preconditioner is defined as

a(t) =

B(t) =

√

1 − e−α2 t .

−1/2

π(θt ) := diag(g ⊙ g t )

(208)

(209)

(210)

Mean-Field Decomposition. By Lemma 7 (Mean-Field Limit of First- and Second-Moment Estimates), consider the evolution q √ mt ⊘ vt → ḡt 1 − e−α1 t ⊘ g ⊙ g t , (211) we add and subtract the mean-field limit q ḡt 1 − e−α1 t ⊘ g ⊙ g t

(212)

from equation (208), then

√ dθt = −B(t) (mt ⊘ vt ) dt q −α1 t = − B(t) ḡt 1 − e ⊘ g ⊙ g t dt | {z } mean-field component

q √ −α1 t − B(t) mt ⊘ vt − ḡt 1 − e ⊘ g ⊙ g t dt. {z } |

(213)

mean-field residual

Since B(t) 1 − e−α1 t = a(t),

(214)

and ḡt ⊘

g ⊙ g t = π(θt )ḡt ,

the first term (mean-field component) in equation (213) becomes q −α1 t −B(t) ḡt 1 − e ⊘ g ⊙ g t dt = −a(t)π(θt )ḡt dt. Martingale Decomposition of Mean-Field Residual.

(215)

(216)

Let

Et [·] := E[· | θt ].

(217)

q √ Rθ,SM (t) := −B(t) Et [mt ⊘ vt ] − ḡt 1 − e−α1 t ⊘ g ⊙ g t

(218)

Define

as the small-manifold residual. Then the second term in equation (213) admits a martingale decomposition q √ −α1 t − B(t) mt ⊘ vt − ḡt 1 − e ⊘ g ⊙ g t dt √ √ = Rθ,SM (t)dt − B(t) [mt ⊘ vt − Et [mt ⊘ vt ]] dt. (219) 31

Preprint

Effective Diffusion.

To analyze the martingale residual √ √ mt ⊘ vt − Et [mt ⊘ vt ] ,

(220)

define m̄t := 1 − e−α1 t ḡt ,

v̄t := g ⊙ g t ,

(221)

and δmt := mt − m̄t , δvt := vt − v̄t . √ A componentwise Taylor expansion of m ⊘ v around (m̄t , v̄t ) gives √ √ 1 3/2 m̄t ⊘ v̄t ⊙ δvt + Qθ (t), mt ⊘ vt = m̄t ⊘ v̄t + π(θt )δmt − 2 where Qθ (t) = O ∥δmt ∥ ∥δvt ∥ + ∥δvt ∥22 .

(222)

(223)

(224)

Taking the conditional mean of equation (223) and subtracting it from equation (223), we obtain √ √ mt ⊘ vt − Et [mt ⊘ vt ] = π(θt ) (δmt − Et [δmt ]) 1 3/2 − m̄t ⊘ v̄t ⊙ (δvt − Et [δvt ]) 2 + Qθ (t) − Et [Qθ (t)]. (225) Variation Residual of First-Moment Estimate. moment estimate

We now analyze the variation residual of first-

δmt − Et [δmt ].

(226)

From Adam’s joint-state SDE, the first-moment component satisfies dmt = −α1 mt dt + α1 gt dt.

(227)

Rearranging equation (227) gives mt dt = gt dt −

1 dmt . α1

(228)

Taking conditional expectations in equation (228) gives Et [mt ]dt = ḡt dt −

1 dEt [mt ]. α1

(229)

Subtracting equation (229) from equation (228) yields (mt − Et [mt ]) dt = (gt − ḡt ) dt −

1 d (mt − Et [mt ]) . α1

(230)

Since δmt − Et [δmt ] = mt − Et [mt ],

(231)

we get (δmt − Et [δmt ]) dt = (gt − ḡt ) dt −

1 d (δmt − Et [δmt ]) . α1

(232)

In continuous-time diffusion scaling, (gt − ḡt ) dt =

p 1/2 η/b Σt dWt .

Substituting equation (233) into equation (232) gives p 1 1/2 (δmt − Et [δmt ]) dt = η/b Σt dWt − d (δmt − Et [δmt ]) . α1 32

(233)

(234)

Preprint

Define dRm (t) := −

1 d (δmt − Et [δmt ]) . α1

(235)

Then (δmt − Et [δmt ]) dt =

1/2

η/b Σt dWt + dRm (t).

(236)

Substituting equation (236) into the first term of equation (225) gives − B(t)π(θt ) (δmt − Et [δmt ]) dt p 1/2 = −B(t) η/b π(θt )Σt dWt − B(t)π(θt )dRm (t).

(237)

Absorbing the sign into Wt , and using the definition 1 1/2 G(θt ) := √ B(t)π(θt )Σt , b

(238)

we obtain −B(t)π(θt ) (δmt − Et [δmt ]) dt = Produce Claims.

√

η G(θt )dWt − B(t)π(θt )dRm (t).

(239)

Combining equations (225) and (239), we get √ √ − B(t) [mt ⊘ vt − Et [mt ⊘ vt ]] dt √ = η G(θt )dWt + dWθ (t),

(240)

dWθ (t) := −B(t)π(θt )dRm (t) i B(t) h 3/2 + m̄t ⊘ v̄t ⊙ (δvt − Et [δvt ]) dt 2 − B(t) [Qθ (t) − Et [Qθ (t)]] dt.

(241)

where

Combining equations (216), (219), and (240), we obtain dθt = −a(t)π(θt )ḡt dt + Rθ,SM (t)dt +

√

η G(θt )dWt + dWθ (t).

Dropping higher-order terms and keeping the leading diffusion order gives √ dθt ≈ −a(t)π(θt )ḡt dt + Rθ,SM (t)dt + η G(θt )dWt .

(242)

(243)

This is the claimed preconditioned decomposition. Remark 6. In particular, in late-stage training with sufficiently large batch size b, the reduction of mt , vt mt → ḡt 1 − e−α1 t , vt → g ⊙ g t , (244) is referred to as the slow-manifold reduction. When the residual Rθ,SM (t) is negligible, then the reduced preconditioned θ-SDE can further be simplified into √ dθt ≈ −a(t)π(θt )ḡt dt + η G(θt )dWt . (245) Remark 7 (Quick Sanity Check). As a sanity check, if π(θt ) = Ip , B(t) → 1, and the residual is ignored at leading order, then 1 1/2 G(θt ) → √ Σt , b

(246)

and the preconditioned θ-SDE reduces to dθt = −ḡt dt +

1/2

η/b Σt dWt ,

which is the standard continuous-time SDE approximation of mini-batch SGD.

(247)

Preprint

A.9

P ROOF : P RECONDITIONED -D ECOMPOSITION OF A DAM ’ S R ADIUS SDE

We decompose Adam’s exact radius SDE into preconditioned and residual terms in Lemma 9. This decomposition is used to derive the reduced radius SDE under the slow-manifold reduction and small-isotropic gradient covariance observations. Its technical correctness is verified in Figure 9, in which the predicted radius dynamics exactly match the theoretical values computed from Lemma 9. ×104

×104 rt2 (with RSM)

G M

rt2

t 22 (measured)

RSM terms (per t)

3 2 1 0

G M

gt 2

a(t) & B(t)

2000 4000 6000 8000 10000 12000 14000 16000

×10 1

0.6 0.4

a(t) = 1 B(t) = 1

0.2

k 2 k 2 /(1

2000 4000 6000 8000 10000 12000 14000 16000

6 5 3

tr(GtGt )

4 tr( t) btr( t 1)

2 1

G M

2000 4000 6000 8000 10000 12000 14000 16000

gt 2

3.5 3.0 2.5 2.0 1.5 1.0 0.5 0.0

k 1)

G M

×108

tr( t)

(1)(t) = 2B RSM (mt vt ) (0)(t) = +2a RSM g (0)(t) + R (1)(t) RSM(t) = RSM SM

0.8

2000 4000 6000 8000 10000 12000 14000 16000

1.0

0.0

4.0 3.5 3.0 2.5 2.0 1.5 1.0 0.5 0.0

t 2 G M

2000 4000 6000 8000 10000 12000 14000 16000

×105

tr(GtGt ) p G M

2000 4000 6000 8000 10000 12000 14000 16000

Figure 9: Preconditioned Radius SDE. This experiment shows the dynamics of the preconditioned radius √ SDE in Lemma 9, and two memorization-regime identities in Lemma 10 where −1/2 and tr(G(θt )G(θt )⊤ ) ≈ p. The predicted radius dynamics tr(π(θt )) ≈ b tr (diag Σt ) match the theoretical values computed from Lemma 9 and Lemma 10, validating the correctness of the decomposition. ⊤ ⊤ Lemma 9 (Preconditioned-Decomposition of Adam’s Radius SDE). Let St = (θt⊤ , m⊤ t , vt ) evolve under Adam’s joint-state SDE, and let rt2 := ∥θt ∥22 . Define √ √ 1 − e−α2 t −α t 2 , B(t) := a(t) := 1 − e 1 − e−α1 t 1 −1/2 1/2 π(θ) := diag(g ⊙ g(θ)) , G(θt ) = √ B(t)π(θt )Σt . (248) b

Then the residual-corrected preconditioned squared-radius dynamics are √ drt2 ≈ −2a(t)θt⊤ π(θt )ḡt + RSM (t) + η tr G(θt )G(θt )⊤ dt + 2 η θt⊤ G(θt )dWt , 34

(249)

Preprint

where the slow-manifold residual is q √ RSM (t) := −2B(t)θt⊤ E [mt ⊘ vt ] − ḡt 1 − e−α1 t ⊘ g ⊙ g t .

(250)

Proof. Let Eθ :=

Ip 0 0

! 0 0 0 0 ∈ R3p×3p 0 0

(251)

⊤ ⊤ be a projector. Since St = (θt⊤ , m⊤ t , vt ) , the squared radius can be written as the joint-state quadratic form

rt2 = ∥θt ∥22 = St⊤ Eθ St .

(252)

2θt 0 0

(253)

Therefore, ∇S rt2 = 2Eθ St =

! ∇2S rt2 = 2Eθ .

⊤ ⊤ Remark 8. Since θt is the θ-component of the joint state St = (θt⊤ , m⊤ t , vt ) , and the components of St are coupled through Adam’s dynamics, we should regard ∥θt ∥22 as a function of the full joint state rather than of an isolated variable θt . Thus, when applying Itô’s lemma to the joint-state SDE, we introduce the projection map Eθ : St 7→ θt and write the squared radius as the bilinear form St⊤ Eθ St .

Applying Itô’s lemma to rt2 = St⊤ Eθ St under Adam’s joint-state SDE gives i 1 h η σ(St )σ(St )⊤ ∇2S rt2 dt. drt2 = (∇S rt2 )⊤ dSt + tr (254) 2 b The joint-state diffusion has zero θ-block, because the Brownian noise enters only through the (m, v) components. Since ∇2S rt2 = 2Eθ only selects the θ-block, the second-order term vanishes: h η i tr σ(St )σ(St )⊤ ∇2S rt2 = 0. (255) b Thus,

Radius Dynamics.

drt2 = (∇S rt2 )⊤ dSt = 2θt⊤ dθt . From the θ-component of Adam’s joint-state SDE, dθt = −B(t) (mt ⊘

(256) √

vt ) dt,

(257)

we obtain the exact induced squared-radius identity drt2 = −2B(t)θt⊤ (mt ⊘

√

vt ) dt.

(258)

Mean-Field Decomposition. By Lemma 7 (Mean-Field Limit of First- and Second-Moment Estimates), consider the mean-field limit q √ mt ⊘ vt → ḡt 1 − e−α1 t ⊘ g ⊙ g t , (259) we add and subtract q ḡt 1 − e−α1 t ⊘ g ⊙ g t

(260)

inside the Adam drift: √ −2B(t)θt⊤ (mt ⊘ vt ) = −2B(t)θt⊤

−α1 t

q ⊘ g ⊙ gt

ḡt 1 − e q √ ⊤ −α1 t − 2B(t)θt mt ⊘ vt − ḡt 1 − e ⊘ g ⊙ gt .

(261)

Preprint

Using the continuous-time bias-correction factor √ B(t) = and the definition a(t) =

√

1 − e−α2 t , 1 − e−α1 t

(262)

1 − e−α2 t , we have B(t) 1 − e−α1 t = a(t).

(263)

Moreover, since −1/2

π(θt ) = diag(g ⊙ g t )

(264)

we have the element-wise identity ḡt ⊘

g ⊙ g t = π(θt )ḡt .

(265)

Therefore, the first term in equation (261) becomes q −2B(t)θt⊤ ḡt 1 − e−α1 t ⊘ g ⊙ g t = −2a(t)θt⊤ π(θt )ḡt .

(266)

Martingale Residual of Mean-Field Decomposition. The second term in equation (261) is exactly the slow-manifold residual: q √ RSM (t) := −2B(t)θt⊤ E [mt ⊘ vt ] − ḡt 1 − e−α1 t ⊘ g ⊙ g t . (267) Combining equations (261)–(267), the exact induced squared-radius has a drift −2a(t)θt⊤ π(θt )ḡt + RSM (t) dt.

(268)

Effective Diffusion. Finally, in the reduced preconditioned diffusion approximation, the mini√ batch gradient fluctuations induce the θ-diffusion term η G(θt )dWt . Applying Itô’s lemma to this diffusion contribution gives the additional correction √ η tr G(θt )G(θt )⊤ dt + 2 η θt⊤ G(θt )dWt . (269) Produce Claims. Adding equation (269) to equation (268) yields √ drt2 ≈ −2a(t)θt⊤ π(θt )ḡt + RSM (t) + η tr G(θt )G(θt )⊤ dt + 2 η θt⊤ G(θt )dWt .

(270)

This proves equation (249). A.10

P ROOF : M EMORIZATION -R EGIME P RECONDITIONER AND E FFECTIVE D IFFUSION I DENTITIES

The validation for Lemma 10 is provided in Figure 9. The experiment shows that, in the memorization regime, the approximated identities −1/2 √ −1/2 π(θt ) := diag(g ⊙ g t ) ⇒ tr(π(θt )) ≈ b tr (diag Σt ) (271) g ⊙ g t ≈ (1/b) diag Σt and 1 1/2 G(θt ) = √ B(t)π(θt )Σt ⇒ tr G(θt )G(θt )⊤ ≈ p b

(272)

hold in the memorization regime. Lemma 10 (Memorization-Regime Preconditioner and Effective Diffusion Identities). Let −1/2

π(θt ) := diag(g ⊙ g t ) 36

(273)

Preprint

and define 1 1/2 G(θt ) := √ B(t)π(θt )Σt , b

√ B(t) :=

1 − e−α2 t . 1 − e−α1 t

(274)

Assume that, in the memorization regime, 1 1 g ⊙ g t = ḡt ⊙ ḡt + diag(Σt ) ≈ diag(Σt ), b b and that the diagonal Adam preconditioner admits the trace-matched scalar approximation 1 1 π(θt ) ≈ s(θt )−1 Ip , := tr π(θ) . s(θ) p

(275)

(276)

Equivalently, under the corresponding scalar covariance approximation, 1 Σt ≈ s(θt )2 Ip . b Then the following two late-stage identities hold: √ −1/2 , tr(π(θt )) ≈ b tr (diag(Σt )) ⊤ tr G(θt )G(θt ) ≈ p.

(277)

(278) (279)

Moreover, G(θt )G(θt )⊤ ≈ Ip

(t → ∞).

(280)

Proof. From equation (275), g ⊙ gt ≈

1 diag(Σt ). b

(281)

Therefore −1/2

π(θt ) = diag(g ⊙ g t ) −1/2 1 ≈ diag diag(Σt ) b √ −1/2 = b diag(diag(Σt )) .

(282)

Taking traces gives tr(π(θt )) ≈

√

−1/2 b tr (diag(Σt )) .

(283)

This proves equation (278). Next, by definition, 1 1/2 G(θt ) = √ B(t)π(θt )Σt . b

(284)

Hence ⊤ B(t)2 1/2 1/2 π(θt )Σt Σt π(θt )⊤ b B(t)2 = π(θt )Σt π(θt )⊤ . b Since π(θt ) is diagonal, it is symmetric. Thus G(θt )G(θt )⊤ =

π(θt )⊤ = π(θt ),

(285)

(286)

and equation (285) becomes G(θt )G(θt )⊤ =

B(t)2 π(θt )Σt π(θt ). b 37

(287)

Preprint

Using the trace-matched scalar approximation π(θt ) ≈ s(θt )−1 Ip ,

(288)

we obtain B(t)2 s(θt )−1 Ip Σt s(θt )−1 Ip b B(t)2 Σt . = b s(θt )2

G(θt )G(θt )⊤ ≈

(289)

By 1 Σt ≈ s(θt )2 Ip , b

(290)

Σt ≈ b s(θt )2 Ip .

(291)

or equivalently,

Substituting equation (291) into equation (289) gives G(θt )G(θt )⊤ ≈

B(t)2 b s(θt )2 Ip b s(θt )2

= B(t)2 Ip .

(292)

Finally, B(t)2 =

1 − e−α2 t

2 →1

(1 − e−α1 t )

(t → ∞).

(293)

Therefore, G(θt )G(θt )⊤ ≈ Ip

(t → ∞).

(294)

tr G(θt )G(θt )⊤ ≈ tr(Ip ) = p.

(295)

Taking traces yields

This proves equations (279) and (280).

Preprint

A.11

P ROOF : R EDUCED M EMORIZATION -R EGIME R ADIUS SDE

We derive the exact closed-form reduced late-stage radius SDE with residual terms in Lemma 11. Its technical correctness is verified in Figure 10, where the predicted radius dynamics match the theoretical values computed from Lemma 11. ×104 t 22 (measured)

rt2 (late-stage radius SDE)

M G

rt2

s( t) = p/tr( t)

3 2 1 0

2000 4000 6000 8000 10000 12000 14000 16000

1.6 1.4 1.2 1.0 0.8 0.6 0.4 0.2 0.0

×10 1

s( t) M G

2000 4000 6000 8000 10000 12000 14000 16000

×104 gt 2

3.5

t 2

3.0

M G

RSM and R

2.5

gt 2

×10 3

2.0 1.5

2 4

1.0

RSM (slow-manifold) R RSM + R

0.5

M G

0.0

2000 4000 6000 8000 10000 12000 14000 16000

Figure 10: Late-Stage Radius SDE. This experiment shows the dynamics of the late-stage radius SDE in Lemma 11. The predicted radius dynamics closely match the theoretical values computed from Lemma 11 without counting residual terms, validating the correctness of the decomposition. In particular, the late-stage residual sum RSM + Rπ is negligible, as hypothesized. Lemma 11 (Reduced Late-Stage Radius SDE with Slow-Manifold, Preconditioned Residual (Restated)). Let rt := ∥θt ∥2 , where rt2 follows the residual-corrected preconditioned radius SDE in equation (249) of Lemma 9 (Preconditioned-Decomposition of Adam’s Radius SDE). In the latestage regime, assume √ 1 − e−α2 t → 1, G(θt )G(θt )⊤ ≈ Ip , (296) B(t) := 1 − e−α1 t so that the residuals from a(t) and the diffusion covariance are neglected. Then the reduced latestage squared-radius SDE is 2λ 2 2 ⊤ √ (r) drt2 ≈ − rt − θt ḡt − λθt + η p + RSM (t) + Rπ (t) dt + 2 η rt dWt , (297) s(θt ) s(θt ) where q √ −α1 t ⊘ g ⊙ gt , E [mt ⊘ vt ] − ḡt 1 − e Rπ (t) := −2θt⊤ π(θt ) − s(θt )−1 Ip ḡt , Z t θ (r) Wt := er (θu )⊤ G(θu )dWu , er (θ) := , ∥θ∥2 0

RSM (t) := −2B(t)θt⊤

(r)

and under G(θt )G(θt )⊤ ≈ Ip , Wt

is a one-dimensional Brownian motion. 39

(298) (299) (300)

Preprint

Remark 9. Moreover, √ RSM (t) + Rπ (t) = −2B(t)θt⊤ E [mt ⊘ vt ] + 2a(t)θt⊤ π(θt )ḡt 2 ⊤ θ ḡt − 2θt⊤ π(θt )ḡt + s(θt ) t √ = −2B(t)θt⊤ E [mt ⊘ vt ] + 2 a(t) − 1 θt⊤ π(θt )ḡt +

2 ⊤ θ ḡt . s(θt ) t

(301)

In particular, in the late-stage limit a(t) → 1, RSM (t) + Rπ (t) ≈ −2B(t)θt⊤ E [mt ⊘

√

vt ] +

2 ⊤ θ ḡt . s(θt ) t

(302)

Proof. Starting from the residual-corrected preconditioned radius SDE in equation (249) of Lemma 9 (Preconditioned-Decomposition of Adam’s Radius SDE), drt2 ≈ −2a(t)θt⊤ π(θt )ḡt + RSM (t) + η tr G(θt )G(θt )⊤ dt √ + 2 η θt⊤ G(θt )dWt . (303) Under a(t) → 1, −2a(t)θt⊤ π(θt )ḡt ≈ −2θt⊤ π(θt )ḡt .

(304)

Decompose the preconditioner as π(θt ) = s(θt )−1 Ip + π(θt ) − s(θt )−1 Ip .

(305)

Substituting this decomposition gives 2 ⊤ θ ḡt + Rπ (t), s(θt ) t

(306)

Rπ (t) = −2θt⊤ π(θt ) − s(θt )−1 Ip ḡt .

(307)

−2θt⊤ π(θt )ḡt = − where

Next, decompose the mean gradient as ḡt = λθt + ḡt − λθt .

(308)

2λ 2 2 ⊤ 2 ⊤ θt ḡt = − rt − θt ḡt − λθt . s(θt ) s(θt ) s(θt )

(309)

Then −

Combining equations (306) and (309), we obtain −2a(t)θt⊤ π(θt )ḡt ≈ −

2λ 2 2 ⊤ r − θ ḡt − λθt + Rπ (t). s(θt ) t s(θt ) t

(310)

Under G(θt )G(θt )⊤ ≈ Ip , η tr G(θt )G(θt )⊤ ≈ η p.

(311)

√ √ 2 η θt⊤ G(θt )dWt = 2 η rt er (θt )⊤ G(θt )dWt √ (r) = 2 η rt dWt .

(312)

Also,

The quadratic variation is ⟨W

(r)

Z t ⟩t =

er (θu )⊤ G(θu )G(θu )⊤ er (θu )du

Z t ≈

er (θu )⊤ er (θu )du = t.

(313)

Preprint

(r)

Thus, by Lévy’s characterization (Karatzas & Shreve, 1991), Wt motion.

is a one-dimensional Brownian

Substituting equations (310), (311), and (312) into equation (303) yields equation (297) 2λ 2 2 ⊤ √ (r) drt2 ≈ − rt − θt ḡt − λθt + η p + RSM (t) + Rπ (t) dt + 2 η rt dWt . (314) s(θt ) s(θt )

Remark 10. Empirical study in Figure 10 shows that RSM (t) + Rπ (t) is negligible during the memorization regime, so that the radius SDE in grokking can be approximated through 2λ 2 2 ⊤ √ (r) 2 drt ≈ − rt − θt ḡt − λθt + η p dt + 2 η rt dWt . (315) s(θt ) s(θt ) We now analyze the cancellation of RSM (t) + Rπ (t). Expanding RSM (t) gives √ RSM (t) = −2B(t)θt⊤ E [mt ⊘ vt ] q ⊤ −α1 t + 2B(t)θt ḡt 1 − e ⊘ g ⊙ gt .

(316)

Using B(t) 1 − e−α1 t = a(t),

q g ⊙ g t = π(θt )ḡt ,

(317)

vt ] + 2a(t)θt⊤ π(θt )ḡt .

(318)

ḡt ⊘

we obtain RSM (t) = −2B(t)θt⊤ E [mt ⊘

√

Adding the explicit expression for Rπ (t), √ RSM (t) + Rπ (t) = −2B(t)θt⊤ E [mt ⊘ vt ] + 2a(t)θt⊤ π(θt )ḡt − 2θt⊤ π(θt ) − s(θt )−1 Ip ḡt √ = −2B(t)θt⊤ E [mt ⊘ vt ] + 2 a(t) − 1 θt⊤ π(θt )ḡt +

2 ⊤ θ ḡt . s(θt ) t

(319)

This proves equation (301). In the late-stage limit a(t), B(t) → 1, the middle term vanishes, and therefore √ 2 ⊤ RSM (t) + Rπ (t) ≈ −2B(t)θt⊤ E [mt ⊘ vt ] + θ ḡt . (320) s(θt ) t Hence the residual sum is small when gradient and its first-moment estimate are near zeros. A.12

P ROOF : S CALING L AW OF M EMORIZATION R ADIUS

Theorem 4 (Scaling Law of Memorization Radius (Restated)). Consider the joint Adam SDE in equation (167) of Lemma 5 (Adam’s Joint-State Continuous-Time SDE Limit) on the joint state ⊤ ⊤ 3p St = (θt⊤ , m⊤ t , vt ) ∈ R , p (321) dSt = µ(St )dt + η/b σ(St )dWt , and assume that the bias-correction factor satisfies B(t) → 1 in the early hitting regime. Let τM (S0 ) := inf{t ≥ 0 : θt (S0 ) ∈ ∂M } be the first-hitting time of the memorization boundary. Define the exit-value function u(S) := E ∥θτM (S0 ) ∥22 S0 = S ,

(322)

(323)

and the mean squared memorization radius ρ2M := ES0 ∼ΘS [u(S0 )] , 41

(324)

Preprint

⊤ ⊤ where ΘS is the initialization distribution of the joint state, with S0 = (θ0⊤ , 0⊤ p , 0p ) .

Let Eθ ∈ R3p×3p be the θ-block projector Eθ :=

Ip 0 0

! 0 0 0 0 , 0 0

(325)

so that for any joint state S = (θ, m, v), ∥θ∥22 = S ⊤ Eθ S.

(326)

Let Φt (S0 ) denote the deterministic Adam flow generated by the drift field µ(S), d Φt (S0 ) = µ(Φt (S0 )), Φ0 (S0 ) = S0 . dt Define the deterministic first-hitting time (0)

(327)

τM (S0 ) := inf{t ≥ 0 : θ(Φt (S0 )) ∈ ∂M }, and the deterministic exit-radius-squared map

(328)

R2 (S0 ) := Φτ (0) (S0 ) (S0 )⊤ Eθ Φτ (0) (S0 ) (S0 ).

(329)

ΣS (S) := σ(S)σ(S)⊤ be the joint-state diffusion covariance. Then, for sufficiently small η/b,

(330)

Let

(0) 2

ρ2M = ρM

(2) + (η/b)cM + (η/b)2 cM + O (η/b)3 ,

(331)

where (0) 2

:= ES0 ∼ΘS R2 (S0 ) , (0) Z τM (S0 ) 1 tr ΣS (Φt (S0 ))∇2S R2 (Φt (S0 )) dt, cM := ES0 ∼ΘS 2 0 (0) Z τM (S0 ) 1 (2) cM := ES0 ∼ΘS tr ΣS (Φt (S0 ))∇2S u1 (Φt (S0 )) dt. 2 0 Here u1 is the first-order perturbation corrector (0) Z τM (S0 ) 1 u1 (S0 ) := tr ΣS (Φt (S0 ))∇2S R2 (Φt (S0 )) dt. 2 0 ρM

(0) 2

(332) (333) (334)

(335)

(2)

The constants ρM , cM , and cM are task-determined and independent of η and b; ℓ2 regular(0) (0) ization coefficient λ enters (ρM )2 through deterministic gradient flow with scaling law (ρM )2 ∝ exp(−O(λ)). Proof. Let τM (S0 ) := inf{t ≥ 0 : θt (S0 ) ∈ ∂M } (336) be the first-hitting time of the memorization boundary for the trajectory {θt } starting from S0 = ⊤ ⊤ (θ0⊤ , 0⊤ p , 0p ) . Bilinear Radius Projector. the θ-block projector

Since the parameter θt is the θ-component of the joint state St , define

! Ip 0 0 Eθ := 0 0 0 ∈ R3p×3p . 0 0 0 Thus, for any joint state S, the squared parameter radius is the joint-state quadratic form r2 (S) := ∥θ∥22 = S ⊤ Eθ S.

(337)

(338)

In particular, ∇S r2 (S) = 2Eθ S,

∇2S r2 (S) = 2Eθ . 42

(339)

Preprint

Define Dirichlet Boundary Problem. We introduce the exit-value function u(S) := E r2 (SτM ) S0 = S = E Sτ⊤M Eθ SτM S0 = S ,

(340)

which is the expected squared parameter norm at the first-hitting time of ∂M , starting from S. We now derive the PDE satisfied by u. The joint Adam SDE has the form p dSt = µ(St ) dt + η/b σ(St ) dWt .

(341)

Therefore its infinitesimal generator is LS f (S) = µ(S)⊤ ∇S f (S) +

1 η tr ΣS (S)∇2S f (S) , 2 b

ΣS (S) := σ(S)σ(S)⊤ .

(342)

Learning rate η and batch size b modulate infinitesimal generator

Remark 11. The structure of the infinitesimal generator LS f (S) shows that this operator contains a term modulated by the learning rate η and the batch size b through the coefficient η/b. This motivates us to treat LS f (S) as an operator expansion with respect to the coefficient η/b, so that we can apply a regular perturbation expansion to study the corresponding expansion of the memorization radius with respect to η/b. Dirichlet Boundary PDE. Since u(S) is the expected boundary value of the process stopped at ∂M , it solves the Dirichlet boundary value problem (Øksendal, 2003) LS u(S) = 0, ⊤

u(S) = S Eθ S,

θ ∈ M ◦,

(343)

θ ∈ ∂M.

(344)

Here equation (343) states that u(St ) is harmonic with respect to the stopped Adam diffusion before hitting the boundary, while equation (344) assigns the squared radius as the boundary payoff. Let ε := η/b. We split the generator into its deterministic and stochastic parts: (0)

(1)

(0)

(1)

LS f := µ⊤ ∇S f,

LS f :=

Thus the Dirichlet PDE becomes (0) (1) LS + εLS u(S) = 0,

θ ∈ M ◦,

LS = LS + εLS ,

1 tr ΣS ∇2S f . 2

(345)

(346)

with boundary condition u(S) = S ⊤ Eθ S, Solving PDE via Regular Perturbation Expansion. expansion in the small parameter ε. Write

θ ∈ ∂M.

(347)

We solve this PDE by a regular perturbation

u(S) = u0 (S) + εu1 (S) + ε2 u2 (S) + O(ε3 ). Substituting equation (348) into equation (346) gives (0) (1) 0 = LS + εLS u0 + εu1 + ε2 u2 + O(ε3 ) (0) (0) (1) (0) (1) = LS u0 + ε LS u1 + LS u0 + ε2 LS u2 + LS u1 + O(ε3 ).

(348)

(349)

Since this identity must hold for all sufficiently small ε, each coefficient of ε must vanish. Therefore, (0)

ε0 :

LS u0 = 0,

(350)

ε1 :

(0) (1) LS u1 = −LS u0 , (0) (1) LS u2 = −LS u1 .

(351)

ε :

(352)

Preprint

The boundary condition is expanded in the same way: u0 + εu1 + ε2 u2 + O(ε3 ) = S ⊤ Eθ S,

θ ∈ ∂M.

(353)

Matching powers of ε on the boundary gives u0 θ∈∂M = S ⊤ Eθ S,

(354)

u1 θ∈∂M = 0,

(355)

u2 θ∈∂M = 0.

(356)

Zero-Order Expansion. We now solve these equations by characteristics. The characteristic (0) curves of LS = µ⊤ ∇S are the deterministic Adam trajectories d Φt (S0 ) = µ(Φt (S0 )), dt

Φ0 (S0 ) = S0 ,

(357)

so that θ̇(0) = −m(0) ⊘

v (0) ,

ṁ(0) = −α1 m(0) − ḡ(θ(0) ) , v̇ (0) = −α2 v (0) − g ⊙ g(θ(0) ) .

(358)

For any smooth function w, along the deterministic flow we have d (0) w(Φt (S0 )) = LS w(Φt (S0 )). dt

(359)

(0)

Let τM (S0 ) be the deterministic hitting time (0)

τM (S0 ) := inf{t ≥ 0 : (Φt (S0 ))θ ∈ ∂M }.

(360)

At order ε0 , equation (350) gives d u0 (Φt (S0 )) = 0. (361) dt Thus u0 is constant along the deterministic trajectory. Evaluating it at the deterministic hitting time and using the boundary condition in equation (354), we obtain u0 (S0 ) = u0 (Φτ (0) (S0 ) (S0 )) M

= Φτ (0) (S0 ) (S0 )⊤ Eθ Φτ (0) (S0 ) (S0 ) M

2 Φτ (0) (S0 ) (S0 ) θ .

(362)

Define the deterministic exit-radius-squared map R2 (S0 ) := Φτ (0) (S0 ) (S0 )⊤ Eθ Φτ (0) (S0 ) (S0 ). M

(363)

Equivalently, R2 (S0 ) =

2 Φτ (0) (S0 ) (S0 ) θ . 2

(364)

Then u0 (S0 ) = R2 (S0 ). First-Order Expansion.

At order ε1 , equation (351) and the transport identity give d (1) u1 (Φt (S0 )) = −LS u0 (Φt (S0 )). dt

(365)

(0)

Integrating from 0 to τM (S0 ) gives (0) Z τM (S0 )

u1 (Φτ (0) (S0 ) (S0 )) − u1 (S0 ) = − M

(1)

LS u0 (Φt (S0 )) dt.

(366)

Preprint

Since u1 = 0 on ∂M , the first term on the left vanishes. Therefore, (0) Z τM (S0 ) (1) LS u0 (Φt (S0 )) dt. u1 (S0 ) =

(367)

(1)

Using the definition of LS and u0 = R2 along the deterministic flow, (0) Z τM (S0 ) 1 u1 (S0 ) = tr ΣS (Φt (S0 ))∇2S R2 (Φt (S0 )) dt. 2 0 Second-Order Expansion.

(368)

At order ε2 , the same argument gives

d (1) u2 (Φt (S0 )) = −LS u1 (Φt (S0 )). dt Since u2 = 0 on ∂M , integrating along the deterministic characteristic yields (0) Z τM (S0 ) (1) u2 (S0 ) = LS u1 (Φt (S0 )) dt.

(369)

(370)

Thus, (0) Z τM (S0 )

u2 (S0 ) = 0

Produce Claims.

1 tr ΣS (Φt (S0 ))∇2S u1 (Φt (S0 )) dt. 2

(371)

Combining the three orders, we have u(S0 ) = R2 (S0 ) + εu1 (S0 ) + ε2 u2 (S0 ) + O(ε3 ).

(372)

Finally, the mean squared memorization radius is obtained by averaging this exit-value function over the initialization distribution: ρ2M := E ∥θτM ∥22 = E Sτ⊤M Eθ SτM = ES0 ∼ΘS [u(S0 )] = ES0 ∼ΘS R2 (S0 ) + εES0 ∼ΘS [u1 (S0 )] + ε2 ES0 ∼ΘS [u2 (S0 )] + O(ε3 ). (373) Define (0) 2

ρM

:= ES0 ∼ΘS R2 (S0 ) ,

cM := ES0 ∼ΘS [u1 (S0 )] , (2)

cM := ES0 ∼ΘS [u2 (S0 )] .

(374)

Substituting these definitions into equation (373) gives (0) 2

ρ2M = ρM

(2)

+ εcM + ε2 cM + O(ε3 ).

(375)

(2) + (η/b)cM + (η/b)2 cM + O (η/b)3 ,

(376)

Since ε = η/b, we obtain (0) 2

ρ2M = ρM which proves the theorem. (0)

(0)

Dependence of ρM on λ under weak task gradient. The leading constant (ρM )2 is determined by the deterministic gradient flow. In the regime where the task gradient is small, ∇Lf∗ (θt ) ≈ 0,

(377)

the deterministic flow is dominated by the ℓ2 drift: dθt ≈ −λθt dt.

(378)

Hence d∥θt ∥22 = 2θt⊤ dθt ≈ −2λ∥θt ∥22 dt. 45

(379)

Preprint

Equivalently, 1 d∥θt ∥22 ≈ −2λ dt. ∥θt ∥22 Integrating along the deterministic memorization path gives (0) ∥θτ (0) ∥22 ≈ ∥θ0 ∥22 exp −2λτM .

(380)

(381)

Therefore, (0) (0) (ρM )2 ∝ exp −O(λτM ) .

(382)

(0)

When τM is weakly dependent on λ in the considered regime, (0)

(ρM )2 ∝ exp(−O(λ)) .

(383)

ρ2M = (ρM )2 + O(η),

(0)

(384)

ρ2M ∝ exp(−O(λ)) .

(385)

Since

the same leading dependence gives

Preprint

A.13

P ROOF : S CALING L AW OF G ENERALIZATION R ADIUS

Theorem 5 (Scaling Law of Generalization Radius, restated). We do not assume a global minizer, for each initialization S0 ∈ ΘS , let λ ⋆ ∗ 2 θ (S0 ) := arg min Lf (θ(S0 )) + ∥θ(S0 )∥2 ∈ G (386) 2 θ(S0 ) denote the regularized local minimizer selected by the trajectory starting from initial state S0 , and Lf∗ (θ(S0 )) is the task loss. The trajectory {St } will be confined in the basin centered at θ⋆ (S0 ). Define (0) 2 ρG := ES0 ∈ΘS ∥θ⋆ (S0 )∥22 . (387) Let τG (S0 ) := inf{t ≥ 0 : θt (S0 ) ∈ ∂G},

(388)

and define the generalization radius by ρ2G := ES0 ∈ΘS E ∥θτG (S0 ) ∥22 | S0 .

(389)

Let D⋆ (S0 ) := G(θ⋆ (S0 ))G(θ⋆ (S0 ))⊤ ,

s⋆ (S0 ) := s(θ⋆ (S0 )),

(390)

where the effective scalar preconditioner at the selected minimizer is defined by 1 s⋆ (S0 )

1 tr π(θ⋆ (S0 )) , p

−1/2

π(θ⋆ (S0 )) = diag(g ⊙ g(θ⋆ (S0 )))

(391)

Then, generalization radius approximately admits, ρ2G ≈

η (0) 2 ρG + cG + O λ

2 η , λ2

(392)

where cG := ES0 ∈ΘS

1 ⋆ ⋆ 2 s (S0 ) tr(D (S0 )) > 0.

(393)

Proof. For a trajectory {St } starting from S0 ∈ ΘS , write θ⋆ := θ⋆ (S0 ),

θt = θ ⋆ + δ t .

(394)

By definition of the local minimizer, ḡ ∗ (θ⋆ ) + λθ⋆ = 0,

(395)

∗

where ḡ (•) denotes the task loss mean. By Lemma 11 (Reduced Late-Stage Radius SDE with SlowManifold, Preconditioned Residual (Restated)), together with Lemma 10 (Memorization-Regime Preconditioner and Effective Diffusion Identities), G(θt )G(θt )⊤ ≈ Ip ,

(396)

RSM (t) + Rπ (t) ≈ 0,

(397)

and the residual

the reduced late-stage squared-radius SDE gives 2 ⊤ ∗ √ (r) drt2 ≈ − θt (ḡt + λθt ) + ηp dt + 2 η rt dWt , s(θt )

(398)

where rt2 = ∥θt ∥22 . 47

(399)

Preprint

Linearization Near Local Minizer.

The Taylor expansion of ḡ ∗ around θ⋆ gives

ḡt∗ = ḡ ∗ (θ⋆ + δt ) = ḡ ∗ (θ⋆ ) + ∇ḡ ∗ (θ⋆ )δt + O ∥δt ∥22 .

(400)

Therefore, ḡt∗ + λθt = ḡ ∗ (θ⋆ ) + ∇ḡ ∗ (θ⋆ )δt + λθ⋆ + λδt + O ∥δt ∥22 = (∇ḡ ∗ (θ⋆ ) + λIp ) δt + O ∥δt ∥22 ,

(401)

where equation (395) was used. Under the dominant-regularization approximation, ∇ḡ ∗ (θ⋆ ) + λIp ≈ λIp ,

(402)

ḡt∗ + λθt ≈ λδt + O ∥δt ∥22

(403)

Using θt = θ⋆ + δt , θt⊤ (ḡt∗ + λθt ) ≈ (θ⋆ + δt )⊤ λδt + O ∥δt ∥22 = λθ⋆⊤ δt + λ∥δt ∥22 + O ∥δt ∥22 . Local Stationary Approximation.

(404)

Under the local stationary approximation,

⋆

s⋆ (S0 ) := s(θ⋆ (S0 )).

s(θt ) = s (S0 ) + O(∥δt ∥2 ) ,

(405)

Taking conditional expectation in equation (398) gives 2 2 η 0≈− ⋆ E θt⊤ (ḡt∗ + λθt ) | S0 + ηp + O 2 , s (S0 ) λ

(406)

h i (r) E rt dWt | S0 = 0.

(407)

since

The local stationary fluctuation is centered, so E[θ⋆⊤ δt | S0 ] = 0.

E[δt | S0 ] = 0,

(408)

Combining equations (404), (406), and (408), 2 2λ η 2 0≈− ⋆ E ∥δt ∥2 | S0 + ηp + O 2 . s (S0 ) λ

(409)

2 η s⋆ (S0 ) η 2 p+O 2 . E ∥δt ∥2 | S0 ≈ λ 2 λ

(410)

Hence

Also, ∥θt ∥22 = ∥θ⋆ + δt ∥22 = ∥θ⋆ ∥22 + 2θ⋆⊤ δt + ∥δt ∥22 .

(411)

Taking conditional expectation gives E ∥θt ∥22 | S0 = ∥θ⋆ (S0 )∥22 + E ∥δt ∥22 | S0 ≈ ∥θ⋆ (S0 )∥22 +

2 η s⋆ (S0 ) η p+O 2 . λ 2 λ

(412)

Solving Generalization Radius. Under the local stationary approximation in the generalization basin, E ∥θτG (S0 ) ∥22 | S0 ≈ E ∥θt ∥22 | S0 . (413) 48

Preprint

Therefore, 2 η η s⋆ (S0 ) 2 ⋆ 2 E ∥θτG (S0 ) ∥2 | S0 ≈ ∥θ (S0 )∥2 + p+O 2 . λ 2 λ

(414)

Averaging over S0 ∈ ΘS , ρ2G := ES0 ∈ΘS E ∥θτG ∥22 | S0 ⋆ 2 ⋆ η s (S0 ) η 2 ≈ ES0 ∈ΘS ∥θ (S0 )∥2 + ES0 ∈ΘS p +O 2 . λ 2 λ

(415)

By definition, (0) 2

ρG

:= ES0 ∈ΘS ∥θ⋆ (S0 )∥22 , ⋆ s (S0 ) cG := ES0 ∈ΘS p . 2

(416) (417)

Thus, (0) 2

ρ2G ≈ ρG

2 η η cG + O 2 , λ λ

(418)

which proves the claimed scaling law. Weak Batch-Size b-Dependence of cG . The weak batch-size dependence of cG enters through the effective scalar preconditioner s⋆ (S0 ). From the definition of cG , ⋆ s (S0 ) cG := ES0 ∈ΘS p . (419) 2 By the late-stage preconditioning identity, 1 1 = tr (π(θ⋆ (S0 ))) , s⋆ (S0 ) p

(420) −1/2

π(θ⋆ (S0 )) = diag (g ⊙ g(θ⋆ (S0 )))

(421)

In the memorization regime, g ⊙ g(θ⋆ (S0 )) ≈

1 diag(Σ(θ⋆ (S0 ))). b

(422)

Therefore, √ −1/2 , π(θ⋆ (S0 )) ≈ b diag (diag(Σ(θ⋆ (S0 )))) √ 1 = O( b), s⋆ (S0 ) 1 ⋆ s (S0 ) = O √ . b

(423) (424) (425)

Using the late-stage identity G(θ)G(θ)⊤ ≈ Ip , the diffusion trace contribution is tr(G(θ)G(θ)⊤ ) ≈ p. Hence ⋆ s (S0 ) 1 cG = ES0 ∈ΘS p =O √ . (426) 2 b

Preprint

A.14

P ROOF : S CALING L AW OF S OLUTION T RANSITION T IME

Theorem 6 (Scaling Law of Solution Transition Time, restated). Assume the reduced late-stage radius SDE in Lemma 11 (Reduced Late-Stage Radius SDE with Slow-Manifold, Preconditioned Residual (Restated)), with conditions in Lemma 10 (Memorization-Regime Preconditioner and Effective Diffusion Identities) G(θt )G(θt )⊤ ≈ Ip ,

(427)

RSM (t) + Rπ (t) ≈ 0.

(428)

and the residual approximation Let ḡ(θ) be the population mean gradient and let s(θ) > 0 denote the scalar preconditioning factor satisfying 1 1 = tr π(θ) . (429) s(θ) p (0)

Let θt (S0 ) be the deterministic late-stage trajectory from ∂M to ∂G starting with initial state S0 , (0) and let τM →G (S0 ) be its deterministic transition time. For each deterministic late-stage trajectory (0) θt (S0 ) from ∂M to ∂G, define its pathwise effective radial contraction rate by " # (0) Z τM (0)⊤ (0) (0) →G (S0 ) θt (S0 ) ḡ(θt (S0 )) − λθt (S0 ) λ 1 + dt. κ(S0 ) := (0) (0) (0) (0) τM →G (S0 ) 0 s(θt (S0 )) s(θt (S0 ))∥θt (S0 )∥22 (430) The scalar s̄ is defined through the ensemble-averaged effective contraction rate λ := ES0 ∈ΘS [κ(S0 )] . s̄

(431)

Let ρM and ρG be the memorization and generalization radii. Assume the asymptotic expansions η (0) (432) ρ2M = (ρM )2 + cM + O (η/b)2 , b 2 η η (0) 2 ρ2G ≈ ρG + cG + O 2 , (433) λ λ where cM is the memorization-radius correction from Theorem 4, and (0) 2 ρG := ES0 ∈ΘS ∥θ⋆ (S0 )∥22 , (434) λ θ⋆ (S0 ) := arg min Lf∗ (θ(S0 )) + ∥θ(S0 )∥22 , (435) 2 θ(S0 ) ⋆ s (S0 ) cG := ES0 ∈ΘS tr(D⋆ (S0 )) , (436) 2 D⋆ (S0 ) := G(θ⋆ (S0 ))G(θ⋆ (S0 ))⊤ .

s⋆ (S0 ) := s(θ⋆ (S0 )), ⋆

(437)

⊤

⋆

Under the late-stage identity G(θ (S0 ))G(θ (S0 )) ≈ Ip , equation (436) reduces to ⋆ s (S0 ) p . cG = ES0 ∈ΘS 2 Then, in iteration time, the expected transition time from ∂M to ∂G satisfies (0) (2) η ρ cτ s̄ cτ ES0 ∈ΘS τM →G ≈ log M + + +O 3 , 2 (0) ηλ bλ λ λ ρG where cM cτ := s̄ , (0) 2(ρM )2 " # 1 cG s̄2 (p − 2) 1 (2) cτ := −s̄ + − (0) . (0) (0) 4 2(ρ )2 (ρ )2 (ρ )2 G

(438)

(439)

(440) (441)

Preprint

(2)

The constants s̄, cτ , and cτ are independent of η and λ in the considered scaling regime; a weak batch-size b-dependence enters through the scalar preconditioning scale s̄ = O( √1b ), and hence (2)

cτ = O( √1b ) and cτ = O( 1b ). Proof. Let rt := ∥θt ∥2 .

(442)

By Lemma 11 (Reduced Late-Stage Radius SDE with Slow-Manifold, Preconditioned Residual (Restated)), under the late-stage observations a(t) → 1,

G(θt )G(θt )⊤ ≈ Ip ,

RSM (t) + Rπ (t) ≈ 0,

the squared-radius process satisfies 2λ 2 2 ⊤ √ (r) 2 drt ≈ − rt − θt ḡt − λθt + η p dt + 2 η rt dWt . s(θt ) s(θt )

(443)

(444)

Effective Deterministic Radius Flow. The leading deterministic radial flow is obtained from equation (444) by keeping the O(1) drift terms and dropping the O(η) Itô correction. Thus, 2λ 2 2 ⊤ drt2 ≈ − rt − θt ḡt − λθt dt. (445) s(θt ) s(θt ) Since drt2 = 2rt drt + O(η)dt,

(446)

the leading deterministic radial drift is 1 λ rt − θt⊤ ḡt − λθt s(θt ) s(θt )rt " # θt⊤ ḡt − λθt λ + =− rt . s(θt ) s(θt )∥θt ∥22

ṙt ≈ −

(447)

(0)

For each deterministic late-stage trajectory θt (S0 ) from ∂M to ∂G, define its pathwise effective radial contraction rate by " # (0) Z τM (0) (0) (0)⊤ →G (S0 ) θt (S0 ) ḡ(θt (S0 )) − λθt (S0 ) λ 1 + dt. κ(S0 ) := (0) (0) (0) (0) τM →G (S0 ) 0 s(θt (S0 )) s(θt (S0 ))∥θt (S0 )∥22 (448) The scalar s̄ is defined through the ensemble-averaged effective contraction rate λ := ES0 ∈ΘS [κ(S0 )] . (449) s̄ Thus, in the averaged effective radial model, the leading deterministic late-stage radius flow is λ (0) = − rt . s̄ Therefore, along the averaged late-stage transition dynamics, (0)

ṙt

−

λ 1 λ rt − θt⊤ ḡt − λθt ≈ − rt . s(θt ) s(θt )rt s̄

(450)

(451)

Effective Radius SDE. We now pass from the squared-radius SDE to the radius SDE. Applying Itô’s lemma to rt = (rt2 )1/2 gives drt =

1 1 dr2 − 3 d⟨r2 ⟩t . 2rt t 8rt 51

(452)

Preprint

From the martingale term in equation (444), d⟨r2 ⟩t = 4ηrt2 dt.

(453)

Substituting equations (444) and (453) into equation (452), and using equation (451), yields λ p 1 √ (r) drt = − rt + η −η dt + η dWt s̄ 2rt 2rt λ p−1 √ (r) = − rt + η (454) dt + η dWt . s̄ 2rt Let κ :=

λ , s̄

ε := η.

(455)

Then equation (454) becomes √ ε(p − 1) (r) drt = −κrt + dt + ε dWt . 2rt

(456)

Mean First-Passage Time PDE. Let T (r) be the expected continuous time for the process in equation (456), initialized at radius r, to hit the absorbing boundary rG , defined as T (r) := E[τG | r0 = r] ,

r > rG .

(457)

The transition starts at the memorization boundary and ends at the generalization boundary: r0 := ρM ,

rG := ρG .

(458)

d ε d2 [•] + [•]. dr 2 dr2

(459)

T (rG ) = 0.

(460)

Its infinitesimal generator is Lr [•] =

−κr +

ε(p − 1) 2r

Then T (r) solves a Dirichlet problem Lr T (r) = −1,

Solving Absorbing Boundary Dirichlet PDE. Equivalently, ε(p − 1) ε ′′ T (r) + −κr + T ′ (r) = −1. 2 2r

(461)

Writing q(r) := T ′ (r),

(462)

2κ p−1 2 q (r) + − r + q(r) = − . ε r ε

(463)

κ I(r) = rp−1 exp − r2 . ε

(464)

I(r)q(r) → 0

(465)

we obtain ′

The integrating factor is

Using as r → ∞,

we obtain 2 1 q(r) = ε I(r)

Z ∞

I(y)dy. r

(466)

Preprint

Since T (rG ) = 0, for an initial radius r0 > rG , Z Z ∞ 2 r0 1 T (r0 ) = I(y)dy dx. ε rG I(x) x

(467)

Define β :=

κ λ = . ε s̄ η

(468)

Then 2

I(r) = rp−1 e−βr ,

(469)

and Z ∞

y p−1 e−βy dy =

1 −p/2 β Γ(p/2, βx2 ), 2

(470)

where Z ∞ Γ(a, z) :=

ta−1 e−t dt.

(471)

Thus, 1 T (r0 ) = β −p/2 ε Small-Learning-Rate Regime.

Z r0

eβx Γ(p/2, βx2 )dx. p−1 rG x

(472)

The learning rate η → 0 is small, so βx2 =

λ 2 x ≫ 1. s̄ η

(473)

Therefore, 2

2 p/2−1 −βx2

Γ(p/2, βx ) = (βx )

p−2 2 −2 1+ + O (βx ) . 2βx2

Substituting equation (474) into equation (472) gives Z r0 1 p−2 1 2 −2 −1 + T (r0 ) = + O (βx ) x dx εβ rG x 2βx3 2 1 r0 1 ε ε(p − 2) 1 = log + . 2 − r 2 + O κ3 κ rG 4κ2 rG 0

(474)

(475)

Substituting κ = λ/s̄ and ε = η yields T (r0 ) =

2 r0 s̄2 (p − 2) 1 1 η s̄ log +η − + O . 2 2 2 λ rG 4λ rG r0 λ3

(476)

Asymptotic Expansions. The transition starts at the memorization radius and ends at the generalization radius. By Theorem 4, η (0) ρ2M = (ρM )2 + cM + O (η/b)2 . (477) b Taking the square root gives η cM (0) r0 = ρM = ρM + (478) + O (η/b)2 . (0) b 2ρ M

Similarly, by Theorem 5, η (0) ρ2G = (ρG )2 + cG + O λ

2 η , λ2

(479)

Preprint

where (0) (ρG )2 = ES0 ∈ΘS ∥θ⋆ (S0 )∥22 ,

(480)

and θ⋆ (S0 ) = arg min

θ(S0 )

Lf∗ (θ(S0 )) +

λ ∥θ(S0 )∥22 . 2

(481)

Therefore, (0)

rG = ρG = ρG +

2 η cG η . + O λ 2ρ(0) λ2

(482)

∆x + O(∆x2 ), x

(483)

Using log(x + ∆x) = log x + we have η cM + O (η/b)2 , b 2(ρ(0) )2 M 2 η η cG (0) + O log rG = log ρG + . λ 2(ρ(0) )2 λ2 G (0)

log r0 = log ρM +

(484) (485)

Hence 2 (0) ρM η cM η cG η r0 − +O 2 . log = log (0) + (0) 2 (0) 2 rG b λ λ ρ 2(ρ ) 2(ρ ) G

(486)

In the explicit O(η) Itô correction term of equation (476), it is sufficient to use the leading asymptotics: 1 1 1 1 − (0) + O(η/b) + O(η/λ). (487) 2 − r2 = (0) 2 rG (ρ ) (ρ )2 0 G

Substituting equations (486) and (487) into equation (476) gives (0)

TM →G =

ρ s̄ log M (0) λ ρG η cM + s̄ bλ 2(ρ(0) )2 ( M " #) η cG s̄2 (p − 2) 1 1 + 2 −s̄ + − (0) (0) (0) λ 4 2(ρG )2 (ρG )2 (ρM )2 2 η +O 3 . λ

(488)

Define cτ := s̄

cM (0)

2(ρM )2

c(2) τ := −s̄

" # s̄2 (p − 2) 1 1 + − (0) . (0) (0) 4 2(ρG )2 (ρG )2 (ρM )2

(489)

2 (0) ρ s̄ η η (2) η log M + c + c + O . τ τ 2 (0) λ bλ λ λ3 ρ

(490)

Then TM →G ≈

Preprint

Computing Number of Iterations. time interpolation is t = ηk:

We compute the number of iterations, since the continuous-

ES0 ∈ΘS [τM →G ] =

TM →G . η

(491)

Dividing equation (490) by η gives (0)

ES0 ∈ΘS [τM →G ] ≈

(2)

η ρ s̄ cτ cτ + O , log M + + (0) ηλ bλ λ2 λ3 ρ

(492)

which proves equation (439). ℓ2 Regularization Coefficient Dependence. In the late-stage radial SDE, the diffusion scale is ε = η and is independent of batch size; the leading first-passage time is obtained by setting ε = 0 and is determined only by the deterministic radial √ drift −λr/s̄. Batch size enters the leading term through the scalar preconditioning scale s̄ = O(1/ b), and the corrections through the memorization asymptotic ρM — whose O(η/b) correction becomes O(1/b) after converting to iteration time √ (2) — and through cτ = O(1/ b) and cτ = O(1/b).

Preprint

3.50

A DDITIONAL R ESULTS : S CALING L AWS OF M ANIFOLD R ADIUS ρ2M ON Z127

×104

×104 mean over seeds ±1 theoretical

3.25

4.4

3.00

×104 mean over seeds ±1 theoretical

3.0

4.2 2 M

2.75 2 M

3.5

mean over seeds ±1 theoretical

2.50

2.5

4.0

2 M

A.15

2.0

3.8

2.25

1.5

3.6

2.00 1.75

1.0

3.4 0.2

0.4

0.6

0.8

1.0

×10 2

100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

×10 3

Figure 11: Scaling Law of Manifold Radius ρ2M on Z127 . We show the scaling law of ρ2M with respect to the learning rate η, batch size b, and ℓ2 regularization coefficient λ on the Z127 task. For each hyperparameter configuration, we train for ten runs. The results show that larger η/b induces stronger diffusion variance, whereas λ does not affect the diffusion variance. We also overlay the theoretical fits.

Preprint

A DDITIONAL R ESULTS : S CALING L AWS OF M ANIFOLD R ADIUS ρ2G ON Z127

×103 5.8

×103 mean over seeds ±1 theoretical

×103 5.0

5.2 5.1 2 G

2 G

5.6

4.5

5.0

5.4

4.0

4.9 5.2

4.8

5.0

4.7

0.2

0.4

mean over seeds ±1 theoretical

5.3

2 G

A.16

0.6

0.8

1.0

×10 2

3.5

mean over seeds ±1 theoretical 100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

×10 3

Figure 12: Scaling Law of Manifold Radius ρ2G on Z127 . We show the scaling law of ρ2G with respect to the learning rate η, batch size b, and ℓ2 regularization coefficient λ on the Z127 task. For each hyperparameter configuration, we train for ten runs. The results show that larger η induces stronger diffusion variance, whereas λ does not affect the diffusion variance. We also overlay the theoretical fits.

Preprint

A DDITIONAL R ESULTS : S CALING L AWS OF S OLUTION T RANSITION T IME ON Z127

×103

×103 7.0

6.5

M G

×103

7.5

mean over seeds ±1 theoretical

6.0

1 0.2

0.4

0.6

0.8

1.0

×10 2

5.0

5 4

5.5

mean over seeds ±1 theoretical

M G

A.17

3 100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

×10 3

Figure 13: Scaling laws of solution transition time on Z127 . We show that the solution transition time τM →G from the memorization manifold M to the generalization manifold G scales with the learning rate η, batch size b, and ℓ2 regularization coefficient λ. For each hyperparameter configuration, we train for ten runs. We also overlay the theoretical fits.

Record · ID 321853 · SHA-256 cbd42249e54b1db4
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.