Information-Theoretic Characterization of the Generalization Error for Iterative Semi-Supervised Learning

In real-life machine learning applications, it is relatively easy and inexpensive to obtain large amounts of unlabelled data, while the number of labelled data examples is usually small due to the high cost of annotating them with true labels. In light of this, semi-supervised learning (SSL) has come to the fore (Chapelle et al., 2006; Zhu, 2008; Van Engelen and Hoos, 2020). SSL makes use of the abundant unlabelled data to augment the performance of learning tasks with few labelled data examples. This has been shown to outperform supervised and unsupervised learning under certain conditions. For example, in a classification problem, the correlation between the additional unlabelled data and the labelled data may help to enhance the accuracy of classifiers. Among the plethora of SSL methods, pseudo-labelling (Lee, 2013) has been observed to be a simple and efficient way to improve the generalization performance empirically. In this paper, we consider the problem of pseudo-labelling a subset of the unlabelled data at each iteration based on the previous output parameter and then refining the model progressively, but we are interested in analysing this procedure theoretically. Our goal in this paper is to understand the impact of pseudo-labelling on the generalization error.

A learning algorithm can be viewed as a randomized map from the training dataset to the output model parameter. The output is highly data-dependent and may suffer from overfitting to the given dataset. In statistical learning theory, the generalization error (gen-error), or generalization bias, is defined as the expected gap between the test and training losses, and is used to measure the extent to which the algorithms overfit to the training data (Russo and Zou, 2016; Xu and Raginsky, 2017; Kawaguchi et al., 2017). In SSL problems, the unlabelled data are expected to improve the generalization performance in a certain manner and thus, it is a worthy endeavor to investigate the behaviour theoretically. Although there exist many works studying the gen-error for supervised learning problems, the gen-error of SSL algorithms is yet to be explored.

Let the instance space be $\mathcal{Z}=\mathcal{X}\times\mathcal{Y}\subset\mathbb{R}^{d+1}$, the model parameter space be $\Theta$ and the loss fucntion be $l:\mathcal{Z}\times\Theta\to\mathbb{R}$, where $d\in\mathbb{N}$. We are given a labelled training dataset $S_{\mathrm{l}}=\{Z_{1},\ldots,Z_{n}\}=\{(X_{i},Y_{i})\}_{i=1}^{n}$ drawn from $\mathcal{Z}$, where each $Z_{i}=(X_{i},Y_{i})$ is independently and identically distributed (i.i.d.) from $P_{Z}=P_{X,Y}\in\mathcal{P}(\mathcal{Z})$. For any $i\in[n]$, $X_{i}$ is a vector of features and $Y_{i}$ is a label indicating the class to which $X_{i}$ belongs. However, in many real-life machine learning applications, we only have a limited number of labelled data while we have access to a large amount of unlabelled data, which are expensive to annotate. Then we can incorporate the unlabelled training data together with the labelled data to improve the performance of the model. This procedure is called semi-supervised learning (SSL). We are given an independent unlabelled training dataset $S_{\mathrm{u}}=\{X^{\prime}_{1},\ldots,X^{\prime}_{\tau m}\},\tau\in\mathbb{N}$, where each $X^{\prime}_{i}$ is generated i.i.d. from $P_{X}\in\mathcal{P}(\mathcal{X})$. Typically, $m\gg n$.

In the following, we consider the iterative self-training with pseudo-labelling in SSL setup, as shown in Figure 1. Let $t\in[0:\tau]$ denote the iteration count. In the initial round ($t=0$), the labelled data $S_{\mathrm{l}}$ are first used to learn an initial model parameter $\theta_{0}\in\Theta$. Next, we split the unlabelled dataset $S_{\mathrm{u}}$ into $\tau$ disjoint equal-size sub-datasets $\{S_{\mathrm{u},k}\}_{k=1}^{\tau}$, where $S_{\mathrm{u},k}=\{X^{\prime}_{(k-1)m+1},\ldots,X^{\prime}_{km}\}$. In each subsequent round $t\in[1:\tau]$, based on $\theta_{t-1}$ trained from the previous round, we use a predictor $f_{\theta_{t-1}}:\mathcal{X}\mapsto\mathcal{Y}$ to assign a pseudo-label $\hat{Y}^{\prime}_{i}$ to the unlabelled sample $X^{\prime}_{i}$ for all $i\in\mathcal{I}_{t}:=\{(t-1)m,(t-1)m+1,\ldots,tm\}$. Let $\hat{S}_{\mathrm{u},t}=\{(X_{i}^{\prime},\hat{Y}_{i}^{\prime})\}_{i\in\mathcal{I}_{t}}$ denote the $t^{\mathrm{th}}$ pseudo-labelled dataset. After pseudo-labelling, both the labelled data $S_{\mathrm{l}}$ and the pseudo-labelled data $\hat{S}_{\mathrm{u},t}$ are used to learn a new model parameter $\theta_{t}$. The procedure is then repeated iteratively until the maximum number of iterations $\tau$ is reached.

