Curvature-Weighted
Gradient Diversity:
A Noise Measure for Geometry-Adaptive SGD

Stochastic Gradient Descent has a dirty secret: its standard convergence analysis treats all parameter directions as equally noisy. It doesn't matter whether you're in a direction with curvature $10^4$ or curvature $1$ — the classic bound charges the full gradient variance $\sigma^2/\mu$ everywhere, taking the worst case over the entire spectrum.

This paper argues that's a missed opportunity. A direction with large curvature $\lambda_k$ forces the optimizer to take a tiny step regardless — the noise-induced displacement ends up proportional to $\eta^2 \sigma^2 / \lambda_k^2$, which is small precisely because the curvature is large. The standard bound doesn't see this.

We introduce Curvature-Weighted Gradient Diversity (CWGD) — a scalar that measures how spread-out per-sample gradients are within a mini-batch, but after rescaling by the inverse square root of the Hessian. High-curvature directions are down-weighted. The result is a tighter proxy for the noise actually experienced by the optimizer — and a provably better learning rate schedule when used as a modulation signal.

The Problem with
Uniform Noise Bounds

The cornerstone result for mini-batch SGD on a $\mu$-strongly-convex, $L$-smooth objective gives a residual floor proportional to $\eta\sigma^2/\mu$. Written out:

The floor $\eta\sigma^2/(2\mu)$ uses $\sigma^2/\mu$ — the full per-sample gradient variance divided by the smallest eigenvalue of the Hessian. This is the same bound regardless of whether the noise falls in a high-curvature direction (where the optimizer already moves carefully) or a low-curvature direction (where it can take long strides).

Why Geometry Matters

Suppose the Hessian $H = \text{diag}(\lambda_1, \ldots, \lambda_d)$ is diagonal with eigenvalues spanning a range (the condition number $\kappa = L/\mu$). The step size constraint $\eta \leq 1/(2L)$ is driven by the largest eigenvalue $L$.

For a direction $k$ with curvature $\lambda_k$, the noise-induced displacement per step is at most $\eta^2 \sigma^2 / \lambda_k^2$. When $\lambda_k$ is large, this is tiny — the curvature is already protecting us. The standard bound, by charging $\sigma^2/\mu$ uniformly, dramatically overestimates the noise contribution of high-curvature directions.

Defining CWGD

Let $f: \mathbb{R}^d \to \mathbb{R}$ be twice differentiable with positive-definite Hessian $H(x) = \nabla^2 f(x)$. Let $\mathcal{B} = \{z_1, \ldots, z_B\}$ be a mini-batch with per-sample gradients $g_i = \nabla f(x; z_i)$.

Geometrically, this is the average squared pairwise distance between per-sample gradients, in the metric defined by the inverse Hessian. Directions that the Hessian makes steep are contracted; directions that are flat are amplified.

The Diagonal Case (What We Compute)

When $H = \text{diag}(\lambda_1, \ldots, \lambda_d)$, the formula simplifies dramatically:

This is a weighted sum of per-coordinate gradient variances, with weights $1/\lambda_k$. High-curvature directions (large $\lambda_k$) are down-weighted. This is exactly what we want: directions where the optimizer already takes small steps contribute less to the noise measure.

Under Isotropic Noise

If $g_i = \nabla f(x) + \varepsilon_i$ with $\varepsilon_i \sim \mathcal{N}(0, \sigma^2 I)$, then $\mathbb{E}[\hat{\sigma}^2_k] = \sigma^2$ for all $k$, and:

This gives us a clean connection to $\sigma^2_{\text{eff}}$: the CWGD is exactly $2d \cdot \sigma^2_{\text{eff}}$. High CWGD means the noise is hitting flat directions hard; low CWGD means the geometry is protecting the optimizer.

Main Theorem:
Tighter Convergence

We use CWGD as a real-time signal to modulate a cosine annealing schedule. When the curvature-weighted diversity is high (noisy region), we pull back the learning rate. When it's low (geometry is protecting us), we let the cosine schedule run freely.

Formally, the modulated learning rate is:

Theorem 4.1 (informal): For strongly-convex quadratics with diagonal Hessian, isotropic noise, and $B \geq 2$, running SGD with the modulated schedule above gives:

The residual floor is smaller than plain cosine (the $\alpha = 0$ case) by a factor of exactly $1/(1+\alpha)$. At the recommended setting $\alpha = 1$, this is a factor of 2.

Two Gaps to Understand

The paper is transparent about two gaps that arise from the theorem:

