Stability-based Generalization Bounds for Variational Inference

TL;DR

This paper develops stability-based generalization bounds for variational inference algorithms, including stochastic gradient descent optimization, providing tighter and more applicable bounds for approximate Bayesian methods in neural networks.

Contribution

It introduces a novel stability-based framework for deriving generalization bounds specifically tailored to approximate Bayesian algorithms like VI, filling a key gap in existing analyses.

Findings

Bounds are non-vacuous on neural networks and datasets

The approach can explain performance differences between Bayesian algorithms

Provides tighter bounds than PAC-Bayes in some cases

Abstract

Variational inference (VI) is widely used for approximate inference in Bayesian machine learning. In addition to this practical success, generalization bounds for variational inference and related algorithms have been developed, mostly through the connection to PAC-Bayes analysis. A second line of work has provided algorithm-specific generalization bounds through stability arguments or using mutual information bounds, and has shown that the bounds are tight in practice, but unfortunately these bounds do not directly apply to approximate Bayesian algorithms. This paper fills this gap by developing algorithm-specific stability based generalization bounds for a class of approximate Bayesian algorithms that includes VI, specifically when using stochastic gradient descent to optimize their objective. As in the non-Bayesian case, the generalization error is bounded by by expected parameter differences on a perturbed dataset. The new approach complements PAC-Bayes analysis and can provide tighter bounds in some cases. An experimental illustration shows that the new approach yields non-vacuous bounds on modern neural network architectures and datasets and that it can shed light on performance differences between variant approximate Bayesian algorithms.