This setup is a classical and widely-used model in the realm of self-training in SSL (Chapelle et al., 2006; Zhu, 2008; Zhu and Goldberg, 2009; Lee, 2013), where in each iteration, only a subset of the unlabelled data are used. Furthermore, as discussed by Arazo et al. (2020), this method is less likely to overfit to incorrect pseudo-labels, compared to using all the unlabelled data in each iteration (also see Figure 10). Under this setup of iterative SSL, during each iteration $t$, our goal is to find a model parameter $\theta_{t}\in\Theta$ that minimizes the population risk with respect to the underlying data distribution

$\displaystyle L_{P_{Z}}(\theta_{t}):=\mathbb{E}_{Z\sim P_{Z}}[l(\theta_{t},Z)].$

(1)

Since $P_{Z}$ is unknown, $L_{P_{Z}}(\theta_{t})$ cannot be computed directly. Hence, we instead minimize the empirical risk. The procedure is termed empirical risk minimization (ERM). For any model parameter $\theta_{t}\in\Theta$, the empirical risk of the labelled data is defined as

$\displaystyle L_{S_{\mathrm{l}}}(\theta_{t}):=\frac{1}{n}\sum_{i=1}^{n}l(\theta_{t},Z_{i}),$

(2)

and for $t\geq 1$, the empirical risk of pseudo-labelled data $\hat{S}_{\mathrm{u},t}$ as

$\displaystyle L_{\hat{S}_{\mathrm{u},t}}(\theta_{t}):=\frac{1}{m}\sum_{i\in\mathcal{I}_{t}}l(\theta_{t},(X_{i}^{\prime},\hat{Y}_{i}^{\prime})).$

(3)

We set $L_{\hat{S}_{\mathrm{u},t}}(\theta_{t})=0$ for $t=0$. For a fixed weight $w\in[0,1]$, the total empirical risk can be defined as the following linear combination of $L_{S_{\mathrm{l}}}(\theta_{t})$ and $L_{\hat{S}_{\mathrm{u},t}}(\theta_{t})$:

$\displaystyle L_{S_{\mathrm{l}},\hat{S}_{\mathrm{u},t}}(\theta_{t})$

$\displaystyle:=wL_{S_{\mathrm{l}}}(\theta_{t})+(1-w)L_{\hat{S}_{\mathrm{u},t}}(\theta_{t}).$

(4)

In the usual case where the algorithm minimizes the average of the empirical training losses, one should set $w=\frac{n}{n+m}$. An SSL algorithm can be characterized by a randomized map from the labelled and unlabelled training data $S_{\mathrm{l}}$, $S_{\mathrm{u}}$ to a model parameter $\theta$ according to a conditional distribution $P_{\theta|S_{\mathrm{l}},S_{\mathrm{u}}}$. Then at each iteration $t$, we can use the sequence of conditional distributions $\{P_{\theta_{k}|S_{\mathrm{l}},S_{\mathrm{u}}}\}_{k=0}^{t}$ with $P_{\theta_{0}|S_{\mathrm{l}},S_{\mathrm{u}}}=P_{\theta_{0}|S_{\mathrm{l}}}$ to represent an iterative SSL algorithm. The generalization error at the $t$-th iteration is defined as the expected gap between the population risk of $\theta_{t}$ and the empirical risk on the training data:

	$\displaystyle\mathrm{gen}_{t}(P_{Z},P_{X},\{P_{\theta_{k}|S_{\mathrm{l}},S_{\mathrm{u}}}\}_{k=0}^{t},\{f_{\theta_{k}}\}_{k=0}^{t-1})$
	$\displaystyle:=\mathbb{E}[L_{P_{Z}}(\theta_{t})-L_{S_{\mathrm{l}},\hat{S}_{\mathrm{u},t}}(\theta_{t})]$		(5)
	$\displaystyle=w\bigg{(}\mathbb{E}_{\theta_{t}}\mathbb{E}_{Z}[l(\theta_{t},Z)]-\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{\theta_{t},Z_{i}}[l(\theta_{t},Z_{i})]\bigg{)}$
	$\displaystyle\quad+(1-w)\bigg{(}\mathbb{E}_{\theta_{t}}\mathbb{E}_{Z}[l(\theta_{t},Z)]-\frac{1}{m}\sum_{i\in\mathcal{I}_{t}}\mathbb{E}_{\theta_{t},X^{\prime}_{i},\hat{Y}^{\prime}_{i}}[l(\theta_{t},(X^{\prime}_{i},\hat{Y}^{\prime}_{i}))]\bigg{)}.$		(6)

