Class-Wise Generalization Error: An Information-Theoretic Analysis

Thank you for your review.

Q: "...provide bounds for the expected generalization over the whole data distribution, which implicitly assumes that the model generalizes uniformly for all the classes."

We agree with the Reviewer that ‘uniformly’ can be interpreted solely from the context of imbalanced datasets. In the context of our paper, we use ‘uniformly’ more generally to the noticeable heterogeneity of generalization among different classes. To avoid this misunderstanding, we have replaced the word uniformly in the text accordingly.

Q: The number of training samples in Figure 1 is the whole sample number of all the classes. To interpret the reason for the different generalization gaps over classes, we need first to consider whether the classes are balanced.

In Figure 1, we use the total training sample and not the samples corresponding to a specific class. In our framework, we resample the entire training data to obtain i.i.d samples generated from the same data distribution P_Z, so that we can study the average class-wise generalization error. For example, the results shown in Figure 1 are the average over multiple random seeds, i.e., average generalization (For CIFAR10 we have 40 experiments, and for the noisy variant 75) and some of these seeds would be more imbalanced than the others, but the randomness is coming from the data distribution P_Z. As CIFAR10 is a balanced dataset, resample would not create significant imbalanced training data, and the total training sample divided by 10 can be interpreted as the number of samples of each class.

We respectfully disagree with the Reviewer that class imbalance is the main cause of the variations in class-wise generalization errors. Although class imbalance might aggravate this phenomenon, such a variation is a general intrinsic property of neural networks and has been observed in multiple literature even with balanced data. For example, Balestriero et al. (2022); Kirichenko et al. (2023) showed that standard data augmentation and regularization techniques, e.g., weight decay and dropout, surprisingly aggravate this class bias yielding in larger class performance variations. In this work, we are interested in studying this phenomenon theoretically and providing insights about it using information-theoretic tools. This shows that even ‘benign’ techniques such as classic data augmentation and regularization approaches exhibit pronounced class dependencies. In other words, while these methods enhance generalization for specific classes, they can harm the performance of other classes, even under class-balanced cases.

Q: How the label noise is introduced?

In our experiments, we treat label noise as a data augmentation technique, i.e., it is added only to the training data. The main point was to show that a slight modification to the learning technique (data augmentation with label noise in this case) dramatically changes the behavior of class-wise generalization. This further motivates the study of this phenomenon, which is the main contribution of this paper.
Note that even if the noise is added to both train and test, we still observe similar behavior, and similar conclusions can be drawn. We have included these results in the revised manuscript (Figure 6 in Appendix C.6).

Q: Assumption 1, Lemma 1, and Theorem 1 seem problematic.

Note that Assumption 1 states that the learned model W and each sample Z_i have the same joint distribution PW,Z. Here, the index i can be interpreted as the order of samples in the training data, and Assumption 1 is saying that swapping the order of any two i.i.d samples $Z_i$ and $Z_j$ in the training data will not change the joint distribution of P_WZ. This is a standard assumption in the MI setting to derive generalization bounds based on individual samples instead of the whole dataset (similar to the analysis by Bu et al. (2019)). And we have rephrased the interpretation of this assumption in the revision.

Using the symmetry property of the joint distribution (Assumption 1), it follows directly that all the terms in the sum in Definition 1 become identical and we obtain the individual-sample-based expression of the class-wise generalization in Lemma 1.

In Theorem 1, the conditional KL divergence term implicitly depends on the number of samples. Note that W is learned from the entire training data with n samples, and PW,X|y denotes the joint distribution of W and a single training data (X,y). The generalization error is upper bounded by the KL divergence between this PW,X|y and the distribution PW PX|Y=y, which replaces one training data with its independent copy. Thus, this KL divergence decreases with $n$.

Q: The theorems are based on tiny modifications on previous MI and CMI bounds, which are not novel to me.

We respectfully disagree with the reviewer. Please refer to our general answer for full details regarding the novelty and the technical challenges in our analysis.

Thank you for your review. In light of our responses to your individual points, please let us know if you have any remaining concerns.

Q: It seems as if all of your results follow straightforwardly by the same technique as you use for Theorem 7

We respectfully disagree with the reviewer that all bounds can be obtained similarly to Theorem 7. While Theorem 1 and Theorem 7 both share the same core proof and can be ‘combined’ as the reviewer suggested, we don’t see how the same logic can be applied to the supersample (CMI) setting to obtain bounds without the indicator functions. If possible, would the reviewer elaborate more on this? Indeed, the indicators are a fundamental factor in the definition of the class-wise generalization error in the CMI setting and hence it is expected that they appear in the bounds.

