Paper deep dive

A Stochastic--Geometric Theory of Scaling Laws in Grokking

Róisín Luo, Christian Gagné, Jonas Ngnawé, Ihsan Ullah, Karyn Morrissey

Year: 2026Venue: arXiv preprintArea: stat.MLType: PreprintEmbeddings: 199

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 96%

Last extracted: 7/5/2026, 3:55:02 AM

Summary

The paper proposes a stochastic-geometric theory to explain 'grokking' (delayed generalization) in neural networks. The authors characterize the reachable solution space as a shell-core topological configuration induced by Adam optimization with weight-shrinkage regularization. In this configuration, random initialization solutions concentrate on an outer spherical shell, which encloses a shell of memorization solutions, which in turn contains a core of generalization solutions. Using stopping-time theory and stochastic differential equations (SDEs), the paper derives scaling laws for the memorization radius, generalization radius, and the transition time from memorization to generalization with respect to learning rate, batch size, and L2 regularization.

Entities (8)

Adam · optimizer · 100%Generalization Manifold · manifold · 100%Grokking · phenomenon · 100%L2 Regularization · regularization · 100%Memorization Manifold · manifold · 100%S5 · taskdataset · 100%Shell-Core Topology · topological-configuration · 100%Z127 · taskdataset · 100%

Relation Signals (4)

Shell-Core Topology → contains → Memorization Manifold

confidence 100% · enclosing another spherical shell of memorization solutions

Memorization Manifold → encloses → Generalization Manifold

confidence 100% · memorization solutions, which in turn contains a core corresponding to the generalization solutions

Adam → induces → Shell-Core Topology

confidence 95% · characterize a shell–core topological configuration of the reachable solution space induced by Adam's optimization dynamics

L2 Regularization → affects → Grokking

