Search papers, labs, and topics across Lattice.
This paper introduces Curvature-Weighted Gradient Diversity (CWGD), a novel measure that adjusts gradient noise by accounting for curvature in the optimization landscape, thereby providing a more accurate representation of effective noise during training. By applying a CWGD-modulated cosine learning-rate schedule, the authors demonstrate a significant reduction in asymptotic optimization error鈥攗p to a factor of two鈥攃ompared to traditional cosine annealing in strongly convex quadratic settings. The proposed CWGD-Cosine method consistently achieves approximately 20% lower final optimization error across various conditions while maintaining low computational overhead, highlighting its practical utility in optimization tasks.
CWGD-Cosine can halve the optimization error floor compared to standard methods, revolutionizing how we approach learning rate schedules in SGD.
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.