Q: One potential way to exploit this would be to actually use the non-y data to construct a more informed prior in PAC-Bayesian bounds similar to Ambroladze et al (2021)

The Reviewer made an interesting possible connection here. Indeed, one potential way to use our analysis is to extend it to PAC-Bayesian bounds and exploit the non-y data in constructing the prior. However, It is beyond the scope of this work, as here we are interested in deriving bounds in the standard MI and CMI settings.

We are appreciative of the Reviewer's valuable suggestion, and we intend to include it in the section on future work, as we believe it is a promising idea worthy of further investigation.

Q: For the sub-task problem, are the target classes assumed to be known?

Yes. As explained at the start of Section 4.2 and in Appendix D.2, it is assumed that the target domain encompasses a specific known subset of the classes encountered during training in the sub-task problem.

Q: Regarding the worst sub-population performance. It is not clear to me why it needs to be rewritten as an information-theoretic bound

The reviewer made a valid point here. The maximum is an upper bound without the need for an information-theoretic bound. In the revised manuscript, we have removed eq 15 and its corresponding discussion. In the original draft, we tried to draw a connection between our results and the empirical findings of Yang et al. (2023), where it is shown empirically that model performance is highly correlated with the worst sub-population performance. As our bound in eq 15 depends on the population achieving the worst KL divergence, it can be loosely interpreted as a justification for the empirical results of Yang et al. (2023).

Q: After 1, you require $n_y < n$. $n_y \leq n$ Isn’t fine?

The case $n_y = n$ refers to the case of having only a single label in the data. In this case, the problem is ill-posed as there is only one class in the data (no learning needed for a classification problem).

Q: The term “overfit” seems a bit unclear (compare “benign overfitting”, e.g.).

In the paper, overfitting refers to the classic phenomenon of having a large generalization error, i.e., a large gap between the empirical risk and the population (test) risk. To avoid this confusion, we now make this clear in the updated manuscript.

Q: I believe the term “supersample” usually refers to the entire set of samples.

A supersample refers to a single pair $(\rmZ_i^+,\rmZ_i^-)$. The total dataset $\rmZ_{[2n]}$ is composed of n supersamples, i.e., a total of 2n samples.


We hope that our response has addressed all your questions and concerns. We kindly ask you to let us know if you have any remaining concerns, and - if we have answered your questions - to reevaluate your score.

I would like to thank the authors for their response, and I have briefly read all the reviews and the corresponding responses. I will read the revised parts more carefully before I discuss with AC and other reviewers (I expected to engage more with the authors before the Review-Author discussion period ends, but there are too many rebuttals submitted today). Below are some quick notes on your response:

Regarding the comparison of the individual bound and Corollary 1, I have no additional questions and would like to share a quick thought that might be helpful: recall the original $I(W;Z)$-based MI bound, by taking a closer look at $W=\mathcal{A}(R,S)$ where random variable $R$ is the randomness of the algorithm $\mathcal{A}$, one can prove a $\mathbb{E}\_{R}\sqrt{I^R(W;Z)}$ based bound, which in some cases can improve over the individual bound, similar bound can be found in [1, Theorem 2], called the "$R$-conditioned bound". In your paper, you focus on $Z=(X,Y)$ in $I(W;Z)$ and obtain the "$Y$-conditioned bound", serving as a "symmetric" technique to the $R$-conditioned bound (to be clear, unlike $R$ that is independent of $Z$, here $Y$ also dependents on $W$ ). This technique of tightening, extracting an implicit random variable from an explicit one in $I(W;S)$, can be traced back to [2], where the authors extract the subset index from $S$. Additionally, it would be more interesting if your class-CMI bound (i.e. Theorem 2) can also recover the individual CMI bound but there might be some difficulties arising from Definition 2. This could be a potential theoretical motivation in addition to your practical motivation drawn from the deep learning scenario.

Regarding Definition 2 and the comparision between Theorem 4 and Theorem 3, the revisions look better to me now and I will check more details in Appendix.

[1] Hellström et al. "New family of generalization bounds using samplewise evaluated CMI." NeurIPS 2022.

[2] Negrea et al. "Information-Theoretic Generalization Bounds for SGLD via Data-Dependent Estimates." NeurIPS 2019.

Thank you for your review. In light of our responses to your individual points, please let us know if you have any remaining concerns.

Q: The empirical results are confined to ResNet50 on CIFAR10/Noisy CIFAR10