When $t=0$ and $w=1$, the definition of the generalization error reduces to that of vanilla supervised learning. Based on this definition, the expected population risk can be decomposed as

$$\mathbb{E}[L_{P_{Z}}(\theta_{t})]=\mathbb{E}[L_{S_{\mathrm{l}},\hat{S}_{\mathrm{u},t}}(\theta_{t})]+\mathrm{gen}_{t},$$

(7)

where the first term on the right-hand side of this equation is what the algorithm minimizes and reflects how well the output hypothesis fits the dataset, and the second term $\mathrm{gen}_{t}$ is used to measure the extent to which the iterative learning algorithm overfits the training data at the $t$-th iteration. To minimize $\mathbb{E}[L_{P_{Z}}(\theta_{t})]$, we need both terms in (7) to be small, but there exists a natural trade-off between them. While the algorithm aims to minimize the empirical risk $\mathbb{E}[L_{S_{\mathrm{l}},\hat{S}_{\mathrm{u},t}}(\theta_{t})]$, studying and controlling $\mathrm{gen}_{t}$ can also help to reduce the population risk $\mathbb{E}[L_{P_{Z}}(\theta_{t})]$, which is the ultimate goal of learning. Instead of focusing on the total generalization error induced during the entire process, we are interested in the following questions. How does $\mathrm{gen}_{t}$ evolve as $t$ increases? Do the unlabelled data examples in $S_{\mathrm{u}}$ help to improve the generalization error?

Inspired by the information-theoretic generalization results in Bu et al. (2020, Theorem 1) and Wu et al. (2020, Theorem 1), we derive an upper bound on the gen-error $\mathrm{gen}_{t}$ in terms of the mutual information between input data samples (either labelled or pseudo-labelled) and the output model parameter $\theta_{t}$, as well as the KL-divergence between the data distribution and the joint distribution of feature vectors and pseudo-labels (cf. Theorem 2.A). Furthermore, by considering the NLL loss function (MacKay, 2003; Goodfellow et al., 2016), we derive the exact characterization for the gen-error $\mathrm{gen}_{t}$ (cf. Theorem 2.B).

Recall that for a given $R>0$, $L$ is an $R$-sub-Gaussian random variable (Vershynin, 2018) if its cumulant generating function $\Lambda_{L}(\lambda):=\log\mathbb{E}[\exp(\lambda(L-\mathbb{E}[L]))]\leq\exp(\lambda^{2}R^{2}/2)$ for all $\lambda\in\mathbb{R}$. If $L$ is $R$-sub-Gaussian, we write this as $L\sim\mathrm{subG}(R)$. Furthermore, let us recall the following somewhat non-standard information quantities (Negrea et al., 2019; Haghifam et al., 2020).

Let $\theta^{(t)}=(\theta_{0},\ldots,\theta_{t})$ for any $t\in[0:\tau]$ and $w=1$ for $t=0$. In iterative SSL, we can characterize the gen-error as shown in Theorem 3 by applying the law of total expectation.

The proof of Theorem 2.A is provided in Appendix A, in which we provide a general upper bound not only applicable to sub-Gaussian loss functions. The proof of Theorem 2.B is provided in Appendix B. Specifically, for NLL loss functions, Theorem 2.B provides an exact characterization of the gen-error at each iteration. This is in stark contrast to most works on information-theoretic generalization error in which only bounds are provided.

