Class-wise Generalization Error: an Information-Theoretic analysis

## Abstract

Abstract Existing generalization theories of supervised learning typically take a holistic approach and provide bounds for the expected generalization over the whole data distribution, which implicitly assumes that the model generalizes similarly for all the classes. In practice, however, there are significant variations in generalization performance among different classes, which cannot be captured by the existing generalization bounds.
In this work, we tackle this problem by theoretically studying the class-generalization error, which quantifies the generalization performance of each individual class.
We derive a novel information-theoretic bound for class-generalization error using the KL divergence, and we further obtain several tighter bounds using the conditional mutual information (CMI), which are significantly easier to estimate in practice.
We empirically validate our proposed bounds in different neural networks and show that they accurately capture the complex class-generalization error behavior. Moreover, we show that the theoretical tools developed in this paper can be applied in several applications beyond this context.

## 1 Introduction

Despite the considerable progress towards a theoretical foundation for neural networks (He & Tao, ), a comprehensive understanding of the generalization behavior of deep learning is still elusive (Zhang et al., ; ). Over the past decade, several approaches have been proposed to uncover and provide a theoretic understanding of the different facets of generalization (Kawaguchi et al., ; He & Tao, ; Roberts et al., ). In particular, multiple tools have been used to characterize the expected generalization error of neural networks, such as VC dimension (Sontag et al., ; Harvey et al., ), algorithmic stability (Bousquet & Elisseeff, ; Hardt et al., ), algorithmic robustness (Xu & Mannor, ; Kawaguchi et al., ), and information-theoretic measures (Xu & Raginsky, ; Steinke & Zakynthinou, ; Wang & Mao, ). However, relying solely on the analysis of the expected generalization over the entire data distribution may not provide a complete picture. One key limitation of the standard expected generalization error is that it does not provide any insight into the class-specific generalization behavior, as it implicitly assumes that the model generalizes similarly for all the classes.

Does the model generalize equally for all classes? To answer this question, we conduct an experiment using deep neural networks, namely ResNet50 (He et al., ) on the CIFAR10 dataset (Krizhevsky et al., ). We plot the standard generalization error along with the class-generalization errors, i.e., the gap between the test error of the samples from the selected class and the corresponding training error, for three different classes of CIFAR10 in Figure [1](#S1.F1) (left). As can be seen, there are significant variants in generalization performance among different classes. For instance, the model overfits the “cats” class, i.e., large generalization error, and generalizes relatively well for the “trucks” class, with a generalization error of the former class consistently 4 times worse than the latter. This suggests that neural networks do not generalize equally for all classes. However, reasoning only with respect to the standard generalization error (red curve) cannot capture this behavior.

Figure: Figure 1: Left: The standard generalization error, i.e., test loss - train loss, and the generalization errors for several classes on CIFAR10 as a function of number of training samples. Right: The standard generalization error, bound proposed by Harutyunyan et al. (), and the generalization errors for several classes on noisy CIFAR10. Experimental details are available in Section [3](#S3).
Refer to caption: extracted/5330490/Figures/introfig4.png

Motivated by these observations, we conduct an additional experiment by introducing label noise ($5\%$) to the CIFAR10 dataset. Results are presented in Figure [1](#S1.F1) (right). Intriguingly, despite the low level of noise, the disparities between the class-wise generalization errors are aggravated, with some classes generalizing up to 8 times worse than others. Further, as shown in this example, different classes do not exhibit the same trend when the number of training samples increases. For instance, unlike the other classes, the generalization error of the “trucks” class increases when more training data is available.
To further illustrate the issue of standard generalization analysis, we plot the information-theoretic generalization bound proposed in Harutyunyan et al. (). Although the bound captures the behavior of the standard generalization error well and can be used to explain the behavior of some classes (e.g., “cat”), it becomes an invalid upper bound for the “trucks” class(^1^11We note that the bound by Harutyunyan et al. () is proposed for the standard generalization error instead of the class-generalization. Here, we plot it only for illustrative purposes.).

When comparing the results on both datasets, it is also worth noting that the generalization error of the same class “trucks” behaves significantly differently on the two datasets. This suggests that class-wise generalization highly depends on factors beyond the class itself, including data distribution, learning algorithm, and the number of training samples. Moreover, in alignment with our findings, Balestriero et al. (); Kirichenko et al. () showed that standard data augmentation and regularization techniques, e.g., weight decay and dropout (Goodfellow et al., ) improve standard average generalization. However, it is surprising to note that these techniques inadvertently exacerbate the performance variations among different classes.

The main conclusion of all the aforementioned observations is that neural networks do not generalize equally for all classes, and their class-wise generalization depends on all ingredients of a learning problem. This paper aims to provide some theoretical understanding of this phenomenon using information-theoretic bounds, as they are both distribution-dependent and algorithm-dependent (Xu & Raginsky, ; Neu et al., ). This makes them an ideal tool to characterize the class-generalization properties of a learning algorithm. A detailed related work discussion is presented in Appendix [A](#A1). Our main contributions are as follows:

- •
We introduce the concept of “class-generalization error,” which quantifies the generalization performance of each individual class. We derive a novel information-theoretic bound for this quantity based on KL divergence (Theorem [1](#Thmtheorem1)). Then, using the super-sample technique by Steinke & Zakynthinou (), we derive various tighter bounds that are significantly easier to estimate and do not require access to the model’s parameters (Theorems [2](#Thmtheorem2),  [3](#Thmtheorem3), and [4](#Thmtheorem4)).
- •
We validate our proposed bounds empirically in different neural networks using CIFAR10 and its noisy variant in Section [3](#S3). We show that the proposed bounds can accurately capture the complex behavior of the class-generalization error behavior in different contexts.
- •
We show that our novel theoretical tools can be applied to the following cases beyond the class-generalization error: (i) study how the class-generalization error affects the standard generalization (Section [4.1](#S4.SS1)); (ii) provide tighter bounds for the subtask problem, where the test data only encompasses a specific subset of the classes encountered during training (Section [4.2](#S4.SS2)); (iii) derive generalization error bounds for learning in the presence of sensitive attributes (Section [4.3](#S4.SS3)).

Notations: We use upper-case letters to denote random variables, e.g., ${\mathbf{Z}}$, and lower-case letters to denote the realization of random variables. $\mathbb{E}_{{\mathbf{Z}}\sim P}$ denotes the expectation of ${\mathbf{Z}}$ over a distribution $P$. Consider a pair of random variables ${\mathbf{W}}$ and ${\mathbf{Z}}=({\mathbf{X}},{\mathbf{Y}})$ with joint distribution $P_{{\mathbf{W}},{\mathbf{Z}}}$. Let $\overline{{\mathbf{W}}}$ be an independent copy of ${\mathbf{W}}$, and $\overline{{\mathbf{Z}}}=(\overline{{\mathbf{X}}},\overline{{\mathbf{Y}}})$ be an independent copy of ${\mathbf{Z}}$, such that $P_{\overline{{\mathbf{W}}},\overline{{\mathbf{Z}}}}=P_{{\mathbf{W}}}\otimes P_
{{\mathbf{Z}}}$. For random variables ${\mathbf{X}}$, ${\mathbf{Y}}$ and ${\mathbf{Z}}$, $I({\mathbf{X}};{\mathbf{Y}})\triangleq D(P_{{\mathbf{X}},{\mathbf{Y}}}\|P_{{
\mathbf{X}}}\otimes P_{{\mathbf{Y}}})$ denotes the mutual information (MI), and $I_{z}({\mathbf{X}};{\mathbf{Y}})\triangleq D(P_{{\mathbf{X}},{\mathbf{Y}}|{
\mathbf{Z}}=z}\|P_{{\mathbf{X}}|{\mathbf{Z}}=z}\otimes P_{{\mathbf{Y}}|{
\mathbf{Z}}=z})$ denotes disintegrated conditional mutual information (CMI), and $\mathbb{E}_{\mathbf{Z}}[I_{{\mathbf{Z}}}({\mathbf{X}};{\mathbf{Y}})]=I({
\mathbf{X}};{\mathbf{Y}}|{\mathbf{Z}})$ is the standard CMI. We will also use the notation ${\mathbf{X}},{\mathbf{Y}}|z$ to simplify ${\mathbf{X}},{\mathbf{Y}}|{\mathbf{Z}}=z$ when it is clear from the context.

## 2 Class-Generalization Error

### 2.1 MI-setting

Typically, in supervised learning, the training set ${\mathbf{S}}=\{({\mathbf{X}}_{i},{\mathbf{Y}}_{i})\}_{i=1}^{n}=\{{\mathbf{Z}}_
{i}\}_{i=1}^{n}$ contains $n$ i.i.d. samples ${\mathbf{Z}}_{i}\in\mathcal{Z}$ generated from the distribution $P_{{\mathbf{Z}}}$. Here, we are interested in the performance of a model with weights $w\in\mathcal{W}$ for data coming from a specific class $y\in\mathcal{Y}$. To this end, we define ${\mathbf{S}}_{y}$ as the subset of ${\mathbf{S}}$ composed of samples only in class $y$. For any model $w\in\mathcal{W}$ and fixed training sets $s$ and $s_{y}$, the class-wise empirical risk can be defined as follows:

$$ $L_{E}(w,s_{y})=\frac{1}{n_{y}}\sum_{(x_{i},y)\in s_{y}}\ell(w,x_{i},y),$ (1) $$

where $n_{y}$ is the size of $s_{y}$ ($n_{y}<n$), and $\ell:\mathcal{W}\times\mathcal{X}\times\mathcal{Y}\to\mathbb{R}_{0}^{+}$ is a non-negative loss function. In addition, the class-wise population risk that quantifies how well $w$ performs on the conditional data distribution $P_{{\mathbf{X}}|{\mathbf{Y}}=y}$ is defined as

$$ $L_{P}(w,P_{{\mathbf{X}}|{\mathbf{Y}}=y})=\mathbb{E}_{P_{{\mathbf{X}}|{\mathbf{ Y}}=y}}[\ell(w,{\mathbf{X}},y)].$ (2) $$

A learning algorithm can be characterized by a randomized mapping from the entire training dataset ${\mathbf{S}}$ to model weights ${\mathbf{W}}$ according to a conditional distribution $P_{{\mathbf{W}}|{\mathbf{S}}}$.
The gap between $L_{P}(w,P_{{\mathbf{X}}|{\mathbf{Y}}=y})$ and $L_{E}(w,s_{y})$ measures how well the trained model ${\mathbf{W}}$ overfits with respect to the data with label $y$,
and the expected class-generalization error is formally defined as follows.

###### Definition 1 .

(class-generalization error)
Given $y\in{\mathcal{Y}}$, the class-generalization error is

$$ $\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})\triangleq\mathbb{E}_{P_{{\mathbf{W}}}}[L_{P}({\mathbf{W}},P_{{ \mathbf{X}}|{\mathbf{Y}}=y})]-\mathbb{E}_{P_{{\mathbf{W}},{\bm{S}}_{y}}}[L_{E} ({\mathbf{W}},{\mathbf{S}}_{y})],$ (3) $$

where $P_{{\mathbf{W}}}$ and $P_{{\mathbf{W}},{\mathbf{S}}_{y}}$ are marginal distributions induced by the learning algorithm $P_{{\mathbf{W}}|{\mathbf{S}}}$ and data generating distribution $P_{{\mathbf{S}}}$.

###### Definition 1 .

#### KL divergence bound

For most learning algorithms used in practice, the index of training samples $i$ will not affect the learned model. Thus, we assume that the learning algorithm is symmetric with respect to each training sample. Formally, this assumption can be expressed as follows:

###### Assumption 1 .

The marginal distribution $P_{{\mathbf{W}},{\mathbf{Z}}_{i}}$ obtained from the joint distribution $P_{{\mathbf{W}},{\mathbf{S}}}$, satisfying
$P_{{\mathbf{W}},{\mathbf{Z}}_{i}}=P_{{\mathbf{W}},{\mathbf{Z}}_{j}}$, $\forall i\neq j$,

Assumption [1](#Thmassumption1) states that the learned model ${\mathbf{W}}$ and each sample ${\mathbf{Z}}_{i}$ has the same joint distribution $P_{{\mathbf{W}},{\mathbf{Z}}}$. Under this assumption, the following lemma simplifies the class-generalization error.

###### Lemma 1 .

Under Assumption [1](#Thmassumption1), the class-generalization error in definition [1](#Thmdefinition1) is given by

$$ $\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})=\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{ \mathbf{X}}}|y}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},y)]- \mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|y}}[\ell({\mathbf{W}},{\mathbf{X}},y)].$ (4) $$

