Minimum Description Length and Generalization Guarantees for Representation Learning

Summary

This paper obtains new generalization bounds of the conditional mutual information-flavor (Steinke and Zakynthinou, 2020), where the bounds depend only on the predictions that the algorithm induces on the supersample. These bounds are phrased more generally in terms of a KL divergence with symmetric priors, analogous to the almost exchangeable priors of Audibert (2004), and are placed in a rate-distortion theoretic perspective. This is further extended to lossy compression and tail bounds. Tighter bounds in terms of a function of the binary entropy are derived, which appear similar to classical PAC-Bayes bounds with the binary KL divergence in the style of Maurer (2004). Throughout, this is related to the representation in the style of the information bottleneck, leading to the main result of the paper: a generalization bound for representation learning algorithms which depends only on the encoder part, allowing for the decoder part to be chosen freely to minimize the training loss without affecting the bound. This allows for priors that are in some senses data-dependent, and the potential benefit of this is demonstrated numerically. **Post-Rebuttal edit**: I read the author rebuttal, which addressed my questions, in particular regarding related literature. I consequently updated the rating (as stated below).

Strengths

— The new results and their relevance for representation learning are well-established. — The theoretical analysis is thorough, covering several possible variants and tightenings of the results. — The presentation is overall very clear and pedagogical, including a very thorough discussion of some related prior work. — The numerical evaluation nicely illustrates the potential benefit of the proposed methods.

Weaknesses

— There are some missing connections to previous work (see questions). — A few of the results appear to be stated incorrectly, and some notation seems undefined (see questions).

Questions

— There appears to be a clear connection to the frameworks of functional and evaluated CMI (f-CMI, e-CMI) of, e.g., Harutyunyan et al (2021) and Hellström and Durisi (2022). Is there a clear distinction between these bounds and your Theorem 1, other than the use of an auxiliary prior (the possibility of which is also noted in the mentioned works)? — The concept of symmetric priors seems identical to the almost exchangeable priors discussed by Audibert (2004). Is there any difference with these? — The bounds in terms of the function $h_D$ appear closely related to bounds with the binary KL divergence on the left-hand side, e.g. Maurer (2004) and Seeger (2002) — the properties of the function appear to be very similar. Is there a difference or are the two approaches equivalent? — Proposition 1: The left-hand side here is a single-draw result, which typically does not lead to bounds in terms of the KL divergence, but instead some form of information density or log-likelihood ratio. This is also possible to derive from your results, but the derivations have to be slightly altered. In the proof of Theorem 4, you instead consider a “PAC-Bayesian” bound which is averaged over the posterior, which naturally leads to bounds in terms of the KL divergence. I assume this is what you meant to write in Proposition 1 as well. As noted by Grunwald et al (2023), such PAC-Bayesian bounds follow immediately from the same derivation as average generalization bounds if they are based on “exponential inequalities” (as yours are). This can be extended to single-draw bounds as well, under more stringent absolute continuity assumptions. Minor: Abstract: IB undefined Intro: $\text{gen}(W,S)$ and the corresponding notation for population and training losses (from e.g. Proposition 1) appear to be undefined Quotation mark lines 34, 234, 33 Line 50: A sentence with only citations. Line 131: KL divergence notation is flipped in LHS of definition Fig. 1 is not on the same page as it is discussed Margin overflow in Eq. (3), lines 579-581, more in appendix Line 242: $\tau*$ instead of $\tau^*$ Lines 266-272: inconsistent conventions for fractions. Can some expressions be moved to equation environments? Maybe this intuition can be expressed more informally and the math can be moved to the appendix where there is space I assume $[a,b]=\{a,a+1,\dots,b\}$ for integers in some expressions, but I do not think this is defined and it conflicts with the standard use for closed intervals Theorem 3: The prior is called $P$ in the theorem statement but $Q$ in the equation. Line 329, 686: Incomplete sentence References: Audibert, “A better variance control for PAC-Bayesian classification”, 2004 Maurer, “A Note on the PAC Bayesian Theorem”, 2004 Harutyunyan et al, "Information-theoretic generalization bounds for black-box learning algorithms”, 2021 Hellstrom and Durisi, “A New Family of Generalization Bounds Using Samplewise Evaluated CMI”, 2022 Seeger, “PAC-Bayesian Generalisation Error Bounds for Gaussian Process Classification”, 2002 Grunwald et al, “Exponential Stochastic Inequality”, 2023

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Perhaps the discussion of limitations can be expanded, but it is not crucial.

