PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.11475v1cs.LGmath.OCstat.ML

TL;DR

Population gradient flow may mispredict finite-batch SGD adaptation after pretraining. In a two-unit ReLU regression, the paper traces the gap to shrinking gate disagreement and proves that sufficiently long pretraining can leave SGD failing over an exponential horizon even when GF recovers linearly in T.

  • Problem

    It is unclear whether population gradient flow predicts how long finite-batch SGD needs to adapt after pretraining, despite its common use for reasoning about neural-network adaptation.

    Source: §1 · p. 1

  • Method

    The paper analyzes a two-unit, bias-free ReLU regression and bounds the cumulative probability that online SGD samples gate-disagreement inputs along the exact recursion.

    Source: §1 · pp. 1–2

  • Results

    GF recovers in time linear in T, whereas warm SGD fails with high probability for a horizon of order e^(c/η) once T ≳ log(b/η), while fixed-T small-step SGD recovers.

    Source: §1 · p. 2, §9 · p. 9

  • Takeaways & Limitations

    The discrepancy is tied to the shrinking mass of inputs that distinguish near-cloned ReLU units, which limits SGD's chance to sample a separating input.

    Source: §1 · p. 2, §9 · p. 9

    The separation is proved for two units; in wider networks other units may fit the target, and the conditions under which differing-response input mass controls risk remain open.

    Source: §9 · p. 9

Abstract

from arXiv · show

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 $\eta$ in both phases instead fails with high probability throughout a horizon of order $e^{c/\eta}$ once $T \gtrsim \log(b/\eta)$, 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 $\delta$ these inputs form a wedge of probability $\delta/\pi$, 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/\eta}$ updates, requires $Nb \gtrsim e^{\lambda T}$ target samples and batch size $b \gtrsim \eta e^{\lambda T}$, where $N$ counts updates and $\lambda$ is the weight decay. In simulations, recovery is approximately a function of the disagreement budget $b\delta/\eta$ and saturates in the horizon.

1. Introduction

The paper asks whether population gradient flow predicts finite-batch adaptation after pretraining and shows that, in a two-unit ReLU model, it can mispredict SGD because pretraining makes the separating inputs rare.

1. Introduction

  • The analysis uses a two-unit, bias-free ReLU network in which a directional source task drives positive proportionality and a radial target rewards separated units.The model makes gate disagreement depend on the angle between the units.Source: §1 · p. 1
  • After source training for time T, target GF recovers in time linear in T, but online SGD fails with high probability for a horizon of order e^(c/η) once T ≳ log(b/η).The guarantee holds uniformly on an explicit initialization set with Gaussian probability above 1.06%.Source: §1 · p. 2
  • At the target clone, the population instability comes entirely from inputs where the ReLU gates disagree, whose wedge probability is δ/π.On other inputs, both units receive the same random linear update, contracting their separation in conditional expectation.Source: §1 · p. 2
  • Source training contracts the angle at rate λ, making disagreement inputs increasingly rare and limiting SGD's chance to sample them before further contraction.The resulting disagreement budget is bounded uniformly in the horizon by a constant multiple of b d/η.Source: §1 · p. 2
  • For each fixed pretraining duration, small-step SGD recovers as gradient flow predicts, so the failure arises in the joint limit of small η and growing T.Figure 1 illustrates a source-SGD checkpoint where target GD separates the units while four target SGD streams sample no disagreement and remain near risk 0.6.Source: §1 · pp. 1–2, Figure 1 · p. 2

2. Related Work

The paper distinguishes its result from prior work on exact clone invariance and stochastic attraction: it bounds how finite-batch SGD separates nearby, distinct neurons when the inputs distinguishing them become rare.

2. Related Work

  • Unlike exact-clone invariance results, this work bounds how long SGD takes to separate nearby, distinct neurons produced by finite pretraining.The bound uses the cumulative probability of sampling inputs that distinguish the neurons.Source: §2 · p. 3
  • Prior stochastic-attraction analyses use local coefficients at invariant sets, whereas this result centers on a destabilizing gate-boundary region that shrinks with neuron separation.At the clone, the averaged fixed-input linearization contracts at rate λ independently of η.Source: §2 · p. 3
  • The paper bounds cumulative sampling probability along the exact online recursion, without a diffusion approximation, to separate SGD from GF in the joint limit of vanishing step size and growing pretraining.The recovery prediction fails once pretraining exceeds order log(b/η).Source: §2 · p. 3
  • Compared with population-GF escape from a bisector saddle, finite-batch SGD from the same checkpoint needs exponentially many samples in source-training time.The cited comparator's population escape law is logarithmic.Source: §2 · p. 3

