Block-Sample MAC-Bayes Generalization Bounds

AI summary

Overview

Research area: Statistical learning theory and information-theoretic generalization bounds, specifically the PAC-Bayes and MAC-Bayes (mean approximately correct) families.

Technical level: Advanced. The paper relies on measure-theoretic probability, Kullback-Leibler divergence, Donsker-Varadhan variational representations, moment generating functions, convexity arguments, and asymptotic order notation.

One-sentence scope: The paper introduces a family of block-sample MAC-Bayes bounds whose divergence terms depend only on subsets ("blocks") of the training data, shows they can be finite where standard PAC-Bayes bounds are vacuous, characterizes the optimal block size order-wise, and proves a negative result on the existence of comparable high-probability versions.

What This Paper Is About

PAC-Bayes bounds control the gap between empirical and population loss with high probability, but they can become vacuous (their right-hand side becomes infinite) for deterministic algorithms and certain priors. This paper builds a generalization of existing MAC-Bayes bounds, which bound the expected generalization error instead, in which the training set is split into blocks and each block contributes its own KL divergence term between the algorithm's conditional distribution and the prior. The goal is to show that this blockwise construction produces finite, and in some regimes tighter, bounds where the original PAC-Bayes bound fails, and to determine how the block size should be chosen.

Key Contributions

1. A block-sample MAC-Bayes bound (Theorem 1). The training set $S$ of length $n$ is partitioned into $J := n/m$ blocks $S_j$ of size $m$, and the bound replaces a single divergence term with the sum $\sum_{j=1}^{J} \mathbb{E}{P{S_j}}D(P_{W|S_j}||Q_W)$, where $P_{W|S_j}$ is the algorithm averaged over all blocks other than $S_j$.

2. Specializations to standard comparator functions (Corollaries 1 and 2). For losses in $[0,1]$ with the Catoni comparator, the paper derives bounds in terms of the binary KL divergence and a square-root generalization bound; for $\sigma^2$-subgaussian losses with the difference comparator, it derives a bound with slightly wider applicability. Substituting the binary KL function directly is shown to be suboptimal because of an extra $\log(2\sqrt{m})/m$ term.

3. Demonstration on a concrete example (Section 4) and a general block-size optimization (Section 5). A Gaussian mean estimation problem under a truncated square loss shows the new bound is finite while the original PAC-Bayes bound is vacuous, and a general analysis relates the growth of the divergence term to the order-wise optimal block size.

4. An impossibility result for high-probability versions (Section 6, Theorem 2). The paper answers negatively whether the block-sample MAC-Bayes bound can be turned into a PAC-Bayes bound of comparable strength, showing that a fast-decaying bound with logarithmic dependence on the error probability is in general not possible.

The authors state that, to the best of their knowledge, this is the first work to present block-based MAC-Bayes bounds in a general framework, optimize the bound, and explore possible PAC-Bayes versions.

Main Findings

• Blockwise divergences can outperform a single divergence term. In the paper's example — $P_Z = \mathcal{N}(\mu,1)$, deterministic algorithm $W := \frac{1}{n}\sum_{i=1}^{n}Z_i$, and loss $\ell(w,z) := K((w-z)^2)$ with $K(x)=x$ on $[0,1)$ and $K(x)=1$ on $[1,\infty)$ — the choice of prior $Q_W := \mathcal{N}(\mu,(n-m)/n^2)$ gives $\mathbb{E}{P_S}D(P{W|S_j}||Q_W) = \frac{m}{2(n-m)}$.

• The resulting bound is finite and decays. Substituting into the square-root bound yields $\mathrm{gen} \leq \frac{1}{2}\sqrt{\frac{1}{2(n-m)}}$, which is $\mathcal{O}(n^{-1/2})$ for any block size other than $m = n$.

• The original PAC-Bayes bound is vacuous here. The case $m = n$ corresponds to the standard PAC-Bayes bound, and the paper reports (with details deferred to Appendix B) that it is vacuous for every choice of prior $Q_W$ in this example.