Summary

The paper advances state-of-the-art on information theoretic generalization guarantees for a class of representational learning problems. The contribution is positioned in the framework of information bottleneck approach and builds on the classical work of Blum and Langford for PAC-MDL bounds. The main result is on non-trivial upper bounds on generalization error of a representational learning algorithm in terms of so-called minimum description length of the label. The result is expressed in terms of KL-divergence of a vector of labels and any arbitrary symmetric prior, and resolves the issues related to use of mutual information measure. The aspects related to block coding and lossy compression are new to this type of results.

Strengths

The writing is adequately clear both in terms of technical precision and intuitive interpretations, where appropriate. The problem is important and the proposed results advance state-of-the-art.

Weaknesses

Perhaps a weakness of the paper is that it relies on a solid understanding of information theoretic background; for e.g., [GK19]. More importantly, Theorem 5 which is main result on generalization bound for representation learning needs a bit more explanation, both in terms of intuitive/conceptual ideas and technical contribution based on work in the previous section. The proposed result – as the authors claim – resolves issues with the information bottleneck method, but the reader gets limited understanding of how the complexity of latent space and validity of any choice of decoder plays out in practical settings. On the other hand, the discussion in Section 2 could be more pointed – with focus on key novelties (which might help save space and leave room for more extensive discussion in Section 3). It also appears that a tabular comparison of generalization bounds proposed in this paper, relative to the ones based on mutual information would clarify some of the issues regarding positioning of the results. Unless I am mistaken, it seems the authors do not formally introduce the key object $\mathbb{E}[gen(S,W)]$.

Questions

1. Better explain the implications of Theorem 5 relative to generalization based on mutual information. Why *should* the encoder structure matter more than information in latent variable. 2. The connection/interpretation to/of geometric compressibility and how the lossy compressibility plays a role here is not very clear. 3. Is possible to have a concrete example where proposed bounds are shown to be superior to classical bounds based on mutual information?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

not applicable.

Summary

**Post-rebuttal** I thank the authors for their response and have raised my score to an accept. Understanding the generalization properties of modern machine learning algorithms is one of the key challenges in statistical learning theory. Focusing on the classification setup, this paper first studies the generalization error in terms of the compressibility of the predicted labels using tools from rate distortion theory. The proposed bounds feature a new KL divergence term that takes into account the structure and "simplicity" of the algorithm. The authors then extend the framework to give a generalization bound for information bottleneck-type representation learning in terms of the compressibility of the latent variable. The authors also empirically demonstrate a practical benefit of their theoretical results for the information bottleneck.

Strengths

