Optimizing Information-theoretical Generalization Bounds via Anisotropic Noise in SGLD

TL;DR

This paper improves generalization bounds for SGLD by optimizing the anisotropic noise structure, showing that the optimal noise covariance aligns with the empirical gradient covariance, validated through theoretical and empirical analysis.

Contribution

It introduces a new information-theoretical bound and derives the optimal noise covariance structure for SGLD, enhancing generalization performance.

Findings

Optimal noise covariance is close to the empirical gradient covariance.

The new bound enables better analysis of noise structure in SGLD.

Empirical results validate the theoretical optimal noise structure.

Abstract

Recently, the information-theoretical framework has been proven to be able to obtain non-vacuous generalization bounds for large models trained by Stochastic Gradient Langevin Dynamics (SGLD) with isotropic noise. In this paper, we optimize the information-theoretical generalization bound by manipulating the noise structure in SGLD. We prove that with constraint to guarantee low empirical risk, the optimal noise covariance is the square root of the expected gradient covariance if both the prior and the posterior are jointly optimized. This validates that the optimal noise is quite close to the empirical gradient covariance. Technically, we develop a new information-theoretical bound that enables such an optimization analysis. We then apply matrix analysis to derive the form of optimal noise covariance. Presented constraint and results are validated by the empirical observations.