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

Abstract: The standard convergence analysis of mini-batch stochastic gradient descent (SGD) models gradient noise using a single variance term that treats all parameter directions equally, ignoring the fact that noise in high-curvature directions has less impact because learning rates are already constrained there. We introduce Curvature-Weighted Gradient Diversity (CWGD), a geometry-aware measure that weights per-sample gradient diversity by the inverse square root of the Hessian, providing a tighter proxy for the effective optimization noise. For strongly convex quadratic objectives with diagonal Hessians and isotropic noise, we prove that a CWGD-modulated cosine learning-rate schedule can reduce the asymptotic optimization error floor by up to a factor of two compared with standard cosine annealing. We implement this idea as CWGD-Cosine using a Hutchinson-based diagonal Hessian estimator that is exact for quadratic objectives. Across a range of condition numbers, batch sizes, and noise structures, CWGD-Cosine consistently achieves approximately 20% lower final optimization error than standard cosine annealing while incurring negligible overhead in the quadratic setting. We also identify and correct a degenerate curvature estimator, analyze the robustness of the proposed estimator, and explicitly discuss the limitations of the method, including Hessian staleness in non-convex optimization. These results establish CWGD as a principled geometry-aware measure of optimization noise and motivate future extensions to more general learning problems.

• The paper introduces CWGD, a noise measure that adapts SGD to the loss landscape geometry to mitigate over-penalization in high-curvature directions.
• It proposes a CWGD-modulated cosine schedule achieving roughly 20% lower final suboptimality than standard cosine annealing in strongly-convex quadratic scenarios.
• Efficient estimation using Hutchinson’s stochastic trace estimator enables robust curvature rescaling, paving the way for enhanced fine-tuning and adaptive learning rates.

Curvature-Weighted Gradient Diversity: Geometry-Adaptive Noise Modulation for SGD

Overview

The paper "Curvature-Weighted Gradient Diversity: A Noise Measure for Geometry-Adaptive SGD Schedules" (2606.30455) introduces Curvature-Weighted Gradient Diversity (CWGD), a theoretically justified and empirically robust noise measure tailored to the geometry of the loss landscape. CWGD is designed to enable geometry-adaptive SGD schedules by quantifying the true effective noise in optimization, particularly by compensating for over-penalization of high-curvature directions by traditional variance-based noise proxies.

The work details the definition, efficient estimation, and analytical properties of CWGD. It proves a reduction of the residual floor in asymptotic convergence bounds for mini-batch SGD when using a CWGD-modulated cosine learning rate schedule and presents consistent 20% improvements in final suboptimality over plain cosine schedules in the strongly-convex quadratic case.

Motivation and Context

Conventional noise analysis for mini-batch SGD applies a uniform scalar variance bound across all parameter directions. This approach inadequately considers that stochastic noise has a different impact depending on curvature. In directions where the objective function is steep (high curvature), SGD takes small steps due to the learning rate constraint ($\eta \leq 1/L$), causing noise-induced displacement to be inherently small. Nevertheless, the standard theoretical bounds do not capitalize on this, leading to suboptimal step size adaptation.

CWGD resolves this by explicitly rescaling mini-batch gradient diversity with the inverse square root of the Hessian, attenuating the contribution of high-curvature directions. Empirically, this leads to a measurement of effective noise significantly more representative of the true convergence behavior.

Definition and Theoretical Properties

CWGD is formally defined for a twice differentiable objective $f$ with positive definite Hessian as follows:

$$CWGD(x, \mathcal{B}) = \frac{2}{B(B-1)} \sum_{1 \leq i < j \leq B} \| H(x)^{-1/2}(g_i - g_j) \|^2$$

where $g_i$ are per-sample gradients in mini-batch $\mathcal{B}$. For a diagonal Hessian, this reduces to a weighted sum of per-coordinate gradient variances, scaling each variance by $1/\lambda_k$.

Under isotropic noise, the expected CWGD is $2\sigma^2 \operatorname{tr}(H^{-1})$, yielding a direct connection to the effective noise floor experienced during optimization. The theoretical improvement over standard scalar variance bounds is quantified by the factor $\rho(\kappa) = d/(\mu \operatorname{tr}(H^{-1}))$, which increases with the condition number $\kappa$.