• Theorem vs. aspiration: The theorem guarantees $2\times$ reduction (at $\alpha=1$). The aspirational $\rho(\kappa)$ ranges from $2.0\times$ to $3.9\times$ across $\kappa \in \{5, 10, 20, 50\}$. For $\kappa > 5$, the theorem is weaker than what a perfectly curvature-adaptive schedule could achieve. Closing this gap requires tracking the full trajectory, not just the asymptotic floor.
• Theorem vs. empirical: The theorem proves $2\times$ floor reduction, yet we observe only $\sim 1.25\times$ (~20%) empirically. This is expected: the floor is asymptotic, experiments measure finite-$T$ loss, and the modulation slows early transient progress. The theorem is correct; the practical gain at finite $T$ is more modest.

The Degenerate Estimator
We Fixed

An earlier version of this work computed a naive CWGD estimate by normalizing coordinate-wise variance by itself:

The Fix: Hutchinson Diagonal Estimator

The correct approach separates the curvature estimate from the diversity estimate. We use $P$ Rademacher probes $v^{(p)} \sim \{\pm 1\}^d$ to estimate $\text{diag}(H)$:

where $[Hv]$ is approximated by the finite-difference $(\nabla f(x + \delta v) - \nabla f(x))/\delta$.

Proposition 5.1: For a diagonal quadratic, this estimator satisfies $\hat{\lambda}_k = \lambda_k$ exactly for any $P \geq 1$ and any $\delta > 0$. This is provable: $[Hv]_k = \lambda_k v_k$ (diagonal Hessian), so $v_k [Hv]_k = \lambda_k v_k^2 = \lambda_k$ (since $v_k \in \{\pm 1\}$). The average over probes is $\lambda_k$ exactly. Empirically, the relative error is less than $10^{-8}$ on diagonal quadratics — confirming the theory.

def cwgd_true_diagonal(gradients: np.ndarray,
                       lambda_diag: np.ndarray,
                       epsilon: float = 1e-12) -> float:
    # Down-weight high-curvature directions
    weights = 1.0 / (np.abs(lambda_diag) + epsilon)
    # Per-coordinate gradient variance within the mini-batch
    var_k   = np.var(gradients, axis=0, ddof=1)
    return float(2.0 * (var_k * weights).sum())

CWGD-Cosine
Algorithm 1

The full algorithm is clean and practical. The Hutchinson diagonal is computed once at initialization (costs $P$ gradient evaluations), then cached for the entire run. The per-step CWGD computation is $O(dB)$ and negligible compared to the SGD update itself.

The cosine envelope provides the global annealing schedule, while the CWGD ratio $r_t = \text{CWGD}_t / \text{CWGD}_0$ provides the per-step adaptation. When $r_t > 1$ (noisier than baseline), the effective learning rate drops. When $r_t < 1$ (quieter region), the schedule is only lightly attenuated.

Implementation: The CWGD-Cosine Scheduler

The Python implementation mirrors Algorithm 1 closely. At construction, it accepts a reference to the Hutchinson-estimated eigenvalues (lambda_ref). The first call with gradients initializes cwgd_ref automatically.

class CWGDCosineScheduler:
    """One-Shot Hutchinson CWGD-Cosine (Algorithm 1)"""
    def __init__(self, eta_max, T, alpha=1.0,
                 lambda_ref=None, cwgd_ref=None,
                 eta_min=0.0, epsilon=1e-8):
        self.eta_max   = eta_max; self.T = T; self.alpha = alpha
        self.lambda_ref = lambda_ref; self.cwgd_ref = cwgd_ref
        self._step = 0; self._lr = eta_max
        self._init = cwgd_ref is not None

    def _cwgd(self, grads):
        if self.lambda_ref is None: return 0.0
        w     = 1.0 / (np.abs(self.lambda_ref) + self.epsilon)
        var_k = np.var(grads, axis=0, ddof=1)
        return float(2.0 * (var_k * w).sum())

    def step(self, gradients=None, cwgd=None):
        self._step += 1
        eta_cos = _cosine(self._step, self.eta_max, self.T)
        if self.alpha == 0.0 or self.lambda_ref is None:
            self._lr = eta_cos; return self._lr

        cwgd_t = self._cwgd(gradients) if gradients is not None else (cwgd or 0.0)
        if not self._init:
            self.cwgd_ref = cwgd_t; self._init = True

        gamma_t    = 1.0 / (1.0 + self.alpha * cwgd_t / self.cwgd_ref)
        self._lr   = eta_cos * gamma_t
        return self._lr

