A Hessian-Aware Stochastic Differential Equation for Modelling SGD

Understand

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.