Slicing Mutual Information Generalization Bounds for Neural Networks

We thank the reviewer for their comments. We hope our answers below address their concerns, and would be happy to discuss if they have further questions.

Answers to questions:

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).
    
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.
    
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).

Indeed, the theoretical bound grows with $d$, but recall this is in no way a problem, since the $d$ here is the small dimension of the slice, not the ambient dimension. Furthermore, observe that $d^{3/2}$ is only slightly worse than linear scaling. We reiterate that MINE is a well-established, empirically highly successful estimator even in very high dimensions; and should not be a cause for concern in this relatively low-dimensional estimation regime. Indeed our empirical results confirm it can be used successfully here - not one of our generalization bounds computed in this way is overly optimistic.

"Restricting the hypothesis space of course leads to better generalization." Agreed, this sums up the field of Learning Theory - yet the field is considered one of the most difficult in theoretical computer science. The challenge, unfortunately, is always going to be identifying a restricted hypothesis space that is useful in practice. Random slicing, we argue, is such a mechanism - in fact one highly amenable to the setting of practically computable information-theoretic generalization bounds. As such, we strongly believe that it will be of interest to the community.

We thank the reviewer for their feedback: we are glad they found our work "original" and "useful", our experiments "convincing", and our paper "clear overall".

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.

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.

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).

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.

Typo: The typo is now fixed in our revision.

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.

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.

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

We thank the reviewer for their positive evaluation of our paper. We appreciate they find our idea "very interesting". Please find our replies below.

Answers to questions:

1- Motivation. We provided a clearer picture of the motivations and key contributions of our work: please see our general answer to all reviewers. We hope our additional explanations addresses the reviewer's question, and we are happy to further clarify any aspect during the discussion period.

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.

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.

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.

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.

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.

We thank the reviewer for finding our contributions "novel" and "interesting", and we appreciate they find the topic "important". We address their questions below. We also refer the reviewer to our general answer for a better picture of the main story and impact of our work.

Answers to questions

1. 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.
    
2. 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.
    
3. 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.
    
4. We evaluated our generalization bounds for projected weights (Figure 5) as well as projected and quantized weights (Figure 3) for MNIST classification with NNs. We also ran an additional experiment in Section C. Interestingly, we see that for appropriate choices of $L$, the bounds for projected and quantized weights are tighter than when only projecting.
    
5. 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.
    
6. 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.
    
7. 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.

Dear Reviewer X27N: Thank you again for your time and your feedback, which have helped us improve our paper. We hope our responses have adequately addressed your questions, and that you will reconsider your score in light of our discussion. As the discussion period ends soon, we would greatly appreciate it if you could let us know any remaining questions or concerns you may have: we are more than willing to provide additional explanations and experiments to clarify these.