Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stability-based Generalization Bounds for Variational Inference

Yadi Wei, Roni Khardon

TL;DR

This paper develops Bayes-stability-based generalization bounds for iterative approximate Bayesian algorithms, including variational inference (VI), when optimized with stochastic gradient descent. It introduces two bound families: TV-based for bounded losses and Wasserstein-based for Lipschitz losses, both reducible to expected parameter differences and expressible as sums of gradient-difference terms along the optimization trajectory. The authors provide a detailed comparison to PAC-Bayes bounds, showing cases where stability-based bounds can be tighter and more data-dependent, and they offer empirical demonstrations on CIFAR-10 with ELBO and DLM that yield non-vacuous bounds and explain performance differences due to data augmentation and algorithm choice. The work highlights a practical pathway to assess and compare generalization in Bayesian neural setups and suggests extensions to broader iterative Bayesian algorithms beyond VI, with limitations mainly tied to SGD-centric analysis. Overall, the approach complements PAC-Bayes analyses and enhances understanding of when approximate Bayesian methods generalize well in deep learning contexts.

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.

Stability-based Generalization Bounds for Variational Inference

TL;DR

This paper develops Bayes-stability-based generalization bounds for iterative approximate Bayesian algorithms, including variational inference (VI), when optimized with stochastic gradient descent. It introduces two bound families: TV-based for bounded losses and Wasserstein-based for Lipschitz losses, both reducible to expected parameter differences and expressible as sums of gradient-difference terms along the optimization trajectory. The authors provide a detailed comparison to PAC-Bayes bounds, showing cases where stability-based bounds can be tighter and more data-dependent, and they offer empirical demonstrations on CIFAR-10 with ELBO and DLM that yield non-vacuous bounds and explain performance differences due to data augmentation and algorithm choice. The work highlights a practical pathway to assess and compare generalization in Bayesian neural setups and suggests extensions to broader iterative Bayesian algorithms beyond VI, with limitations mainly tied to SGD-centric analysis. Overall, the approach complements PAC-Bayes analyses and enhances understanding of when approximate Bayesian methods generalize well in deep learning contexts.

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.

Paper Structure

This paper contains 15 sections, 7 theorems, 38 equations, 3 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Generalization Bounds through Bayes Stability
  5. Bayes Stability
  6. Bounds on Expected Parameter Differences
  7. Discussion: Stability vs. PAC-Bayes Bounds
  8. Experimental Illustration
  9. Expansion Rate
  10. Generalization bounds: ELBO with data augmentation and random labels.
  11. Generalization bounds: ELBO vs. DLM.
  12. Conclusion and Future Work
  13. An explanation why the Proof of Li does not Apply to Variational Inference
  14. Omitted Proofs
  15. An Extreme Scenario for PAC-Bayes Bounds

Key Result

Lemma 2

$\text{TV}\left(\mathbb{E}_{P(X)}[P(Y|X)], \mathbb{E}_{P(X)}[Q(Y|X)] \right) \leq \mathbb{E}_{P(X)}[\text{TV}(P(Y|X), Q(Y|X))]$.

Figures (3)

  • Figure 1: Cumulative expansion rates under various conditions. The left panel displays expansion rates with and without data augmentation, comparing cases with random labels (50% random, labeled as 0.5) and without random labels (labeled as 0.0). The right panel shows expansion rates across different algorithms with data augmentation and no random labels. The shaded areas represent the standard deviation across 10 runs.
  • Figure 2: Generalization error and bounds.
  • Figure 3: Generalization error and bounds for ELBO and DLM with data augmentation and no random labels.

Theorems & Definitions (15)

  • Lemma 2
  • proof
  • Lemma 3: Bayes-Stability 1
  • proof
  • Lemma 4
  • proof
  • Definition 5
  • Lemma 6: Bayes-Stability 2
  • proof
  • Lemma 7
  • ...and 5 more