Lemma [1](#Thmlemma1) shows that, similar to the standard generalization error (Xu & Raginsky, ; Bu et al., ; Zhou et al., ), the class-wise generalization error can be expressed as the difference between the loss evaluated
under the joint distribution and the product-of-marginal distribution. The key difference is that both expectations are taken with respect to conditional distributions (${\mathbf{Y}}=y$).

The following theorem provides a bound for the class-generalization error in Definition [1](#Thmdefinition1).

###### Theorem 1 .

For $y\in{\mathcal{Y}}$, assume Assumption [1](#Thmassumption1) holds and the loss $\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},y)$ is $\sigma_{y}$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}|\overline{{
\mathbf{Y}}}=y}$, then the class-generalization error of class $y$ in Definition [1](#Thmdefinition1) can be bounded as:

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\sqrt{2\sigma_{y}^{2}D(P_{{\mathbf{W}},{\mathbf{X}}|y}||P_{{ \mathbf{W}}}\otimes P_{{\mathbf{X}}|{\mathbf{Y}}=y})}.$ (5) $$

The full proof is given in Appendix [B.1](#A2.SS1), which utilizes Lemma [1](#Thmlemma1) and Donsker-Varadhan’s variational representation of the KL divergence.

Theorem [1](#Thmtheorem1) shows that the class-generalization error can be bounded using a class-dependent conditional KL divergence. Theorem [1](#Thmtheorem1) implies that classes with a lower conditional KL divergence between the conditional joint distribution and the product of the marginal can generalize better. To our best knowledge, the bound in Theorem [1](#Thmtheorem1) is the first label-dependent bound that can explain the variation of generalization errors among the different classes.

We note that our class-generalization error bound is obtained by considering the generalization gap of each individual sample with label $y$. This approach, as shown in Bu et al. (); Zhou et al. (); Harutyunyan et al. (), yields tighter bounds using the mutual information (MI) between an individual sample and the output of the learning algorithm, compared to the conventional bounds relying on the MI of the total training set and the algorithm’s output (Xu & Raginsky, ).

###### Assumption 1 .

###### Lemma 1 .

###### Theorem 1 .

### 2.2 CMI-setting

One limitation of the proposed bound in Theorem [1](#Thmtheorem1) is that it can be vacuous or even intractable to estimate in practice, as the bound involves high dimensional entities, i.e., model weights ${\mathbf{W}}$ and dataset ${\mathbf{S}}_{y}$.
The conditional mutual information (CMI) framework, as pioneered by Steinke & Zakynthinou (), has been shown in recent studies (Zhou et al., ; Wang & Mao, ) to offer tighter bounds on generalization error. Remarkably, CMI bounds are always finite even if the weights ${\mathbf{W}}$ are high dimensional and continuous.

In this section, we extend our analysis using the CMI framework. In particular, we assume that there are $n$ super-samples ${\mathbf{Z}}_{[2n]}=({\mathbf{Z}}^{\pm}_{1},\cdots,{\mathbf{Z}}^{\pm}_{n})\in{
\mathcal{Z}}^{2n}$ i.i.d generated from $P_{{\mathbf{Z}}}$. The training data ${\mathbf{S}}=({\bm{Z}}^{{\mathbf{U}}_{1}}_{1},{\mathbf{Z}}^{{\mathbf{U}}_{2}}_
{2},\cdots,{\mathbf{Z}}^{{\mathbf{U}}_{n}}_{n})$ are selected from ${\mathbf{Z}}_{[2n]}$, where ${\mathbf{U}}=({\mathbf{U}}_{1},\cdots,{\mathbf{U}}_{n})\in\{-1,1\}^{n}$ is the selection vector composed of $n$ independent Rademacher random variables. Intuitively, ${\mathbf{U}}_{i}$ selects sample ${{\mathbf{Z}}^{{\mathbf{U}}_{i}}_{i}}$ from ${{\mathbf{Z}}^{\pm}_{i}}$ to be used in training, and the remaining one ${{\mathbf{Z}}^{-{\mathbf{U}}_{i}}_{i}}$ is for the test.

For a specific class $y\in\mathcal{Y}$, let $n^{y}=nP({\mathbf{Y}}=y)$, the number of supersamples n scaled with the probability of class $y$. We define the class-generalization error in the CMI setting as follows:

###### Definition 2 .

(super-sample-based class-generalization error)
For any $y\in{\mathcal{Y}}$, the class-generalization error is defined as

$$ $\small\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W} }|{\mathbf{S}}})\triangleq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y }}\sum_{i=1}^{n}\mathbb{E}_{{\mathbf{U}}_{i},{\bm{W}}|{\mathbf{Z}}_{[2n]}}\big {[}\mathds{1}_{\{Y_{i}^{-U_{i}}=y\}}\ell({\bm{W}},{\mathbf{Z}}^{-{\mathbf{U}}_ {i}}_{i})-\mathds{1}_{\{Y_{i}^{U_{i}}=y\}}\ell({\mathbf{W}},{\mathbf{Z}}^{{ \mathbf{U}}_{i}}_{i})\big{]}\Big{]},$ (6) $$

where $\mathds{1}_{\{a=b\}}$ is the indicator function, returning 1 when $a=b$ and zero otherwise.

Similar to Definition [1](#Thmdefinition1), the class-generalization error in Definition [2](#Thmdefinition2) measures the expected error gap between the training set and the test set relative to one specific class $y$. Compared to the standard generalization error definition typically used in the super-sample setting (Steinke & Zakynthinou, ; Zhou et al., ), we highlight two key differences: (i) our class-wise generalization error involves indicator functions to consider only samples belonging to a specific class $y$. (ii) Our generalization error is normalized by $\smash{n_{{\mathbf{Z}}_{[2n]}}^{y}/2}$, which is half of the total samples in ${\mathbf{Z}}_{[2n]}$ with label $y$. In contrast, the risks are averaged over $n$ samples in the standard super-sample setting.

###### Definition 2 .

#### Class-CMI bound.

The following theorem provides a bound for the super-sample-based class-generalization error using the disintegrated conditional mutual information between ${\mathbf{W}}$ and the selection variable ${\mathbf{U}}_{i}$ conditioned on super-sample ${\mathbf{Z}}_{[2n]}$.

###### Theorem 2 (class-CMI) .

Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, then the class-generalization error for class $y$ in Definition [2](#Thmdefinition2) can be bounded as

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}({\mathbf{W}};{\mathbf{U}}_{i} )}\Big{]}.$ $$

The proof is based on Donsker-Varadhan’s variational representation of the KL divergence and Hoeffding’s Lemma (Hoeffding, ). The key idea of the proof is specifying the function in Donsker-Varadhan’s variational representation to match the definition of class-generalization error. Then, using the fact that for a fixed realization $z_{[2n]}$, $\mathds{1}_{\{y^{{\mathbf{U}}_{i}}=y\}}\ell({\mathbf{W}},z^{{\mathbf{U}}_{i}}_
{i})-\mathds{1}_{\{y^{-{\mathbf{U}}_{i}}=y\}}\ell({\bm{W}},z^{-{\mathbf{U}}_{i
}}_{i})={\mathbf{U}}_{i}(\mathds{1}_{\{y_{i}^{-}=y\}}\ell({\mathbf{W}},z^{-}_{
i})-\mathds{1}_{\{y_{i}^{+}=y\}}\ell({\mathbf{W}},z^{+}_{i}))$ coupled with Hoeffding’s Lemma yields the desired result. The full proof is provided in Appendix [B.2](#A2.SS2).

Theorem [2](#Thmtheorem2) provides a bound of the class-generalization error with explicit dependency on the model weights ${\mathbf{W}}$. It implies that the class-generalization error depends on how much information the random selection reveals about the weights when at least one of the two samples of $z_{i}^{\pm}$ corresponds to the class of interest $y$. Links between overfitting and memorization in the weights have been established in Zhang et al. (); Arpit et al. (); Chatterjee (). Here, we also see that if the model parameters ${\mathbf{W}}$ memorize the random selection ${\mathbf{U}}$, the CMI and the class-generalization error will exhibit high values.

###### Theorem 2 (class-CMI) .

#### Class-f-CMI bound.

While the bound in Theorem [2](#Thmtheorem2) is always finite as ${\mathbf{U}}_{i}$ is binary, evaluating $I_{{\mathbf{Z}}_{[2n]}}({\mathbf{W}};{\mathbf{U}}_{i})$ in practice can be challenging, especially when dealing with high-dimensional ${\mathbf{W}}$ as in deep neural networks. One way to overcome this issue is by considering the predictions of the model $f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})$ instead of the model weights ${\mathbf{W}}$, as proposed by Harutyunyan et al. (). Here, we re-express the loss function $\ell$ based on the prediction $\hat{y}=f_{w}(x)$ as $\ell(w,x,y)=\ell(\hat{y},y)=\ell(f_{w}(x),y)$. Throughout the rest of the paper, we use these two expressions of loss functions interchangeably when it is clear from the context.

In the following theorem, we bound the class-generalization error based on the disintegrated CMI between the model prediction $f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})$ and the random selection, i.e., $I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_
{i})$.

###### Theorem 3 .

(class-f-CMI)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$ is bounded, then the class-generalization error for class $y$ in Definition [2](#Thmdefinition2) can be bounded as

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}} ^{\pm}_{i});{\mathbf{U}}_{i})}\Big{]}.$ $$

The main benefit of the class-f-CMI bound in Theorem [3](#Thmtheorem3), compared to all previously presented bounds, lies in the evaluation of the CMI term involving a typically low-dimensional random variable $f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})$ and a binary random variable ${\mathbf{U}}_{i}$. For example, in the case of binary classification, $f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})$ will be a pair of two binary variables, which enables us to estimate the class-f-CMI bound efficiently and
accurately, as will be shown in Section [3](#S3).

###### Remark 1 .

In contrast to the bound in Theorem [2](#Thmtheorem2), the bound in Theorem [3](#Thmtheorem3) does not require access to the model parameters ${\mathbf{W}}$. It only requires the model output $f(\cdot)$, which makes it suitable even for non-parametric approaches and black-box algorithms.

###### Remark 2 .

Both our bounds in Theorems [2](#Thmtheorem2) and [3](#Thmtheorem3) involve the term $\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+
}=y\}})$, which filters out the CMI terms where neither sample $Z_{i}^{+}$ nor $Z_{i}^{-}$ corresponds to the class $y$. However, this term does not require both samples $Z_{i}^{\pm}$ to belong to class $y$. 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, this term is non-zero and the information from both samples of the pair contributes to the bound ($I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{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$.

###### Theorem 3 .

###### Remark 1 .

###### Remark 2 .

#### Class- Δ y ⁢ L subscript Δ 𝑦 𝐿 \Delta_{y}L roman_Δ start_POSTSUBSCRIPT italic_y end_POSTSUBSCRIPT italic_L -CMI bound.

The term $\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+
}=y\}})$ is a proof artifact and makes the bound inconveniently loose. One solution to overcome this is to consider a new random variable $\Delta_{y}{\mathbf{L}}_{i}$ based on the indicator function and the loss, i.e., $\Delta_{y}{\mathbf{L}}_{i}\triangleq\mathds{1}_{\{y_{i}^{-}=y\}}\ell(f_{{
\mathbf{W}}}({\mathbf{X}}_{i})^{-},y^{-}_{i})-\mathds{1}_{\{y_{i}^{+}=y\}}\ell
(f_{{\mathbf{W}}}({\mathbf{X}}_{i})^{+},y^{+}_{i})$. As shown in Wang & Mao (); Hellström & Durisi (), using the difference of the loss functions on ${\mathbf{Z}}_{i}^{\pm}$ instead of the model output yields tighter generalization bounds for the standard generalization error. We note that $\Delta_{y}{\mathbf{L}}_{i}$ is fundamentally different than ${\mathbf{L}}_{i}$ introduced in Wang & Mao (). ${\mathbf{L}}_{i}$ simply represents the difference in loss functions, while $\Delta_{y}{\mathbf{L}}_{i}$ can be interpreted as a weighted sum of class-dependent losses.
The following Theorem provides a bound based on the CMI between this newly introduced variable and the random selection.

###### Theorem 4 .

(class-$\Delta_{y}L$-CMI)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$, then the class-generalization error of class $y$ defined in [2](#Thmdefinition2) can be bounded as

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}} _{i})}\Big{]}.$ (7) $$

Moreover, we have the $\Delta_{y}L$-CMI bound is always tighter than the class-$f$-CMI bound in Theorem [3](#Thmtheorem3), and the latter is always tighter than the class-CMI bound in Theorem [2](#Thmtheorem2).

Unlike the bound in Theorem [3](#Thmtheorem3), the bound in Theorem [4](#Thmtheorem4) does not directly rely on the model output $f(\cdot)$. Instead, it only requires the loss values for both ${\mathbf{Z}}_{i}^{\pm}$ to compute $\Delta_{y}{\mathbf{L}}_{i}$.

Intuitively, the difference between two weighted loss values, $\Delta_{y}{\mathbf{L}}_{i}$, reveals much less information about the selection process ${\mathbf{U}}_{i}$ compared to the pair $f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})$. Another key advantage of the bound in Theorem [4](#Thmtheorem4) compared to Theorem [3](#Thmtheorem3) is
that computing the CMI term $I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})$ is simpler, given that $\Delta_{y}{\mathbf{L}}_{i}$ is a one-dimensional scalar, as opposed to the two-dimensional $f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})$.