Experiments:
Synthetic Quadratics

All experiments use synthetic strongly-convex quadratics where theory applies, making the results cleanly interpretable. The setup:

Main Results

CWGD-Cosine consistently achieves ~20% lower final suboptimality than plain cosine annealing across all tested condition numbers ($p < 10^{-4}$ in every case). EigenCurve underperforms at high $\kappa$ because its fixed warm phase doesn't adapt to specific noise structure. LANTON tracks cosine closely in this setting.

Notice the observed improvement ratio $\rho_{\text{obs}}$ is stable at $1.25$–$1.26\times$ across all $\kappa$. This is consistent with the theory: the $1/(1+\alpha)$ floor reduction gives $2\times$ asymptotically, but finite-$T$ effects and transient slowing reduce the realized gain to ~20%.

Ablations:
What Moves the Needle

Modulation Strength α

Sweeping $\alpha \in \{0.0, 0.25, 0.50, 0.75, 1.00\}$ at $\kappa = 20$ (10 runs each) shows final loss decreasing monotonically in $\alpha$, confirming Corollary 4.3. $\alpha = 0$ recovers plain cosine. $\alpha = 1.0$ gives the best finite-training performance.

Batch Size

Varying $B \in \{8, 16, 32, 64\}$ at $\kappa = 20$, $\alpha = 1.0$ (10 runs). The improvement is stable at 20–24% across all batch sizes, with all $p < 10^{-3}$. This is expected from theory: CWGD scales with gradient variance within the mini-batch, and $\rho(\kappa)$ depends only on Hessian structure.

Noise Structure

We test whether the benefit extends to aligned noise $\varepsilon_i \sim \mathcal{N}(0, \sigma^2 H^\gamma)$ with $\gamma \in \{1.0, 1.5\}$, modeling the empirical observation that gradient noise in deep networks often scales with curvature. CWGD-Cosine improves by 23–24% in both aligned-noise cases, slightly more than under isotropic noise. Intuition: when high-curvature directions also have more noise, down-weighting them in the CWGD signal is especially beneficial.

Estimator Robustness
to Non-Diagonal Hessians

Our Hutchinson estimator assumes a diagonal Hessian. What happens when this is violated? We test by rotating the Hessian by angle $\theta$ (making it non-diagonal) and measuring the relative error in the resulting $\widehat{\text{CWGD}}$.

For the diagonal case ($\theta = 0$), error is below $10^{-8}$, confirming exact recovery on quadratics. Even at $\kappa = 50$ and $\theta = 45°$, the relative error is only $0.81\%$ — well within an acceptable range for a schedule modulator. The diagonal approximation is adequate for moderate off-diagonal structure.

Compute-Normalised Comparison

For a perfectly fair comparison, we run extended cosine annealing at $T = 4020$ steps (to match the 20 extra Hutchinson gradient evaluations) against CWGD-Cosine at $T = 4000$ steps.

Honest
Limitations

The paper is unusually transparent about scope and failure modes. This is worth highlighting — it's a model for how to write limitations sections.

1. Scope restricted to diagonal strongly-convex quadratics. Theorem 4.1 requires $H$ to be diagonal, the loss to be quadratic, and noise to be isotropic. None of these hold for real neural networks. Extending to non-diagonal Hessians requires either a full (expensive) Hutchinson estimator or an explicit bound on off-diagonal effects.
2. Hessian staleness in nonconvex settings. On a nonconvex model (e.g., a small MLP), the Hessian changes quickly during training. A Hutchinson estimate at step $t_0$ can be nearly uncorrelated with $H(x_{t_0 + \Delta})$ after a few hundred steps. In nonconvex experiments, this leads to CWGD-Cosine underperforming plain cosine by 2–14%. The fix is an online EMA of per-coordinate gradient second moments (as Adam does) — but this changes the curvature proxy and hasn't been analyzed.
3. Theory-practice gap. The theoretical $\rho(\kappa) \in [2.0, 3.9]$ is substantially larger than the observed ~1.25×. The gap comes from transient behaviour, the cosine envelope interacting non-trivially with the modulator, and asymptotic steady-state approximations in the proof.
4. Statistical scope. All experiments use $d = 50$. Results at larger $d$ (e.g., $d = 10^3$) might differ, particularly for the estimator, which accumulates error with $d$ when the Hessian is non-diagonal.

Walking Through
the Code

