Paper: Slicing Mutual Information Generalization Bounds for Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The target learning scenarios seems too restrictive for modern deep-learning models. The subspace mechanism severely restricts the available number of parameters in the neural network, making them hardly applicable in modern learning tasks. The Lipschitz condition is also hardly satisfied in conventional network architectures. As seen in Figure 7, the training accuracy cannot reach 80% on CIFAR10, even for a binary classification task.

Author response:
1. It seems like there was some unfortunate confusion regarding Figure 7: the $y$-axis reports the error as measured by the binary cross-entropy loss on the test set. The corresponding training accuracy can reach 100\% after 100 epochs (for $\lambda=0$). Note that the size of the training dataset is $2000$ to make generalization harder. We make this clearer in the revision. The remaining concerns raised are fortunately not an obstacle. First, the random compression scheme we study has found tangible success in prior recent work: we refer the reviewer to our general answer for detailed explanations. Second, regarding *"the Lipschitz condition is hardly satisfied in conventional network architectures"*: feedforward NNs based on common choices of activations (e.g., ReLU) can easily be made Lipschitz with respect to the weights by simply constraining the maximum norm of the data input and weights (Bethune et al., 2023). We also emphasize that the Lipschitz condition is a common strategy to control generalization error bounds that contain a distortion term (Sefidgaran et al., 2022; Neyshabur et al., 2017: see our response to Reviewer qRo1).

Reviewer comment (Theory):
Although the dimensionality is reduced in the presented bounds, this improvement hardly solves the tractability problem. The key mutual information I(W’;S) is still computationally intractable for greater d values, e.g. d > 20. For empirical analysis, the authors adopt a neural network-based mutual information estimator (MINE), whose accuracy lacks theoretical guarantees and is thus questionable for high-dimensional variables. The viability of applying MINE for theoretical analysis is doubtful.

Author response:
2. *"The viability of applying MINE for theoretical analysis is doubtful"*: On the contrary, theoretical guarantees on the convergence rate of MINE are well studied and strong, see e.g. Theorem 2 of Goldfeld et al. (2022) which studies MINE in the context of random slicing and shows a clear error bound decaying as the dimension $k$ of the slices (our $d$) decrease. We will add a remark in the revision pointing out this fact. Also, see Remark 2 in Goldfeld et al. (2022) for further discussion of relaxation of some of the assumptions of MINE, which enables optimizing over a larger class of distribution in lower-dimensional spaces. While  the results in Goldfeld et al. (2022) are stated for 3 layer MINE networks for simplicity, the extension to deep networks is given in followup works if the reviewer is interested. Note the regime discussed by McAllester & Stratos (2020) (see the discussion in our Introduction) is not relevant for us since it applies to the case of large MI -- in that setting, the generalization error bound would be very large and vacuous anyway. We added these elements in our paper to clarify why and when our approach is beneficial for estimating MI.

Reviewer comment (Novelty):
The derivation of Theorem 3.1 and 3.2 seems trivial for me: given a fixed Theta, the trainable set of parameters then becomes W’ instead of W. Then Theorem 3.1 and 3.2 directly follows by Xu and Bu’s results by replacing I(W;S) with I(W’;S), or replacing I(W;Z_i) with I(W’;Z_i). The current contribution of the paper may not be sufficient to be published in ICLR.

Author response:
3. While our paper has an important theoretical component, the main motivation and message of our paper are practical (cf. our general response to all reviewers). Given the growing interest for NNs trained on random subspaces (Li et al., 2018; Lotfi et al., 2022), our work seeks to derive insights and build theoretical foundations on their generalization ability. We do not need our proof techniques to be difficult or of independent interest to achieve this goal. In fact, we would argue there is strength in simplicity. That said, we emphasize that our theoretical results in Section 3.2, which result from applying our generic bounds in specific settings, are nontrivial. Even for Gaussian mean estimation and linear regression, the derivation of the generalization error and the bound on its cumulant-generating function is not straightforward, e.g. the loss follows a more complex distribution (generalizated chi-square) than in the original space (Sections A.3 and A.4).

