SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For any sample  $z = (x,y)$  with label  $y = k$  and for any iteration  $t \geqslant 1$ , let  $r \in [M]$  denote the index of the category-dependent prior (center) which is the closest to  $\mathcal{N}(\mu_x, \mathrm{diag}(\sigma_x^2))$  in terms of the KL-divergence
$$
r _ {z} ^ {(t)} = \underset {r \in [ M ]} {\arg \min } D _ {K L} \left[ \mathcal {N} \left(\mu_ {x}, \operatorname {d i a g} \left(\sigma_ {x} ^ {2}\right)\right) \| \mathcal {N} \left(\bar {\mu} _ {k, r} ^ {(t)}, \bar {\sigma} _ {k, r} ^ {(t)} ^ {2}\right) \right].
$$
Suppose that the picked mini-batch at iteration  $t$  is  $\mathcal{B}_t = \{z_1,\dots ,z_b\}$ . For each  $k\in [K]$  and  $r\in [M]$ , let
$$
\mathcal {I} _ {k, r} ^ {(t)} = \left\{z _ {i} \colon i \in [ b ], y _ {i} = k, r _ {z _ {i}} ^ {(t)} = r \right\},
$$
and let  $b_{k,r} = |\mathcal{I}_{k,r}|$ . Now, if  $b_{k,r} \neq 0$ , update  $\bar{\mu}_{k,r}^{(t)}$  and  $\bar{\sigma}_{k,r}^{(t)}$  as
$$
\bar {\mu} _ {k, r} ^ {(t)} := (1 - \alpha b _ {k, r}) \bar {\mu} _ {k, r} ^ {(t - 1)} + \alpha \sum_ {z _ {i} \in \mathcal {T} _ {k, r}} \mu_ {x _ {i}},
$$
$$
\bar {\sigma} _ {k, r, j} ^ {(t)} := \sqrt {(1 - \alpha b _ {k , r}) \bar {\sigma} _ {k , r , j} ^ {(t - 1) ^ {2}} + \alpha \sum_ {z _ {i} \in \mathcal {I} _ {k , r}} \sigma_ {x _ {i} , j} ^ {2}}, \quad j = 1, \ldots , m,
$$
where  $\alpha \in [0,1]$  is a coefficient that smoothens the evolution of the mean and variance, and also implicitly takes into account the effect of the ghost dataset  $S^{\prime}$  appearing in the bounds of Theorems 4 and 7. The optimal value of  $\alpha$ , among others, depends on the mini-batch size  $b$  and the number of centers  $M$ .<sup>8</sup>
Finally, the considered objective function at iteration  $t$  in the Lossless CDVIB approach is
$$
\begin{array}{l} \frac {1}{b} \sum_ {i = 1} ^ {b} \left\{ \right.\beta D _ {K L} \Bigg [ \mathcal {N} \Big (\mu_ {x _ {i}}, \mathrm {d i a g} (\sigma_ {x _ {i}} ^ {2}) \Big) \| \mathcal {N} \Big (\bar {\mu} _ {y _ {i}, r _ {z _ {i}} ^ {(t)}}, \mathrm {d i a g} (\bar {\sigma} _ {y _ {i}, r _ {z _ {i}} ^ {(t)}} ^ {(2)}) \Big) \Bigg ] \\ \left. - \mathbb {E} _ {U _ {i} \sim P _ {U | X, W _ {e}} \left(U _ {i} \mid x _ {i}, W _ {e}\right)} \left[ \log P _ {\hat {Y} | U, W _ {d}} \left(y _ {i} \mid U _ {i}, W _ {d}\right) \right] \right\}. \tag {14} \\ \end{array}
$$
D.3 Lossy CDVIB objective function
Inspired by the lossy compression and bounds introduced in our work, we consider the MDL of the "perturbed" latent variable, while passing the un-perturbed latent variable to the decoder. More precisely, as before, we consider the log loss for evaluation of the relevance of  $U$  in the decoder, i.e.,
$$
\mathbb {E} _ {U _ {i} \sim P _ {U | X, W _ {e}} (U _ {i} | x _ {i}, W _ {e})} \left[ \log P _ {\hat {Y} | U, W _ {d}} \left(y _ {i} | U _ {i}, W _ {d}\right) \right].
$$
For the regularizer, we first consider the perturbed  $U$  as
$$
\hat {U} = U + Z _ {2} = \mu_ {X} + Z _ {2} + \sigma_ {X} Z _ {1} = \hat {U} _ {1} + \hat {U} _ {2}, \tag {15}
$$
where  $Z_{1}$  and  $Z_{2}$  are independently drawn from the same distribution  $\mathcal{N}(\mathbf{0}_m,\mathrm{I}_m)$ . Note that we chose
$$
\hat {U} _ {1} := \mu_ {X} + Z _ {2}, \quad \hat {U} _ {2} := \sigma_ {X} Z _ {1}. \tag {16}
$$
Hence, given  $(X,W_{e})$ ,  $\hat{U}_{1}\sim \mathcal{N}(\mu_{X},\mathrm{I}_{m})$  is independent from  $\hat{U}_2\sim \mathcal{N}(\mathbf{0}_m,\mathrm{diag}(\sigma_X^2))$ . Let us define two sets of priors  $\mathcal{Q}_1\coloneqq \{Q_{1,k,r}\}_{k\in [K],r\in [M]}$  over  $\hat{U}_1$  and  $\mathcal{Q}_2\coloneqq \{Q_{2,k,r}\}_{k\in [K],r\in [M]}$  over  $\hat{U}_2$ . Next, for each  $i\in [b]$ , we select two priors  $Q_{1,i}\in \mathcal{Q}_1$  and  $Q_{2,i}\in \mathcal{Q}_2$ , in a manner that will become clear in the following. Denote the induced prior for  $\hat{U} = \hat{U}_1 + \hat{U}_2$ , where  $\hat{U}_1\sim Q_{1,i}$  and  $\hat{U}_2\sim Q_{2,i}$  by  $Q_{i}$ . Then, the KL divergence of  $P_{\hat{U}|X,W_e}(\hat{U}|X,W_e)$  and  $Q_{i}$ , can be upper bounded by
$$
\begin{array}{l} D _ {K L} \Bigg [ P _ {\hat {U} | X, W _ {e}} (\hat {U} | X, W _ {e}) \| Q _ {i} \Bigg ] \\ \leqslant D _ {K L} \left[ P _ {\hat {U} _ {1} | X, W _ {e}} (\hat {U} _ {1} | X, W _ {e}) \| Q _ {1, i} \right] + D _ {K L} \left[ P _ {\hat {U} _ {2} | X, W _ {e}} (\hat {U} _ {2} | X, W _ {e}) \| Q _ {2, i} \right] \\ = D _ {K L} \left[ \mathcal {N} \left(\mu_ {X}, \mathrm {I} _ {m}\right) \| Q _ {1, i} \right] + D _ {K L} \left[ \mathcal {N} \left(\mathbf {0} _ {m}, \operatorname {d i a g} \left(\sigma_ {X} ^ {2}\right)\right) \| Q _ {2, i} \right]. \tag {17} \\ \end{array}
$$
It can be shown that this choice of prior for  $M = 1$  satisfies the type-III symmetry property and for  $M > 1$  is an approximation of a prior that satisfies such a symmetry.
We use this upper bound for our regularizer. To make things more formal, for each label  $k \in [K]$ , we consider  $2M$  priors,  $M \in \mathbb{N}$ ,
$$
Q _ {1, k, r} ^ {(t)} = \mathcal {N} \Big (\bar {\mu} _ {k, r} ^ {(t)}, \mathrm {I} _ {m} \Big), \quad r \in [ M ], k \in [ K ],
$$
$$
Q _ {2, k, r} ^ {(t)} = \mathcal {N} \left(\mathbf {0} _ {m}, \operatorname {d i a g} \left(\bar {\sigma} _ {k, r} ^ {(t)} ^ {2}\right)\right), \quad r \in [ M ], k \in [ K ],
$$
Similar to lossless CDVIB, firstly, the vectors  $\bar{\mu}_{k,r}^{(t)}$  and  $\bar{\sigma}_{k,r}^{(t)}^2$  are initialized respectively as  $m$ -dimensional vectors of zeros and ones. Next, these initialized vectors are updated at each iteration using a procedure defined below.
For any sample  $z = (x,y)$  with label  $y = k$  and for any iteration  $t \geqslant 1$ , let  $r \in [M]$  denote the index of the category-dependent prior (center) which has the smallest distance to  $(Q_{1,k,r}^{(t)}, Q_{2,k,r}^{(t)})$  in a sense that minimizes the RHS of (17). In other words,  $r_z^{(t)}$  is equal to
$$
\underset {r \in [ M ]} {\arg \min } \Bigg \{D _ {K L} \bigg [ \mathcal {N} (\mu_ {x}, \mathrm {I} _ {m}) \| \mathcal {N} \Big (\bar {\mu} _ {k, r} ^ {(t)}, \mathrm {I} _ {m} \Big) \bigg ] + D _ {K L} \bigg [ \mathcal {N} (\mathbf {0} _ {m}, \operatorname {d i a g} (\sigma_ {x} ^ {2})) \| \mathcal {N} \Big (\mathbf {0} _ {m}, \operatorname {d i a g} \left(\bar {\sigma} _ {k, r} ^ {(t)} ^ {2}\right) \Big) \bigg ] \Bigg \}.
$$
Then, the mean and variances of the priors are updated exactly similarly to Lossless CDVIB. For completeness, we repeat this procedure here. Suppose that the picked mini-batch at iteration  $t$  is  $\mathcal{B}_t = \{z_1,\dots ,z_b\}$ . For each  $k\in [K]$  and  $r\in [M]$ , let
$$
\bar {\mu} _ {k, r} ^ {(t)} := (1 - \alpha b _ {k, r}) \bar {\mu} _ {k, r} ^ {(t - 1)} + \alpha \sum_ {z _ {i} \in \mathcal {I} _ {k, r}} \mu_ {x _ {i}},
$$
where  $\alpha \in [0,1]$  is a coefficient that smoothens the evolution of the mean and variance, and also implicitly takes into account the effect of the ghost dataset  $S^{\prime}$  appearing in the bounds of Theorems 4 and 7. The optimal value of  $\alpha$ , among others, depends on the mini-batch size  $b$  and the number of centers  $M$ .
Finally, the considered objective function at iteration  $t$  in the Lossy CDVIB approach is
$$
\begin{array}{l} \frac {1}{b} \sum_ {i = 1} ^ {b} \left\{\beta D _ {K L} \left[ \mathcal {N} (\mu_ {x _ {i}}, \mathrm {I} _ {m}) \| \mathcal {N} \left(\bar {\mu} _ {y _ {i}, r _ {z _ {i}} ^ {(t)}}, \mathrm {I} _ {m}\right) \right] + \beta D _ {K L} \left[ \mathcal {N} (\mathbf {0} _ {m}, \mathrm {d i a g} (\sigma_ {x _ {i}} ^ {2})) \| \mathcal {N} \left(\mathbf {0} _ {m}, \mathrm {d i a g} \left(\bar {\sigma} _ {y _ {i}, r _ {z _ {i}} ^ {(t)}} ^ {(t)}\right) ^ {2}\right) \right] \right. \\ \left. - \mathbb {E} _ {U _ {i} \sim P _ {U | X, W _ {e}} \left(U _ {i} \mid x _ {i}, W _ {e}\right)} \left[ \log P _ {\hat {Y} | U, W _ {d}} \left(y _ {i} \mid U _ {i}, W _ {d}\right) \right] \right\}. \tag {18} \\ \end{array}
$$
D.4 Datasets
In the experiments, we used CIFAR10  $\mathrm{[KH^{+}09]}$ . The full dataset was split into a training set with 50,000 labeled images and a validation set with 10,000 labeled images, all of them of size  $32\times 32\times 3$ . The input images were scaled to have most values between 0 and 1 before being fed to the network.
D.5 Architecture details
The model architecture considered in our experiments is detailed in Table 1. The encoder part of our two-step prediction model is a convolutional network consisting of four convolutional layers followed by two linear layers. We use max-pooling and a LeakyReLU activation function with a negative slope coefficient equal to 0.1. The encoder takes as input re-scaled images and produces parameters  $\mu_x$  and variance  $\sigma_x^2$  of the latent variable of dimension  $m = 64$ . The latent samples are generated using the reparameterization trick of [KW14]. Next, the produced latent samples are processed by a decoder consisting of one linear layer with a softmax activation function. The decoder outputs a soft class prediction.
Similarly, as in [DKSV20], our evaluated encoder is complex enough in order to make it close to "a universal function approximator". On the other hand, we use a simple decoder as in [AFDM17] in order to reduce spurious regularization introduced by the high decoder's complexity and hence to highlight the benefits of our regularizer in terms of generalization performance.
Table 1: The model architecture used in experiments. The convolutional layers are parametrized respectively by the number of input channels, the number of output channels, and the filter size. The linear layers are defined by their input and output sizes.
Encoder	Encoder cont'd	Encoder cont'd
Number	Layer	Number	Layer	Number	Layer
1	Conv2D(3,8,5)	6	Conv2D(16,16,3)	11	LeakyReLU(0.1)
2	Conv2D(3,8,5)	7	LeakyReLU(0.1)	12	Linear(256,128)
3	LeakyReLU(0.1)	8	MaxPool(2,2)	Decoder
4	MaxPool(2,2)	9	Flatten	1	Linear(64,10)
5	Conv2D(8,16,3)	10	Linear(1024,256)	2	Softmax
D.6 Implementation and training details
Our prediction model was trained using PyTorch  $\mathrm{[PGM^{+}19]}$  and a GPU Tesla P100 with CUDA 11.0. All weights were initialized using the default PyTorch Xavier initialization scheme [GB10] with all biases initialized to zero. The Adam optimizer [KB15]  $(\beta_{1} = 0.5, \beta_{2} = 0.999)$  was used with an initial learning rate of  $10^{-4}$  and an exponential decay of 0.97. The batch size was equal to 128 throughout the whole experiment. The code used in the experiments is available at https://github.com/PiotrKrasnowski/MDL_and_Generalization_Guarantees_for_Representation_Learning.
During the training phase, we jointly trained the encoder and the decoder parts for 200 epochs using either the standard Gaussian prior of the traditional VIB objective function or our CDVIB objective functions. As in [AFDM17], we generated one latent sample per image during training and 12 samples during testing.
D.7 Numerical findings
Figure 3 displays the training and test performance of our model for a wide range of parameters  $\beta$ . It could be noticed that the best test accuracy with our Lossy CDVIB objective function is  $65\%$  (for  $\beta = 0.01$ ), which is about  $2.5\%$  better than the best test accuracy for VIB ( $62.5\%$  for  $\beta = 0.005$ ). In the case of Lossless CDVIB, the best achieved test accuracy is  $64\%$  for  $\beta = 1e - 5$ .
(b) Log-likelihood.
Figure 3: Test and train performances of our two-step prediction model trained using the standard VIB prior, the "lossless" CDVIB prior, and the "lossy" CDVIB prior, both with  $M = 5$ . The plots show the average over 5 runs and  $95\%$  bootstrap confidence intervals.
E.1 General proof techniques
Most of our proofs contain two main steps.
Technically speaking, the first step uses Donsker-Varadhan's variational representation lemma, to change the measure. As an example, in proof of Theorem 1, the first step gives
$$
\begin{array}{l} \lambda \mathbb {E} _ {S, W} [ \operatorname {g e n} (S, W) ] \\ \leqslant D _ {K L} \left(\mu_ {Y} ^ {\otimes 2 n} \mathbb {E} _ {S ^ {\prime}, S, W} \left[ P _ {\hat {Y} | X, W} ^ {\otimes 2 n} \left(\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} \mid X ^ {n}, X ^ {\prime n}, W\right) \right] \| \mu_ {Y} ^ {\otimes 2 n} \mathbf {Q} \left(\hat {Y} ^ {n}, \hat {Y} _ {i} ^ {\prime n} \mid Y ^ {n}, Y ^ {\prime n}\right)\right) \\ + \log \mathbb {E} _ {\boldsymbol {Y} ^ {n}, \boldsymbol {Y} ^ {\prime n}, \hat {\boldsymbol {Y}} ^ {n}, \hat {\boldsymbol {Y}} ^ {\prime n} \sim \mu_ {\boldsymbol {Y}} ^ {\otimes 2 n} \mathbf {Q} (\hat {\boldsymbol {Y}} ^ {n}, \hat {\boldsymbol {Y}} _ {i} ^ {\prime n} | \boldsymbol {Y} ^ {n}, \boldsymbol {Y} ^ {\prime n})} \left[ e ^ {\frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{\boldsymbol {Y} _ {i} ^ {\prime} \neq \hat {\boldsymbol {Y}} _ {i} ^ {\prime} \}} - \mathbb {1} _ {\{\boldsymbol {Y} _ {i} \neq \hat {\boldsymbol {Y}} _ {i} \}}\right)} \right], \\ \end{array}
$$
that changes the measure from
$$
\mu_ {Y} ^ {\otimes 2 n} \mathbb {E} _ {S ^ {\prime}, S, W} \Big [ P _ {\hat {Y} | X, W} ^ {\otimes 2 n} (\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | X ^ {n}, X ^ {\prime n}, W) \Big ],
$$
$$
\mu_ {Y} ^ {\otimes 2 n} \mathbf {Q} (\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}).
$$
This has a particular meaning and intuition in our fixed-size compressibility approach. Indeed, one can show that there exists a proper sequence of compression books, with size  $|\hat{y}_m| \leqslant e^{mR}$ , where
$$
R = D _ {K L} \left( \right.\mu_ {Y} ^ {\otimes 2 n} P _ {\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}} \left\| \right. \mu_ {Y} ^ {\otimes 2 n} \mathbf {Q} \left(\hat {Y} ^ {n}, \hat {Y} _ {i} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}\right)\left. \right).
$$
This means that by this step, we fix a suitable sequence of compression books such that with high probability (that goes to 1 as  $m \to \infty$ ), one can find the sequence of predicted labels in this codebook. Using Donsker-Varadhan's change of measure is a shortcut to this, as previously explained in [SGRS22, Appendix B.1.] in the context of hypothesis compression. Then, we consider the union bound over all elements of this codebook in the next step. For tail bounds, similar interpretations hold, but this time with variable-size compressibility notion, as shown in [SZ23], in the context of hypothesis compression.
The second step can be seen as bounding the generalization error for every element of such codebook. This step is achieved in quite different manners for the proofs of different results of the paper, as follows in the coming sections. Note that in this step, the predicted labels have distributions according to the prior  $\mathbf{Q}(\hat{Y}^n,\hat{Y}_i^{\prime n}|Y^n,Y^{\prime n})$ , rather than the one induced by the learning algorithm, i.e.,  $P_{\hat{Y}\mid X,W}^{\otimes 2n}(\hat{Y}^n,\hat{Y}^{\prime n}|X^n,X^{\prime n},W)$ .
E.2 Proof of Theorem 1
E.2.1 Part i.
Proof. We first show that
$$
\inf  _ {\mathbf {Q} \in \mathcal {Q} _ {i}} \mathbb {E} _ {\mathbf {Y}, \mathbf {Y} ^ {\prime}} \left[ D _ {K L} \left(\mathbb {E} _ {\mathbf {X} ^ {\prime}, \mathbf {X}, W} \left[ P _ {\hat {Y} | X, W} ^ {\otimes 2 n} (\hat {\mathbf {Y}}, \hat {\mathbf {Y}} ^ {\prime} | \mathbf {X}, \mathbf {X} ^ {\prime}, W) \right] \| \mathbf {Q}\right) \right] = I (\mathbf {J}; \hat {Y} ^ {2 n} | Y ^ {2 n}). \tag {19}
$$
Let  $\mathcal{Q}_i^\prime$ , be the set of conditional priors  $\mathbf{Q}'$  that can be written as
$$
\mathbf {Q} ^ {\prime} \left(\hat {\mathbf {Y}}, \hat {\mathbf {Y}} ^ {\prime} \mid \mathbf {Y}, \mathbf {Y} ^ {\prime}\right) = \mathbb {E} _ {\mathbf {J}} \left[ \mathbf {Q} _ {1} ^ {\prime} \left(\hat {Y} _ {\mathbf {J}} ^ {2 n}, \hat {Y} _ {\mathbf {J} ^ {c}} ^ {2 n} \mid Y _ {\mathbf {J}} ^ {2 n}, Y _ {\mathbf {J} ^ {c}} ^ {2 n}\right) \right], \tag {20}
$$
for some arbitrary distribution  $\mathbf{Q}_1'$ . Here  $Y^{2n}$  and  $\hat{Y}^{2n}$  are the concatenations of the vectors  $(\mathbf{Y},\mathbf{Y}')$  and  $(\hat{\mathbf{Y}},\hat{\mathbf{Y}}')$ , respectively. It is easy to verify that  $\mathcal{Q}_i = \mathcal{Q}_i'$ . Hence, by denoting  $P(\hat{\mathbf{Y}},\hat{\mathbf{Y}}'|\mathbf{Y},\mathbf{Y}') \coloneqq \mathbb{E}_{\mathbf{X}',\mathbf{X},W}\left[P_{\hat{Y}|X,W}^{\otimes 2n}(\hat{\mathbf{Y}},\hat{\mathbf{Y}}'|_{\mathbf{X}},\mathbf{X}',W)\right]$  we can write
$$
\begin{array}{l} \mathrm {L H S} = \inf  _ {\mathbf {Q} \in \mathcal {Q} _ {i}} \mathbb {E} _ {\mathbf {Y}, \mathbf {Y} ^ {\prime}} \left[ D _ {K L} \left(P (\hat {\mathbf {Y}}, \hat {\mathbf {Y}} ^ {\prime} | \mathbf {Y}, \mathbf {Y} ^ {\prime}) \| \mathbf {Q}\right) \right] \\ = \inf _ {\mathbf {Q} ^ {\prime} \in \mathcal {Q} _ {i} ^ {\prime}} \mathbb {E} _ {\mathbf {Y}, \mathbf {Y} ^ {\prime}} \left[ D _ {K L} \left(P (\hat {\mathbf {Y}}, \hat {\mathbf {Y}} ^ {\prime} | \mathbf {Y}, \mathbf {Y} ^ {\prime}) \| \mathbf {Q} ^ {\prime}\right) \right] \\ \end{array}
$$
$$
\begin{array}{l} = \inf  _ {\mathbf {Q} ^ {\prime} \in \mathcal {Q} _ {i} ^ {\prime}} \mathbb {E} _ {Y ^ {2 n}} \mathbb {E} _ {\mathbf {J}} \left[ D _ {K L} \left(P \left(\hat {Y} _ {\mathbf {J}} ^ {2 n}, \hat {Y} _ {\mathbf {J} ^ {c}} ^ {2 n} \mid Y _ {\mathbf {J}} ^ {2 n}, Y _ {\mathbf {J} ^ {c}} ^ {2 n}\right) \| \mathbf {Q} ^ {\prime}\right) \right] \\ = \inf  _ {\mathbf {Q} _ {1} ^ {\prime}} \mathbb {E} _ {Y ^ {2 n}} \mathbb {E} _ {\mathbf {J}} \left[ D _ {K L} \left(P \left(\hat {Y} _ {\mathbf {J}} ^ {2 n}, \hat {Y} _ {\mathbf {J} ^ {c}} ^ {2 n} \mid Y _ {\mathbf {J}} ^ {2 n}, Y _ {\mathbf {J} ^ {c}} ^ {2 n}\right) \| \mathbb {E} _ {\mathbf {J}} \left[ \mathbf {Q} _ {1} ^ {\prime} \left(\hat {Y} _ {\mathbf {J}} ^ {2 n}, \hat {Y} _ {\mathbf {J} ^ {c}} ^ {2 n} \mid Y _ {\mathbf {J}} ^ {2 n}, Y _ {\mathbf {J} ^ {c}} ^ {2 n}\right) \right]\right) \right] \\ = I \left(\mathbf {J}; \hat {Y} ^ {2 n} \mid Y ^ {2 n}\right). \\ \end{array}
$$
This completes the proof of showing the equality in (2).
Now, we proceed to prove the upper bound. Let  $\mathbf{Q}$  be any fixed type-I symmetric conditional prior on  $(\hat{Y}^n,\hat{Y}^{\prime n})$  given  $(Y^n,Y^{\prime n})$ . We show that
$$
\mathbb {E} _ {S, W} [ \operatorname {g e n} (S, W) ] \leqslant \sqrt {\frac {2 \mathbb {E} _ {Y ^ {n} , Y ^ {\prime n}} \left[ \right. D _ {K L} \left( \right.\mathbb {E} _ {X ^ {\prime n} , X ^ {n} , W} \left[ P _ {\hat {Y} | X , W} ^ {\otimes 2 n} \left(\hat {Y} ^ {n} , \hat {Y} ^ {\prime n} \mid X ^ {n} , X ^ {\prime n} , W\right)\right]\left\| \right. \mathbf {Q}\left. \right)\left. \right]}{n}}, \tag {21}
$$
where  $Y^{n}, Y'^{m} \sim \mu_{Y}^{\otimes 2n}$  and  $X'^{m}, X^n, W \sim P_{X'^{m}|Y'^{m}}P_{X^n,W|Y^n}$ . For ease of notation, denote
$$
P _ {\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}} := \mathbb {E} _ {X ^ {\prime n}, X ^ {n}, W} \Big [ P _ {\hat {Y} | X, W} ^ {\otimes 2 n} (\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | X ^ {n}, X ^ {\prime n}, W) \Big ],
$$
where again  $X'^n, X^n, W \sim P_{X'^n|Y'^n}P_{X^n,W|Y^n}$ . We start the proof by applying the change of measure using Donsker-Varadhan's variational representation (step  $(a)$  below):
$$
\begin{array}{l} \lambda \mathbb {E} _ {S, W} [ \operatorname {g e n} (S, W) ] \\ = \mathbb {E} _ {S, S ^ {\prime}, W, \hat {Y} ^ {n}, \hat {Y} ^ {\prime n} \sim P _ {S ^ {\prime}} P _ {S, W}} P _ {\hat {Y} | X, W} ^ {\otimes 2 n} (\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | X ^ {n}, X ^ {\prime n}, W) \left[ \frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{Y _ {i} ^ {\prime} \neq \hat {Y} _ {i} ^ {\prime} \}} - \mathbb {1} _ {\{Y _ {i} \neq \hat {Y} _ {i} \}}\right) \right] \\ = \mathbb {E} _ {Y ^ {n}, Y ^ {\prime n}, \hat {Y} ^ {n}, \hat {Y} ^ {\prime n} \sim \mu_ {Y} ^ {\otimes 2 n} P _ {\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}}} \left[ \frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{Y _ {i} ^ {\prime} \neq \hat {Y} _ {i} ^ {\prime} \}} - \mathbb {1} _ {\{Y _ {i} \neq \hat {Y} _ {i} \}}\right) \right] \\ \stackrel {(a)} {\leqslant} D _ {K L} \left(\mu_ {Y} ^ {\otimes 2 n} P _ {\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}} \| \mu_ {Y} ^ {\otimes 2 n} \mathbf {Q} \left(\hat {Y} ^ {n}, \hat {Y} _ {i} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}\right)\right) \\ + \log \mathbb {E} _ {Y ^ {n}, Y ^ {\prime n}, \hat {Y} ^ {n}, \hat {Y} ^ {\prime n} \sim \mu_ {Y} ^ {\otimes 2 n}} \mathbf {Q} (\hat {Y} ^ {n}, \hat {Y} _ {i} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}) \left[ e ^ {\frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{Y _ {i} ^ {\prime} \neq \hat {Y} _ {i} ^ {\prime} \}} - \mathbb {1} _ {\{Y _ {i} \neq \hat {Y} _ {i} \}}\right)} \right]. \tag {22} \\ \end{array}
$$
Now, we bound the second term. Consider the notation  $\mathfrak{Y}_{i,j} \in \mathcal{Y}^{2n}$  and  $\hat{\mathfrak{Y}}_{i,j} \in \mathcal{Y}^{2n}$  for  $i \in [n]$  and  $j \in \{1,2\}$ . Denote  $j^c \coloneqq \mathbb{1}_{\{j = 1\}} + 1$ . Furthermore, for every  $i \in [n]$ , let  $K_i$  be a random variable that takes values uniformly over  $\{1,2\}$  — The variables  $\{K_i\}$  are assumed to be mutually independent. Let the complement variable  $K_i^c$  be equal to 2 if  $K_{i,j} = 1$  and 1 otherwise.
$$
\begin{array}{l} \log \mathbb {E} _ {Y ^ {n}, Y ^ {\prime n}, \hat {Y} ^ {n}, \hat {Y} ^ {\prime n} \sim \mu_ {Y} ^ {\otimes 2 n}} \mathbf {Q} (\hat {Y} ^ {n}, \hat {Y} _ {i} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}) \left[ e ^ {\frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{Y _ {i} ^ {\prime} \neq \hat {Y} _ {i} ^ {\prime} \}} - \mathbb {1} _ {\{Y _ {i} \neq \hat {Y} _ {i} \}}\right)} \right] \\ \stackrel {(a)} {=} \log \mathbb {E} _ {\mathfrak {Y} ^ {n \times 2}, \hat {\mathfrak {Y}} ^ {n \times 2}, K ^ {n} \sim P _ {Y} ^ {\otimes 2 n}} \mathbf {Q} (\hat {\mathfrak {Y}} ^ {n \times 2} | \mathfrak {Y} ^ {n \times 2}) \operatorname {U n i f} (1, 2) ^ {\otimes n} \left[ e ^ {\frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{\mathfrak {Y} _ {K _ {i}, i} \neq \hat {\mathfrak {Y}} _ {K _ {i}, i} \}} - \mathbb {1} _ {\{\mathfrak {Y} _ {K _ {i} ^ {c}, i} \neq \hat {\mathfrak {Y}} _ {K _ {i} ^ {c}, i} \}}\right)} \right] \\ = \log \mathbb {E} _ {\mathfrak {Y} ^ {n \times 2}, \hat {\mathfrak {Y}} ^ {n \times 2} \sim P _ {Y} ^ {\otimes 2 n}} \mathbf {Q} (\hat {\mathfrak {Y}} ^ {n \times 2} | \mathfrak {Y} ^ {n \times 2}) \mathbb {E} _ {K ^ {n} \sim \mathrm {U n i f} (1, 2) ^ {\otimes n}} \left[ e ^ {\frac {\lambda}{n} \sum_ {i \in [ n ]} \left(\mathbb {1} _ {\{\mathfrak {Y} _ {K _ {i}, i} \neq \hat {\mathfrak {Y}} _ {K _ {i}, i} \}} - \mathbb {1} _ {\{\mathfrak {Y} _ {K _ {i} ^ {c}, i} \neq \hat {\mathfrak {Y}} _ {K _ {i} ^ {c}, i} \}}\right)} \right] \\ \stackrel {(b)} {\leqslant} \log \left(\frac {e ^ {\lambda / n} + e ^ {- \lambda / n}}{2}\right) ^ {n} \\ \leqslant \frac {\lambda^ {2}}{2 n}, \tag {23} \\ \end{array}
$$
where  $(a)$  is concluded by symmetry of  $\mathbf{Q}$  and  $(b)$  by inequality  $\frac{e^x + e^{-x}}{2} \leqslant e^{x^2 / 2}$ . Combining this with (22) yield
$$
\mathbb {E} _ {S, W} [ \operatorname {g e n} (S, W) ] \leqslant \frac {1}{\lambda} D _ {K L} \left( \right.\mu_ {Y} ^ {\otimes 2 n} P _ {\hat {Y} ^ {n}, \hat {Y} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n}} \left\| \right. \mu_ {Y} ^ {\otimes 2 n} \mathbf {Q} (\hat {Y} ^ {n}, \hat {Y} _ {i} ^ {\prime n} | Y ^ {n}, Y ^ {\prime n})\left. \right) + \frac {\lambda}{2 n}.
$$
$$
\lambda := \sqrt {2 n D _ {K L} \left(\mu_ {Y} ^ {\otimes 2 n} P _ {\hat {Y} ^ {n} , \hat {Y} ^ {\prime n} | Y ^ {n} , Y ^ {\prime n}} \left\lVert \mu_ {Y} ^ {\otimes 2 n} \mathbf {Q} (\hat {Y} ^ {n} , \hat {Y} _ {i} ^ {\prime n} | Y ^ {n} , Y ^ {\prime n})\right)\right)},
$$
completes the proof.
E.2.2 Part ii.
Proof. The proof of this proposition follows similarly as proof of Theorem 5 (with  $\epsilon = 0$ ), and avoided for brevity. Note that similar to Theorem 5, by noting that with probability at least  $(1 - \delta)$ ,  $\hat{\mathcal{L}}(S', W) \geqslant \mathcal{L}(W) - \sqrt{\log(1/\delta)/(2n)}$ , one can also establish a tail bound on  $\operatorname{gen}(S, W)$ .
E.3 Proof of Lemma 1