In addition to CIFAR10//Noisy CIFAR10 results, we also provide results using the MNIST dataset in Appendix C.5, which is a similar digits classification task as SVHN.

Q: Regarding Definition 2: Is the class-generalization error defined in Eq.(6) equivalent to Eq.(3) in Definition 1?

The Reviewer makes a valid point here. Indeed if we use $n^y_{Z_{[2n]}} /2$, taking expectation over Y for eq. (6) does not correspond to the classic standard generalization. To overcome this issue and keep corollary 2, we change the normalization factor $n^y_{Z_{[2n]}} /2$ in Definition 2 to $n P(\rmY=y)$. This does not change the proof of all the results and the standard generalization bound in corollary 2 can be obtained by taking an expectation over $P_Y$ directly.
In practice, $P(\rmY=y)$ is unknown and needs to be estimated using the available data. One simple way is to use $\frac{n^y_{Z_{[2n]}}}{2n}$ as an estimate $P(\rmY=y)$, which is exactly what we did in the empirical sections. Note that our empirical results validate the proposed bounds for this estimated class-wise generalization error. We have made the corresponding changes in the revised manuscript.

Q: Regarding Corollary 1: Is it tighter than the classic individual mutual information bound in Bu et al. (2019)

Good point. Indeed, as shown by the Reviewer, the bound in Corollary 1 is tighter than the bound in Bu et al. (2019) if the sub-Gaussian constant does not depend on Y. This presents an additional contribution to the class-wise analysis. We sincerely appreciate the Reviewer for pointing this out, and we have integrated this result into the revised manuscript. Please refer to Appendix C.4 for further details.

Q: Could you elaborate more on the last sentence in Remark 2..

In the case that one sample in the pair (Z− i , Z+ i ) is from class y and the other is from a different class, the term in the bound is non-zero and the information from both samples of the pair contributes to the bound (I_Z[2n] (fW(X± i ); U_i)). From this perspective, samples from other classes ($\neq y$) can still affect these bounds, leading to loose bounds on the generalization error of class y. In the revised manuscript, we now rewrite Remark 2 to clarify this point.

Q: Therefore, it is uncertain whether Theorem 4 is tighter than Theorem 3 or not

Thanks for pointing out this point. Theorem 4 is tighter even if the bound in Theorem 3 has these indicator. This can be seeing that for a fixed Z_[2n], we have 4 different possible cases for each term in the sum of the bounds. We have included the detailed proof as part of Theorem 4 proof in the revised paper, please refer.

Q: Hrayr Harutyunyan, et al. (2020) is a related work that could be included

Thanks for the reference. We have added it to the related work section of the revised paper.

Q: In Theorem 8, you may use the widely accepted term for the loss pair-based CMI, namely evaluated CMI or e-CMI

Thanks for pointing this out. We have changed the notation in the updated manuscript and adopted the same notation for the loss-based bound as in Steinke & Zakynthinou (2020).

Q: Minor comments

Thanks for pointing this out. We have now fixed the mentioned points in the revised manuscript.


We hope that our response has addressed all your questions and concerns. We kindly ask you to let us know if you have any remaining concerns, and - if we have answered your questions - to reevaluate your score.

Thank you for your comments. We provide answers to individual points below.

Q: No new technical challenge

We respectfully disagree. Please refer to the comprehensive general response provided earlier. As pioneers in the exploration of this tool in such problems, we believe that novelty should not be evaluated solely based on technical complexity.

Q: Previous works on computing class-wise generalization bounds

To the best of our knowledge, this is the first work that considers such a setting, i.e., data with particular attributes, which is one of the key novelties of this paper. Moreover, as shown in Section 4, the developed bounds are useful beyond the scope of class-wise generalization and can provide insights in various contexts.

Q: Information-theoretic bounds are natural for this setting as they depend on both the algorithm and the data

Thanks for pointing this out. It is possible that stability or hypothesis-based bounds can be adopted to capture class-wise generalization. However, we choose information-theoretic bounds because those bounds do not fully characterize all the aspects that could influence the generalization error of a supervised learning problem. For example, VC dimension-based bounds depend only on the hypothesis class, and algorithmic stability-based bounds only exploit the properties of the learning algorithm. As a consequence, both methods fail to capture the fact that generalization error depends strongly on the interplay between the hypothesis class, learning algorithm, and the underlying data-generating distribution. In contrast, information-theoretic bounds incorporate all of these factors into consideration, providing a more holistic understanding of the class-wise generalization problem.

We are happy to answer any further questions you may have. If we have addressed your concerns, we kindly ask that you may consider raising your score.