Reviewer comment (Theory):
Most of the bounds derived by the author can be understood as a bias-variance tradeoff, with some terms that decrease with $n$ (the *variance*), and some others that do not decrease with $n$ (the *bias*). Therefore, it is a bit upsetting to not see a clear discussion of this. For example, in the rate-distortions bounds (theorems 3.3 and 3.4) the left term does not depend on $n$. What are the practical implications of this? Maybe it is worth adding a discussion on this topic, and to underline the terms of your bounds with `\underbrace`.

Author response:
**Bias-variance tradeoff:** This is an insightful connection to make, indeed the bias-variance tradeoff is spiritually very similar to the rate-distortion tradeoff. That said, the concept of rate-distortion is the more applicable concept here since we are using the framework of compression and information content. As noted, the first term will be the distortion (c.f. bias) which will decrease as $d$ increases, but will not go away with $n$. The second term is the rate (c.f. variance) which will decay with both $n$ and $d$. We will add in a more extended discussion of this tradeoff and how it guides the choice of $d$; we initially had to omit it for space reasons.

Reviewer comment (Experiments):
Theoretical results show a dependency of the bound on the number of levels in quantization. However, in none of the experiments is the number of levels of quantization kept constant in the experiment of Fig 3. It would be interesting to measure the generation bounds when the number of levels vary, compared to the (empirical) generalization error.

Author response:
**Quantization:** As suggested by the reviewer, we added an additional experiment to illustrate how the number of quantization levels influences the empirical generalization error and our bound. Our results and interpretation are provided in Section C, and confirm our theoretical analysis.

Reviewer comment (Theory):
The motivation behind introducing $\Theta$ is unclear. It *seems* that sampling an element from the Grassmannian $\mathcal{M}$ would be enough. It seems that it is not crucial that $w=\Theta w'$ for the theoretical results to hold. What *seems* to matter is that $w$ is effectively low-rank, i.e., lives on a low-dimensional manifold $\mathcal{M}$ which is sampled independently from the data, such that $\mathbb{P}(\mathcal{M}|Z)=\mathbb{P}(\mathcal{M})$ so that expectation can be taken over $\mathbb{P}(\mathcal{M})$ in the bound. The fact that $\Theta$ is orthogonal does not seem to be necessary as long as $w$ is sampled from $\mathcal{M}$. Is there a motivation behind using orthogonal matrices $\Theta$? Getting rid of it would allow removing the `scipy.linalg.orth` in the experiments.

Author response:
**Motivation behind the orthogonal projector:** The reviewer's intuition makes sense: our theoretical results could be extended to more general low-rank representations of $W$. We actually stated Lemma A.1 so that it is generic enough (i.e., not restricted to $w = \Theta w'$) to facilitate such extensions. 
That being said, we believe it is important to consider orthonormal projectors $\Theta$ for several reasons. First, prior work showed that this is an interesting setting in deep learning, e.g., to efficiently compress NNs while maintaining a reasonable performance, or to tighten PAC-Bayesian generalization bounds: see paragraph "Sliced neural networks" in Section 1 or our paper. Then, assuming $\Theta^\top \Theta = I_d$ simplifies some derivations when computing the generalization error and its bound, e.g., for Gaussian mean estimation (Section A.3). Last but not least, orthonormality is advantageous for computational purposes, as it prevents poor conditioning when training on $\mathrm{W}_{\Theta, d}$ (Lotfi et al., 2022).