The result of Theorem [4](#Thmtheorem4) provides an interesting interpretation of the class-generalization error in a communication setting. In particular, for fixed $z_{[2n]}$, $P_{\Delta_{y}{\mathbf{L}}_{i}|{\mathbf{U}}_{i}}$ specifies a memoryless channel with input ${\mathbf{U}}_{i}$ and output $\Delta_{y}{\mathbf{L}}_{i}$. Then, the $\Delta_{y}L$-CMI term can be interpreted as “the rate of reliable communication” (Shannon, ; Gallager, ) achievable with the Rademacher random variable ${\mathbf{U}}_{i}$ over this channel $P_{\Delta_{y}{\mathbf{L}}_{i}|{\mathbf{U}}_{i}}$. From this perspective, Theorem [4](#Thmtheorem4) suggests that classes with *lower* rates of reliable communication generalize better.

###### Theorem 4 .

## 3 Empirical Evaluations

In this section, we conduct empirical experiments to evaluate the effectiveness of our class-wise generalization error bounds. As mentioned earlier, The bounds in Section [2.2](#S2.SS2) are significantly easy to estimate in practical scenarios. Here, we evaluate the error bounds in Theorems [3](#Thmtheorem3) and  [4](#Thmtheorem4) for deep neural networks.

We follow the same experimental settings in Harutyunyan et al. (), i.e., we fine-tune a ResNet-50 (He et al., ) on the CIFAR10 dataset (Krizhevsky et al., ) (Pretrained on ImageNet (Deng et al., )). Moreover, to understand how well our bounds perform in a more challenging situation and to further highlight their effectiveness, we conduct an additional experiment with a noisy variant (5% label noise) of CIFAR10. The experimental details are provided in Appendix [C.1](#A3.SS1).

The class-wise generalization error of two classes from CIFAR10 “trucks” and “cats”, along with the bounds in Theorems [3](#Thmtheorem3) and  [4](#Thmtheorem4) are presented in the first two columns of Figure [2](#S3.F2). The results on all the 10 classes for both datasets are presented in Appendix [C.3](#A3.SS3).

Figure: Figure 2: Experimental results of class-generalization error and our bounds in Theorems [3](#Thmtheorem3) and [4](#Thmtheorem4) for the class of “trucks” (left) and “cats” (middle) in CIFAR10 (top) and noisy CIFAR10 (bottom), as we increase the number of training samples. In the right column, we provide the scatter plots between the bound in Theorem [4](#Thmtheorem4) and the class-generalization error of the different classes for CIFAR10 (top) and noisy CIFAR10 (bottom).
Refer to caption: extracted/5330490/Figures/Cifar10/cifar10classtrucksgenerrorplot.png

Figure [2](#S3.F2) shows that both bounds can capture the behavior of the class-generalization error. As expected, the class-$\Delta_{y}L$-CMI is consistently tighter and more stable compared to the class-$f$-CMI bound for all the different scenarios. For CIFAR10 in Figure [2](#S3.F2) (top), as we increase the number of training samples, the “trucks” class has a relatively constant class-generalization error, while the “cats” class has a large slope at the start and then a steady incremental decrease.
For both classes, the class-$\Delta_{y}L$-CMI precisely predicts the behavior of class-generalization error.

The results on noisy CIFAR10 in Figure [2](#S3.F2) (bottom) and the results in Appendix [C.3](#A3.SS3) are consistent with these observations. Notably, the “trucks” generalization error decreases for CIFAR10 and increases for noisy CIFAR10 with respect to the number of samples. Moreover, the class-generalization error of “cat” is worse than “trucks” in CIFAR10, but the opposite is true for the noisy CIFAR10. All these complex behaviors of class-generalization errors are successfully captured by the class-$\Delta_{y}L$-CMI bound.

As can be seen from the left and middle plots in Figure [2](#S3.F2), the class-$\Delta_{y}L$-CMI bound scales proportionally to the true class-generalization error, i.e., higher class-$\Delta_{y}L$-CMI bound value corresponds to higher class-generalization error. To further highlight this dependency, we plot in Figure [2](#S3.F2) (right) the scatter plot between the different class-generalization errors and their corresponding class-$\Delta_{y}L$-CMI bound values for the different classes in CIFAR10 (top) and Noisy CIFAR10 (bottom) under different number of samples. We observe that our bound is linearly correlated with the class-generalization error and can efficiently predict its behavior. We also note that $f$-CMI bound exhibits a similar behavior as shown in Appendix [C.3](#A3.SS3). This suggests that one can use our bound to predict which classes will generalize better than others. Additional results are provided in Appendix [C.5](#A3.SS5).

## 4 Other applications

Besides enabling us to study class-wise generalization errors, the tools developed in this work can also be used to provide theoretical insights into several other applications. In this Section, we explore several use cases for the developed tools.

### 4.1 From class-generalization error to standard generalization error

In this subsection, we study the connection between the standard expected generalization error and the class-generalization error. We extend the bounds presented in Section [2](#S2) into class-dependent expected generalization error bounds.

First, we notice that taking the expectation over $P_{\mathbf{Y}}$ for the class-generalization error defined in [1](#Thmdefinition1) yields the standard expected generalization error.
Thus, we obtain the first class-dependent bound for the standard generalization error by taking the expectation of $y\sim P_{{\mathbf{Y}}}$ in Theorem [1](#Thmtheorem1).

###### Corollary 1 .

Assume that for every $y\in{\mathcal{Y}}$, the loss $\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},y)$ is $\sigma_{y}$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}|\overline{{
\mathbf{Y}}}=y}$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{Y^{\prime}}\bigg{[}\sqrt{2\sigma_{{\mathbf{Y}}^{ \prime}}^{2}D(P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{X}}|{\mathbf{Y} }={\mathbf{Y}}^{\prime}}||P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}}|{\mathbf{Y}} ={\mathbf{Y}}^{\prime}})}\bigg{]}.$ (8) $$

We note that in the case where sub-gaussian parameter $\sigma_{y}=\sigma$ is independent of $y$, we can further show that the bound in [1](#Thmcorollary1) is tighter than the individual sample bound in Bu et al. (). The proof is available in Appendix [C.4](#A3.SS4).
We also extend the results in Theorems [2](#Thmtheorem2),  [3](#Thmtheorem3), and [4](#Thmtheorem4) into standard generalization bounds. For example, in Corollary [2](#Thmcorollary2), we provide such an extension of Theorem [4](#Thmtheorem4).

###### Corollary 2 .

Assume that the loss $\ell(\hat{y},y)\in[0,1]$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\bm{Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]} }\Big{[}\frac{1}{n^{\mathbf{Y}}}\sum_{i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}( \Delta_{{\mathbf{Y}}}{\mathbf{L}}_{i};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ (9) $$

To the best of our knowledge, Corollaries [1](#Thmcorollary1) and [2](#Thmcorollary2) are the first generalization bounds to provide explicit dependency on the labels. While prior bounds (Steinke & Zakynthinou, ; Harutyunyan et al., ; Wang & Mao, ; ) are tight and efficient to estimate, they fail to provide any insights on how different classes affect the standard generalization error. The results presented here address this gap and provide explicit label-dependent bounds. It shows that the generalization error of an algorithm scales as the expectation of the bound depending only on a single class.
Moreover, in classification with $m$ classes, the bounds become a sum of each class-generalization error weighted by the probability of the class, i.e., $P({\mathbf{Y}}=y)$. This shows that classes with a higher occurrence probability affect the standard generalization error more. From this perspective, our results can also provide insights into developing algorithms with better generalization capabilities by focusing on the class-generalization error. For example, one can employ augmentation techniques targeted at the classes with higher class-generalization error bounds to attenuate their respective class error and thus improve the standard generalization of the model.

###### Corollary 1 .

###### Corollary 2 .

### 4.2 Sub-task problem

Here, “subtask problem” refers to a specific case of distribution shift in supervised learning, where the training data generated from the source domain $P_{{\mathbf{X}},{\mathbf{Y}}}$ consists of multiple classes, while the test data for the target domain $Q_{{\mathbf{X}},{\mathbf{Y}}}$ only encompasses a specific known subset of the classes encountered during training. This problem is motivated by the situation where a large model has been trained on numerous classes, potentially over thousands, but is being utilized in a target environment where only a few classes, observed during training, exist. By tackling the problem as a standard domain adaptation task, the generalization error of the subtask problem can be bounded as follows:

$$ $\overline{\mathrm{gen}}_{Q,E_{P}}\triangleq\mathbb{E}_{P_{{\mathbf{W}},{ \mathbf{S}}}}[L_{Q}({\mathbf{W}})-L_{E}({\mathbf{W}},{\mathbf{S}})]\leq\sqrt{2 \sigma^{2}D(Q_{{\mathbf{X}},{\mathbf{Y}}}\|P_{{\mathbf{X}},{\mathbf{Y}}})}+ \sqrt{2\sigma^{2}I({\mathbf{W}};{\mathbf{S}})},$ (10) $$

where $L_{Q}(w)=L_{P}(w,Q_{{\mathbf{X}},{\mathbf{Y}}})$ denotes the population risk of $w$ under distribution $Q_{{\mathbf{X}},{\mathbf{Y}}}$.
The full details of deriving such bound are available in appendix [D.2](#A4.SS2). We note that (Wu et al., ) further tightens the result in equation [10](#S4.E10), but these bounds are all based on the KL divergence $D(Q_{{\mathbf{X}},{\mathbf{Y}}}\|P_{{\mathbf{X}},{\mathbf{Y}}})$ for any generic distribution shift problem and do not leverage the fact that the target task is encapsulated in the source task.

Obtaining tighter generalization error bounds for the subtask problem is straightforward using our class-wise generalization bounds. In fact, the generalization error bound of the subtask can be obtained by summing the class-wise generalization over the space of the subtask classes ${\mathcal{A}}$. Formally, by taking the expectation of ${\mathbf{Y}}\sim Q_{{\mathbf{Y}}}$, we obtain the following notion of the subtask generalization error:

$$ $\overline{\mathrm{gen}}_{Q,E_{Q}}\triangleq\mathbb{E}_{Q_{\mathbf{Y}}}\big{[} \overline{\mathrm{gen}}_{\mathbf{Y}}\big{]}=\mathbb{E}_{P_{{\mathbf{W}},{ \mathbf{S}}}}[L_{Q}(w)-L_{E_{Q}}({\mathbf{W}},{\mathbf{S}})],$ (11) $$

where $L_{E_{Q}}(w,S)=\frac{1}{n_{{\mathcal{A}}}}\sum_{y_{i}\in{\mathcal{A}}}\ell(w,x
_{i},y_{i})$ is the empirical risk relative to the target domain $Q$, and $n_{{\mathcal{A}}}$ is the number of samples in $S$ such that their labels $y_{i}\in{\mathcal{A}}$.
We are interested in deriving generalization bounds for $\overline{\mathrm{gen}}_{Q,E_{Q}}$, as it only differs from $\overline{\mathrm{gen}}_{Q,E_{P}}$ by the difference in the empirical risk $L_{E_{Q}}({\mathbf{W}},{\mathbf{S}})-L_{E}({\mathbf{W}},{\mathbf{S}})$, which can be computed easily in practice.

Using Jensen’s inequality, we have $|\overline{\mathrm{gen}}_{Q,E_{Q}}|=|\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf{Y
}}}\big{[}\overline{\mathrm{gen}}_{\mathbf{Y}}\big{]}|\leq\mathbb{E}_{{\mathbf
{Y}}\sim Q_{\mathbf{Y}}}\big{[}|\overline{\mathrm{gen}}_{\mathbf{Y}}|\big{]}$. Thus, we can use the results from Section [2](#S2) to obtain tighter bounds. For example, using Theorem [4](#Thmtheorem4), we can obtain the corresponding subtask generalization error bound in Theorem [6](#Thmtheorem6).

###### Theorem 5 .

(subtask-CMI)
Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, then the subtask generalization error defined in [11](#S4.E11) can be bounded as

$$ $|\overline{\mathrm{gen}}_{Q,E_{Q}}|\leq\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf {Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{\mathbf{Y}}} \sum_{i=1}^{n}\sqrt{2\max(\mathds{1}_{{\mathbf{Y}}_{i}^{-}={\mathbf{Y}}}, \mathds{1}_{{\mathbf{Y}}_{i}^{+}={\mathbf{Y}}})I_{{\mathbf{Z}}_{[2n]}}({ \mathbf{W}};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ $$

###### Theorem 6 .

(subtask-$\Delta~{}L_{y}$-CMI)
Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, Then the subtask generalization error defined in [11](#S4.E11) can be bounded as

$$ $|\overline{\mathrm{gen}}_{Q,E_{Q}}|\leq\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf {Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{\mathbf{Y}}} \sum_{i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}(\Delta_{{\mathbf{Y}}}{\mathbf{L}} _{i};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ $$

Similarly, we can extend Theorems [3](#Thmtheorem3) and [2](#Thmtheorem2) to use the model’s output or weights instead of $\Delta_{y}{\mathbf{L}}_{i}$.

###### Remark 3 .

Existing distribution shift bounds, e.g., the bound in equation [10](#S4.E10), typically depend on some measure that quantifies the discrepancy between the target and domain distributions, e.g., KL divergence. However, as can be seen in Theorem [6](#Thmtheorem6), the proposed bounds here are discrepancy-independent and involve only a quantity analog to the mutual information term in equation [10](#S4.E10).

###### Theorem 5 .

###### Theorem 6 .

###### Remark 3 .

### 4.3 Generalization certificates with sensitive attributes

One main concern hindering the use of machine learning models in high-stakes applications is the potential biases on sensitive attributes such as gender and skin color (Mehrabi et al., ; Barocas et al., ). Thus, it is critical not only to reduce sensitivity to such attributes but also to be able to provide guarantees on the fairness of the models (Holstein et al., ; Rajkomar et al., ). One aspect of fairness is that the machine learning model should generalize equally well for each minor group with different sensitive attributes (Barocas et al., ; Williamson & Menon, ).

By tweaking the definition of our class-generalization error, we show that the theoretical tools developed in this paper can be used to obtain bounds for attribute-generalization errors. Suppose we have a random variable ${\mathbf{T}}\in{\mathcal{T}}$ representing a sensitive feature. One might be interested in studying the generalization of the model for the sub-population with the attribute ${\mathbf{T}}=t$. Inspired by our class-generalization, we define the attribute-generalization error as follows:

###### Definition 3 .

(attribute-generalization error)
Given $t\in{\mathcal{T}}$, the attribute-generalization error is defined as follows:

$$ $\overline{\mathrm{gen}_{t}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})=\mathbb{E}_{P_{{\mathbf{W}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}} =t}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{Z}}})]-\mathbb{E}_{P_{{ \mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}}[\ell({ \mathbf{W}},{\mathbf{Z}})].$ (12) $$

By exchanging ${\mathbf{X}}$ and ${\mathbf{Y}}$ with ${\mathbf{Z}}$ and ${\mathbf{T}}$ in Theorem [1](#Thmtheorem1), respectively, we can show the following bound for the attribute-generalization error.

###### Theorem 7 .

Given $t\in{\mathcal{T}}$, assume that the loss $\ell({\mathbf{W}},{\mathbf{Z}})$ is $\sigma$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{Z}}}}$, then the attribute-generalization error of the sub-population ${\mathbf{T}}=t$, can be bounded as follows:

$$ $|\overline{\mathrm{gen}_{t}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\sqrt{2\sigma^{2}D(P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{ \mathbf{Z}}|{\mathbf{T}}=t}||P_{{{\mathbf{W}}}}\otimes P_{{{\mathbf{Z}}}|{{ \mathbf{T}}}=t})}.$ (13) $$

We note extending our results to the CMI settings is also straightforward. Using the attribute generalization, we can show that the standard expected generalization error can be bounded as follows:

###### Corollary 3 .

Assume that the loss $\ell({\mathbf{W}},{\mathbf{Z}})$ is $\sigma$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{Z}}}}$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{T}}^{\prime}}\Big{[}\sqrt{2\sigma^{2}D( P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}={\mathbf{T}} ^{\prime}}||P_{{{\mathbf{W}}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}={\mathbf{T} }^{\prime}})}\Big{]}.$ (14) $$

