Tighter CMI-Based Generalization Bounds via Stochastic Projection and Quantization

AI summary

Overview

Research area: Statistical learning theory — specifically information-theoretic generalization bounds, conditional mutual information (CMI), and lossy compression/dimension reduction.

Technical level: Advanced. The paper builds on rate-distortion theory, CMI frameworks, convex optimization classes (SCO, CLB, CSL), and Johnson-Lindenstrauss dimension reduction.

Scope (one sentence): The paper introduces a new CMI-based generalization bound that combines stochastic projection with lossy compression and shows it remains meaningful (decaying as 𝒪(1/√n)) on problem instances where classic MI and CMI bounds were recently shown to fail.

What This Paper Is About

Classic information-theoretic bounds on generalization error — both mutual information (MI) and conditional mutual information (CMI) bounds — were recently shown by Attias et al. [2024] and Livni [2023] to become vacuous or fail to describe the correct generalization behavior on carefully constructed Stochastic Convex Optimization (SCO) instances, leading some to question whether information-theoretic bounds are useful at all. This paper argues that these limitations are not inherent to the CMI framework: by adding a suitable stochastic projection and a lossy compression step, the authors derive new CMI bounds that handle those same counterexamples and decay properly with the training-set size n.

Key Contributions

1. A new CMI-based bound. The authors introduce stochastic projection together with lossy compression into the CMI framework, producing a new bound that is stated to be generally tighter than the CMI bounds of [12].

2. Resolution of known counterexamples. In sharp contrast with classic CMI bounds — which fail on the CLB, CSL, and SCO instances of [43, 46] and may not even decay with n — the new bound yields meaningful results decaying as 𝒪(1/√n).

3. Extension beyond the convex case. Applied to generalized linear stochastic (non-convex) optimization problems in the appendices, the bounds remain non-vacuous, but with a slower decay of 𝒪(1/n^(1/4)) rather than the 𝒪(1/√n) obtained in the convex case studied in [50].

4. A result on memorization. Using stochastic projection and lossy compression, the authors show that for any learning algorithm 𝒜 that memorizes the training data, there exists an alternate algorithm à with comparable generalization error that does not memorize the training data for any data distribution — partly showing memorization is not necessary for good generalization in SCO.

5. Subspace training analysis (appendices). The general bound is applied to subspace training with SGD or SGLD, deriving new bounds based on the differential entropy of Gaussian mixture distributions, where that entropy depends on the gradient difference for training and test datasets, the noise power, the learning rate, and the uncertainty of the training dataset's index within the super-dataset.

Main Findings

• Main theorem (Theorem 1): For any learning algorithm 𝒜: 𝒵ⁿ → 𝒲 with 𝒲 ⊆ ℝᴰ, and for every ε, every d ∈ ℕ, and every projected model quantization set Ŵ ⊆ ℝᵈ, the generalization error is bounded by an infimum over Markov kernels P_{Ŵ|ΘᵀW} and projection distributions P_Θ of 𝔼[√(2 Δℓ_ŵ(S̃,Θ)/n · CMI^Θ(S̃, Â))] + ε, subject to a distortion criterion bounding the expected gap between the original model's and the projected-back model's generalization error by ε.

• CLB counterexample resolved (Theorem 3): For the problem instance 𝒫_cvx^(D) with loss ℓ_c(z,w) = −L⟨w,z⟩, the new bound gives gen(μ,𝒜) ≤ 8LR/√n, whereas the classic CMI bound of equation 6 scales as Θ(LR) and does not decay with n. For ε-learner algorithms with sample complexity Θ(L²R²/ε²), this gives gen(μ,𝒜) = 𝒪(ε).

• One-dimensional projection suffices there: For Theorem 3, projection onto a one-dimensional space (d = 1) is enough, using the Johnson-Lindenstrauss dimension reduction algorithm, with an extra lossy compression step (adding independent noise in the lower-dimensional space) to control the induced distortion.

• CSL counterexample resolved (Proposition 1): For the instance 𝒫_scvx^(D) with loss ℓ_sc(z,w) = −L_c⟨w,z⟩ + (λ/2)‖w‖², the new bound gives gen(μ,𝒜) ≤ 8L_c R/√n. Setting L_c = R = λ = 1 and sample complexity N(ε,δ) = c/ε for ERM, this gives gen(μ,𝒜) = 𝒪(√ε). This contrasts with the earlier finding that CMI(μ,𝒜_n) = Ω(1/ε) for that instance.

• Larger dimensions may be needed elsewhere: For generalized linear stochastic optimization problems, the appendix shows d = Θ(√n) may be needed; in Section 5 and Appendix C, projections with d = n^(2r−1) (r < 1) and d = Θ(log n) are used.

• Non-convex cost: Applying the bounds to generalized linear stochastic (non-convex) optimization, the decay is 𝒪(1/n^(1/4)) rather than 𝒪(1/√n) in the convex case.