Reviewer comment (Experiments):
The experimental setup of Section 4.2 does not match the theorems. The theorems assume that the train set $S_n$ is sampled from $\mu$. However, as written in the appendix, classes are hidden during training, leading to Out Of Distribution (OOD) test sets compared to the train sets. This is not the Empirical Risk Minimization (ERM) paradigm anymore, but rather evaluation under distribution shift. Therefore, the generalization error reported cannot be compliant with the bound. For example, the marginal probabilities of the $y=1$ class fall from $\frac{1}{4}$ to $\frac{1}{5}$ for MNIST, and from $\frac{1}{3}$ to $\frac{1}{4}$ for CIFAR-10. The relevance of this setting is questionable, and sticking with the ERM paradigm would be better.

Author response:
**Experimental setup of sec 4.2:** This is a fair point: we added the experiment suggested by the reviewer (see Section C). Instead of hiding a class in the training dataset to make generalization harder, we reduced the size of the dataset. Our plot further illustrates our rate-distortion bounds and is line with our intuition that approximately-compressible NNs (i.e., with smaller distortion term) generalize better.

Reviewer comment (Writing):
Typo: In page 1 the loss is defined as $\ell:Z\times W\rightarrow\mathbb{R}$, but everywhere else it is used as $\ell(w,z_i)$.

Author response:
**Typo:** The typo is now fixed in our revision.

Reviewer comment (Theory):
Regarding bounded binary cross entropy, you write (p9): *The role of $\epsilon>0$ is to make $\ell$ bounded thus meet the conditions of Theorems 3.3 and 3.4.* However, if no constraint is put on $f(w,X)$ then $f(w,X)+\epsilon$ can take any value in $(-\infty,+\infty)$. Unless $f$ is lower bounded. In the appendix (p21) $w_i$, $b_i$ and $x$ are bounded, which seems to bound $f$, but the bound on $f$ is not given and the link with boundedness of $f$ is not explicitly drawn. Can you clarify?

Author response:
**Bounded binary cross entropy:** We indeed forgot to explicitly mention that $f(w,X)$ needs to be lower-bounded for $\ell$ (the binary cross-entropy) to be bounded -- and as correctly pointed out by the reviewer, this is the case in our experiments (Sections 4.2 and B.3). The hyperparameter $\varepsilon$ can thus be used as a way to adjust the bound on $\ell$ hence the generalization bound, if desired. We thank the reviewer for raising this point and we clarified it in our revised paper: see Sections 4.2 and B.3.

Reviewer comment (Theory):
Regarding unbounded MI, you write: *for instance if $W ′$ is a deterministic function of $S_n$ given $\Theta$ then $I_{\Theta}(W ′; S_n) = +\infty$.* Since Mutual Information is bounded by the entropy of the variables involved, $I_{\Theta}(W ′; S_n)\leq H(S_n)$, do you mean that $H(S_n)$ is unbounded?