The result of Corollary [3](#Thmcorollary3) shows that the average generalization error is upper-bounded by the expectation over the attribute-wise generalization. This shows that it is possible to improve the overall generalization by reducing the generalization of each population relative to the sensitive attribute.

###### Definition 3 .

###### Theorem 7 .

###### Corollary 3 .

## 5 Conclusion & Future Work

This paper studied the deep learning generalization puzzle of the noticeable heterogeneity of overfitting among different classes by introducing and exploring the concept of “class-generalization error”. To our knowledge, we provided the first rigorous generalization bounds for this concept in the MI and CMI settings. We also empirically strengthened the findings with supporting experiments validating the efficiency of the proposed bounds. Furthermore, we demonstrated the versatility of our theoretical tools in providing tight bounds for various contexts.

Overall, our goal is to understand generalization in deep learning through the lens of information theory, which motivates future work on how to prevent high class-generalization error variability and ensure ’equal’ generalization among the different classes. Other possible future research endeavors focus on obtaining tighter bounds for the class-generalization error and studying this concept in different contexts beyond supervised learning, e.g., transfer and self-supervised learning.

## Appendix A Appendix: Related work

Information-theoretic generalization error bounds: Information-theoretic bounds have attracted a lot of attention recently to characterize the generalization of learning algorithms (Neu et al., ; Wang et al., ; Aminian et al., ; Wu et al., ; Wang & Mao, ; Modak et al., ; Wang & Mao, ; Shui et al., ; Wang et al., ; Alabdulmohsin, ). In the supervised learning context, several standard generalization error bounds have been proposed to based on different tools, e.g., KL divergence (Zhou et al., ), Wasserstein distance (Rodríguez Gálvez et al., ), mutual information between the samples and the weights (Xu & Raginsky, ; Bu et al., ). Recently, it was shown that tighter generalization bounds could be obtained based on the conditional mutual information (CMI) setting (Steinke & Zakynthinou, ; Zhou et al., ). Based on this framework, Harutyunyan et al. () derived $f$-CMI bounds based on the model output. In Hellström & Durisi (), tighter bounds have been obtained based on the CMI of the loss function, which is further tightened by Wang & Mao () using the $\delta L$ CMI.

Class-dependent analysis:
Incorporating label information in generalization analysis is not entirely new (He & Su, ; Chen et al., ; Deng et al., ). For example, in He & Su (), the questions “When and how does the update of weights of neural networks using induced gradient at an example impact the prediction at another example?” have been extensively studied, and it was observed that the impact is significant if the two samples are from the same class. In Balestriero et al. (); Kirichenko et al. (); Bitterwolf et al. (); Lee et al. (), it has been showed that while standard data augmentation techniques (Goodfellow et al., ) help improve overall performance, it yields lower performance on minority classes.
From a theoretical perspective, (Deng et al., ) noticed that in uniform stability context, the sensitivity of neural networks depends highly on the label information and thus proposed the concept of “Locally Elastic Stability” to derive tighter algorithmic stability generalization bounds.
In Tishby et al. (), an information bottleneck principle is proposed which states that an optimal feature map simultaneously minimizes its mutual information with the feature distribution and maximizes its mutual information with the label distribution, thus incorporating the class information. In Tishby & Zaslavsky (); Saxe et al. (); Kawaguchi et al. (), this principle was used to explain the generalization of neural networks.

## Appendix B Appendix: Proofs of the Theorems in Section 2

This appendix includes the missing proofs of the results presented in the main text in Section [2](#S2).

###### Lemma 2 .

Let ${\mathbf{X}}$ be a bounded random variable, i.e., ${\mathbf{X}}\in[a,b]$ almost surely. If $\mathbb{E}[{\mathbf{X}}]=0$, then ${\mathbf{X}}$ is $(b-a)$-subgaussian and we have:

$$ $\mathbb{E}[e^{\lambda{\bm{X}}}]\leq e^{\frac{\lambda^{2}(b-a)^{2}}{8}},\>\> \forall\lambda\in{\mathbb{R}}.$ (15) $$

Next, we need the following Lemma [3](#Thmlemma3) as a main tool to prove Theorem [2](#Thmtheorem2), [3](#Thmtheorem3), and [8](#Thmtheorem8) in Section [2.2](#S2.SS2) in the CMI setting.

###### Lemma 3 .

Consider the CMI setting. Let ${\mathbf{V}}\in\mathcal{V}$ be a random variable, possibly depending on ${\bm{W}}$. For any function g that can be written as $g({\mathbf{V}},{\mathbf{U}}_{i},z_{[2n]})=\mathds{1}_{\{y^{{\mathbf{U}}_{i}}=y
\}}h({\mathbf{V}},z^{{\mathbf{U}}_{i}}_{i})-\mathds{1}_{\{y^{-{\mathbf{U}}_{i}
}=y\}}h({\bm{V}},z^{-{\mathbf{U}}_{i}}_{i})$ such that $h\in[0,1]$ is a bounded function, we have

$$ $\mathbb{E}_{{\mathbf{V}};{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[g({ \mathbf{V}},{\mathbf{U}}_{i},z_{[2n]})]\leq\sqrt{2\max(\mathds{1}_{\{y_{i}^{-} =y\}},\mathds{1}_{\{y_{i}^{+}=y\}})I_{z_{[2n]}}({\mathbf{V}};{\mathbf{U}}_{i})}.$ (16) $$

###### Proof.

Let $(\overline{{\mathbf{V}}},\overline{{\bm{U}}_{i}})$ be an independent copy of $({\mathbf{V}},{\mathbf{U}}_{i})$. The disintegrated mutual information $I_{z_{[2n]}}({\mathbf{V}};{\mathbf{U}}_{i})$ is equal to:

$$ $I_{z_{[2n]}}({\mathbf{U}}_{i};{\mathbf{V}})=D\big{(}P_{{\mathbf{V}},{\mathbf{U }}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}\|P_{{\mathbf{V}}|{\mathbf{Z}}_{[2n]}=z_{[ 2n]}}P_{{\mathbf{U}}_{i}}\big{)},$ (17) $$

Thus, by the Donsker–Varadhan variational representation of KL
divergence, $\forall\lambda\in{\mathbb{R}}$ and for every function $g$, we have

$$ $I_{z_{[2n]}}({\mathbf{V}};{\mathbf{U}}_{i})\geq\lambda\mathbb{E}_{{\mathbf{V}} ,{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[g({\mathbf{V}},{\mathbf{U}}_{i },z_{[2n]})]-\log\mathbb{E}_{\overline{{\mathbf{V}}},\overline{{\mathbf{U}}}_{ i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{{\mathbf{V}}}, \overline{{\mathbf{U}}}_{i},z_{[2n]})}].$ (18) $$

Next, let $g({\mathbf{V}},{\mathbf{U}}_{i},z_{[2n]})=\mathds{1}_{\{y^{{\mathbf{U}}_{i}}=y
\}}h({\mathbf{V}},z^{{\mathbf{U}}_{i}}_{i})-\mathds{1}_{\{y^{-{\mathbf{U}}_{i}
}=y\}}h({\bm{V}},z^{-{\mathbf{U}}_{i}}_{i})$. It is easy to see that $g(\overline{{\mathbf{V}}},\overline{{\mathbf{U}}}_{i},z_{[2n]})$ can be rewritten as follows:

$$ $g(\overline{{\mathbf{V}}},\overline{{\mathbf{U}}}_{i},z_{[2n]})=\overline{{ \mathbf{U}}}_{i}(\mathds{1}_{\{y_{i}^{-}=y\}}h(\overline{{\mathbf{V}}},z^{-}_{ i})-\mathds{1}_{\{y_{i}^{+}=y\}}h(\overline{{\mathbf{V}}},z^{+}_{i})).$ (19) $$

Thus, we have

$$ $\log\mathbb{E}_{\overline{{\mathbf{V}}},\overline{{\mathbf{U}}}_{i}|{\mathbf{Z }}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{{\mathbf{V}}},\overline{{\mathbf{U} }}_{i},z_{[2n]})}]=\log\mathbb{E}_{\overline{{\mathbf{V}}},\overline{{\mathbf{ U}}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[e^{\lambda\overline{{\mathbf{U}}}_{i}( \mathds{1}_{\{y_{i}^{-}=y\}}h(\overline{{\mathbf{V}}},z^{-}_{i})-\mathds{1}_{ \{y_{i}^{+}=y\}}h(\overline{{\mathbf{V}}},z^{+}_{i}))}].$ (20) $$

Note that $\mathbb{E}_{\overline{{\mathbf{U}}}_{i}}[\overline{{\mathbf{U}}}_{i}(\mathds{1
}_{\{y_{i}^{-}=y\}}h(\overline{{\mathbf{V}}},z^{-}_{i})-\mathds{1}_{\{y_{i}^{+
}=y\}}h(\overline{{\mathbf{V}}},z^{+}_{i}))]=0$ and $\overline{{\mathbf{U}}}_{i}\in\{-1,+1\}$. Thus, using Hoeffding’s Lemma, we have

$$ $\log\mathbb{E}_{\overline{{\mathbf{V}}},\overline{{\mathbf{U}}}_{i}|{\mathbf{Z }}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{{\mathbf{V}}},\overline{{\mathbf{U} }}_{i},z_{[2n]})}]\leq\log\mathbb{E}_{\overline{{\mathbf{V}}}|{\mathbf{Z}}_{[2 n]}=z_{[2n]}}[e^{\frac{\lambda^{2}}{2}\big{(}\mathds{1}_{\{y_{i}^{-}=y\}}h( \overline{{\mathbf{V}}},z^{-}_{i})-\mathds{1}_{\{y_{i}^{+}=y\}}h(\overline{{ \mathbf{V}}},z^{+}_{i})\big{)}^{2}}].$ (21) $$

Next, as $h\in[0,1]$, $\Big{|}\mathds{1}_{\{y_{i}^{-}=y\}}h(\overline{{\mathbf{V}}},z^{-}_{i})-
\mathds{1}_{\{y_{i}^{+}=y\}}h(\overline{{\mathbf{V}}},z^{+}_{i})\Big{|}\leq
\max(\mathds{1}_{\{y_{i}^{-}=y\}},\mathds{1}_{\{y_{i}^{+}=y\}})$.
Thus,

$$ $\log\mathbb{E}_{\overline{{\mathbf{V}}},\overline{{\mathbf{U}}}_{i}|{\mathbf{Z }}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{{\mathbf{V}}},\overline{{\mathbf{U} }}_{i},z_{[2n]})}]\leq\frac{\lambda^{2}}{2}\max(\mathds{1}_{\{y_{i}^{-}=y\}}, \mathds{1}_{\{y_{i}^{+}=y\}})^{2}=\frac{\lambda^{2}}{2}\max(\mathds{1}_{\{y_{i }^{-}=y\}},\mathds{1}_{\{y_{i}^{+}=y\}}).$ (22) $$

Replacing in equation [18](#A2.E18), we have

$$ $I_{z_{[2n]}}({\mathbf{V}};{\mathbf{U}}_{i})\geq\lambda\mathbb{E}_{{\mathbf{V}} ,{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[\mathds{1}_{y^{{\mathbf{U}}_{i }}=y}h({\mathbf{V}},z^{{\mathbf{U}}_{i}}_{i})-\mathds{1}_{\{y^{-{\mathbf{U}}_{ i}}=y\}}h({\bm{W}},z^{-{\mathbf{U}}_{i}}_{i})]\\ -\frac{\lambda^{2}}{2}\max(\mathds{1}_{\{y_{i}^{-}=y\}},\mathds{1}_{\{y_{i}^{+ }=y\}}).$ (23) $$

So, $\forall\lambda\in{\mathbb{R}}$

$$ $\frac{\lambda^{2}}{2}\max(\mathds{1}_{\{y_{i}^{-}=y\}},\mathds{1}_{\{y_{i}^{+} =y\}})-\lambda\mathbb{E}_{{\mathbf{V}};{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_ {[2n]}}[g({\mathbf{V}},{\mathbf{U}}_{i},z_{[2n]})]+I_{z_{[2n]}}({\mathbf{V}};{ \mathbf{U}}_{i})\geq 0.$ (24) $$

The equation [24](#A2.E24) is a non-negative parabola with respect to $\lambda$. Thus its discriminant must be non-positive. This implies

$$ $\mathbb{E}_{{\mathbf{V}};{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[g({ \mathbf{V}},{\mathbf{U}}_{i},z_{[2n]})]\leq\sqrt{2\max(\mathds{1}_{\{y_{i}^{-} =y\}},\mathds{1}_{\{y_{i}^{+}=y\}})I_{z_{[2n]}}({\mathbf{V}};{\mathbf{U}}_{i})}.$ (25) $$

∎

###### Lemma 2 .

###### Lemma 3 .

###### Proof.

### B.1 Proof of Theorem 1

Theorem [1](#Thmtheorem1) (restated) For $y\in{\mathcal{Y}}$, assume Assumption [1](#Thmassumption1) holds and the loss $\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},y)$ is $\sigma_{y}$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}|\overline{{
\mathbf{Y}}}=y}$, then the class-generalization error of class $y$ in Definition [1](#Thmdefinition1) can be bounded as:

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\sqrt{2\sigma_{y}^{2}D(P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf {Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}}|{\mathbf{Y}}=y})}.$ (26) $$

###### Proof.

From lemma [1](#Thmlemma1), we have

$$ $\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})=\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{ \mathbf{X}}}|{\mathbf{Y}}=y}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{ X}}},y)]-\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}}[\ell({ \mathbf{W}},{\mathbf{X}},y)].$ (27) $$

Using the Donsker–Varadhan variational representation of the relative entropy, we have