confidence 90% · derives grokking scaling laws for the learning rate, batch size, and ℓ2 regularization coefficient

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:Delayed generalization (\ie~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 $\ell_2$ regularization coefficient, which are further validated through experiments and shown to recover results from prior literature.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

198,295 characters extracted from source content.

Expand or collapse full text

A Stochastic–Geometric Theory of Scaling Laws in Grokking Róisín Luo Research Ireland – Centre for Research Training in AI J.E. Cairnes School of Business & Economics University of Galway &Christian Gagné Université Laval Canada–CIFAR AI Chair Mila – Québec AI Institute &Jonas Ngnawé Université Laval Mila – Qué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 Correspondence to: roisincrtai@gmail.com Abstract 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 _2 regularization coefficient, which are further validated through experiments and shown to recover results from prior literature. 1 Introduction 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 and 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. (a) Weight-Shrinkage Grokking (b) Trajectories w/ S5S_5 (c) Trajectories w/ ℤ127Z_127 Figure 1: Optimization-Induced Shell–Core Topology and Empirical Evidence. Left (1(a)): grokking dynamics under weight-shrinkage regularization induce a shell–core topology in parameter space, where initialization points θ0∼(0,σ2Ip) _0 (0,σ^2I_p) concentrate on a hyperspherical shell Θ with radius ρΘ≈σp _ ≈σ p and thickness σ/2σ/ 2, enclosing the memorization shell M∖GM G and the generalization core G; the radii ρG _G, ρM _M, and ρΘ _ are determined by the task and optimization configuration. Middle–right (1(b)–1(c)): empirical evidence on S5S_5 and ℤ127Z_127 using isometric manifold learning (e.g., MDS) to visualize ten optimization parameter trajectories, where θ0 _0, θτM _ _M, and θτG _ _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 ℝpR^p. Empirical findings consistently suggest that training dynamics in the grokking regime exhibit three-stage 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 _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 1(a). 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 1(b)–1(c) provide empirical evidence consistent with this optimization-induced topological configuration by visualizing optimization trajectories via isometric manifold learning. This optimizer-induced topological configuration, as illustrated in Figure 1(a), 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 and Ba, 2015) designed to analyze grokking dynamics. Building on this formulation, Section 4.2 derives the scaling law for the 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. 2 Notations Data and Network. Let Dtr:=(xi,yi)i=1nD_tr:=\(x_i,y_i)\_i=1^n denote a finite training set of size n, where each pair (xi,yi)∈×(x_i,y_i) ×Y is drawn i.i.d. from the empirical distribution P^ P_XY, with underlying true distribution PP_XY. Let fθ:→f_θ:X be a model parameterized by θ∈ℝpθ ^p. Throughout, we use XtX_t to denote the time-indexed value of a variable X at time t; for example, θt _t denotes the parameter at time t. Instance, Batch, Empirical, and True Loss. Let ℓf(x,y;θ) _f(x,y;θ) denote the instance loss evaluated at a point (x,y)(x,y). Let ℒf(ξ;θ) L_f(ξ;θ) denote the batch loss evaluated on a mini-batch ξ=(xξi,yξi)i=1bξ=\(x_ _i,y_ _i)\_i=1^b of size b. Let ℒ^f(θ):=ξ∼Dtr[ℒf(ξ;θ)] L_f(θ):=E_ξ D_tr\! [ L_f(ξ;θ) ] denote the empirical loss evaluated over the training set DtrD_tr. Let ℒf(θ) L_f(θ) denote the true loss with respect to the underlying data distribution PP_XY. Initialization, Memorization, and Generalization Manifolds. We study the solution manifolds reachable under the dynamics induced by Adam (Kingma and Ba, 2015). The ambient parameter space is ℝpR^p endowed with the Euclidean metric. Let Θ denote the set of initialization solutions under some initialization distribution PΘP_ . Let M:=θ∈ℝp∣ℒ^f(θ)<ϵ∧pAdam(θ∣θ0∈Θ)>0M:=\θ ^p L_f(θ)<ε\; \;p_Adam(θ _0∈ )>0\ denote the set of memorization solutions under Adam’s dynamics with a loss tolerance ϵε. Similarly, let G:=θ∈ℝp∣ℒf(θ)<ϵ∧pAdam(θ∣θ0∈Θ)>0G:=\θ ^p L_f(θ)<ε\; \;p_Adam(θ _0∈ )>0\ denote the set of generalization solutions. Throughout, ∂S∂ S and S∘S denote the boundary and interior of a set S, respectively; ∂M∂ M is the outer boundary of M∖GM G, and ∂G∂ G is its inner boundary. Let ρM2:=θ0∈Θ,θ∈∂M[‖θ‖22∣θ0] _M^2:=E_ _0∈ ,θ∈∂ M\! [\|θ\|_2^2 _0 ] and ρG2:=θ0∈Θ,θ∈∂G[‖θ‖22∣θ0] _G^2:=E_ _0∈ ,θ∈∂ G\! [\|θ\|_2^2 _0 ] denote the mean squared radii of the memorization and generalization manifolds, respectively. 3 Initialization Concentration We first demonstrate the topological concentration of network-parameter initialization under normal and uniform distributions. Consider that θ0=(θ0(1),⋯,θ0(k))∈ℝp _0=( _0^(1),·s, _0^(k)) ^p consists of k sub-vectors such that θ0(j)∈ℝpj _0^(j) ^p_j where ∑j=1kpj=p _j=1^kp_j=p. Define effective dimension p~ p such that 1p~:=1k∑j=1k1pj 1 p:= 1k _j=1^k 1p_j. For coordinate-wise initialization [θ0(j)]i∼(0,σj2)[ _0^(j)]_i (0, _j^2), the initialization radius square ρΘ2(θ0)=‖θ0‖22=∑j=1k(∑i=1pj|[θ0(j)]i|2)∼∑j=1kσj2χpj2, _ ^2( _0)=\| _0\|_2^2= _j=1^k ( _i=1^p_j\; |[ _0^(j)]_i |^2 ) _j=1^k _j^2 _p_j^2, (1) follows a sum of scaled chi-square distributions, where each σj2χpj2 _j^2 _p_j^2 has pjp_j degrees of freedom. Hence, the radius of Θ admits the closed-form expression σ~2:=1p∑j=1kpjσj2,θ0[ρΘ2(θ0)]=σ~2p=O(k),Varθ0[ρΘ2(θ0)]=2∑j=1kpjσj4=O(kp~), σ^2:= 1p _j=1^kp_j _j^2, _ _0 [ _ ^2( _0) ]\;=\; σ^2p=O(k), _ _0 [ _ ^2( _0) ]=2 _j=1^kp_j _j^4=O( k p), under a coordinate scaling σj=O(1/pj) _j=O(1/ p_j) (LeCun et al., 1998; Glorot and Bengio, 2010). This suggests that initialization solutions concentrate on a thin annulus of radius O(k)O(k) with thickness O(k/p~)O( k/ p) in parameter space. Similar concentration also holds for coordinate-wise uniform initialization distribution [θ0(j)]i∼(−εj,εj)[ _0^(j)]_i (- _j, _j), ε~2:=1p∑j=1kpjεj2,θ0[ρΘ2(θ0)]=pε~23=O(k),Varθ0[ρΘ2(θ0)]=445∑j=1kpjεj4=O(kp~), ^2:= 1p _j=1^kp_j _j^2, _ _0 [ _ ^2( _0) ]\;=\; p ^23=O(k), _ _0 [ _ ^2( _0) ]= 445 _j=1^kp_j _j^4=O( k p), under a coordinate scaling εj=O(1/pj) _j=O(1/ p_j). Proofs for the normal and uniform concentration results are provided in Appendices A.3 and A.4, respectively. 4 Manifold Radius and Solution Transition Analysis 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 _M, defined as the parameter norm at the first-hitting time of the empirical-loss boundary ∂M∂ M, via perturbation method (Bender and 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 _G, defined as the parameter norm at the first-hitting time of the generalization boundary ∂G∂ G, and the solution transition time [τM→G]E[ _M→ G], defined as the mean time required to transition from memorization boundary ∂M∂ M to generalization boundary ∂G∂ G. A high-level sketch for theoretical analysis framework is illustrated as in Figure 6 of the appendix. 4.1 Adam’s Joint-State Continuous-Time SDE Limit We consider Adam dynamics in the small learning rate and large batch size regime η→0,b→∞η→ 0,\,b→∞, with exponential moving average (EMA) coefficients (β1,β2)∈(0,1)2( _1, _2)∈(0,1)^2 and an ℓ2 _2 regularizer R(θ)=12‖θ‖22R(θ)= 12\|θ\|_2^2 with coefficient λ>0λ>0. Let gt:=∇[ℒf∗(ξt;θt)+λR(θt)]g_t:=∇[ L^*_f( _t; _t)+λ R( _t)] denote the regularized mini-sample gradient flow, with per-sample mean g¯t:=s[∇ℓf(s;θt)] g_t:=E_s[∇ _f(s; _t)] and covariance Σt:=Covs[∇ℓf(s;θt)] _t:=Cov_s[∇ _f(s; _t)]. For sufficiently large batch size b, the mini-batch gradient admits the distributional limit gt∼(g¯t,Σt/b)g_t ( g_t, _t/b ). We refer to this setting as coupled weight decay, since the ℓ2 _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 and Hutter, 2019). Adam’s Parameter Update Rules. Let θ∈ℝpθ ^p be model’s parameters. Let m,v∈ℝpm,v ^p be the first- and second-moment estimates, respectively. Let m^,v m, v be the m,vm,v with bias corrections. Adam’s discrete update rules at iteration k are given as mk+1=β1mk+(1−β1)gk,vk+1=β2vk+(1−β2)(gk⊙gk), m_k+1= _1\,m_k+(1- _1)\,g_k, v_k+1= _2\,v_k+(1- _2)\,(g_k g_k), m^k+1=mk+1/(1−β1k+1),v^k+1=vk+1/(1−β2k+1), m_k+1=m_k+1\,/\,(1- _1^k+1), v_k+1=v_k+1\,/\,(1- _2^k+1), θk+1=θk−ηm^k+1⊘(v^k+1+ε 1p), _k+1= _k-η\, m_k+1 ( v_k+1+ \,1_p ), with initial values θ0∈Θ,m0=0 _0∈ ,m_0=0, v0=0v_0=0, where ⊙ is element-wise product and ⊘ is element-wise division. m m and v v are Adam’s bias corrections, which compensate the warm-up underestimation of (mk,vk)(m_k,v_k) caused by the zero initialization m0=v0=0m_0=v_0=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)∈ℝ3pS_t:=( _t,m_t,v_t) ^3p and using the continuous-time interpolation t=ηkt=η k, we state the resulting joint-state SDE for Adam in Lemma 1. The proof sketch is illustrated in Figure 7. The proof is provided in Appendix A.5. The technical correctness is verified through an induced radius SDE with Itô’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⊤,mt⊤,vt⊤)⊤∈ℝ3pS_t:=( _t ,m_t ,v_t ) ^3p denote the joint state of Adam, where θt _t, mtm_t, and vtv_t are the parameters, first-moment estimates, and second-moment estimates, respectively. Taking the continuous-time interpolation t=ηkt=η k, the discrete Adam updates admit the following Itô SDE limit on the joint state dSt=μ(St)dt+ηbσ(St)dWt, S_t\;=\;μ(S_t)\, t\;+\; ηb\,σ(S_t)\, W_t, (2) where dWt∈ℝp W_t ^p denotes the infinitesimal increment of a Wiener process adapted to the filtration generated by the mini-batch sampling process ξt\ _t\, with drift factor μ(St)∈ℝ3pμ(S_t) ^3p and diffusion factor σ(St)∈ℝ3p×pσ(S_t) ^3p× p μ(St)=(−B(t)mt⊘vt−α1(mt−g¯t)−α2(vt−g⊙g¯t)),σ(St)=(0α1Σt1/2α2Dt),B(t):= 1−e−α2t 1−e−α1t, μ(S_t)\;=\; pmatrix-\,B(t)\,m_t v_t\\[2.0pt] -\, _1\, (m_t- g_t )\\[2.0pt] -\, _2\, (v_t- g g_t ) pmatrix, σ(S_t)\;=\; pmatrix0\\[2.0pt] _1\, _t^1/2\\[2.0pt] _2\,D_t pmatrix, B(t)\;:=\; \,1-e^- _2t\,\,1-e^- _1t\,, (3) where αi:=(1−βi)/η _i:=(1- _i)/η and B(t)B(t) is the bias-correction factor in continuous time and g⊙g¯t:=g¯t⊙g¯t+1bdiag(Σt)∈ℝp,Dt:= 2diag(g¯t)Σt1/2∈ℝp×p. g g_t\;:=\; g_t g_t\;+\; 1b\,diag( _t) ^p, D_t\;:=\;2\,diag ( g_t )\, _t^1/2 ^p× p. (4) In particular, to decompose this SDE into a tractable form, we define π(θt):=diag(g⊙g¯t)−1/2∈ℝp×pπ( _t):=diag\! ( g g_t )^-1/2 ^p× p as Adam’s preconditioner. We write θt(θ0),θt(S0) _t( _0), _t(S_0) to denote the Adam’s θ-evolution starting from initial states θ0∈Θ,S0∈ΘS _0∈ ,S_0∈ _S, respectively, where ΘS _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 continuous-time treatment of Adam that explicitly accounts for stochasticity remains, to the best of our knowledge, comparatively less developed. For example, the joint (m,v)(m,v) deterministic limit of Adam appears in (Da Silva and 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 Memorization Radius Starting from the initial state S0=(θ0⊤,0p⊤,0p⊤)⊤S_0=( _0 ,0_p ,0_p ) , Adam’s joint-state SDE (2) generates a stochastic trajectory St\S_t\ adapted to the filtration generated by mini-batch sampling. The first time at which θt _t reaches the memorization boundary ∂M∂ M defines the stopping time (Karatzas and Shreve, 1991) and mean squared memorization radius as τM(S0):=inft≥0:θt(S0)∈∂M,andρM2:=S0∈ΘS[‖θτM(S0)‖22], _M(S_0):= \t≥ 0: _t(S_0)∈∂ M\, _M^2:=E_S_0∈ _S [\| _ _M(S_0)\|_2^2 ], (5) respectively. This defines a Dirichlet exit problem on the state StS_t, with absorbing boundary ∂M∂ M and the infinitesimal generator (i.e., a differential operator) (Dynkin, 1965) induced by Adam’s SDE in Lemma 1 as: ℒS[∙]:=μ(S)⊤∇S[∙]+12ηbtr(ΣS(S)∇S2[∙]), _S[ ]\;:=\;μ(S)^\! _S[ ]\;+\; 12\, ηb\,tr\! ( _S(S)\, _S^2\,[ ]\, ), (6) which contains an operator modulated by a coefficient η/bη/b with ΣS(S):=σ(S)σ(S)⊤ _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 and Orszag, 1999) in ε=η/b =η/b to obtain Theorem 1. The proof is provided in Appendix A.12. Experimental verification is provided for S5S_5 in Figure 2 and for ℤ127Z_127 in Appendix A.15. Figure 2: Scaling Laws of Manifold Radius ρM2 _M^2 on S5S_5. We show the scaling law of ρM2 _M^2 with respect to the learning rate η, batch size b, and ℓ2 _2 regularization coefficient λ on the S5S_5 task. For each hyperparameter configuration, we train for ten runs. The results show that larger η/bη/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 ρM2 _M^2 with respect to the learning rate η exhibits a U-shaped curve. An additional experimental results for ℤ127Z_127 are provided 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η/b as ρM2=(ρM(0))2+ηbcM+(ηb)2cM(2)+O((ηb)3), _M^2\;=\;( _M^(0))^2\;+\; ηb\,c_M\;+\; ( ηb )^\!2\!c_M^(2)\;+\;O\! (\! ( ηb )^\!3 ), (7) where (ρM(0))2( _M^(0))^2, cMc_M, cM(2)c_M^(2) are task-determined constants independent of η and b, as defined in Appendix A.12. The quadratic regularizer λ2‖θt‖22 λ2\| _t\|_2^2 induces deterministic contraction. The leading constant (ρM(0))2( _M^(0))^2 depends on the deterministic gradient flow ∇ℒf(θt)∇ L_f( _t), when the task gradient ∇ℒf∗(θt)∇ L_f^*( _t) is small, λ enters (ρM(0))2( _M^(0))^2 through d‖θt‖22=2θt⊤dθt≈−2λ‖θt‖22dt⇒1(ρM(0))2d(ρM(0))2=−2λdt \| _t\|_2^2=2 _t _t≈-2λ\| _t\|_2^2\, t 1( _M^(0))^2 ( _M^(0))^2=-2λ\, t. Thus, (ρM(0))2∝exp⁡(−O(λ))( _M^(0))^2 (-O(λ)), and hence ρM2∝exp⁡(−O(λ)) _M^2 (-O(λ)). Remark 2. In the radius-shrinkage regime, diffusion first shifts the hitting point inward, giving cM<0c_M<0, while boundary curvature produces a stabilizing second-order correction, giving cM(2)>0c_M^(2)>0. Our theory predicts that ρM2 _M^2 is U-shaped in η/bη/b with two regimes: a sub-linear decrease ρM2≈(ρM(0))2+(η/b)cM _M^2≈( _M^(0))^2+(η/b)c_M for η/b≪|cM|/(2cM(2))η/b |c_M|/(2c_M^(2)), and a sub-quadratic increase ρM2≈(ρM(0))2+(η/b)2cM(2) _M^2≈( _M^(0))^2+(η/b)^2c_M^(2) for η/b≳|cM|/(2cM(2))η/b |c_M|/(2c_M^(2)), with minimum at (η/b)⋆=|cM|/(2cM(2))(η/b) =|c_M|/(2c_M^(2)). 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 Late-Stage Radius Dynamics To characterize the manifold geometry, we derive the radius SDE induced by Adam’s SDE (2) via Itô’s lemma under the late-stage reduction observations. Empirically, as shown in Figure 3, once the trajectory St\S_t\ reaches the memorization manifold M (i.e., memorization regime), the late-stage grokking dynamics exhibit two reduction properties. First, refer to equation (3), the first- and second-moment estimates mtm_t and vtv_t rapidly converge to g¯t(1−e−α1t) g_t (1-e^- _1t ) and g⊙g¯t g g_t, respectively, at an exponential rate exp⁡(−O(t)) (-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 _t, become small and isotropic at an exponential rate exp⁡(−O(t)) (-O(t)). We refer to this property as small-isotropic gradient covariance, formalized in Observation 2. (a) Reduction on (mt,vt)(m_t,v_t) (b) Gradient noise magnitude (c) Gradient noise isotropy Figure 3: Observations in Grokking Dynamics. The markers τM _M and τG _G denote the first-hitting times of ∂M∂ M and ∂G∂ G, respectively, and the y-axis is shown on a logarithmic scale. Left (3(a)) shows the exponential reduction of moment estimates; (mt,vt)(m_t,v_t) rapidly converges to (g¯t(1−e−α1t),g⊙g¯t(1−e−α2t)) ( g_t(1-e^- _1t), g g_t(1-e^- _2t) ), with exponential decay exp⁡(−O(t)) (-O(t)). Middle–right (3(b)–3(c)) shows the exponential reduction of the gradient noise; the mean and variance of the singular values of the gradient covariance Σt _t decay as exp⁡(−O(t)) (-O(t)), indicating small and isotropic gradient fluctuations in late-stage dynamics. Observation 1 (Slow-Manifold Reduction on (mt,vt)(m_t,v_t)). If (mt,vt)(m_t,v_t) varies slowly as t→∞t→∞, the first- and second-moment estimates mt,vtm_t,v_t rapidly converge to its mean-field limit mt→g¯t(1−e−α1t)→g¯tandvt→g⊙g¯t(1−e−α2t)→g⊙g¯t, m_t\;→\; g_t\, (1-e^- _1t )→ g_t v_t\;→\; g g_t(1-e^- _2t)→ g g_t, (8) respectively, at an exponential rate exp⁡(−O(t)) (-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 _t rapidly becomes small and approximately isotropic as t→∞t→∞. With this isotropy and by equation (4), the Adam’s preconditioner is therefore approximated by π(θ)≈s(θ)−1Ipπ(θ)≈ s(θ)^-1I_p, as t→∞t→∞, where the state-dependent scalar s(θ)>0s(θ)>0 is defined by 1s(θ):=1ptr(π(θ)),π(θ):=diag(g⊙g¯(θ))−1/2, 1s(θ):= 1ptr (π(θ) ), π(θ):=diag\! ( g g(θ) )^-1/2, (9) where p is the parameter dimension. In particular, let G(θ):=1bB(t)π(θ)Σ(θ)1/2G(θ):= 1 bB(t)π(θ) (θ)^1/2 be the preconditioned diffusion factor, which admits G(θ)G(θ)⊤≈IpG(θ)G(θ) ≈ I_p. See Appendix A.10. Under Observations 1–2, Lemma 2 gives the approximate closed-form dynamics of rt2=‖θt‖22r_t^2=\| _t\|_2^2. 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‖22r_t^2:=\| _t\|_2^2. In Adam’s late-stage regime underlying grokking, assume that the diffusion covariance becomes isotropic G(θt)G(θt)⊤≈IpG( _t)G( _t) ≈ I_p as t→∞t→∞, then the reduced late-stage squared-radius SDE admits drt2≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ηp+ℛSM(t)+ℛπ(t)]dt+2ηrtdWt(r), r_t^2≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+η\,p+R_SM(t)+R_π(t) ] t+2 η\,r_t\, W_t^(r), (10) where slow-manifold residual ℛSM(t)R_SM(t) and preconditioner residual ℛπ(t)R_π(t) are ℛSM(t) _SM(t) :=−2B(t)θt⊤[[mt⊘vt]−g¯t(1−e−α1t)⊘g⊙g¯t], :=-2B(t) _t [E [m_t v_t ]- g_t (1-e^- _1t ) g g_t ], (11) ℛπ(t) _π(t) :=−2θt⊤[π(θt)−s(θt)−1Ip]g¯t, :=-2 _t [π( _t)-s( _t)^-1I_p ] g_t, (12) respectively, and a one-dimensional Wiener process Wt(r):=∫0ter(θu)⊤G(θu)dWuW_t^(r):= _0^te_r( _u) G( _u) W_u with er(θ):=θ‖θ‖2e_r(θ):= θ\|θ\|_2 by Lévy’s characterization: ⟨Wt(r)⟩t=t W_t^(r) _t=t. Empirically, in the memorization regime, the residual sum ℛSM(t)+ℛπ(t)R_SM(t)+R_π(t) becomes negligible. Hence, we set ℛSM(t)+ℛπ(t)=0R_SM(t)+R_π(t)=0 in the subsequent analysis, see Appendix A.11. 4.4 Generalization Radius Figure 4: Scaling Laws of Manifold Radius ρG2 _G^2 on S5S_5. We show the scaling laws of ρG2 _G^2 with respect to the learning rate η, batch size b, and ℓ2 _2 regularization coefficient λ on the S5S_5 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 ℤ127Z_127 are provided in Appendix A.16. Let SτMS_ _M denote the state at which the optimization trajectory first reaches ∂M∂ M. We treat SτMS_ _M as the initial state for the subsequent first-hitting problem (Karatzas and Shreve, 1991) on ∂G∂ G. Define the stopping time for first hitting ∂G∂ G and the corresponding mean squared generalization radius as τG(S0):=inft≥τM:θt(S0)∈∂G,ρG2:=S0∈ΘS[‖θτG(S0)‖22]. _G(S_0):= \t≥ _M: _t(S_0)∈∂ G\, _G^2:=E_S_0∈ _S [\| _ _G(S_0)\|_2^2 ]. (13) Once the trajectory enters ∂G∂ G, the drift and diffusion in equation (2) equilibrate, and Adam’s SDE approaches its equilibrium. The center of this equilibrium starting from S0S_0 is the regularized local minimizer θ⋆(S0):=arg⁡minθ(S0)⁡ℒf∗(θ(S0))+λ2‖θ(S0)‖22, θ (S_0):= _θ(S_0) L^*_f(θ(S_0))+ λ2\|θ(S_0)\|_2^2, (14) where ℒf∗(θ(S0)) L^*_f(θ(S_0)) is task loss. This yields an equilibrium-distribution problem for the reduced late-stage dynamics around θ⋆(S0)θ (S_0), initialized from the post-memorization state SτM(S0)S_ _M(S_0). Linearizing the reduced dynamics around θ⋆(S0)θ (S_0) solves Theorem 2. The proof is provided in Appendix A.13. Experimental verification is provided for S5S_5 in Figure 4 and for ℤ127Z_127 in Appendix A.16. Theorem 2 (Scaling Law of Generalization Radius). The mean squared generalization manifold radius admits the asymptotic expansion ρG2≈(ρG(0))2+ηλcG+O(η2λ2), _G^2\;≈\; ( _G^(0) )^2\;+\; ηλ\,c_G\;+\;O\! ( η^2λ^2 ), (15) with (ρG(0))2:=S0∈ΘS[‖θ⋆(S0)‖22]( _G^(0))^2:=E_S_0∈ _S [\|θ (S_0)\|_2^2 ] a landscape constant and cG>0c_G>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)c_G=O( 1 b). 4.5 Solution Transition Time Starting from S0∈ΘS_0∈ _S, the transition time on the trivial annulus M∖GM G is given as τM→G(S0):=τG(S0)−τM(S0), _M→ G(S_0)\;:=\; _G(S_0)- _M(S_0), (16) where τM(S0) _M(S_0) and τG(S0) _G(S_0) are the stopping times that the trajectory first hits ∂M∂M and ∂G∂G, respectively, defined in Sections 4.2–4.4. Taking θτM(S0)∈∂M _ _M(S_0)∈∂ M as the initial condition (Theorem 1), we formulate the transition from ∂M∂ M to ∂G∂ G as a mean first-passage problem (Karatzas and Shreve, 1991) for the radial process rt:=‖θt‖2r_t:=\| _t\|_2 on the annulus M∖GM 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 S5S_5 in Figure 5 and for ℤ127Z_127 in Appendix A.17. Theorem 3 (Scaling Law of Solution Transition Time). The mean number of optimizer iterations between the first hit at ∂M∂ M and the first hit at ∂G∂ G admits the asymptotic expansion τM→G:=S0∈ΘS[τM→G(S0)]≈s¯ηλlog⁡ρM(0)ρG(0)+cτbλ+cτ(2)λ2+O(ηλ3), _M→ G:=E_S_0∈ _S [ _M→ G(S_0) ]\;≈\; sη\,λ\, \! _M^(0) _G^(0)\;+\; c_τb\,λ\;+\; c_τ^(2)λ^2\;+\;O\! ( ηλ^3 ), (17) where s¯ s is determined by the harmonic mean of the effective deterministic gradient flow, while cτc_τ and cτ(2)c_τ^(2) collect the trajectory corrections, as defined in Appendix A.14. All three are task-determined constants independent of η and λ; a weak batch-size b-dependence enters through the scalar preconditioning scale s¯=O(1b) s=O( 1 b), and hence cτ=O(1b)c_τ=O( 1 b) and cτ(2)=O(1b)c_τ^(2)=O( 1b). Figure 5: Scaling laws of solution transition time on S5S_5. We show that the solution transition time τM→G _M→ G from the memorization manifold M to the generalization manifold G scales with the learning rate η, batch size b, and ℓ2 _2 regularization coefficient λ. For each hyperparameter configuration, we train for ten runs. We also overlay the theoretical fits. An additional experimental results for ℤ127Z_127 are provided in Figure 13. 5 Validation We validate the three scaling laws in Section 5.1 by directly measuring ρM2 _M^2, ρG2 _G^2, and τM→G _M→ G on the group-theoretic learning task S5S_5 and modular arithmetic tasks over ℤ127Z_127, and further compare them with published results from the grokking literature in Section 5.2. 5.1 Experimental Validation Experimental Setting. We use a two-layer MLP: each input token is embedded into 256256 dimensions, the two embeddings are concatenated, passed through a width-512512 ReLU hidden layer, and projected to |||Y| output logits, where ||=120|Y|=120 for S5S_5 and 127127 for ℤ127Z_127. The embedding layers are initialized from a standard normal distribution, while the linear layers use Kaiming initialization. Training is performed in 3232-bit floating-point precision. We find that 1616-bit precision can introduce numerical instability in grokking, suggesting that grokking dynamics are sensitive to numerical precision errors. We use base hyperparameters η=10−3η=10^-3, (β1,β2)=(0.9,0.999)( _1, _2)=(0.9,0.999), ℓ2 _2 regularization coefficient λ=10−4λ=10^-4, and batch size 100100. 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 SnS_n and a modular arithmetic learning task on ℤpZ_p. For SnS_n, each input is an ordered pair of permutations (f,g)∈Sn2(f,g)∈ S_n^2, and the model is trained to predict their group product f∘gf g under permutation composition. For ℤpZ_p, each input is a pair (a,b)∈ℤp2(a,b) _p^2, and the model is trained to predict the modular sum a+bmodpa+b 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 ρM2 _M^2 and ρG2 _G^2 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 _M→ G with respect to the same hyperparameters, as stated in Theorem 3. Additional results on ℤ127Z_127 are provided in Appendices A.15, A.16, and A.17. Overall, the experiments support the theoretical predictions; in particular, ρM2 _M^2 exhibits the predicted U-shaped dependence on η/bη/b. Discussions and Limitations. The scaling laws originate from how η, b, and λ modulate Adam’s SDE: η and b module the stochastic diffusion by scaling the gradient covariance with a factor η/b η/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η/b regime. 5.2 Literature Validation Theorem 1 characterizes the scaling of the memorization radius ρM _M with respect to η/bη/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 _G, showing that it decomposes into the landscape-dependent term (ρG(0))2=θ0∈Θ[‖θ⋆(θ0)‖22] ( _G^(0) )^2=E_ _0∈ [\|θ ( _0)\|_2^2 ] and an O(η/(bλ))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). Theorem 3 recovers the inverse dependence of the grokking delay on the regularization coefficient, including the scaling τ∝1/(ηλ)τ 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. 6 Conclusions 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 É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 Québec (calculquebec.ca) and the Digital Research Alliance of Canada (alliancecan.ca). Partial computational support was also provided by Taighde É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 Université Laval and Mila – Québec AI Institute. For the purpose of open access, the author has applied a C BY public copyright licence to any Author Accepted Manuscript version arising from this submission. References C. M. Bender and S. A. Orszag (1999) Advanced mathematical methods for scientists and engineers i: asymptotic methods and perturbation theory. 1 edition, Springer New York, New York, NY. Note: Originally published by McGraw Hill, 1978 External Links: Document, ISBN 978-0-387-98931-0 Cited by: §4.2, §4. E. M. Compagnoni, T. Liu, R. Islamov, F. N. Proske, A. Orvieto, and A. Lucchi (2025) Adaptive methods through the lens of SDEs: theoretical insights on the role of noise. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: Remark 1. A. B. Da Silva and M. Gazeau (2020) A general system of differential equations to model first-order adaptive algorithms. J. Mach. Learn. Res. 21 (1). External Links: ISSN 1532-4435 Cited by: Remark 1. E. B. Dynkin (1965) Markov processes: volume i. 1st edition, Grundlehren der mathematischen Wissenschaften, Springer Berlin, Heidelberg, Berlin, Heidelberg. External Links: Document, ISBN 978-3-662-23320-7 Cited by: §1, §4.2. X. Glorot and Y. Bengio (2010) Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, Y. W. Teh and M. Titterington (Eds.), Proceedings of Machine Learning Research, Vol. 9, Chia Laguna Resort, Sardinia, Italy, p. 249–256. External Links: Link Cited by: §3. H. Hwang and Y. Park (2026) Intrinsic task symmetry drives generalization in algorithmic tasks. In Proceedings of the Forty-Third International Conference on Machine Learning, External Links: 2603.01968, Link Cited by: §1. I. Karatzas and S. E. Shreve (1991) Brownian motion and stochastic calculus. 2nd edition, Graduate Texts in Mathematics, Springer New York, New York, NY. External Links: Document, ISBN 978-0-387-97655-6 Cited by: §A.11, §4.2, §4.4, §4.5. D. P. Kingma and J. Ba (2015) Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, Y. Bengio and Y. LeCun (Eds.), External Links: Link Cited by: §1, §2. T. Kumar, B. Bordelon, S. J. Gershman, and C. Pehlevan (2024) Grokking as the transition from lazy to rich training dynamics. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: §1. Y. LeCun, L. Bottou, G. B. Orr, and K. Müller (1998) Efficient backprop. In Neural Networks: Tricks of the Trade, This Book is an Outgrowth of a 1996 NIPS Workshop, Berlin, Heidelberg, p. 9–50. External Links: ISBN 3540653112 Cited by: §3. Q. Li, C. Tai, and W. E (2017) Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning, D. Precup and Y. W. Teh (Eds.), Proceedings of Machine Learning Research, Vol. 70, p. 2101–2110. External Links: Link Cited by: Remark 1. Z. Liu, O. Kitouni, N. Nolte, E. J. Michaud, M. Tegmark, and M. Williams (2022) 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. External Links: ISBN 9781713871088 Cited by: §1, §5.2. I. Loshchilov and F. Hutter (2019) Decoupled weight decay regularization. In International Conference on Learning Representations, External Links: Link Cited by: §4.1. R. Luo, J. McDermott, C. Gagné, Q. Sun, and C. O’Riordan (2025) Optimization-induced dynamics of Lipschitz continuity in neural networks. External Links: 2506.18588, Link Cited by: Remark 1. S. Malladi, K. Lyu, A. Panigrahi, and S. Arora (2022) 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. External Links: ISBN 9781713871088 Cited by: Remark 1. S. Mandt, M. D. Hoffman, and D. M. Blei (2017) Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research 18 (134), p. 1–35. External Links: Link Cited by: Remark 1. N. Nanda, L. Chan, T. Lieberum, J. Smith, and J. Steinhardt (2023) Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §1, §1. P. J. T. Notsawo, G. Dumas, and G. Rabusseau (2026) Grokking finite-dimensional algebra. In Proceedings of the Forty-Third International Conference on Machine Learning, External Links: 2602.19533, Link Cited by: §1. B. Øksendal (2003) Stochastic differential equations: an introduction with applications. 6 edition, Universitext, Springer Berlin, Heidelberg, Berlin, Heidelberg. External Links: Document, ISBN 978-3-540-04758-2 Cited by: §A.12, §A.6, §1. A. Power, Y. Burda, H. Edwards, I. Babuschkin, and V. Misra (2022) Grokking: generalization beyond overfitting on small algorithmic datasets. External Links: 2201.02177, Link Cited by: §1, §1. D. Stander, Q. Yu, H. Fan, and S. Biderman (2024) Grokking group multiplication with cosets. In Proceedings of the 41st International Conference on Machine Learning, ICML’24. Cited by: §1. V. Varma, R. Shah, Z. Kenton, J. Kramár, and R. Kumar (2023) Explaining grokking through circuit efficiency. External Links: 2309.02390, Link Cited by: §5.2. Z. Zhong, Z. Liu, M. Tegmark, and J. Andreas (2023) The clock and the pizza: two stories in mechanistic explanation of neural networks. In Thirty-seventh Conference on Neural Information Processing Systems, External Links: Link Cited by: §1. Appendix A Appendix 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. Adam’s Joint-State SDE (Lemma 5) dSt=μ(St)dt+η/bσ(St)dWt,St=(θt,mt,vt) S_t=μ(S_t) t+ η/b\,σ(S_t) W_t, S_t=( _t,m_t,v_t) Induced Radius SDE (Lemma 6) drt2=−2B(t)θt⊤(mt⊘vt)dt r_t^2=-2B(t) _t (m_t v_t ) t Memorization-Regime Identities (Lemma 10) Late-stage identities tr(πt)≈btr[(diagΣ(θt))−1/2], ( _t)≈ b\,tr\! [ (diag\, ( _t) )^-1/2 ], tr(G(θt)G(θt)⊤)≈p \! (G( _t)G( _t) )≈ p Preconditioned-Decomposition of Adam’s θ-SDE (Lemma 8) dθt _t ≈−a(t)π(θt)g¯(θt)dt+ℛθ,SM(t)dt ≈-a(t)π( _t) g( _t) t+R_θ,SM(t) t +ηG(θt)dWt + η\,G( _t) W_t Scaling Law of Memorization Radius (Theorem 4) Joint-state first-exit analysis for ρM2 _M^2 Mean-Field Limit of (mt,vt)(m_t,v_t) (Lemma 7) mt→g¯t(1−e−α1t),vt→g⊙g¯t m_t→ g_t(1-e^- _1t),\ v_t→ g g_t Preconditioned-Decomposition of Adam’s Radius SDE (Lemma 9) Deriving the preconditioned-decomposition of radius SDE for analyzing manifold radius precisely. drt2≈[−2a(t)θt⊤π(θt)g¯t+ℛSM(t)+ηtr(G(θt)G(θt)⊤)]dt+2ηθt⊤G(θt)dWt r_t^2≈ [-2a(t) _t π( _t) g_t+R_SM(t)+η\,tr\! (G( _t)G( _t) ) ] t+2 η\, _t G( _t) W_t Reduced Memorization-Regime Radius SDE (Lemma 11) Use memorization-regime identities to reduce the preconditioned radius dynamics drt2≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ηp+ℛSM(t)+ℛπ(t)]dt+2ηrtdWt(r) r_t^2≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+η\,p+R_SM(t)+R_π(t) ] t+2 η\,r_t\, W_t^(r) Scaling Law of Generalization Radius (Theorem 5) Stationary late-stage radial analysis for ρG2 _G^2 Scaling Law of Solution Transition Time (Theorem 6) First-passage analysis from ∂M ∂ M to ∂G ∂ G Figure 6: High-Level Sketch of Theoretical Analysis Framework. Table of Appendix Contents: A.1 Experimental Settings: Learning Tasks Group-theoretic learning task on SnS_n. Let SnS_n denote the symmetric group on n elements. Each group element f∈Snf∈ S_n is a permutation sending its index set [n]:=1,2,…,n[n]:=\1,2,…,n\ to a fixed enumeration of all n!n! permutations. Let f[i]f[i] denote the indexed permutation i↦ji j from an index i∈[n]i∈[n] to a target j∈[n]j∈[n]. For example, given two permutations f,g∈S5f,g∈ S_5: f=(12345↓25143)g=(12345↓31524), f= pmatrix1&2&3&4&5\\ & & & & \\ 2&5&1&4&3 pmatrix g= pmatrix1&2&3&4&5\\ & & & & \\ 3&1&5&2&4 pmatrix, the permutations are given by f[1]:1↦2f[1]:1 2, f[2]:2↦5f[2]:2 5, etc. The product of f and g is defined by permutation composition (f∘g)[i]=g[f[i]](f g)[i]=g [f[i] ] given by f∘g=(12345↓14325). f g= pmatrix1&2&3&4&5\\ & & & & \\ 1&4&3&2&5 pmatrix. The learning task is to predict the group product f∘gf g given the pair (f,g)(f,g). The dataset consists of all ordered pairs (f,g)∈Sn2(f,g)∈ S_n^2, with a randomly sampled subset used for training and the remaining pairs used for evaluation. Modular arithmetic learning task on ℤpZ_p. Let ℤpZ_p denote the cyclic group of integers modulo p. Each input consists of a pair (a,b)∈ℤp2(a,b) _p^2, and the task is to predict their modular sum a+bmodpa+b p. As in the group-theoretic task, training is performed on a randomly selected subset of all p2p^2 input–output pairs, while evaluation is conducted on the full set. A.2 Additional Results: Adam–Induced Shell–Core Radius & Stopping-Time Concentration Table 1: Adam–Induced Shell–Core Radius & Stopping-Time Concentration on Learning Task S5S_5. The experiment shows the concentration of Adam–induced shell–core radii and stopping times on the learning task S5S_5, across 10 seeds, with learning rate η=10−3,(β1,β2)=(0.9,0.999)η=10^-3,( _1, _2)=(0.9,0.999), ℓ2 _2 regularization coefficient 10−410^-4, and batch size 100100. Seed (S0S_0) ρ0(S0) _0(S_0) ρM(S0) _M(S_0) τM(S0) _M(S_0) ρG(S0) _G(S_0) τG(S0) _G(S_0) τM→G(S0) _M→ G(S_0) 1 249.27 186.07 1006 71.39 14000 12994 2 247.09 184.59 1009 69.22 14850 13841 3 247.87 185.19 1007 68.84 12150 11143 4 248.13 185.16 1009 69.67 13950 12941 5 248.85 185.64 1008 70.72 12500 11492 6 248.44 186.61 939 69.34 12100 11161 7 247.98 185.05 1009 69.50 15450 14441 8 248.03 185.19 1009 68.33 13850 12841 9 247.95 186.23 937 67.80 13050 12113 10 248.88 185.59 1004 70.15 12750 11746 [⋅]E[·] 248.25 185.53 993.7 69.50 13465.0 12471.3 Std[⋅]Std[·] 0.60 0.59 27.9 1.02 1083.3 1071.5 Table 2: Adam–Induced Shell–Core Radius & Stopping-Time Concentration on Learning Task ℤ127Z_127. The experiment shows the concentration of Adam–induced shell–core radii and stopping times on the learning task S5S_5, across 10 seeds, with learning rate η=10−3,(β1,β2)=(0.9,0.999)η=10^-3,( _1, _2)=(0.9,0.999), ℓ2 _2 regularization coefficient 10−410^-4, and batch size 100100. Seed (S0S_0) ρ0(S0) _0(S_0) ρM(S0) _M(S_0) τM(S0) _M(S_0) ρG(S0) _G(S_0) τG(S0) _G(S_0) τM→G(S0) _M→ G(S_0) 1 256.25 186.57 1041 71.47 8000 6959 2 254.24 183.95 1115 72.40 8600 7485 3 254.76 184.10 1121 71.35 8300 7179 4 255.47 184.79 1118 71.32 8300 7182 5 256.10 185.09 1119 70.61 8400 7281 6 255.76 184.74 1121 71.06 8450 7329 7 255.09 184.36 1116 72.32 8400 7284 8 255.45 184.88 1114 72.21 8500 7386 9 255.46 185.97 1047 72.66 7650 6603 10 256.06 186.53 1042 73.44 7800 6758 [⋅]E[·] 255.46 185.10 1095.4 71.88 8240.0 7144.6 Std[⋅]Std[·] 0.60 0.90 34.2 0.82 299.8 269.4 A.3 Proof: Concentration of Normal Initialization Lemma 3 (Gaussian Initialization Concentration). Let θ0=(θ0(1),…,θ0(k))∈ℝp,θ0(j)∈ℝpj,∑j=1kpj=p, _0= ( _0^(1),…, _0^(k) ) ^p, _0^(j) ^p_j, _j=1^kp_j=p, (18) where the subvectors are independent and [θ0(j)]i∼(0,σj2)[ _0^(j)]_i (0, _j^2). Define the effective coordinate variance σ~2 σ^2 and the effective dimension parameter p~ p by σ~2:=1p∑j=1kpjσj2,1p~:=1k∑j=1k1pj. σ^2:= 1p _j=1^kp_j _j^2, 1 p:= 1k _j=1^k 1p_j. (19) Then the squared initialization radius satisfies [‖θ0‖22]=σ~2p,Var[‖θ0‖22]=2∑j=1kpjσj4. [\| _0\|_2^2 ]= σ^2p, [\| _0\|_2^2 ]=2 _j=1^kp_j _j^4. (20) In particular, for initialization schemes used for numerical stability, each block-wise coordinate scale is often chosen so that σj=O(pj−1/2) _j=O(p_j^-1/2). Under this scaling, σ~2=O(kp),[‖θ0‖22]=O(k),Var[‖θ0‖22]=O(kp~). σ^2=O\! ( kp ), [\| _0\|_2^2 ]=O(k), [\| _0\|_2^2 ]=O\! ( k p ). (21) Proof. We consider a Gaussian initialization with block-wise coordinate scales. Write θ0=(θ0(1),…,θ0(k)),θ0(j)=σjZ(j),Z(j)∼(0,Ipj), _0= ( _0^(1),…, _0^(k) ), _0^(j)= _jZ^(j), Z^(j) (0,I_p_j), (22) where the subvectors Z(j)Z^(j) are independent. Then ‖θ0‖22=∑j=1k‖θ0(j)‖22=∑j=1kσj2‖Z(j)‖22, \| _0\|_2^2= _j=1^k\| _0^(j)\|_2^2= _j=1^k _j^2\|Z^(j)\|_2^2, (23) and it suffices to study the distribution of each ‖Z(j)‖22\|Z^(j)\|_2^2. Let Rj:=‖Z(j)‖22. R_j:=\|Z^(j)\|_2^2. (24) The random variable RjR_j follows a chi-square distribution with pjp_j degrees of freedom, whose density is fRj(r)=12pj2Γ(pj2)rpj2−1e−r/2,r≥0. f_R_j(r)= 12 p_j2 \! ( p_j2 )r p_j2-1e^-r/2, r≥ 0. (25) Computing [‖θ0‖22]E [\| _0\|_2^2 ]. The expected squared Euclidean norm of the initialization is therefore [‖θ0‖22] [\| _0\|_2^2 ] =∑j=1kσj2[Rj] = _j=1^k _j^2\,E[R_j] (26) =∑j=1kσj2∫0∞rfRj(r)dr = _j=1^k _j^2 _0^∞rf_R_j(r)\, r (27) =∑j=1kσj22pj2Γ(pj2)∫0∞rpj2e−r/2dr. = _j=1^k _j^22 p_j2 \! ( p_j2 ) _0^∞r p_j2e^-r/2\, r. (28) Applying the change of variables u=r/2u=r/2 yields ∫0∞rpj2e−r/2dr=2pj2+1Γ(pj2+1). _0^∞r p_j2e^-r/2\, r=2 p_j2+1 \! ( p_j2+1 ). (29) Substituting back, we obtain [‖θ0‖22] [\| _0\|_2^2 ] =∑j=1k2σj2Γ(pj2+1)Γ(pj2) = _j=1^k2 _j^2 \! ( p_j2+1 ) \! ( p_j2 ) (30) =∑j=1kpjσj2 = _j=1^kp_j _j^2 (31) =σ~2p. = σ^2p. (32) Computing Var[‖θ0‖22]Var [\| _0\|_2^2 ]. We next compute the fluctuations of the squared initialization radius. Since the subvectors are independent and ‖θ0‖22=∑j=1kσj2Rj\| _0\|_2^2= _j=1^k _j^2R_j, it follows that Var[‖θ0‖22]=∑j=1kσj4Var[Rj]. [\| _0\|_2^2 ]= _j=1^k _j^4Var[R_j]. (33) To compute Var[Rj]Var[R_j], we first compute the second moment: [Rj2] [R_j^2] =∫0∞r2fRj(r)dr = _0^∞r^2f_R_j(r)\, r (34) =12pj2Γ(pj2)∫0∞rpj2+1e−r/2dr. = 12 p_j2 \! ( p_j2 ) _0^∞r p_j2+1e^-r/2\, r. (35) Applying again the change of variables u=r/2u=r/2 gives ∫0∞rpj2+1e−r/2dr=2pj2+2Γ(pj2+2). _0^∞r p_j2+1e^-r/2\, r=2 p_j2+2 \! ( p_j2+2 ). (36) Thus, [Rj2] [R_j^2] =4Γ(pj2+2)Γ(pj2) =4 \! ( p_j2+2 ) \! ( p_j2 ) (37) =4(pj2+1)(pj2) =4 ( p_j2+1 ) ( p_j2 ) (38) =pj(pj+2). =p_j(p_j+2). (39) Therefore, Var[Rj] [R_j] =[Rj2]−([Rj])2 =E[R_j^2]- (E[R_j] )^2 (40) =pj(pj+2)−pj2 =p_j(p_j+2)-p_j^2 (41) =2pj. =2p_j. (42) Consequently, Var[‖θ0‖22] [\| _0\|_2^2 ] =∑j=1kσj4Var[Rj] = _j=1^k _j^4Var[R_j] (43) =2∑j=1kpjσj4. =2 _j=1^kp_j _j^4. (44) 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 σj=O(1pj),σj2=O(1pj). _j=O\! ( 1 p_j ), _j^2=O\! ( 1p_j ). (45) Substituting this scaling into the definition of the effective coordinate variance gives σ~2=1p∑j=1kpjσj2=O(kp). σ^2= 1p _j=1^kp_j _j^2=O\! ( kp ). (46) Therefore, [‖θ0‖22] [\| _0\|_2^2 ] =∑j=1kpjσj2=σ~2p=O(k), = _j=1^kp_j _j^2= σ^2p=O(k), (47) Var[‖θ0‖22] [\| _0\|_2^2 ] =2∑j=1kpjσj4=O(∑j=1k1pj). =2 _j=1^kp_j _j^4=O\! ( _j=1^k 1p_j ). (48) Using the definition 1p~=1k∑j=1k1pj, 1 p= 1k _j=1^k 1p_j, (49) we obtain Var[‖θ0‖22]=O(kp~). [\| _0\|_2^2 ]=O\! ( k p ). (50) ∎ A.4 Proof: Concentration of Uniform Initialization Lemma 4 (Uniform Initialization Concentration). Let θ0=(θ0(1),…,θ0(k))∈ℝp,θ0(j)∈ℝpj,∑j=1kpj=p, _0= ( _0^(1),…, _0^(k) ) ^p, _0^(j) ^p_j, _j=1^kp_j=p, (51) where the subvectors are independent and the coordinates of each subvector satisfy [θ0(j)]i∼(−εj,εj)[ _0^(j)]_i (- _j, _j). Define the effective coordinate variance ε~2 ^2 and the effective dimension parameter p~ p by ε~2:=1p∑j=1kpjεj2,1p~:=1k∑j=1k1pj. ^2:= 1p _j=1^kp_j _j^2, 1 p:= 1k _j=1^k 1p_j. (52) Then the squared initialization radius satisfies [‖θ0‖22]=ε~2p3,Var(‖θ0‖22)=445∑j=1kpjεj4. [\| _0\|_2^2]= ^2p3, (\| _0\|_2^2)= 445 _j=1^kp_j _j^4. (53) Moreover, by the central limit theorem, ‖θ0‖22≈(ε~2p3,445∑j=1kpjεj4),p→∞. \| _0\|_2^2\; d≈\;N\! ( ^2p3, 445 _j=1^kp_j _j^4 ), p→∞. (54) In particular, for initialization schemes with block-wise coordinate half-widths εj=O(pj−1/2) _j=O(p_j^-1/2), the squared initialization radius satisfies ε~2=O(kp),[‖θ0‖22]=O(k),Var(‖θ0‖22)=O(kp~). ^2=O\! ( kp ), [\| _0\|_2^2]=O(k), (\| _0\|_2^2)=O\! ( k p ). (55) For example, a common numerically stable block-wise choice is εj=pj−1/2 _j=p_j^-1/2, in which case ‖θ0‖22≈(k3,4k45p~). \| _0\|_2^2\; d≈\;N\! ( k3, 4k45 p ). (56) Thus uniform initialization concentrates on a thin hyperspherical shell with squared radius k/3k/3 and squared-radius variance 4k/(45p~)4k/(45 p). Proof. For each coordinate Xi(j):=[θ0(j)]i∼(−εj,εj)X_i^(j):=[ _0^(j)]_i (- _j, _j), the density of Xi(j)X_i^(j) is fXi(j)(x)=12εjf_X_i^(j)(x)= 12 _j for x∈[−εj,εj]x∈[- _j, _j]. Then the second moment is computed as [(Xi(j))2]=∫−εjεjx2fXi(j)(x)dx=12εj∫−εjεjx2dx=εj23. [(X_i^(j))^2]= _- _j _jx^2f_X_i^(j)(x)\, x= 12 _j _- _j _jx^2\, x= _j^23. (57) The fourth moment is computed as [(Xi(j))4]=∫−εjεjx4fXi(j)(x)dx=12εj∫−εjεjx4dx=εj45. [(X_i^(j))^4]= _- _j _jx^4f_X_i^(j)(x)\, x= 12 _j _- _j _jx^4\, x= _j^45. (58) Defining Si(j):=(Xi(j))2S_i^(j):=(X_i^(j))^2, we have [Si(j)]=[(Xi(j))2]E[S_i^(j)]=E[(X_i^(j))^2] and [(Si(j))2]=[(Xi(j))4]E[(S_i^(j))^2]=E[(X_i^(j))^4], so that Var((Xi(j))2)=Var(Si(j)) ((X_i^(j))^2)=Var(S_i^(j)) =[(Si(j))2]−[Si(j)]2 =E[(S_i^(j))^2]-E[S_i^(j)]^2 (59) =[(Xi(j))4]−([(Xi(j))2])2 =E[(X_i^(j))^4]- (E[(X_i^(j))^2] )^2 (60) =εj45−(εj23)2 = _j^45- ( _j^23 )^2 (61) =4εj445. = 4 _j^445. (62) Now set R2=∑j=1k∑i=1pj(Xi(j))2=∑j=1k∑i=1pjSi(j). R^2= _j=1^k _i=1^p_j(X_i^(j))^2= _j=1^k _i=1^p_jS_i^(j). (63) Since the random variables Si(j)S_i^(j) are independent, we obtain [R2] [R^2] =∑j=1kpj[Si(j)]=13∑j=1kpjεj2=ε~2p3, = _j=1^kp_j\,E[S_i^(j)]= 13 _j=1^kp_j _j^2= ^2p3, (64) Var(R2) (R^2) =∑j=1kpjVar(Si(j))=445∑j=1kpjεj4. = _j=1^kp_j\,Var(S_i^(j))= 445 _j=1^kp_j _j^4. (65) Moreover, by the central limit theorem, R2≈(∑j=1kpj[Si(j)],∑j=1kpjVar(Si(j)))=(ε~2p3,445∑j=1kpjεj4),p→∞. R^2 d≈N\! ( _j=1^kp_j\,E[S_i^(j)], _j=1^kp_j\,Var(S_i^(j)) )=N\! ( ^2p3, 445 _j=1^kp_j _j^4 ), p→∞. (66) For numerically stable uniform initialization, the coordinate half-width of each block is often chosen to scale as εj=O(1pj). _j=O\! ( 1 p_j ). (67) Substituting this scaling into the preceding identities gives ε~2=1p∑j=1kpjεj2=O(kp), ^2= 1p _j=1^kp_j _j^2=O\! ( kp ), (68) and therefore [‖θ0‖22]=ε~2p3=O(k),Var(‖θ0‖22)=445∑j=1kpjεj4=O(∑j=1k1pj). [\| _0\|_2^2]= ^2p3=O(k), (\| _0\|_2^2)= 445 _j=1^kp_j _j^4=O\! ( _j=1^k 1p_j ). (69) Using the definition 1p~=1k∑j=1k1pj, 1 p= 1k _j=1^k 1p_j, (70) we obtain Var(‖θ0‖22)=O(kp~). (\| _0\|_2^2)=O\! ( k p ). (71) For example, taking the common block-wise choice εj=1pj, _j= 1 p_j, (72) we obtain R2≈(k3,4k45p~). R^2 d≈N\! ( k3, 4k45 p ). (73) ∎ A.5 Proof: Adam’s Closed-Form Continuous-Time SDE Limit Modeling Mini-Batch Gradient as Wiener Process by Central-Limit Theorem gk=∇ℒf(θk)−1bΣ(θk)1/2ζk,ζk∼(0,Ip)∈ℝp g_k=∇ L_f( _k)- 1 b ( _k)^1/2 _k, _k (0,I_p) ^p Adam’s Discrete Update Rule mk+1=β1mk+(1−β1)gk,vk+1=β2vk+(1−β2)(gk⊙gk),m^k+1=mk+11−β1k+1,v^k+1=vk+11−β2k+1,θk+1=θk−ηm^k+1⊘v^k+1. aligned m_k+1&= _1m_k+(1- _1)g_k,&v_k+1&= _2v_k+(1- _2)(g_k g_k),\\ m_k+1&= m_k+11- _1^k+1,& v_k+1&= v_k+11- _2^k+1,\\ _k+1&= _k-η\, m_k+1 v_k+1. aligned Continuous-Time Interpolation (Discrete–Continuous Scaling) t=ηk,βi=1−ηαi,η→0 t=η k, _i=1-η _i, η→ 0 Continuous-Time Limit of First-Moment Estimate dmt=−α1(mt−g¯t)dt+η/bα1Σt1/2dWt. aligned m_t&=- _1(m_t- g_t) t\\ & + η/b\, _1 _t^1/2 W_t. aligned Continuous-Time Limit of Second-Moment Estimate dvt=−α2(vt−g⊙g¯t)dt+η/bα2DtdWt. aligned v_t&=- _2(v_t- g g_t) t\\ & + η/b\, _2D_t W_t. aligned Expansion of Element-Wise Squared Gradient g⊙g¯t:=g¯t⊙g¯t+1bdiagΣt,Dt:=2diag(g¯t)Σt1/2. aligned g g_t&:= g_t g_t+ 1bdiag\, _t,\\ D_t&:=2\,diag( g_t) _t^1/2. aligned Continuous-Time Limit of Parameter Dynamics dθt=−B(t)(mt⊘vt)dt,whereB(t)=1−e−α2t1−e−α1t. aligned _t&=-B(t) (m_t v_t ) t,\\[4.2679pt] where~B(t)&= 1-e^- _2t1-e^- _1t. aligned Joint-State SDE dSt=μ(St)dt+η/bσ(St)dWt,St=(θt,mt,vt),μ(St)=(−B(t)mt⊘vt−α1(mt−g¯t)−α2(vt−g⊙g¯t)),σ(St)=(0α1Σt1/2α2Dt). aligned S_t&=μ(S_t) t+ η/b\,σ(S_t) W_t, S_t=( _t,m_t,v_t),\\ μ(S_t)&= pmatrix-B(t)m_t v_t\\ - _1(m_t- g_t)\\ - _2(v_t- g g_t) pmatrix, σ(S_t)= pmatrix0\\ _1 _t^1/2\\ _2D_t pmatrix. aligned Sanity Check via Lemma 6: (Induced Radius SDE). (Derived via Itô’s Lemma) drt2=−2B(t)θt⊤(mt⊘vt)dt. r_t^2=-2B(t) _t (m_t v_t ) t. 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 continuous-time 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)S_t:=(m_t,v_t, _t). Lemma 5 (Adam’s Joint-State Continuous-Time SDE Limit). For a network f parameterized with θ∈ℝpθ ^p, let ℓf(s;θ) _f(s;θ) be the loss for sample s, let ℒf(ξ;θ) L_f(ξ;θ) be the mini-batch loss for batch ξ, and let ℒf(θ) L_f(θ) be the population loss over true dataset distribution. Let θk,θt _k, _t denote the parameter at iteration k and time t, respectively. Discrete-Time Update Rules. Consider Adam with small learning rate η, large batch size b, and coefficients (β1,β2)( _1, _2). Let Sk:=(θk⊤,mk⊤,vk⊤)⊤∈ℝ3p,m0=0p,v0=0p, S_k:=( _k ,m_k ,v_k ) ^3p, m_0=0_p, v_0=0_p, (74) denote the discrete-time representation of the parameter, first-moment, and second-moment estimates at iteration k. For a batch ξk _k at step k, the discrete Adam update rules with the bias-correction step are: gk g_k =∇ℒf(ξk;θk) =∇ L_f( _k; _k) (mini-batch gradient) (mini-batch gradient) (75) mk+1 m_k+1 =β1mk+(1−β1)gk = _1m_k+(1- _1)g_k (first-moment estimate) (first-moment estimate) (76) vk+1 v_k+1 =β2vk+(1−β2)(gk⊙gk) = _2v_k+(1- _2)(g_k g_k) (second-moment estimate) (second-moment estimate) (77) m^k+1 m_k+1 =mk+11−β1k+1 = m_k+11- _1^\,k+1 (bias-corrected first-moment estimate) (bias-corrected first-moment estimate) (78) v^k+1 v_k+1 =vk+11−β2k+1 = v_k+11- _2^\,k+1 (bias-corrected second-moment estimate) (bias-corrected second-moment estimate) (79) θk+1 _k+1 =θk−ηh(m^k+1,v^k+1) = _k-η\,h( m_k+1, v_k+1) (bias-corrected parameter update), (bias-corrected parameter update), (80) where h(m^k+1,v^k+1)h( m_k+1, v_k+1) is the per-coordinate Adam update direction, defined as h(m^,v^)=m^⊘(v^+ϵp)=[m^(1)v^(1)+ϵ⋮m^(p)v^(p)+ϵ],h( m, v)\;=\; m ( v+ 1_p)\;=\; bmatrix m^(1) v^(1)+ε\\[4.0pt] \\ m^(p) v^(p)+ε bmatrix, (81) ⊘ 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 _k, the mini-batch gradient gk:=∇ℒf(ξk;θk)g_k:=∇ L_f( _k; _k) is an empirical average of b i.i.d. per-sample gradient gk(s):=∇ℓf(s;θk)g_k(s):=∇ _f(s; _k). Hence, by the central limit theorem, as batch size b→∞b→∞ gk=1b∑i=1bgk(si)⟶(g¯k,1bΣk),g¯k:=s[gk(s)],Σk:=Covs[gk(s)], g_k= 1b _i=1^bg_k(s_i) d N\! ( g_k, 1b _k ), g_k:=E_s[g_k(s)], _k:=Cov_s[g_k(s)], (82) where g¯k g_k and Σk _k are the per-sample gradient expectation and covariance at step k, respectively. Continuous-Time SDE Limit. Under the continuous-time interpolation t=ηkt=η k, with the conditions of sufficiently small learning rate η→0η→ 0 and large batch size b→∞b→∞, the discrete Adam updates admit the Itô SDE limit dSt=μ(St)dt+ηbσ(St)dWt, S_t=μ(S_t) t+ ηb\,σ(S_t) W_t, (83) where Wt∈ℝpW_t ^p is a Wiener process adapted to the filtration generated by the mini-batch sampling process ξt\ _t\. The drift and diffusion factors are μ(St)=(−B(t)mt⊘vt−α1(mt−g¯t)−α2(vt−g⊙g¯t)),σ(St)=(0α1Σt1/2α2Dt), μ(S_t)= pmatrix-\,B(t)\,m_t v_t\\[2.0pt] -\, _1 (m_t- g_t )\\[2.0pt] -\, _2 (v_t- g g_t ) pmatrix, σ(S_t)= pmatrix0\\[2.0pt] _1 _t^1/2\\[2.0pt] _2D_t pmatrix, (84) with B(t):=1−e−α2t1−e−α1t,αi:=1−βiη,i=1,2, B(t):= 1-e^- _2t1-e^- _1t, _i:= 1- _iη, i=1,2, (85) and g⊙g¯t:=g¯t⊙g¯t+1bdiag⁡(Σt),Dt:=2diag⁡(g¯t)Σt1/2. g g_t:= g_t g_t+ 1bdiag( _t), D_t:=2\,diag( g_t) _t^1/2. (86) Proof. We derive the continuous-time limit of the Adam optimizer starting from its discrete update rules. For a network f parameterized with θ∈ℝpθ ^p, let ℓf(s;θ) _f(s;θ) be the loss for sample s, let ℒf(ξ;θ) L_f(ξ;θ) be the mini-batch loss for batch ξ, and let ℒf(θ) L_f(θ) be the population loss over true dataset distribution. Let θk,θt _k, _t denote the parameter at iteration k and time t, respectively. Mini-Batch Gradient as Wiener Process. Let gk(s):=∇ℓf(s;θk)g_k(s):=∇ _f(s; _k) be the per-sample gradient for sample s with the mean and covariance g¯k=s[gk(s)],Σk=Covs[gk(s)], g_k=E_s[g_k(s)], _k=Cov_s [g_k(s) ], (87) respectively. For a batch ξk:=sii=1b _k:=\s_i\_i=1^b of size b, the gradient at time k is given by: gk g_k =1b∑i=1bgk(si), = 1b _i=1^bg_k(s_i), (88) by central limit theorem, so that gkg_k distributionally converge to gk⟶(g¯k,1bΣk), g_k d N\! ( g_k, 1b _k ), (89) as batch size is sufficiently large b→∞b→∞. We therefore can model mini-batch gradient gkg_k as a Wiener process gk=g¯k−1bΣk1/2ζk,(Σk1/2)(Σk1/2)⊤=Σk, g_k= g_k- 1 b _k^1/2 _k, ( _k^1/2)( _k^1/2) = _k, (90) where ζk∼(0,Ip) _k (0,I_p) is Wiener process adapted to the filtration generated from mini-batch sampling sequence ξk\ _k\. Because gradient descent uses negative gradient, we take a negative sign on ζk _k for simplifying the later analysis algebraically. Discrete Adam with Bias Correction. Let η>0η>0 be sufficiently small learning rate, β1,β2∈(0,1) _1, _2∈(0,1) be exponential decay rates, and ϵ>0ε>0 be numerical stability constant. Let θt∈ℝp _t ^p, mtm_t, vtv_t be parameter, first-moment, and second-moment estimates, with the initialization m0=0m_0=0 and v0=0v_0=0. Let k=0,1,2,…k=0,1,2,… be discrete iteration time points. The discrete Adam update rules with the bias-correction step written explicitly are: g¯k g_k :=∇[ℒf(θk)]=∇[ℒf∗(θk)+λ2‖θk‖22] :=∇[ L_f( _k)]=∇[ L^*_f( _k)+ λ2\| _k\|_2^2] (coupled ℓ2-regularizer) (coupled $ _2$-regularizer) (91) gk g_k =g¯k−1bΣk1/2ζk = g_k- 1 b _k^1/2 _k (mini-batch gradient) (mini-batch gradient) (92) mk+1 m_k+1 =β1mk+(1−β1)gk = _1m_k+(1- _1)g_k (first-moment estimate) (first-moment estimate) (93) vk+1 v_k+1 =β2vk+(1−β2)(gk⊙gk) = _2v_k+(1- _2)(g_k g_k) (second-moment estimate) (second-moment estimate) (94) m^k+1 m_k+1 =mk+11−β1k+1 = m_k+11- _1^\,k+1 (bias-corrected first-moment estimate) (bias-corrected first-moment estimate) (95) v^k+1 v_k+1 =vk+11−β2k+1 = v_k+11- _2^\,k+1 (bias-corrected second-moment estimate) (bias-corrected second-moment estimate) (96) θk+1 _k+1 =θk−ηh(m^k+1,v^k+1) = _k-η\,h( m_k+1, v_k+1) (bias-corrected parameter update), (bias-corrected parameter update), (97) where λ is the regularization coefficient, and ℒf(θk) L_f( _k) is the population coupled loss and ℒf∗(θk) L^*_f( _k) is the population task loss, then the per-coordinate Adam update direction is h(m^,v^)=m^⊘(v^+ϵp)=[m^(1)v^(1)+ϵ⋮m^(p)v^(p)+ϵ],h( m, v)\;=\; m ( v+ 1_p)\;=\; bmatrix m^(1) v^(1)+ε\\[4.0pt] \\ m^(p) v^(p)+ε bmatrix, (98) ⊘ 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 _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 _2 penalty is instead decoupled and applied as a standalone update term, θk+1=θk−ηh(m^k+1,v^k+1)−ηλθk, _k+1= _k-η\,h( m_k+1, v_k+1)-η\,λ\, _k, (99) where the resulting optimizer is referred to as AdamW. Continuous–Time Interpolation. Taking continuous-time interpolation in discrete iteration k by: t=kη,t=kη, (100) therefore: dt=limΔt→0Δt=limΔt→0η=η. t= _ t→ 0 t= _ t→ 0η=η. (101) Wiener Increment in Gradient. Let gk(s):=∇ℓf(s;θk)g_k(s):=∇ _f(s; _k) be the per-sample gradient for sample s with the mean and covariance g¯k=s[gk(s)],Σk=Covs[gk(s)], g_k=E_s[g_k(s)], _k=Cov_s [g_k(s) ], (102) respectively. For a batch ξk:=sii=1b _k:=\s_i\_i=1^b of size b, the gradient at time k is given by: gk g_k =1b∑i=1bgk(si), = 1b _i=1^bg_k(s_i), (103) by central limit theorem, so that gkg_k distributionally converge to gk⟶(g¯k,1bΣk), g_k d N\! ( g_k, 1b _k ), (104) as batch size is sufficiently large b→∞b→∞. We therefore can model mini-batch gradient gkg_k as a Wiener process gk=g¯k−1bΣk1/2ζk,(Σk1/2)(Σk1/2)⊤=Σk, g_k= g_k- 1 b _k^1/2 _k, ( _k^1/2)( _k^1/2) = _k, (105) where ζk∼(0,Ip) _k (0,I_p) is Wiener process adapted to the filtration generated from mini-batch sampling sequence ξk\ _k\. Because gradient descent uses negative gradient, we take a negative sign on ζk _k for simplifying the later analysis algebraically. For the adapted noise vector ζk∼(0,Ip)∈ℝp _k (0,I_p) ^p at iteration k (time t=kηt=kη) corresponds to the Wiener increment over [t,t+η][t,t+η]: ζk=1ηζkη=1ηζkΔt=1ηΔWk, _k= 1 η _k η= 1 η _k t= 1 η W_k, (106) where Δt=η t=η and the increment follows the property of Brownian motion: ΔWk=ζkΔt∼(0,ΔtIp). W_k= _k t (0, t\;I_p). (107) Remark 4 (Sanity-Check Reference: Continuous SDE for SGD). Therefore, for SGD, the continuous-time SDE limit is, with t=kηt=kη and Δt=η t=η: dθt _t =limΔt→0(θt+Δt−θt) = _ t→ 0 ( _t+ t- _t ) (108) =limΔt→0[−η(g¯k−1bΣk1/2ζk)] = _ t→ 0 [-η ( g_k- 1 b _k^1/2 _k ) ] (109) =limΔt→0[−ηg¯k+ηbΣk1/2ηζk] = _ t→ 0 [-η g_k+ ηb _k^1/2 η _k ] (110) =limΔt→0[−ηg¯k+ηbΣk1/2ΔWk] = _ t→ 0 [-η g_k+ ηb _k^1/2 W_k ] (111) =−g¯tdt+ηbΣt1/2dWt. =- g_t t+ ηb _t^1/2 W_t. (112) Discrete Increment of First-Moment Estimate. From equation (93): mk+1 m_k+1 =β1mk+(1−β1)gk = _1m_k+(1- _1)g_k (113) mk+1−mk m_k+1-m_k =β1mk+(1−β1)gk−mk = _1m_k+(1- _1)g_k-m_k (114) =(β1−1)mk+(1−β1)gk =( _1-1)m_k+(1- _1)g_k (115) =−(1−β1)mk+(1−β1)gk =-(1- _1)m_k+(1- _1)g_k (116) =(1−β1)(gk−mk) =(1- _1)(g_k-m_k) (117) =−(1−β1)(mk−gk) =-(1- _1)(m_k-g_k) (118) Substituting equation (105): gk=g¯k−1bΣk1/2ζk g_k= g_k- 1 b _k^1/2 _k (119) into equation (118) yields: mk+1−mk m_k+1-m_k =−(1−β1)[mk−g¯k+1bΣk1/2ζk] =-(1- _1) [m_k- g_k+ 1 b _k^1/2 _k ] (120) =−(1−β1)[mk−g¯k]−(1−β1)1bΣk1/2ζk =-(1- _1)[m_k- g_k]-(1- _1) 1 b _k^1/2 _k (121) =−(1−β1)[mk−g¯k]−1−β1bΣk1/2ζk. =-(1- _1)[m_k- g_k]- 1- _1 b _k^1/2 _k. (122) Continuous-Time Limit of First-Moment Estimate. Use Δt=η t=η and ΔWk=Δζk W_k= _k, dm m =limΔt→0Δm = _ t→ 0 m (123) =limΔt→0(mk+1−mk) = _ t→ 0(m_k+1-m_k) (124) =limΔt→0[−(1−β1)[mk−g¯k]−1−β1bΣk1/2ζk] = _ t→ 0 [-(1- _1)[m_k- g_k]- 1- _1 b _k^1/2 _k ] (125) =limΔt→0[−(1−β1)[mk−g¯k]1ηΔt−1−β1bΣk1/21ηΔtζk] = _ t→ 0 [-(1- _1)[m_k- g_k] 1η t- 1- _1 b _k^1/2 1 η t _k ] (126) =−1−β1η[mt−g¯k]dt−1−β1bηΣk1/2dWt. =- 1- _1η[m_t- g_k] t- 1- _1 bη _k^1/2 W_t. (127) Expansion of Element-Wise Squared Gradient. To derive the continuous-time limit of the second-moment estimate dv v. We need to expand gk⊙gkg_k g_k of equation (96) into a tractable form. Consider: gk=g¯k−1bΣk1/2ζk, g_k= g_k- 1 b _k^1/2 _k, (128) we expand the element-wise quadratic product gk⊙gkg_k g_k by: gk⊙gk g_k g_k =[g¯k−1bΣk1/2ζk]⊙[g¯k−1bΣk1/2ζk] = [ g_k- 1 b _k^1/2 _k ] [ g_k- 1 b _k^1/2 _k ] (129) =g¯k⊙g¯k−2(g¯k⊙1bΣk1/2ζk)+1b(Σk1/2ζk)⊙(Σk1/2ζk) = g_k g_k-2 ( g_k 1 b _k^1/2 _k )+ 1b ( _k^1/2 _k ) ( _k^1/2 _k ) (130) =[g¯k]⊙2−2bg¯k⊙[Σk1/2ζk]+1b(Σk1/2ζk)⊙2. = [ g_k ] 2- 2 b g_k [ _k^1/2 _k ]+ 1b( _k^1/2 _k) 2. (131) Using the fundamental identity, for two vectors u and v: u⊙v=diag(u)v, u v=diag(u)v, (132) then the Hadamard cross product can be written as: g¯k⊙[Σk1/2ζk]=diag(g¯k)Σk1/2ζk. g_k [ _k^1/2 _k ]=diag( g_k) _k^1/2 _k. (133) We write gk⊙gkg_k g_k as: gk⊙gk=g¯k⊙2−2bdiag(g¯k)[Σk1/2ζk]+1b(Σk1/2ζk)⊙2. g_k g_k= g_k 2- 2 bdiag( g_k) [ _k^1/2 _k ]+ 1b( _k^1/2 _k) 2. (134) Statistics of Element-Wise Squared Gradient. Fix θk _k over all batches, the element-wise squared gradient gk⊙gkg_k g_k admits a mean ξk[gk⊙gk∣θk] _ _k[g_k g_k _k] =[g¯k⊙2−2bdiag⁡(g¯k)Σk1/2ζk+1b(Σk1/2ζk)⊙2] =E [ g_k 2- 2 bdiag\! ( g_k ) _k^1/2 _k+ 1b ( _k^1/2 _k ) 2 ] =g¯k⊙2−2bdiag⁡(g¯k)[Σk1/2ζk]+1b[(Σk1/2ζk)⊙2] = g_k 2- 2 bdiag\! ( g_k )E\! [ _k^1/2 _k ]+ 1bE [ ( _k^1/2 _k ) 2 ] =g¯k⊙2+1bdiag⁡(Σk), = g_k 2+ 1bdiag\! ( _k ), (135) and a covariance Covξk(gk⊙gk∣θk) _ _k(g_k g_k _k) =Cov(−2bdiag⁡(g¯k)Σk1/2ζk+1b(Σk1/2ζk)⊙2) =Cov (- 2 bdiag\! ( g_k ) _k^1/2 _k+ 1b ( _k^1/2 _k ) 2 ) =4bdiag⁡(g¯k)Σkdiag⁡(g¯k)+1b2Cov((Σk1/2ζk)⊙2) = 4bdiag\! ( g_k ) _kdiag\! ( g_k )+ 1b^2Cov ( ( _k^1/2 _k ) 2 ) =4bdiag⁡(g¯k)Σkdiag⁡(g¯k)+1b2[2(Σk⊙Σk)] = 4bdiag\! ( g_k ) _kdiag\! ( g_k )+ 1b^2 [2 ( _k _k ) ] =4bdiag⁡(g¯k)Σkdiag⁡(g¯k)+2b2(Σk⊙Σk), = 4bdiag\! ( g_k ) _kdiag\! ( g_k )+ 2b^2 ( _k _k ), (136) where the last identity follows componentwise from [Cov((Σk1/2ζk)⊙2)]ij [Cov ( ( _k^1/2 _k ) 2 ) ]_ij =[(Σk1/2ζk)i2(Σk1/2ζk)j2]−[(Σk1/2ζk)i2][(Σk1/2ζk)j2] =E [ ( _k^1/2 _k )_i^2 ( _k^1/2 _k )_j^2 ]-E [ ( _k^1/2 _k )_i^2 ]E [ ( _k^1/2 _k )_j^2 ] =[(Σk)ii(Σk)jj+2(Σk)ij2]−(Σk)ii(Σk)jj = [( _k)_i( _k)_j+2( _k)_ij^2 ]-( _k)_i( _k)_j =2(Σk)ij2=2(Σk⊙Σk)ij. =2( _k)_ij^2=2( _k _k)_ij. (137) Continuous-Time Limit of Second-Moment Estimate. From equation (96): vk+1−vk v_k+1-v_k =β2vk+(1−β2)(gk⊙gk)−vk = _2v_k+(1- _2)(g_k g_k)-v_k (138) =−(1−β2)vk+(1−β2)(gk⊙gk) =-(1- _2)v_k+(1- _2)(g_k g_k) (139) =−(1−β2)[vk−gk⊙gk], =-(1- _2)[v_k-g_k g_k], (140) substituting equation (134) yields: vk+1−vk= v_k+1-v_k= −(1−β2)[vk−g¯k⊙2+2bdiag(g¯k)Σk1/2ζk−1b(Σk1/2ζk)⊙2] -(1- _2) [v_k- g_k 2+ 2 bdiag( g_k) _k^1/2 _k- 1b( _k^1/2 _k) 2 ] (141) Expanding yields: vk+1−vk= v_k+1-v_k= −(1−β2)[vk−g¯k⊙2]−2(1−β2)bdiag(g¯k)Σk1/2ζk+1−β2b(Σk1/2ζk)⊙2. -(1- _2) [v_k- g_k 2 ]- 2(1- _2) bdiag( g_k) _k^1/2 _k+ 1- _2b( _k^1/2 _k) 2. (142) Consider the diffusion terms: 2(1−β2)bdiag(g¯k)Σk1/2ζk 2(1- _2) bdiag( g_k) _k^1/2 _k =2(1−β2)bdiag(g¯k)Σk1/2ΔWkη = 2(1- _2) bdiag( g_k) _k^1/2 W_k η (143) =2(1−β2)bηdiag(g¯k)Σk1/2ΔWk, = 2(1- _2) bηdiag( g_k) _k^1/2 W_k, (144) and: [Σk1/2ζk]i2 [ _k^1/2 _k ]_i^2 =[Σk1/2ΔWkη]i2 = [ _k^1/2 W_k η ]_i^2 (145) =(∑j=1p(Σk)ij1/2ΔWk(j)η)2 = ( _j=1^p( _k)_ij^1/2 W_k^(j) η )^2 (146) =1η∑j=1p∑h=1p(Σk)ij1/2(Σk)ih1/2ΔWk(j)ΔWk(h) = 1η _j=1^p _h=1^p( _k)_ij^1/2( _k)_ih^1/2 W_k^(j) W_k^(h) (147) =1η∑j=1p∑h=1p(Σk)ij1/2(Σk)ih1/2δjhΔt = 1η _j=1^p _h=1^p( _k)_ij^1/2( _k)_ih^1/2 _jh t (148) =1η∑j=1p[(Σk)ij1/2]2Δt = 1η _j=1^p[( _k)_ij^1/2]^2\, t (149) =Δtη(Σk)ii=Δtηdiag(Σk)i, = tη\,( _k)_i\;=\; tη\,diag ( _k )_i\,, (150) so (Σ1/2ζk)⊙2=(Δt/η)diag(Σk)( ^1/2 _k) 2=( t/η)\,diag( _k). Take limits Δt→dt,ΔWk→dWt t→ t, W_k→ W_t, then the continuous-time limit of Δv=vk+1−vk v=v_k+1-v_k is: dvt= v_t= −1−β2η[vt−g¯k⊙2]dt - 1- _2η [v_t- g_k 2 ] t −2(1−β2)bηdiag(g¯t)Σt1/2dWt+1−β2ηbdiag(Σt)dt. - 2(1- _2) bη\,diag( g_t)\, _t^1/2 W_t+ 1- _2η b\,diag ( _t ) t. (151) Continuous-Time Limit of Parameter Dynamics. From equation (97) (with ε dropped), the discrete update uses the bias-corrected moments m^k=mk/(1−β1k),v^k=vk/(1−β2k), m_k=m_k/(1- _1^k), v_k=v_k/(1- _2^k), (152) then θk+1=θk−ηm^k⊘v^k=θk−η1−β2k1−β1kmk⊘vk. _k+1\;=\; _k-η\, m_k v_k\;=\; _k-η\, 1- _2^k1- _1^k\;m_k v_k. (153) Under the joint scaling η→0η→ 0 with (1−βi)/η=αi(1- _i)/η= _i fixed and t=ηkt=η k, 1−βik= 1−(1−αiη)t/η→η→0 1−e−αit, 1- _i^k\;=\;1-(1- _iη)^t/η\; [η→ 0]\;1-e^- _it, 1−β2k1−β1k→η→01−e−α2t1−e−α1t=:B(t), 1- _2^k1- _1^k\; [η→ 0]\; 1-e^- _2t1-e^- _1t=:B(t), (154) with B(t)B(t) defined in equation (85). The continuous-time θ-dynamics is therefore dθt=−B(t)mt⊘vtdt, _t\;=\;-\,B(t)\,m_t v_t\, t, (155) which is time-inhomogeneous through B(t)B(t). Joint-State SDE. To simplify discussion, by equations (134), (135), (136), gt⊙gt=g¯t⊙g¯t−2bdiag(g¯t)[Σt1/2ζt]+1b(Σt1/2ζt)⊙2, g_t g_t= g_t g_t- 2 bdiag( g_t) [ _t^1/2 _t ]+ 1b( _t^1/2 _t) 2, (156) ξt[gt⊙gt∣θt]=g¯t⊙g¯t+1bdiag⁡(Σt), _ _t[g_t g_t _t]= g_t g_t+ 1bdiag\! ( _t ), (157) Covξt(gt⊙gt∣θt)=4bdiag⁡(g¯t)Σtdiag⁡(g¯t)+2b2(Σt⊙Σt), _ _t(g_t g_t _t)= 4bdiag\! ( g_t ) _tdiag\! ( g_t )+ 2b^2 ( _t _t ), (158) we set g⊙g¯t:=[gt⊙gt∣θk]=g¯t⊙g¯t+1bdiag(Σt), g g_t\;:=\;E[g_t g_t _k]= g_t g_t\;+\; 1b\,diag\,( _t), (159) Dt:= 2diag(g¯t)Σt1/2,1bDtDt⊤=Cov[gk⊙gk∣θk]−O(1b2), D_t\;:=\;2\,diag ( g_t )\, _t^1/2, 1bD_tD_t =Cov[g_k g_k _k]-O( 1b^2), (160) where DtD_t is the diffusion factor of the leading cross-term in gk⊙gkg_k g_k; the sub-leading term has Wick’s variance at a scale by O(1/b2)O(1/b^2) and is dropped if b→∞b→∞. Set αi:=(1−βi)/η, _i:=(1- _i)/η, (161) and combining equations (155), (127), (151), dθt _t =−B(t)mt⊘vtdt, =-\,B(t)\,m_t v_t\, t, (162) dmt m_t =−1−β1η[mt−g¯t]dt−1−β1bηΣt1/2dWt =- 1- _1η [m_t- g_t ] t- 1- _1 bη _t^1/2 W_t =−α1(mt−g¯t)dt−ηbα1Σt1/2dWt, =- _1(m_t- g_t) t- ηb\, _1 _t^1/2 W_t, (163) dvt v_t =−1−β2η[vt−g¯t⊙g¯t]dt =- 1- _2η [v_t- g_t g_t ] t −2(1−β2)bηdiag⁡(g¯t)Σt1/2dWt+1−β2ηbdiag⁡(Σt)dt =\;- 2(1- _2) bη\,diag( g_t)\, _t^1/2 W_t+ 1- _2η b\,diag ( _t ) t =−α2[vt−g¯t⊙g¯t−1bdiag⁡(Σt)]dt−ηbα2DtdWt =- _2 [v_t- g_t g_t- 1bdiag( _t) ] t- ηb\, _2D_t W_t =−α2(vt−g⊙g¯t)dt−ηbα2DtdWt, =- _2 (v_t- g g_t ) t- ηb\, _2D_t W_t, (164) so that: d(θtmtvt)=(−B(t)mt⊘vt−α1(mt−g¯t)−α2(vt−g⊙g¯t))dt+ηb(0p×pα1Σt1/2α2Dt)dWt. pmatrix _t\\ m_t\\ v_t pmatrix= pmatrix-B(t)\,m_t v_t\\ - _1\,(m_t- g_t)\\ - _2\,(v_t- g g_t) pmatrix t+ ηb pmatrix0_p× p\\ _1\, _t^1/2\\ _2\,D_t pmatrix W_t. (165) Write the joint-state as St:=(θtmtvt)∈ℝ3p, S_t:= pmatrix _t\\ m_t\\ v_t pmatrix ^3p, (166) then Adam’s optimization is characterized by dSt=μ(St)dt+ηbσ(St)dWt, S_t\;=\;μ(S_t)\, t\;+\; ηb\,σ(S_t)\, W_t, (167) with WtW_t a p-dimensional Wiener process and μ(St)=(−B(t)mt⊘vt−α1(mt−g¯t)−α2(vt−g⊙g¯t))∈ℝ3p,σ(St)=(0p×pα1Σt1/2α2Dt)∈ℝ3p×p.μ(S_t)\;=\; pmatrix-B(t)\,m_t v_t\\ - _1\,(m_t- g_t)\\ - _2\,(v_t- g g_t) pmatrix ^3p, σ(S_t)\;=\; pmatrix0_p× p\\ _1\, _t^1/2\\ _2\,D_t pmatrix ^3p× p. (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 _1=0, so that mt=gtm_t=g_t, and ignore the second-moment state vtv_t equivalently, take β2=0 _2=0 and replace vt v_t by p1_p. Then Adam’s θ-update reduces to the SGD update θk+1=θk−ηgk, _k+1= _k-η g_k, (169) and its continuous-time interpolation gives the standard SGD SDE dθt=−g¯tdt+η/bΣt1/2dWt. _t=- g_t t+ η/b\, _t^1/2 W_t. (170) A.6 Proof: Induced Radius SDE (a) Problem S5S_5 (b) Problem ℤ127Z_127 Figure 8: Sanity Check with Radius SDE. We use the induced radius SDE (Lemma 6), derived via Itô’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 Itô’s lemma (Øksendal, 2003) to the quadratic form rt2=St⊤EStr_t^2=S_t^\! E\,S_t with the θ-block projector E∈ℝ3p×3pE ^3p× 3p defined in equation (176) below must reproduce the Itô decomposition stated in Lemma 6; empirical verification on Adam runs is reported in Figure 8. Lemma 6 (Induced Radius SDE). Let St=(θt⊤,mt⊤,vt⊤)⊤∈ℝ3pS_t=( _t ,m_t ,v_t ) ^3p evolve according to Adam’s joint-state SDE dSt=μ(St)dt+ηbσ(St)dWt, S_t=μ(S_t) t+ ηb\,σ(S_t) W_t, (171) where μ(St)=(−B(t)mt⊘vt−α1(mt−g¯t)−α2(vt−g⊙g¯t)),σ(St)=(0α1Σt1/2α2Dt), μ(S_t)= pmatrix-\,B(t)\,m_t v_t\\ -\, _1(m_t- g_t)\\ -\, _2(v_t- g g_t) pmatrix, σ(S_t)= pmatrix0\\ _1 _t^1/2\\ _2D_t pmatrix, (172) Wt∈ℝpW_t ^p is a Wiener process, and B(t)B(t) is the Adam bias-correction factor. Define rt2:=‖θt‖22, r_t^2:=\| _t\|_2^2, (173) then rt2r_t^2 admits the dynamics drt2=−2B(t)θt⊤(mt⊘vt)dt. r_t^2=-2B(t) _t (m_t v_t ) t. (174) This is not an ordinary differential equation (ODE) as the mtm_t and vtv_t are stochastic through its dependence on the joint state StS_t. Proof. Recall from Lemma 5 the joint-state SDE on St=(θt⊤,mt⊤,vt⊤)⊤∈ℝ3pS_t=( _t ,m_t ,v_t ) ^3p, dSt=μ(St)dt+η/bσ(St)dWt, S_t\;=\;μ(S_t)\, t\;+\; η/b\,σ(S_t)\, W_t, (175) with μ(St)μ(S_t) and σ(St)σ(S_t) as in equations (168). Define a projector E∈ℝ3p×3pE ^3p× 3p E:=(Ip0p0p0p0p0p0p0p0p), E\;:=\; pmatrixI_p&0_p&0_p\\ 0_p&0_p&0_p\\ 0_p&0_p&0_p pmatrix, (176) where each block is p×p× p. Then E is symmetric and idempotent, E⊤=E^\! =E and E2=E^2=E, and for S=(θ⊤,m⊤,v⊤)⊤S=(θ ,m ,v ) the squared parameter norm is the quadratic form r2=‖θ‖22=S⊤ES. r^2\;=\;\|θ\|_2^2\;=\;S^\! E\,S. (177) Apply Itô’s lemma to f(S):=S⊤ESf(S):=S^\! E\,S under equation (175). The gradient and Hessian of f are ∇Sf(S)= 2ES,∇S2f(S)= 2E, _Sf(S)\;=\;2\,E\,S, _S^2f(S)\;=\;2\,E, (178) so drt2= 2St⊤Eμ(St)dt+ 2η/bSt⊤Eσ(St)dWt+ηbtr(σ(St)⊤Eσ(St))dt. r_t^2\;=\;2\,S_t^\! E\,μ(S_t)\, t\;+\;2 η/b\,S_t^\! E\,σ(S_t)\, W_t\;+\; ηb\,tr\! (σ(S_t)^\! E\,σ(S_t) )\, t. (179) Each contraction with E keeps only the θ-block: St⊤Eμ(St) S_t^\! E\,μ(S_t)\; =θt⊤μθ(St)=−B(t)θt⊤(mt⊘vt), =\; _t^\! \, _θ(S_t)\;=\;-\,B(t)\, _t^\! \! (m_t v_t ), (180) St⊤Eσ(St) S_t^\! E\,σ(S_t)\; =θt⊤σθ(St)= 0, =\; _t^\! \, _θ(S_t)\;=\;0, (181) tr(σ(St)⊤Eσ(St)) \! (σ(S_t)^\! E\,σ(S_t) )\; =tr(σθ(St)⊤σθ(St))= 0. =\;tr\! ( _θ(S_t)^\! _θ(S_t) )\;=\;0. (182) Equations (181)–(182) use σθ≡0 _θ≡ 0. Substituting equations (180)–(182) into equation (179) yields drt2=− 2B(t)θt⊤(mt⊘vt)dt, r_t^2\;=\;-\,2\,B(t)\, _t^\! \! (m_t v_t )\, t, (183) which is equation (174). ∎ A.7 Proof: Mean-Field Limit of First- and Second-Moment Estimates Lemma 7 (Mean-Field Limit of First- and Second-Moment Estimates). Let mtm_t and vtv_t be Adam’s first- and second-moment states, let gtg_t be the mini-batch gradient at time t, let g¯t:=[gt∣θt] g_t:=E[g_t _t] be the mean of gtg_t, and let g⊙g¯t:=[gt⊙gt∣θt] g g_t:=E[g_t g_t _t] be the mean of gt⊙gtg_t g_t. Assume that g¯t g_t and g⊙g¯t g g_t vary slowly in the long-term dynamics as t→∞t→∞. Then, under initialization m0=0,v0=0m_0=0,v_0=0, the first- and second-moment estimates admit the mean-field limit mt m_t →g¯t(1−e−α1t),vt→g⊙g¯t(1−e−α2t), → g_t (1-e^- _1t ), v_t→ g g_t (1-e^- _2t ), (184) so that mt⊘vt m_t v_t →g¯t(1−e−α1t)⊘g⊙g¯t(1−e−α2t) → g_t (1-e^- _1t ) g g_t(1-e^- _2t) ≈g¯t(1−e−α1t)⊘g⊙g¯t. ≈ g_t (1-e^- _1t ) g g_t. (185) Proof. We start from the stochastic moment dynamics of the joint-state Adam SDE, dmt m_t =−α1(mt−g¯t)dt−ηbα1Σt1/2dWt, =- _1(m_t- g_t) t- ηb\, _1 _t^1/2 W_t, (186) dvt v_t =−α2(vt−g⊙g¯t)dt−ηbα2DtdWt. =- _2 (v_t- g g_t ) t- ηb\, _2D_t W_t. (187) The mean-field approximation replaces the stochastic moment dynamics by their conditional mean dynamics. Since the stochastic terms are martingale increments, [dWt∣θt,mt,vt] [ W_t _t,m_t,v_t] =0. =0. (188) Taking conditional expectation in (186) and (187) gives [dmt∣θt,mt,vt] [ m_t _t,m_t,v_t] =−α1(mt−g¯t)dt, =- _1(m_t- g_t) t, (189) [dvt∣θt,mt,vt] [ v_t _t,m_t,v_t] =−α2(vt−g⊙g¯t)dt. =- _2 (v_t- g g_t ) t. (190) Thus, under the mean-field dynamics, dmt m_t =−α1(mt−g¯t)dt, =- _1(m_t- g_t) t, (191) dvt v_t =−α2(vt−g⊙g¯t)dt. =- _2 (v_t- g g_t ) t. (192) Because g¯t g_t varies slowly on the mtm_t relaxation time scale, we freeze g¯t g_t when solving equation (191). Hence dmtdt m_t t =−α1mt+α1g¯t. =- _1m_t+ _1 g_t. (193) Multiplying by eα1te _1t gives eα1tdmtdt+α1eα1tmt e _1t m_t t+ _1e _1tm_t =α1eα1tg¯t = _1e _1t g_t ddt(eα1tmt) t (e _1tm_t ) =α1eα1tg¯t. = _1e _1t g_t. (194) Integrating from 0 to t, eα1tmt−m0 e _1tm_t-m_0 =∫0tα1eα1ug¯tdu = _0^t _1e _1u g_t\, u =g¯t(eα1t−1). = g_t (e _1t-1 ). (195) Using m0=0m_0=0, we obtain mt m_t =g¯t(1−e−α1t). = g_t (1-e^- _1t ). (196) Similarly, because g⊙g¯t g g_t varies slowly on the vtv_t relaxation time scale, we freeze g⊙g¯t g g_t when solving (192). Hence dvtdt v_t t =−α2vt+α2g⊙g¯t. =- _2v_t+ _2 g g_t. (197) Multiplying by eα2te _2t gives eα2tdvtdt+α2eα2tvt e _2t v_t t+ _2e _2tv_t =α2eα2tg⊙g¯t = _2e _2t g g_t ddt(eα2tvt) t (e _2tv_t ) =α2eα2tg⊙g¯t. = _2e _2t g g_t. (198) Integrating from 0 to t, eα2tvt−v0 e _2tv_t-v_0 =∫0tα2eα2ug⊙g¯tdu = _0^t _2e _2u g g_t\, u =g⊙g¯t(eα2t−1). = g g_t (e _2t-1 ). (199) Therefore, vt v_t =g⊙g¯t+(v0−g⊙g¯t)e−α2t. = g g_t+ (v_0- g g_t )e^- _2t. (200) Using v0=0v_0=0, this becomes vt v_t =g⊙g¯t(1−e−α2t). = g g_t (1-e^- _2t ). (201) Combining equation (196) and (201), we get mt⊘vt m_t v_t =g¯t(1−e−α1t)⊘g⊙g¯t(1−e−α2t). = g_t (1-e^- _1t ) g g_t(1-e^- _2t). (202) For large t on the vtv_t relaxation time scale, 1−e−α2t 1-e^- _2t decays faster than 1−e−α1t1-e^- _1t, so that mt⊘vt m_t v_t ≈g¯t(1−e−α1t)⊘g⊙g¯t. ≈ g_t (1-e^- _1t ) g g_t. (203) ∎ A.8 Proof: Preconditioned-Decomposition of Adam’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⊤,mt⊤,vt⊤)⊤S_t=( _t ,m_t ,v_t ) evolve under Adam’s joint-state SDE. Define a(t):=1−e−α2t,B(t):=1−e−α2t1−e−α1t, a(t):= 1-e^- _2t, B(t):= 1-e^- _2t1-e^- _1t, (204) π(θt):=diag(g⊙g¯t)−1/2,G(θt)=1bB(t)π(θt)Σt1/2. π( _t):=diag\! ( g g_t )^-1/2, G( _t)= 1 bB(t)π( _t) _t^1/2. (205) where π(θt)π( _t) is referred to as the preconditioner of Adam’s SDE and G(θt)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¯tdt+ℛθ,SM(t)dt+ηG(θt)dWt, _t≈-a(t)π( _t) g_t t+R_θ,SM(t) t+ η\,G( _t) W_t, (206) where ℛθ,SM(t):=−B(t)[[mt⊘vt]−g¯t(1−e−α1t)⊘g⊙g¯t], _θ,SM(t):=-B(t) [E [m_t v_t ]- g_t (1-e^- _1t ) 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=p,π(θt)=IpB(t)=1,m_t=g_t,v_t=1_p,π( _t)=I_p. Proof. We decompose the θ-component of Adam’s joint-state SDE dθt=−B(t)(mt⊘vt)dt, _t=-B(t) (m_t v_t ) t, (208) where B(t)=1−e−α2t1−e−α1t,a(t)=1−e−α2t. B(t)= 1-e^- _2t1-e^- _1t, a(t)= 1-e^- _2t. (209) The SDE preconditioner is defined as π(θt):=diag(g⊙g¯t)−1/2. π( _t):=diag\! ( g g_t )^-1/2. (210) Mean-Field Decomposition. By Lemma 7 ( (Mean-Field Limit of First- and Second-Moment Estimates).), consider the evolution mt⊘vt→g¯t(1−e−α1t)⊘g⊙g¯t, m_t v_t→ g_t (1-e^- _1t ) g g_t, (211) we add and subtract the mean-field limit g¯t(1−e−α1t)⊘g⊙g¯t g_t (1-e^- _1t ) g g_t (212) from equation (208), then dθt _t =−B(t)(mt⊘vt)dt =-B(t) (m_t v_t ) t =−B(t)[g¯t(1−e−α1t)⊘g⊙g¯t]⏟mean-field componentdt =- B(t) [ g_t (1-e^- _1t ) g g_t ]_mean-field component t −B(t)[mt⊘vt−g¯t(1−e−α1t)⊘g⊙g¯t]⏟mean-field residualdt. - B(t) [m_t v_t- g_t (1-e^- _1t ) g g_t ]_mean-field residual t. (213) Since B(t)(1−e−α1t)=a(t), B(t) (1-e^- _1t )=a(t), (214) and g¯t⊘g⊙g¯t=π(θt)g¯t, g_t g g_t=π( _t) g_t, (215) the first term (mean-field component) in equation (213) becomes −B(t)[g¯t(1−e−α1t)⊘g⊙g¯t]dt=−a(t)π(θt)g¯tdt. -B(t) [ g_t (1-e^- _1t ) g g_t ] t=-a(t)π( _t) g_t t. (216) Martingale Decomposition of Mean-Field Residual. Let t[⋅]:=[⋅∣θt]. _t[·]:=E[· _t]. (217) Define ℛθ,SM(t):=−B(t)[t[mt⊘vt]−g¯t(1−e−α1t)⊘g⊙g¯t] _θ,SM(t):=-B(t) [E_t\! [m_t v_t ]- g_t (1-e^- _1t ) g g_t ] (218) as the small-manifold residual. Then the second term in equation (213) admits a martingale decomposition −B(t)[mt⊘vt−g¯t(1−e−α1t)⊘g⊙g¯t]dt -B(t) [m_t v_t- g_t (1-e^- _1t ) g g_t ] t =ℛθ,SM(t)dt−B(t)[mt⊘vt−t[mt⊘vt]]dt. =R_θ,SM(t) t-B(t) [m_t v_t-E_t\! [m_t v_t ] ] t. (219) Effective Diffusion. To analyze the martingale residual mt⊘vt−t[mt⊘vt], m_t v_t-E_t\! [m_t v_t ], (220) define m¯t:=(1−e−α1t)g¯t,v¯t:=g⊙g¯t, m_t:= (1-e^- _1t ) g_t, v_t:= g g_t, (221) and δmt:=mt−m¯t,δvt:=vt−v¯t. δ m_t:=m_t- m_t, δ v_t:=v_t- v_t. (222) A componentwise Taylor expansion of m⊘vm v around (m¯t,v¯t)( m_t, v_t) gives mt⊘vt m_t v_t =m¯t⊘v¯t+π(θt)δmt−12(m¯t⊘v¯t3/2)⊙δvt+θ(t), = m_t v_t+π( _t)δ m_t- 12 ( m_t v_t^3/2 ) δ v_t+Q_θ(t), (223) where θ(t)=O(‖δmt‖‖δvt‖+‖δvt‖22). _θ(t)=O\! (\|δ m_t\|\,\|δ v_t\|+\|δ v_t\|_2^2 ). (224) Taking the conditional mean of equation (223) and subtracting it from equation (223), we obtain mt⊘vt−t[mt⊘vt] m_t v_t-E_t\! [m_t v_t ] =π(θt)(δmt−t[δmt]) =π( _t) (δ m_t-E_t[δ m_t] ) −12(m¯t⊘v¯t3/2)⊙(δvt−t[δvt]) - 12 ( m_t v_t^3/2 ) (δ v_t-E_t[δ v_t] ) +θ(t)−t[θ(t)]. +Q_θ(t)-E_t[Q_θ(t)]. (225) Variation Residual of First-Moment Estimate. We now analyze the variation residual of first-moment estimate δmt−t[δmt]. δ m_t-E_t[δ m_t]. (226) From Adam’s joint-state SDE, the first-moment component satisfies dmt=−α1mtdt+α1gtdt. m_t=- _1m_t t+ _1g_t t. (227) Rearranging equation (227) gives mtdt=gtdt−1α1dmt. m_t t=g_t t- 1 _1 m_t. (228) Taking conditional expectations in equation (228) gives t[mt]dt=g¯tdt−1α1dt[mt]. _t[m_t] t= g_t t- 1 _1 _t[m_t]. (229) Subtracting equation (229) from equation (228) yields (mt−t[mt])dt=(gt−g¯t)dt−1α1d(mt−t[mt]). (m_t-E_t[m_t] ) t= (g_t- g_t ) t- 1 _1 (m_t-E_t[m_t] ). (230) Since δmt−t[δmt]=mt−t[mt], δ m_t-E_t[δ m_t]=m_t-E_t[m_t], (231) we get (δmt−t[δmt])dt=(gt−g¯t)dt−1α1d(δmt−t[δmt]). (δ m_t-E_t[δ m_t] ) t= (g_t- g_t ) t- 1 _1 (δ m_t-E_t[δ m_t] ). (232) In continuous-time diffusion scaling, (gt−g¯t)dt=η/bΣt1/2dWt. (g_t- g_t ) t= η/b\, _t^1/2 W_t. (233) Substituting equation (233) into equation (232) gives (δmt−t[δmt])dt=η/bΣt1/2dWt−1α1d(δmt−t[δmt]). (δ m_t-E_t[δ m_t] ) t= η/b\, _t^1/2 W_t- 1 _1 (δ m_t-E_t[δ m_t] ). (234) Define dℛm(t):=−1α1d(δmt−t[δmt]). _m(t):=- 1 _1 (δ m_t-E_t[δ m_t] ). (235) Then (δmt−t[δmt])dt=η/bΣt1/2dWt+dℛm(t). (δ m_t-E_t[δ m_t] ) t= η/b\, _t^1/2 W_t+ _m(t). (236) Substituting equation (236) into the first term of equation (225) gives −B(t)π(θt)(δmt−t[δmt])dt -B(t)π( _t) (δ m_t-E_t[δ m_t] ) t =−B(t)η/bπ(θt)Σt1/2dWt−B(t)π(θt)dℛm(t). =-B(t) η/b\,π( _t) _t^1/2 W_t-B(t)π( _t) _m(t). (237) Absorbing the sign into WtW_t, and using the definition G(θt):=1bB(t)π(θt)Σt1/2, G( _t):= 1 bB(t)π( _t) _t^1/2, (238) we obtain −B(t)π(θt)(δmt−t[δmt])dt=ηG(θt)dWt−B(t)π(θt)dℛm(t). -B(t)π( _t) (δ m_t-E_t[δ m_t] ) t= η\,G( _t) W_t-B(t)π( _t) _m(t). (239) Produce Claims. Combining equations (225) and (239), we get −B(t)[mt⊘vt−t[mt⊘vt]]dt -B(t) [m_t v_t-E_t\! [m_t v_t ] ] t =ηG(θt)dWt+dθ(t), = η\,G( _t) W_t+ _θ(t), (240) where dθ(t) _θ(t) :=−B(t)π(θt)dℛm(t) :=-B(t)π( _t) _m(t) +B(t)2[(m¯t⊘v¯t3/2)⊙(δvt−t[δvt])]dt + B(t)2 [ ( m_t v_t^3/2 ) (δ v_t-E_t[δ v_t] ) ] t −B(t)[θ(t)−t[θ(t)]]dt. -B(t) [Q_θ(t)-E_t[Q_θ(t)] ] t. (241) Combining equations (216), (219), and (240), we obtain dθt=−a(t)π(θt)g¯tdt+ℛθ,SM(t)dt+ηG(θt)dWt+dθ(t). _t=-a(t)π( _t) g_t t+R_θ,SM(t) t+ η\,G( _t) W_t+ _θ(t). (242) Dropping higher-order terms and keeping the leading diffusion order gives dθt≈−a(t)π(θt)g¯tdt+ℛθ,SM(t)dt+ηG(θt)dWt. _t≈-a(t)π( _t) g_t t+R_θ,SM(t) t+ η\,G( _t) W_t. (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,vtm_t,v_t mt→g¯t(1−e−α1t),vt→g⊙g¯t, m_t→ g_t (1-e^- _1t ), v_t→ g g_t, (244) is referred to as the slow-manifold reduction. When the residual ℛθ,SM(t)R_θ,SM(t) is negligible, then the reduced preconditioned θ-SDE can further be simplified into dθt≈−a(t)π(θt)g¯tdt+ηG(θt)dWt. _t≈-a(t)π( _t) g_t t+ η\,G( _t) W_t. (245) Remark 7 (Quick Sanity Check). As a sanity check, if π(θt)=Ipπ( _t)=I_p, B(t)→1B(t)→ 1, and the residual is ignored at leading order, then G(θt)→1bΣt1/2, G( _t)→ 1 b _t^1/2, (246) and the preconditioned θ-SDE reduces to dθt=−g¯tdt+η/bΣt1/2dWt, _t=- g_t t+ η/b\, _t^1/2 W_t, (247) which is the standard continuous-time SDE approximation of mini-batch SGD. A.9 Proof: Preconditioned-Decomposition of Adam’s Radius 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. 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 tr(π(θt))≈btr[(diagΣt)−1/2]tr(π( _t))≈ b\,tr\! [ (diag\, _t )^-1/2 ] and tr(G(θt)G(θt)⊤)≈ptr(G( _t)G( _t) )≈ p. The predicted radius dynamics 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⊤,mt⊤,vt⊤)⊤S_t=( _t ,m_t ,v_t ) evolve under Adam’s joint-state SDE, and let rt2:=‖θt‖22r_t^2:=\| _t\|_2^2. Define a(t):=1−e−α2t,B(t):=1−e−α2t1−e−α1t a(t):= 1-e^- _2t, B(t):= 1-e^- _2t1-e^- _1t π(θ):=diag(g⊙g¯(θ))−1/2,G(θt)=1bB(t)π(θt)Σt1/2. π(θ):=diag\! ( g g(θ) )^-1/2, G( _t)= 1 bB(t)π( _t) _t^1/2. (248) Then the residual-corrected preconditioned squared-radius dynamics are drt2≈[−2a(t)θt⊤π(θt)g¯t+ℛSM(t)+ηtr(G(θt)G(θt)⊤)]dt+2ηθt⊤G(θt)dWt, r_t^2≈ [-2a(t) _t π( _t) g_t+R_SM(t)+η\,tr\! (G( _t)G( _t) ) ] t+2 η\, _t G( _t) W_t, (249) where the slow-manifold residual is ℛSM(t):=−2B(t)θt⊤[[mt⊘vt]−g¯t(1−e−α1t)⊘g⊙g¯t]. _SM(t):=-2B(t) _t [E [m_t v_t ]- g_t (1-e^- _1t ) g g_t ]. (250) Proof. Let Eθ:=(Ip00000000)∈ℝ3p×3p E_θ:= pmatrixI_p&0&0\\ 0&0&0\\ 0&0&0 pmatrix ^3p× 3p (251) be a projector. Since St=(θt⊤,mt⊤,vt⊤)⊤S_t=( _t ,m_t ,v_t ) , the squared radius can be written as the joint-state quadratic form rt2=‖θt‖22=St⊤EθSt. r_t^2=\| _t\|_2^2=S_t E_θS_t. (252) Therefore, ∇Srt2=2EθSt=(2θt00),∇S2rt2=2Eθ. _Sr_t^2=2E_θS_t= pmatrix2 _t\\ 0\\ 0 pmatrix, _S^2r_t^2=2E_θ. (253) Remark 8. Since θt _t is the θ-component of the joint state St=(θt⊤,mt⊤,vt⊤)⊤S_t=( _t ,m_t ,v_t ) , and the components of StS_t are coupled through Adam’s dynamics, we should regard ‖θt‖22\| _t\|_2^2 as a function of the full joint state rather than of an isolated variable θt _t. Thus, when applying Itô’s lemma to the joint-state SDE, we introduce the projection map Eθ:St↦θtE_θ:S_t _t and write the squared radius as the bilinear form St⊤EθStS_t E_θS_t. Radius Dynamics. Applying Itô’s lemma to rt2=St⊤EθStr_t^2=S_t E_θS_t under Adam’s joint-state SDE gives drt2 r_t^2 =(∇Srt2)⊤dSt+12tr[(ηbσ(St)σ(St)⊤)∇S2rt2]dt. =( _Sr_t^2) S_t+ 12tr [ ( ηbσ(S_t)σ(S_t) ) _S^2r_t^2 ] t. (254) The joint-state diffusion has zero θ-block, because the Brownian noise enters only through the (m,v)(m,v) components. Since ∇S2rt2=2Eθ _S^2r_t^2=2E_θ only selects the θ-block, the second-order term vanishes: tr[(ηbσ(St)σ(St)⊤)∇S2rt2]=0. [ ( ηbσ(S_t)σ(S_t) ) _S^2r_t^2 ]=0. (255) Thus, drt2 r_t^2 =(∇Srt2)⊤dSt =( _Sr_t^2) S_t =2θt⊤dθt. =2 _t _t. (256) From the θ-component of Adam’s joint-state SDE, dθt=−B(t)(mt⊘vt)dt, _t=-B(t) (m_t v_t ) t, (257) we obtain the exact induced squared-radius identity drt2=−2B(t)θt⊤(mt⊘vt)dt. r_t^2=-2B(t) _t (m_t v_t ) t. (258) Mean-Field Decomposition. By Lemma 7 ( (Mean-Field Limit of First- and Second-Moment Estimates).), consider the mean-field limit mt⊘vt→g¯t(1−e−α1t)⊘g⊙g¯t, m_t v_t→ g_t (1-e^- _1t ) g g_t, (259) we add and subtract g¯t(1−e−α1t)⊘g⊙g¯t g_t (1-e^- _1t ) g g_t (260) inside the Adam drift: −2B(t)θt⊤(mt⊘vt) -2B(t) _t (m_t v_t ) =−2B(t)θt⊤[g¯t(1−e−α1t)⊘g⊙g¯t] =-2B(t) _t [ g_t (1-e^- _1t ) g g_t ] −2B(t)θt⊤[mt⊘vt−g¯t(1−e−α1t)⊘g⊙g¯t]. -2B(t) _t [m_t v_t- g_t (1-e^- _1t ) g g_t ]. (261) Using the continuous-time bias-correction factor B(t)=1−e−α2t1−e−α1t, B(t)= 1-e^- _2t1-e^- _1t, (262) and the definition a(t)=1−e−α2ta(t)= 1-e^- _2t, we have B(t)(1−e−α1t)=a(t). B(t) (1-e^- _1t )=a(t). (263) Moreover, since π(θt)=diag(g⊙g¯t)−1/2, π( _t)=diag\! ( g g_t )^-1/2, (264) we have the element-wise identity g¯t⊘g⊙g¯t=π(θt)g¯t. g_t g g_t=π( _t) g_t. (265) Therefore, the first term in equation (261) becomes −2B(t)θt⊤[g¯t(1−e−α1t)⊘g⊙g¯t] -2B(t) _t [ g_t (1-e^- _1t ) g g_t ] =−2a(t)θt⊤π(θt)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: ℛSM(t):=−2B(t)θt⊤[[mt⊘vt]−g¯t(1−e−α1t)⊘g⊙g¯t]. _SM(t):=-2B(t) _t [E [m_t v_t ]- g_t (1-e^- _1t ) g g_t ]. (267) Combining equations (261)–(267), the exact induced squared-radius has a drift [−2a(t)θt⊤π(θt)g¯t+ℛSM(t)]dt. [-2a(t) _t π( _t) g_t+R_SM(t) ] t. (268) Effective Diffusion. Finally, in the reduced preconditioned diffusion approximation, the mini-batch gradient fluctuations induce the θ-diffusion term ηG(θt)dWt η\,G( _t) W_t. Applying Itô’s lemma to this diffusion contribution gives the additional correction ηtr(G(θt)G(θt)⊤)dt+2ηθt⊤G(θt)dWt. η\,tr\! (G( _t)G( _t) ) t+2 η\, _t G( _t) W_t. (269) Produce Claims. Adding equation (269) to equation (268) yields drt2≈[−2a(t)θt⊤π(θt)g¯t+ℛSM(t)+ηtr(G(θt)G(θt)⊤)]dt+2ηθt⊤G(θt)dWt. r_t^2≈ [-2a(t) _t π( _t) g_t+R_SM(t)+η\,tr\! (G( _t)G( _t) ) ] t+2 η\, _t G( _t) W_t. (270) This proves equation (249). ∎ A.10 Proof: Memorization-Regime Preconditioner and Effective Diffusion Identities The validation for Lemma 10 is provided in Figure 9. The experiment shows that, in the memorization regime, the approximated identities π(θt):=diag(g⊙g¯t)−1/2g⊙g¯t≈(1/b)diagΣt⇒tr(π(θt))≈btr[(diagΣt)−1/2] . array[]cπ( _t):=diag\! ( g g_t )^-1/2\\ g g_t≈(1/b)\,diag\, _t array \ (π( _t))≈ b\,tr\! [ (diag\, _t )^-1/2 ] (273) and G(θt)=1bB(t)π(θt)Σt1/2⇒tr(G(θt)G(θt)⊤)≈p G( _t)= 1 bB(t)π( _t) _t^1/2 \! (G( _t)G( _t) )≈ p (274) hold in the memorization regime. Lemma 10 (Memorization-Regime Preconditioner and Effective Diffusion Identities). Let π(θt):=diag(g⊙g¯t)−1/2, π( _t):=diag\! ( g g_t )^-1/2, (275) and define G(θt):=1bB(t)π(θt)Σt1/2,B(t):=1−e−α2t1−e−α1t. G( _t):= 1 bB(t)π( _t) _t^1/2, B(t):= 1-e^- _2t1-e^- _1t. (276) Assume that, in the memorization regime, g⊙g¯t=g¯t⊙g¯t+1bdiag(Σt)≈1bdiag(Σt), g g_t= g_t g_t+ 1bdiag( _t)≈ 1bdiag( _t), (277) and that the diagonal Adam preconditioner admits the trace-matched scalar approximation π(θt)≈s(θt)−1Ip,1s(θ):=1ptr(π(θ)). π( _t)≈ s( _t)^-1I_p, 1s(θ):= 1ptr (π(θ) ). (278) Equivalently, under the corresponding scalar covariance approximation, 1bΣt≈s(θt)2Ip. 1b _t≈ s( _t)^2I_p. (279) Then the following two late-stage identities hold: tr(π(θt)) (π( _t)) ≈btr[(diag(Σt))−1/2], ≈ b\,tr\! [ (diag( _t) )^-1/2 ], (280) tr(G(θt)G(θt)⊤) \! (G( _t)G( _t) ) ≈p. ≈ p. (281) Moreover, G(θt)G(θt)⊤≈Ip(t→∞). G( _t)G( _t) ≈ I_p (t→∞). (282) Proof. From equation (277), g⊙g¯t≈1bdiag(Σt). g g_t≈ 1bdiag( _t). (283) Therefore π(θt) π( _t) =diag(g⊙g¯t)−1/2 =diag\! ( g g_t )^-1/2 ≈diag(1bdiag(Σt))−1/2 \! ( 1bdiag( _t) )^-1/2 =bdiag(diag(Σt))−1/2. = b\,diag\! (diag( _t) )^-1/2. (284) Taking traces gives tr(π(θt)) (π( _t)) ≈btr[(diag(Σt))−1/2]. ≈ b\,tr\! [ (diag( _t) )^-1/2 ]. (285) This proves equation (280). Next, by definition, G(θt)=1bB(t)π(θt)Σt1/2. G( _t)= 1 bB(t)π( _t) _t^1/2. (286) Hence G(θt)G(θt)⊤ G( _t)G( _t) =B(t)2bπ(θt)Σt1/2(Σt1/2)⊤π(θt)⊤ = B(t)^2bπ( _t) _t^1/2 ( _t^1/2 ) π( _t) =B(t)2bπ(θt)Σtπ(θt)⊤. = B(t)^2bπ( _t) _tπ( _t) . (287) Since π(θt)π( _t) is diagonal, it is symmetric. Thus π(θt)⊤=π(θt), π( _t) =π( _t), (288) and equation (287) becomes G(θt)G(θt)⊤=B(t)2bπ(θt)Σtπ(θt). G( _t)G( _t) = B(t)^2bπ( _t) _tπ( _t). (289) Using the trace-matched scalar approximation π(θt)≈s(θt)−1Ip, π( _t)≈ s( _t)^-1I_p, (290) we obtain G(θt)G(θt)⊤ G( _t)G( _t) ≈B(t)2b(s(θt)−1Ip)Σt(s(θt)−1Ip) ≈ B(t)^2b (s( _t)^-1I_p ) _t (s( _t)^-1I_p ) =B(t)2bs(θt)2Σt. = B(t)^2b\,s( _t)^2 _t. (291) By 1bΣt≈s(θt)2Ip, 1b _t≈ s( _t)^2I_p, (292) or equivalently, Σt≈bs(θt)2Ip. _t≈ b\,s( _t)^2I_p. (293) Substituting equation (293) into equation (291) gives G(θt)G(θt)⊤ G( _t)G( _t) ≈B(t)2bs(θt)2bs(θt)2Ip ≈ B(t)^2b\,s( _t)^2b\,s( _t)^2I_p =B(t)2Ip. =B(t)^2I_p. (294) Finally, B(t)2=1−e−α2t(1−e−α1t)2→1(t→∞). B(t)^2= 1-e^- _2t (1-e^- _1t )^2→ 1 (t→∞). (295) Therefore, G(θt)G(θt)⊤≈Ip(t→∞). G( _t)G( _t) ≈ I_p (t→∞). (296) Taking traces yields tr(G(θt)G(θt)⊤)≈tr(Ip)=p. \! (G( _t)G( _t) ) (I_p)=p. (297) This proves equations (281) and (282). ∎ A.11 Proof: Reduced Memorization-Regime Radius 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. 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 ℛSM+ℛπR_ SM+R_π is negligible, as hypothesized. Lemma 11 (Reduced Late-Stage Radius SDE with Slow-Manifold, Preconditioned Residual (Restated)). Let rt:=‖θt‖2r_t:=\| _t\|_2, where rt2r_t^2 follows the residual-corrected preconditioned radius SDE in equation (249) of Lemma 9 ( (Preconditioned-Decomposition of Adam’s Radius SDE).). In the late-stage regime, assume B(t):=1−e−α2t1−e−α1t→1,G(θt)G(θt)⊤≈Ip, B(t):= 1-e^- _2t1-e^- _1t→ 1, G( _t)G( _t) ≈ I_p, (298) so that the residuals from a(t)a(t) and the diffusion covariance are neglected. Then the reduced late-stage squared-radius SDE is drt2≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ηp+ℛSM(t)+ℛπ(t)]dt+2ηrtdWt(r), r_t^2≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+η\,p+R_SM(t)+R_π(t) ] t+2 η\,r_t\, W_t^(r), (299) where ℛSM(t) _SM(t) :=−2B(t)θt⊤[[mt⊘vt]−g¯t(1−e−α1t)⊘g⊙g¯t], :=-2B(t) _t [E [m_t v_t ]- g_t (1-e^- _1t ) g g_t ], (300) ℛπ(t) _π(t) :=−2θt⊤[π(θt)−s(θt)−1Ip]g¯t, :=-2 _t [π( _t)-s( _t)^-1I_p ] g_t, (301) Wt(r) W_t^(r) :=∫0ter(θu)⊤G(θu)dWu,er(θ):=θ‖θ‖2, := _0^te_r( _u) G( _u) W_u, e_r(θ):= θ\|θ\|_2, (302) and under G(θt)G(θt)⊤≈IpG( _t)G( _t) ≈ I_p, Wt(r)W_t^(r) is a one-dimensional Brownian motion. Remark 9. Moreover, ℛSM(t)+ℛπ(t) _SM(t)+R_π(t) =−2B(t)θt⊤[mt⊘vt]+2a(t)θt⊤π(θt)g¯t =-2B(t) _t E [m_t v_t ]+2a(t) _t π( _t) g_t −2θt⊤π(θt)g¯t+2s(θt)θt⊤g¯t -2 _t π( _t) g_t+ 2s( _t) _t g_t =−2B(t)θt⊤[mt⊘vt]+2(a(t)−1)θt⊤π(θt)g¯t+2s(θt)θt⊤g¯t. =-2B(t) _t E [m_t v_t ]+2 (a(t)-1 ) _t π( _t) g_t+ 2s( _t) _t g_t. (303) In particular, in the late-stage limit a(t)→1a(t)→ 1, ℛSM(t)+ℛπ(t)≈−2B(t)θt⊤[mt⊘vt]+2s(θt)θt⊤g¯t. _SM(t)+R_π(t)≈-2B(t) _t E [m_t v_t ]+ 2s( _t) _t g_t. (304) Proof. Starting from the residual-corrected preconditioned radius SDE in equation (249) of Lemma 9 ( (Preconditioned-Decomposition of Adam’s Radius SDE).), drt2 r_t^2 ≈[−2a(t)θt⊤π(θt)g¯t+ℛSM(t)+ηtr(G(θt)G(θt)⊤)]dt ≈ [-2a(t) _t π( _t) g_t+R_SM(t)+η\,tr\! (G( _t)G( _t) ) ] t +2ηθt⊤G(θt)dWt. +2 η\, _t G( _t) W_t. (305) Under a(t)→1a(t)→ 1, −2a(t)θt⊤π(θt)g¯t≈−2θt⊤π(θt)g¯t. -2a(t) _t π( _t) g_t≈-2 _t π( _t) g_t. (306) Decompose the preconditioner as π(θt)=s(θt)−1Ip+[π(θt)−s(θt)−1Ip]. π( _t)=s( _t)^-1I_p+ [π( _t)-s( _t)^-1I_p ]. (307) Substituting this decomposition gives −2θt⊤π(θt)g¯t -2 _t π( _t) g_t =−2s(θt)θt⊤g¯t+ℛπ(t), =- 2s( _t) _t g_t+R_π(t), (308) where ℛπ(t)=−2θt⊤[π(θt)−s(θt)−1Ip]g¯t. _π(t)=-2 _t [π( _t)-s( _t)^-1I_p ] g_t. (309) Next, decompose the mean gradient as g¯t=λθt+(g¯t−λθt). g_t=λ _t+ ( g_t-λ _t ). (310) Then −2s(θt)θt⊤g¯t - 2s( _t) _t g_t =−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt). =- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t ). (311) Combining equations (308) and (311), we obtain −2a(t)θt⊤π(θt)g¯t≈−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ℛπ(t). -2a(t) _t π( _t) g_t≈- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+R_π(t). (312) Under G(θt)G(θt)⊤≈IpG( _t)G( _t) ≈ I_p, ηtr(G(θt)G(θt)⊤)≈ηp. η\,tr\! (G( _t)G( _t) )≈η\,p. (313) Also, 2ηθt⊤G(θt)dWt 2 η\, _t G( _t) W_t =2ηrter(θt)⊤G(θt)dWt =2 η\,r_te_r( _t) G( _t) W_t =2ηrtdWt(r). =2 η\,r_t\, W_t^(r). (314) The quadratic variation is ⟨W(r)⟩t W^(r) _t =∫0ter(θu)⊤G(θu)G(θu)⊤er(θu)du = _0^te_r( _u) G( _u)G( _u) e_r( _u) u ≈∫0ter(θu)⊤er(θu)du=t. ≈ _0^te_r( _u) e_r( _u) u=t. (315) Thus, by Lévy’s characterization (Karatzas and Shreve, 1991), Wt(r)W_t^(r) is a one-dimensional Brownian motion. Substituting equations (312), (313), and (314) into equation (305) yields equation (299) drt2≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ηp+ℛSM(t)+ℛπ(t)]dt+2ηrtdWt(r). r_t^2≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+η\,p+R_SM(t)+R_π(t) ] t+2 η\,r_t\, W_t^(r). (316) ∎ Remark 10. Empirical study in Figure 10 shows that ℛSM(t)+ℛπ(t)R_SM(t)+R_π(t) is negligible during the memorization regime, so that the radius SDE in grokking can be approximated through drt2≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ηp]dt+2ηrtdWt(r). r_t^2≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+η\,p ] t+2 η\,r_t\, W_t^(r). (317) We now analyze the cancellation of ℛSM(t)+ℛπ(t)R_SM(t)+R_π(t). Expanding ℛSM(t)R_SM(t) gives ℛSM(t) _SM(t) =−2B(t)θt⊤[mt⊘vt] =-2B(t) _t E [m_t v_t ] +2B(t)θt⊤[g¯t(1−e−α1t)⊘g⊙g¯t]. +2B(t) _t [ g_t (1-e^- _1t ) g g_t ]. (318) Using B(t)(1−e−α1t)=a(t),g¯t⊘g⊙g¯t=π(θt)g¯t, B(t) (1-e^- _1t )=a(t), g_t g g_t=π( _t) g_t, (319) we obtain ℛSM(t)=−2B(t)θt⊤[mt⊘vt]+2a(t)θt⊤π(θt)g¯t. _SM(t)=-2B(t) _t E [m_t v_t ]+2a(t) _t π( _t) g_t. (320) Adding the explicit expression for ℛπ(t)R_π(t), ℛSM(t)+ℛπ(t) _SM(t)+R_π(t) =−2B(t)θt⊤[mt⊘vt]+2a(t)θt⊤π(θt)g¯t =-2B(t) _t E [m_t v_t ]+2a(t) _t π( _t) g_t −2θt⊤[π(θt)−s(θt)−1Ip]g¯t -2 _t [π( _t)-s( _t)^-1I_p ] g_t =−2B(t)θt⊤[mt⊘vt]+2(a(t)−1)θt⊤π(θt)g¯t+2s(θt)θt⊤g¯t. =-2B(t) _t E [m_t v_t ]+2 (a(t)-1 ) _t π( _t) g_t+ 2s( _t) _t g_t. (321) This proves equation (303). In the late-stage limit a(t),B(t)→1a(t),B(t)→ 1, the middle term vanishes, and therefore ℛSM(t)+ℛπ(t)≈−2B(t)θt⊤[mt⊘vt]+2s(θt)θt⊤g¯t. _SM(t)+R_π(t)≈-2B(t) _t E [m_t v_t ]+ 2s( _t) _t g_t. (322) Hence the residual sum is small when gradient and its first-moment estimate are near zeros. A.12 Proof: Scaling Law of Memorization Radius 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 St=(θt⊤,mt⊤,vt⊤)⊤∈ℝ3pS_t=( _t ,m_t ,v_t ) ^3p, dSt=μ(St)dt+η/bσ(St)dWt, S_t=μ(S_t) t+ η/b\,σ(S_t) W_t, (323) and assume that the bias-correction factor satisfies B(t)→1B(t)→ 1 in the early hitting regime. Let τM(S0):=inft≥0:θt(S0)∈∂M _M(S_0):= \t≥ 0: _t(S_0)∈∂ M\ (324) be the first-hitting time of the memorization boundary. Define the exit-value function u(S):=[∥θτM(S0)∥22|S0=S], u(S):=E\! [\| _ _M(S_0)\|_2^2\, |\,S_0=S ], (325) and the mean squared memorization radius ρM2:=S0∼ΘS[u(S0)], _M^2:=E_S_0 _S\! [u(S_0) ], (326) where ΘS _S is the initialization distribution of the joint state, with S0=(θ0⊤,0p⊤,0p⊤)⊤S_0=( _0 ,0_p ,0_p ) . Let Eθ∈ℝ3p×3pE_θ ^3p× 3p be the θ-block projector Eθ:=(Ip00000000), E_θ:= pmatrixI_p&0&0\\ 0&0&0\\ 0&0&0 pmatrix, (327) so that for any joint state S=(θ,m,v)S=(θ,m,v), ‖θ‖22=S⊤EθS. \|θ\|_2^2=S E_θS. (328) Let Φt(S0) _t(S_0) denote the deterministic Adam flow generated by the drift field μ(S)μ(S), ddtΦt(S0)=μ(Φt(S0)),Φ0(S0)=S0. t _t(S_0)=μ( _t(S_0)), _0(S_0)=S_0. (329) Define the deterministic first-hitting time τM(0)(S0):=inft≥0:θ(Φt(S0))∈∂M, _M^(0)(S_0):= \t≥ 0:θ( _t(S_0))∈∂ M\, (330) and the deterministic exit-radius-squared map R2(S0):=ΦτM(0)(S0)(S0)⊤EθΦτM(0)(S0)(S0). R^2(S_0):= _ _M^(0)(S_0)(S_0) E_θ _ _M^(0)(S_0)(S_0). (331) Let ΣS(S):=σ(S)σ(S)⊤ _S(S):=σ(S)σ(S) (332) be the joint-state diffusion covariance. Then, for sufficiently small η/bη/b, ρM2=ρM(0) 2+(η/b)cM+(η/b)2cM(2)+O((η/b)3), _M^2= _M^(0)\,2+(η/b)c_M+(η/b)^2c_M^(2)+O ((η/b)^3 ), (333) where ρM(0) 2 _M^(0)\,2 :=S0∼ΘS[R2(S0)], :=E_S_0 _S\! [R^2(S_0) ], (334) cM c_M :=S0∼ΘS∫0τM(0)(S0)12tr(ΣS(Φt(S0))∇S2R2(Φt(S0)))dt, :=E_S_0 _S _0 _M^(0)(S_0) 12\,tr\! ( _S( _t(S_0)) _S^2R^2( _t(S_0)) ) t, (335) cM(2) c_M^(2) :=S0∼ΘS∫0τM(0)(S0)12tr(ΣS(Φt(S0))∇S2u1(Φt(S0)))dt. :=E_S_0 _S _0 _M^(0)(S_0) 12\,tr\! ( _S( _t(S_0)) _S^2u_1( _t(S_0)) ) t. (336) Here u1u_1 is the first-order perturbation corrector u1(S0):=∫0τM(0)(S0)12tr(ΣS(Φt(S0))∇S2R2(Φt(S0)))dt. u_1(S_0):= _0 _M^(0)(S_0) 12\,tr\! ( _S( _t(S_0)) _S^2R^2( _t(S_0)) ) t. (337) The constants ρM(0) 2 _M^(0)\,2, cMc_M, and cM(2)c_M^(2) are task-determined and independent of η and b; ℓ2 _2 regularization coefficient λ enters (ρM(0))2( _M^(0))^2 through deterministic gradient flow with scaling law (ρM(0))2∝exp⁡(−O(λ))( _M^(0))^2 \! (-O(λ) ). Proof. Let τM(S0):=inft≥0:θt(S0)∈∂M _M(S_0):= \t≥ 0: _t(S_0)∈∂ M\ (338) be the first-hitting time of the memorization boundary for the trajectory θt\ _t\ starting from S0=(θ0⊤,0p⊤,0p⊤)⊤S_0=( _0 ,0_p ,0_p ) . Bilinear Radius Projector. Since the parameter θt _t is the θ-component of the joint state StS_t, define the θ-block projector Eθ:=(Ip00000000)∈ℝ3p×3p. E_θ:= pmatrixI_p&0&0\\ 0&0&0\\ 0&0&0 pmatrix ^3p× 3p. (339) Thus, for any joint state S, the squared parameter radius is the joint-state quadratic form r2(S):=‖θ‖22=S⊤EθS. r^2(S):=\|θ\|_2^2=S E_θS. (340) In particular, ∇Sr2(S)=2EθS,∇S2r2(S)=2Eθ. _Sr^2(S)=2E_θS, _S^2r^2(S)=2E_θ. (341) Define Dirichlet Boundary Problem. We introduce the exit-value function u(S):=[r2(SτM)|S0=S]=[SτM⊤EθSτM|S0=S], u(S):=E\! [r^2(S_ _M)\, |\,S_0=S ]=E\! [S_ _M E_θS_ _M\, |\,S_0=S ], (342) which is the expected squared parameter norm at the first-hitting time of ∂M∂ M, starting from S. We now derive the PDE satisfied by u. The joint Adam SDE has the form dSt=μ(St)dt+η/bσ(St)dWt. S_t=μ(S_t)\, t+ η/b\,σ(S_t)\, W_t. (343) Therefore its infinitesimal generator is ℒSf(S)=μ(S)⊤∇Sf(S)+12ηbtr(ΣS(S)∇S2f(S)),ΣS(S):=σ(S)σ(S)⊤. _Sf(S)=μ(S) _Sf(S)+ 12\, to4.97pt to13.19pt \> -3.4722pt to0.0pt @beginscope @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke @setlinewidth @invoke to0.0pt @beginscope @invoke @beginscope @invoke @beginscope @invoke @transformcm1.00.00.01.0-2.48265pt0.0pt @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke $ ηb$ @invoke @endscope @invoke @endscope @invoke @endscope to0.0pt @discardpath @invoke @endscope to4.12pt to9.24pt -2.43054pt to0.0pt @beginscope @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke @setlinewidth @invoke to0.0pt @beginscope @invoke @beginscope @invoke @beginscope @invoke @transformcm1.00.00.01.0-2.06009pt0.0pt @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke $ ηb$ @invoke @endscope @invoke @endscope @invoke @endscope to0.0pt @discardpath @invoke @endscope to3.72pt to6.6pt -1.73608pt to0.0pt @beginscope @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke @setlinewidth @invoke to0.0pt @beginscope @invoke @beginscope @invoke @beginscope @invoke @transformcm1.00.00.01.0-1.85939pt0.0pt @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke $ ηb$ @invoke @endscope @invoke @endscope @invoke @endscope to0.0pt @discardpath @invoke @endscope to3.72pt to6.6pt -1.73608pt to0.0pt @beginscope @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke @setlinewidth @invoke to0.0pt @beginscope @invoke @beginscope @invoke @beginscope @invoke @transformcm1.00.00.01.0-1.85939pt0.0pt @invoke pgfstrokecolorrgb0,0,0 @color@rgb@stroke000 @invoke @color@rgb@fill000 @invoke $ ηb$ @invoke @endscope @invoke @endscope @invoke @endscope to0.0pt @discardpath @invoke @endscope \,tr\! ( _S(S) _S^2f(S) ), _S(S):=σ(S)σ(S) . (344) Learning rate η and batch size bbmodulate infinitesimal generator Remark 11. The structure of the infinitesimal generator ℒSf(S)L_Sf(S) shows that this operator contains a term modulated by the learning rate η and the batch size b through the coefficient η/bη/b. This motivates us to treat ℒSf(S)L_Sf(S) as an operator expansion with respect to the coefficient η/bη/b, so that we can apply a regular perturbation expansion to study the corresponding expansion of the memorization radius with respect to η/bη/b. Dirichlet Boundary PDE. Since u(S)u(S) is the expected boundary value of the process stopped at ∂M∂ M, it solves the Dirichlet boundary value problem (Øksendal, 2003) ℒSu(S) _Su(S) =0, =0, θ∈M∘, θ∈ M , (345) u(S) u(S) =S⊤EθS, =S E_θS, θ∈∂M. θ∈∂ M. (346) Here equation (345) states that u(St)u(S_t) is harmonic with respect to the stopped Adam diffusion before hitting the boundary, while equation (346) assigns the squared radius as the boundary payoff. Let ε:=η/b :=η/b. We split the generator into its deterministic and stochastic parts: ℒS=ℒS(0)+εℒS(1),ℒS(0)f:=μ⊤∇Sf,ℒS(1)f:=12tr(ΣS∇S2f). _S=L_S^(0)+ _S^(1), _S^(0)f:=μ _Sf, _S^(1)f:= 12\,tr\! ( _S _S^2f ). (347) Thus the Dirichlet PDE becomes (ℒS(0)+εℒS(1))u(S)=0,θ∈M∘, (L_S^(0)+ _S^(1) )u(S)=0, θ∈ M , (348) with boundary condition u(S)=S⊤EθS,θ∈∂M. u(S)=S E_θS, θ∈∂ M. (349) Solving PDE via Regular Perturbation Expansion. We solve this PDE by a regular perturbation expansion in the small parameter ε . Write u(S)=u0(S)+εu1(S)+ε2u2(S)+O(ε3). u(S)=u_0(S)+ u_1(S)+ ^2u_2(S)+O( ^3). (350) Substituting equation (350) into equation (348) gives 0 0 =(ℒS(0)+εℒS(1))(u0+εu1+ε2u2+O(ε3)) = (L_S^(0)+ _S^(1) ) (u_0+ u_1+ ^2u_2+O( ^3) ) =ℒS(0)u0+ε(ℒS(0)u1+ℒS(1)u0)+ε2(ℒS(0)u2+ℒS(1)u1)+O(ε3). =L_S^(0)u_0+ (L_S^(0)u_1+L_S^(1)u_0 )+ ^2 (L_S^(0)u_2+L_S^(1)u_1 )+O( ^3). (351) Since this identity must hold for all sufficiently small ε , each coefficient of ε must vanish. Therefore, ε0: ^0: ℒS(0)u0=0, _S^(0)u_0=0, (352) ε1: ^1: ℒS(0)u1=−ℒS(1)u0, _S^(0)u_1=-L_S^(1)u_0, (353) ε2: ^2: ℒS(0)u2=−ℒS(1)u1. _S^(0)u_2=-L_S^(1)u_1. (354) The boundary condition is expanded in the same way: u0+εu1+ε2u2+O(ε3)=S⊤EθS,θ∈∂M. u_0+ u_1+ ^2u_2+O( ^3)=S E_θS, θ∈∂ M. (355) Matching powers of ε on the boundary gives u0|θ∈∂M u_0 |_θ∈∂ M =S⊤EθS, =S E_θS, (356) u1|θ∈∂M u_1 |_θ∈∂ M =0, =0, (357) u2|θ∈∂M u_2 |_θ∈∂ M =0. =0. (358) Zero-Order Expansion. We now solve these equations by characteristics. The characteristic curves of ℒS(0)=μ⊤∇SL_S^(0)=μ _S are the deterministic Adam trajectories ddtΦt(S0)=μ(Φt(S0)),Φ0(S0)=S0, t _t(S_0)=μ( _t(S_0)), _0(S_0)=S_0, (359) so that θ˙(0) θ^(0) =−m(0)⊘v(0), =-m^(0) v^(0), m˙(0) m^(0) =−α1(m(0)−g¯(θ(0))), =- _1 (m^(0)- g(θ^(0)) ), v˙(0) v^(0) =−α2(v(0)−g⊙g¯(θ(0))). =- _2 (v^(0)- g g(θ^(0)) ). (360) For any smooth function w, along the deterministic flow we have ddtw(Φt(S0))=ℒS(0)w(Φt(S0)). tw( _t(S_0))=L_S^(0)w( _t(S_0)). (361) Let τM(0)(S0) _M^(0)(S_0) be the deterministic hitting time τM(0)(S0):=inft≥0:(Φt(S0))θ∈∂M. _M^(0)(S_0):= \t≥ 0:( _t(S_0))_θ∈∂ M\. (362) At order ε0 ^0, equation (352) gives ddtu0(Φt(S0))=0. tu_0( _t(S_0))=0. (363) Thus u0u_0 is constant along the deterministic trajectory. Evaluating it at the deterministic hitting time and using the boundary condition in equation (356), we obtain u0(S0) u_0(S_0) =u0(ΦτM(0)(S0)(S0)) =u_0( _ _M^(0)(S_0)(S_0)) =ΦτM(0)(S0)(S0)⊤EθΦτM(0)(S0)(S0) = _ _M^(0)(S_0)(S_0) E_θ _ _M^(0)(S_0)(S_0) =‖(ΦτM(0)(S0)(S0))θ‖22. = \| ( _ _M^(0)(S_0)(S_0) )_θ \|_2^2. (364) Define the deterministic exit-radius-squared map R2(S0):=ΦτM(0)(S0)(S0)⊤EθΦτM(0)(S0)(S0). R^2(S_0):= _ _M^(0)(S_0)(S_0) E_θ _ _M^(0)(S_0)(S_0). (365) Equivalently, R2(S0)=‖(ΦτM(0)(S0)(S0))θ‖22. R^2(S_0)= \| ( _ _M^(0)(S_0)(S_0) )_θ \|_2^2. (366) Then u0(S0)=R2(S0)u_0(S_0)=R^2(S_0). First-Order Expansion. At order ε1 ^1, equation (353) and the transport identity give ddtu1(Φt(S0))=−ℒS(1)u0(Φt(S0)). tu_1( _t(S_0))=-L_S^(1)u_0( _t(S_0)). (367) Integrating from 0 to τM(0)(S0) _M^(0)(S_0) gives u1(ΦτM(0)(S0)(S0))−u1(S0)=−∫0τM(0)(S0)ℒS(1)u0(Φt(S0))dt. u_1( _ _M^(0)(S_0)(S_0))-u_1(S_0)=- _0 _M^(0)(S_0)L_S^(1)u_0( _t(S_0))\, t. (368) Since u1=0u_1=0 on ∂M∂ M, the first term on the left vanishes. Therefore, u1(S0)=∫0τM(0)(S0)ℒS(1)u0(Φt(S0))dt. u_1(S_0)= _0 _M^(0)(S_0)L_S^(1)u_0( _t(S_0))\, t. (369) Using the definition of ℒS(1)L_S^(1) and u0=R2u_0=R^2 along the deterministic flow, u1(S0)=∫0τM(0)(S0)12tr(ΣS(Φt(S0))∇S2R2(Φt(S0)))dt. u_1(S_0)= _0 _M^(0)(S_0) 12\,tr\! ( _S( _t(S_0)) _S^2R^2( _t(S_0)) )\, t. (370) Second-Order Expansion. At order ε2 ^2, the same argument gives ddtu2(Φt(S0))=−ℒS(1)u1(Φt(S0)). tu_2( _t(S_0))=-L_S^(1)u_1( _t(S_0)). (371) Since u2=0u_2=0 on ∂M∂ M, integrating along the deterministic characteristic yields u2(S0)=∫0τM(0)(S0)ℒS(1)u1(Φt(S0))dt. u_2(S_0)= _0 _M^(0)(S_0)L_S^(1)u_1( _t(S_0))\, t. (372) Thus, u2(S0)=∫0τM(0)(S0)12tr(ΣS(Φt(S0))∇S2u1(Φt(S0)))dt. u_2(S_0)= _0 _M^(0)(S_0) 12\,tr\! ( _S( _t(S_0)) _S^2u_1( _t(S_0)) )\, t. (373) Produce Claims. Combining the three orders, we have u(S0)=R2(S0)+εu1(S0)+ε2u2(S0)+O(ε3). u(S_0)=R^2(S_0)+ u_1(S_0)+ ^2u_2(S_0)+O( ^3). (374) Finally, the mean squared memorization radius is obtained by averaging this exit-value function over the initialization distribution: ρM2 _M^2 :=[‖θτM‖22]=[SτM⊤EθSτM]=S0∼ΘS[u(S0)] :=E\! [\| _ _M\|_2^2 ]=E\! [S_ _M E_θS_ _M ]=E_S_0 _S\! [u(S_0) ] =S0∼ΘS[R2(S0)]+εS0∼ΘS[u1(S0)]+ε2S0∼ΘS[u2(S0)]+O(ε3). =E_S_0 _S\! [R^2(S_0) ]+ _S_0 _S\! [u_1(S_0) ]+ ^2E_S_0 _S\! [u_2(S_0) ]+O( ^3). (375) Define ρM(0) 2 _M^(0)\,2 :=S0∼ΘS[R2(S0)], :=E_S_0 _S\! [R^2(S_0) ], cM c_M :=S0∼ΘS[u1(S0)], :=E_S_0 _S\! [u_1(S_0) ], cM(2) c_M^(2) :=S0∼ΘS[u2(S0)]. :=E_S_0 _S\! [u_2(S_0) ]. (376) Substituting these definitions into equation (375) gives ρM2=ρM(0) 2+εcM+ε2cM(2)+O(ε3). _M^2= _M^(0)\,2+ c_M+ ^2c_M^(2)+O( ^3). (377) Since ε=η/b =η/b, we obtain ρM2=ρM(0) 2+(η/b)cM+(η/b)2cM(2)+O((η/b)3), _M^2= _M^(0)\,2+(η/b)c_M+(η/b)^2c_M^(2)+O ((η/b)^3 ), (378) which proves the theorem. ∎ Dependence of ρM(0) _M^(0) on λ under weak task gradient. The leading constant (ρM(0))2( _M^(0))^2 is determined by the deterministic gradient flow. In the regime where the task gradient is small, ∇ℒf∗(θt) ∇ L_f^*( _t) ≈0, ≈ 0, (379) the deterministic flow is dominated by the ℓ2 _2 drift: dθt _t ≈−λθtdt. ≈-λ _t\, t. (380) Hence d‖θt‖22 \| _t\|_2^2 =2θt⊤dθt =2 _t _t ≈−2λ‖θt‖22dt. ≈-2λ\| _t\|_2^2\, t. (381) Equivalently, 1‖θt‖22d‖θt‖22 1\| _t\|_2^2 \| _t\|_2^2 ≈−2λdt. ≈-2λ\, t. (382) Integrating along the deterministic memorization path gives ‖θτM(0)‖22 \| _ _M^(0)\|_2^2 ≈‖θ0‖22exp⁡(−2λτM(0)). ≈\| _0\|_2^2 \! (-2λ _M^(0) ). (383) Therefore, (ρM(0))2 ( _M^(0))^2 ∝exp⁡(−O(λτM(0))). \! (-O(λ _M^(0)) ). (384) When τM(0) _M^(0) is weakly dependent on λ in the considered regime, (ρM(0))2 ( _M^(0))^2 ∝exp⁡(−O(λ)). \! (-O(λ) ). (385) Since ρM2=(ρM(0))2+O(η), _M^2=( _M^(0))^2+O(η), (386) the same leading dependence gives ρM2 _M^2 ∝exp⁡(−O(λ)). \! (-O(λ) ). (387) A.13 Proof: Scaling Law of Generalization Radius Theorem 5 (Scaling Law of Generalization Radius, restated). We do not assume a global minizer, for each initialization S0∈ΘS_0∈ _S, let θ⋆(S0):=arg⁡minθ(S0)⁡ℒf∗(θ(S0))+λ2‖θ(S0)‖22∈G θ (S_0):= _θ(S_0) \ L^*_f(θ(S_0))+ λ2\|θ(S_0)\|_2^2 \∈ G (388) denote the regularized local minimizer selected by the trajectory starting from initial state S0S_0, and ℒf∗(θ(S0)) L^*_f(θ(S_0)) is the task loss. The trajectory St\S_t\ will be confined in the basin centered at θ⋆(S0)θ (S_0). Define (ρG(0))2:=S0∈ΘS[‖θ⋆(S0)‖22]. ( _G^(0) )^2:=E_S_0∈ _S [\|θ (S_0)\|_2^2 ]. (389) Let τG(S0):=inft≥0:θt(S0)∈∂G, _G(S_0):= \t≥ 0: _t(S_0)∈∂ G\, (390) and define the generalization radius by ρG2:=S0∈ΘS[[‖θτG(S0)‖22∣S0]]. _G^2:=E_S_0∈ _S [E [\| _ _G(S_0)\|_2^2 S_0 ] ]. (391) Let s⋆(S0):=s(θ⋆(S0)),D⋆(S0):=G(θ⋆(S0))G(θ⋆(S0))⊤, s (S_0):=s(θ (S_0)), D (S_0):=G(θ (S_0))G(θ (S_0)) , (392) where the effective scalar preconditioner at the selected minimizer is defined by 1s⋆(S0)=1ptr(π(θ⋆(S0))),π(θ⋆(S0))=diag(g⊙g¯(θ⋆(S0)))−1/2. 1s (S_0)= 1ptr (π(θ (S_0)) ), π(θ (S_0))=diag\! ( g g(θ (S_0)) )^-1/2. (393) Then, generalization radius approximately admits, ρG2≈(ρG(0))2+ηλcG+O(η2λ2), _G^2≈ ( _G^(0) )^2+ ηλ\,c_G+O\! ( η^2λ^2 ), (394) where cG:=S0∈ΘS[12s⋆(S0)tr(D⋆(S0))]>0. c_G:=E_S_0∈ _S [ 12s (S_0)\,tr\! (D (S_0) ) ]>0. (395) Proof. For a trajectory St\S_t\ starting from S0∈ΘS_0∈ _S, write θ⋆ θ :=θ⋆(S0), :=θ (S_0), θt _t =θ⋆+δt. =θ + _t. (396) By definition of the local minimizer, g¯∗(θ⋆)+λθ⋆=0, g^*(θ )+λθ =0, (397) where g¯∗(∙) g^*( ) denotes the task loss mean. By Lemma 11 ( (Reduced Late-Stage Radius SDE with Slow-Manifold, Preconditioned Residual (Restated)).), together with Lemma 10 ( (Memorization-Regime Preconditioner and Effective Diffusion Identities).), G(θt)G(θt)⊤≈Ip, G( _t)G( _t) ≈ I_p, (398) and the residual ℛSM(t)+ℛπ(t)≈0, _SM(t)+R_π(t)≈ 0, (399) the reduced late-stage squared-radius SDE gives drt2≈[−2s(θt)θt⊤(g¯t∗+λθt)+ηp]dt+2ηrtdWt(r), r_t^2≈ [- 2s( _t) _t ( g^*_t+λ _t )+η p ] t+2 η\,r_t\, W_t^(r), (400) where rt2=‖θt‖22. r_t^2=\| _t\|_2^2. (401) Linearization Near Local Minizer. The Taylor expansion of g¯∗ g^* around θ⋆θ gives g¯t∗ g^*_t =g¯∗(θ⋆+δt) = g^*(θ + _t) =g¯∗(θ⋆)+∇g¯∗(θ⋆)δt+O(‖δt‖22). = g^*(θ )+∇ g^*(θ ) _t+O\! (\| _t\|_2^2 ). (402) Therefore, g¯t∗+λθt g^*_t+λ _t =g¯∗(θ⋆)+∇g¯∗(θ⋆)δt+λθ⋆+λδt+O(‖δt‖22) = g^*(θ )+∇ g^*(θ ) _t+λθ +λ _t+O\! (\| _t\|_2^2 ) =(∇g¯∗(θ⋆)+λIp)δt+O(‖δt‖22), = (∇ g^*(θ )+λ I_p ) _t+O\! (\| _t\|_2^2 ), (403) where equation (397) was used. Under the dominant-regularization approximation, ∇g¯∗(θ⋆)+λIp ∇ g^*(θ )+λ I_p ≈λIp, ≈λ I_p, (404) g¯t∗+λθt g^*_t+λ _t ≈λδt+O(‖δt‖22). ≈λ _t+O\! (\| _t\|_2^2 ). (405) Using θt=θ⋆+δt _t=θ + _t, θt⊤(g¯t∗+λθt) _t ( g^*_t+λ _t ) ≈(θ⋆+δt)⊤[λδt+O(‖δt‖22)] ≈(θ + _t) [λ _t+O\! (\| _t\|_2^2 ) ] =λθ⋆⊤δt+λ‖δt‖22+O(‖δt‖22). =λθ _t+λ\| _t\|_2^2+O\! (\| _t\|_2^2 ). (406) Local Stationary Approximation. Under the local stationary approximation, s(θt) s( _t) =s⋆(S0)+O(‖δt‖2), =s (S_0)+O\! (\| _t\|_2 ), s⋆(S0) s (S_0) :=s(θ⋆(S0)). :=s(θ (S_0)). (407) Taking conditional expectation in equation (400) gives 0 0 ≈−2s⋆(S0)[θt⊤(g¯t∗+λθt)∣S0]+ηp+O(η2λ2), ≈- 2s (S_0)E [ _t ( g^*_t+λ _t ) S_0 ]+η p+O\! ( η^2λ^2 ), (408) since [rtdWt(r)∣S0]=0. [r_t\, W_t^(r) S_0 ]=0. (409) The local stationary fluctuation is centered, so [δt∣S0] [ _t S_0] =0, =0, [θ⋆⊤δt∣S0] [θ _t S_0] =0. =0. (410) Combining equations (406), (408), and (410), 0 0 ≈−2λs⋆(S0)[‖δt‖22∣S0]+ηp+O(η2λ2). ≈- 2λs (S_0)E [\| _t\|_2^2 S_0 ]+η p+O\! ( η^2λ^2 ). (411) Hence [‖δt‖22∣S0]≈ηλs⋆(S0)2p+O(η2λ2). [\| _t\|_2^2 S_0 ]≈ ηλ s (S_0)2p+O\! ( η^2λ^2 ). (412) Also, ‖θt‖22 \| _t\|_2^2 =‖θ⋆+δt‖22 =\|θ + _t\|_2^2 =‖θ⋆‖22+2θ⋆⊤δt+‖δt‖22. =\|θ \|_2^2+2θ _t+\| _t\|_2^2. (413) Taking conditional expectation gives [‖θt‖22∣S0] [\| _t\|_2^2 S_0 ] =‖θ⋆(S0)‖22+[‖δt‖22∣S0] =\|θ (S_0)\|_2^2+E [\| _t\|_2^2 S_0 ] ≈‖θ⋆(S0)‖22+ηλs⋆(S0)2p+O(η2λ2). ≈\|θ (S_0)\|_2^2+ ηλ s (S_0)2p+O\! ( η^2λ^2 ). (414) Solving Generalization Radius. Under the local stationary approximation in the generalization basin, [‖θτG(S0)‖22∣S0] [\| _ _G(S_0)\|_2^2 S_0 ] ≈[‖θt‖22∣S0]. [\| _t\|_2^2 S_0 ]. (415) Therefore, [‖θτG(S0)‖22∣S0] [\| _ _G(S_0)\|_2^2 S_0 ] ≈‖θ⋆(S0)‖22+ηλs⋆(S0)2p+O(η2λ2). ≈\|θ (S_0)\|_2^2+ ηλ s (S_0)2p+O\! ( η^2λ^2 ). (416) Averaging over S0∈ΘS_0∈ _S, ρG2 _G^2 :=S0∈ΘS[[‖θτG‖22∣S0]] :=E_S_0∈ _S [E [\| _ _G\|_2^2 S_0 ] ] ≈S0∈ΘS[‖θ⋆(S0)‖22]+ηλS0∈ΘS[s⋆(S0)2p]+O(η2λ2). _S_0∈ _S [\|θ (S_0)\|_2^2 ]+ ηλE_S_0∈ _S [ s (S_0)2p ]+O\! ( η^2λ^2 ). (417) By definition, (ρG(0))2 ( _G^(0) )^2 :=S0∈ΘS[‖θ⋆(S0)‖22], :=E_S_0∈ _S [\|θ (S_0)\|_2^2 ], (418) cG c_G :=S0∈ΘS[s⋆(S0)2p]. :=E_S_0∈ _S [ s (S_0)2p ]. (419) Thus, ρG2≈(ρG(0))2+ηλcG+O(η2λ2), _G^2≈ ( _G^(0) )^2+ ηλc_G+O\! ( η^2λ^2 ), (420) which proves the claimed scaling law. Weak Batch-Size b-Dependence of cGc_G. The weak batch-size dependence of cGc_G enters through the effective scalar preconditioner s⋆(S0)s (S_0). From the definition of cGc_G, cG c_G :=S0∈ΘS[s⋆(S0)2p]. :=E_S_0∈ _S [ s (S_0)2p ]. (421) By the late-stage preconditioning identity, 1s⋆(S0) 1s (S_0) =1ptr(π(θ⋆(S0))), = 1ptr (π(θ (S_0)) ), (422) π(θ⋆(S0)) π(θ (S_0)) =diag(g⊙g¯(θ⋆(S0)))−1/2. =diag ( g g(θ (S_0)) )^-1/2. (423) In the memorization regime, g⊙g¯(θ⋆(S0)) g g(θ (S_0)) ≈1bdiag(Σ(θ⋆(S0))). ≈ 1bdiag( (θ (S_0))). (424) Therefore, π(θ⋆(S0)) π(θ (S_0)) ≈bdiag(diag(Σ(θ⋆(S0))))−1/2, ≈ b\,diag (diag( (θ (S_0))) )^-1/2, (425) 1s⋆(S0) 1s (S_0) =O(b), =O( b), (426) s⋆(S0) s (S_0) =O(1b). =O\! ( 1 b ). (427) Using the late-stage identity G(θ)G(θ)⊤≈IpG(θ)G(θ) ≈ I_p, the diffusion trace contribution is tr(G(θ)G(θ)⊤)≈ptr(G(θ)G(θ) )≈ p. Hence cG=S0∈ΘS[s⋆(S0)2p]=O(1b). c_G=E_S_0∈ _S [ s (S_0)2p ]=O\! ( 1 b ). (428) ∎ A.14 Proof: Scaling Law of Solution Transition Time 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, G( _t)G( _t) ≈ I_p, (429) and the residual approximation ℛSM(t)+ℛπ(t)≈0. _SM(t)+R_π(t)≈ 0. (430) Let g¯(θ) g(θ) be the population mean gradient and let s(θ)>0s(θ)>0 denote the scalar preconditioning factor satisfying 1s(θ)=1ptr(π(θ)). 1s(θ)= 1ptr (π(θ) ). (431) Let θt(0)(S0) _t^(0)(S_0) be the deterministic late-stage trajectory from ∂M∂ M to ∂G∂ G starting with initial state S0S_0, and let τM→G(0)(S0) _M→ G^(0)(S_0) be its deterministic transition time. For each deterministic late-stage trajectory θt(0)(S0) _t^(0)(S_0) from ∂M∂ M to ∂G∂ G, define its pathwise effective radial contraction rate by κ(S0):=1τM→G(0)(S0)∫0τM→G(0)(S0)[λs(θt(0)(S0))+θt(0)⊤(S0)(g¯(θt(0)(S0))−λθt(0)(S0))s(θt(0)(S0))‖θt(0)(S0)‖22]dt. κ(S_0):= 1 _M→ G^(0)(S_0) _0 _M→ G^(0)(S_0) [ λs( _t^(0)(S_0))+ _t^(0) (S_0) ( g( _t^(0)(S_0))-λ _t^(0)(S_0) )s( _t^(0)(S_0))\| _t^(0)(S_0)\|_2^2 ] t. (432) The scalar s¯ s is defined through the ensemble-averaged effective contraction rate λs¯:=S0∈ΘS[κ(S0)]. λ s:=E_S_0∈ _S [κ(S_0) ]. (433) Let ρM _M and ρG _G be the memorization and generalization radii. Assume the asymptotic expansions ρM2 _M^2 =(ρM(0))2+ηbcM+O((η/b)2), =( _M^(0))^2+ ηbc_M+O ((η/b)^2 ), (434) ρG2 _G^2 ≈(ρG(0))2+ηλcG+O(η2λ2), ≈ ( _G^(0) )^2+ ηλc_G+O\! ( η^2λ^2 ), (435) where cMc_M is the memorization-radius correction from Theorem 4, and (ρG(0))2 ( _G^(0) )^2 :=S0∈ΘS[‖θ⋆(S0)‖22], :=E_S_0∈ _S [\|θ (S_0)\|_2^2 ], (436) θ⋆(S0) θ (S_0) :=arg⁡minθ(S0)⁡ℒf∗(θ(S0))+λ2‖θ(S0)‖22, := _θ(S_0) \ L^*_f(θ(S_0))+ λ2\|θ(S_0)\|_2^2 \, (437) cG c_G :=S0∈ΘS[s⋆(S0)2tr(D⋆(S0))], :=E_S_0∈ _S [ s (S_0)2tr\! (D (S_0) ) ], (438) s⋆(S0) s (S_0) :=s(θ⋆(S0)),D⋆(S0):=G(θ⋆(S0))G(θ⋆(S0))⊤. :=s(θ (S_0)), D (S_0):=G(θ (S_0))G(θ (S_0)) . (439) Under the late-stage identity G(θ⋆(S0))G(θ⋆(S0))⊤≈IpG(θ (S_0))G(θ (S_0)) ≈ I_p, equation (438) reduces to cG=S0∈ΘS[s⋆(S0)2p]. c_G=E_S_0∈ _S [ s (S_0)2p ]. (440) Then, in iteration time, the expected transition time from ∂M∂ M to ∂G∂ G satisfies S0∈ΘS[τM→G]≈s¯ηλlog⁡ρM(0)ρG(0)+cτbλ+cτ(2)λ2+O(ηλ3), _S_0∈ _S [ _M→ G ]≈ sη\,λ \! _M^(0) _G^(0)+ c_τb\,λ+ c_τ^(2)λ^2+O\! ( ηλ^3 ), (441) where cτ c_τ :=s¯cM2(ρM(0))2, := s\, c_M2( _M^(0))^2, (442) cτ(2) c_τ^(2) :=−s¯cG2(ρG(0))2+s¯2(p−2)4[1(ρG(0))2−1(ρM(0))2]. :=- s\, c_G2( _G^(0))^2+ s^2(p-2)4 [ 1( _G^(0))^2- 1( _M^(0))^2 ]. (443) The constants s¯ s, cτc_τ, and cτ(2)c_τ^(2) are independent of η and λ in the considered scaling regime; a weak batch-size b-dependence enters through the scalar preconditioning scale s¯=O(1b) s=O( 1 b), and hence cτ=O(1b)c_τ=O( 1 b) and cτ(2)=O(1b)c_τ^(2)=O( 1b). Proof. Let rt:=‖θt‖2. r_t:=\| _t\|_2. (444) 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,ℛSM(t)+ℛπ(t)≈0, a(t)→ 1, G( _t)G( _t) ≈ I_p, _SM(t)+R_π(t)≈ 0, (445) the squared-radius process satisfies drt2≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)+ηp]dt+2ηrtdWt(r). r_t^2≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t )+η\,p ] t+2 η\,r_t\, W_t^(r). (446) Effective Deterministic Radius Flow. The leading deterministic radial flow is obtained from equation (446) by keeping the O(1)O(1) drift terms and dropping the O(η)O(η) Itô correction. Thus, drt2 r_t^2 ≈[−2λs(θt)rt2−2s(θt)θt⊤(g¯t−λθt)]dt. ≈ [- 2λs( _t)r_t^2- 2s( _t) _t ( g_t-λ _t ) ] t. (447) Since drt2=2rtdrt+O(η)dt, r_t^2=2r_t\, r_t+O(η) t, (448) the leading deterministic radial drift is r˙t r_t ≈−λs(θt)rt−1s(θt)rtθt⊤(g¯t−λθt) ≈- λs( _t)r_t- 1s( _t)r_t _t ( g_t-λ _t ) =−[λs(θt)+θt⊤(g¯t−λθt)s(θt)‖θt‖22]rt. =- [ λs( _t)+ _t ( g_t-λ _t )s( _t)\| _t\|_2^2 ]r_t. (449) For each deterministic late-stage trajectory θt(0)(S0) _t^(0)(S_0) from ∂M∂ M to ∂G∂ G, define its pathwise effective radial contraction rate by κ(S0):=1τM→G(0)(S0)∫0τM→G(0)(S0)[λs(θt(0)(S0))+θt(0)⊤(S0)(g¯(θt(0)(S0))−λθt(0)(S0))s(θt(0)(S0))‖θt(0)(S0)‖22]dt. κ(S_0):= 1 _M→ G^(0)(S_0) _0 _M→ G^(0)(S_0) [ λs( _t^(0)(S_0))+ _t^(0) (S_0) ( g( _t^(0)(S_0))-λ _t^(0)(S_0) )s( _t^(0)(S_0))\| _t^(0)(S_0)\|_2^2 ] t. (450) The scalar s¯ s is defined through the ensemble-averaged effective contraction rate λs¯:=S0∈ΘS[κ(S0)]. λ s:=E_S_0∈ _S [κ(S_0) ]. (451) Thus, in the averaged effective radial model, the leading deterministic late-stage radius flow is r˙t(0)=−λs¯rt(0). r_t^(0)=- λ sr_t^(0). (452) Therefore, along the averaged late-stage transition dynamics, −λs(θt)rt−1s(θt)rtθt⊤(g¯t−λθt)≈−λs¯rt. - λs( _t)r_t- 1s( _t)r_t _t ( g_t-λ _t )≈- λ sr_t. (453) Effective Radius SDE. We now pass from the squared-radius SDE to the radius SDE. Applying Itô’s lemma to rt=(rt2)1/2r_t=(r_t^2)^1/2 gives drt=12rtdrt2−18rt3d⟨r2⟩t. r_t= 12r_t r_t^2- 18r_t^3 r^2 _t. (454) From the martingale term in equation (446), d⟨r2⟩t=4ηrt2dt. r^2 _t=4η r_t^2\, t. (455) Substituting equations (446) and (455) into equation (454), and using equation (453), yields drt r_t =[−λs¯rt+ηp2rt−η12rt]dt+ηdWt(r) = [- λ sr_t+η p2r_t-η 12r_t ] t+ η\, W_t^(r) =[−λs¯rt+ηp−12rt]dt+ηdWt(r). = [- λ sr_t+η p-12r_t ] t+ η\, W_t^(r). (456) Let κ:=λs¯,ε:=η. κ:= λ s, :=η. (457) Then equation (456) becomes drt=[−κrt+ε(p−1)2rt]dt+εdWt(r). r_t= [-κ r_t+ (p-1)2r_t ] t+ \, W_t^(r). (458) Mean First-Passage Time PDE. Let T(r)T(r) be the expected continuous time for the process in equation (458), initialized at radius r, to hit the absorbing boundary rGr_G, defined as T(r):=[τG∣r0=r],r>rG. T(r):=E\! [ _G r_0=r ], r>r_G. (459) The transition starts at the memorization boundary and ends at the generalization boundary: r0 r_0 :=ρM, := _M, rG r_G :=ρG. := _G. (460) Its infinitesimal generator is ℒr[∙]=(−κr+ε(p−1)2r)ddr[∙]+ε2d2dr2[∙]. _r[ ]= (-κ r+ (p-1)2r ) r[ ]+ 2 ^2 r^2[ ]. (461) Then T(r)T(r) solves a Dirichlet problem ℒrT(r)=−1,T(rG)=0. _rT(r)=-1, T(r_G)=0. (462) Solving Absorbing Boundary Dirichlet PDE. Equivalently, ε2T′(r)+(−κr+ε(p−1)2r)T′(r)=−1. 2T (r)+ (-κ r+ (p-1)2r )T (r)=-1. (463) Writing q(r):=T′(r), q(r):=T (r), (464) we obtain q′(r)+(−2κεr+p−1r)q(r)=−2ε. q (r)+ (- 2κ r+ p-1r )q(r)=- 2 . (465) The integrating factor is I(r)=rp−1exp⁡(−κεr2). I(r)=r^p-1 \! (- κ r^2 ). (466) Using I(r)q(r)→0as r→∞, I(r)q(r)→ 0 r→∞, (467) we obtain q(r)=2ε1I(r)∫r∞I(y)dy. q(r)= 2 1I(r) _r^∞I(y) y. (468) Since T(rG)=0T(r_G)=0, for an initial radius r0>rGr_0>r_G, T(r0)=2ε∫rGr01I(x)∫x∞I(y)dydx. T(r_0)= 2 _r_G^r_0 1I(x) _x^∞I(y) y\, x. (469) Define β:=κε=λs¯η. β:= κ = λ s\,η. (470) Then I(r)=rp−1e−βr2, I(r)=r^p-1e^-β r^2, (471) and ∫x∞yp−1e−βy2dy=12β−p/2Γ(p/2,βx2), _x^∞y^p-1e^-β y^2 y= 12β^-p/2 (p/2,β x^2), (472) where Γ(a,z):=∫z∞ta−1e−tdt. (a,z):= _z^∞t^a-1e^-t t. (473) Thus, T(r0)=1εβ−p/2∫rGr0eβx2xp−1Γ(p/2,βx2)dx. T(r_0)= 1 β^-p/2 _r_G^r_0 e^β x^2x^p-1 (p/2,β x^2) x. (474) Small-Learning-Rate Regime. The learning rate η→0η→ 0 is small, so βx2=λs¯ηx2≫1. β x^2= λ s\,ηx^2 1. (475) Therefore, Γ(p/2,βx2)=(βx2)p/2−1e−βx2[1+p−22βx2+O((βx2)−2)]. (p/2,β x^2)=(β x^2)^p/2-1e^-β x^2 [1+ p-22β x^2+O\! ((β x^2)^-2 ) ]. (476) Substituting equation (476) into equation (474) gives T(r0) T(r_0) =1εβ∫rGr0[1x+p−22βx3+O((βx2)−2x−1)]dx = 1 β _r_G^r_0 [ 1x+ p-22β x^3+O\! ((β x^2)^-2x^-1 ) ] x =1κlog⁡r0rG+ε(p−2)4κ2[1rG2−1r02]+O(ε2κ3). = 1κ r_0r_G+ (p-2)4κ^2 [ 1r_G^2- 1r_0^2 ]+O\! ( ^2κ^3 ). (477) Substituting κ=λ/s¯κ=λ/ s and ε=η =η yields T(r0)=s¯λlog⁡r0rG+ηs¯2(p−2)4λ2[1rG2−1r02]+O(η2λ3). T(r_0)= sλ r_0r_G+η s^2(p-2)4λ^2 [ 1r_G^2- 1r_0^2 ]+O\! ( η^2λ^3 ). (478) Asymptotic Expansions. The transition starts at the memorization radius and ends at the generalization radius. By Theorem 4, ρM2=(ρM(0))2+ηbcM+O((η/b)2). _M^2=( _M^(0))^2+ ηbc_M+O ((η/b)^2 ). (479) Taking the square root gives r0=ρM=ρM(0)+ηbcM2ρM(0)+O((η/b)2). r_0= _M= _M^(0)+ ηb c_M2 _M^(0)+O ((η/b)^2 ). (480) Similarly, by Theorem 5, ρG2=(ρG(0))2+ηλcG+O(η2λ2), _G^2=( _G^(0))^2+ ηλc_G+O\! ( η^2λ^2 ), (481) where (ρG(0))2=S0∈ΘS[‖θ⋆(S0)‖22], ( _G^(0))^2=E_S_0∈ _S [\|θ (S_0)\|_2^2 ], (482) and θ⋆(S0)=arg⁡minθ(S0)⁡ℒf∗(θ(S0))+λ2‖θ(S0)‖22. θ (S_0)= _θ(S_0) \ L^*_f(θ(S_0))+ λ2\|θ(S_0)\|_2^2 \. (483) Therefore, rG=ρG=ρG(0)+ηλcG2ρG(0)+O(η2λ2). r_G= _G= _G^(0)+ ηλ c_G2 _G^(0)+O\! ( η^2λ^2 ). (484) Using log⁡(x+Δx)=log⁡x+Δx+O(Δx2), (x+ x)= x+ xx+O( x^2), (485) we have log⁡r0 r_0 =log⁡ρM(0)+ηbcM2(ρM(0))2+O((η/b)2), = _M^(0)+ ηb c_M2( _M^(0))^2+O ((η/b)^2 ), (486) log⁡rG r_G =log⁡ρG(0)+ηλcG2(ρG(0))2+O(η2λ2). = _G^(0)+ ηλ c_G2( _G^(0))^2+O\! ( η^2λ^2 ). (487) Hence log⁡r0rG r_0r_G =log⁡ρM(0)ρG(0)+ηbcM2(ρM(0))2−ηλcG2(ρG(0))2+O(η2λ2). = _M^(0) _G^(0)+ ηb c_M2( _M^(0))^2- ηλ c_G2( _G^(0))^2+O\! ( η^2λ^2 ). (488) In the explicit O(η)O(η) Itô correction term of equation (478), it is sufficient to use the leading asymptotics: 1rG2−1r02=1(ρG(0))2−1(ρM(0))2+O(η/b)+O(η/λ). 1r_G^2- 1r_0^2= 1( _G^(0))^2- 1( _M^(0))^2+O(η/b)+O(η/λ). (489) Substituting equations (488) and (489) into equation (478) gives TM→G T_M→ G =s¯λlog⁡ρM(0)ρG(0) = sλ _M^(0) _G^(0) +ηbλs¯cM2(ρM(0))2 + ηbλ s c_M2( _M^(0))^2 +ηλ2−s¯cG2(ρG(0))2+s¯2(p−2)4[1(ρG(0))2−1(ρM(0))2] + ηλ^2 \- s c_G2( _G^(0))^2+ s^2(p-2)4 [ 1( _G^(0))^2- 1( _M^(0))^2 ] \ +O(η2λ3). +O\! ( η^2λ^3 ). (490) Define cτ c_τ :=s¯cM2(ρM(0))2, := s\, c_M2( _M^(0))^2, cτ(2) c_τ^(2) :=−s¯cG2(ρG(0))2+s¯2(p−2)4[1(ρG(0))2−1(ρM(0))2]. :=- s\, c_G2( _G^(0))^2+ s^2(p-2)4 [ 1( _G^(0))^2- 1( _M^(0))^2 ]. (491) Then TM→G≈s¯λlog⁡ρM(0)ρG(0)+ηbλcτ+ηλ2cτ(2)+O(η2λ3). T_M→ G≈ sλ _M^(0) _G^(0)+ ηbλc_τ+ ηλ^2c_τ^(2)+O\! ( η^2λ^3 ). (492) Computing Number of Iterations. We compute the number of iterations, since the continuous-time interpolation is t=ηkt=η k: S0∈ΘS[τM→G]=TM→Gη. _S_0∈ _S[ _M→ G]= T_M→ Gη. (493) Dividing equation (492) by η gives S0∈ΘS[τM→G]≈s¯ηλlog⁡ρM(0)ρG(0)+cτbλ+cτ(2)λ2+O(ηλ3), _S_0∈ _S[ _M→ G]≈ sη\,λ _M^(0) _G^(0)+ c_τb\,λ+ c_τ^(2)λ^2+O\! ( ηλ^3 ), (494) which proves equation (441). ℓ2 _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 =0 and is determined only by the deterministic radial drift −λr/s¯-λ r/ s. Batch size enters the leading term through the scalar preconditioning scale s¯=O(1/b) s=O(1/ b), and the corrections through the memorization asymptotic ρM _M — whose O(η/b)O(η/b) correction becomes O(1/b)O(1/b) after converting to iteration time — and through cτ=O(1/b)c_τ=O(1/ b) and cτ(2)=O(1/b)c_τ^(2)=O(1/b). ∎ A.15 Additional Results: Scaling Laws of Manifold Radius ρM2 _M^2 on ℤ127Z_127 Figure 11: Scaling Law of Manifold Radius ρM2 _M^2 on ℤ127Z_127. We show the scaling law of ρM2 _M^2 with respect to the learning rate η, batch size b, and ℓ2 _2 regularization coefficient λ on the ℤ127Z_127 task. For each hyperparameter configuration, we train for ten runs. The results show that larger η/bη/b induces stronger diffusion variance, whereas λ does not affect the diffusion variance. We also overlay the theoretical fits. A.16 Additional Results: Scaling Laws of Manifold Radius ρG2 _G^2 on ℤ127Z_127 Figure 12: Scaling Law of Manifold Radius ρG2 _G^2 on ℤ127Z_127. We show the scaling law of ρG2 _G^2 with respect to the learning rate η, batch size b, and ℓ2 _2 regularization coefficient λ on the ℤ127Z_127 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. A.17 Additional Results: Scaling Laws of Solution Transition Time on ℤ127Z_127 Figure 13: Scaling laws of solution transition time on ℤ127Z_127. We show that the solution transition time τM→G _M→ G from the memorization manifold M to the generalization manifold G scales with the learning rate η, batch size b, and ℓ2 _2 regularization coefficient λ. For each hyperparameter configuration, we train for ten runs. We also overlay the theoretical fits.

Instant research discovery