Rare Gate Disagreements Can Limit Plasticity: When Gradient Flow Mispredicts Finite-Batch SGD

Ruoyu Zhao; Mingxuan Zhang; Jianbo Dai; Jiaqi Wu; Chenyu Zhu; Tong Che

cs.LGOct 8, 2026

Rare Gate Disagreements Can Limit Plasticity: When Gradient Flow Mispredicts Finite-Batch SGD

Authors: Ruoyu Zhao, Mingxuan Zhang, Jianbo Dai, Jiaqi Wu, Chenyu Zhu, Tong Che

Organizations: City University of Hong Kong · Microsoft · Copula Lab · NVIDIA Research

Abstract

Population gradient flow is a common tool for reasoning about how neural networks adapt, including after pretraining. We show that it can mispredict finite-batch stochastic gradient descent (SGD) qualitatively, and we trace the discrepancy to a specific mechanism. In a two-unit ReLU regression, a source task drives the two neurons toward positive proportionality and a target task rewards separating them. After source training for time $T$ , gradient flow recovers on the target in time linear in $T$ . Online SGD with batch size $b$ and step size $η$ in both phases instead fails with high probability throughout a horizon of order $e^{c/η}$ once $T \gtrsim \log(b/η)$ , uniformly on an explicit set of initializations with Gaussian probability above one percent. For each fixed $T$ , small-step SGD still recovers, so the failure requires the joint limit of small steps and long pretraining. At the target clone, the population instability is carried entirely by inputs on which the two ReLU gates disagree. For units at angle $δ$ these inputs form a wedge of probability $δ/π$ , and weight decay shrinks the angle exponentially during pretraining. On every other input both units receive the same random linear update, which contracts their separation in conditional expectation. Bounding the cumulative probability of sampling the wedge along the exact online recursion, without a diffusion approximation, shows that recovery with fixed probability from an identical source-gradient-flow checkpoint, within $e^{c/η}$ updates, requires $Nb \gtrsim e^{λT}$ target samples and batch size $b \gtrsim ηe^{λT}$ , where $N$ counts updates and $λ$ is the weight decay. In simulations, recovery is approximately a function of the disagreement budget $bδ/η$ and saturates in the horizon.

Figures & tables

Figure 1: Mechanism on one trajectory (

T=400

\eta=10^{-3}

b=16

). (a) Two gates at angle

\delta

disagree on wedges of probability

\delta/\pi

. (b) Source training contracts the angle at rate

\lambda

. (c,d) From the source-SGD checkpoint, target GD (blue; Euler steps of GF) separates the units and crosses risk

0.4

near the GF delay

\lambda T/(2/\pi-\lambda)

(arrow), while four target SGD streams (orange) sample no disagreement and stay at risk

{\approx}0.6

. (c) shares the vertical axis of (b); dashed: rates of Section 4 ; shading:

\Gamma<\pi\lambda

, under one expected disagreement (Poisson approximation).

Figure 2: Identical checkpoints. Target hitting times from shared source-SGD checkpoints (

T\in\{100,200,300,400\}

\eta=10^{-3}

b=16

). GD (blue) follows

A^{-1}\log(1/\delta_{T})

plus the median offset (dashed); SGD (orange) rarely recovers by

H=200

(top strip) once

\Gamma_{T}<\pi\lambda

(shaded).

Figure 3: Recovery is organized by the disagreement budget. (a) Recovery probability is approximately a function of

\Gamma=b\delta_{0}/\eta

and lies at or below the Poisson reference

1-e^{-\Gamma/(\pi\lambda)}

(dashed) up to sampling error. Orange: constructed checkpoints with 95% Wilson intervals (filled:

\eta=0.01

); diamonds: checkpoints from Gaussian initializations. (b) Recovery saturates in the horizon (solid), far below the frozen-angle reference (dotted); three of the five long-horizon settings are shown.

	source SGD		source GD
$T$	target SGD	target GD	target SGD	target GD
0	64	64	64	64
100	64	64	64	64
200	59	64	61	64
300	20	64	19	64
400	7	64	7	64

Table 1: Identical checkpoints. Number of the 64 saved initializations from which each target optimizer reaches risk

0.4

H=200

(

\eta=10^{-3}

b=16

). The source-GD columns are the protocol of Corollary 5 .

$b$	$\delta_{0}$	$\Gamma$	rec. $t\leq 60$	rec. $t\leq 200$	rec. $t\leq 450$	$1-e^{-\Gamma/(\pi\lambda)}$	frozen, $t=200$
16	$10^{-5}$	0.016	0.070	0.078	0.078	0.097	0.639
64	$10^{-5}$	0.064	0.227	0.250	0.250	0.335	0.983
1	$10^{-3}$	0.100	0.305	0.336	0.336	0.471	0.998
16	$10^{-4}$	0.160	0.578	0.633	0.633	0.639	1.000
64	$10^{-4}$	0.640	0.969	0.984	0.984	0.983	1.000

Table 2: Long horizons. Recovered fractions by target time

t

in the five constructed settings (

\eta=0.01

, 128 streams each), the Poisson reference, and the frozen-angle reference

1-e^{-\Gamma t/\pi}

t=200

Figure 4: Joint dependence on source time and step size. (a) Warm recovery declines with

T

, later for smaller