$$ $D(P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{ \mathbf{X}}|{\mathbf{Y}}=y})\geq\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|{ \mathbf{Y}}=y}}[\lambda\ell({\mathbf{W}},{\mathbf{X}},y)]\\ -\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}} |\overline{{\mathbf{Y}}}=y}}[e^{\lambda\ell(\overline{{\mathbf{W}}},\overline{ {\mathbf{X}}},y)}],\forall\lambda\in{\mathbb{R}}.$ (28) $$

On the other hand, we have:

$$ $\displaystyle\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{ \mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}\Big{[}e^{\lambda\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{X}}},y)-\lambda\mathbb{E}[\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{X}}},y)]}\Big{]}$ $\displaystyle=\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{ {\mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}\Big{[}e^{\lambda\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{X}}},y)}e^{-\lambda\mathbb{E}[\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{X}}},y)]\big{)}}\Big{]}$ $\displaystyle=\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{ {\mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}[e^{\lambda\ell(\overline{{\mathbf{W} }},\overline{{\mathbf{X}}},y)}]-\lambda\mathbb{E}_{P_{\overline{{\mathbf{W}}}} \otimes P_{\overline{{\mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}[\ell(\overline{ {\mathbf{W}}},\overline{{\mathbf{X}}},y)].$ $$

Using the sub-gaussian assumption, we have

$$ $\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}| \overline{{\mathbf{Y}}}=y}}[e^{\lambda\ell(\overline{{\mathbf{W}}},\overline{{ \mathbf{X}}},y)}]\leq\lambda\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{ \overline{{\mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}(\ell(\overline{{\mathbf{W} }},\overline{{\mathbf{X}}},y))+\frac{\lambda^{2}\sigma_{y}^{2}}{2}.$ (29) $$

By replacing in equation [28](#A2.E28), we have

$$ $D(P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{ \mathbf{X}}|{\mathbf{Y}}=y})\geq\lambda\big{(}\mathbb{E}_{P_{{\mathbf{W}},{ \mathbf{X}}|{\mathbf{Y}}=y}}[\ell({\mathbf{W}},{\mathbf{X}},y)]-\\ \mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}| \overline{{\mathbf{Y}}}=y}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X} }},y)]\big{)}-\frac{\lambda^{2}\sigma_{y}^{2}}{2}.$ (30) $$

Thus,

$$ $D(P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{ \mathbf{X}}|{\mathbf{Y}}=y})-\lambda(\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}| {\mathbf{Y}}=y}}[\ell({\mathbf{W}},{\mathbf{X}},y)]-\mathbb{E}_{P_{\overline{{ \mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}[ \ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},y)])\\ +\lambda^{2}\sigma_{y}^{2}\geq 0,\forall\lambda\in{\mathbb{R}}.$ (31) $$

Equation equation [31](#A2.E31) is a non-negative parabola with respect to $\lambda$, which implies its discriminant must be non-positive. Thus,

$$ $|\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}}[\ell({\mathbf{W}},{ \mathbf{X}},y)]-\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{ \mathbf{X}}}|\overline{{\mathbf{Y}}}=y}}[\ell(\overline{{\mathbf{W}}}, \overline{{\mathbf{X}}},y)]|\leq\sqrt{2\sigma_{y}^{2}D(P_{{\mathbf{W}},{ \mathbf{X}}|{\mathbf{Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}}|{\mathbf{ Y}}=y})}.$ (32) $$

This completes the proof.
∎

###### Proof.

### B.2 Proof of Theorem 2

Theorem [2](#Thmtheorem2) (restated)
Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, then the class-generalization error for class $y$ in Definition [2](#Thmdefinition2) can be bounded as

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}({\mathbf{W}};{\mathbf{U}}_{i} )}\Big{]}.$ (33) $$

###### Proof.

Using Lemma [3](#Thmlemma3) with ${\mathbf{V}}={\mathbf{W}}$ and $h({\mathbf{V}},z)=\ell({\mathbf{W}},z)$ in, we have

$$ $\mathbb{E}_{{\mathbf{W}};{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[g({ \mathbf{W}},{\mathbf{U}}_{i},z_{[2n]})]\leq\sqrt{2\max(\mathds{1}_{\{y_{i}^{-} =y\}},\mathds{1}_{\{y_{i}^{+}=y\}})I_{z_{[2n]}}({\mathbf{W}};{\mathbf{U}}_{i})},$ (34) $$

where $g({\mathbf{W}},{\mathbf{U}}_{i},z_{[2n]})=\mathds{1}_{\{y^{{\mathbf{U}}_{i}}=y
\}}\ell({\mathbf{W}},z^{{\mathbf{U}}_{i}}_{i})-\mathds{1}_{\{y^{-{\mathbf{U}}_
{i}}=y\}}\ell({\bm{W}},z^{-{\mathbf{U}}_{i}}_{i})$. Thus, by summing over the different terms in Definition [2](#Thmdefinition2) and taking expectation over ${\mathbf{Z}}_{[2n]}$,

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}({\mathbf{W}};{\mathbf{U}}_{i} )}\Big{]}.$ (35) $$

∎

###### Proof.

### B.3 Proof of Theorem 3

Theorem [3](#Thmtheorem3) (restated) Assume that the loss $\ell({\mathbf{W}},{\mathbf{X}},y)\in[0,1]$. The class-generalization error of class $y$, as defined in [2](#Thmdefinition2), can be bounded as follows:

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}} ^{\pm}_{i});{\mathbf{U}}_{i})}\Big{]}.$ (36) $$

###### Proof.

Similar to the proof of Theorem  [2](#Thmtheorem2).
Using Lemma [3](#Thmlemma3) with ${\mathbf{V}}=f_{{\mathbf{W}}}(x^{\pm}_{i})$ and $h({\mathbf{V}},z_{i})=\ell(f_{{\mathbf{W}}}(x_{i}),y_{i})$ in, we have

$$ $\mathbb{E}_{f_{{\mathbf{W}}}(x^{\pm}_{i});{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]} =z_{[2n]}}[g(f_{{\mathbf{W}}}({\mathbf{X}}_{i}),{\mathbf{U}}_{i},z_{[2n]})] \leq\sqrt{2\max(\mathds{1}_{\{y_{i}^{-}=y\}},\mathds{1}_{\{y_{i}^{+}=y\}})I_{z _{[2n]}}(f_{{\mathbf{W}}}(x^{\pm}_{i});{\mathbf{U}}_{i})}.$ (37) $$

Thus taking expectation with respect to ${\mathbf{Z}}_{[2n]}$ yields the desired result

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}} ^{\pm}_{i});{\mathbf{U}}_{i})}\Big{]}.$ (38) $$

∎

###### Proof.

### B.4 Extra bound of class-generalization error using the Loss pair 𝐋 i ± subscript superscript 𝐋 plus-or-minus 𝑖 {\mathbf{L}}^{\pm}_{i} bold_L start_POSTSUPERSCRIPT ± end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT :

###### Theorem 8 .

(class-e-CMI)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$. The class-generalization error of class $y$, as defined in [2](#Thmdefinition2), can be bounded as follows:

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}({\mathbf{L}}^{\pm}_{i};{ \mathbf{U}}_{i})}\Big{]}.$ (39) $$

###### Proof.

Similar to the proof of Theorems  [2](#Thmtheorem2) and [3](#Thmtheorem3).
Using Lemma [3](#Thmlemma3) with ${\mathbf{V}}={\mathbf{L}}^{\pm}_{i}$ and $h({\mathbf{V}},z_{i})={\mathbf{L}}_{i}$ in, we have

$$ $\mathbb{E}_{{\mathbf{L}}^{\pm}_{i};{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n ]}}[g(f_{{\mathbf{W}}}({\mathbf{X}}_{i}),{\mathbf{U}}_{i},z_{[2n]})]\leq\sqrt{ 2\max(\mathds{1}_{\{y_{i}^{-}=y\}},\mathds{1}_{\{y_{i}^{+}=y\}})I_{z_{[2n]}}({ \mathbf{L}}^{\pm}_{i};{\mathbf{U}}_{i})}.$ (40) $$

Thus taking expectation with respect to ${\mathbf{Z}}_{[2n]}$ yields the desired result

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{ \mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}({\mathbf{L}}^{\pm}_{i};{ \mathbf{U}}_{i})}\Big{]}.$ (41) $$

∎

###### Theorem 8 .

###### Proof.

### B.5 Proof of Theorem 4

Theorem [4](#Thmtheorem4) (restated)
($\Delta_{y}L$-CMI)
$\Delta_{y}{\mathbf{L}}_{i}\triangleq\mathds{1}_{\{y_{i}^{-}=y\}}\ell(f_{{
\mathbf{W}}}({\mathbf{X}}_{i})^{-},y^{-}_{i})-\mathds{1}_{\{y_{i}^{+}=y\}}\ell
(f_{{\mathbf{W}}}({\mathbf{X}}_{i})^{+},y^{+}_{i})$. Assume that the loss $\ell(\hat{y},y)\in[0,1]$. The class-generalization error of class $y$, as defined in [2](#Thmdefinition2), can be bounded as follows:

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}} _{i})}\Big{]}.$ (42) $$

###### Proof.

First, we notice that for a fixed realization $z_{[2n]}$, $\mathds{1}_{\{y^{{\mathbf{U}}_{i}}=y\}}\ell({\mathbf{W}},z^{{\mathbf{U}}_{i}}_
{i})-\mathds{1}_{\{y^{-{\mathbf{U}}_{i}}=y\}}\ell({\bm{W}},z^{-{\mathbf{U}}_{i
}}_{i})={\mathbf{U}}_{i}(\mathds{1}_{\{y_{i}^{-}=y\}}\ell({\mathbf{W}},z^{-}_{
i})-\mathds{1}_{\{y_{i}^{+}=y\}}\ell({\mathbf{W}},z^{+}_{i}))={\bm{U}}_{i}
\Delta_{y}{\mathbf{L}}_{i}$.

Next, let $(\overline{\Delta_{y}{\mathbf{L}}}_{i},\overline{{\mathbf{U}}}_{i})$ be an independent copy of $(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})$. Using the Donsker–Varadhan variational representation of KL
divergence, we have $\forall\lambda\in{\mathbb{R}}$ and for every function $g$

$$ $I_{z_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})\geq\lambda\mathbb{E} _{\Delta_{y}{\mathbf{L}}_{i},{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[g( \Delta_{y}{\mathbf{L}}_{i},{\mathbf{U}}_{i},z_{[2n]})]\\ -\log\mathbb{E}_{\overline{\Delta_{y}{\mathbf{L}}}_{i},\overline{{\mathbf{U}}} _{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{\Delta_{y}{\mathbf{L }}}_{i},\overline{{\mathbf{U}}}_{i},z_{[2n]})}].$ (43) $$

Next, let $g(\Delta_{y}{\mathbf{L}}_{i},{\mathbf{U}}_{i},z_{[2n]})={\bm{U}}_{i}\Delta_{y}
{\mathbf{L}}_{i}$. Thus, we have

$$ $\log\mathbb{E}_{\overline{\Delta_{y}{\mathbf{L}}}_{i},\overline{{\mathbf{U}}}_ {i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{\Delta_{y}{\mathbf{L} }}_{i},\overline{{\mathbf{U}}}_{i},z_{[2n]})}]=\log\mathbb{E}_{\overline{ \Delta_{y}{\mathbf{L}}}_{i},\overline{{\mathbf{U}}}_{i}|{\mathbf{Z}}_{[2n]}=z_ {[2n]}}[e^{\lambda\overline{{\mathbf{U}}}_{i}\overline{\Delta_{y}{\mathbf{L}}} _{i}}].$ (44) $$

We have $\mathbb{E}_{\overline{{\mathbf{U}}}_{i}}[\overline{{\mathbf{U}}}_{i}\overline{
\Delta_{y}{\mathbf{L}}}_{i}]=0$ and $\overline{{\mathbf{U}}}_{i}\in\{-1,+1\}$. Thus, using Hoeffding’s Lemma, we have

$$ $\log\mathbb{E}_{\overline{{\mathbf{L}}}^{\pm}_{i},\overline{{\mathbf{U}}}_{i}| {\mathbf{Z}}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{\Delta_{y}{\mathbf{L}}}_{ i},\overline{{\mathbf{U}}}_{i},z_{[2n]})}]\leq\\ \log\mathbb{E}_{\overline{\Delta_{y}{\mathbf{L}}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[ 2n]}}[e^{\frac{\lambda^{2}}{2}\overline{\Delta_{y}{\mathbf{L}}}_{i}^{2}}].$ (45) $$

Next, as $\ell\in[0,1]$, it follows that $\Delta_{y}{\mathbf{L}}_{i}\in[-1,1]$. Thus, $|\overline{\Delta_{y}{\mathbf{L}}}_{i}|\leq 1$.
Thus,

$$ $\log\mathbb{E}_{\overline{\Delta_{y}{\mathbf{L}}}_{i},\overline{{\mathbf{U}}}_ {i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[e^{\lambda g(\overline{\Delta_{y}{\mathbf{L} }}_{i},\overline{{\mathbf{U}}}_{i},z_{[2n]})}]\leq\frac{\lambda^{2}}{2}.$ (46) $$

Replacing in equation [43](#A2.E43), we have

$$ $I_{z_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})\geq\lambda\mathbb{E} _{\Delta_{y}{\mathbf{L}}_{i},{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[{ \mathbf{U}}_{i}\Delta_{y}{\mathbf{L}}_{i}]-\frac{\lambda^{2}}{2}.$ (47) $$

So $\forall\lambda\in{\mathbb{R}}$

$$ $\frac{\lambda^{2}}{2}-\lambda\mathbb{E}_{\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U }}_{i}|{\mathbf{Z}}_{[2n]}=z_{[2n]}}[g(\Delta_{y}{\mathbf{L}}_{i},{\mathbf{U}} _{i},z_{[2n]})]+I_{z_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})\geq 0.$ (48) $$

The equation [48](#A2.E48) is a non-negative parabola with respect to $\lambda$. Thus, its discriminant must be non-positive. This implies

$$ $\mathbb{E}_{\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i}|{\mathbf{Z}}_{[2n]}=z_ {[2n]}}[g(f_{{\mathbf{W}}}({\mathbf{X}}_{i}),{\mathbf{U}}_{i},z_{[2n]})]\leq \sqrt{2I_{z_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})}.$ (49) $$