In contrast to Bu et al. (2020, Theorem 1) and Wu et al. (2020, Theorem 1) which pertain to supervised learning, Theorem 3 characterizes the gen-error at each iteration during the pseudo-labelling and training process. Note that the quantities in Theorem 2.B satisfy $\mathrm{\Delta h}^{(i)}_{\theta_{t}}=I_{\theta^{(t-1)}}^{(i)}+D(P_{Z}\|p_{\theta_{t}})-D(P_{Z_{i}|\theta_{t}}\|p_{\theta_{t}})$ and $\widetilde{\mathrm{\Delta h}}^{\prime(i)}_{\theta^{(t)}}=I_{\theta^{(t-1)}}^{\prime(i)}+D_{\theta^{(t-1)}}(P_{X_{i}^{\prime},\hat{Y}_{i}^{\prime}}\|p_{\theta_{t}})-D_{\theta^{(t-1)}}(P_{X_{i}^{\prime},\hat{Y}_{i}^{\prime}|\theta_{t}}\|p_{\theta_{t}})$. Thus, it is plausible that the upper bound based on $I_{\theta^{(t-1)}}^{(i)}$ and $I_{\theta^{(t-1)}}^{\prime(i)}$ in Theorem 2.A can help to understand and control the exact gen-error. Intuitively, the mutual information between the individual input data sample $Z_{i}$ and the output model parameter $\theta_{t}$ in Theorem 2.A and the cross-entropy divergences $\mathrm{\Delta h}^{(i)}_{\theta_{t}}$, $\widetilde{\mathrm{\Delta h}}^{\prime(i)}_{\theta^{(t)}}$ in Theorem 2.B both measure the extent to which the algorithm is sensitive to each data example at each iteration $t$. The KL-divergence between the underlying $P_{Z}$ and pseudo-labelled distribution $P_{X_{i}^{\prime},\hat{Y}_{i}^{\prime}}$ in Theorem 2.A and the cross-entropy divergence $\mathrm{\Delta h}^{\prime(i)}_{\theta^{(t)}}$ in Theorem 2.B measure how effectively the pseudo-labelling process works. As $n\to\infty$ and $m\to\infty$, we show that the mutual information (as well as $\mathrm{\Delta h}^{(i)}_{\theta_{t}}$ and $\widetilde{\mathrm{\Delta h}}^{\prime(i)}_{\theta^{(t)}}$) vanishes but the divergences $D_{\theta^{(t-1)}}^{\prime(i)}$ and $\mathrm{\Delta h}^{\prime(i)}_{\theta^{(t)}}$ do not, which reflects the impact of pseudo-labelling on the gen-error.

In iterative learning algorithms, by applying the law of total expectation and conditioning the information-theoretic quantities on the output model parameters $\theta^{(t-1)}=\{\theta_{1},\ldots,\theta_{t-1}\}$ from previous iterations, we are able to calculate the gen-error iteratively. In the next section, we apply the exact iterated gen-error in Theorem 2.B to a classification problem under a specific generative model—the bGMM. This simple model allows us to derive a tractable characterization on the gen-error as a function of iteration number $t$ that we can compute numerically.

We now particularize the iterative semi-supervised classification setup to the bGMM. We evaluate (9) to understand the effect of multiple self-training rounds on the gen-error.

In Section 4.3, it is shown that for difficult classification problems with large class conditional variance, the gen-error increases after using pseudo-labelled data. The reason is that the learned initial parameter $\bm{\theta}_{0}$ can only generate low-accurate pseudo-labels and thus the pseudo-labelled data cannot help improve the generalization performance. In this section, we prove that by adding regularization to the loss function, we can mitigate the undesirable increase of gen-error across the pseudo-labelling iterations.

Since $\mathrm{gen}_{0}$ in (22) does not depend on data variance $\sigma^{2}$, here we focus on subsequent iterations $t\in[1:\tau]$. By considering the $\ell_{2}$ regularization (i.e., adding $\frac{\lambda}{2}\|\bm{\theta}\|_{2}^{2}$ to (15)), we obtain the new parameters (cf. (16)) as follows

$\displaystyle\bm{\theta}_{t}^{\mathrm{reg}}=\frac{\bm{\theta}_{t}}{1+\sigma^{2}\lambda},\quad\forall\,t\in[1:\tau].$

(26)

The derivations are provided in Appendix H. By applying Theorem 3, the following theorem provides a characterization for the gen-error of the case with regularization at each iteration $t$ for $m$ large enough. Let $\mathrm{gen}_{t}^{\mathrm{reg}}$ denotes the gen-error of the case with regularization and we drop the fixed quantities $(P_{\mathbf{Z}},P_{\mathbf{X}},\{P_{\bm{\theta}_{k}|S_{\mathrm{l}},S_{\mathrm{u}}}\}_{k=0}^{t},\{f_{\bm{\theta}_{k}}\}_{k=0}^{t-1})$ for notational simplicity.