\eta

(light to dark: increasing

\eta

). (b) Binned by

\Gamma_{T}=b\delta_{T}/\eta

, all step sizes fall near the Poisson reference (dashed); bins of width

0.5

\log_{10}\Gamma_{T}

with at least 8 records.

Figure 5: Weight-decay and parameter-noise controls from constructed checkpoints (

b=16

\eta=0.01

, 128 runs per setting). (a,b) With target weight decay, recovery (solid) plateaus at or below the Poisson reference (dashed); without it, it keeps growing. Dotted: fraction that has sampled a disagreement. (c) Independent parameter noise raises recovery with

s

(solid); shared noise barely does (dash-dotted);

\lambda=0.05

\delta_{0}=10^{-5}

; black: no noise.

Explore similar work

Sep 30, 2026cs.LG

Awakening of the Buddha: Subspace Learning During Population-Loss Plateaus

Population loss can remain nearly constant while a neural network learns a substantially more predictive representation. We establish this separation for two-layer ReLU and leaky-ReLU networks trained on Gaussian inputs by simultaneous fixed-step population gradient descent on all parameters. For structured additive teachers whose links are positive mixtures of Gaussian-damped cubics in

H^1(γ)

, we give explicit conditions under which small IID Gaussian initialization yields a high-probability guarantee: at a checkpoint during a high-loss plateau, minimum alignment between the rank-

r

teacher subspace and the leading

r

-dimensional eigenspace of the predictor's average gradient outer product (AGOP) increases by at least

1/2

, and the minimum refit MSE under unchanged coefficient budgets decreases by more than

0.399

, both relative to initialization. The same trajectory subsequently attains a trained loss below every value in the plateau window. A complementary result treats unequal-weight cubic teachers and small additive Sobolev perturbations using projected-feature refits. For SwiGLU networks with an exactly fitted intercept, we prove leading-AGOP alignment during a loss plateau at fixed width and dimension as Gaussian initialization vanishes, for square-integrable teachers with nonzero Hermite content of degree one, two, or three. A rank-one cubic specialization also gives simultaneous unrestricted-refit gains at a prescribed width. An approximation lower bound further shows that certain interaction targets retain nonzero error when ridge neurons are restricted to shared orthogonal axes within the teacher subspace. Population-moment experiments with ReLU students across 21 teachers and 50 initializations per teacher complement the analysis.

Akash Kumar

Department of Computer Science & Engineering University of California-San Diego

Representation Learning ReLU Neural Networks Neural Network Training Dynamics

Ask PDF

Sep 30, 2026cs.LG

Exact information accounting for SGD methods

As an alternative to the standard geometric analyses, we give an exact, information-theoretic analysis of stochastic gradient descent (SGD) and its variants. We show that a preconditioned SGD step is the posterior-mean update of a Gaussian Bayes model, and that its one-step regret splits into an intrinsic-time cost and a change in comparator information. The split extends to an identity for the objective itself. Convex convergence, strict-saddle-point escape, the link between flatness and generalization, the standard learning-rate schedules, adaptive optimizers, and the noisy, momentum, heavy-tailed, and gradient-free variants of SGD each correspond to a term or a special case of this identity. We measure its terms on synthetic and real training runs. On real networks it attributes the slack of classical convergence bounds to the terms their derivations drop and separates optimizers that reach the same training loss. That separation follows the number and consistency of their steps. Its relation to which of them generalizes better differs between networks. For gradient-free SGD the identity determines how a curvature preconditioner should enter the update. The sharpness-based generalization certificate it yields, with a data-independent isotropic prior, is vacuous at network scale unless the curvature spectrum is nearly flat across all parameters.

Akshay Balsubramani

Neural Network Generalization Stochastic Optimization Convergence Adaptive Gradient Methods Stochastic Gradient Descent

Ask PDF

May 21, 2026cs.LG

Why SGD is not Brownian Motion: A New Perspective on Stochastic Dynamics

Stochastic Gradient Descent (SGD) is commonly modeled as a Langevin process, assuming that minibatch noise acts as Brownian motion. However, this approximation relies on a continuous-time limit and a sqrt(eta) noise scaling that does not match the discrete SGD update at finite learning rate. In this work, we propose an alternative formulation of SGD as deterministic dynamics in a fluctuating loss landscape induced by minibatch sampling. Starting directly from the discrete update, we derive a master equation for the parameter distribution and obtain a discrete Fokker--Planck equation that differs from the standard Langevin form at order eta^2. Using this framework, we analyze SGD dynamics near critical points of the loss. We show that the behavior decomposes along the eigenbasis of the mean Hessian into qualitatively distinct regimes. In particular, nearly-flat directions do not admit a stationary distribution: the variance grows over time, corresponding to effective diffusion along valleys with a coefficient proportional to the learning rate. We provide empirical evidence supporting these predictions on neural network models in computer vision and natural language processing, observing a clear qualitative separation between confined and diffusive modes.

Igor Ignashin, Anna Radovskaya, Andrew Semenov +7

Basic Research of Artificial Intelligence Laboratory (BRAIn Lab) · P.N. Lebedev Physical Institute of the Russian Academy of Sciences · Innopolis University

Neural Network Training Dynamics Gradient Descent Dynamics Stochastic Gradient Descent

Ask PDF