• Small blocks are optimal in this example; the bound is not highly sensitive to $m$. Figure 1 (plotted for $\mu = 1/2$) shows $m=1$ is optimal, but the bound is not overly sensitive to suboptimal choices of $m$ as long as $m \neq n$. The suboptimal bounds derived from the binary-KL specialization and the subgaussian specialization also decay reasonably when $m$ is adjusted appropriately.

• The optimal block size follows a dichotomy governed by $\gamma$. Under the assumption $\mathbb{E}{P_S}D(P{W|S_j}||Q_W) \leq \frac{\mathcal{O}(m^{\gamma})}{\Theta(n)}$ for some $\gamma \geq 0$ (which holds with $\gamma = 1$ in the paper's example, provided $m \neq n$), Table 1 gives convergence rates for $\mathrm{gen}$ when $m = \Theta(n^{\alpha})$ with $\alpha \in [0,1]$. For $\gamma < 1$ any constant block size is optimal, while for $\gamma > 1$ the block size should grow linearly with $n$ — a regime that includes $m = n$ and $m = n/2$.

• A refinement using the empirical loss. If additionally $\mathbb{E}_{P_S}\hat{L}(W,S) = \mathcal{O}(n^{-\varepsilon})$ for $\varepsilon > 0$, Table 1's second row provides a strictly better rate of the form $\mathcal{O}(n^{\min(\gamma-1,0)-1}) + \mathcal{O}(n^{-\varepsilon}\log n)$.

• No comparable high-probability version exists in general. The paper formulates the desired PAC-Bayes form as $P_S(\mathbb{E}{P{W|S}}d(\hat{L}(W,S),L(W)) \leq A_n + B_n f(1/\delta)) \geq 1-\delta$ with $A_n \to 0$ and $B_n \to 0$. It proves that one cannot establish such a bound which both (a) vanishes at a rate faster than $\mathcal{O}(1/\log n)$ whenever the proposed MAC-Bayes bound vanishes at rate $\mathcal{O}(n^{-1/2})$, and (b) exhibits a logarithmic dependence on the permitted error probability.

• A computational caveat. Like other information-theoretic bounds (the paper cites Bu et al. (2020) and Negrea et al. (2019)), the new bounds depend on the distribution of the training data. This can make them tighter than distribution-independent bounds when partial statistical knowledge is available, but potentially non-computable when it is not; Section 7 discusses these limitations in more detail.

Methodology in Plain English

The authors start from the standard setup of statistical learning: data $S=(Z_1,\dots,Z_n)$ drawn i.i.d. from an unknown distribution $P_Z$, a loss $\ell(w,z)$, an empirical loss $\hat{L}$, and a population loss $L$. They split $S$ into $J = n/m$ consecutive blocks of equal size $m$.

The key technical move is to bound the generalization error block by block. Because the comparator function $d$ is assumed jointly convex, the average over all blocks can be pushed inside $d$ using Jensen's inequality, and the average over the data outside block $j$ can be folded into the algorithm, giving the averaged conditional distribution $P_{W|S_j} := \mathbb{E}{P{S_1},\dots,P_{S_{j-1}},P_{S_{j+1}},\dots,P_{S_J}}P_{W|S}$. A change of measure from $P_{W|S_j}$ to the prior $Q_W$ then uses the Donsker-Varadhan variational representation of KL divergence, followed by Fubini-Tonelli and another Jensen step for the logarithm. The result is a bound whose divergence term is a sum of $J$ per-block quantities rather than one global KL term.

For the concrete example, the authors compute the KL divergence in closed form for Gaussian distributions, which yields the simple rate $\frac{m}{2(n-m)}$ and the resulting square-root bound. For the general analysis, they state an assumption on how the per-block divergence grows with $m$ and $n$, then optimize over block sizes $m = \Theta(n^{\alpha})$ and tabulate the resulting convergence orders. Finally, for the impossibility question they construct a learning scenario and show that any high-probability bound of the desired shape must sacrifice either the decay rate or the logarithmic dependence on $\delta$.

Why This Matters

Impact on research. The work extends the MAC-Bayes line (traced back to Alquier (2006) and Catoni (2007), and named by Grunwald et al. (2021)) with a blockwise construction inspired by the individual-sample information-theoretic bound of Bu et al. (2020). It gives a general recipe rather than a single bound, since Theorems 1 can be instantiated with different comparator functions, and it delineates a limit on how far the approach can be pushed into the high-probability regime — a boundary result relevant to the broader literature on information-theoretic generalization bounds.

Real-world applications (as framed by the paper; the authors stress that substantial additional research is needed before applying the bound to complex algorithms, and demonstrate only on a simple example):

• Image classification with neural networks, the motivating example in the introduction, where $S$ is a sequence of images with class labels and $W$ is the set of weights and biases.
• Training via stochastic gradient descent, whose (possibly stochastic) mapping from training set to hypothesis is modeled as $P_{W|S}$.
• Gaussian mean estimation under a truncated square loss (the paper's numerical example, a loss type the authors note has been analyzed and applied to practical ML algorithms by Le & Zach (2021)).
• General distribution-dependent settings where partial statistical knowledge about the training data is available, where the authors argue the bound can potentially beat distribution-independent alternatives.

Industry relevance. Tight generalization bounds are proposed in the paper as an alternative to relying on held-out test sets, which shrink the training set, and whose repeated use for algorithm tuning technically requires fresh independent data. Where such bounds can be computed, they address a practical bottleneck in model selection and validation. Because the current bound requires distributional knowledge and is demonstrated only on a toy problem, its near-term industrial relevance is prospective rather than demonstrated.

Future Directions

1. Verify the divergence-growth assumption for realistic algorithms. Assumption (15) on the growth of $\mathbb{E}{P_S}D(P{W|S_j}||Q_W)$ has to be checked case by case; the authors state that substantial extra work is needed to apply the bounds to complex algorithms such as deep networks.

2. Develop computable versions of the bound. The bounds depend on the training data distribution and may not be computable when such knowledge is unavailable (Section 7), so finding tractable distribution-independent surrogates or partial-knowledge schemes is an open problem.

3. Determine whether block-sample PAC-Bayes versions are possible under restrictions. The impossibility result rules out a broad class of high-probability bounds of the form in (16); the natural follow-up is to ask whether narrower settings, or bounds with different functional forms, can recover the blockwise tightness.

4. Reconcile with the related impossibility result of Harutyunyan et al. (2022). That result resembles Theorem 2 but concerns a different system setup and intersects with the present work only at block size $m=1$; for $m>1$ it is not applicable, leaving scope for further comparison.

Target Audience

Researchers and graduate students in statistical learning theory, information theory, and PAC-Bayes analysis who are already comfortable with KL divergences, comparator functions, and generalization bounds. The paper is less suited to practitioners seeking a ready-to-deploy tool, given that it demonstrates the approach on one simple Gaussian mean estimation example and explicitly notes that solving the practical issues for complex algorithms requires substantial additional research.

Authors’ abstract

We present a family of novel block-sample MAC-Bayes bounds (mean approximately correct). While PAC-Bayes bounds (probably approximately correct) typically give bounds for the generalization error that hold with high probability, MAC-Bayes bounds have a similar form but bound the expected generalization error instead. The family of bounds we propose can be understood as a generalization of an expectation version of known PAC-Bayes bounds. Compared to standard PAC-Bayes bounds, the new bounds contain divergence terms that only depend on subsets (or \emph{blocks}) of the training data. The proposed MAC-Bayes bounds hold the promise of significantly improving upon the tightness of traditional PAC-Bayes and MAC-Bayes bounds. This is illustrated with a simple numerical example in which the original PAC-Bayes bound is vacuous regardless of the choice of prior, while the proposed family of bounds are finite for appropriate choices of the block size. We also explore the question whether high-probability versions of our MAC-Bayes bounds (i.e., PAC-Bayes bounds of a similar form) are possible. We answer this question in the negative with an example that shows that in general, it is not possible to establish a PAC-Bayes bound which (a) vanishes with a rate faster than $\mathcal{O}(1/\log n)$ whenever the proposed MAC-Bayes bound vanishes with rate $\mathcal{O}(n^{-1/2})$ and (b) exhibits a logarithmic dependence on the permitted error probability.