3. Setting and Main Results

The paper studies online SGD and gradient flow in a two-unit ReLU regression with directional source and radial target tasks. It establishes a linear-in-T GF recovery delay alongside a high-probability, exponentially long SGD failure window after sufficient pretraining.

3. Setting and Main Results

  • The setting uses inputs uniform on a disk, a directional source task, a radial target task, and independent fresh batches for online SGD.The disk preserves the population risk and GF equations of Gaussian inputs with matching second moment while bounding sample gradients.Source: §3 · p. 3
  • Proportional units have target risk at least 0.5947, whereas target GF from the specified initialization set converges to risk below 0.27.Reaching risk 0.4 therefore requires leaving a neighborhood of the positively proportional family.Source: §3 · p. 4
  • Target GF has a linear-in-T warm recovery delay, while warm SGD fails throughout an exponential horizon once pretraining is sufficiently long.For small η, the stated window includes T between C0 log(1/η) and c0/η and horizon H = e^(c1 T), with conditional probability at least 0.99.Source: §3 · p. 4
  • The proof's source-phase control combines contraction of projective separation with bounds on unfavorable allocation drift and excursions of shared parameters.Any γ < λ is admissible, at the cost of larger burn-in and constants and a smaller confinement rate and step-size threshold.Source: §3 · p. 4
  • The failure guarantee applies on an explicit initialization set with Gaussian probability above 1%.The conditional confidence can be raised to any level below one by taking smaller ε.Source: §3 · p. 4
  • For every fixed pretraining duration, warm SGD reaches the target risk threshold in bounded time with probability at least 1 − C e^(−c/η).The constants and step-size threshold are independent of batch size but are not uniform in pretraining duration.Source: §3 · p. 4

4. The Instability Lives on the Gate Boundary

The target population instability at the proportional clone is carried entirely by inputs where the two ReLU gates disagree. Because this wedge shrinks with neuron angle, finite-batch updates on the common-gate region instead contract separation.

4. The Instability Lives on the Gate Boundary

  • The target field's destabilizing contribution is carried entirely by the gate-disagreement region, whose probability is proportional to neuron angle.For rotationally uniform inputs, two units at angle δ disagree on a wedge of probability δ/π.Source: §4 · p. 5
  • At λ = 1/20, the boundary contribution q ≈ 0.637 reverses the common-gate rate −0.05.The figure shows both rates along actual trajectories.Source: §4 · p. 5
  • Averaging fixed-input sample linearizations predicts contraction, while the population linearization predicts escape because the moving gate boundary breaks the interchange.The population rate is recovered only through samples from a region of probability δ/π, and both clone-local criteria miss pretraining duration.Source: §4 · p. 5
  • Each neuron activates on half the distribution at every angle, yet the probability their gates differ vanishes as δ/π.Activation frequency therefore misses the loss of distinct responses: common sampled updates rescale the shared output but rarely separate the units.Source: §4 · p. 5
  • Source gradient flow shrinks the angle exponentially, while target flow grows it at rate A after a bounded amplitude transient.The source and target clone amplitudes differ, and the target task-switch transient adjusts shared amplitude while preserving proportionality.Source: §4 · p. 5

5. Identical Checkpoints and Sample Sizes

Starting target gradient flow and target SGD from the same source-gradient-flow checkpoint isolates the effect of target sampling. The paper derives necessary sample and batch-size conditions for fixed-probability SGD recovery and gives a joint scaling where recovery probability vanishes.

5. Identical Checkpoints and Sample Sizes

  • From the identical source-GF checkpoint, fixed-probability recovery within N updates requires b min{N, η^-1} ≥ c_p e^λT.Target gradient flow from that checkpoint reaches risk 0.4 at the time given by equation (3).Source: §5 · pp. 5–6
  • The necessary sample and batch conditions are not sufficient, since sampling a gate disagreement need not trigger escape.The sample condition bounds the total number of target samples, while the batch condition can persist as N grows when later batches meet a smaller wedge.Source: §5 · p. 6
  • For fixed b and T = c_T log(1/η) with c_T > 1/λ, SGD recovery probability tends to zero on horizons H = η^-c_H, while GF recovers in O(log(1/η)) time.The conclusion holds uniformly on K for fixed c_H > 0.Source: §5 · p. 6