In terms of originality, the generalization guarantee in terms of the compressibility of the representation that takes into account the complexity and structure of the encoder class appear novel. Also unlike existing mutual information-based bounds, the new bounds are not vacuous for deterministic algorithms. The work builds on prior work on PAC-MDL bounds from Blum and Langford (2003). The latter used a bits-back coding argument, while the tools developed in this paper rely on rate distortion theory and block coding arguments building on prior work [SGRS22] and [SZ23]. The authors also demonstrate a practical benefit of their theoretical results by comparing the performance of the vanilla VIB and a `"category-dependent" version, with the latter performing better in terms of generalization error and population risk. Overall, the writing is clear though the presentation can be improved in certain parts.

Weaknesses

The presentation in the main text can be improved in certain parts, e.g., giving more intuition for the proof techniques; see below. Also, in its current form, the paper seems to end abruptly. A separate Conclusions section summarizing the main contributions along with scope for future work / citing open problems can be very useful.

Questions

* What is a high-level intuition for variable-size compressibility for the tail bounds? Likewise, for block covering does compressing multiple instances simultaneously allow for smaller distortion levels? It would definitely help the reader if some of the high-level proof ideas are laid out more clearly either in the main text or supplementary.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

I do not foresee any potential negative societal impact of this work.

Summary

This submission considers supervised learning problems and derives two types of generalization gap bounds. The first type states that generalization gap can be upper bounded by approximately $\sqrt{\mathrm{KL}(P || Q) / n}$, where $P$ is the average (wrt data and training stochasticity) of predicted label distribution on $n$ training examples and $n$ test examples conditioned on the labels of the $2n$ examples, while $Q$ is a fixed **symmetric** prior distribution for predicted labels on $2n$ examples given $2n$ labels. $Q$ is symmetric when $Q(\hat{y_1},\ldots,\hat{y_{2n}} | y_1,\ldots,y_{2n}) = Q(\hat{y_{\pi(1)}},\ldots,\hat{y_{\pi(2n)}} | y_{\pi(1)},\ldots,y_{\pi(2n)})$ for any permutation $\pi$. In a sense, the symmetricity means that $Q$ puts a distribution on $\hat{\mathcal{Y}}^{2n}$ given $y_1,\ldots,y_{2n}$ without "knowing" which $n$ of the $2n$ examples belong to the training set. The second type of generalization bounds are designed for classifiers which have a representation part (an encoder) and a classification head. The main term in these type of bounds is also of the form $\sqrt{\mathrm{KL}(P || Q) / n}$, but where representations play the role of predicted labels. Again there are $n$ training, $n$ testing examples, and a symmetric prior $Q$ but now over representations. Notably, this type of bounds do not depend on the classification head, they only depend on the encoder. For both type of generalization bounds, both in-expectation and in-probability bounds are derived. Furthermore, the authors derive tighter bounds where the mismatch between empirical and population risks is measured by a convex function $h:[0,1] \times [0,1] \rightarrow [0,2]$ instead of the $|x-y|$ function. Finally, to avoid KL divergence becoming infinite when working with continuous labels or continuous representations, the authors derive bounds where one seeks a "quantized" algorithm that has a similar generalization gap but non-infinite KL divergence.

Strengths

**Strength #1: Relevance.** Understanding how some properties of learned representations affect generalization is an important and highly relevant direction for the NeurIPS community. The authors present a good summary of the criticism of information-bottleneck compression ideas in explaining generalization and propose to consider instead the joint compressibility of training and test representations. **Strength #2: Originality.** To my best knowledge, the bounds that depend on representation complexity (results of Section 3) are novel. In some sense, these results suggest a new type of information bottleneck for representation learning. **Weaker strength: Soundness.** The results of this submission are sound as far as I checked. However, in some places I couldn't follow the proofs very carefully as some steps were not detailed enough or the notation was confusing (please see my comments below).

Weaknesses

**Weakness #1: Hard to judge the significance.** While the overall direction of obtaining generalization guarantees via joint compressibility is interesting and this submission makes a good step towards exploring that direction, it is hard to judge how significant are the main results of Section 3. - [Estimation] I don't see a good way of estimating the bound of Theorem 5. In the case of continuous representations, how can one take an infimum over learning algorithms that have similar generalization gap? - [Experiments] The experiments do not help to judge whether the KL term in the bound can be successfully used as a regularization. * [E1] Line 608, "As in the VIB approach, inevitably, we only restrict the optimization to the training set.": This seems to go against the spirit of the approach employed throughout the paper. Namely, as I understand the symmetricity of $\boldsymbol{Q}$ can be interpreted as $\boldsymbol{Q}$ not being aware of the train-test set split. Forcing the encoder to be close to such a distribution $\boldsymbol{Q}$ forces the encoder to behave similarly for seen and unseen examples. From this one gets a generalization guarantee. - [E2] Even if one splits an original training set into a new training set and a ghost set to follow the logic of the main text, the KL regularization term will probably be fitted by "cheating", in the sense that representations of the new training set and ghost set examples will be different from those of unseen examples. The nature of ghost examples (to be not used during the training) makes it impossible to use the KL divergence term of Theorem 5 as a regularization. * [E3] Lines 612-617: $\boldsymbol{Q}$ will not be symmetric in this construction, as it is constructed using statistics of *training data*. Consequently, the KL divergence can become optimistically small (when training examples are "memorized"). - [E4] Lines 667-669, "It can be observed that our scheme outperforms VIB in terms of both generalization error and population risk, for all values of the Lagrange multiplier $\beta$": Since the two objectives are different one should not compare then for a fixed $\beta$. Instead, we can consider tuning $\beta$ for each of them independently to minimize test error. We see that in that case there is no significant difference between the two approaches. Similarly, there is no significant difference if we compare generalization gaps for a fixed level of small test error. - [Applications] The submission would benefit significantly from an analytical evaluation/bounding of the bound of Theorem 5 in some tractable settings (maybe a two-layer neural network on Gaussian data). This would allow to see if the bound is powerful enough to produce meaningful results in some settings. **Weakness #2: Clarity & Presentation.** This submission would benefit from improvements in terms of clarity and presentation. It was hard for me to read and understand the details. There are many small mistakes, the notation is cumbersome, and some parts of derivations are not detailed enough. Please see more detailed comments and suggestions below. **Weakness #3: Missing important related work.** This submission does not cite well the literature of information-theoretic generalization bounds initiated by the works Russo and Zou [1] and Xu and Raginsky [2], and developed further by many others. This line of work is relevant for multiple reasons. 1. The general technique of deriving bounds in this work and in the mentioned line of work is the same (using Donsker-Varadhan inequality and bounding a cumulant-generating function). 2. The general idea of measuring how different a learned function acts on training and testing examples has been used in many follow-ups on [3]. 3. The idea of symmetric priors appears in [4] under the name "almost exchangeable data-dependent priors". In their bound there is also a term of the form $\mathrm{KL}(P||Q)$, where $P$ is the output distribution (over hypotheses) after training on $n$ training examples, while $Q$ is a symmetric prior over hypotheses given $2n$ examples. **Important:** when considering in-expectation bounds, the optimal $Q$ is the average of output distributions over all possible $2^n$ training/test splits. The KL divergence for this optimal prior becomes the conditional mutual information (CMI) $I(W; J | \tilde{Z}_0, \tilde{Z}_1)$, where $J \sim \mathrm{Uniform}\\{0,1\\}^n$ is the training-test split, $W$ is the output of the learning algorithm when splitting with $J$, and $\tilde{Z}_0$ and $\tilde{Z}_1$ are the $n$ training and $n$ testing examples respectively. 4. While in works [2], [3], and [4] the bounds depend on some kind of information captured by the output of a learning algorithm (i.e., the output hypothesis), in [5] there are bounds that measure information in *predicted labels* in the random subsample setting of Steinke and Zakynthinou [3]. Given the important part of the last point, the result of Theorem 1 can be seen as an f-CMI bound written differently. The only minor difference is that in f-CMI bounds there is conditioning on $\tilde{Z}_0$ and $\tilde{Z}_1$ instead of $\tilde{Y}_0$ and $\tilde{Y}_1$ (i.e., only the label parts), but it is easy to see that the proof does not use the input parts ($\tilde{X}_0$ and $\tilde{X}_1$). 5. It is also worth to mention that CMI (or f-CMI or e-CMI) bounds are usually improved by measuring information with individual pairs separately (variations of the technique of Bu et al. [6]), and by conditioning on an individual pair of train/test examples (e.g., [7]). These two techniques together produce f-CMI bounds of form $\sum_{i=1}^n \sqrt{I(\hat{Y_{0,i}}, \hat{Y_{1,i}}; J_i | Y_{0,i}, Y_{1,i})}$ that can be much tighter than the bound of Theorem 1 (the equivalent of the f-CMI bound when measuring information in all predictions at once). **References** [1] D. Russo and J. Zou. How much does your data exploration overfit? controlling bias via information usage. IEEE Transactions on Information Theory, 2019. [2] A. Xu and M. Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. NeurIPS 2017. [3] T. Steinke and L. Zakynthinou. Reasoning About Generalization via Conditional Mutual Information. COLT 2020. [4] P. Grunwald, T. Steinke, L. Zakynthinou. Pac-bayes, Mac-bayes and conditional mutual information: Fast rate bounds that handle general VC classes. COLT 2021. [5] H. Harutyunyan, M. Raginsky, GV. Steeg, and A. Galstyan. Information-theoretic generalization bounds for black-box learning algorithms. NeurIPS 2021. [6] Y. Bu, S. Zou, and V. V. Veeravalli. Tightening mutual information-based bounds on generalization error. IEEE Journal on Selected Areas in Information Theory, 2020. [7] B. Rodríguez-Gálvez, G. Bassi, R. Thobaben, and M. Skoglund. On random subset generalization error bounds and the stochastic gradient Langevin dynamics algorithm. IEEE ITW 2021.

Questions

1. Lines 45-48: What exact result is being referred here? The mutual information between two variables cannot be seen directly as some sort of description length (just consider two variables with 0 mutual information). 2. Section 2.1 Compressibility framework: The concepts introduced in this section are not used in the main text (specifically the block quantities). I recommend moving this subsection to Appendix C. This way it becomes easier to read the main text and the already complicated notation does not become more complicated. The only idea that I recommend keeping from this section is that if one can encode $\hat{Y}^{2n}$ with short code on average, then one should expect good generalization. Similarly, the sentence of lines 209-211 should be moved too. 3. Theorem 1: $\text{gen}(S,W)$ should be defined. 4. Equation (3): The training algorithm can act differently on training examples with different indices. It is not necessarily true that the RHS of eq. (3) can be written as $n$ times the quantity for a single example. 5. Theorem 2: $\mathcal{L}(W)$ and $\hat{\mathcal{L}}(S,W)$ should be defined. 6. Line 295: $\boldsymbol{Q}$ instead of $P$ for the prior. 7. Theorem 4: Why is the idea of a posterior $Q$ introduced here, can't it be just $P_{W|S}$ and $P_{\hat{W}|S}$? As it has been done before this point, the letter $Q$ would be good to use only for priors. In the equation after line 311, it should be $\hat{W} \sim P_{\hat{W}|S}$ instead of $\hat{W} \sim Q$. Also, in eq. (6) it would be good to write it as a difference of two expectations. In that case it will be clear what distributions $W$ and $\hat{W}$ follow. 8. Line 802: It should be $\boldsymbol{Q}$ (the symmetric prior) instead of $Q$ (the posterior over $W$). Similarly in all $\Delta(S,Q)$ terms. Also $\Delta(S,Q)$ should be $\Delta(S,S',Q)$ as it depends on $S'$ too. 9. Line 350: Why is the property of preserving labeling important for the permutation $\pi$ in the definition of a symmetric prior? I recommend adding more details of the transition at line 839 where the symmetricity of $\boldsymbol{Q}$ is used. 10. Line 831-832, "it can be easily verified that this distribution is symmetric": Which distribution is referred here? 11. Line 833: It should be $k : [n] \rightarrow [n]$. 12. Line 834: $||\hat{p_s} - \hat{p_{s'}}||_1 \in [0, 2]$, so $T$ can become negative. It should be divided by 2. 13. Lines 833-834: The construction of the permutation $k$ is not always possible. It might not be possible to map the *first* $T$ samples in $Y$ to another $T$ equal samples in $Y'$. 14. The technique of taking an infimum over close enough "quantized" algorithms is a step that can be applied no matter whether the bound depends on joint compressibility or $\hat{I}(X; U)$ or something else. Therefore, I feel that one should not claim advantage over other results in terms of being applicable also to continuous variables when the advantage is a result of applying this general technique (not due to something inherent to the approach). 15. It would be good to have a small discussion on the fact that representation learning has many goals (e.g., learning representations that are "useful" for many tasks). Learning representations that guarantee small generalization gap for a *given* task is not of high importance in common representation learning settings. I don't think that one would be happy to find out that after a heavy pretraining round with a pretraining task $Y$ the network learned a representation $R=\phi(X)$ such that $R=g(Y)$ for some function $g$.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The limitations are addressed adequately.

Reviewer response

Thank you for the detailed response. The inclusion of a discussion on the overlooked related work and the refinements of presentation improve this submission. I have adjusted the score accordingly.