Understanding Stability of Outer Nesterov Momentum on Stochastic Pseudo-Gradients
Summary
该论文研究了分布式训练中外部Nesterov动量在随机伪梯度上的稳定性,推导了随机二次模型上的精确均方稳定性判据,揭示了内层批量大小、本地步数与工作节点平均对训练稳定性的联合影响,并在神经网络上验证了理论结论。
View Cached Full Text
Cached at: 10/19/26, 07:09 AM
# Understanding Stability of Outer Nesterov Momentum on Stochastic Pseudo-Gradients Source: [https://arxiv.org/html/2610.22623](https://arxiv.org/html/2610.22623) Nishant SharmaAffiliation:Chandar Research LabAffiliation:Mila – Quebec AI InstituteAffiliation:Université de MontréalEmail:[nishant\.sharma@mila\.quebec](mailto:[email protected])Ganesh TalluriAffiliation:Arizona State UniversityEmail:[talluri\.ganeshsai@gmail\.com](mailto:[email protected])Sarath ChandarAffiliation:Chandar Research LabAffiliation:Mila – Quebec AI InstituteAffiliation:Polytechnique MontréalAffiliation:Canada CIFAR AI ChairEmail:[sarath\.chandar@mila\.quebec](mailto:[email protected])Hiroki Naganuma††thanks:This work was conducted while Hiroki Naganuma was at Mila\.Affiliation:Mila – Quebec AI InstituteEmail:[naganuma\.hiroki@mila\.quebec](mailto:[email protected]) ###### Abstract Large neural networks are commonly trained across many workers, and communication between them is a major cost\. Workers therefore take several local optimization steps before sending their net parameter displacements, orpseudo\-gradients, to an outer optimizer, which often applies Nesterov momentum\. We study how the outer momentum, the number of local steps and the averaging over workers control the training dynamics, and in particular their stability\. For random quadratics, we derive an exact mean\-square stability criterion in any dimension, which reveals that the inner batch size determines whether increasing the number of local steps eventually stabilizes or destabilizes the outer iteration\. Moreover, while averaging the pseudo\-gradients of independent workers can make larger outer momentum values admissible, the number of workers required for stability can grow exponentially with the number of local steps\. Finally, we show that the mean\-square dynamics governing stability can be probed directly at intermediate neural\-network checkpoints using shared\-minibatch perturbations\. We test our theoretical findings across networks, and measure their consequences for training loss\. Figure 1:Mean\-square stable regions in\(α,β\)\(\\alpha,\\beta\)\.For inner batch sizebband inner\-batch thresholdbcb\_\{c\}\(as defined in Section[5\.2](https://arxiv.org/html/2610.22623#S5.SS2)\), more local steps enlarge the mean\-square stable region whenb\>bcb\>b\_\{c\}\(left\), and shrink it whenb<bcb<b\_\{c\}\(middle\)\. Averaging the pseudo\-gradient overWWworkers enlarges the mean\-square stable region \(right\), but never past the mean boundary \(gray\), which averaging leaves unchanged\.## 1Introduction Deep neural networks, and large language models \(LLMs\) in particular, are now trained on many workers, often spread across compute nodes or even data centers\. In such distributed training, each worker holds a copy of the model, and gradients are exchanged at every step to keep the copies in sync\. This communication has become one of the main drivers of training costs at scale\([Douillard et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib14)\)\. A common solution is to let each worker train locally for𝑲\\boldsymbol\{K\}optimization steps and send only its displacement, thepseudo\-gradient, to an outer update\([McMahan et al\., 2017](https://arxiv.org/html/2610.22623#bib.bib36)\)\. This outer update often applies Nesterov momentum\([Nesterov, 1983](https://arxiv.org/html/2610.22623#bib.bib41)\), which has been found to accelerate training on one or many workers\([Kallusky et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib25);[Agarwala, 2026](https://arxiv.org/html/2610.22623#bib.bib1)\)and to outperform SGD and Adam as the outer optimizer\([Douillard et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib14)\)\. Concretely, at outer steptt, each of𝑾\\boldsymbol\{W\}workers starts from the shared parametersθt\\theta\_\{t\}and runsKKsteps of minibatch stochastic gradient descent \(SGD\) with batch sizebband inner learning rateη\\eta\(itsinner trajectory\)\. The outer optimizer averages the workers’ pseudo\-gradients and applies Nesterov momentum to the average, with outer learning rateα\\alphaand outer momentumβ\\beta\(Algorithm[1](https://arxiv.org/html/2610.22623#alg1)\)\. Algorithm 1Two\-phase optimization with outer Nesterov momentum1:start θ0\\theta\_\{0\}, inner rate η\\eta, inner batch size bb, outer rate α\>0\\alpha\>0, outer momentum 0≤β<10\\leq\\beta<1, number of local steps KK, number of workers WW 2: u−1←0u\_\{\-1\}\\leftarrow 0 3:for t=0,1,2,…t=0,1,2,\\ldotsdo 4:for w=1,…,Ww=1,\\ldots,Wdo⊳\\trianglerightin parallel 5:runKKinner SGD steps at rateη\\etawith batch sizebbfromθt\\theta\_\{t\}on fresh minibatches, ending atθt,K\(w\)\\theta\_\{t,K\}^\{\(w\)\} 6:endfor 7: Pt←θt−1W∑w=1Wθt,K\(w\)P\_\{t\}\\leftarrow\\theta\_\{t\}\-\\frac\{1\}\{W\}\\sum\_\{w=1\}^\{W\}\\theta\_\{t,K\}^\{\(w\)\}⊳\\trianglerightpseudo\-gradient 8: ut←βut−1\+Ptu\_\{t\}\\leftarrow\\beta u\_\{t\-1\}\+P\_\{t\}⊳\\trianglerightbuffer updated first 9: θt\+1←θt−α\(βut\+Pt\)\\theta\_\{t\+1\}\\leftarrow\\theta\_\{t\}\-\\alpha\\bigl\(\\beta u\_\{t\}\+P\_\{t\}\\bigr\)⊳\\trianglerightuses the updated buffer 10:endfor 1Innerphase2Outerstep Several familiar training methods fit this two\-phase template\. Local SGD and federated averaging\([McMahan et al\., 2017](https://arxiv.org/html/2610.22623#bib.bib36);[Stich, 2019](https://arxiv.org/html/2610.22623#bib.bib46);[Lin et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib33)\)are the simplest case,α=1\\alpha=1andβ=0\\beta=0, in which the outer step just averages the workers’ parameters\. Lookahead\([Zhang et al\., 2019](https://arxiv.org/html/2610.22623#bib.bib57)\)keepsβ=0\\beta=0but takesα<1\\alpha<1on a single worker, and SNOO\([Kallusky et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib25)\)adds the outer Nesterov step, still on a single worker\. DiLoCo\([Douillard et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib14)\)and, more recently, MuLoCo\([Thérien et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib47)\)apply the outer Nesterov step across workers using AdamW\([Kingma and Ba, 2015](https://arxiv.org/html/2610.22623#bib.bib28);[Loshchilov and Hutter, 2019](https://arxiv.org/html/2610.22623#bib.bib34)\)and Muon\([Jordan et al\., 2024](https://arxiv.org/html/2610.22623#bib.bib24)\)as inner optimizers, respectively\. SlowMo\([Wang et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib51)\)also adds momentum to the outer step, adaptive federated optimizers\([Reddi et al\., 2021](https://arxiv.org/html/2610.22623#bib.bib43)\)use an adaptive outer optimizer, semi\-synchronous variants target high\-latency links\([Naganuma et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib40)\), and asynchronous ones tolerate stragglers and hardware failures\([Douillard et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib16)\)\. These methods, which have been scaled to language models with billions of parameters\([Charles et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib8);[Douillard et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib15)\), have many interacting hyperparameters, such as the inner and outer optimizers, the outer momentum, the learning rates, the batch size and the number of local stepsKK, and today they are set by tuning\. SNOO, for example, finds that outer momentum0\.90\.9is more sensitive toKKthan smaller values, and that suboptimal outer hyperparameters can substantially degrade performance\([Kallusky et al\., 2025](https://arxiv.org/html/2610.22623#bib.bib25)\)\. Among these choices, the number of local stepsKKis special, because it influences the variance of the pseudo\-gradient \(Section[3\.2](https://arxiv.org/html/2610.22623#S3.SS2)\), and this variance grows with the error left by earlier outer steps, so it does not average out over the outer steps\. As a result, the error can shrink on average while its mean square grows without bound, so an analysis of the mean alone can wrongly predict stability, as in the region between the green and grey boundaries of Figure[1](https://arxiv.org/html/2610.22623#S0.F1)\. Moreover, the stability of this two\-phase setup is less understood than that of standard SGD\([Wu et al\., 2018](https://arxiv.org/html/2610.22623#bib.bib54);[Ma and Ying, 2021](https://arxiv.org/html/2610.22623#bib.bib35);[Granziol et al\., 2022](https://arxiv.org/html/2610.22623#bib.bib21);[Velikanov et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib49);[Ziyin et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib58);[Mulayoff and Michaeli, 2024](https://arxiv.org/html/2610.22623#bib.bib39)\), because the outer optimizer acts on a stochasticKK\-step pseudo\-gradient rather than on a single stochastic gradient\. In this work, we give a theory that explains how changing the outer momentum, the number of local steps and the averaging over workers controls the training dynamics, and in particular their mean\-square stability, which provides both a deeper mechanistic understanding of these methods and principled guidance for choosing their hyperparameters\. Our contributions are as follows\. 1. 1\.For random quadratics, we derive a necessary and sufficient condition for the mean\-square stability of outer Nesterov momentum on pseudo\-gradients fromKKstochastic inner SGD steps \(Theorem[1](https://arxiv.org/html/2610.22623#Thmtheorem1)\)\. 2. 2\.From it we derive four implications: \(i\) at a fixedα\\alpha, the noise budget can only shrink asβ\\betagrows \(Proposition[2](https://arxiv.org/html/2610.22623#Thmtheorem2)\); \(ii\) the inner batch size decides whether largeKKis eventually stable or unstable \(Proposition[3](https://arxiv.org/html/2610.22623#Thmtheorem3); Figure[1](https://arxiv.org/html/2610.22623#S0.F1)\); \(iii\) averagingWWworkers divides the variance only byWW, and the number of workers needed for stability can grow exponentially inKK\(Corollary[4](https://arxiv.org/html/2610.22623#Thmtheorem4); Figure[1](https://arxiv.org/html/2610.22623#S0.F1)\); \(iv\) for fixed total batch size budget, splitting it across workers cannot improve mean\-square stability \(Proposition[5](https://arxiv.org/html/2610.22623#Thmtheorem5)\)\. 3. 3\.We measure stability at an intermediate neural\-network checkpoint without knowing a minimizer, by tracking whether a small difference between two runs that share their minibatches grows or shrinks\. We apply this measurement at several checkpoints of an MLP and a ResNet\-20, both trained on CIFAR\-10, and of Pythia language models from 70m to 1b parameters, and test the trends that the four implications predict \(Figure[5](https://arxiv.org/html/2610.22623#S6.F5)\)\. We also measure the consequences of instability for the training loss \(Section[6\.3](https://arxiv.org/html/2610.22623#S6.SS3); Figure[5](https://arxiv.org/html/2610.22623#S6.F5)\)\. ## 2Related Work [Agarwala \(2026\)](https://arxiv.org/html/2610.22623#bib.bib1)gives a theory of two\-phase optimizers, deriving the second\-moment dynamics of Lookahead and LA\-DiLoCo on high\-dimensional linear regression, and shows in a deterministic analysis that outer Nesterov momentum accelerates convergence on an effective spectrum set by the inner optimizer\.[Topollai et al\. \(2026\)](https://arxiv.org/html/2610.22623#bib.bib48)show in language\-model pretraining that the stable range of outer learning rates and momentum values changes with the communication period, and use deterministic mode dynamics to motivate restarts of the outer momentum that widen this range\.[Khaled et al\. \(2025\)](https://arxiv.org/html/2610.22623#bib.bib26)prove convergence guarantees for local SGD with an outer learning rate, outer momentum and outer acceleration\. However, their worst\-case guarantees depend on the gradient noise only through a uniform variance bound, so two problems with the same bound can receive the same guarantee even when one is stable and the other is not \(illustrated in Appendix\)\. Mathematically, our analysis is closest to[Mulayoff and Michaeli \(2024\)](https://arxiv.org/html/2610.22623#bib.bib39), who derive an exact mean\-square threshold for minibatch SGD\. They connect the mean\-square stability of SGD to stationary covariance and expected loss\. We extend this perspective to outer momentum acting on an entire stochastic trajectory\. While their analysis is around a minimum, through shared\-minibatch perturbations we also investigate these dynamics at intermediate neural\-network checkpoints\. Their threshold continues a line of work that uses linear stability to characterize which minima SGD can reach\([Wu et al\., 2018](https://arxiv.org/html/2610.22623#bib.bib54);[Ma and Ying, 2021](https://arxiv.org/html/2610.22623#bib.bib35);[Wu et al\., 2022](https://arxiv.org/html/2610.22623#bib.bib55)\)\. Stochastic momentum is known to be sensitive to multiplicative noise\([Gupta et al\., 2024](https://arxiv.org/html/2610.22623#bib.bib22)\)\. To our knowledge, however, no exact mean\-square criterion was available for outer momentum on a stochastic pseudo\-gradient\. Our criterion is the spectral radius of a map on second moments, which is the form of mean\-square stability for jump\-linear systems\([Costa et al\., 2005](https://arxiv.org/html/2610.22623#bib.bib11)\)\. A separate line of work studies how training behaves at its stability boundary\. Under full\-batch gradient descent, the sharpness of a neural network typically rises to2/η2/\\etaand hovers there, at the edge of stability\([Cohen et al\., 2021](https://arxiv.org/html/2610.22623#bib.bib10);[Damian et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib12)\)\. Minibatch SGD, with and without momentum, trains at an analogous edge of stochastic stability\([Andreyev and Beneventano, 2026](https://arxiv.org/html/2610.22623#bib.bib4);[Liao et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib32);[Andreyev et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib5)\)\. Among two\-phase optimizers,[Khandelwal et al\. \(2026\)](https://arxiv.org/html/2610.22623#bib.bib27)observe edge\-of\-stability dynamics in federated averaging, the caseα=1\\alpha=1andβ=0\\beta=0\. Whether training with outer momentum also operates at the edge of stability is, to our knowledge, open \(Section[7](https://arxiv.org/html/2610.22623#S7)\)\. ## 3Setup Here, we make the setup from the introduction precise\. With a single worker per outer step \(averaging overWWworkers in Section[5\.3](https://arxiv.org/html/2610.22623#S5.SS3)\), the pseudo\-gradient comes fromKKsteps that start atθ\\thetaand end atθK,ω\\theta\_\{K,\\omega\}, whereω\\omegadenotes their minibatchesξ0,…,ξK−1\\xi\_\{0\},\\ldots,\\xi\_\{K\-1\}, each ofbbsamples\. Written as a function of the starting pointθ\\theta, PK,ω\(θ\)=θ−θK,ω,P\_\{K,\\omega\}\(\\theta\)=\\theta\-\\theta\_\{K,\\omega\},\(1\)we call this map fromθ\\thetato its displacement ablock, and writePωP\_\{\\omega\}for it whenKKis fixed \(Figure[2](https://arxiv.org/html/2610.22623#S3.F2)\(a\)\)\. At outer steptt, the pseudo\-gradient is this map evaluated at the current parameters,Pt=Pωt\(θt\)P\_\{t\}=P\_\{\\omega\_\{t\}\}\(\\theta\_\{t\}\), where the minibatchesωt\\omega\_\{t\}are drawn independently at every outer step, and the outer update of Algorithm[1](https://arxiv.org/html/2610.22623#alg1)is ut=βut−1\+Pt,θt\+1=θt−α\(βut\+Pt\),u\_\{t\}=\\beta u\_\{t\-1\}\+P\_\{t\},\\qquad\\theta\_\{t\+1\}=\\theta\_\{t\}\-\\alpha\\,\(\\beta u\_\{t\}\+P\_\{t\}\),\(2\)with the momentum buffer started atu−1=0u\_\{\-1\}=0\. Throughout,α\>0\\alpha\>0and0≤β<10\\leq\\beta<1\. Figure 2:The pseudo\-gradient ofKKinner SGD steps, and the block operator\.\(a\) Fromθt\\theta\_\{t\}the inner trajectory takesKKSGD steps on minibatchesξ0,…,ξK−1\\xi\_\{0\},\\ldots,\\xi\_\{K\-1\}and ends atθK,ω\\theta\_\{K,\\omega\}\. The pseudo\-gradientPω\(θt\)P\_\{\\omega\}\(\\theta\_\{t\}\)is the displacement fromθK,ω\\theta\_\{K,\\omega\}toθt\\theta\_\{t\}\. The outer step movesθt\\theta\_\{t\}by−αPω\(θt\)\-\\alpha P\_\{\\omega\}\(\\theta\_\{t\}\)toθt\+1\\theta\_\{t\+1\}\. \(b\) The block operatorDK,ωD\_\{K,\\omega\}maps a displacementδ\\deltaof the starting point to the changeDK,ωδD\_\{K,\\omega\}\\deltaof the pseudo\-gradient\. Eachgreenarrow is the changeDK,ωδD\_\{K,\\omega\}\\deltafor one draw of the minibatchesω\\omega\. The draws have meanD¯Kδ\\bar\{D\}\_\{K\}\\delta, and theshadedellipse marks their covariance\. Mean\-square stability depends onDK,ωD\_\{K,\\omega\}only throughD¯K\\bar\{D\}\_\{K\}and the second moment𝖲K\\mathsf\{S\}\_\{K\}of Section[3\.2](https://arxiv.org/html/2610.22623#S3.SS2)\.Initially, through Section[3\.1](https://arxiv.org/html/2610.22623#S3.SS1)we will assume that the minibatch losses share a minimizerθ∗\\theta^\{\*\}\(aninterpolating minimumin the terminology of[Mulayoff and Michaeli 2024](https://arxiv.org/html/2610.22623#bib.bib39)\) and as a part of the stability analysis, we track the erroret=θt−θ∗e\_\{t\}=\\theta\_\{t\}\-\\theta^\{\*\}\. With a shared minimizer the error obeys a homogeneous recurrence, so stability means that the error vanishes\. Section[6\.1](https://arxiv.org/html/2610.22623#S6.SS1)drops this assumption, which makes the minimizer of the full loss aregular minimumin the terminology of[Mulayoff and Michaeli \(2024\)](https://arxiv.org/html/2610.22623#bib.bib39)\. The error then no longer vanishes, but, importantly, the same stability criterion still decides whether it stays bounded\. ### 3\.1Block operator and stability We model the minibatch losses as quadratics with a common minimizer, ℓξ\(θ\)=12\(θ−θ∗\)⊤Hξ\(θ−θ∗\),Hξ=Hξ⊤⪰0,\\begin\{gathered\}\\ell\_\{\\xi\}\(\\theta\)=\\tfrac\{1\}\{2\}\(\\theta\-\\theta^\{\*\}\)^\{\\top\}H\_\{\\xi\}\(\\theta\-\\theta^\{\*\}\),\\\\ H\_\{\\xi\}=H\_\{\\xi\}^\{\\top\}\\succeq 0,\\end\{gathered\}\(3\)whereHξH\_\{\\xi\}is the Hessian of the minibatch lossℓξ\\ell\_\{\\xi\}\. Since an SGD step at rateη\\etasends the erroreeto\(I−ηHξ\)e\(I\-\\eta H\_\{\\xi\}\)e, afterKKsteps the error is multiplied byMK,ω=QξK−1⋯Qξ0M\_\{K,\\omega\}=Q\_\{\\xi\_\{K\-1\}\}\\cdots Q\_\{\\xi\_\{0\}\}withQξ=I−ηHξQ\_\{\\xi\}=I\-\\eta H\_\{\\xi\}, so the pseudo\-gradient is linear in the error, PK,ω\(θ\)=DK,ω\(θ−θ∗\),DK,ω=I−MK,ω\.\\begin\{gathered\}P\_\{K,\\omega\}\(\\theta\)=D\_\{K,\\omega\}\(\\theta\-\\theta^\{\*\}\),\\\\ D\_\{K,\\omega\}=I\-M\_\{K,\\omega\}\.\\end\{gathered\}\(4\)We call the mapDK,ωD\_\{K,\\omega\}theblock operator\. In other words, the inner trajectory replaces the curvature of the loss with the block operator, on which the outer Nesterov step acts\. In the deterministic setting its spectrum is the effective spectrum of[Agarwala \(2026\)](https://arxiv.org/html/2610.22623#bib.bib1)\. Note that the block operator is random, because it depends on the minibatchesω\\omega, although not onθ\\theta, and forK≥2K\\geq 2it is in general not symmetric, so the pseudo\-gradient need not be the gradient of any loss\. Since it does not depend onθ\\theta, displacing the starting point byδ\\deltachanges the pseudo\-gradient by exactlyDK,ωδD\_\{K,\\omega\}\\delta\(Figure[2](https://arxiv.org/html/2610.22623#S3.F2)\(b\)\), which is the reading that Section[6\.1](https://arxiv.org/html/2610.22623#S6.SS1)carries over to networks\. Substituting \([4](https://arxiv.org/html/2610.22623#S3.E4)\) into \([2](https://arxiv.org/html/2610.22623#S3.E2)\) gives linear recurrence forzt=\(et,ut−1\)z\_\{t\}=\(e\_\{t\},u\_\{t\-1\}\), the error and the momentum buffer that enter steptt, zt\+1=𝖳\(DK,ωt\)zt,𝖳\(D\)=\[I−α\(1\+β\)D−αβ2IDβI\]\.\\begin\{gathered\}z\_\{t\+1\}=\\mathsf\{T\}\(D\_\{K,\\omega\_\{t\}\}\)\\,z\_\{t\},\\\\ \\mathsf\{T\}\(D\)=\\begin\{bmatrix\}I\-\\alpha\(1\+\\beta\)D&\-\\alpha\\beta^\{2\}I\\\\ D&\\beta I\\end\{bmatrix\}\.\\end\{gathered\}\(5\) ###### Assumption 1\(i\.i\.d\. block operators\) The block operatorsDt:=DK,ωtD\_\{t\}:=D\_\{K,\\omega\_\{t\}\},t≥0t\\geq 0, of different outer steps are jointly independent and identically distributed \(i\.i\.d\.\), and independent of the initial statez0z\_\{0\}\. Both have finite second moments,𝔼‖Dt‖F2<∞\\mathbb\{E\}\\left\\lVert D\_\{t\}\\right\\rVert\_\{F\}^\{2\}<\\inftyand𝔼‖z0‖2<∞\\mathbb\{E\}\\left\\lVert z\_\{0\}\\right\\rVert^\{2\}<\\infty\. In the model, the assumption holds whenever every outer step draws itsKKminibatches afresh\.111Independence of the minibatches within a block is used for the block moments of Section[3\.2](https://arxiv.org/html/2610.22623#S3.SS2)\. Propositions[3](https://arxiv.org/html/2610.22623#Thmtheorem3)and[5](https://arxiv.org/html/2610.22623#Thmtheorem5)and Corollary[4](https://arxiv.org/html/2610.22623#Thmtheorem4)\(b\) use these moments\. ###### Definition 1\(Stability\) The recurrence \([5](https://arxiv.org/html/2610.22623#S3.E5)\) is*mean stable*if𝔼\[zt\]→0\\mathbb\{E\}\[z\_\{t\}\]\\to 0, and*mean\-square stable*if𝔼‖zt‖2→0\\mathbb\{E\}\\left\\lVert z\_\{t\}\\right\\rVert^\{2\}\\to 0, for every initial state admitted by Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1)\. As one would expect, mean\-square stability implies mean stability, because‖𝔼\[zt\]‖≤\(𝔼‖zt‖2\)1/2\\left\\lVert\\mathbb\{E\}\[z\_\{t\}\]\\right\\rVert\\leq\(\\mathbb\{E\}\\left\\lVert z\_\{t\}\\right\\rVert^\{2\}\)^\{1/2\}, but not conversely \(Figure[6](https://arxiv.org/html/2610.22623#A2.F6)in Appendix[B](https://arxiv.org/html/2610.22623#A2)gives an example\)\.222If everyHξH\_\{\\xi\}vanishes on a direction, every block leaves it unchanged, and an error along it neither grows nor decays\. We therefore carry out the analysis on the orthogonal complement of such common flat directions\. ### 3\.2Block moments and howKKandbbenter them Define two averages of the block operator, its meanD¯K=𝔼\[DK,ω\]\\bar\{D\}\_\{K\}=\\mathbb\{E\}\[D\_\{K,\\omega\}\]and its second moment𝖲K=𝔼\[DK,ω⊗DK,ω\]\\mathsf\{S\}\_\{K\}=\\mathbb\{E\}\[D\_\{K,\\omega\}\\otimes D\_\{K,\\omega\}\], with⊗\\otimesthe Kronecker product\. We will refer to them as theblock moments\. As Section[4](https://arxiv.org/html/2610.22623#S4)shows, the stability criterion depends on the block operator only through these block moments\. In the model, both are built from the mean and the second moment of a single inner step \(derivation in Appendix\)\. If theKKone\-step maps of a block are i\.i\.d\. withQ¯=𝔼\[Qξ\]\\bar\{Q\}=\\mathbb\{E\}\[Q\_\{\\xi\}\]and𝒬=𝔼\[Qξ⊗Qξ\]\\mathcal\{Q\}=\\mathbb\{E\}\[Q\_\{\\xi\}\\otimes Q\_\{\\xi\}\], then D¯K=I−Q¯K,𝖲K=I⊗I−I⊗Q¯K−Q¯K⊗I\+𝒬K\.\\begin\{gathered\}\\bar\{D\}\_\{K\}=I\-\\bar\{Q\}^\{K\},\\\\ \\mathsf\{S\}\_\{K\}=I\\otimes I\-I\\otimes\\bar\{Q\}^\{K\}\-\\bar\{Q\}^\{K\}\\otimes I\+\\mathcal\{Q\}^\{K\}\.\\end\{gathered\}\(6\)Since\(Qξ⊗Qξ\)vec\(Y\)=vec\(QξYQξ\)\(Q\_\{\\xi\}\\otimes Q\_\{\\xi\}\)\\operatorname\{vec\}\(Y\)=\\operatorname\{vec\}\(Q\_\{\\xi\}YQ\_\{\\xi\}\), the matrix𝒬\\mathcal\{Q\}acts on a symmetric matrixYYexactly as the mapY↦𝔼\[QξYQξ\]Y\\mapsto\\mathbb\{E\}\[Q\_\{\\xi\}YQ\_\{\\xi\}\]does, and we use both forms\. Equation \([6](https://arxiv.org/html/2610.22623#S3.E6)\) is precisely where the number of local steps enters\. The mean block operator is governed byQ¯K\\bar\{Q\}^\{K\}alone, whereas the second moment also contains𝒬K\\mathcal\{Q\}^\{K\}, and the two powers vary at different rates inKK, so an analysis using onlyD¯K\\bar\{D\}\_\{K\}does not account for how the noise changes withKK\. #### The batch size\. Forbbsamples drawn with replacement from sample curvatureshih\_\{i\}with meanH¯\\bar\{H\}, so thatHξH\_\{\\xi\}is the average of thebbsampledhih\_\{i\}, the one\-step meanQ¯=I−ηH¯\\bar\{Q\}=I\-\\eta\\bar\{H\}does not depend onbband the one\-step second\-moment map splits as 𝒬b=𝒬∞\+1b𝒞,𝒬∞\(Y\)=Q¯YQ¯,𝒞\(Y\)=η2𝔼i\[\(hi−H¯\)Y\(hi−H¯\)\],\\begin\{gathered\}\\mathcal\{Q\}\_\{b\}=\\mathcal\{Q\}\_\{\\infty\}\+\\tfrac\{1\}\{b\}\\,\\mathcal\{C\},\\qquad\\mathcal\{Q\}\_\{\\infty\}\(Y\)=\\bar\{Q\}Y\\bar\{Q\},\\\\ \\mathcal\{C\}\(Y\)=\\eta^\{2\}\\,\\mathbb\{E\}\_\{i\}\\bigl\[\(h\_\{i\}\-\\bar\{H\}\)\\,Y\\,\(h\_\{i\}\-\\bar\{H\}\)\\bigr\],\\end\{gathered\}\(7\)with𝒬∞\\mathcal\{Q\}\_\{\\infty\}the large\-batch limit and𝒞\\mathcal\{C\}positive333A linear map on symmetric matrices is*positive*if it preserves positive semidefiniteness, andℳ⪯𝒩\\mathcal\{M\}\\preceq\\mathcal\{N\}means that𝒩−ℳ\\mathcal\{N\}\-\\mathcal\{M\}is positive\.\(derivation in Appendix\)\. The batch therefore enters only through the noise term, as the factor1/b1/b\. ## 4The Stability Boundary ### 4\.1Second\-moment map and the block noise Under Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1), the first two moments ofzt\+1z\_\{t\+1\}depend only on those ofztz\_\{t\}\. WithDKD\_\{K\}drawn from the common distribution of theDtD\_\{t\}, define thesecond\-moment mapon symmetric matrices ℒK\(X\)=𝔼\[𝖳\(DK\)X𝖳\(DK\)⊤\]\.\\mathcal\{L\}\_\{K\}\(X\)=\\mathbb\{E\}\\bigl\[\\mathsf\{T\}\(D\_\{K\}\)\\,X\\,\\mathsf\{T\}\(D\_\{K\}\)^\{\\top\}\\bigr\]\.\(8\)Then, withA=𝖳\(D¯K\)A=\\mathsf\{T\}\(\\bar\{D\}\_\{K\}\)themean transitionandΣt=𝔼\[ztzt⊤\]\\Sigma\_\{t\}=\\mathbb\{E\}\[z\_\{t\}z\_\{t\}^\{\\top\}\]the second\-moment matrix ofztz\_\{t\}, 𝔼\[zt\+1\]=A𝔼\[zt\],Σt\+1=ℒK\(Σt\),\\mathbb\{E\}\[z\_\{t\+1\}\]=A\\,\\mathbb\{E\}\[z\_\{t\}\],\\qquad\\Sigma\_\{t\+1\}=\\mathcal\{L\}\_\{K\}\(\\Sigma\_\{t\}\),\(9\)henceΣt=ℒKt\(Σ0\)\\Sigma\_\{t\}=\\mathcal\{L\}\_\{K\}^\{t\}\(\\Sigma\_\{0\}\)\(proof in Appendix\)\. Crucially, the second\-moment map splits as ℒK\(X\)=AXA⊤\+𝖡𝒱K\(𝖤X𝖤⊤\)𝖡⊤,𝒱K\(Y\)=𝔼\[\(DK−D¯K\)Y\(DK−D¯K\)⊤\],𝖡=\[−α\(1\+β\)II\],𝖤=\[I0\],\\begin\{gathered\}\\mathcal\{L\}\_\{K\}\(X\)=AXA^\{\\top\}\+\\mathsf\{B\}\\,\\mathcal\{V\}\_\{K\}\(\\mathsf\{E\}X\\mathsf\{E\}^\{\\top\}\)\\,\\mathsf\{B\}^\{\\top\},\\\\ \\mathcal\{V\}\_\{K\}\(Y\)=\\mathbb\{E\}\\bigl\[\(D\_\{K\}\-\\bar\{D\}\_\{K\}\)\\,Y\\,\(D\_\{K\}\-\\bar\{D\}\_\{K\}\)^\{\\top\}\\bigr\],\\\\ \\mathsf\{B\}=\\begin\{bmatrix\}\-\\alpha\(1\+\\beta\)I\\\\ I\\end\{bmatrix\},\\qquad\\mathsf\{E\}=\\begin\{bmatrix\}I&0\\end\{bmatrix\},\\end\{gathered\}\(10\)with𝒱K\\mathcal\{V\}\_\{K\}theblock noise, the covariance map of the block operator\. A mean\-stability analysis replacesDKD\_\{K\}by its meanD¯K\\bar\{D\}\_\{K\}, so it keeps only the first term of \([10](https://arxiv.org/html/2610.22623#S4.E10)\), which propagates the second moment through the mean transitionAA\. The second term adds, at every outer step, the block noise applied to the current error\. Because this term is a positive map, it can only raise thespectral radiusρ\\rho, soρ\(ℒK\)≥ρ\(A\)2\\rho\(\\mathcal\{L\}\_\{K\}\)\\geq\\rho\(A\)^\{2\}, with equality when there is no block noise\. A mean\-stability analysis can therefore be wrong only by predicting stability for an unstable setting \(example in Appendix\)\. ### 4\.2The mean condition and the noise budget By \([10](https://arxiv.org/html/2610.22623#S4.E10)\),ρ\(ℒK\)\\rho\(\\mathcal\{L\}\_\{K\}\)mixes the mean dynamics and the block noise in a single number\. In Theorem[1](https://arxiv.org/html/2610.22623#Thmtheorem1)we separate them into a condition on the mean transition and a budget for the noise\. Whenρ\(A\)<1\\rho\(A\)<1, let 𝒢α,β\(Y\)=∑t≥0𝖤At𝖡Y𝖡⊤\(A⊤\)t𝖤⊤\\mathcal\{G\}\_\{\\alpha,\\beta\}\(Y\)=\\sum\_\{t\\geq 0\}\\mathsf\{E\}A^\{t\}\\mathsf\{B\}\\;Y\\;\\mathsf\{B\}^\{\\top\}\(A^\{\\top\}\)^\{t\}\\mathsf\{E\}^\{\\top\}\(11\)be thenoise gain\. If the block noise injects a covarianceYYat one outer step,𝒢α,β\(Y\)\\mathcal\{G\}\_\{\\alpha,\\beta\}\(Y\)is the second moment of the error that the mean dynamics build from it over all later steps\. It depends on the distribution of the block operator only through its meanD¯K\\bar\{D\}\_\{K\}\. The error in turn feeds the noise of the next block through𝒱K\\mathcal\{V\}\_\{K\}, so we callρ\(𝒱K∘𝒢α,β\)\\rho\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta\}\)theloop gain\. For minibatch SGD,[Agarwala and Pennington \(2024\)](https://arxiv.org/html/2610.22623#bib.bib2)split stability in the same way, into a condition on the deterministic dynamics and a noise kernel norm below one, under a high\-dimensional approximation of the dynamics\. Theorem[1](https://arxiv.org/html/2610.22623#Thmtheorem1)makes this split exact forKK\-step blocks with outer momentum\. We call the reciprocal of the loop gain,1/ρ\(𝒱K∘𝒢α,β\)1/\\rho\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta\}\), thenoise budgetof\(α,β\)\(\\alpha,\\beta\)\. It is the factor by which the block noise could grow before the criterion fails, so the setting is stable when the budget exceeds one, and a larger noise gain means a smaller budget\. ###### Theorem 1\(Exact stability criterion and noise budget\) Under Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1):\(i\)the recurrence \([5](https://arxiv.org/html/2610.22623#S3.E5)\) is mean\-square stable if and only ifρ\(ℒK\)<1\\rho\(\\mathcal\{L\}\_\{K\}\)<1, equivalentlyρ\(𝔼\[𝖳\(DK\)⊗𝖳\(DK\)\]\)<1\\rho\(\\mathbb\{E\}\[\\mathsf\{T\}\(D\_\{K\}\)\\otimes\\mathsf\{T\}\(D\_\{K\}\)\]\)<1;\(ii\)ρ\(ℒK\)<1\\rho\(\\mathcal\{L\}\_\{K\}\)<1if and only ifρ\(A\)<1\\rho\(A\)<1andρ\(𝒱K∘𝒢α,β\)<1;\\rho\\bigl\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta\}\\bigr\)<1;\(12\)\(iii\)ρ\(A\)<1\\rho\(A\)<1if and only if, for every eigenvalueλ\\lambdaofD¯K\\bar\{D\}\_\{K\}, the2×22\\times 2matrix𝖳\(λ\)\\mathsf\{T\}\(\\lambda\)of \([5](https://arxiv.org/html/2610.22623#S3.E5)\) has spectral radius below one; when the eigenvalues ofD¯K\\bar\{D\}\_\{K\}are real it reads0<αλ<2\(1\+β\)1\+2β\.0<\\alpha\\lambda<\\frac\{2\(1\+\\beta\)\}\{1\+2\\beta\}\.\(13\) Mean stability is therefore necessary but not sufficient, because the block noise must also fit within the noise budget that the mean dynamics set\. Note that the moment recurrences \([9](https://arxiv.org/html/2610.22623#S4.E9)\), the split \([10](https://arxiv.org/html/2610.22623#S4.E10)\) and Theorem[1](https://arxiv.org/html/2610.22623#Thmtheorem1)use nothing aboutDKD\_\{K\}beyond Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1), so they hold for any i\.i\.d\. block operators, of which the quadratic model is one case\. Part \(i\) still holds if we consider only starts with an empty buffer,u−1=0u\_\{\-1\}=0, as in the algorithm \(proof in Appendix\)\. In one dimension the criterion has a closed form \(Appendix\)\. #### Small outer steps\. Under the mean condition and for a symmetricD¯K\\bar\{D\}\_\{K\}, the criterion withWWworkers reads, to leading order inα\\alpha, ακK\(1−β\)W<1,κK=ρ\(𝒱K∘𝒮−1\),𝒮\(X\)=D¯KX\+XD¯K,\\begin\{gathered\}\\frac\{\\alpha\\,\\kappa\_\{K\}\}\{\(1\-\\beta\)\\,W\}<1,\\\\ \\kappa\_\{K\}=\\rho\\bigl\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{S\}^\{\-1\}\\bigr\),\\qquad\\mathcal\{S\}\(X\)=\\bar\{D\}\_\{K\}X\+X\\bar\{D\}\_\{K\},\\end\{gathered\}\(14\)which has the form of the noise scale of heavy\-ball SGD\([Smith and Le, 2018](https://arxiv.org/html/2610.22623#bib.bib44)\), withα\\alphaas the learning rate,WWas the batch size andκK\\kappa\_\{K\}as the training\-set size \(Appendix\)\. AtK=1K=1and\(α,β\)=\(1,0\)\(\\alpha,\\beta\)=\(1,0\)the criterion reduces to that of[Mulayoff and Michaeli \(2024\)](https://arxiv.org/html/2610.22623#bib.bib39)\. At\(α,β\)=\(1,0\)\(\\alpha,\\beta\)=\(1,0\)and anyKKthe outer step continues the inner trajectory, and the criterion becomes the mean\-square stability of plain SGD observed everyKKsteps\. Figure 3:Mechanism behind each implication, on the model of Figure[1](https://arxiv.org/html/2610.22623#S0.F1)\.\(i\) The noise budget shrinks asβ\\betagrows \(Proposition[2](https://arxiv.org/html/2610.22623#Thmtheorem2)\)\. \(ii\) The block noise grows withKKbelowbcb\_\{c\}and dies out above it \(Proposition[3](https://arxiv.org/html/2610.22623#Thmtheorem3)\)\. \(iii\) Averaging divides the loop gain byWW\(Corollary[4](https://arxiv.org/html/2610.22623#Thmtheorem4)\)\. \(iv\) Splitting a fixed total batch raisesρ\(ℒK,W\)\\rho\(\\mathcal\{L\}\_\{K,W\}\)\(Proposition[5](https://arxiv.org/html/2610.22623#Thmtheorem5)\)\.Greenmarks stable settings andredmarks unstable ones\. ## 5Implications forβ\\beta,KK&WW ### 5\.1Outer momentum shrinks the noise budget While the momentum buffer can accelerate training by summing past pseudo\-gradients, it also accumulates their block noise\. Proposition[2](https://arxiv.org/html/2610.22623#Thmtheorem2)shows how stability depends onβ\\betaat fixedα\\alpha, for any distribution of the block operator and in any dimension \(Figure[3](https://arxiv.org/html/2610.22623#S4.F3)\(i\)\)\. ###### Proposition 2\(Momentum monotonicity\) Fixα\\alphaand the distribution of the block operator, and letβ≤β′\\beta\\leq\\beta^\{\\prime\}with\(α,β′\)\(\\alpha,\\beta^\{\\prime\}\)mean stable\. Then\(α,β\)\(\\alpha,\\beta\)is mean stable,𝒢α,β⪯𝒢α,β′\\mathcal\{G\}\_\{\\alpha,\\beta\}\\preceq\\mathcal\{G\}\_\{\\alpha,\\beta^\{\\prime\}\}, andρ\(𝒱K∘𝒢α,β\)≤ρ\(𝒱K∘𝒢α,β′\)\\rho\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta\}\)\\leq\\rho\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta^\{\\prime\}\}\)for every block noise𝒱K\\mathcal\{V\}\_\{K\}\. In particular, mean\-square stability at\(α,β′\)\(\\alpha,\\beta^\{\\prime\}\)implies mean\-square stability at\(α,β\)\(\\alpha,\\beta\)\. Consequently, if a setting is unstable because of its block noise, loweringβ\\betacan only help, and so can adding workers \(Section[5\.3](https://arxiv.org/html/2610.22623#S5.SS3)\)\. Note that the noise budget is monotone inβ\\beta, whereasρ\(ℒK\)\\rho\(\\mathcal\{L\}\_\{K\}\)itself need not be \(Figurein Appendix\)\. ### 5\.2Local steps and the inner batch size Practitioners prefer the number of local steps to be as large as stability allows, since longer blocks reduce the frequency of communication\([Stich, 2019](https://arxiv.org/html/2610.22623#bib.bib46);[Wang and Joshi, 2019](https://arxiv.org/html/2610.22623#bib.bib50);[Douillard et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib14);[Topollai et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib48)\)\. By \([6](https://arxiv.org/html/2610.22623#S3.E6)\) the block noise is𝒱K=𝒬K−𝒬∞K\\mathcal\{V\}\_\{K\}=\\mathcal\{Q\}^\{K\}\-\\mathcal\{Q\}\_\{\\infty\}^\{K\}, with𝒬=𝒬b\\mathcal\{Q\}=\\mathcal\{Q\}\_\{b\}and𝒬∞\\mathcal\{Q\}\_\{\\infty\}as in \([7](https://arxiv.org/html/2610.22623#S3.E7)\), while the mean block operator saturates,D¯K→I\\bar\{D\}\_\{K\}\\to Iwhenρ\(Q¯\)<1\\rho\(\\bar\{Q\}\)<1\. Whether the block noise dies out or grows withKKis therefore decided byρ\(𝒬\)\\rho\(\\mathcal\{Q\}\)alone\. Since𝒬k\\mathcal\{Q\}^\{k\}propagates the second moment of a displacement throughkkSGD steps,ρ\(𝒬\)<1\\rho\(\\mathcal\{Q\}\)<1is mean\-square stability of the inner steps themselves\([Mulayoff and Michaeli, 2024](https://arxiv.org/html/2610.22623#bib.bib39)\)\. ###### Proposition 3\(Long blocks and the inner\-batch threshold\) Let the inner steps of a block be i\.i\.d\., letρ\(Q¯\)<1\\rho\(\\bar\{Q\}\)<1, and fix\(α,β\)\(\\alpha,\\beta\)\.\(i\)Ifρ\(𝒬\)<1\\rho\(\\mathcal\{Q\}\)<1and\(α,β\)\(\\alpha,\\beta\)satisfies the mean condition \([13](https://arxiv.org/html/2610.22623#S4.E13)\) atλ=1\\lambda=1, thenρ\(ℒK\)→ρ\(𝖳\(I\)\)2<1\\rho\(\\mathcal\{L\}\_\{K\}\)\\to\\rho\(\\mathsf\{T\}\(I\)\)^\{2\}<1, and the recurrence is mean\-square stable for all largeKK\.\(ii\)Ifρ\(𝒬\)\>1\\rho\(\\mathcal\{Q\}\)\>1, thenρ\(ℒK\)≥γ2\(ρ\(𝒬\)K−ρ\(Q¯\)2K\)→∞\\rho\(\\mathcal\{L\}\_\{K\}\)\\geq\\gamma^\{2\}\\bigl\(\\rho\(\\mathcal\{Q\}\)^\{K\}\-\\rho\(\\bar\{Q\}\)^\{2K\}\\bigr\)\\to\\inftywithγ=α\(1\+β\)\\gamma=\\alpha\(1\+\\beta\), and the recurrence is mean\-square unstable for all largeKKat every outer rate and momentum\.\(iii\)With minibatches of sizebbdrawn with replacement as in \([7](https://arxiv.org/html/2610.22623#S3.E7)\),ρ\(𝒬b\)<1⟺b\>bc:=ρ\(𝒞∘\(ℐ−𝒬∞\)−1\),\\begin\{split\}&\\rho\(\\mathcal\{Q\}\_\{b\}\)<1\\\\ &\\quad\\Longleftrightarrow\\quad b\>b\_\{c\}:=\\rho\\bigl\(\\mathcal\{C\}\\circ\(\\mathcal\{I\}\-\\mathcal\{Q\}\_\{\\infty\}\)^\{\-1\}\\bigr\),\\end\{split\}\(15\)so \(i\) applies whenb\>bcb\>b\_\{c\}and \(ii\) whenb<bcb<b\_\{c\}\. That is, in case \(i\) the blocks behave in the limit like one that lands exactly on the minimizer,DK=ID\_\{K\}=I\. Figure[3](https://arxiv.org/html/2610.22623#S4.F3)\(ii\) shows the block noise in the two branches, and the left and middle panels of Figure[1](https://arxiv.org/html/2610.22623#S0.F1)show their stable regions\.[Mulayoff and Michaeli \(2024\)](https://arxiv.org/html/2610.22623#bib.bib39)showed that larger inner batches can only help the inner steps, and we show that the thresholdbcb\_\{c\}grows linearly withη\\etaat small inner rates \(proof in Appendix\)\. Given a mean\-stable outer step, whether long enough blocks are stable is thus decided by the inner optimizer alone, so to communicate less often one should raise the per\-worker batch size above the thresholdbcb\_\{c\}or lower the inner rate\. ### 5\.3Averaging over workers Besides the inner batch size, averaging over workers also reduces the block noise, because averaging independent pseudo\-gradients keeps their mean and shrinks only their spread around it \(Figure[3](https://arxiv.org/html/2610.22623#S4.F3)\(iii\)\)\. ###### Corollary 4\(Averaging over workers\) Let the pseudo\-gradient of each outer step be the average overWWworkers whose inner trajectories start from the same point, are independent given that point, and have block operators as in Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1)\.\(a\)The averaged block operator has meanD¯K\\bar\{D\}\_\{K\}and block noise𝒱K/W\\mathcal\{V\}\_\{K\}/W, andρ\(ℒK,W\)\\rho\(\\mathcal\{L\}\_\{K,W\}\)is non\-increasing inWW\. The averaged recurrence is mean\-square stable if and only if the mean condition holds andW\>ρ\(𝒱K∘𝒢α,β\)W\>\\rho\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta\}\), soW⋆=⌊ρ\(𝒱K∘𝒢α,β\)⌋\+1W^\{\\star\}=\\lfloor\\rho\(\\mathcal\{V\}\_\{K\}\\circ\\mathcal\{G\}\_\{\\alpha,\\beta\}\)\\rfloor\+1workers are needed and enough\. If the mean condition fails, no number of workers suffices\.\(b\)Under the hypotheses of Proposition[3](https://arxiv.org/html/2610.22623#Thmtheorem3)\(ii\), mean\-square stability withWWworkers requires W\>γ2\(ρ\(𝒬\)K−ρ\(Q¯\)2K\)W\>\\gamma^\{2\}\\bigl\(\\rho\(\\mathcal\{Q\}\)^\{K\}\-\\rho\(\\bar\{Q\}\)^\{2K\}\\bigr\)withγ=α\(1\+β\)\\gamma=\\alpha\(1\+\\beta\), a bound that grows exponentially inKK\. By Proposition[2](https://arxiv.org/html/2610.22623#Thmtheorem2)the countW⋆W^\{\\star\}never falls asβ\\betarises, so the largest admissible momentum at a fixed rate is non\-decreasing inWW\.[Agarwala \(2026\)](https://arxiv.org/html/2610.22623#bib.bib1)analyzes averaging over workers in the zero\-momentum case, and Tablein Appendixchecks the countW⋆W^\{\\star\}against the second\-moment map\. ### 5\.4Fixed total batch More workers at a fixed per\-worker batch size cost more gradient computation, so for a fairer comparison we hold the total batchW×bW\\times bfixed and ask whether splitting it across workers helps \(Figure[3](https://arxiv.org/html/2610.22623#S4.F3)\(iv\)\)\. ###### Proposition 5\(Fixed total batch\) Let the inner steps of a block be i\.i\.d\. with minibatches drawn with replacement as in \([7](https://arxiv.org/html/2610.22623#S3.E7)\), and letWWdividebb\. At fixed\(α,β,K\)\(\\alpha,\\beta,K\),WWworkers with batchb/Wb/Weach have the same mean block operator as one worker with batchbband at least as much block noise,1W𝒱K\(b/W\)⪰𝒱K\(b\)\\tfrac\{1\}\{W\}\\mathcal\{V\}\_\{K\}\(b/W\)\\succeq\\mathcal\{V\}\_\{K\}\(b\), with equality atK=1K=1\. In particular,ρ\(ℒK,W\(b/W\)\)≥ρ\(ℒK\(b\)\)\\rho\\bigl\(\\mathcal\{L\}\_\{K,W\}\(b/W\)\\bigr\)\\geq\\rho\\bigl\(\\mathcal\{L\}\_\{K\}\(b\)\\bigr\)\. Expanding𝒬bK\\mathcal\{Q\}\_\{b\}^\{K\}by \([7](https://arxiv.org/html/2610.22623#S3.E7)\) writes the block noise as a sum over powers of1/b1/bwith positive coefficients, so splitting is nearly free while per\-worker batches stay large \(proof in Appendix\), and since workers run in parallel, splitting the batch makes each outer step faster\.[Agarwala \(2026\)](https://arxiv.org/html/2610.22623#bib.bib1)makes the same comparison in the isotropic case without momentum\. ## 6Testing the Implications on Neural Networks We now test the implications of Section[5](https://arxiv.org/html/2610.22623#S5)on neural networks\. A network’s minibatch losses have different minimizers, and its training rarely runs to convergence\. We therefore linearize the pseudo\-gradient of each block at a checkpoint and measure instability with a quantity that makes no reference to a minimizer\. ### 6\.1The block Jacobian To connect our model to a network, we begin by asking how the pseudo\-gradient of a block changes in response to a parameter perturbationδ\\delta\. The first\-order expansion is PK,ω\(θ\+δ\)=PK,ω\(θ\)\+Jω\(θ\)δ\+rω\(δ\),P\_\{K,\\omega\}\(\\theta\+\\delta\)=P\_\{K,\\omega\}\(\\theta\)\+J\_\{\\omega\}\(\\theta\)\\,\\delta\+r\_\{\\omega\}\(\\delta\),\(16\)where the remainderrω\(δ\)r\_\{\\omega\}\(\\delta\)is of order‖δ‖2\\left\\lVert\\delta\\right\\rVert^\{2\}wherever the pseudo\-gradient is smooth inθ\\theta\. We callJω\(θ\)J\_\{\\omega\}\(\\theta\)theblock Jacobian\. The chain rule gives Jω\(θ\)=I−∏k=K−10\[I−η∇2ℓξk\(θk,ω\)\]\.J\_\{\\omega\}\(\\theta\)=I\-\\prod\_\{k=K\-1\}^\{0\}\\bigl\[I\-\\eta\\,\\nabla^\{2\}\\ell\_\{\\xi\_\{k\}\}\(\\theta\_\{k,\\omega\}\)\\bigr\]\.\(17\) Each Hessian is taken at the point the inner iterate visits, so the block Jacobian depends on the minibatchesω\\omegaand, in general, on the starting pointθ\\theta\. In the model of Section[3\.1](https://arxiv.org/html/2610.22623#S3.SS1)the block JacobianJω\(θ\)J\_\{\\omega\}\(\\theta\)becomes the block operatorDK,ωD\_\{K,\\omega\}at everyθ\\theta\. Theorem[1](https://arxiv.org/html/2610.22623#Thmtheorem1), Proposition[2](https://arxiv.org/html/2610.22623#Thmtheorem2)and Corollary[4](https://arxiv.org/html/2610.22623#Thmtheorem4)\(a\) require only that the block Jacobian does not depend onθ\\thetaand that blocks are i\.i\.d\. across outer steps\. We refer to this assumption as thefrozen linearization\. In practice, we obtain this property by taking the block Jacobian of every outer step at a checkpointθ0\\theta\_\{0\}, with the minibatches drawn afresh for each worker at every outer step \(as outlined in Algorithm[2](https://arxiv.org/html/2610.22623#alg2)\)\. Note that each inner trajectory still runs on the actual loss fromθ0\\theta\_\{0\}, and each Hessian in \([17](https://arxiv.org/html/2610.22623#S6.E17)\) is taken at the point it visits, so the inner steps are not replaced by a quadratic model of the loss atθ0\\theta\_\{0\}\. Figure 4:Measured amplification on the CIFAR MLP at step120,000120\{,\}000\.Every panel showslog10Aθ\(20\)\\log\_\{10\}A\_\{\\theta\}\(20\), the amplification after2020outer steps, against the outer momentum in \(i\) and against the number of local steps in \(ii\) to \(iv\)\. Solid lines areAθA\_\{\\theta\}measured along actual runs, and dotted lines its prediction under the frozen linearization\. Arrows \(↑\\uparrow\) mark points where some runs diverged\.Figure 5:Spread cost on the CIFAR MLP at step120,000120\{,\}000\.The spread cost \([20](https://arxiv.org/html/2610.22623#S6.E20)\) of3232independent runs \(from different seeds\) after2020outer steps, at the settings of Figure[5](https://arxiv.org/html/2610.22623#S6.F5)\. A cross \(×\\times\) on the top edge marks a setting whose runs’ average ends with a loss more than0\.10\.1above the checkpoint’s\. Panel \(ii\) shows the spread cost of each inner batch relative toK=1K=1\. Under the frozen linearization, we can interpret the block Jacobian at the checkpoint as the block operator of Section[3\.1](https://arxiv.org/html/2610.22623#S3.SS1), and so we writeDω=Jω\(θ0\)D\_\{\\omega\}=J\_\{\\omega\}\(\\theta\_\{0\}\)\. Since the minibatches are drawn afresh at every outer step, theDt=Jωt\(θ0\)D\_\{t\}=J\_\{\\omega\_\{t\}\}\(\\theta\_\{0\}\)are i\.i\.d\., so Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1)holds whenever they have finite second moments\. The pseudo\-gradient then becomes affine inθ\\theta,PK,ωt\(θ\)=pt\+Dt\(θ−θ0\)P\_\{K,\\omega\_\{t\}\}\(\\theta\)=p\_\{t\}\+D\_\{t\}\(\\theta\-\\theta\_\{0\}\), as in the regular\-minimum case of Section[3](https://arxiv.org/html/2610.22623#S3), with anoffsetpt=PK,ωt\(θ0\)p\_\{t\}=P\_\{K,\\omega\_\{t\}\}\(\\theta\_\{0\}\)that need not have mean zero becauseθ0\\theta\_\{0\}is not a minimizer\. Importantly, the displacement of a run fromθ0\\theta\_\{0\}obeys the forced recurrence \(Appendix\) zt\+1=𝖳\(Dt\)zt\+𝖡ptz\_\{t\+1\}=\\mathsf\{T\}\(D\_\{t\}\)\\,z\_\{t\}\+\\mathsf\{B\}\\,p\_\{t\}\(18\)exactly\. The offset leavesℒK\\mathcal\{L\}\_\{K\}unchanged, so by Propositionthe covariance ofθt\\theta\_\{t\}across independent streams of minibatches, which we call thespreadof a run, stays bounded at a noise floor whenρ\(ℒK\)<1\\rho\(\\mathcal\{L\}\_\{K\}\)<1and generically grows without bound otherwise\. #### Shared\-minibatch perturbations\. Consider two runs of \([18](https://arxiv.org/html/2610.22623#S6.E18)\) that see the same minibatches but start from different points\. Because their minibatches are the same, both runs receive the same offsetptp\_\{t\}, and subtracting one recurrence from the other removes it\. Their differenceztz\_\{t\}therefore obeyszt\+1=𝖳\(Dt\)ztz\_\{t\+1\}=\\mathsf\{T\}\(D\_\{t\}\)\\,z\_\{t\}, which is \([5](https://arxiv.org/html/2610.22623#S3.E5)\)\. Unlike the spread, this difference decays exactly whenρ\(ℒK\)<1\\rho\(\\mathcal\{L\}\_\{K\}\)<1\. We call the differenceete\_\{t\}in their parameters ashared\-minibatch perturbationand measure stability from it \(Section[6\.2](https://arxiv.org/html/2610.22623#S6.SS2)\)\. At an actual network, the difference follows \([5](https://arxiv.org/html/2610.22623#S3.E5)\) only to first order \(Appendix\)\. ### 6\.2Amplification and spread cost We read the stability at a checkpoint from theamplificationof shared\-minibatch perturbations, and then measure its consequences on training loss through thespread cost\. #### Amplification\. We define theamplificationafterTTouter steps as Aθ\(T\)=𝔼‖eT‖2‖e0‖2,A\_\{\\theta\}\(T\)=\\frac\{\\mathbb\{E\}\\,\\\|e\_\{T\}\\\|^\{2\}\}\{\\\|e\_\{0\}\\\|^\{2\}\},\(19\)whereete\_\{t\}is the shared\-minibatch perturbation and the expectation is over streams of minibatches, and we say that a setting amplifies whenAθ\(T\)\>1A\_\{\\theta\}\(T\)\>1\. Following how a small perturbation grows is a common way to study the stability of training viewed as a dynamical system\([Wu et al\., 2018](https://arxiv.org/html/2610.22623#bib.bib54);[Dieuleveut et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib13);[Chemnitz and Engel, 2025](https://arxiv.org/html/2610.22623#bib.bib9);[Li and Xie, 2025](https://arxiv.org/html/2610.22623#bib.bib31)\)\. Notably,[Altıntaş et al\. \(2025\)](https://arxiv.org/html/2610.22623#bib.bib3)measure how far two such runs diverge in parameter distance and in functional similarity\. In contrast, under the frozen linearization our amplification decays exactly whenρ\(ℒK\)<1\\rho\(\\mathcal\{L\}\_\{K\}\)<1, and we test whether it changes with the outer momentum, the number of local steps, the inner batch size and the number of workers in the way the four implications predict\. As in Section[3\.1](https://arxiv.org/html/2610.22623#S3.SS1), we restrict to the complement of the common flat directions, on which every minibatch Hessian vanishes\. We therefore takee0=Jω′\(θ0\)ge\_\{0\}=J\_\{\\omega^\{\\prime\}\}\(\\theta\_\{0\}\)\\,gfor a random vectorggand a separate blockω′\\omega^\{\\prime\}, which lies in that complement, with an empty buffer \(Appendix\)\. Without ever formingℒK\\mathcal\{L\}\_\{K\}\(a map on2p×2p2p\\times 2pmatrices for a network withppparameters\), we estimate the amplification by propagatingete\_\{t\}with Hessian\-vector products \(HVPs\), along actual runs for the measured values and by Algorithm[2](https://arxiv.org/html/2610.22623#alg2)for the predictions under the frozen linearization\. Each Hessian\-vector product costs about as much as a gradient evaluation\([Pearlmutter, 1994](https://arxiv.org/html/2610.22623#bib.bib42)\), so trackingete\_\{t\}adds about one gradient’s worth of computation per inner step\. Sinceete\_\{t\}evolves linearly, we rescale it to unit length after every outer step and accumulate the logarithm of the scale, as for Lyapunov exponents\. #### Mean or mean\-square instability\. Varying the number of workers lets us tell whether an amplifying setting is unstable in its mean dynamics or only in mean square\. This works because averaging over workers divides the block noise byWWbut leaves the mean transition unchanged \(Corollary[4](https://arxiv.org/html/2610.22623#Thmtheorem4)\)\. Moreover, under the frozen linearization the amplification is non\-increasing inWWfor everyTT\(Appendix\)\. Hence, if a setting amplifies with one worker but contracts after many outer steps once enough workers are averaged, this is consistent with stable mean dynamics and mean\-square instability with one worker\. On the other hand, a setting that still amplifies at the largestWWwe try is either mean unstable or needs more workers\. #### Spread cost\. The change in the mean training loss ofSSindependent runs\([Fort et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib18);[Frankle et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib19)\)splits exactly into two parts\. The first, the averaged\-model change, is driven by the offset even in a stable setting, so a mean\-square instability can raise the loss only through the second, the spread cost\. Withμ^t\\hat\{\\mu\}\_\{t\}the average of the runs,L¯t\\bar\{L\}\_\{t\}their mean loss andLLthe full\-data loss, L¯t−L\(θ0\)=L\(μ^t\)−L\(θ0\)⏟averaged\-model change\+L¯t−L\(μ^t\)⏟spread cost\.\\bar\{L\}\_\{t\}\-L\(\\theta\_\{0\}\)=\\underbrace\{L\(\\hat\{\\mu\}\_\{t\}\)\-L\(\\theta\_\{0\}\)\}\_\{\\text\{averaged\-model change\}\}\+\\underbrace\{\\bar\{L\}\_\{t\}\-L\(\\hat\{\\mu\}\_\{t\}\)\}\_\{\\text\{spread cost\}\}\.\(20\)The spread cost is the loss due to the disagreement between the runs, and replacing the runs by their average would recover exactly this amount\. To second order it is a curvature\-weighted trace of their covariance, which Propositionpropagates \(Appendix\)\. We say that the runs left the checkpoint’s basin when their average ends with a loss more than0\.10\.1above the checkpoint’s\. ### 6\.3Empirical results We conduct experiments at checkpoints of an MLP and a ResNet\-20, both trained by plain SGD on CIFAR\-10\([Krizhevsky, 2009](https://arxiv.org/html/2610.22623#bib.bib29)\), and of Pythia\-70m \(Appendix\), and at one checkpoint of Pythia\-160m, Pythia\-410m and Pythia\-1b \(Appendix\)\. Figures[5](https://arxiv.org/html/2610.22623#S6.F5)and[5](https://arxiv.org/html/2610.22623#S6.F5)show a single checkpoint, the CIFAR MLP at step120,000120\{,\}000, and experiments with different checkpoints and other models are in Appendicesto\. The 95% bootstrap confidence intervals over the 48 minibatch streams \(from different seeds\) have a median width of 0\.027 inlog10Aθ\(20\)\\log\_\{10\}A\_\{\\theta\}\(20\)\. We take the checkpoints after the early phase of training, as it has been shown that linear approximations about an early checkpoint do not track the actual training well\([Fort et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib18);[Meterez et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib37)\)\. Figure[5](https://arxiv.org/html/2610.22623#S6.F5)is consistent with each of the four implications\. Asβ\\betagrows at a fixed outer rate, a setting turns from contracting to amplifying at most once \(i\)\. AsKKgrows, a setting contracts more at inner batches88and1616and amplifies more at batches11and22\(ii\), matching the two behaviors of Proposition[3](https://arxiv.org/html/2610.22623#Thmtheorem3)\. More workers never raise the amplification \(iii\), and splitting a fixed total batch did not lower it in any setting of Figure[5](https://arxiv.org/html/2610.22623#S6.F5)\(iv\)\. The setting of panel \(iii\) amplifies with one worker and contracts with two or more, so this instability appears to be mean\-square \(Section[6\.2](https://arxiv.org/html/2610.22623#S6.SS2)\)\. As with the amplification, the spread cost rises withβ\\beta, falls withWWand is not lowered by splitting a fixed total batch \(Figure[5](https://arxiv.org/html/2610.22623#S6.F5)\)\. AsKKgrows, it first rises at every inner batch, because the noise that a block injects accumulates over itsKKsteps \(Proposition\)\. The batches then separate as in Proposition[3](https://arxiv.org/html/2610.22623#Thmtheorem3)\. The spread cost keeps rising atb≤4b\\leq 4and turns down atb=8b=8and1616\. On ResNet\-20, on the Pythia models and at other checkpoints, the amplification is also largely consistent with the implications \(Appendix\)\. ## 7Discussion Our results show that the inner trajectory acts on the outer iteration as a random operator, the block operator, and that once the mean condition holds, stability is decided by whether the noise of this operator fits within the noise budget of the outer step\. The number of local steps, the inner batch size and the number of workers change the block operator, while the outer learning rate and momentum shape the noise budget, so the effect of any one hyperparameter depends on the others\. Stability can also be read at a checkpoint of a network from short continuations of training that share their minibatches\. In distributed training these hyperparameters are often constrained, since limited bandwidth sets a minimum number of local steps, the memory of each worker caps the inner batch size, and hardware failures can reduce the number of workers during training\. Under such constraints, the checkpoint measurement suggests a way to check whether a planned change keeps training stable before committing to it\. However, the spread cost need not follow the trends of the amplification at every checkpoint \(Appendix\), as it also depends on the noise that each block injects and on the curvature of the loss \(Proposition; Appendix\)\. #### Extensions\. Our criterion also restricts the minima that two\-phase optimizers can settle into, since the outer iteration cannot settle in mean square at a minimum where the criterion fails, and certain desirable properties of such minima might be valuable to the community\([Wu and Su, 2023](https://arxiv.org/html/2610.22623#bib.bib53)\)\. In this direction,[Mulayoff and Michaeli \(2024\)](https://arxiv.org/html/2610.22623#bib.bib39)study the dynamical stability of optimization methods in the vicinity of the minima of the loss, and there is also a rich literature on which minima are dynamically accessible to an optimization algorithm\([Wu et al\., 2018](https://arxiv.org/html/2610.22623#bib.bib54)\)\. Although it has yet to be shown, a two\-phase system with outer momentum may also display catapults\([Lewkowycz et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib30)\)or train at the edge of stability, as gradient descent\([Cohen et al\., 2021](https://arxiv.org/html/2610.22623#bib.bib10)\), SGD\([Andreyev and Beneventano, 2026](https://arxiv.org/html/2610.22623#bib.bib4)\), SGD with momentum\([Andreyev et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib5)\)and federated averaging\([Khandelwal et al\., 2026](https://arxiv.org/html/2610.22623#bib.bib27)\)do\. #### Limitations\. The mean\-square analysis focuses on stability rather than on the rate of convergence, and it does not capture the fact that training can, in the longer term, recover from a local instability\([Lewkowycz et al\., 2020](https://arxiv.org/html/2610.22623#bib.bib30);[Damian et al\., 2023](https://arxiv.org/html/2610.22623#bib.bib12)\)\. Our scope is limited to SGD for the inner steps, whereas modern two\-phase methods can use AdamW or Muon inside the block\. Such settings are especially difficult to study when an adaptive inner optimizer carries state across blocks and preconditions the response, so its block Jacobian is not the product \([17](https://arxiv.org/html/2610.22623#S6.E17)\) and Assumption[1](https://arxiv.org/html/2610.22623#Thmassumption1)then has to be re\-examined\. ## Acknowledgments This research was enabled in part by compute resources provided by Mila \(mila\.quebec\), and by support provided by Calcul Québec \(calculquebec\.ca\), Compute Ontario \(computeontario\.ca\), the BC DRI Group and the Digital Research Alliance of Canada \(alliancecan\.ca\)\. We also thank the Masason Foundation for providing compute resources\. Sarath Chandar is supported by the Canada CIFAR AI Chairs program, the Canada Research Chair in Lifelong Machine Learning, and an NSERC Discovery Grant\. ## References - Agarwala \(2026\)Atish Agarwala\.High dimensional theory of two\-phase optimizers, 2026\.URL[https://arxiv\.org/abs/2603\.26954](https://arxiv.org/abs/2603.26954)\. - Agarwala and Pennington \(2024\)Atish Agarwala and Jeffrey Pennington\.High dimensional analysis reveals conservative sharpening and a stochastic edge of stability, 2024\.URL[https://arxiv\.org/abs/2404\.19261](https://arxiv.org/abs/2404.19261)\. - Altıntaş et al\. \(2025\)Gül Sena Altıntaş, Devin Kwok, Colin Raffel, and David Rolnick\.The butterfly effect: Neural network training trajectories are highly sensitive to initial conditions\.In*Proceedings of the 42nd International Conference on Machine Learning*, volume 267 of*Proceedings of Machine Learning Research*, pages 1314–1342\. PMLR, 2025\.arXiv:2506\.13234\. - Andreyev and Beneventano \(2026\)Arseniy Andreyev and Pierfrancesco Beneventano\.Edge of stochastic stability: Revisiting the edge of stability for SGD\.In*Advances in Neural Information Processing Systems*, volume 39, 2026\. - Andreyev et al\. \(2026\)Arseniy Andreyev, Advikar Ananthkumar, Marc Walden, Tomaso Poggio, and Pierfrancesco Beneventano\.Momentum further constrains sharpness at the edge of stochastic stability\.In*Proceedings of the 43rd International Conference on Machine Learning*, volume 306 of*Proceedings of Machine Learning Research*, pages 2831–2871, 2026\. - Berman and Plemmons \(1994\)Abraham Berman and Robert J\. Plemmons\.*Nonnegative Matrices in the Mathematical Sciences*\.SIAM, 1994\.doi:10\.1137/1\.9781611971262\. - Biderman et al\. \(2023\)Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal\.Pythia: A suite for analyzing large language models across training and scaling\.In*Proceedings of the 40th International Conference on Machine Learning*, 2023\. - Charles et al\. \(2025\)Zachary Charles, Gabriel Teston, Lucio Dery, Keith Rush, Nova Fallen, Zachary Garrett, Arthur Szlam, and Arthur Douillard\.Communication\-efficient language model training scales reliably and robustly: Scaling laws for DiLoCo\.In*Advances in Neural Information Processing Systems*, 2025\.URL[https://arxiv\.org/abs/2503\.09799](https://arxiv.org/abs/2503.09799)\. - Chemnitz and Engel \(2025\)Dennis Chemnitz and Maximilian Engel\.Characterizing dynamical stability of stochastic gradient descent in overparameterized learning\.*Journal of Machine Learning Research*, 26\(134\):1–46, 2025\.arXiv:2407\.20209\. - Cohen et al\. \(2021\)Jeremy M\. Cohen, Simran Kaur, Yuanzhi Li, J\. Zico Kolter, and Ameet Talwalkar\.Gradient descent on neural networks typically occurs at the edge of stability\.In*International Conference on Learning Representations*, 2021\.URL[https://arxiv\.org/abs/2103\.00065](https://arxiv.org/abs/2103.00065)\. - Costa et al\. \(2005\)Oswaldo Luiz Valle Costa, Marcelo Dutra Fragoso, and Ricardo Paulino Marques\.*Discrete\-Time Markov Jump Linear Systems*\.Probability and Its Applications\. Springer London, 2005\.doi:10\.1007/b138575\.URL[https://link\.springer\.com/book/10\.1007/b138575](https://link.springer.com/book/10.1007/b138575)\. - Damian et al\. \(2023\)Alex Damian, Eshaan Nichani, and Jason D\. Lee\.Self\-stabilization: The implicit bias of gradient descent at the edge of stability\.In*International Conference on Learning Representations*, 2023\.URL[https://arxiv\.org/abs/2209\.15594](https://arxiv.org/abs/2209.15594)\. - Dieuleveut et al\. \(2020\)Aymeric Dieuleveut, Alain Durmus, and Francis Bach\.Bridging the gap between constant step size stochastic gradient descent and Markov chains\.*Annals of Statistics*, 2020\.URL[https://arxiv\.org/abs/1707\.06386](https://arxiv.org/abs/1707.06386)\. - Douillard et al\. \(2023\)Arthur Douillard, Qixuan Feng, Andrei A\. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc’Aurelio Ranzato, Arthur Szlam, and Jiajun Shen\.DiLoCo: Distributed low\-communication training of language models, 2023\.URL[https://arxiv\.org/abs/2311\.08105](https://arxiv.org/abs/2311.08105)\. - Douillard et al\. \(2025\)Arthur Douillard, Yani Donchev, Keith Rush, Satyen Kale, Zachary Charles, Gabriel Teston, Zachary Garrett, Jiajun Shen, Ross McIlroy, David Lacey, Alexandre Ramé, Arthur Szlam, Marc’Aurelio Ranzato, and Paul Barham\.Streaming DiLoCo with overlapping communication\.In*Conference on Language Modeling*, 2025\.URL[https://arxiv\.org/abs/2501\.18512](https://arxiv.org/abs/2501.18512)\. - Douillard et al\. \(2026\)Arthur Douillard, Keith Rush, Yani Donchev, Zachary Charles, Nova Fallen, Ayush Dubey, Ionel Gog, Josef Dean, Blake Woodworth, Zachary Garrett, Nate Keating, Jenny Bishop, Henry Prior, Edouard Yvinec, Arthur Szlam, Marc’Aurelio Ranzato, and Jeff Dean\.Decoupled DiLoCo for resilient distributed pre\-training, 2026\.URL[https://arxiv\.org/abs/2604\.21428](https://arxiv.org/abs/2604.21428)\. - Duren \(1970\)Peter L\. Duren\.*Theory ofHpH^\{p\}Spaces*, volume 38 of*Pure and Applied Mathematics*\.Academic Press, New York, 1970\. - Fort et al\. \(2020\)Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M\. Roy, and Surya Ganguli\.Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel\.In*Advances in Neural Information Processing Systems*, volume 33, pages 5850–5861, 2020\.arXiv:2010\.15110\. - Frankle et al\. \(2020\)Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin\.Linear mode connectivity and the lottery ticket hypothesis\.In*Proceedings of the 37th International Conference on Machine Learning*, volume 119 of*Proceedings of Machine Learning Research*, pages 3259–3269\. PMLR, 2020\. - Gao et al\. \(2020\)Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy\.The Pile: An 800GB dataset of diverse text for language modeling\.*arXiv preprint arXiv:2101\.00027*, 2020\. - Granziol et al\. \(2022\)Diego Granziol, Stefan Zohren, and Stephen Roberts\.Learning rates as a function of batch size: A random matrix theory approach to neural network training\.*Journal of Machine Learning Research*, 2022\.URL[https://www\.jmlr\.org/papers/v23/20\-1258\.html](https://www.jmlr.org/papers/v23/20-1258.html)\. - Gupta et al\. \(2024\)Kanan Gupta, Jonathan W\. Siegel, and Stephan Wojtowytsch\.Nesterov acceleration despite very noisy gradients\.In*Advances in Neural Information Processing Systems*, volume 37, pages 20694–20744, 2024\.doi:10\.52202/079017\-0653\.URL[https://proceedings\.neurips\.cc/paper\_files/paper/2024/hash/24d2dd6dc9b79116f8ebc852ddb9dc94\-Abstract\-Conference\.html](https://proceedings.neurips.cc/paper_files/paper/2024/hash/24d2dd6dc9b79116f8ebc852ddb9dc94-Abstract-Conference.html)\. - He et al\. \(2016\)Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun\.Deep residual learning for image recognition\.In*Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition*, pages 770–778, 2016\. - Jordan et al\. \(2024\)Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein\.Muon: An optimizer for hidden layers in neural networks, 2024\.URL[https://kellerjordan\.github\.io/posts/muon/](https://kellerjordan.github.io/posts/muon/)\. - Kallusky et al\. \(2025\)Dominik Kallusky, Vinay Rao, Vishal Nandavanam, and Hao\-Jun Michael Shi\.SNOO: Step\-KKNesterov outer optimizer—the surprising effectiveness of Nesterov momentum applied to pseudo\-gradients, 2025\.URL[https://arxiv\.org/abs/2510\.15830](https://arxiv.org/abs/2510.15830)\. - Khaled et al\. \(2025\)Ahmed Khaled, Satyen Kale, Arthur Douillard, Chi Jin, Rob Fergus, and Manzil Zaheer\.Understanding outer optimizers in local SGD: Learning rates, momentum, and acceleration\.In*Advances in Neural Information Processing Systems*, volume 38, pages 111163–111216, 2025\. - Khandelwal et al\. \(2026\)Anant Khandelwal, Michael Crawshaw, and Mingrui Liu\.How edge of stability hinders SCAFFOLD in federated optimization, 2026\.URL[https://arxiv\.org/abs/2608\.25873](https://arxiv.org/abs/2608.25873)\. - Kingma and Ba \(2015\)Diederik P\. Kingma and Jimmy Ba\.Adam: A method for stochastic optimization\.In*3rd International Conference on Learning Representations \(ICLR\)*, San Diego, CA, USA, 2015\.URL[https://arxiv\.org/abs/1412\.6980](https://arxiv.org/abs/1412.6980)\. - Krizhevsky \(2009\)Alex Krizhevsky\.Learning multiple layers of features from tiny images\.Technical report, University of Toronto, 2009\. - Lewkowycz et al\. \(2020\)Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl\-Dickstein, and Guy Gur\-Ari\.The large learning rate phase of deep learning: The catapult mechanism, 2020\.URL[https://arxiv\.org/abs/2003\.02218](https://arxiv.org/abs/2003.02218)\. - Li and Xie \(2025\)Xiang Li and Qiaomin Xie\.Coupling\-based convergence diagnostic and stepsize scheme for stochastic gradient descent\.In*Proceedings of the AAAI Conference on Artificial Intelligence*, volume 39, pages 18494–18501, 2025\.doi:10\.1609/aaai\.v39i17\.34035\.arXiv:2412\.11341\. - Liao et al\. \(2026\)Fangshuo Liao, Afroditi Kolomvaki, and Anastasios Kyrillidis\.SGD at the edge of stability: The stochastic sharpness gap\.In*Advances in Neural Information Processing Systems*, volume 39, 2026\. - Lin et al\. \(2020\)Tao Lin, Sebastian U\. Stich, Kumar Kshitij Patel, and Martin Jaggi\.Don’t use large mini\-batches, use local SGD\.In*International Conference on Learning Representations*, 2020\.URL[https://arxiv\.org/abs/1808\.07217](https://arxiv.org/abs/1808.07217)\. - Loshchilov and Hutter \(2019\)Ilya Loshchilov and Frank Hutter\.Decoupled weight decay regularization\.In*International Conference on Learning Representations \(ICLR\)*, 2019\.URL[https://arxiv\.org/abs/1711\.05101](https://arxiv.org/abs/1711.05101)\. - Ma and Ying \(2021\)Chao Ma and Lexing Ying\.On linear stability of SGD and input\-smoothness of neural networks\.In*Advances in Neural Information Processing Systems*, 2021\.URL[https://arxiv\.org/abs/2105\.13462](https://arxiv.org/abs/2105.13462)\. - McMahan et al\. \(2017\)Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas\.Communication\-efficient learning of deep networks from decentralized data\.In*Proceedings of the 20th International Conference on Artificial Intelligence and Statistics*, volume 54 of*Proceedings of Machine Learning Research*, pages 1273–1282, 2017\.URL[https://proceedings\.mlr\.press/v54/mcmahan17a\.html](https://proceedings.mlr.press/v54/mcmahan17a.html)\. - Meterez et al\. \(2026\)Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade, and Alex Damian\.A defense of the quadratic model, 2026\.URL[https://arxiv\.org/abs/2607\.21716](https://arxiv.org/abs/2607.21716)\. - Mohammadi et al\. \(2021\)Hesameddin Mohammadi, Meisam Razaviyayn, and Mihailo R\. Jovanović\.Robustness of accelerated first\-order algorithms for strongly convex optimization problems\.*IEEE Transactions on Automatic Control*, 66\(6\):2480–2495, 2021\.arXiv:1905\.11011\. - Mulayoff and Michaeli \(2024\)Rotem Mulayoff and Tomer Michaeli\.Exact mean square linear stability analysis for SGD\.In*Proceedings of Thirty Seventh Conference on Learning Theory*, volume 247 of*Proceedings of Machine Learning Research*, pages 3915–3969, 2024\.URL[https://proceedings\.mlr\.press/v247/mulayoff24a\.html](https://proceedings.mlr.press/v247/mulayoff24a.html)\. - Naganuma et al\. \(2025\)Hiroki Naganuma, Xinzhi Zhang, Man\-Chung Yue, Ioannis Mitliagkas, Russell J\. Hewett, Philipp A\. Witte, and Yin Tat Lee\.Pseudo\-asynchronous local SGD: Robust and efficient data\-parallel training\.*Transactions on Machine Learning Research*, 2025\.URL[https://openreview\.net/forum?id=8VTrvS5vN7](https://openreview.net/forum?id=8VTrvS5vN7)\. - Nesterov \(1983\)Yurii Nesterov\.A method of solving a convex programming problem with convergence rateO\(1/k2\)O\(1/k^\{2\}\)\.*Soviet Mathematics Doklady*, 27\(2\):372–376, 1983\. - Pearlmutter \(1994\)Barak A\. Pearlmutter\.Fast exact multiplication by the Hessian\.*Neural Computation*, 1994\.URL[https://doi\.org/10\.1162/neco\.1994\.6\.1\.147](https://doi.org/10.1162/neco.1994.6.1.147)\. - Reddi et al\. \(2021\)Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H\. Brendan McMahan\.Adaptive federated optimization\.In*International Conference on Learning Representations*, 2021\.URL[https://arxiv\.org/abs/2003\.00295](https://arxiv.org/abs/2003.00295)\. - Smith and Le \(2018\)Samuel L\. Smith and Quoc V\. Le\.A Bayesian perspective on generalization and stochastic gradient descent\.In*International Conference on Learning Representations*, 2018\. - Smith et al\. \(2018\)Samuel L\. Smith, Pieter\-Jan Kindermans, Chris Ying, and Quoc V\. Le\.Don’t decay the learning rate, increase the batch size\.In*International Conference on Learning Representations*, 2018\. - Stich \(2019\)Sebastian U\. Stich\.Local SGD converges fast and communicates little\.In*International Conference on Learning Representations*, 2019\.URL[https://arxiv\.org/abs/1805\.09767](https://arxiv.org/abs/1805.09767)\. - Thérien et al\. \(2025\)Benjamin Thérien, Xiaolong Huang, Aaron Defazio, Irina Rish, and Eugene Belilovsky\.MuLoCo: Muon is a practical inner optimizer for DiLoCo\.*arXiv preprint arXiv:2505\.23725*, 2025\. - Topollai et al\. \(2026\)Kristi Topollai, Allan Ma, Tolga Dimlioglu, Sui Jiet Tay, and Anna Choromanska\.Outer\-momentum restarting in high\-dimensional two\-phase optimization, 2026\.URL[https://arxiv\.org/abs/2605\.28585](https://arxiv.org/abs/2605.28585)\. - Velikanov et al\. \(2023\)Maksim Velikanov, Denis Kuznedelev, and Dmitry Yarotsky\.A view of mini\-batch SGD via generating functions: conditions of convergence, phase transitions, benefit from negative momenta\.In*International Conference on Learning Representations*, 2023\. - Wang and Joshi \(2019\)Jianyu Wang and Gauri Joshi\.Adaptive communication strategies to achieve the best error\-runtime trade\-off in local\-update SGD\.In*Proceedings of Machine Learning and Systems*, 2019\.URL[https://arxiv\.org/abs/1810\.08313](https://arxiv.org/abs/1810.08313)\. - Wang et al\. \(2020\)Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat\.SlowMo: Improving communication\-efficient distributed SGD with slow momentum\.In*International Conference on Learning Representations*, 2020\.URL[https://arxiv\.org/abs/1910\.00643](https://arxiv.org/abs/1910.00643)\. - Wortsman et al\. \(2022\)Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo\-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt\.Model soups: averaging weights of multiple fine\-tuned models improves accuracy without increasing inference time\.In*Proceedings of the 39th International Conference on Machine Learning*, volume 162 of*Proceedings of Machine Learning Research*, pages 23965–23998\. PMLR, 2022\. - Wu and Su \(2023\)Lei Wu and Weijie J\. Su\.The implicit regularization of dynamical stability in stochastic gradient descent\.In*Proceedings of the International Conference on Machine Learning*, 2023\.URL[https://arxiv\.org/abs/2305\.17490](https://arxiv.org/abs/2305.17490)\. - Wu et al\. \(2018\)Lei Wu, Chao Ma, and Weinan E\.How SGD selects the global minima in over\-parameterized learning: A dynamical stability perspective\.In*Advances in Neural Information Processing Systems*, 2018\.URL[https://papers\.neurips\.cc/paper/8049\-how\-sgd\-selects\-the\-global\-minima\-in\-over\-parameterized\-learning\-a\-dynamical\-stability\-perspective\.pdf](https://papers.neurips.cc/paper/8049-how-sgd-selects-the-global-minima-in-over-parameterized-learning-a-dynamical-stability-perspective.pdf)\. - Wu et al\. \(2022\)Lei Wu, Mingze Wang, and Weijie J\. Su\.The alignment property of SGD noise and how it helps select flat minima: A stability analysis\.In*Advances in Neural Information Processing Systems*, 2022\.URL[https://arxiv\.org/abs/2207\.02628](https://arxiv.org/abs/2207.02628)\. - Wu and He \(2018\)Yuxin Wu and Kaiming He\.Group normalization\.In*Proceedings of the European Conference on Computer Vision*, pages 3–19, 2018\. - Zhang et al\. \(2019\)Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E\. Hinton\.Lookahead optimizer:kksteps forward, 1 step back\.In*Advances in Neural Information Processing Systems*, volume 32, 2019\.URL[https://proceedings\.neurips\.cc/paper/2019/hash/90fd4f88f588ae64038134f1eeaa023f\-Abstract\.html](https://proceedings.neurips.cc/paper/2019/hash/90fd4f88f588ae64038134f1eeaa023f-Abstract.html)\. - Ziyin et al\. \(2023\)Liu Ziyin, Botao Li, Tomer Galanti, and Masahito Ueda\.Type\-II saddles and probabilistic stability of stochastic gradient descent, 2023\.URL[https://arxiv\.org/abs/2303\.13093](https://arxiv.org/abs/2303.13093)\. ###### Contents of the Appendix 1. [1Introduction](https://arxiv.org/html/2610.22623#S1) 2. [2Related Work](https://arxiv.org/html/2610.22623#S2) 3. [3Setup](https://arxiv.org/html/2610.22623#S3)1. [3\.1Block operator and stability](https://arxiv.org/html/2610.22623#S3.SS1) 2. [3\.2Block moments and howKKandbbenter them](https://arxiv.org/html/2610.22623#S3.SS2) 4. [4The Stability Boundary](https://arxiv.org/html/2610.22623#S4)1. [4\.1Second\-moment map and the block noise](https://arxiv.org/html/2610.22623#S4.SS1) 2. [4\.2The mean condition and the noise budget](https://arxiv.org/html/2610.22623#S4.SS2) 5. [5Implications forβ\\beta,KK&WW](https://arxiv.org/html/2610.22623#S5)1. [5\.1Outer momentum shrinks the noise budget](https://arxiv.org/html/2610.22623#S5.SS1) 2. [5\.2Local steps and the inner batch size](https://arxiv.org/html/2610.22623#S5.SS2) 3. [5\.3Averaging over workers](https://arxiv.org/html/2610.22623#S5.SS3) 4. [5\.4Fixed total batch](https://arxiv.org/html/2610.22623#S5.SS4) 6. [6Testing the Implications on Neural Networks](https://arxiv.org/html/2610.22623#S6)1. [6\.1The block Jacobian](https://arxiv.org/html/2610.22623#S6.SS1) 2. [6\.2Amplification and spread cost](https://arxiv.org/html/2610.22623#S6.SS2) 3. [6\.3Empirical results](https://arxiv.org/html/2610.22623#S6.SS3) 7. [7Discussion](https://arxiv.org/html/2610.22623#S7) 8. [References](https://arxiv.org/html/2610.22623#bib) 9. [AThe Two\-Phase Optimizer with Nesterov Momentum](https://arxiv.org/html/2610.22623#A1) 10. [BMean and Mean\-Square Stability](https://arxiv.org/html/2610.22623#A2) 11. [CNotation and Conventions](https://arxiv.org/html/2610.22623#A3) ## Appendix AThe Two\-Phase Optimizer with Nesterov Momentum Algorithm[1](https://arxiv.org/html/2610.22623#alg1)gives one run of the two\-phase optimizer \([2](https://arxiv.org/html/2610.22623#S3.E2)\) withWWworkers\. Most of Sections[3](https://arxiv.org/html/2610.22623#S3)and[4](https://arxiv.org/html/2610.22623#S4)treat the caseW=1W=1, and Section[5\.3](https://arxiv.org/html/2610.22623#S5.SS3)the general case\. At each outer step the buffer is updated first, and the outer step then uses the updated buffer\. This makes the outer step Nesterov momentum, since heavy\-ball momentum would instead step along the buffer alone,θt\+1=θt−αut\\theta\_\{t\+1\}=\\theta\_\{t\}\-\\alpha u\_\{t\}\. Algorithm[2](https://arxiv.org/html/2610.22623#alg2)gives the frozen linearization of Section[6\.1](https://arxiv.org/html/2610.22623#S6.SS1), which produces the dotted lines in Figure[5](https://arxiv.org/html/2610.22623#S6.F5)and in the appendix figures\. It is Algorithm[1](https://arxiv.org/html/2610.22623#alg1)applied to a perturbationete\_\{t\}, with every inner trajectory started at the checkpointθ0\\theta\_\{0\}instead of the current iterate\. Algorithm 2The frozen linearization at a checkpointθ0\\theta\_\{0\}1:checkpoint θ0\\theta\_\{0\}, perturbation e0e\_\{0\}, inner rate η\\eta, inner batch size bb, outer rate α\\alpha, outer momentum β\\beta, number of local steps KK, number of workers WW, number of outer steps TT 2: u−1←0u\_\{\-1\}\\leftarrow 0 3:for t=0,…,T−1t=0,\\ldots,T\-1do 4:for w=1,…,Ww=1,\\ldots,Wdo⊳\\trianglerightper worker 5:runKKinner SGD steps at rateη\\etawith batch sizebbfromθ0\\theta\_\{0\}on fresh minibatchesωt\(w\)\\omega\_\{t\}^\{\(w\)\} 6: dt\(w\)←Jωt\(w\)\(θ0\)etd\_\{t\}^\{\(w\)\}\\leftarrow J\_\{\\omega\_\{t\}^\{\(w\)\}\}\(\\theta\_\{0\}\)\\,e\_\{t\}⊳\\triangleright\([17](https://arxiv.org/html/2610.22623#S6.E17)\),KKHVPs 7:endfor 8: dt←1W∑w=1Wdt\(w\)d\_\{t\}\\leftarrow\\frac\{1\}\{W\}\\sum\_\{w=1\}^\{W\}d\_\{t\}^\{\(w\)\} 9: ut←βut−1\+dtu\_\{t\}\\leftarrow\\beta u\_\{t\-1\}\+d\_\{t\} 10: et\+1←et−α\(βut\+dt\)e\_\{t\+1\}\\leftarrow e\_\{t\}\-\\alpha\\bigl\(\\beta u\_\{t\}\+d\_\{t\}\\bigr\) 11:endfor 1Linearizedinner phase2Outerstep ## Appendix BMean and Mean\-Square Stability Here the erroret=θt−θ∗e\_\{t\}=\\theta\_\{t\}\-\\theta^\{\*\}is the displacement from the minimizerθ∗\\theta^\{\*\}of the random quadratic model of Section[3\.1](https://arxiv.org/html/2610.22623#S3.SS1), andzt=\(et,ut−1\)z\_\{t\}=\(e\_\{t\},u\_\{t\-1\}\)stacks it with the momentum buffer\. By Definition[1](https://arxiv.org/html/2610.22623#Thmdefinition1), the recurrence is mean stable if𝔼\[zt\]→0\\mathbb\{E\}\[z\_\{t\}\]\\to 0and mean\-square stable if𝔼‖zt‖2→0\\mathbb\{E\}\\\|z\_\{t\}\\\|^\{2\}\\to 0\. Since𝔼‖zt‖2=‖𝔼\[zt\]‖2\+trCov\(zt\)\\mathbb\{E\}\\\|z\_\{t\}\\\|^\{2\}=\\\|\\mathbb\{E\}\[z\_\{t\}\]\\\|^\{2\}\+\\operatorname\{tr\}\\operatorname\{Cov\}\(z\_\{t\}\), mean\-square stability also requires the fluctuations ofztz\_\{t\}around its mean to vanish\. Figure 6:Mean stable but mean\-square unstable\.The mean‖𝔼\[zt\]‖2\\\|\\mathbb\{E\}\[z\_\{t\}\]\\\|^\{2\}and the second moment𝔼‖zt‖2\\mathbb\{E\}\\\|z\_\{t\}\\\|^\{2\}for the linear\-regression model above \(η=2\.5\\eta=2\.5,b=1b=1,K=4K=4,α=1\\alpha=1,β=0\.5\\beta=0\.5\), starting from a unit error with an empty buffer\. Both curves are computed exactly from the moment recurrences \([9](https://arxiv.org/html/2610.22623#S4.E9)\)\. The dotted line marks the initial value\.Figure[6](https://arxiv.org/html/2610.22623#A2.F6)shows a setting in which the mean ofztz\_\{t\}converges while its second moment diverges\. The model used is linear regression in eight dimensions with unit\-norm features, so each sample loss has curvature one along its own feature, while the average curvature is about1/81/8in every direction \(Appendix\)\. With the inner rateη=2\.5\\eta=2\.5, a step on a single sample therefore overshoots along that sample’s feature, while the average step shrinks the error by a factor of about0\.690\.69\. The pseudo\-gradient thus points towards the minimizer on average, but its fluctuations are proportional to the current error; so they compound over the outer steps while the mean vanishes\. The divergence of the second moment does not come from rare runs, since every one of4,0004\{,\}000simulated runs diverges\. In the terms of Section[4](https://arxiv.org/html/2610.22623#S4), we see that the mean condition holds \(ρ\(A\)≈0\.37\\rho\(A\)\\approx 0\.37\) while the block noise does not fit within the noise budget \(ρ\(ℒK\)≈3\.9\\rho\(\\mathcal\{L\}\_\{K\}\)\\approx 3\.9, Theorem[1](https://arxiv.org/html/2610.22623#Thmtheorem1)\(ii\)\)\. An analysis of the mean dynamics alone would therefore call this setting stable\. ## Appendix CNotation and Conventions The table below lists the notation of the paper and Appendixlists the terms the paper defines\. Symbols used inside a single proof are defined there\.
Similar Articles
Edge of Stability Selectively Shapes Learning Across the Data Distribution
MIT researchers show that the edge of stability (EoS) in neural network training is not merely a global optimization phenomenon but selectively redistributes learning across subsets of the training distribution, amplifying progress on some data groups while suppressing others. They identify two key conditions governing this allocation: gradient alignment with the top Hessian eigenvector and sustained non-vanishing gradient magnitude.
Theoretical and Practical Advances in Hierarchical Bayesian Deep Learning Methodology
该论文提出了多层粒子边缘Metropolis-Hastings(MLPMMH)方法,通过伪边缘推断与跨网络宽度截断层级的耦合SMC相结合,实现分层贝叶斯深度学习中参数与超参数的联合推断,并在理论上证明了O(ε²)的均方误差与O(ε⁻²)的计算复杂度,同时给出可扩展的近似版本S-MLPMMH用于MNIST实验。
The Road Taken: The Role of Optimizers at the Edge of Stability
This paper re-examines the edge of stability in deep learning optimization, proposing a new formulation based on directional Hessian and gradient-alignment score for more accurate predictions and diagnostic tools.
Almost Sure Convergence Analysis of Stochastic Gradient Methods with Clipping and Additive Noise
The paper proves almost sure convergence of stochastic gradient descent with clipping and additive noise, including momentum variants, under smoothness and bounded gradient noise assumptions, providing theoretical foundations for stable training in convex and nonconvex settings.
Generalization Dynamics of LM Pre-training (4 minute read)
该论文提出"模式跳变"(mode-hopping)概念,指出LM预训练中平滑下降的损失与稳定基准分数会掩盖模型答题方式的突变(如OLMo3-32B在算术提示上81%→0%→81.7%的反复),并通过六项廉价行为测试排除了评估指标非线性与局部优化波动等简单解释。