Author response:
**Unbounded MI:** The inequality $I^\Theta(W';S_n) \leq H(S_n)$ does not hold in our general setting, where $S_n$ is a continuous random variable. By eq.(3), we can only say that $I^\Theta(W';S_n)=h^\Theta(S_n)-h^\Theta(S_n|W')$, where $h^\Theta(S_n)\triangleq - \int p(x|\Theta'=\Theta) \log p(x|\Theta'=\Theta) \mathrm{d}x$ and $h^\Theta(S_n|W') \triangleq - \int p(s_n|w', \Theta'=\Theta) \log p(s_n,w'|\Theta'=\Theta) \mathrm{d}s_n \mathrm{d}w'$. Since $p(s_n|w', \Theta'=\Theta)$ denotes a probability density function, it can take values greater than 1, thus $h^\Theta(S_n | W')$ can be negative. Additionally, $I(W';S_n)=+\infty$ if $W'=g(X)$ with $g$ a smooth, non-constant and deterministic function. We clarified this point in Section 3.1 of our revised paper.

Reviewer comment (Theory):
In the appendix, can you clarify how you went from step (39) to step (40) on p15? Similarly, how did you go from steps (42) to (43)?

Author response:
**Appendix:** We added more explanations on *"step (39) to step (40)"* and *"step (42) to step (43)"*: see Section A.3.

Reviewer comment (Writing):
Intrinsic Dimension: Throughout your paper, there is a recurring mention of the intrinsic dimension. Could you clarify what you mean by "intrinsic dimension"?

Author response:
2- **Intrinsic Dimension.** We refer by "intrinsic dimension" the notion described in Section 2 of (Li et al., 2018), since our setting is largely inspired by this work. We added this clarification in the revised manuscript to resolve any ambiguity. Briefly, the "intrinsic dimension" is the slice dimension $d$ below which the test performance of the sliced neural network begins to degrade significantly.

Reviewer comment (Theory):
Comparison between your work and "functional" or "evaluated" MI bounds: In the existing literature, there are generalization bounds that do not suffer from the curse of dimensionality. It would be valuable to understand how your proposed bounds in this paper compare to the results presented in the following papers:  
-- Harutyunyan, Hrayr, et al. "Information-theoretic generalization bounds for black-box learning algorithms." Advances in Neural Information Processing Systems 34 (2021): 24670-24682.  
-- Haghifam, Mahdi, et al. "Understanding generalization via leave-one-out conditional mutual information." 2022 IEEE International Symposium on Information Theory (ISIT). IEEE, 2022.

Author response:
3- **Comparison between your work and "functional" or "evaluated" MI bounds.** We thank the reviewer for sharing these two references: a discussion on how their strategies compare to our work is provided in Section 2, paragraph "Conditional MI generalization bounds". We will add the second reference (Hafigham et al., 2022) to our manuscript (the first one, Harutyunyan et al., 2021, was already cited). We will be happy to complement this discussion to address any questions the reviewer may still have.

Reviewer comment (Novelty):
Different Compression Methods: Your paper discusses modifying learning algorithms to improve information-theoretic generalization bounds. This approach reminds me of adding Gaussian noise to the output of a learning algorithm and trading off information with distortion, which has been explored in other works, such as:  
-- Neyshabur, Behnam, Srinadh Bhojanapalli, and Nathan Srebro. "A pac-bayesian approach to spectrally-normalized margin bounds for neural networks." arXiv preprint arXiv:1707.09564 (2017).  
-- Dziugaite, Gintare Karolina, and Daniel M. Roy. "Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data." arXiv preprint arXiv:1703.11008 (2017).  
Could you outline the main differences between your approach and the approach based on adding Gaussian noise?

Author response:
4- **Different Compression Methods.** The two references provided by the reviewer use an approach that is in essence different from ours: generalization error bounds for NNs are derived by *adding a random perturbation to the weights* and using *PAC-Bayes theory*, whereas we *project the weights onto lower-dimensional random subspaces* and use *information theory*. We view these two compression approaches to be orthogonal yet complementary, in the sense that the two compression methods address two very different intuitions about the structure of neural network weights (low rank structure vs. smoothness), but could easily be combined in practice (this would be an interesting direction for future work). One similarity with our work is that neural networks are *Lipschitz w.r.t the weights* to *bound a distortion term*: Lemma 2 in Neyshabur et al. (2017) applies to NNs with ReLU activations and weights with bounded norms; we implemented the same NNs to illustrate our rate-distortion bounds, since these are Lipschitz w.r.t the weights, as we show in our Section B.3. We will add this discussion in our paper.

Reviewer comment (Theory):
PostHoc explanation of low dimensional compressibility: This paper is motivated by the observation that many neural networks are highly compressible. However, it seems that the method based on random projection might not capture low-dimensional compressibility adequately since the subspace might be data-dependent. Additionally, one limitation of your paper is that it's not clear how to choose the dimension for projection (d). Could you provide insights or potential solutions to address these concerns?

Author response:
5- **PostHoc explanation of low dimensional compressibility.** The use of random projections is a simple approach to (a) check in practice if neural networks are highly compressible (this is often the case as revealed by prior work, e.g., Li et al., 2018), (b) understand the impact of the subspace dimension on the generalization capacity (this is our contribution). In fact, learning $\Theta$ in a data-dependent manner would entirely defeat the purpose of our approach, as these information-theoretic generalization bounds require the accounting for any parameters that are data-dependent. The empirical success of random $\Theta$ was what inspired us to write this paper, as the fact we can ignore a random $\Theta$ in the mutual information computation is a massive savings.

That said, we admit that one limitation of this method is the need to find the intrinsic dimension: the common method to do so is by trial-and-error (Li et al., 2018; Lotfi et al., 2022). Our generalization bounds improve on this situation, as our rate-distortion bounds in particular provide a generalization bound that can be optimized in terms of $d$, trading off the cost of larger $d$ and the benefit of lower distortion. To speed up the search of the intrinsic dimension, one can run the training for different $d$ in parallel, or use search algorithms with favorable time complexity.

Reviewer comment (Novelty):
Connections to Pruning Methods: Can you explain how your findings are connected to pruning techniques used in neural networks?

Author response:
6- **Connections to Pruning Methods.** This is an interesting question. First, we note that training on $W_{\Theta, d}$ is a more effective and flexible method for model compression than pruning, since $\Theta$ can be dense (see Lotfi et al., 2022). Then, while our theorems hold for any distribution $P_\Theta$ on the Stiefel manifold, we believe our proof techniques could be adapted to study more general low-rank representations of the parameter space (as implied by our Lemma A.1 and pointed out by reviewer KfBv). In that sense, our findings can be extended to gain more insights on the generalization ability of neural networks compressed via *random pruning*. However, pruning methods that remove weights in a *data-dependent manner* (such as magnitude-based pruning) do not fit our framework, since only data-independent dimension reduction approaches simplify the mutual information bounds.

Reviewer comment (Theory):
Overall, the bounds also become increasingly loose as the dimension increases.

Author response:
The reviewer is correct in saying that our bounds "become increasingly [large] as the dimension increases", but this is precisely an important take-home message of our work: **reducing the dimension of the training space (through random projections) tightens information-theoretic bounds on the generalization error**. This feature is then leveraged to produce tighter generalization bounds for *approximately-compressible NNs*. In that sense, we disagree with *"this limits the value and applicability of the bounds"*. In fact, it provides a theoretical confirmation of the practitioner-driven observed empirical success of random projection, and is in that sense certainly relevant practically. Indeed, our rate distortion bounds motivated a further modification of the random projection training procedure, allowing for increased flexibility. While a full exploration of the empirical properties of that algorithm is beyond the scope of this first paper, the initial results we show are encouraging. Finally, generalization bounds themselves do have practical utility, in that they can be used to guide difficult architecture choices and in the future inform "scaling laws" (e.g. in our case, the tradeoff between random projection and distortion). Scaling laws generally are often used practically, e.g. the Chinchilla scaling laws which enabled much of the recent growth in transformer architectures for LLMs.  

We are not sure how to correctly understand what "converse bounds" means. If the reviewer refers to lower bounds on the generalization error, we unfortunately do not have an immediate answer on how to address this -- in general, obtaining lower bounds on generalization error is an extremely nontrivial area of learning theory, and matching upper and lower bounds are exceedingly rare in the literature.

Reviewer comment (Theory):
In the footnote 3, the authors mentioned that their bound can be shown theoretically to be tighter for the two settings (mean estimation and linear regression) using data processing inequality. Can this detailed discussion be included and made clearer in the main text? Also, can it be shown theoretically that the bound is tighter in a more general setting (e.g., with neural networks)?

Author response:
In Gaussian mean estimation, $W' = \Theta W$ where $W$ is the solution of the original problem (the weights are optimized on $\mathbb{R}^D$). By the data processing inequality, $I(W';Z_i) \leq I(W;Z_i)$. Therefore, we obtain tighter bounds as compared to strategies which ignore the existence of an intrinsic dimension. The same reasoning applies to linear regression, since $W' = \phi(\Theta,X)W$. In general settings such as neural networks, the question remains open: the relation between $W'$ and $W$ is unknown, thus we cannot conclude using data processing inequality. We added this discussion in Section 3.2 to clarify footnote 3.