6. Proof Overview

The proof separates source and target phases, tracking projective separation during source training and bounding target escape through an auxiliary common-gate process. It combines these estimates with a cumulative disagreement hazard and confinement of shared parameters.

6. Proof Overview

  • During source SGD, positive proportionality is preserved by every sample update, and the projective separation contracts near the source minima.The allocation also moves by at most Cηd_n per step, making its total motion summable after burn-in.Source: §6 · p. 6
  • In the target phase, an auxiliary process applies the same gate and random matrix to both units and matches actual SGD until the first disagreement.Its projective separation contracts in conditional expectation.Source: §6 · p. 6
  • The summed disagreement hazard is at most CbΔ_0 min{N, 4/(λη)}, while shared-parameter confinement contributes C(N + 1)e^-c/η.Confinement is established over blocks of fixed physical time.Source: §6 · p. 6
  • At the checkpoint, Δ_0 ≤ C d_T; combining the source bound with the target estimate yields the theorem, while fresh-start success follows from finite-time tracking of target GF with a strict risk margin..Source: §6 · p. 6

7. Proof of the Target Obstruction

The target obstruction bounds recovery near the proportional family by tracking an auxiliary process that contracts separation and estimating when real SGD can depart from it. The proof also confines shared parameters and translates remaining in the tube into risk above 0.4.

7. Proof of the Target Obstruction

  • Near the proportional family, target SGD reaches risk 0.4 only through a gate disagreement, a projective or allocation exit, or an exponentially rare amplitude-and-balance excursion.The amplitude-and-balance confinement is handled by a block argument.Source: §7 · p. 6

7.1. An Exact Common-Gate Process

  • The common-gate auxiliary process applies one gate to both units and couples exactly to ordinary SGD until a gate disagreement or tube exit.When the gates agree on every batch, the auxiliary update is exactly the ordinary update.Source: §7.1 · p. 7

7.2. Exact Finite-Step Projective Contraction

  • Within a compact tube, the common-gate mean dynamics contract transverse projective separation for every batch size when the step size is sufficiently small.The estimates use bounded samples and constants independent of batch size.Source: §7.2 · p. 7

7.3. Confinement of the Shared Parameters

  • The task switch adjusts shared amplitude along a fixed-duration bridge into a common tube around k = 0 and B = B_t.Confinement of amplitude and balance over fixed-physical-time blocks supplies the bound used in the obstruction.Source: §7.3 · p. 8

7.4. Stopping, Gate Hazards, and Actual Population Risk

  • Fresh batches have conditional disagreement probability at most b C_ang ||ẽ_n||/π, so contraction makes the cumulative hazard saturate in N.The geometric sum caps expected disagreeing samples at order bΔ_0/(λη); states remaining in the tube have risk above 0.4.Source: §7.4 · p. 8

8. Simulations

The simulations show that target SGD recovery from shared pretrained checkpoints depends on the remaining gate-disagreement budget and can stop accumulating over long horizons. Target GD still recovers from the tested checkpoints.

Identical checkpoints

  • At T = 400, 57 of 64 source-SGD checkpoints had Γ_T < 0.04; none sampled a disagreement, and no run recovered after t = 100, whereas GD recovered by about t = 40.The seven recoveries had Γ_T > 1 and near-opposite initializations (initial angle above 2.5); each had one unit with amplitude below 10^-8.Source: §8 · p. 9

Recovery across disagreement budgets

  • Recovery probability approximately follows the disagreement budget Γ = bδ_0/η; at δ_0 = 10^-4 and η = 0.01, increasing b from 1 to 64 raised recovery from 8% to 94%.The Poisson reference captures the transition's location and shape, and lies inside or above every cell's 95% interval.Source: §8 · p. 9

Saturation over long horizons

  • In five settings with Γ from 0.016 to 0.64, no recovery occurred after about t = 112 through t = 450, despite a frozen-angle reference above 0.63 by t = 200 in every setting.Observed recovery saturates in the horizon, unlike the frozen-angle reference.Source: §8 · p. 9

