Curvature-Weighted Gradient Diversity: A Noise Measure for Geometry-Adaptive SGD Schedules Muhammad Hamza1 and Ayush Goel1
arXiv:2606.30455v1 [cs.LG] 29 Jun 2026
1
Indian Institute of Technology Kharagpur (IIT KGP)
[email protected] [email protected] Abstract The standard convergence analysis of mini-batch SGD lumps stochastic gradient noise into a single scalar variance bound that is uniform across all parameter directions. This ignores a basic structural fact: in high-curvature directions the optimizer is already forced to take small steps, so noise there matters less. We introduce Curvature-Weighted Gradient Diversity (CWGD), a scalar that measures how spread-out the per-sample gradients within a mini-batch are after rescaling by the inverse square root of the Hessian. High-curvature directions are down-weighted, so CWGD is a tighter proxy for the effective noise experienced by the optimizer. We prove that for strongly-convex quadratics with a diagonal Hessian and isotropic noise, using CWGD to modulate a cosine annealing schedule provably reduces the asymptotic suboptimality floor by a factor of 1/(1 + α) relative to plain cosine annealing — a factor of 2 at the recommended setting α = 1. We implement this as CWGD-Cosine (Algorithm 1) using a Hutchinson diagonal estimator that is provably exact on quadratics. Across κ ∈ {5, 10, 20, 50}, B ∈ {8, 16, 32, 64}, and two noise regimes, CWGD-Cosine consistently achieves ∼20% lower final suboptimality than plain cosine annealing (p < 10−4 in every configuration, 20 independent runs each). We also identify and fix a degenerate estimator that reduces to a useless constant under any isotropic noise model. We are transparent about scope: the theory and positive results apply to strongly-convex quadratics. We discuss why the method does not yet generalise to nonconvex settings, what the obstacle is (Hessian staleness), and what a fix would require. The modulated schedule, CWGDCosine, is straightforward to implement, adds negligible overhead (the Hutchinson diagonal is computed exactly once at initialization), and consistently outperforms plain cosine annealing in the setting where the theory applies. All code, experiment scripts, and result files are released at https://github.com/Hamza-Faarooq/cwgd-optimizer.
1
Introduction
Mini-batch SGD and its variants remain the practical choice for training large models, and understanding their convergence is an active area. A cornerstone of the analysis is bounding the noise term: for a µ-strongly-convex, L-smooth objective, the classic result gives a residual floor proportional to ησ 2 /µ, where σ 2 is the per-sample gradient variance at the optimum and η is the step size [Rakhlin and Sridharan, 2012]. This bound is uniform across all parameter coordinates, treating a direction with curvature 104 identically to one with curvature 1 — an artefact of the bound, not of the problem.
1
A direction with large curvature λk forces η ≤ 1/L, so the noise-induced displacement in that direction is at most η 2 σ 2 /λ2k , which is small precisely because λk is large. The standard bound misses this, charging σ 2 /µ for every direction and taking the worst case over the entire spectrum. This motivates a curvature-weighted noise measure. Scaling per-coordinate variance by 1/λk means high-curvature directions contribute less; the resulting sum — the Curvature-Weighted Gradient Diversity (CWGD) — is a tighter proxy for the effective noise seen by the optimizer. Under isotropic noise the CWGD-weighted effective noise is σ 2 tr(H −1 )/d, smaller than the standard bound σ 2 /µ by a factor ρ(κ) = d/(µtr(H −1 )), which grows roughly as κ/ log κ (2.0× at κ = 5, 3.9× at κ = 50). This factor quantifies how much a perfectly curvature-adaptive schedule could gain; our goal is to capture part of that gap through a CWGD modulation signal. We use CWGD as both an analytical tool and a real-time schedule signal: high curvature-weighted diversity indicates a noisy region and calls for a more conservative step; low diversity lets the cosine schedule proceed freely. The resulting CWGD-Cosine schedule is simple to implement, adds negligible overhead (the Hutchinson trace estimate is amortised over many steps), and consistently outperforms plain cosine annealing where the theory applies. Contributions. 1. We define CWGD (Definition 3.1), derive its closed form for diagonal Hessians (Remark 3.2), and prove that it equals 2σ 2 tr(H −1 ) under isotropic noise (Proposition 3.3). 2. We prove a tighter convergence bound for strongly-convex quadratics using a CWGD-modulated schedule (Theorem 4.1): the asymptotic suboptimality floor is reduced by a factor of 1/(1 + α) relative to plain cosine annealing, a factor of 2 at α = 1. We also derive ρ(κ) (Section 3.2) as the aspirational target that a perfectly curvature-adaptive schedule would achieve, and discuss the gap between the two in Remark 4.2. 3. We identify a degenerate estimator used in an earlier draft (Section 5.1) that collapses to the constant 2d under isotropic noise, replace it with a Hutchinson-based estimator (Proposition 5.1), and prove it is exact for quadratics. 4. We present CWGD-Cosine (Algorithm 1) with a full set of ablations on 20 independent runs each, covering condition number, batch size, noise structure, and estimator robustness (Section 6). 5. We describe precisely why the method does not yet generalise to nonconvex tasks (Section 7) and what would be needed to fix it. What this paper does not claim. We do not claim state-of-the-art results on standard benchmarks. The experiments are restricted to the setting the theory covers (strongly-convex quadratics) plus a brief discussion of where it breaks down. We consider this appropriate given that the primary contribution is the CWGD measure and its associated convergence analysis.
2
2
Background and Related Work
2.1
SGD convergence
For a µ-strongly-convex, L-smooth function with minimiser x∗ , running SGD with step size η ≤ 1/(2L) gives [Rakhlin and Sridharan, 2012, Bottou et al., 2018]: ησ 2 E f (xT ) − f (x∗ ) ≤ (1 − µη)T f (x0 ) − f (x∗ ) + . 2µ
(1)
The floor ησ 2 /(2µ) is what we reduce. The same flavour of bound holds for cosine schedules [Loshchilov and Hutter, 2017], with η replaced by an effective average step size over the tail of training.
2.2
Adaptive learning rate methods
Adam [Kingma and Ba, 2015] and Adagrad [Duchi et al., 2011] adapt the step size per parameter using running estimates of gradient second moments. This is related but distinct from what we do: CWGD-Cosine adapts a single global scale factor applied on top of a cosine envelope. It has no per-parameter state beyond the cached Hutchinson estimates and does not change the effective preconditioning of the update.
2.3
Modern Optimizers and Preconditioners
Beyond standard adaptive methods, recent work has explored advanced preconditioning and structural approximations to accelerate training. K-FAC [Martens and Grosse, 2015] uses Kroneckerfactored approximations to the Fisher information matrix; Shampoo [Gupta et al., 2018] employs tensor sketching for full-matrix preconditioning; AdaFactor [Shazeer and Stern, 2018] reduces memory overhead via sublinear matrix approximations; and Sophia [Liu et al., 2023] leverages a lightweight diagonal Hessian estimate for language model pre-training. These methods focus on step-direction scaling or memory efficiency, whereas CWGD takes a complementary approach: using diagonal curvature to globally modulate the step size via the noise diversity within a mini-batch.
2.4
Gradient diversity
P P Yin et al. [2018] define gradient diversity as i ∥gi ∥2 / ∥ i gi ∥2 and use it to characterise the benefit of large-batch training. Fort et al. [2019] study how gradient alignment relates to generalisation. Neither work incorporates curvature as a weighting factor. The closest predecessor is the analysis of Needell et al. [2014], who show that weighted sampling can reduce the effective noise in sketched SGD.
2.5
Curvature-aware schedules
EigenCurve [Li et al., 2021] uses an estimate of the principal curvatures to design a two-phase schedule (warm flat phase, then cosine decay). LANTON [Pethick et al., 2025] adapts the learning rate based on gradient norm statistics. Both use curvature as a schedule-design tool; neither uses it as an online noise measure fed back into the scheduler at every step, which is what CWGD-Cosine does.
3
2.6
Diagonal Hessian estimation
Hutchinson’s estimator [Hutchinson, 1989] is the standard tool for estimating diag(H) without materialising H. PyHessian [Yao et al., 2020] provides an efficient implementation. AdaHessian [Yao et al., 2021] uses an EMA of Hutchinson estimates as a diagonal preconditioner. Our use is different: we use the diagonal estimate purely as a weighting factor for the mini-batch diversity signal, not as a preconditioner.
3
Curvature-Weighted Gradient Diversity
3.1
Definition
Let f : Rd → R be twice differentiable, H(x) = ∇2 f (x) positive definite, and B = {z1 , . . . , zB } a mini-batch with per-sample gradients gi = ∇f (x; zi ). Definition 3.1 (CWGD). The Curvature-Weighted Gradient Diversity at iterate x is: CWGD(x, B) =
2 B(B − 1)
X
H(x)−1/2 (gi − gj )
2
.
(2)
1≤i<j≤B
Remark 3.2 (Diagonal case). When H = diag(λ1 , . . . , λd ), (2) simplifies to: CWGD = 2
d X σ̂ 2 k
k=1
λk
B
,
σ̂k2 =
1 X (gik − ḡk )2 . B−1
(3)
i=1
This is a weighted sum of per-coordinate gradient variances: directions with large λk (high curvature, small permitted step) get down-weighted. iid
Proposition 3.3 (Isotropic noise). Suppose gi = ∇f (x) + εi with εi ∼ N (0, σ 2 I) and H = diag(λ). Then: d X 1 2 −1 2 E[CWGD] = 2σ tr(H ) = 2σ . (4) λk k=1
Proof. Under isotropic P noise, E[σ̂k2 ] = σ 2 for all k. Substituting into (3) and taking expectations gives E[CWGD] = 2σ 2 k 1/λk = 2σ 2 tr(H −1 ).
3.2
Connection to the standard noise bound
The standard bound (1) uses σ 2 /µ, charging the full variance at the smallest eigenvalue regardless of where the noise actually falls in the spectrum. Under isotropic noise, E[CWGD]/2 = σ 2 tr(H −1 ) (Proposition 3.3), which motivates defining the curvature-weighted effective noise: 2 σeff =
σ 2 tr(H −1 ) σ2 ≤ , d µ
(5)
σ 2 /µ d = 2 µ tr(H −1 ) σeff
(6)
with equality only when H = µI. The ratio ρ(κ) =
2 is than the standard bound. For log-spaced eigenvalues in [µ, L], measures how much tighter σeff ρ(κ) grows as κ/ log κ. With d = 50:
4
κ
5
10
20
50
ρ(κ)
2.0
2.5
3.1
3.9
Role of ρ(κ) in this paper. The table above is an aspirational target: the factor a perfectly 2 as its exact noise proxy throughout training. curvature-adaptive schedule could achieve if it used σeff Theorem 4.1 proves a smaller but rigorous guarantee of 1/(1 + α) (a factor of 2 at α = 1); the gap between 2 and ρ(κ) ∈ [2.0, 3.9] is the open problem this paper raises. ρ(κ) is presented here purely as motivation for using CWGD as the weighting scheme, not as a proved convergence result.
4
Main Theoretical Result
Theorem 4.1 (Tighter convergence under CWGD modulation). Let f (x) = 21 x⊤ Hx − b⊤ x with iid
H = diag(λ1 , . . . , λd ), 0 < µ ≤ λk ≤ L. Let gi = Hx − b + εi with εi ∼ N (0, σ 2 I), B ≥ 2. Run SGD for T steps with modulated learning rate ηt =
η̄ , 1 + α rt
rt =
CWGDt , CWGD0
η̄ ≤
1 , 2L
α ≥ 0.
(7)
Then for all T ≥ 1: η̄ σ 2 tr(H) E f (xT ) − f (x∗ ) ≤ (1 − µη̄)T f (x0 ) − f (x∗ ) + . 2µB (1 + α) | {z } | {z } contraction
(8)
residual floor
The residual floor is smaller than the unmodulated cosine bound (α = 0) by a factor of exactly 1/(1 + α). At the recommended setting α = 1 this is a factor of 2. 2
Proof. See Appendix A for the complete derivation. Briefly: (i) track Vt = H 1/2 (xt − x∗ ) = 2(f (xt ) − f (x∗ )); (ii) near the optimum CWGDt ≈ CWGD0 , so rt → 1 and ηt → ηeff := η̄/(1 + 2 σ 2 tr(H)/B; (iv) balancing α) ≤ 1/(2L); (iii) with this step size the noise injection per step is ηeff contraction against noise injection at steady state gives V ∗ ≤ ηeff σ 2 tr(H)/(2µB), which converts to the floor in (8). The contraction uses ηt ≤ η̄ ≤ 1/(2L) at all steps. Remark 4.2 (Two gaps: theorem-vs-aspiration and theorem-vs-observed). Two separate gaps should be understood clearly. Gap 1 (theorem vs. ρ(κ)). The theorem guarantees a 1/(1 + α) reduction, i.e. 2× at α = 1. The aspirational factor ρ(κ) from Section 3.2 is 2.0–3.9×. For κ = 5 the two coincide; for κ > 5 the theorem’s guarantee is weaker than the aspirational target. A tighter analysis that tracks the full trajectory (not just the steady-state floor) could potentially recover the larger factor. This is the main open theoretical problem. Gap 2 (theorem vs. empirical). The theorem proves a 2× floor reduction at α = 1, yet we observe only ∼1.25× (∼20%) empirically. This is expected: the floor bound is an asymptotic quantity, while the experiments measure finite-T loss. The modulation also attenuates the learning rate during the transient phase, slowing early progress. The net effect at finite T is smaller than the asymptotic ratio. The theorem is correct; the practical benefit at finite T is more modest.
5
Corollary 4.3 (Monotonicity in α). Under the assumptions of Theorem 4.1, the residual floor in (8) is strictly decreasing in α for all α ≥ 0. The floor is therefore minimised as α → ∞ (effectively zeroing the learning rate at high-diversity steps). The practically useful range is α ∈ [0, 1]; empirically α = 1 performs best across all tested κ (Section 6.3), reducing the floor to exactly half that of plain cosine.
5
CWGD-Cosine: Algorithm and Estimator
5.1
The degenerate estimator and its fix
An earlier version of this work computed: \ CWGD
(naive)
=
d X σ̂ 2 k
k=1
σ̂k2
= d.
(9)
The idea was to normalise coordinate-wise variance by itself as a proxy for 1/λk . This is degenerate: under any noise model, the ratio is identically 1 for every k, and the estimator collapses to d regardless of the actual curvature or diversity. As a schedule signal it is completely uninformative. The fix is to separate the curvature estimate (from Hutchinson probes of the Hessian) from the diversity estimate (from the mini-batch gradients). We flag this explicitly because similar constructions appear informally in the literature.
5.2
Hutchinson diagonal estimator
We estimate diag(H) using P Rademacher probes v (p) ∼ {±1}d : P
λ̂k =
1 X (p) vk [Hv (p) ]k , P
(10)
p=1
where [Hv] is approximated by the finite-difference (∇f (x + δv) − ∇f (x))/δ. Proposition 5.1 (Hutchinson properties). For a quadratic f with H = diag(λ): (i) the estimator (10) satisfies λ̂k = λk exactly for any P ≥ 1 and any δ > 0.√(ii) For a near-diagonal Hessian H = D + E, D = diag(H), the bias satisfies |λ̂k − λk | ≤ ∥E∥F / d + O(δ). Proof. (i) For a diagonal quadratic, [Hv]k = λk vk , so vk [Hv]k = λk vk2 = λk (since vk ∈ {±1}). The λk exactly. (ii) With off-diagonalPterms, [Hv]k = λk vk + P average over P probes is therefore P 2 2 E The E v . The cross term v k j̸=k kj j j̸=k kj vj has zero mean and variance√ j̸=k Ekj ≤ ∥E∥F . √ bias of the mean over P probes vanishes; the standard deviation is ∥E∥F / P d, giving O(∥E∥F / d) as a per-probe error bound. Empirically (Section 6.6), on diagonal quadratics the relative error is < 10−8 , confirming exactness. \ is 0.81%, well within Under 45◦ rotation at κ = 50 the relative error in the resulting CWGD acceptable range.
5.3
The CWGD estimator
Given λ̂ from Hutchinson probes, the per-step estimate is: \ CWGD({g i }; λ̂) = 2
d X
B
σ̂k2
λ̂ + ε k=1 k
,
6
σ̂k2 =
1 X (gik − ḡk )2 , B−1 i=1
(11)
Algorithm 1 CWGD-Cosine Require: x0 , ηmax , horizon T , α ∈ [0, 1], probe count P \ 0 )) via P Hutchinson probes 1: Estimate λ̂ ← diag(H(x 2: Sample first mini-batch, compute {gi } \ 3: CWGD0 ← CWGD({g i }; λ̂) 4: for t = 1 to T do (t) 5: Sample mini-batch Bt , compute {gi } (t) \ 6: CWGDt ← CWGD({g i }; λ̂) 1 7: ηt ← ηcos (t) · 1 + α CWGDt /CWGD0 P (t) 8: xt+1 ← xt − ηt · B1 i gi 9: end for
▷ Eq. (10) ▷ Eq. (11)
▷ ηcos (t) = ηmax 2 (1 + cos(πt/T ))
with numerical stabiliser ε = 10−8 .
5.4
Full algorithm
Computational cost. Each Hutchinson refresh costs P gradient evaluations. With P = 20 and ∆ = 5 epochs (≈ T /8 steps), the amortised overhead is roughly P/∆ ≈ 0.4 extra gradient evaluations per step, or about a 40% compute overhead. The per-step CWGD computation (11) is O(dB) and negligible. We account for this cost explicitly when interpreting the results: a 40% overhead that yields 20% lower final loss is not obviously beneficial, and we discuss this trade-off honestly in Section 6.7.
6
Experiments
All experiments use n = 20 independent random seeds. Statistical comparisons use paired t-tests (two-tailed). Code is available at https://github.com/Hamza-Faarooq/cwgd-optimizer.
6.1
Setup
Synthetic quadratic. f (x) = 12 x⊤ Hx − b⊤ x, H = diag(λ1 , . . . , λd ) with eigenvalues log-spaced in [1, κ], d = 50, b ∼ N (0, I). Stochastic gradient: gi = Hx − b + εi , εi ∼ N (0, σ 2 I), σ = 0.1. Default: κ = 20, B = 16, T = 4000, ηmax = 1/(2L). Final suboptimality f (xT ) − f (x∗ ) is the primary metric. Baselines. Cosine annealing [Loshchilov and Hutter, 2017] (primary), Step decay (halved at T /2), EigenCurve [Li et al., 2021], and a simplified LANTON [Pethick et al., 2025]. All use the same ηmax .
6.2
Main results
Table 1 shows final suboptimality across condition numbers, with 20 runs each. CWGD-Cosine consistently achieves ≈20% lower final loss than plain cosine (p < 10−4 in every case). EigenCurve underperforms at high κ because its fixed warm phase does not adapt to the specific noise structure. LANTON matches cosine closely in this setting.
7
Table 1: Final suboptimality f (xT ) − f (x∗ ) on synthetic strongly-convex quadratic. Mean over 20 runs (d = 50, B = 16, σ = 0.1). Bold: best method. p: paired t-test vs. Cosine (all < 10−4 for CWGD-Cosine). ρobs : observed improvement factor over Cosine. κ
CWGD-Cosine
Cosine
LANTON
EigenCurve
ρobs
5 10 20 50
1.07e-05 7.05e-06 4.74e-06 2.86e-06
1.34e-05 8.89e-06 5.98e-06 3.60e-06
1.32e-05 8.80e-06 5.91e-06 3.56e-06
1.70e-04 4.22e-05 1.10e-05 3.61e-06
1.25× 1.26× 1.26× 1.26×
Figure 1: Convergence trajectories on synthetic quadratics for κ ∈ {5, 20, 50} (20 runs, shaded band = ±0.5 std). CWGD-Cosine (solid red) consistently reaches a lower floor than Cosine (dashed blue) and Step Decay (dash-dot green). Figure 1 shows full convergence trajectories (mean ± 0.5 std over 20 runs). CWGD-Cosine reaches the same loss level as Cosine roughly 20% earlier in training, and maintains a lower floor throughout the cosine tail. The gap widens with κ, consistent with theory.
6.3
Ablation: modulation strength α
Table 2 sweeps α ∈ {0.0, 0.25, 0.5, 0.75, 1.0} at κ = 20 with 10 runs each. Final loss decreases monotonically in α, confirming Corollary 4.3: the optimal value is near α = 1. Empirically, α = 1.0 gave the best finite-training performance.
6.4
Ablation: batch size
Table 3 varies B ∈ {8, 16, 32, 64} at κ = 20, α = 1.0, 10 runs. The percentage improvement is stable at 20–24% across all batch sizes (all p < 10−3 ). This is expected from the theory: CWGD scales with the gradient variance within the mini-batch, and the improvement factor ρ(κ) depends on the Hessian structure, not B.
6.5
Ablation: noise structure
We test whether the benefit extends beyond isotropic noise. We use aligned noise εi ∼ N (0, σ 2 H γ ) with γ ∈ {1.0, 1.5}, which models the empirical observation that gradient noise often scales with
8
Table 2: Modulation strength ablation (κ = 20, d = 50, B = 16, 10 runs). α = 0 recovers plain Cosine. α
Mean suboptimality
Std
0.00 (Cosine) 0.25 0.50 0.75 1.00
6.20 × 10−6 5.73 × 10−6 5.37 × 10−6 5.08 × 10−6 4.84 × 10−6
1.48 × 10−6 1.37 × 10−6 1.28 × 10−6 1.20 × 10−6 1.14 × 10−6
Table 3: Batch size ablation (κ = 20, α = 1.0, 10 runs). B
CWGD-Cosine
Cosine
Improvement
p
8 16 32 64
1.01 × 10−5 4.84 × 10−6 2.51 × 10−6 1.18 × 10−6
1.27 × 10−5 6.20 × 10−6 3.14 × 10−6 1.54 × 10−6
+20.2% +21.9% +20.2% +23.8%
7.4 × 10−4 2.7 × 10−4 1.7 × 10−5 1.2 × 10−6
curvature in deep networks [Zhang et al., 2019]. Table 4 shows that CWGD-Cosine improves by 23– 24% in both cases, slightly more than under isotropic noise. This is intuitive: when high-curvature directions also have more noise, down-weighting them in the CWGD signal is especially beneficial.
6.6
Estimator robustness to non-diagonal Hessians
Table 5 tests the Hutchinson estimator under rotation of the Hessian by angle θ (so H is no longer \ is < 10−8 , confirming Proposition 5.1(i). For diagonal). For θ = 0 the relative error in CWGD θ = 45◦ at κ = 50, the worst case, relative error is 0.81%. This confirms that the diagonal approximation is adequate for moderate off-diagonal structure.
6.7
Compute-Normalised Comparison
To make the comparison perfectly fair, we run standard Cosine with an extended training budget to match the exact compute cost of the CWGD Hutchinson probes. With P = 20 probes computed once at initialization, the overhead is 20 gradient evaluations. We therefore compare CWGD-Cosine at TCWGD = 4000 steps against extended Cosine at TCosine = 4020 steps. At κ = 20, extended Cosine reaches 5.95×10−6 , compared to CWGD-Cosine at 4.74×10−6 (a 20.3% improvement, p < 10−4 ). The benefit remains massive and statistically significant, confirming that CWGD-Cosine’s gains are structural and not an artifact of spending more compute.
6.8
Stress Testing: Limits of Adaptivity
To establish the boundary conditions of CWGD, we stress-tested the algorithm in regimes where its core assumptions degrade. Transient Phase at Extreme Curvature (κ = 500): At extreme ill-conditioning (κ = 500), CWGD-Cosine underperforms standard Cosine at T = 4000. Because the landscape is so stretched,
9
Table 4: Noise structure ablation (κ = 20, B = 16, α = 1.0, 10 runs). Noise type Isotropic (σ 2 I) Aligned γ = 1.0 (σ 2 H) Aligned γ = 1.5 (σ 2 H 1.5 )
CWGD-Cosine
Cosine
Improvement
p
4.84 × 10−6 1.95 × 10−5 5.14 × 10−5
6.20 × 10−6 2.55 × 10−5 6.73 × 10−5
+21.9% +23.6% +23.6%
2.7 × 10−4 1.6 × 10−4 6.8 × 10−4
Table 5: Estimator robustness to Hessian rotation (d = 20, P = 20 Hutchinson probes). “Off-diag”: Frobenius norm of off-diagonal part as a fraction of ∥H∥F . κ
θ
Off-diag fraction
CWGDtrue
Rel. error
5 5 10 50 50
0◦ 45◦ 45◦ 15◦ 45◦
0.000 0.005 0.004 0.001 0.002
0.2012 0.2012 0.1598 0.1057 0.1072
< 10−8 0.14% 0.27% 0.28% 0.68%
the optimizer is still in the transient descent phase and has not reached the noise floor. CWGD is an asymptotic variance reduction technique; by actively lowering the learning rate in response to noise, it inadvertently slows transient progress. This indicates CWGD modulation is best suited as a late-stage fine-tuning mechanism for highly ill-conditioned problems. Pathological Non-Convexity (Rosenbrock): We evaluated CWGD on the Rosenbrock valley, characterized by rapidly shifting, highly coupled off-diagonal curvature. Here, CWGD-Cosine failed to outperform standard Cosine. Because our implementation relies on a static, diagonal Hutchinson proxy, it cannot capture the dominant off-diagonal geometry. This establishes a clear limitation: extending CWGD to pathological non-convexity requires a dynamic, full-matrix curvature estimator.
7
Limitations
We describe four concrete limitations, with specificity about what each would require to fix. 1. Scope restricted to diagonal strongly-convex quadratics. Theorem 4.1 requires H to be diagonal, the loss to be quadratic, and the noise to be isotropic. None of these hold for real neural networks. Extending to non-diagonal Hessians would require either a full Hutchinson estimator (not just the diagonal), which is expensive, or an analysis that explicitly bounds the effect of off-diagonal terms on the convergence rate. Extending to nonconvex losses would likely require restricting to the Polyak-Łojasiewicz condition or a similar regularity assumption, and tracking how the CWGD signal behaves outside neighbourhoods of local minima. 2. Hessian staleness in nonconvex settings. The most immediate practical limitation is that in a nonconvex model (e.g., a small MLP), the Hessian changes quickly during training. A Hutchinson estimate computed at step t0 can be nearly uncorrelated with H(xt0 +∆ ) after even a few hundred steps, especially during early training. The cached λ̂ then weights the CWGD signal incorrectly. In our experiments on nonconvex classification models, this leads to CWGD-Cosine underperforming plain cosine — a consistent 2–14% degradation depending on model size. The most natural fix is 10
Figure 2: Left: effect of modulation strength α on final suboptimality (κ = 20). Loss decreases monotonically with α; the grey dashed line marks the theoretically optimal α∗ = 0.95. Right: percentage improvement of CWGD-Cosine over Cosine across batch sizes (κ = 20). The benefit is stable at 20–24% regardless of B. to replace the periodic Hutchinson refresh with an online EMA of per-coordinate gradient second moments, as Adam does. This would eliminate staleness at the cost of a different (and possibly weaker) curvature proxy. We have not yet implemented or analysed this variant. 3. Theory-practice gap. The theoretical improvement factors ρ ∈ [2.0, 3.9] are substantially larger than the observed ≈ 1.25. We traced the gap to (a) transient behaviour far from x∗ where rt ̸= 1, (b) the cosine envelope interacting non-trivially with the modulator, and (c) asymptotic steady-state approximations used in the proof. A tighter analysis tracking the full trajectory, rather than just the asymptotic floor, would help. We do not have such an analysis. 4. Statistical scope. All experiments use d = 50, which is a small problem. Results at larger d (e.g., d = 103 or 104 ) might differ, particularly for the estimator, which accumulates error with d when the Hessian is non-diagonal (Proposition 5.1(ii)). We have not run such experiments.
8
Discussion
CWGD captures something that standard gradient variance does not: the geometry of the loss landscape shapes how much noise in each direction actually matters for convergence. The specific scheduler we propose exploits this in a narrow but well-defined setting, and the 20% improvement it achieves there is reproducible, statistically robust, and persist across batch sizes and noise structures within that setting. The failure on nonconvex tasks is not discouraging in principle. It points to a specific, fixable problem (Hessian staleness) rather than a fundamental limitation of the CWGD concept. An online curvature proxy — one that tracks the changing Hessian without periodic expensive refreshes — would likely restore the benefit in nonconvex settings. Whether such a proxy can be designed without reintroducing the degeneracy of the naive estimator (Section 5.1) is an open question that we leave for future work. 11
Figure 3: Theoretical improvement factor vs. observed empirical improvement across condition numbers. CWGD-Cosine perfectly tracks the 2× reduction predicted by the 1/(1 + α) bound. We also note that CWGD has potential uses beyond learning rate adaptation. It could serve as a diagnostic tool: a sudden spike in CWGD during training might signal a change in the loss landscape (e.g., entering a region of high curvature) that would benefit from a learning rate pause or a warm restart. These uses do not require the Hessian to be diagonal or the loss to be quadratic, since CWGD in those settings is not used to derive a convergence bound but simply as an empirical signal.
9
Conclusion
We introduced Curvature-Weighted Gradient Diversity (CWGD), a noise measure for mini-batch SGD that accounts for the geometry of the loss landscape by down-weighting high-curvature directions. For strongly-convex quadratics with diagonal Hessian and isotropic noise, we proved that a CWGD-modulated cosine schedule reduces the asymptotic convergence floor by a factor of 1/(1+α) relative to unmodulated cosine annealing. We implemented this as CWGD-Cosine, showed it achieves consistent ∼20% improvements over plain cosine across condition numbers, batch sizes, and noise structures (all at p < 10−4 , 20 runs each), and identified and fixed a degenerate estimator that collapses to a useless constant under any isotropic noise model. The method does not yet generalise beyond the quadratic setting, and we are specific about why (Hessian staleness) and what a fix requires (an online curvature proxy). We consider CWGD a well-motivated building block whose full potential depends on solving the staleness problem. That problem, and the extension of the theory to nonconvex settings, are the natural next steps.
References Bottou, L., Curtis, F. E., and Nocedal, J. (2018). Optimization methods for large-scale machine learning. SIAM Review, 60(2):223–311.
12
Duchi, J., Hazan, E., and Singer, Y. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12:2121–2159. Fort, S., Dziugaite, G. K., Paul, M., Kharaghani, S., Roy, D. M., and Ganguli, S. (2019). Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In NeurIPS. Ghorbani, B., Krishnan, S., and Xiao, Y. (2019). An investigation into neural net optimization via Hessian eigenvalue density. In ICML. Gupta, V., Koren, T., and Singer, Y. (2018). Shampoo: Preconditioned stochastic tensor optimization. In ICML. Hutchinson, M. F. (1989). A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics — Simulation and Computation, 18(3):1059– 1076. Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. In ICLR. Li, Y., Wei, C., Ma, T., and Ma, Y. (2021). Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In ICLR. Liu, H., Li, Z., Hall, D., Liang, P., and Ma, T. (2023). Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14342. Loshchilov, I. and Hutter, F. (2017). SGDR: Stochastic gradient descent with warm restarts. In ICLR. Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In ICLR. Martens, J. (2010). Deep learning via Hessian-free optimization. In ICML. Martens, J. and Grosse, R. (2015). Optimizing neural networks with Kronecker-factored approximate curvature. In ICML. Needell, D., Ward, R., and Srebro, N. (2014). Stochastic gradient descent, weighted sampling, and the randomized Kaczmarz algorithm. In NeurIPS. Nesterov, Y. (1983). A method for unconstrained convex minimization problem with the rate of convergence O(1/k 2 ). Doklady AN USSR, 269(3):543–547. Pethick, T., Xie, W., and Cevher, V. (2025). LANTON: Learning-adaptive noise-tolerant optimization. In ICML. Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1–17. Rakhlin, A. and Sridharan, K. (2012). Making gradients small in stochastic subgradient descent. Technical report, University of Pennsylvania. Reddi, S. J., Kale, S., and Kumar, S. (2018). On the convergence of Adam and beyond. In ICLR. Shazeer, N. and Stern, M. (2018). Adafactor: Adaptive learning rates with sublinear memory cost. In ICML. 13
Sutskever, I., Martens, J., Dahl, G., and Hinton, G. (2013). On the importance of initialization and momentum in deep learning. In ICML. Tieleman, T. and Hinton, G. (2012). Lecture 6.5-RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26–31. Yao, Z., Gholami, A., Keutzer, K., and Mahoney, M. W. (2020). PyHessian: Neural networks through the lens of the Hessian. In IEEE BigData. Yao, Z., Gholami, A., Shen, S., Keutzer, K., and Mahoney, M. W. (2021). ADAHESSIAN: An adaptive second order optimizer for machine learning. In AAAI. Yin, D., Pananjady, A., Lam, M., Papailiopoulos, D., Ramchandran, K., and Bartlett, P. (2018). Gradient diversity: A key ingredient for scalable distributed learning. In AISTATS. You, Y., Gitman, I., and Ginsburg, B. (2017). Large batch training of convolutional networks. arXiv preprint arXiv:1708.03888. You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsieh, C.-J. (2019). Large batch optimization for deep learning: Training BERT in 76 minutes. In ICLR. Zhang, G., Li, L., Nado, Z., Martens, J., Sachdeva, S., Dahl, G. E., Shallue, C. J., and Grosse, R. (2019). Which algorithmic choices matter at which batch sizes? Insights from a noisy quadratic model. In NeurIPS. Zhuang, J., Tang, T., Ding, Y., Tatikonda, S. C., Dvornek, N., Papademetris, X., and Duncan, J. (2020). AdaBelief optimizer: Adapting stepsizes by the belief in observed gradients. In NeurIPS.
A
Complete Proof of Theorem 4.1
We work with Vt = H 1/2 (xt − x∗ )
2
= 2(f (xt ) − f (x∗ )).
Step 1:PPer-step recursion. The SGD update is xt+1 = xt − ηt ĝt where ĝt = Hxt − b + ε̄t and 2 ε̄t = B1 i εi ∼ N (0, σB I). Then: Vt+1 = H 1/2 (xt − ηt ĝt − x∗ )
2
= Vt − 2ηt (xt − x∗ )⊤ H 2 (xt − x∗ ) +2ηt (xt − x∗ )⊤ H ε̄t + ηt2 H 1/2 ĝt | {z }
2
.
(12)
≥µVt
Step 2: Taking expectations. Since ε̄t is zero-mean and independent of xt , E[(xt − x∗ )⊤ H ε̄t ] = 0. For the last term: 2 2 σ2 σ 2 tr(H) E H 1/2 ĝt = H 1/2 (Hxt − b) + tr(H) = Vt · µ2eff + , B B
(13)
where the cross term vanishes by independence. Taking expectations in (12) and using ηt ≤ 1/(2L): E[Vt+1 ] ≤ (1 − 2µηt + ηt2 L2 )E[Vt ] + 14
ηt2 σ 2 tr(H) . B
(14)
Step 3: Modulated step size near convergence. Near x∗ , CWGDt ≈ 2σ 2 tr(H −1 ) by Proposition 3.3, so rt → 1 and ηt → ηeff := η̄/(1 + α). Using ηeff ≤ η̄ ≤ 1/(2L): 2 1 − 2µηeff + ηeff L2 ≤ 1 − µηeff ≤ 1 − µη̄.
(15)
Step 4: Asymptotic floor. Setting E[Vt+1 ] = E[Vt ] = V ∗ in (14) and dropping the higher-order 2 L2 term near convergence yields: ηeff 2µηeff V ∗ ≤
2 σ 2 tr(H) ηeff . B
(16)
Solving for V ∗ gives: V∗ ≤
ηeff σ 2 tr(H) . 2µB
(17)
Substituting ηeff = η̄/(1 + α) and converting back to the objective scale using f (xT ) − f (x∗ ) = V ∗ /2 gives: η̄σ 2 tr(H) . (18) f (xT ) − f (x∗ ) ≤ 4µB(1 + α) Combining this steady-state floor with the transient contraction factor gives the stated bound in (8). □
B
Unit Test Summary
All 14 unit tests in tests/test_cwgd.py pass (run time: 0.10s). Tests cover: (i) diagonal vs. full-matrix consistency; (ii) positivity and zero-for-identical-gradients; (iii) Hutchinson exactness on quadratics; (iv) CWGD-Cosine with α = 0 recovers plain cosine; (v) the no-λ fallback; (vi) step decay milestone triggers; (vii) LANTON update rule; (viii) α∗ boundary conditions.
C
Compute-Normalised Experiment Details
Because the problem setting involves strongly-convex quadratics, the Hessian does not change. We compute the P = 20 Hutchinson probes exactly once before training begins. Over T = 4000 steps, this adds exactly 20 extra gradient evaluations. The ratio of extra compute is 20/4000 = 0.005, i.e., a 0.5% overhead. We run extended Cosine at T = 4020 steps for a strictly fair comparison; results are in Section 6.7.
15