Reviewer comment (Experiments):
Is it possible to conduct empirical evaluation of the linear regression generalization bound?

Author response:
The empirical evaluation of the linear regression generalization bound requires finding the infimum in eq. (69). The solution is not available in closed form and can be approximated in practice using numerical solvers.

Reviewer comment (Experiments):
Is it possible to empirically evaluate how varying the number of projection matrices and number of samples of $(W’,Z_i)$ affects the bound?

Author response:
The generalization error and our bounds decrease with increasing number of training samples: see Figures 1 and 2 for illustration. Increasing the number of samples $(W', Z_i)$ will improve the accuracy of the estimates of $I^\Theta(W';Z_i)$ given by MINE. Finally, increasing the number of random projections for a fixed $d$ has little practical impact on our bounds. This is consistent with prior work (Li et al, 2018), which showed that the test performance does not vary much accross multiple projection matrices for a fixed $d$, while the choice of $d$ has a greater impact on the quality of solutions. We will add empirical results illustrating our conclusions.

Reviewer comment (Experiments):
For Section 4.2, is it possible to vary the generalization error by changing the number of samples (similar to Figure 2), and evaluate the rate-distortion bound accordingly?

Author response:
We provided more insights on the impact of $n$ on the rate-distortion bound: see paragraph "Bias-variance tradeoff" in our answer to Reviewer KfBv. We will include the experiment suggested by the reviewer.

