[Paper Review] On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)

Zhiyuan Li, Sadhika Malladi+1 more|arXiv (Cornell University)|Feb 24, 2021

Neural Networks and Applications46 references18 citations

TL;DR

This paper provides theoretical and empirical validation for modeling Stochastic Gradient Descent (SGD) with Itô Stochastic Differential Equations (SDEs), proposing an efficient simulation method (SVAG) that provably approximates the SDE. It establishes a necessary condition for the SDE approximation to hold—namely, that gradient norm must be smaller than noise variance at equilibrium—and demonstrates that this condition predicts the failure of the linear scaling rule in deep networks with normalization layers.

ABSTRACT

It is generally recognized that finite learning rate (LR), in contrast to infinitesimal LR, is important for good generalization in real-life deep nets. Most attempted explanations propose approximating finite-LR SGD with Ito Stochastic Differential Equations (SDEs), but formal justification for this approximation (e.g., (Li et al., 2019)) only applies to SGD with tiny LR. Experimental verification of the approximation appears computationally infeasible. The current paper clarifies the picture with the following contributions: (a) An efficient simulation algorithm SVAG that provably converges to the conventionally used Ito SDE approximation. (b) A theoretically motivated testable necessary condition for the SDE approximation and its most famous implication, the linear scaling rule (Goyal et al., 2017), to hold. (c) Experiments using this simulation to demonstrate that the previously proposed SDE approximation can meaningfully capture the training and generalization properties of common deep nets.

Motivation & Objective

• To formally justify the use of Itô SDEs as a continuous-time approximation of finite-learning-rate SGD, which is currently heuristic.
• To address the lack of computational feasibility in empirically verifying SDE approximations due to the need for fine-grained time integration.
• To provide a theoretically grounded, testable necessary condition for the SDE approximation to be valid in deep learning settings.
• To investigate the failure of the widely used linear scaling rule (LSR) at large learning rates and batch sizes, particularly in normalization-based networks.
• To demonstrate that the SDE approximation can meaningfully capture training dynamics and generalization behavior of real-world deep networks.

Proposed method

• Proposes SVAG (Stochastic Variance Amplified Gradient), a novel numerical method that simulates the Itô SDE approximation of SGD with provable weak order-1 convergence.
• Uses a continuous-time SDE model: dX_t = -∇L(X_t)dt + (ηΣ(X_t))^{1/2} dW_t, where Σ(X_t) is the gradient noise covariance.
• Derives a necessary condition for the SDE approximation to hold: at equilibrium, the squared gradient norm must be less than the noise variance (‖∇L‖² < Var(∇L_γ)).
• Empirically tests the SDE approximation by comparing SVAG trajectories to actual SGD using deep networks on CIFAR-10, CIFAR-100, and SVHN.
• Applies the condition to predict failure of the linear scaling rule by measuring G_t/N_t (gradient-to-noise ratio) over training phases.
• Uses GroupNorm instead of BatchNorm in NGD experiments to better match gradient noise covariance with SGD.

Experimental results

Research questions

• RQ1 Under what conditions is the Itô SDE approximation of finite-learning-rate SGD mathematically valid?
• RQ2 Can a computationally efficient simulation method like SVAG reliably reproduce the dynamics of SGD and validate the SDE approximation?
• RQ3 Why does the linear scaling rule fail at large learning rates and batch sizes, and can this be predicted by a theoretical condition?
• RQ4 Does the SDE approximation accurately reflect the generalization and training behavior of real deep networks?
• RQ5 Is the gradient norm-to-noise variance ratio a reliable predictor of SDE approximation validity and linear scaling rule failure?

Key findings

• SVAG converges to the Itô SDE approximation and closely tracks SGD trajectories in practice, validating the SDE model as a meaningful approximation.
• The necessary condition for SDE validity—‖∇L‖² < Var(∇L_γ)—is empirically observed to be a strong predictor of linear scaling rule failure.
• The linear scaling rule fails when the gradient-to-noise ratio G_t/N_t exceeds a threshold (C²=2), and this threshold aligns with empirical failure points.
• Experiments on CIFAR-10, CIFAR-100, and SVHN show that SVAG and SGD exhibit nearly identical training and test curves, confirming the SDE approximation's fidelity.
• NGD with matched noise covariance closely replicates SGD performance, further validating the SDE framework when noise structure is preserved.
• The condition G_t/N_t ≈ C² is not only necessary but also close to sufficient for predicting the failure of the linear scaling rule across diverse architectures and datasets.