Non-Vacuous Generalization Bounds: Can Rescaling Invariances Help?

TL;DR

The paper tackles the difficulty of obtaining non-vacuous PAC-Bayes generalization bounds for overparameterized ReLU networks by introducing a rescaling-invariant lifted representation via a measurable lift $\psi$ so that ${f}_w={g}(\psi(w))$. It establishes validity of lifted KL-based bounds through a lifted Donsker–Varadhan change-of-measure and derives a chain of inequalities comparing lifted, stochastic, deterministic, and original divergences to show potential tightening when symmetries are collapsed. Because exact lifted KLs (notably with path+sign lifts) and stochastic rescaling infima are intractable, the authors propose a tractable deterministic rescaling proxy that reduces to a globally convergent one-sided optimization under Gaussian priors and can be solved with a simple block coordinate descent scheme. Empirically, this deterministic rescaling yields substantially tighter PAC-Bayes bounds on MNIST and CIFAR-10, often turning vacuous guarantees non-vacuous, and highlights the potential for invariant priors and lifted analyses to improve generalization guarantees in deep networks.

Abstract

A central challenge in understanding generalization is to obtain non-vacuous guarantees that go beyond worst-case complexity over data or weight space. Among existing approaches, PAC-Bayes bounds stand out as they can provide tight, data-dependent guarantees even for large networks. However, in ReLU networks, rescaling invariances mean that different weight distributions can represent the same function while leading to arbitrarily different PAC-Bayes complexities. We propose to study PAC-Bayes bounds in an invariant, lifted representation that resolves this discrepancy. This paper explores both the guarantees provided by this approach (invariance, tighter bounds via data processing) and the algorithmic aspects of KL-based rescaling-invariant PAC-Bayes bounds.

Figures (2)

• Figure 1: Impact of deterministic rescaling on PAC-Bayes bounds. Left (MNIST): MLPs with varying hidden-layer widths. Right (CIFAR-10): CNN with varying $(\sigma_{\text{prior}}, \sigma_{\text{posterior}})$. Circles: original bounds; diamonds: bounds optimized over deterministic rescaling (which is an upper bound on the lifted $D_{\mathrm{KL}}$ by \ref{['eq:chain-intro']}). The red dashed line marks the non-vacuous threshold ( $<1$ ).
• Figure 2: The information funnel $\mathcal{W} \to \mathcal{Z} \to \mathcal{F}$. Weight-space symmetries (e.g., rescaling) could be collapsed by the lift $\psi$, and the induced map to function space ${f}$ further compresses information. Divergences (e.g., KL) are expected to decrease along this chain, motivating the use of lifted-space bounds.

Theorems & Definitions (4)

• Proposition 1: McAllester’s bound in lifted space
• Definition 2
• Lemma 3: Function and lift invariance under stochastic rescaling
• Theorem 4: Optimized deterministic rescaling for zero-mean Gaussian priors