Reviewer comment (Writing):
The writing and details provided in the paper could be improved to make the paper more readable and useful. For example, the paper contains a multitude of bounds but does not provide a unifying perspective. It would be good to have more discussion on the bounds and when to apply each one.

Author response:
The key message of our paper is the observation that using random compression tightens information-theoretic generalization bounds for neural networks. As suggested by the reviewer, we added a paragraph in Section 3.1 explaining when to apply each of our bounds. We hope this clarifies our contributions and will be happy to take into account further feedback on that point.

Reviewer comment (Theory):
Overall, the bounds become increasingly loose as the dimension increases, which limits the value and applicability of the bounds in general. Is it possible to tighten the bounds and say something about converse bounds?

Author response:
The reviewer is correct in saying that our bounds "become increasingly [large] as the dimension increases", but this is precisely an important take-home message of our work: **reducing the dimension of the training space (through random projections) tightens information-theoretic bounds on the generalization error**. This feature is then leveraged to produce tighter generalization bounds for *approximately-compressible NNs*. In that sense, we disagree with *"this limits the value and applicability of the bounds"*. In fact, it provides a theoretical confirmation of the practitioner-driven observed empirical success of random projection, and is in that sense certainly relevant practically. Indeed, our rate distortion bounds motivated a further modification of the random projection training procedure, allowing for increased flexibility. While a full exploration of the empirical properties of that algorithm is beyond the scope of this first paper, the initial results we show are encouraging. Finally, generalization bounds themselves do have practical utility, in that they can be used to guide difficult architecture choices and in the future inform "scaling laws" (e.g. in our case, the tradeoff between random projection and distortion). Scaling laws generally are often used practically, e.g. the Chinchilla scaling laws which enabled much of the recent growth in transformer architectures for LLMs.  

We are not sure how to correctly understand what "converse bounds" means. If the reviewer refers to lower bounds on the generalization error, we unfortunately do not have an immediate answer on how to address this -- in general, obtaining lower bounds on generalization error is an extremely nontrivial area of learning theory, and matching upper and lower bounds are exceedingly rare in the literature.