The Colab notebook reproduces every experiment in the paper in about 15 minutes on a free CPU instance. Here's how the pieces fit together.

The Synthetic Quadratic Environment

The problem generator creates a diagonal Hessian with eigenvalues log-spaced in $[1, \kappa]$, a random linear term $b \sim \mathcal{N}(0, I)$, and the closed-form optimum $x^* = b/\text{eigs}$. The stochastic gradient adds isotropic (or aligned) noise to the true gradient.

def make_quadratic(d, kappa, seed=0):
    # Log-spaced eigenvalues: κ measures the condition number
    eigs   = np.logspace(0, np.log10(kappa), d)
    rng    = np.random.RandomState(seed)
    b      = rng.randn(d)
    x_star = b / eigs   # Closed-form optimum for diagonal quadratic
    return np.diag(eigs), b, x_star, eigs

def sample_grads(x, H_diag, b, sigma, B, rng,
                  noise_type='isotropic', gamma_noise=1.0):
    true_g = H_diag * x - b          # True gradient (diagonal H)
    noise  = rng.randn(B, len(x)) * sigma
    if noise_type == 'aligned':
        # Noise scales with curvature: mimics deep network behaviour
        noise *= (np.abs(H_diag) ** (gamma_noise / 2.0))[None, :]
    return true_g[None, :] + noise

The SGD Runner

A single, unified runner handles all schedulers. LANTON gets special treatment since it needs to update its internal noise estimate before stepping. Everything else just passes the current gradients to scheduler.step(gradients=grads).

def run_sgd(H_diag, b, x_star, scheduler, sigma, B, T, seed, ...):
    rng    = np.random.RandomState(seed)
    x      = np.random.RandomState(seed + 9999).randn(len(b)) * 0.5
    f_star = f_val(x_star, H_diag, b)
    losses = np.empty(T)

    for t in range(T):
        grads = sample_grads(x, H_diag, b, sigma, B, rng, ...)
        if is_lanton:
            scheduler.update_noise(grads)    # LANTON needs EMA of noise
            lr = scheduler.step()
        else:
            lr = scheduler.step(gradients=grads)   # CWGD reads grads here
        x = x - lr * grads.mean(axis=0)  # Mini-batch gradient step
        losses[t] = max(f_val(x, H_diag, b) - f_star, 1e-20)
    return losses

Running All Experiments

The notebook runs every ablation in sequence. The total runtime on a free-tier Colab CPU is around 15 minutes. All results are saved as JSON files and plots as high-resolution PNGs.

print("Running Main Results...");       main_results     = run_main_experiment()
print("Running Compute Normalization..."); compute_results  = run_compute_normalised()
print("Running Alpha Ablation...");        alpha_results    = run_alpha_ablation()
print("Running Batch Ablation...");        batch_results    = run_batch_ablation()
print("Running Noise Ablation...");        noise_results    = run_noise_ablation()
print("Running Estimator Ablation...");    estimator_results= run_estimator_ablation()
print("Running Convergence Curves...");    curve_data       = run_convergence_curves()

# Generate all 3 paper figures as high-res PNGs
plot_convergence({int(k): v for k, v in curve_data.items()})
plot_ablations(alpha_results, batch_results)
plot_theory_vs_observed(main_results)

Complete
Proofs

Every formal claim in the paper has a full proof. We reproduce all four here — Proposition 3.3, Theorem 4.1 (the complete Appendix A derivation), Proposition 5.1, and Corollary 4.3 — with step-by-step exposition so the mathematical flow is clear.

Discussion &
Future Directions

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. This is not a small correction — the improvement factor $\rho(\kappa)$ scales as $\kappa/\log\kappa$, potentially 3–4× better than standard bounds for moderately ill-conditioned problems.

The specific scheduler proposed here exploits this in a narrow but well-defined setting, and the 20% improvement it achieves is reproducible, statistically robust, and persists across batch sizes and noise structures within that setting. The overhead is minimal: 0.5% extra compute for a one-time Hutchinson pass.

Beyond Learning Rate Adaptation

CWGD has potential uses beyond schedule modulation. A sudden spike in CWGD during training might signal a change in the loss landscape — entering a region of high curvature — that would benefit from a learning rate pause or a warm restart. This diagnostic use doesn't require the Hessian to be diagonal or the loss to be quadratic.

The Path to Nonconvex Settings

The most natural fix for nonconvex tasks is replacing 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. Designing this proxy without reintroducing the degeneracy of the naive estimator remains an open problem.