A Hessian-Aware Stochastic Differential Equation for Modelling SGD
Fuente:
arXiv
Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866910981073928192 |
|---|---|
| author | Li, Xiang Shen, Zebang Zhang, Liang He, Niao |
| author_facet | Li, Xiang Shen, Zebang Zhang, Liang He, Niao |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_18373 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | A Hessian-Aware Stochastic Differential Equation for Modelling SGD Li, Xiang Shen, Zebang Zhang, Liang He, Niao Machine Learning Optimization and Control 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. |
| title | A Hessian-Aware Stochastic Differential Equation for Modelling SGD |
| topic | Machine Learning Optimization and Control |
| url | https://arxiv.org/abs/2405.18373 |