Technology Category
Reasoning under Uncertainty: Stochastic OptimizationSearch and Optimization: Learning to SearchMachine Learning: Learning with Manifolds
A Hessian-Aware Stochastic Differential Equation for Modelling SGD
🤖 AI Summary
Existing continuous-time SDE approximations of SGD fail to accurately characterize its escape dynamics from stationary points—especially local minima—exhibiting significant deviation on quadratic objectives. To address this, we propose the Hessian-Aware Stochastic Differential Equation (HA-SDE), the first SDE framework that jointly incorporates local Hessian information into both drift and diffusion terms, enabling precise modeling of SGD’s local dynamics near stationary points. Theoretically, HA-SDE exactly reproduces the SGD iterate distribution in the quadratic case—achieving zero-order optimal approximation error—and yields highly accurate escape probabilities and trajectory predictions. Moreover, it substantially weakens dependence on higher-order smoothness constants of the objective. This work establishes a tighter, geometrically informed continuous-time benchmark for analyzing SGD’s generalization mechanisms and optimization dynamics.
Technology Category
Application Category
📝 Abstract
Continuous-time approximation of Stochastic Gradient Descent (SGD) is a crucial tool to study its escaping behaviors from stationary points. However, existing stochastic differential equation (SDE) models fail to fully capture these behaviors, even for simple quadratic objectives. Built on a novel stochastic backward error analysis framework, we derive the Hessian-Aware Stochastic Modified Equation (HA-SME), an SDE that incorporates Hessian information of the objective function into both its drift and diffusion terms. Our analysis shows that HA-SME matches the order-best approximation error guarantee among existing SDE models in the literature, while achieving a significantly reduced dependence on the smoothness parameter of the objective. Further, for quadratic objectives, under mild conditions, HA-SME is proved to be the first SDE model that recovers exactly the SGD dynamics in the distributional sense. Consequently, when the local landscape near a stationary point can be approximated by quadratics, HA-SME is expected to accurately predict the local escaping behaviors of SGD.
Problem
Research questions and friction points this paper is trying to address.
Model SGD escaping behaviors accurately with Hessian-aware SDE
Reduce approximation error dependence on objective smoothness
Characterize local SGD dynamics near stationary points precisely
Innovation
Methods, ideas, or system contributions that make the work stand out.
Hessian-Aware Stochastic Modified Equation (HA-SME)
Incorporates Hessian into drift and diffusion
Order-best approximation error guarantee
🔎 Similar Papers