The proof of Theorem 4 is provided in Appendix H. From (27), we observe that as $\lambda$ increases, the gen-error decreases. In Figure 8, we first empirically show that regularization can help mitigate the increase of gen-error during SSL iterations for hard-to-distinguish classes by comparing the empirical gen-error under $\lambda=0,0.1,0.5,2$ when $\sigma=3$ and $d=2$. Then in Figure 8, we plot the theoretical gen-error in (27) versus $\lambda$ when $t=1$ for the cases with small and large variances, i.e., $\sigma^{2}=0.6^{2}$ and $\sigma^{2}=3^{2}$. We also compare the theoretical results with the empirical gen-errors, which turn out to corroborate the theoretical ones. For the case with smaller variance, the improvement on gen-error is barely visible as $\lambda$ increases. For the case with larger variance, the decrease of the gen-error is more pronounced, which implies that $\ell_{2}$-regularization can effectively mitigate the impact on the gen-error induced by large class overlapping and pseudo-labels with low accuracy.

The adept reader might naturally wonder why one would not set $\lambda\to\infty$ in (27), which results in the gen-error tending to zero, which presumably is a desirable phenomenon. However, ultimately, what we wish to control is the expected population risk, which, according to (7), is the sum of the expected empirical risk on the training data and the gen-error. Even if the gen-error is zero, the expected population risk might be large. Hence, as $\lambda$ increases, we see a tradeoff between the gen-error and the empirical risk.

To further illustrate that our theory is indeed behind the empirical behaviour of the iterative self-training with pseudo-labelling, in this section, we conduct experiments on real-world benchmark datasets, which demonstrates that our theoretical results on the bGMM example can also reflect the training dynamics on more realistic real-world tasks. The code to reproduce all the experiments can be found at https://github.com/HerianHe/GenErrorSSL_2022.git.

Recall that in the bGMM example, a higher standard deviation $\sigma$ represents a higher in-class variance, larger class-overlap, and consequently higher difficulty in classification. By a whitening argument, this also holds for bGMMs with non-isotropic covariance matrices. In our experiments on real-world data, we use the difficulty level of classification to mimic different in-class variances of bGMM. We pick two easy-to-distinguish class pairs (“automobile” and “truck”, “horse” and “ship”) from the CIFAR-10 dataset (Krizhevsky, 2009) as an analogy to bGMM with small in-class variance, and one difficult-to-distinguish class pair (“cat” and “dog”) from the same dataset as an analogy to bGMM with large in-class variance. Furthermore, to extend the analogy to multi-class classification, we conduct experiments on the 10-class MNIST dataset to gain more intuition.

We train deep neural networks (DNNs) via an iterative self-learning strategy (under the same setting as Figure 1) to perform binary and multi-class classification. In the first iteration, we only use a few labelled data examples to initialize the DNN with a sufficient number of epochs. In the subsequent iterations, we first sample a subset of unlabelled data and generate pseudo-labels for them via the model trained from the previous iteration. Then we update the model for a small number of epochs with both the labelled and pseudo-labelled data.

Experimental settings: For binary classification, we collect pairs of classes of images, i.e., “automobile” and “truck”, “horse” and “ship”, and “cat” and “dog” from the CIFAR10 (Krizhevsky, 2009) dataset. In this dataset, each class has $5000$ images for training and $1000$ images for testing. For each selected pair of classes, we manually divide the $10000$ training images into two sets: the labelled training set with $500$ images and the unlabelled training set with $9500$ images. We train a convolutional neural network, ResNet-10 (He et al., 2016), and use stochastic gradient descent (SGD) optimizer to minimize the cross-entropy loss. In PyTorch, the cross-entropy loss is defined as the negative logarithm of the output softmax probability corresponding to the true class, which is analogous to the NLL of the data under the parameters of the neural network. For the task with $\ell_{2}$-regularization, we train the neural network by setting different weight decay parameters (equivalent to $\lambda\times$ learning rate). In each pseudo-labelling iteration, we sample $2500$ unlabelled images. The complete training procedure lasts for $50$ self-training iterations.

We further validate our theoretical contributions on a multi-class classification problem in which we train a ResNet-6 model with the cross-entropy loss to perform $10$-class handwritten digits classification on the MNIST (LeCun et al., 1998) dataset. We sample $51000$ images from the training set, which contains $6000$ images for each of the ten classes. We divide them into two sets, i.e., a labelled training set with $1000$ images and an unlabelled set with $50000$ images. The optimizer and training iterations follow those in the aforementioned binary classification tasks without regularization.