Thus taking expectation with respect to ${\mathbf{Z}}_{[2n]}$ yields the desired result

$$ $|\overline{\mathrm{gen}_{y}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{y}}\sum_{ i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}} _{i})}\Big{]}.$ (50) $$

Proof that the $\Delta_{y}L$-CMI is always tighter than the class-$f$-CMI bound in Theorem [3](#Thmtheorem3), and the latter is always tighter than the class-CMI bound in Theorem [2](#Thmtheorem2):

Due to the data processing inequality , we have ${\mathbf{U}}\to{\mathbf{W}}\to f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})\to
\Delta_{y}{\mathbf{L}}_{i}$. It then follows directly that the class-$f$-CMI bound is always tighter than the class-CMI bound. Moreover, we can also show that for a the $\Delta_{y}L$-CMI is always tighter than the class-$f$-CMI bound. For a fixed ${\mathbf{Z}}_{[2n]}$, we have 4 different possible cases for each term in the sum:

- •
i) If $y_{i}^{-}\neq y$ and $y_{i}^{+}\neq y$: In this case, $\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+
}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{
\mathbf{U}}_{i})=0I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_
{i});{\mathbf{U}}_{i})=0$. On the other hand, we have $\Delta_{y}{\mathbf{L}}_{i}=0$ so $I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})=I_{{
\mathbf{Z}}_{[2n]}}(0;{\mathbf{U}}_{i})=0\leq 0=\max(\mathds{1}_{\{{\mathbf{Y}
}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+}=y\}})I_{{\mathbf{Z}}_{[2n]}}
(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_{i})$
- •
ii),If $y_{i}^{-}=y$ and $y_{i}^{+}=y$: In this case, $\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+
}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{
\mathbf{U}}_{i})=I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{
i});{\mathbf{U}}_{i})$. Due to the data processing inequality, $I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})\leq I_{{
\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_{i})
=\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{
+}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{
\mathbf{U}}_{i})$
- •
iii) if ,If $y_{i}^{+}\neq y$ and $y_{i}^{-}=y$, In this case, $\max(\mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+
}=y\}})I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{
\mathbf{U}}_{i})=I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{
i});{\mathbf{U}}_{i})$ and $\Delta_{y}{\mathbf{L}}_{i}={\mathbf{L}}_{i}^{+}$. As ${\mathbf{W}}\to f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i})\to{\mathbf{L}}_{i}^{+}$ is a Markov chain, using the data processing inequality, we have $I_{{\mathbf{Z}}_{[2n]}}({\mathbf{L}}_{i}^{+};{\mathbf{U}}_{i})\leq I_{{\mathbf
{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_{i})$ and thus $I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})\leq I_{{
\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_{i})$
- •
iv) If $y_{i}^{+}=y$ and $y_{i}^{-}\neq y$, this corresponds to the same as iii) by changing the + and -.

As for all the different cases, we have

$$ $I_{{\mathbf{Z}}_{[2n]}}(\Delta_{y}{\mathbf{L}}_{i};{\mathbf{U}}_{i})\leq\max( \mathds{1}_{\{{\mathbf{Y}}_{i}^{-}=y\}},\mathds{1}_{\{{\mathbf{Y}}_{i}^{+}=y\} })I_{{\mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U} }_{i})$ (51) $$

$\Delta_{y}L$-CMI is always tighter than the class-$f$-CMI bound.

∎

###### Proof.

## Appendix C Additional Empirical results

### C.1 Experiment Setup

Here, we fully describe the experimental setup used in the main body of the paper.
We use the exact same setup as in Harutyunyan et al. (), where the code is publicly available(^2^22[https://github.com/hrayrhar/f-CMI/tree/master](https://github.com/hrayrhar/f-CMI/tree/master)). For every number of training data $n$, we draw $m_{1}$ of the random variable ${\mathbf{Z}}_{[2n]}$, i.e., we select $m_{1}$ different 2n samples from the original dataset of size $h>2n$.Then, for each $z_{[2n]}$, we draw $m_{2}$ different train/test splits, i.e., $m_{2}$ random realizations of ${\mathbf{U}}$. In total, we have $m_{1}m_{2}$ experiments. We report the mean and standard deviation on the $m_{1}$ results. For the CIFER10 experiments, we select $m_{1}=2$ and $m_{2}=20$. For its noisy variant, we select $m_{1}=5$ and $m_{2}=15$. For both datasets, we use ResNet50 pre-trained on ImageNet. The training is conducted for 40 epochs using SGD with a learning rate of 0.01 and a batch size of 256.

In practice, $P({\mathbf{Y}}=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({\mathbf{Y}}=y)$. Thus, $n^{y}=nP({\mathbf{Y}}=y)$used in Section [2.2](#S2.SS2), can be estimated with $\frac{2}{n_{z_{[2n]}}^{y}}$. Note that our empirical results validate the proposed bounds for this estimated class-wise generalization error.

### C.2 full class-generalization error vs. standard generalization results on CIFAR10

As a supplement to Figure [1](#S1.F1) (left), we plot the standard generalization error along with the class-generalization error of all the classes of CIFAR10 in Figure [3](#A3.F3). Similar to the observations highlighted in Section [1](#S1), we notice the significant variability in generalization performance across different classes.

Figure: Figure 3: The standard generalization error and the generalization error relative for all classes on CIFAR10 as a function of the number of training data.
Refer to caption: extracted/5330490/Figures/Cifar10/cifar10standardvsclasswiseerrorplot.png

### C.3 Numerical results with our bounds for all classes of CIFAR10 and its noisy variant

In Figure [4](#A3.F4), we present the results of the empirical evaluation of our bounds on all the classes of CIFAR10. Moreover, we generate the scatter plot between the class-generalization error and the class-$f$-CMI bound. We note that similar to the class-$\Delta L_{y}$ results in Figure [2](#S3.F2), our bound scales linearly with the error. The results of noisy CIFAR10, presented in Figure [5](#A3.F5), are also consistent with these findings.

Figure: Figure 4: Class-wise generalization on the 10 classes of CIFAR10 and the scatter plot between class-generalization error and the class-$f$-CMI bound in Theorem [3](#Thmtheorem3).
Refer to caption: extracted/5330490/Figures/Cifar10/cifar10classairplanesgenerrorplot.png

Figure: Figure 5: Class-wise generalization on the 10 classes of noisy CIFAR10 (clean validation) and the scatter plot between class-generalization error and the class-$f$-CMI bound in Theorem [3](#Thmtheorem3).
Refer to caption: extracted/5330490/Figures/noisyCifar05/noisycifar10classairplanesgenerrorplot.png

Figure: Figure 6: Class-wise generalization on the 10 classes of noisy CIFAR10 (noise added to both train and validation) and the scatter plot between class-generalization error and the class-$f$-CMI bound in Theorem [3](#Thmtheorem3).
Refer to caption: extracted/5330490/Figures/NoisyCIFAR10both/noisycifar10classairplanesgenerrorplot.png

### C.4 Corollary 1 compared to the bound in Bu et al. ( 2020 )

In the case of standard loss sub-gaussianity assumption, i.e., $\sigma_{y}=\sigma$ is independent of $y$, it is possible to show that the bound in [1](#Thmcorollary1) is tighter than the bound in Bu et al. ().
This is due to the fact that

$$ $\mathbb{E}_{P_{Y}}\sqrt{2\sigma^{2}D(P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}} }||P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}}|{\mathbf{Y}}})}=\mathbb{E}_{P_{Y}} \sqrt{2\sigma^{2}\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}}}\frac{ P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}}P_{{\bm{Y}}}}{P_{{\mathbf{W}}} \otimes P_{{\mathbf{X}}|{\mathbf{Y}}}P_{{\bm{Y}}}}}\\ =\mathbb{E}_{P_{Y}}\sqrt{2\sigma^{2}\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|{ \mathbf{Y}}}}\frac{P_{{\mathbf{W}},{\mathbf{X}},{\mathbf{Y}}}}{P_{{\mathbf{W}} }\otimes P_{{\mathbf{X}},{\mathbf{Y}}}}}\leq\sqrt{2\sigma^{2}I({\mathbf{W}};{ \mathbf{Z}})}$ (52) $$

where the last inequality comes from Jensen’s inequality. This shows that class-wise analysis can be used to derive tighter generalization bounds.

### C.5 Extra results: Expected recall & specificity generalization

In the special case of binary classification with the 0-1 loss, the class-generalization errors studied within this paper correspond to generalization in terms of recall and specificity:
$expected\;recall-empirical\;recall=gen_{p}$, where $p$ is the positive class.
$expected\;specificity-empirical\;specificity=gen_{n}$, where $n$ is the negative class.

Standard Generalization bounds (Xu & Raginsky, ; Harutyunyan et al., ) provide theoretical certificates for learning algorithms regarding classification error/accuracy. However, in several ML applications, e.g., an imbalanced binary dataset, accuracy/error are not considered good performance metrics. For example, consider the binary classification problem of detecting a rare cancer type. While analyzing the model’s generalization error is important, in this case, we might be more interested in understanding generalization in terms of recall, as that is more critical in this case. Standard theoretical bounds (Neyshabur et al., ; Wu et al., ; Wang & Mao, ) do not provide any insights for such metrics.

The developed tools in this paper can be used to close this gap and allow us to understand generalization for recall and specificity theoretically.

We conduct an experiment of binary MNIST (digit 4 vs. digit 9), similar to Harutyunyan et al. (). $m_{1}$ and $m_{2}$ discussed in Section [C.1](#A3.SS1) are selected to be $m_{1}=5$ and $m_{2}=30$. Empirical results for this particular case are presented in Figure [7](#A3.F7). As can be seen in the Figure, our bounds efficiently estimate the expected recall and specificity errors.

Figure: Figure 7: recall generalization error (left) and specificity generalization error (right) for the binary classification 4vs9 from MNIST. The digit 4 is considered the positive class.
Refer to caption: extracted/5330490/Figures/MNIST/mnistclass0generrorplot.png

## Appendix D Extra analysis for Section 4 : other applications

### D.1 Full details of Section 4.1 : FROM CLASS-GENERALIZATION ERROR TO STANDARD GENERALIZATION ERROR

Corollary [1](#Thmcorollary1) (restated)
Assume that for every $y\in{\mathcal{Y}}$, the loss $\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},y)$ is $\sigma_{y}$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{X}}}|\overline{{
\mathbf{Y}}}=y}$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{Y}\sqrt{2\sigma_{{ \mathbf{Y}}}^{2}D(P_{{\mathbf{W}},{\mathbf{X}}_{i}|{\mathbf{Y}}_{i}=y}||P_{{ \mathbf{W}}}\otimes P_{{\mathbf{X}}_{i}|{\mathbf{Y}}_{i}=y})}.$ (53) $$

###### Proof.

The generalization error can be written as

$$ $\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{\mathbf {S}}})=\frac{1}{n}\sum_{i=1}^{n}\big{(}\mathbb{E}_{{\mathbf{W}},\overline{{ \mathbf{Z}}}}[\ell({\mathbf{W}},\overline{{\mathbf{Z}}})]-\mathbb{E}_{{\mathbf {W}},{\mathbf{Z}}_{i}}[\ell({\mathbf{W}},{\mathbf{Z}}_{i})]\big{)}.$ (54) $$

As the loss $\ell$ is $\sigma_{\mathbf{Y}}$ sub-gaussian, using Theorem [1](#Thmtheorem1), we have

$$ $\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}}[\ell({\mathbf{W}},{ \mathbf{X}},{\mathbf{Y}})]-\mathbb{E}_{P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}} |{\mathbf{Y}}=y}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}}, \overline{{\mathbf{Y}}})]\leq\sqrt{2\sigma_{y}^{2}D(P_{{\mathbf{W}},{\mathbf{X }}|{\mathbf{Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}}|{\mathbf{Y}}=y})}.$ (55) $$

Taking the expectation over ${\mathbf{Y}}$ in both sides, we have

$$ $\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{X}},{\mathbf{Y}}}}[\ell({\mathbf{W}},{ \mathbf{X}},{\mathbf{Y}})]-\mathbb{E}_{P_{{\mathbf{W}}}\otimes P_{{\mathbf{X}} ,{\mathbf{Y}}}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{X}}},\overline {{\mathbf{Y}}})]\leq\mathbb{E}_{Y}\sqrt{2\sigma_{{\mathbf{Y}}}^{2}D(P_{{ \mathbf{W}},{\mathbf{X}}|{\mathbf{Y}}=y}||P_{{\mathbf{W}}}\otimes P_{{\mathbf{ X}}|{\mathbf{Y}}=y})}.$ (56) $$

