Paper: Class-Wise Generalization Error: An Information-Theoretic Analysis (ICLR-2024)

Reviewer comment (Theory):
"...provide bounds for the expected generalization over the whole data distribution, which implicitly assumes that the model generalizes uniformly for all the classes." I am afraid that I cannot agree with this statement. Considering the expected generalization error over the whole data distribution means measuring the model **on average** over all the classes, not **uniformly**, especially when the classes are imbalanced.

Author response:
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.

Reviewer comment (Evaluation):
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 the website of the CIFAR 10 dataset: "...The training batches contain the remaining images in random order, but some training batches **may contain more images from one class than another**. Between them, the training batches contain exactly 5000 images from each class." So, could the authors also plot the class-wise gap w.r.t the sample numbers for each class?

Author response:
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.

Reviewer comment (Experiments):
Introducing label noise (5%) for each class will cause an imbalanced sample number for each class, which means the class-wise sample numbers in the right plot must be different (in terms of the given label not true label). That's why we observe a different trend. How the label noise is introduced? Have you introduced the noise on both train data and test data? If only on the train data, there is a distribution shift. If on both, it's the above-mentioned class imbalance.

Author response:
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).

Reviewer comment (Novelty):
The theorems are based on tiny modifications on previous MI and CMI bounds, which are not novel to me.

Author response:
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.

Reviewer comment (Theory):
It seems as if all of the results follow straightforwardly by the same technique as used for Theorem 7. Considering a distribution \(P_Z\) on a space \(\mathcal Z = \mathcal T \times \mathcal X\), where \(\mathcal T\) is the label space or a sensitive attribute, standard techniques for information-theoretic bounds with the joint distribution \(P_{W\vert Z}P_{Z\vert T=t}\) could be applied to avoid the machinery with indicator functions used for the CMI derivations.

Author response:
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.

Reviewer comment (Theory):
One potential way to exploit the above observation would be to actually use the non-\(t\) data to construct a more informed prior. For example, using the data-split method in PAC-Bayesian bounds, as in “Tighter PAC-Bayes bounds” by Ambroladze et al or “On The role of Data in PAC-Bayes Bounds” by Dziugaite et al, could be combined with the CMI approach (or standard MI bounds) so that the prior can depend on data from \(P_{Z\vert T\neq t}\).

Author response:
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.

Reviewer comment (Experiments):
For the sub-task problem, it is unclear whether the target classes are assumed to be known.

Author response:
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.

Reviewer comment (Theory):
Regarding the worst sub-population performance: since \(\text{gen} = E_T \text{gen}_T \leq \max_t  \text{gen}_t\), it is not clear why it needs to be re-written as an information-theoretic bound to establish a statement similar to (15). Furthermore, it is questionable whether this really motivates that the two are closely correlated; while the maximum is clearly an upper bound, it is not evident why this implies correlation.

Author response:
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).

Reviewer comment (Theory):
After 1, the requirement \(n_y < n\) seems unnecessary; isn’t \(n_y \leq n\) sufficient?

Author response:
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).

Reviewer comment (Theory):
The term “overfit” seems a bit unclear (compare “benign overfitting”), and it is uncertain whether it necessarily implies a high population loss.

Author response:
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.

Reviewer comment (Writing):
The term “supersample” is believed to usually refer to the entire set of \(2n\) samples, but it seems it is used here to refer only to a single pair of samples. This usage should be clarified.

Author response:
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.

Reviewer comment (Experiments):
One main limitation is that the empirical results are confined to ResNet50 on CIFAR10/Noisy CIFAR10. Adding more experimental settings (e.g., SVHN) would further show the effeteness of the proposed bounds.

Author response:
>**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.

Reviewer comment (Theory):
Regarding Definition 2: Is the class-generalization error defined in Eq.(6) equivalent to Eq.(3) in Definition 1? If so, could you rigorously demonstrate their equivalence? To my understanding, it is natural to see that Eq.(3) is the same as the following expression:
$$
\mathbb{E}\_{\bf Z}\left[\mathbb{E}\_{U,W|{\bf Z}}\left[\sum\_{i=1}^n\left(\frac{\mathbb{1}\_{\{Y\_i^{-U\_i}=y\}}}{\sum\_{i=1}^n \mathbb{1}\_{\{Y\_i^{-U\_i}=y\}}}\ell(W,Z^{-U\_i}\_i)-\frac{\mathbb{1}\_{\{Y\_i^{U\_i}=y\}}}{\sum\_{i=1}^n\mathbb{1}\_{\{Y\_i^{U\_i}=y\}}}\ell(W,Z^{U\_i}\_i)\right)\right]\right].
$$
By definition, we know that $\sum\_{i=1}^n\mathbb{1}\_{\{Y\_i^{-U\_i}=y\}}+\sum\_{i=1}^n\mathbb{1}\_{\{Y_i^{U_i}=y\}}=n^y_{\bf Z}$, and I think on average we also have $\mathbb{E}\_{U}\left[\sum\_{i=1}^n\mathbb{1}\_{\{Y_i^{-U_i}=y\}}\right]=\mathbb{E}\_{U}\left[\sum_{i=1}^n\mathbb{1}\_{\{Y_i^{U_i}=y\}}\right]=n^y_{\bf Z}/2$. However, this may still not be sufficient to derive Eq. (6). This also raises the question of whether we can use the class-dependent CMI bounds to bound the standard generalization error, i.e. is Corollary 2 a bound for standard generalization error or some other generalization notion?