Weight decay

  • With target weight decay λ ∈ {0.01, 0.05}, no first disagreement was observed after about t = 240 and recovery plateaued through H = 1000; without target decay, recovery kept accumulating.The comparison attributes saturation to the contraction supplied by weight decay.Source: §8 · p. 9

9. Conclusion and Future Work

The conclusion attributes the gap between population instability and finite-batch adaptation to shrinking gate disagreements and bounded chances of sampling them. It also identifies limits on how far the two-unit result extends and summarizes the proof's tracking and confinement tools.

9. Conclusion and Future Work

  • In the two-unit model, SGD recovers as η → 0 at each fixed T but fails once T ≳ log(b/η); longer runs cannot replace larger batches over horizons of order e^(c/η).Perturbations acting separately on each unit restore disagreement mass.Source: §9 · pp. 9–10
  • The proof establishes the result for two units; in wider networks, other units may fit the target, so a clone pair need not block learning.For bias-free units with rotation-invariant inputs, two gates disagree with probability δ/π in any dimension; the role of that differing-response mass in risk remains open.Source: §9 · p. 9

A. Finite-Block Tracking and Confinement

  • A single concentration estimate is used for source entry and confinement, the target bridge and confinement, and fresh success.The setup writes each SGD step as population-field drift plus centered batch-gradient error; only the population field is required to be Lipschitz across gates.Source: A. Finite-Block Tracking and Confinement · p. 12

B. Source SGD: Projective Contraction

  • In a compact positive-head source neighborhood, the projective separation contracts in conditional expectation at a rate γ that can be chosen at least λ/2.The local one-step estimate accounts for gate-disagreement corrections without an additive floor.Source: B. Source SGD: Projective Contraction · p. 12, §B.1 · p. 13

B.2. Allocation and Normal Confinement

  • Source allocation exit before normal-tube exit has probability at most C d_0/(γm), while normal exit by time T is bounded by C(1 + T/η)e^(-c/η).These bounds support confinement near the source clone manifold.Source: §B.2 · p. 14

B.3. Entry from the Fixed Gaussian Initialization Region

  • A finite source burn-in sends initializations in K into the inner source tube, with SGD entry probability at least 1 − C e^(-c_burn/η).The burn-in is chosen to make the allocation-exit probability fit a specified allowance.Source: §B.3 · p. 14

C. Target Confinement and the Bridge

  • The target proof first tracks a fixed-duration bridge into a stationary inner tube, then applies block confinement to control excursions.The bridge failure probability is at most C e^(-c/η).Source: C. Target Confinement and the Bridge · p. 15

D. Proof of the Two-Phase Theorem

The proof combines population-flow results with finite-time SGD tracking and confinement to establish the two-phase theorem. In the model, source flow converges to clone minima, while target flow separates the units and reaches low risk.

D. Proof of the Two-Phase Theorem

  • The two-phase proof combines source allocation bounds and target confinement to bound warm-SGD recovery, with necessary sample and batch sizes using e^(γT) in place of e^(λT).The fresh-training argument uses a common finite horizon over the compact initialization set.Source: D. Proof of the Two-Phase Theorem · p. 15

E. Population Equations and Regularity

  • The population equations and regularity analysis bound the trajectories and Hessian along the considered paths, including at coincident directions.These properties support the source- and target-flow convergence arguments.Source: E. Population Equations and Regularity · pp. 16–17

F. Global Source Convergence

  • Every source gradient-flow trajectory from D converges to the clone manifold S, whose points are global source minimizers and normally attracting.On interior compact segments, the five normal decay rates are B_s, 1/2, λ, 2λ, and 2λ.Source: F. Global Source Convergence · p. 17

G. Global Target Convergence and the Risk Gap

  • Target gradient flow from every point of D separates the units, with δ → π and amplitudes ω_1, ω_2 → 2A, reaching limiting risk R_good.A positively proportional pair has radial target risk at least 0.5947152, above the 0.4 threshold.Source: G. Global Target Convergence and the… · p. 18

H. Source Contraction and Target Escape Time

Source training contracts the angle between the units, and target gradient flow needs time to escape near-proportionality. The appendix establishes these bounds uniformly over a compact initialization set.

