A Hessian-aware stochastic differential equation for modelling SGD

• 成果类型: Article; Early Access
• 署名作者: Li, Xiang; Shen, Zebang; Zhang, Liang; He, Niao
• 署名单位: Swiss Federal Institutes of Technology Domain; ETH Zurich
• 刊物名称: MATHEMATICAL PROGRAMMING
• ISSN/ISSBN: 0025-5610; 1436-4646
• DOI: 10.1007/s10107-026-02333-5
• 发表日期: 2026-02-23
• 关键词: Stochastic gradient descent; Continuous-Time Approximation; stochastic differential equation; Escaping Behaviors; DYNAMICAL-SYSTEMS; error

摘要

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 achieves the order-best approximation error guarantee among existing SDE models in the literature, while significantly reducing the dependence on the smoothness parameter of the objective. Empirical experiments on neural network-based loss functions further validate this improvement. 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 provides a more precise characterization of the local escaping behaviors of SGD. With the enhanced approximation guarantee, we further conduct an escape time analysis using HA-SME, showcasing how it can be employed to analytically study the escaping behavior of SGD for general function classes.