CWGD-Modulated Cosine Scheduling and Main Results

A key contribution is the construction of the CWGD-Cosine schedule, where the step size is globally modulated by the normalized CWGD signal:

$$\eta_t = \frac{\bar{\eta}}{1 + \alpha r_t}, \quad r_t = CWGD_t/CWGD_0$$

Theoretical analysis for strongly-convex quadratics with diagonal Hessian and isotropic noise proves that this modulation strictly reduces the asymptotic residual floor by a factor $f$0 over standard cosine annealing, with the recommended setting $f$1 resulting in a two-fold reduction.

Figure 1: Convergence trajectories on synthetic quadratics for $f$2; CWGD-Cosine consistently accelerates convergence and lowers the residual floor compared to standard cosine and step decay schedules.

The gap between this guaranteed improvement and the theoretical optimum grows with increasing $f$3 but is nontrivial for moderate condition numbers ($f$4).

Efficient Estimation and Algorithmic Implementation

CWGD estimation leverages the Hutchinson stochastic trace estimator to obtain diagonal Hessian elements, enabling an efficient and unbiased computation of the curvature weights. Correct use requires separation of curvature and diversity estimation, as naive normalization collapses to an uninformative constant under isotropic noise.

The estimator's robustness to small off-diagonal structure is empirically validated, with relative errors below 1% under moderate non-diagonality and exactness for diagonal quadratic losses.

Empirical Evaluation

CWGD-Cosine is rigorously evaluated on synthetic strongly-convex quadratics across varying condition numbers, batch sizes, and noise structures. It consistently achieves approximately 20% lower final suboptimality than standard cosine annealing, regardless of batch size or isotropic versus aligned noise.

Figure 2: Left: final suboptimality monotonically decreases with increasing modulation strength $f$5; Right: the percentage improvement of CWGD-Cosine over Cosine remains stable across varying batch sizes.

Figure 3: Theoretical improvement factor vs. empirical improvement across condition numbers; CWGD-Cosine tracks the predicted $f$6 reduction at $f$7.

In compute-normalized comparisons—where the wall-clock cost of Hessian estimation is exactly matched by additional gradient steps in the plain cosine baseline—the superiority of CWGD-Cosine persists, confirming that gains are not attributable to increased computation.

Stress testing highlights that the benefit is specific to settings where the optimizer is near the noise floor; in very ill-conditioned or nonconvex configurations, or in the presence of significant Hessian staleness, the merits of the approach are diminished.

Limitations

The main theoretical and empirical scope is restricted to strongly-convex diagonal quadratics under isotropic or moderately aligned noise. In nonconvex problems, the Hessian's rapid evolution degrades CWGD's efficacy due to the staleness of curvature estimates. Furthermore, the theory-practice gap between residual floor reduction and observed improvements is explained by transient dynamics not captured in steady-state analyses.

Key challenges for generalization involve efficient, low-bias online curvature estimation and tracking, especially for off-diagonal and dynamically shifting Hessian structures. The estimator's statistical properties in high dimensionality ($f$8) and large-scale neural network settings remain uncharacterized.

Implications and Future Directions

The results establish CWGD as a principled, geometry-aware noise measure capable of informing adaptive SGD schedules with provable and reproducible advantage within its regime of validity. Practically, CWGD could serve as a late-phase fine-tuning scheduler or as a diagnostic for sudden landscape shifts during optimization.

Extending CWGD's utility to nonconvex or large-scale models mandates the development of online, robust curvature proxies free of staleness and estimator degeneracy. Integrating CWGD-like mechanisms with scalable second-order or adaptive optimizers (e.g., in Adam- or Shampoo-style frameworks) presents a natural direction. Additional theoretical work is required to close the trajectory-level gap between attainable and realized noise reduction factors.

Conclusion

CWGD, together with its application in CWGD-Cosine scheduling, advances the analytical and practical understanding of how geometry and noise interact in stochastic first-order optimization. By moving beyond naive variance averaging and explicitly embedding curvature information into noise measurement and learning rate adaptation, this work sets the stage for more nuanced, data-efficient, and robust optimization strategies in convex machine learning problems. Addressing Hessian staleness and estimator generalization will be critical steps toward broader adoption and impact in deep learning contexts.