Leveraging Flatness to Improve Information-Theoretic Generalization Bounds for SGD

AI summary

Overview

• Research area: Statistical learning theory, specifically information-theoretic (IT) generalization bounds and the flatness bias of Stochastic Gradient Descent (SGD).
• Technical level: Advanced. The paper works with mutual information, sub-Gaussianity, Hessians of empirical and population losses, and minimax excess risk on convex-Lipschitz-Bounded (CLB) problems.
• Scope: The paper derives a new information-theoretic generalization bound for SGD that explicitly and fully depends on loss-landscape flatness, introduces an "omniscient trajectory" auxiliary-trajectory technique, and shows that technique also improves minimax rates for Gradient Descent on CLB problems.

What This Paper Is About

Deep networks trained by SGD generalize well despite over-parameterization, and IT bounds (PAC-Bayesian, mutual-information) are attractive because they are data- and algorithm-dependent, so they can exploit algorithmic properties. The authors observe that existing IT bounds for SGD do not adequately leverage SGD's flatness bias: when flatness is improved by lowering batch size, the true generalization error decreases while the representative bound of Wang and Mao (2022) increases, and the bound is numerically loose. The goal is to build an IT bound for SGD whose dependence on flatness is explicit, correctly directional, and numerically tight.

Key Contributions

1. A new information-theoretic generalization bound for SGD that better leverages the algorithm's flatness bias and is numerically tighter than prior bounds.
2. A statement of how the direction of flatness affects generalization: models generalize better when the large-variance directions of the final weight covariance have small local curvatures in the loss landscape.
3. The "omniscient trajectory" technique, a flexible construction in which the auxiliary trajectory depends on all random variables in the training process, enabling instance-level optimization inside the expectation.
4. Application of the technique to Gradient Descent's minimax excess risk on convex-Lipschitz-Bounded problems, improving representative IT bounds' Ω(1) rates to O(1/√n), which also implies a by-pass of memorization-generalization trade-offs.

Main Findings

• Existing bounds misalign with observed generalization. Figure 1 shows Wang and Mao (2022)'s bound for ResNet-18 on CIFAR-10 under varied flatness: as batch size decreases, the actual generalization error decreases while the bound increases. Their bound splits into a trajectory term and a flatness term; the flatness term captures generalization to some extent, but the trajectory term does not, causing the misalignment. The authors attribute this to the trajectory term's dependence on gradient variance, an implicit measure of flatness that requires extra conditions (for example, near-zero losses) to approximate Hessian traces.
• Two drawbacks in the prior anisotropic bound. Proposition 8 of Neu et al. (2021) optimizes the Gaussian covariance Σ outside the expectation over the output weight W, so Σ cannot adapt to individual W instances, making it suboptimal. It is also hard to estimate: this resembles empirical risk minimization over a Σ living in a Θ(d²)-dimensional manifold, and with only a few sampled W instances (each requiring training a deep model) the estimate suffers severe negative bias from overfitting. The authors state that, to their knowledge, Proposition 2 lacks numerical results.
• The new bound's flatness condition is directional. Intuitively, the algorithm generalizes well when output weights are flat and the flatness aligns with the covariance of output weights, meaning variance is low along sharp directions and high along flat directions, as illustrated in Figures 2(b) and 2(c). The paper states in Section 3.4 that this alignment better leverages flatness than similar alignment notions in Wang and Mao (2024b) and Wang et al. (2021a).
• Numerical tightness. Evaluated on ResNet-18 trained by CIFAR-10. When varying batch size, the bound aligns well with the actual generalization error, and the bound is reported to be only a few percentages looser than the truth across hyperparameters.
• Improved minimax rate for GD. On convex-Lipschitz-Bounded problems, the technique improves representative IT bounds' Ω(1) rates to O(1/√n).
• Memorization-generalization trade-off bypassed. The paper argues that although accurate learners themselves memorize a lot, they are quite close to some oracle learners that memorize little.
• Form of the optimized bound. Theorem 2 decomposes the bound into three expectations corresponding to the penalty, flatness, and trajectory terms, with parameters C = (3/2)((R²/n)|tr(H̃_flat)|)^(1/3) and ΔG derived from ΔH(W_T), Ĥ_S(W_T), and the gradient gap J, and a residual r = O(d²σ⁴) from a second-order approximation (the form of σ is given in Eq. B.48).

Methodology in Plain English

The starting point is the mutual-information bound of Xu and Raginsky (2017), where the generalization error is bounded by 2R²I(W;S)/n under an R-sub-Gaussian loss assumption. Mutual information for SGD can be infinite, so the field uses auxiliary trajectories: perturbed copies of the training trajectory whose mutual information with the data is easier to bound, at the cost of a penalty term measuring the loss difference between the original and perturbed output weights. The paper's Insight 1 is that the ideal perturbation should have large projections onto flat directions (to reduce the trajectory term) and small projections onto sharp directions (to keep the penalty small).

