[Paper Review] Generalization Bounds of SGLD for Non-convex Learning: Two Theoretical Viewpoints

Wenlong Mou, Liwei Wang+2 more|arXiv (Cornell University)|Jul 19, 2017

Stochastic Gradient Optimization Techniques23 references55 citations

TL;DR

The paper derives two algorithm-dependent generalization bounds for SGLD in non-convex learning, using stability and PAC-Bayesian approaches, with bounds that do not explicitly depend on model dimension and rely on aggregated step sizes.

ABSTRACT

Algorithm-dependent generalization error bounds are central to statistical learning theory. A learning algorithm may use a large hypothesis space, but the limited number of iterations controls its model capacity and generalization error. The impacts of stochastic gradient methods on generalization error for non-convex learning problems not only have important theoretical consequences, but are also critical to generalization errors of deep learning. In this paper, we study the generalization errors of Stochastic Gradient Langevin Dynamics (SGLD) with non-convex objectives. Two theories are proposed with non-asymptotic discrete-time analysis, using Stability and PAC-Bayesian results respectively. The stability-based theory obtains a bound of $O\left(\frac{1}{n}L\sqrt{βT_k} ight)$, where $L$ is uniform Lipschitz parameter, $β$ is inverse temperature, and $T_k$ is aggregated step sizes. For PAC-Bayesian theory, though the bound has a slower $O(1/\sqrt{n})$ rate, the contribution of each step is shown with an exponentially decaying factor by imposing $\ell^2$ regularization, and the uniform Lipschitz constant is also replaced by actual norms of gradients along trajectory. Our bounds have no implicit dependence on dimensions, norms or other capacity measures of parameter, which elegantly characterizes the phenomenon of "Fast Training Guarantees Generalization" in non-convex settings. This is the first algorithm-dependent result with reasonable dependence on aggregated step sizes for non-convex learning, and has important implications to statistical learning aspects of stochastic gradient methods in complicated models such as deep learning.

Motivation & Objective

• Understand how stochastic gradient Langevin dynamics (SGLD) affects generalization in non-convex learning.
• Provide non-asymptotic, algorithm-dependent bounds using two theoretical frameworks: stability and PAC-Bayes.
• Show that bounds can be dimension-free and depend on aggregated step sizes rather than parameter norms.
• Connect theory to practical implications for deep learning training where non-convexity and stochasticity are prominent.

Proposed method

• Model the learning objective as regularized empirical risk F_n(w) = (1/n) sum_i f_i(w) + R(w).
• Analyze the SGLD update w_{k+1} = w_k - eta_k g_hat_k(w) + sqrt(2 eta_k / beta) N(0, I_d).
• Employ two analytical frameworks: uniform stability (leading to fast O(1/n) rates) and PAC-Bayesian theory (leading to O(1/√n) rates with trajectory-adaptive terms).
• Relate discrete-time SGLD to a continuous-time Langevin equation and its Fokker-Planck description to bound distributional changes via Hellinger distance and KL divergence.
• Emphasize that the resulting bounds are independent of parameter dimension and rely on aggregated step sizes and gradient norms along the trajectory.

Experimental results

Research questions

• RQ1 How does SGLD influence the generalization error in non-convex learning settings?
• RQ2 Can we obtain non-asymptotic, algorithm-dependent generalization bounds for SGLD using stability and PAC-Bayesian techniques?
• RQ3 Do the bounds depend on aggregated step sizes rather than model dimension or parameter norms, and how do gradient norms along the trajectory affect them?
• RQ4 What are the trade-offs between stability-based and PAC-Bayesian bounds for non-convex stochastic optimization?

Key findings

• Stability-based bound yields an O(1/n) rate, scaling with L, beta, and the square root of the accumulated step sizes.
• PAC-Bayesian bound yields an O(1/√n) rate, with an exponentially decaying factor across iterations and dependence on gradient norms along the trajectory.
• Continuous-time Langevin analysis provides an O(L C sqrt(beta T)/(sqrt{2} n)) bound for the idealized case, highlighting the role of aggregated time T.
• Discrete-time SGLD stability analyses show that with random data sampling, the squared Hellinger distance across neighboring datasets can be controlled, leading to favorable generalization bounds.
• Bounds do not explicitly depend on the dimension of the parameter space or on norms of the parameters, supporting the intuition of “Fast Training Guarantees Generalization” in non-convex settings.