• Beyond the framework: The paper notes that well-known techniques such as random-subset and individual-sample methods, or fast-rate improvements in small-empirical-risk regimes, can be applied to Theorem 1 to obtain improved bounds; a single-datum version of Theorem 1 is provided in the supplements.

Methodology in Plain English

The authors start from the CMI "super-sample" construction: 2n data points are drawn, a vector J of n Bernoulli(1/2) variables selects which points form the training set, and the rest act as a ghost/test set. CMI measures how much the trained model reveals about the membership vector J given the super-sample.

Rather than bounding the CMI of the original high-dimensional model W ∈ ℝᴰ, the authors route the model through a two-step pipeline:

1. Stochastic projection: multiply W by a random matrix Θ ∈ ℝ^(D×d) with d ≪ D, obtained independently of the data. This forces the model into a much lower-dimensional space, making the CMI term small. For the CLB example they use the Johnson-Lindenstrauss algorithm.
2. Lossy compression (quantization): apply a kernel P_{Ŵ|ΘᵀW} that maps the projected model to a compressed model Ŵ ∈ Ŵ ⊆ ℝᵈ — effectively adding noise, reminiscent of lossy source coding.

The catch is that projecting and compressing changes the model's generalization error. The authors therefore add a distortion criterion: the expected difference between the generalization error of the original model W and the projected-back model ΘŴ must be at most ε. They show this distortion term can be analyzed with classic concentration inequalities, since it only needs to account for the randomness of the projection matrix, the quantization noise, and the discrepancy between the empirical measure of the sample and the true distribution μ — not for statistical dependencies between the data and the model. The final bound sums the low-dimensional CMI term and this distortion term.

Why This Matters

This work reasserts the value of information-theoretic generalization bounds at a moment when their utility had been questioned. It separates genuine limitations of the CMI framework from artifacts of high dimensionality, showing that a low-dimensional CMI on a projected and compressed model can capture the correct generalization behavior even in over-parameterized regimes where D grows with n (e.g., D = Ω(n⁴ log n) in the studied counterexample). It also contributes to the debate about memorization, showing that memorization is not necessary for good generalization in SCO.

Real-world applications (domains the paper's setting touches):

• Over-parameterized model training, where the hypothesis dimension grows with the number of data samples.
• Subspace training algorithms, including training via SGD or SGLD, analyzed in the appendices.
• Stochastic convex optimization problems such as linear prediction with inner-product losses.
• Strongly convex and Lipschitz-bounded learning problems, relevant to regularized estimation.

Industry relevance: The methods matter for practitioners who need reliable generalization guarantees in settings where model dimensions dwarf the dataset, and for the subfield studying privacy- and memorization-related properties of trained models, since CMI and membership inference are used together in the memorization analysis.

Future Directions

• How far the projection-plus-quantization approach can be pushed toward 𝒪(1/√n) decay in non-convex settings, given the current 𝒪(1/n^(1/4)) rate for generalized linear stochastic optimization.
• Choosing the projection dimension d systematically: the paper uses d = 1 for one counterexample, d = Θ(√n) for generalized linear stochastic optimization, and d = n^(2r−1) or Θ(log n) elsewhere, but the general trade-off between CMI and distortion is left open.
• Combining the new bound with random-subset, individual-sample, and fast-rate techniques to obtain further improvements, as the authors suggest is possible.
• Extending the auxiliary non-memorizing algorithm construction beyond SCO to establish more broadly when memorization is or is not necessary for good generalization.

Target Audience

Researchers in statistical learning theory, information theory, and optimization theory, particularly those working on information-theoretic generalization bounds (MI, CMI, chained CMI, evaluated/f-CMI), rate-distortion compression approaches to learning theory, and the theory of memorization and membership inference. The paper assumes familiarity with CMI frameworks, SCO problem classes (CLB, CSL), and concentration inequalities.

Authors’ abstract

In this paper, we leverage stochastic projection and lossy compression to establish new conditional mutual information (CMI) bounds on the generalization error of statistical learning algorithms. It is shown that these bounds are generally tighter than the existing ones. In particular, we prove that for certain problem instances for which existing MI and CMI bounds were recently shown in Attias et al. [2024] and Livni [2023] to become vacuous or fail to describe the right generalization behavior, our bounds yield suitable generalization guarantees of the order of $\mathcal{O}(1/\sqrt{n})$, where $n$ is the size of the training dataset. Furthermore, we use our bounds to investigate the problem of data "memorization" raised in those works, and which asserts that there are learning problem instances for which any learning algorithm that has good prediction there exist distributions under which the algorithm must "memorize" a big fraction of the training dataset. We show that for every learning algorithm, there exists an auxiliary algorithm that does not memorize and which yields comparable generalization error for any data distribution. In part, this shows that memorization is not necessary for good generalization.