Prior work used independent Gaussian perturbations — isotropic in Neu et al. (2021) and with a shared covariance Σ in their Proposition 8 — which cannot adapt to individual output-weight instances. The authors invoke a simple principle: the minimum of an expectation is at least the expectation of the minimum. Moving the optimization inside the expectation makes it instance-specific, tightens the bound, and removes overfitting as a source of negative bias when estimating the bound by averaging per-instance optima. To achieve this interchange, they let the auxiliary perturbation depend on all random variables in the training process (the training set, internal randomness, gradients, and the trajectory itself), producing the omniscient trajectory; an SGLD-like trajectory is then built on top of it to bound the mutual information. They simplify the T-step perturbation to a single function ΔG (Corollary B.1), then optimize ΔG by second-order approximation, where Hessians appear, and solve an optimization problem formed by the output weights and the Hessians. Experiments control flatness by varying learning rate and batch size on ResNet-18 and CIFAR-10, per prior findings that learning rate and batch size control flatness.

Why This Matters

• Impact on research: The paper attacks a documented limitation of information-theoretic generalization theory — the Ω(1) lower bounds for GD or any accurate learner on some CLB problems highlighted by Livni (2024), Haghifam et al. (2023), and Attias et al. (2024) — and shows its technique reaches an O(1/√n) minimax rate. It also connects flatness to memorization-generalization trade-offs.
• Real-world applications: The paper reports no real-world application studies; its only empirical evaluation is ResNet-18 on CIFAR-10. Domains where a numerically tight, flatness-aware generalization bound could plausibly matter, though not evaluated here, include:
  • Training of large vision models, where batch size and learning rate are used to control flatness.
  • Model selection and hyperparameter tuning when validation data is scarce or expensive.
  • Safety-critical deployments where a reliable pre-deployment generalization estimate is required.
  • Comparing training algorithms on the basis of predicted rather than observed generalization.
• Industry relevance: Because the bound tracks actual generalization error under varied flatness and is reported to be close to the truth, it is a candidate tool for choosing batch size and learning rate and for comparing training algorithms. The authors also stress that unbiased estimation matters for evaluating and comparing bounds and for designing algorithms, which is a practical concern given the cost of training many deep models.

Future Directions

• Scale the evaluation beyond ResNet-18 on CIFAR-10 to larger architectures and datasets; the paper reports no such experiments.
• Quantify the computational and data cost of estimating the bound, which requires training multiple models to obtain the output-weight covariance; the paper does not report this cost.
• Tighten or eliminate the residual r = O(d²σ*⁴) and the third-order differentiability and bounded-residual assumptions used in Theorem 2.
• Combine the omniscient trajectory technique with other bound variants (chaining, individual-sample, conditional mutual information, rate-distortion, Wasserstein) beyond the combinations shown, and extend its empirical validation to SGLD, discretized SDE, momentum, and Adam, which the paper's abstract algorithm form covers but the experiments do not test.

Target Audience

Learning theorists and graduate researchers working on generalization bounds, PAC-Bayesian and mutual-information analyses, and optimization theory for stochastic algorithms. It is also relevant to practitioners who study loss-landscape flatness, batch-size and learning-rate effects, and who need a numerically usable generalization estimate rather than a purely asymptotic statement. Readers need background in information theory, Hessian-based flatness, and convex optimization to follow the derivations.

Authors’ abstract

Information-theoretic (IT) generalization bounds have been used to study the generalization of learning algorithms. These bounds are intrinsically data- and algorithm-dependent so that one can exploit the properties of data and algorithm to derive tighter bounds. However, we observe that although the flatness bias is crucial for SGD's generalization, these bounds fail to capture the improved generalization under better flatness and are also numerically loose. This is caused by the inadequate leverage of SGD's flatness bias in existing IT bounds. This paper derives a more flatness-leveraging IT bound for the flatness-favoring SGD. The bound indicates the learned models generalize better if the large-variance directions of the final weight covariance have small local curvatures in the loss landscape. Experiments on deep neural networks show our bound not only correctly reflects the better generalization when flatness is improved, but is also numerically much tighter. This is achieved by a flexible technique called "omniscient trajectory". When applied to Gradient Descent's minimax excess risk on convex-Lipschitz-Bounded problems, it improves representative IT bounds' $Ω(1)$ rates to $O(1/\sqrt{n})$. It also implies a by-pass of memorization-generalization trade-offs.