H. Source Contraction and Target Escape Time

  • Source separation decays as δ_s(T) = Θ(e^(-λT)), while target gradient flow takes time linear in T to reach a fixed small separation δ*.At that separation, both amplitudes lie in a fixed compact positive interval; the remaining time to cross the risk threshold is uniformly bounded.Source: H. Source Contraction and Target… · pp. 19–20

I. The Compact Gaussian Basin and Initial Risk

  • The compact Gaussian initialization basin D_0 consists of positive-head, angularly separated initializations with bounded amplitudes; its probability is bounded using the Gaussian parameter distribution.The explicit compact subset K stays away from zero radii and angular branch cuts.Source: §I · p. 20

J. Experimental Details

The experiments show that warm recovery declines with source-training time, while smaller step sizes delay that decline; recovery fractions across step sizes align when binned by Γ_T = bδ_T/η. Checkpoint and long-horizon tests also connect recovery to sampled gate disagreements.

Identical checkpoints

  • The experiments use shared target input streams across source-optimizer checkpoints and check target risk every ten updates in the identical-checkpoint study.The target stream depends only on initialization and source-training time.Source: J. Experimental Details · p. 21

Constructed checkpoints

  • Every observed recovery in the constructed-checkpoint runs follows a sampled disagreement, but sampling one does not always lead to recovery.Among runs that sample a disagreement, the recovered fraction ranges from 0.81 to 1.Source: J. Experimental Details · p. 21

Long horizons

  • In five long-horizon settings, the last recovery occurs at target times 76.1, 68.1, 103.7, 112.1, and 64.2, respectively.The runs that do not recover continue to show angle decay; higher-precision reruns preserve recovery and first-disagreement times.Source: J. Experimental Details · p. 21

Trajectory

  • At a representative checkpoint with angle 1.67 × 10^-8, target GD’s fitted log-angle slope is 0.5861, whereas target SGD’s is −0.0500.The fitted source-training slope is also −0.0500.Source: J. Experimental Details · p. 21

Joint dependence on T and η

  • At every fixed T, the recovered fraction increases as η decreases, up to sampling fluctuations.The interpolated half-recovery source times are 241.7, 257.9, 265.6, 284.4, and 295.0 across the five step sizes.Source: J. Experimental Details · p. 22
  • Across step sizes, warm recovery fractions fall near one curve when binned by Γ_T = bδ_T/η, while recovery declines with T and later for smaller η.The plotted bins use widths of 0.5 in log10 Γ_T and include at least eight records.Source: Figure 4 · p. 22, J. Experimental Details · p. 22

K. Weight-Decay and Perturbation Controls

Weight-decay and perturbation controls connect recovery to the supply of gate disagreements: decay limits that supply, while independent parameter noise raises recovery in the tested setting. The accompanying notation and method entries specify the model, training parameters, and quantities used to describe the dynamics.

K. Weight-Decay and Perturbation Controls

  • With λ = 0.01 and 0.05, recovery for δ0 = 10^-5 plateaus at 0.29 and 0.04, respectively, and no first disagreement is observed after t ≈ 240 through H = 1000.Without target weight decay, recovery continues growing; population GD reaches risk 0.4 by t ≈ 26 in every setting.Source: K. Weight-Decay and Perturbation… · p. 22, Figure 5 · p. 23
  • At H = 1000, independent noise raises recovery from 0.04 without noise to 0.23, 0.89, and 1.00 for s = 10^-7, 10^-6, and 10^-5, while shared noise yields 0.04, 0.05, and 0.10.The independent-noise runs continue to gain recovery with time.Source: K. Weight-Decay and Perturbation… · pp. 22–23, Figure 5 · p. 23

Model and training

  • The model uses a two-unit, bias-free ReLU network with source label σ(x1), target label ∥x∥, weight decay λ, batch size b, and step size η.Inputs are uniform on the disk, and SGD uses fresh inputs at every update.Source: Model and training · p. 23

Geometry, initialization and risks

  • The angle δ between hidden vectors determines disagreement probability δ/π, while Γ denotes the disagreement budget bδ/η and Γ_T its value after source time T.The source-GF angle is described as δ_s(T) = C_s e^-λT(1 + o(1)).Source: Geometry, initialization and risks · p. 24

Source SGD (Appendix B)

Mechanism and proof overview (Sections 4–6)

Loading 2610.11475v1…