Applying equation [56](#A4.E56) on each term of equation [54](#A4.E54) for each ${\mathbf{Z}}_{i}$ completes the proof.

∎

With similar proofs, we can also extend the results in Theorems [2](#Thmtheorem2),  [3](#Thmtheorem3), and [4](#Thmtheorem4) into standard generalization bounds by taking the expectation over $y\sim P_{{\mathbf{Y}}}$. In Corollaries [4](#Thmcorollary4),  [5](#Thmcorollary5),  [6](#Thmcorollary6), and [2](#Thmcorollary2), we provide such an extension of Theorems [2](#Thmtheorem2),  [3](#Thmtheorem3),  [8](#Thmtheorem8), and [4](#Thmtheorem4), respectively.

Corollary [2](#Thmcorollary2)(restated)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\bm{Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]} }\Big{[}\frac{1}{n^{\mathbf{Y}}}\sum_{i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}( \Delta_{{\mathbf{Y}}}{\mathbf{L}}_{i};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ (57) $$

###### Corollary 4 .

(extra result)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\bm{Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]} }\Big{[}\frac{1}{n^{\mathbf{Y}}}\sum_{i=1}^{n}\sqrt{2\max(\mathds{1}_{{\mathbf {Y}}_{i}^{-}={\mathbf{Y}}},\mathds{1}_{{\mathbf{Y}}_{i}^{+}={\mathbf{Y}}})I_{{ \mathbf{Z}}_{[2n]}}({\mathbf{W}};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ (58) $$

###### Corollary 5 .

(extra result)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\bm{Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]} }\Big{[}\frac{1}{n^{\mathbf{Y}}}\sum_{i=1}^{n}\sqrt{2\max(\mathds{1}_{{\mathbf {Y}}_{i}^{-}={\mathbf{Y}}},\mathds{1}_{{\mathbf{Y}}_{i}^{+}={\mathbf{Y}}})I_{{ \mathbf{Z}}_{[2n]}}(f_{{\mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_{i}) }\Big{]}\bigg{]}.$ (59) $$

###### Corollary 6 .

(extra result)
Assume that the loss $\ell(\hat{y},y)\in[0,1]$, then

$$ $|\overline{\mathrm{gen}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\mathbb{E}_{{\bm{Y}}}\Bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]} }\Big{[}\frac{1}{n^{\mathbf{Y}}}\sum_{i=1}^{n}\sqrt{2\max(\mathds{1}_{{\mathbf {Y}}_{i}^{-}={\mathbf{Y}}},\mathds{1}_{{\mathbf{Y}}_{i}^{+}={\mathbf{Y}}})I_{{ \mathbf{Z}}_{[2n]}}({\mathbf{L}}^{\pm}_{i};{\mathbf{U}}_{i})}\Big{]}\Bigg{]}.$ (60) $$

###### Proof.

###### Corollary 4 .

###### Corollary 5 .

###### Corollary 6 .

### D.2 Full details of the sub-task problem

Consider a supervised learning problem where
the machine learning model $f_{\mathbf{W}}(\cdot)$, parameterized with $w\in{\mathcal{W}}$, is obtained with a training dataset $S$ consisting of $n$ i.i.d samples $z_{i}=(x_{i},\!y_{i})\in\!\mathcal{X}\!\times\!\mathcal{Y}\triangleq\mathcal{Z}$ generated from distribution $P_{{\mathbf{X}}{\mathbf{Y}}}$. The quality of the model with parameter $w$ is evaluated with a loss function $\ell:{\mathcal{W}}\times{\mathcal{Z}}\rightarrow{\mathbb{R}}^{+}$.

For any $w\in{\mathcal{W}}$, the population risk is defined as follows

$$ $L_{P}(w)=\mathbb{E}_{P_{{\mathbf{X}},{\mathbf{Y}}}}[\ell(w,{\mathbf{X}},{ \mathbf{Y}})].$ (61) $$

and the empirical risk is:

$$ $L_{E_{P}}(w,S)=\frac{1}{n}\sum_{i=1}^{n}\ell(w,x_{i},y_{i}).$ (62) $$

Here, we are interested in the subtask problem, which is a special case of distribution shift, i.e., the test performance of the model $w$ is evaluated using a specific subset of classes ${\mathcal{A}}\subset{\mathcal{Y}}$ of the source distribution $P_{{\mathbf{X}}{\mathbf{Y}}}$. Thus, the target distribution $Q_{{\mathbf{X}}{\mathbf{Y}}}$ is defined as
$Q_{{\mathbf{X}}{\mathbf{Y}}}(x,y)=\frac{P_{{\mathbf{X}}{\mathbf{Y}}}(x,y)
\mathds{1}_{\{y\in{\mathcal{A}}\}}}{P_{{\mathbf{Y}}}(y\in{\mathcal{A}})}$.
The population risk on the target domain $Q$ of the subtask problem is

$$ $L_{Q}(w)=\mathbb{E}_{Q_{{\mathbf{X}},{\mathbf{Y}}}}[\ell(w,{\mathbf{X}},{ \mathbf{Y}})].$ (63) $$

A learning algorithm can be modeled as a randomized mapping from the training set $S$ onto a model parameter $w\in\mathcal{W}$ according to the conditional distribution $P_{{\mathbf{W}}|S}$.
The expected generalization error on the subtask problem is the difference between the population risk of $Q$ and the empirical risk evaluated using all samples from $S$:

$$ $\overline{\mathrm{gen}}_{Q,E_{P}}=\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{S}}}}[L _{Q}({\mathbf{W}})-L_{E_{P}}({\mathbf{W}},{\mathbf{S}})],$ (64) $$

where the expectation is taken over the joint distribution $P_{{\mathbf{W}},{\mathbf{S}}}=P_{{\mathbf{W}}|S}\otimes P_{{\mathbf{Z}}}^{n}$.

The generalization error defined above can be decomposed as follows:

$$ $\overline{\mathrm{gen}}_{Q,E_{P}}=\mathbb{E}_{P_{{\mathbf{W}}}}[L_{Q}({\mathbf {W}})-L_{P}({\mathbf{W}})]\\ +\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{S}}}}[L_{P}({\mathbf{W}})-L_{E_{P}}({ \mathbf{W}},{\mathbf{S}})].$ (65) $$

The first term quantifies the gap of the population risks in two different domains, and the second term is the source domain generalization error. Assuming that loss is $\sigma$-subgaussian under $P_{{\mathbf{Z}}}$, it is shown in Wang & Mao () that
the first term can be bounded using the KL divergence between $P$ and $Q$:

$$ $\mathbb{E}_{P_{{\mathbf{W}}}}[L_{Q}(w)-L_{P}(w)]\leq\sqrt{2\sigma^{2}D(Q\|P)}.$ (66) $$

The second term can be bounded using the standard mutual information approach in Xu & Raginsky () as

$$ $\mathbb{E}_{P_{{\mathbf{W}},{\mathbf{S}}}}[L_{P}({\mathbf{W}})-L_{E_{P}}({ \mathbf{W}},{\mathbf{S}})]\leq\sqrt{2\sigma^{2}I({\mathbf{W}};{\mathbf{S}})}.$ (67) $$

Thus, the generalization error of the subtask problem can be bounded as follows:

$$ $\overline{\mathrm{gen}}_{Q,E_{P}}\leq\sqrt{2\sigma^{2}D(Q\|P)}+\sqrt{2\sigma^{ 2}I({\mathbf{W}};{\mathbf{S}})}.$ (68) $$

Obtaining tighter generalization error bounds for the subtask problem is straightforward using our class-wise generalization bounds. In fact, the generalization error bound of the subtask can be obtained by taking the expectation of ${\mathbf{Y}}\sim Q_{{{\mathbf{Y}}}}$.