Author response:
>**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.

Reviewer comment (Theory):
Regarding Corollary 1: Is it tighter than the classic individual mutual information bound in Bu et al. (2019) when $\sigma_Y$ does not depend on $Y$ (e.g., using loss boundedness instead)? I think this might be accurate due to the following:
$$
\mathbb{E}\_{P_Y}\sqrt{D(P_{W,X|Y}||P_{W}\otimes P_{X|Y})}=\mathbb{E}\_{P_Y}\sqrt{\mathbb{E}\_{P_{W,X|Y}}\log\frac{P_{W,X|Y}P_{Y}}{P_{W}\otimes P_{X|Y}P_{Y}}}=\mathbb{E}\_{P_Y}\sqrt{\mathbb{E}\_{P_{W,X|Y}}\log\frac{P_{W,X,Y}}{P_{W}\otimes P_{X,Y}}}\leq \sqrt{I(W;Z)},
$$
where the last inequality is by Jensen's inequality. If this is correct, you could explicitly state that the individual bound can be recovered from your Corollary 1 or Theorem 1.

Author response:
>**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.

Reviewer comment (Theory):
Could you elaborate more on the last sentence in Remark 2, namely "Therefore, samples from other classes can still affect these bounds, leading to loose bounds on the generalization error of class y"? As $\max(\mathbb{1}\_{\{Y_i^{-U_i}=y\}}, \mathbb{1}\_{\{Y_i^{U_i}=y\}})=\mathbb{1}\_{\{Y_i^{-U_i}=y \\;{\rm or}\\; Y_i^{U_i}=y\}}$. I understand the other parts in Remark 2. However, I don't understand why samples from other classes contribute to the looseness of the bounds.

Author response:
>**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.

Reviewer comment (Theory):
My previous question also raises another concern regarding the comparison between Theorem 4 and Theorem 3: I agree that $I_{\bf Z}(\Delta_y L_i; U_i)\leq I_{\bf Z}(f_W(X_i^{\pm}); U_i)$; however, I would like to point out that we also have $\max(\mathbb{1}\_{\{Y_i^{-U_i}=y\}}, \mathbb{1}\_{\{Y_i^{U_i}=y\}})I_{\bf Z}(f_W(X_i^{\pm}); U_i) \leq I_{\bf Z}(f_W(X_i^{\pm}); U_i)$ (doesn't this imply that $\max(\mathbb{1}\_{\{Y_i^{-U_i}=y\}}, \mathbb{1}\_{\{Y_i^{U_i}=y\}})$ makes Theorem 3 tighter instead of looser?) Therefore, it is uncertain whether Theorem 4 is tighter than Theorem 3 or not. The empirical results suggest that Theorem 4 is tighter, but I hope the authors could clarify why this is expected.

Author response:
>**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.

Reviewer comment (Novelty):
There is a related work that could be included: Hrayr Harutyunyan, et al. "Improving generalization by controlling label-noise information in neural network weights." ICML 2020. In that work, they decompose the mutual information term by chain rule: $I(W;Z)=I(W;X)+I(W;Y|X)$, and using $I(W;Y|X)$ as a regularization term. This might be the first work to incorporate label information into information-theoretic bounds.

Author response:
>**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.

Reviewer comment (Writing):
In Theorem 8 in the Appendix, you may use the widely accepted term for the loss pair-based CMI, namely *evaluated CMI* or *e-CMI*, which was initially introduced in the original CMI paper (see Section 6.2 in the arxiv version of Steinke \& Zakynthinou (2020)).

Author response:
> **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).

Reviewer comment (Writing):
When the authors or the publication are not included in the sentence, the citation should be in parenthesis using **\citep{}** instead of **\citet{}** or **\cite{}**. Most of citations in this paper are not in parenthesis while the authors or the publication are not part of the sentence. For example, in the first sentence of introduction, He \& Tao (2020) should be (He \& Tao (2020)) by using **\citep{}**.

Author response:
>**Q**: Minor comments 

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

Reviewer comment (Novelty):
There do not appear to be any new technical challenges that emerged due to the additional conditioning, and the authors have not mentioned any such challenges.

Author response:
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.

Reviewer comment (Novelty):
Have there been any previous works on computing class-wise generalization bounds, or works on generalization bounds for a subset of a dataset with a particular attribute?

Author response:
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.

Reviewer comment (Theory):
The authors state that information theoretic bounds are natural for this setting as they depend on both the algorithm and the data. It would be useful to discuss this in more detail and explain why it would be hard to obtain class-wise generalization bounds for other types such as stability-based or hypothesis-based bounds. Couldn’t differing properties of the subclasses be captured in some way to get different generalization bounds for different classes?

Author response:
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.