Temperature is All You Need for Generalization in Langevin Dynamics and other Markov Processes

TL;DR

This paper demonstrates that the generalization gap during training with Langevin dynamics depends primarily on temperature and initial loss, not on training duration or model complexity, using a thermodynamic perspective.

Contribution

It provides a novel, simple bound on the generalization gap for Langevin dynamics that is independent of training time, mixing, and model properties, based on thermodynamic principles.

Findings

Generalization gap bound depends on temperature and initial loss

No dependence on training time or model dimensionality

Bound holds for any Markov process with Gibbs stationary distribution

Abstract

We analyze the generalization gap (gap between the training and test errors) when training a potentially over-parametrized model using a Markovian stochastic training algorithm, initialized from some distribution $\theta_0 \sim p_0$. We focus on Langevin dynamics with a positive temperature $\beta^{-1}$, i.e. gradient descent on a training loss $L$ with infinitesimal step size, perturbed with $\beta^{-1}$-variances Gaussian noise, and lightly regularized or bounded. There, we bound the generalization gap, at any time during training, by $\sqrt{(\beta\mathbb{E} L (\theta_0) + \log(1/\delta))/N}$ with probability $1-\delta$ over the dataset, where $N$ is the sample size, and $\mathbb{E} L (\theta_0) =O(1)$ with standard initialization scaling. In contrast to previous guarantees, we have no dependence on either training time or reliance on mixing, nor a dependence on dimensionality, gradient norms, or any other properties of the loss or model. This guarantee follows from a general analysis of any Markov process-based training that has a Gibbs-style stationary distribution. The proof is surprisingly simple, once we observe that the marginal distribution divergence from initialization remains bounded, as implied by a generalized second law of thermodynamics.