Using Jensen’s inequality, we have $|\overline{\mathrm{gen}}_{Q,E_{Q}}|=|\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf{Y
}}}\big{[}\overline{\mathrm{gen}}_{\mathbf{Y}}\big{]}|\leq\mathbb{E}_{{\mathbf
{Y}}\sim Q_{\mathbf{Y}}}\big{[}|\overline{\mathrm{gen}}_{\mathbf{Y}}|\big{]}$. Thus, we can use the results from Section [2](#S2) to obtain tighter bounds.

Theorem [5](#Thmtheorem5) (subtask-CMI) (restated)
Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, then the subtask generalization error defined in [11](#S4.E11) can be bounded as

$$ $|\overline{\mathrm{gen}}_{Q,E_{Q}}|\leq\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf {Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{\mathbf{Y}}} \sum_{i=1}^{n}\sqrt{2\max(\mathds{1}_{{\mathbf{Y}}_{i}^{-}={\mathbf{Y}}}, \mathds{1}_{{\mathbf{Y}}_{i}^{+}={\mathbf{Y}}})I_{{\mathbf{Z}}_{[2n]}}({ \mathbf{W}};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ $$

Similarly, we can extend the result of class-$\Delta L_{y}$-CMI, with similar proof, as follows:

Theorem [6](#Thmtheorem6) (subtask-$\Delta L_{y}$-CMI) (restated)
Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, Then the subtask generalization error defined in [11](#S4.E11) can be bounded as

$$ $|\overline{\mathrm{gen}}_{Q,E_{Q}}|\leq\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf {Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{\mathbf{Y}}} \sum_{i=1}^{n}\sqrt{2I_{{\mathbf{Z}}_{[2n]}}(\Delta_{{\mathbf{Y}}}{\mathbf{L}} _{i};{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ $$

Similarly, we can also extend the result of Theorem [8](#Thmtheorem8) to the subtask as follows:

###### Theorem 9 .

(subtask-e-CMI) (extra result)
Assume that the loss $\ell(w,x,y)\in[0,1]$ is bounded, then the subtask generalization error defined in [11](#S4.E11) can be bounded as

$$ $|\overline{\mathrm{gen}}_{Q,E_{Q}}|\leq\mathbb{E}_{{\mathbf{Y}}\sim Q_{\mathbf {Y}}}\bigg{[}\mathbb{E}_{{\mathbf{Z}}_{[2n]}}\Big{[}\frac{1}{n^{\mathbf{Y}}} \sum_{i=1}^{n}\sqrt{2\max(\mathds{1}_{{\mathbf{Y}}_{i}^{-}={\mathbf{Y}}}, \mathds{1}_{{\mathbf{Y}}_{i}^{+}={\mathbf{Y}}})I_{{\mathbf{Z}}_{[2n]}}(f_{{ \mathbf{W}}}({\mathbf{X}}^{\pm}_{i});{\mathbf{U}}_{i})}\Big{]}\bigg{]}.$ $$

###### Theorem 9 .

### D.3 Extra details of Section 4.3 : Generalization certificates with sensitive attributes

Theorem [7](#Thmtheorem7) (restated) Given $t\in{\mathcal{T}}$, assume that the loss $\ell({\mathbf{W}},{\mathbf{Z}})$ is $\sigma$ sub-gaussian under $P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{Z}}}}$, then the attribute-generalization error of the sub-population ${\mathbf{T}}=t$, as defined in [3](#Thmdefinition3), can be bounded as follows:

$$ $|\overline{\mathrm{gen}_{t}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})|\leq\sqrt{2\sigma^{2}D(P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{ \mathbf{Z}}|{\mathbf{T}}=t}||P_{{\mathbf{W}}}\otimes P_{{\mathbf{Z}}|\overline {{\mathbf{T}}}=t})}.$ (69) $$

###### Proof.

We have

$$ $\overline{\mathrm{gen}_{t}}(P_{{\mathbf{X}},{\mathbf{Y}}},P_{{\mathbf{W}}|{ \mathbf{S}}})=\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{ \mathbf{Z}}}|{\mathbf{T}}=t}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{ Z}}})]-\mathbb{E}_{P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{ \mathbf{T}}=t}}[\ell({\mathbf{W}},{\mathbf{Z}})].$ (70) $$

Using the Donsker–Varadhan variational representation of the relative entropy, we have

$$ $D(P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}||P_{{ \mathbf{W}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t})\geq\mathbb{E}_{P_{{ \mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}}[\lambda\ell( {\mathbf{W}},{\mathbf{Z}})]\\ -\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{Z}}} |\overline{{\mathbf{T}}}=t}}[e^{\lambda\ell(\overline{{\mathbf{W}}},\overline{ {\mathbf{Z}}})}],\forall\lambda\in{\mathbb{R}}.$ (71) $$

On the other hand, we have:

$$ $\displaystyle\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{ \mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}\Big{[}e^{\lambda\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{Z}}})-\lambda\mathbb{E}[\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{Z}}})]}\Big{]}$ $\displaystyle=\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{ {\mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}\Big{[}e^{\lambda\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{Z}}})}e^{-\lambda\mathbb{E}[\ell(\overline{{ \mathbf{W}}},\overline{{\mathbf{Z}}})]\big{)}}\Big{]}$ $\displaystyle=\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{ {\mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}[e^{\lambda\ell(\overline{{\mathbf{W} }},\overline{{\mathbf{Z}}})}]-\lambda\mathbb{E}_{P_{\overline{{\mathbf{W}}}} \otimes P_{\overline{{\mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}[\ell(\overline{ {\mathbf{W}}},\overline{{\mathbf{Z}}})].$ $$

Using the sub-gaussian assumption, we have

$$ $\log\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{Z}}}| \overline{{\mathbf{T}}}=t}}[e^{\lambda\ell(\overline{{\mathbf{W}}},\overline{{ \mathbf{Z}}})}]\leq\lambda\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{ \overline{{\mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}(\ell(\overline{{\mathbf{W} }},\overline{{\mathbf{Z}}}))+\frac{\lambda^{2}\sigma^{2}}{2}.$ (72) $$

By replacing in equation [71](#A4.E71), we have

$$ $D(P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}||P_{{ \mathbf{W}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t})\geq\lambda\big{(}\mathbb{ E}_{P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}}[\ell ({\mathbf{W}},{\mathbf{Z}})]-\\ \mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{\overline{{\mathbf{Z}}}| \overline{{\mathbf{T}}}=t}}[\ell(\overline{{\mathbf{W}}},\overline{{\mathbf{Z} }})]\big{)}-\frac{\lambda^{2}\sigma}{2}.$ (73) $$

Thus, we have:

$$ $D(P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}||P_{{ \mathbf{W}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t})-\lambda(\mathbb{E}_{P_{{ \mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}}[\ell({ \mathbf{W}},{\mathbf{Z}})]-\mathbb{E}_{P_{\overline{{\mathbf{W}}}}\otimes P_{ \overline{{\mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}[\ell(\overline{{\mathbf{W} }},\overline{{\mathbf{Z}}})])\\ +\lambda^{2}\sigma^{2}\geq 0,\forall\lambda\in{\mathbb{R}}.$ (74) $$

equation [74](#A4.E74) is a non-negative parabola with respect to $\lambda$. Thus, its discriminant must be non-positive. This implies

$$ $|\mathbb{E}_{P_{{\mathbf{W}}|{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}} =t}}[\ell({\mathbf{W}},{\mathbf{Z}})]-\mathbb{E}_{P_{\overline{{\mathbf{W}}}} \otimes P_{\overline{{\mathbf{Z}}}|\overline{{\mathbf{T}}}=t}}[\ell(\overline{ {\mathbf{W}}},\overline{{\mathbf{Z}}})]|\leq\sqrt{2\sigma^{2}D(P_{{\mathbf{W}} |{\mathbf{Z}}}\otimes P_{{\mathbf{Z}}|{\mathbf{T}}=t}||P_{{\mathbf{W}}}\otimes P _{{\mathbf{Z}}|{\mathbf{T}}=t})}.$ (75) $$

This completes the proof.
∎

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: Major questions/comments:

1. Definition 1 looks like an adaptation of Eq 3 in Bu et al., 2020. In particular, the difference of two terms given by the population loss and the empirical loss remains. The difference is that Def 1 in the present work considers conditioning on a specific class y. Beyond this generalization to per-class analysis, **could the authors please highlight contributions compared to Bu et al., 2020 (and also Xu & Raginsky, 2017, Zhou et al., 2020)?** I am asking as a non-specialist reader.

2. Theorem 1 in the present work is similar to Theorem 1 of Xu & Raginsky, 2017, except that Xu & Raginsky, 2017 considers no conditioning on a specific class y. Could the authors comment on the difference? 

3. The idea of “supersample” used in Def 3 seems to be inspired by Steinke & Zakynthinou 2020 (see Sec 1.2 therein). Could the authors comment on the new parts in the present work wrt to the original supersample idea?


4. The notation $P(\mathbf{W}|\mathbf{S})$ is used in the paper to denote the distribution of learned weights $\mathbf{W}$ given a training sample $S$. In Def 1, the defined class-wise generalization error involves $P_{\mathbf{W}, \mathbf{S}_y}$. This presumably means  $P(\mathbf{W} | \mathbf{S}_y)  P(\mathbf{S}_y)$. What does $P(\mathbf{W} | \mathbf{S}_y)  $ mean? How can a classifier learn from only data from one class $y$?

5. L210. How does Def 2 depend on $P(y)$? Likewise how does Def 3 *not* depend on $P(y)$?

6. L272: “if model parameters W memorize the random selection U, the CMI and the class-generalization error will be large”. I thought that U is a vector of Rademacher random variables introduced solely for analysis. Do you actually empirically sample U and use it to train W?


Minor:
* L1454 in the appendix. Log is missing when expanding the KL divergence.

* Minor suggestion. I think CIFAR10 can be considered a solved problem (much like MNIST). I am aware that the paper has experimental results on CIFAR100 in the appendix. Test accuracy on CIFAR 100 of state-of-the-art models can be almost 100% these days. On the empirical side, it would be good to consider one more experimental setting where the model lacks capacity to tackle the problem. CIFAR 10 is too easy for a model class like ResNet 50. For instance, considering a ResNet 8 on ImageNet would be a good setting where the model struggles to learn well due to the lack of capacity. It would be good to check whether the proposed bounds hold in this setting as well. 
(Note that this is not a request for more experiments for the purposes of evaluating the paper. It is a suggestion to help strengthen this work. The authors need not try to conduct this experiment during the rebuttal. It’s not a major request.)
soundness: 3
strengths: This paper is theoretical in nature. The idea of pursuing generalization errors for each class is claimed to be new. If so (to be verified with other reviewers), then this research direction is original, The proposed information theoretical bounds appear to generalize the bounds studied in  Bu et al., 2020, Xu & Raginsky, 2017, Zhou et al., 2020 and other works, which mostly study the standard generalization bounds (i.e., averaged over all classes, and are not class-wise).

Empirically it is verified that the results in Theorem 3 and Theorem 4 provide valid class-wise generalization bounds in both standard CIFAR 10 and a noisy variant of CIFAR 10. This demonstrates that the contributes bounds are widely applicable to several settings.
summary: This paper proposes characterizing classification generalization errors of each class separately. The motivation is that neural networks do not generalize equally for all classes. The paper demonstrates this phenomenon empirically on CIFAR10 and CIFAR100, and also notes that class-wise generalization depends on factors beyond the class itself. Main contributions are several variants of per-class information-theoretic generalization bounds. The paper notes that the mutual information between the model and the class data can be used to characterize (or upper bound) the class generalization error.
weaknesses: (Recommended actions are clearly listed in the Questions section.)


*Presentation:*

Writing can be much improved. In particular, many theoretical results are not accompanied with a justification, or a practical implication. Writing appears to be tailored to specialists of this subject area. Sections and writing flow in the paper can also be improved. For instance, Theorem 1 is introduced only to be superseded by Sec 2.2 which, at the beginning of the section, states a limitation of Theorem 1. Def 2 (a definition of class-wise generalization errors) is introduced only to be superseded by Def 3 (see reason at L210). It is hard for me to form a coherent story. Same for Theorem 2 (class-CMI) which is introduced only to be commented on before introducing Theorem 3 that the bound in Theorem 2 is hard to evaluate (see L275). 

*Novelty:*

The contributed results, while generalizing bounds in existing results, are also heavily based on them. This raises a concern about the novelty. For instance, Def 1 looks like an adaptation of Eq 3 in Bu et al., 2020.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1.Motivation: Even though the authors reveal label-dependency of generalization bound, in addition to mentioned in the Sec 4.1, can you discuss how different labels affect the standard generalization error? Or can you explain the scenarios where class generalization must be considered?
2.Essentiality: The derivation of Theorem 1 and Lemma 2 starts from KL divergence (or CMI) and does not exceed this scope. Do these conclusions reveal any more fundamental facts?
3.Proof Details: (i) Regarding Equation 30, is there an inversion of the positive and negative signs? (ii) In response to the issue mentioned in Remark 2, the authors introduce a novel loss random variable. However, the bound presented in Equation 47 is also applicable to the expression preceding Equation 31. Consequently, in Theorem 3, the term$\max\left(\mathbb{1}_{\left\{y_{i}^{-}=y\right\}},\mathbb{1}_{\left\{y_{i}^{+}=y\right\}}\right) $can be omitted, similar to the approach taken in Theorem 4. Do these facts have any impact on your findings?
soundness: 3
strengths: Originality: The authors introduce a new concept - "class-generalization error", to address the issue that the existing generalization boundaries cannot capture the varying generalization performance across different classes.
Quality: The authors have provided detailed proofs for each conclusion and conducted comprehensive numerical experiments.
Significance: The authors not only extend their theoretical findings to derive bounds for diverse scenarios but also shed light on label-dependent generalization bounds.
summary: In order to address the issue that the existing generalization boundaries cannot capture the varying generalization performance across different classes, the authors introduce the concept of "class-generalization error," and derive bounds based on KL-divergence and super-sample technique, validating their bounds in different networks.

Their theoretical tools extend to: (i) Deriving class-dependent generalization error bounds affecting standard generalization and tightening existing bounds; (ii) Providing tight bounds for subtask problems with test data subsets; (iii) Deriving bounds for learning with sensitive attributes.
weaknesses: The notation in the formulas is not concise enough. For instance, in Eq 23., the parameters of $\ell$ are in the order of w, x, y, whereas in Eq 54, they have become z, w. In addition, some formulas are not aesthetically pleasing, such as [] in Eq11.

Some questions about the proof are discussed in Question part.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: — Could you address the comment above regarding Appendix C.2? I acknowledge that I may have missed something.

— How does the approach with class-specific gradient noise in Appendix D.7 compare with an approach that adds gradient noise for every class?

— There are some results in the PAC-Bayesian literature on generalization bounds for the confusion matrix (e.g. Morvant et al, “PAC-Bayesian Generalization Bound on Confusion Matrix for Multi-Class Classification”). How does your work relate to this literature?

— Lemma 2: Does $V$ depend only on $W$ here or on $W$ and $z_{[2n]}$?
soundness: 3
strengths: The paper begins by motivating the problem at hand in a very nice way, illustrating that classwise error discrepancies occur and do not behave in ways that are immediately obvious. The presented results are evaluated thoroughly, and interesting extensions are discussed. While the core of the derivations rely on standard techniques, there are some technical intricacies to be dealt with, and the present application is new.
summary: In this paper, the generalization performance of learning algorithms are studied on a classwise level—i.e., bounds are provided between the population loss for samples on a class and the training loss for samples from that class. These bounds are obtained using information-theoretic tools, i.e. the KL divergence and conditional mutual information. Empirical evaluations are performed to verify that classwise discrepancies exist in real applications, and that the proposed bounds correlate with them. Several extensions and connections to related problems (subtask, sensitive attributes) are presented.
weaknesses: Some more discussion on the topic of imbalanced data sets and how this affects classwise generalization error would be of interest. For instance, explicitly stating that the classes in CIFAR10 are balanced could be useful for the reader.

The discussion in Appendix C.2 does not seem to be accurate. Unless I’m mistaken, the following appears to be perfectly valid:

$$ E_{\hat Z_{[2n]}, U, W}[ \frac{1}{n} \sum_i \ell(\hat Z_i^{-U_i}, W)  - \ell(\hat Z_i^{U_i}, W)] = \frac{1}{n} \sum_i  \big( E_{\hat Z_{i}, U, W}[ \ell(\hat Z_i^{-U_i}, W)]  - E_{\hat Z_{i}, U, W}[ \ell(\hat Z_i^{U_i}, W)]  \big) $$
$$ = \big( E_{\hat Z_{1}, U, W}[ \ell(\hat Z_1^{-U_1}, W)]  - E_{\hat Z_{1}, U, W}[ \ell(\hat Z_1^{U_1}, W)]  \big) =  \big( E_{\bar Z, \bar W}[ \ell(\bar Z, \bar W)]  - E_{Z, W}[ \ell(Z, W)]  \big) $$

which is the generalization gap. First step was linearity of expectation, second used the symmetry of the algorithm + iid (despite the altered process, the sample pairs are iid), and the last one used that $WW$ and $hat Z_1^{-U_1}$ are independent. The fact that the generalization gap is linear in the test/train losses mean that the $y$-specificity of the pairs is marginalized out.

Now, if one were to consider bounds where the left-hand-side was given in terms of a function of train/test loss (e.g., the binary KL divergence as in the Maurer-Langford-Seeger bound), this would no longer hold.

Minor:

Should $n_y <n $ on line 136 be $\leq$?

“$n$ super-samples” on line 193: in the original paper from Steinke and Zakynthinou, the term “supersample” was used to describe the full set of $2n$ samples.

Corollary 2 and Theorem 5 (e.g.): Would be nice to adjust the size of brackets.

Line 1070: “requieres”

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: - In the example of "truck", the authors show that their bound captures the behavior of the noisy case. Is it possible to explain the increase in the error w.r.t samples using the proposed bound?
- In Remark 3, the authors mention that their bound is discrepancy-independent. However, the reason for this is not their bound, but the change of the definition from $L_E(W, S)$ to $L_{E_Q}(W, S)$. I would suggest the authors rephrase this part and perhaps break it down into two factors, one due to the change in the definition (still dependent on the discrepancy), and one the bound (discrepancy-independent) to avoid being misleading.
- In [1], the authors provide a general framework for information-theoretic bounds which saves a lot of proof as long as the authors can show some regularity of the distributions. Since the authors' technique of proof follows the same techniques in the general generalization bounds, have authors considered using this framework?


[1] Chained Generalisation Bounds, Clerico et al. COLT 2022
soundness: 3
strengths: The paper is well-written, and easy to follow. The authors have motivated their work nicely both at the beginning of the paper and the end by providing specific examples of some scenarios in which the class-wise generalization bound can be insightful like the sub-task problem as a specific case of distribution shift. The idea of caring about each class independently is new and in some cases tightens the pre-existing bounds.
summary: The paper studies the class-generalization error bound as opposed to the more traditional case of expected generalization over the whole data distribution. The authors provide a new class-specific generalization definition and give multiple information-theoretic bounds based on the KL divergence (more classic), and the super-sample technique (maybe somewhat newer). They analyze the tightness of their bounds in some experiments on CIFAR10 and CIFAR100 showing the failure of the general generalization bounds and the success of their class-specific approach. Finally, the authors provide some examples of the applications in which the provided bounds can be useful and add new insights.
weaknesses: The technical contribution of the paper is limited. While the definition of class-wise generalization bound is new, the results are heavily based on the previous works [Xu and Raginksy 2017, Zhou et al 2022, Harutyunyan 2021, Steinke and Zakynthinou 2020, Clerico et al 2022, Wang and Mao 2023], and follow the same flow of proof, without any particular novelty. It would be useful if the authors mention the technical challenges their definition might propose to the problem and how it differs from the challenges one needs to overcome while taking the general generalization bounds.
The insight of some bounds being tighter than others again is well-studied in previous works. However, I agree with the authors that the tightness of class-wise generalization as opposed to general generalization bound in Corollary 1 is new.

Meta-review: This submission is a borderline case. It is a pure learning theory paper. It proposed to analyze the class-wise generalization error from an information-theoretic point of view. After the rebuttal, all four reviewers gave a rating of 6 that is "marginally above the acceptance threshold". During the internal discussions, none of four positive reviewers showed up, even after I cued them, suggesting that they were not truly supportive. In particular, I think the review comments and the review ratings do not match, for example,
> The technical contribution of the paper is limited ... the results are heavily based on the previous works ... follow the same flow of proof, without any particular novelty ... The insight of some bounds being tighter than others again is well-studied in previous works.

Based on my own experience as an author and AC in these years, this comment looks more like in a review with a rating of 3 rather than 6. Note that two reviewers questioned the novelty and two reviewers questioned the clarity or presentation. I quickly checked it and found it not easy to follow or to appreciate the merits (except the motivation which was quite nice and easy to follow). Given the current version of this work, even if we accept it, it will not make reasonable impact that it should have due to its poor presentation. Therefore, I don't think we should accept it for publication before the authors significantly improve the current version.

Decision: Reject