Information-Theoretic Generalization Bounds for Iterative Semi-Supervised Learning

1 INTRODUCTION . In real-life machine learning applications , it is relatively easy and cheap 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 & 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 et al. , 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 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 . In SSL problems , the unlabelled data are expected to improve the generalization performance in a certain manner and thus , it is worthwhile to investigate the behaviour theoretically . In this paper , we leverage results in Bu et al . ( 2020 ) ; Wu et al . ( 2020 ) to derive an informationtheoretic generalization error bound at each iteration for iterative SSL . We state our main theoretical contribution informally as follows . Theorem [ Informal ] For a d-variate binary Gaussian mixture model ( bGMM ) in which each component has variance 2 , the generalization error across the different semi-supervised training iterations |gent| can be bounded with high probability as follows : |gent| / const · E hq G ( F ( t 1 ) ( ↵ ) ) i , ( 1 ) where ↵ represents the correlation between the optimal and estimated parameter vectors , F ( t ) is the iterated composition of the function F ( sketched in Figure 3 ) , and G ( sketched in Figure 5 ) represents the KL-divergence between the pseudo-labelled and true data distributions . As shown in Figure 1 , the upper bound is monotonically decreasing in the iteration count t and converges at around t = 2 with a sufficiently large amount of unlabelled data . In Section 4 , we also show that when the number of labelled data or the variance is large enough , using the unlabelled data does not help to significantly reduce the generalization error across iterations t. The behaviour of the empirical generalization error for the bGMM coincides with the upper bound . The results suggest that the proposed upper bound serves as a useful guide to understand how the generalization error changes across the semi-supervised training iterations and it can be used to establish conditions under which unlabelled data can help in terms of generalization . Experimental results on the MNIST and CIFAR datasets corroborate the phenomena for the bGMM that with few labelled data and abundant unlabelled data , the generalization error decreases quickly in the early pseudo-labelling iterations and saturates thereafter . For a more extensive literature review , please refer to Appendix A . 2 PROBLEM SETUP . Let the instance space be Z = X⇥Y ⇢ Rd+1 , the model parameter space be ⇥ and the loss fucntion be l : Z ⇥ ⇥ ! R , where d 2 N. We are given a labelled training dataset Sl = { Z1 , . . . , Zn } = { ( Xi , Yi ) } ni=1 drawn from Z , where each Zi = ( Xi , Yi ) is independently and identically distributed ( i.i.d . ) from PZ = PX , Y 2 P ( Z ) and Xi is i.i.d . from PX 2 P ( X ) . For any i 2 [ n ] , Xi is a vector of features and Yi is a label indicating the class to which Xi 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 Su = { X 01 , . . . , X 0⌧m } , ⌧ 2 N , where each X 0i is i.i.d . generated from PX 2 P ( X ) . Typically , m n. In the following , we consider the iterative self-training with pseudo-labelling in SSL setup , as shown in Figure 2 . Let t 2 [ 0 : ⌧ ] denote the iteration counter . In the initial round ( t = 0 ) , the labelled data Sl are first used to learn an initial model parameter ✓0 2 ⇥ . Next , we split the unlabelled dataset Su into ⌧ disjoint equal-size sub-datasets { Su , k } ⌧k=1 , where Su , k = { X 0 ( k 1 ) m+1 , . . . , X 0km } . In each subsequent round t 2 [ 1 : ⌧ ] , based on ✓t 1 trained from the previous round , we use a predictor f✓t 1 : X 7 ! Y to assign a pseudo-label Ŷ 0i to the unlabelled sample X 0i for all i 2 [ ( t 1 ) m+ 1 : tm ] : = { ( t 1 ) m , ( t 1 ) m+1 , . . . , tm } . Let Ŝu , t = { ( X 0i , Ŷ 0i ) } tmi= ( t 1 ) m+1 denote the tth pseudolabelled dataset . After pseudo-labelling , both the labelled data Sl and the pseudo-labelled data Ŝu , t are used to learn a new model parameter ✓t . The procedure is then repeated iteratively until the maximum number of iterations ⌧ is reached . Under the setup of iterative SSL , during each iteration t , our goal is to find a model parameter ✓t 2 ⇥ that minimizes the population risk with respect to the underlying data distribution LPZ ( ✓t ) : = EZ⇠PZ [ l ( ✓t , Z ) ] . ( 2 ) Since PZ is unknown , LPZ ( ✓t ) can not be computed directly . Hence , we instead minimize the empirical risk . The procedure is termed empirical risk minimization ( ERM ) . For any model parameter ✓t 2 ⇥ , the empirical risk of the labelled data is defined as LSl ( ✓t ) : = 1 n nX i=1 l ( ✓t , Zi ) , ( 3 ) and for t 1 , the empirical risk of pseudo-labelled data Ŝu , t as LŜu , t ( ✓t ) : = 1 m tmX i= ( t 1 ) m+1 l ( ✓t , ( X 0 i , Ŷ 0 i ) ) . ( 4 ) We set LŜu , t ( ✓t ) = 0 for t = 0 . For a fixed weight w 2 [ 0 , 1 ] , the total empirical risk can be defined as the following linear combination of LSl ( ✓t ) and LŜu , t ( ✓t ) : LSl , Ŝu , t ( ✓t ) : = wLSl ( ✓t ) + ( 1 w ) LŜu , t ( ✓t ) . ( 5 ) An SSL algorithm can be characterized by a randomized map from the labelled and unlabelled training data Sl , Su to a model parameter ✓ according to a conditional distribution P✓|Sl , Su . Then at each iteration t , we can use the sequence of conditional distributions { P✓k|Sl , Su } tk=0 with P✓0|Sl , Su = P✓0|Sl 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 ✓t and the empirical risk on the training data : gent ( PZ , PX , { P✓k|Sl , Su } t k=0 , { f✓k } t 1k=0 ) : = E [ LPZ ( ✓t ) LSl , Ŝu , t ( ✓t ) ] ( 6 ) = w ✓ E✓t [ EZ [ l ( ✓t , Z ) | ✓t ] ] 1 n nX i=1 E✓t , Zi [ l ( ✓t , Zi ) ] ◆ + ( 1 w ) ✓ E✓t [ EZ [ l ( ✓t , Z ) | ✓t ] ] 1 m tmX i= ( t 1 ) m+1 E✓t , X0i , Ŷ 0i [ l ( ✓t , ( X 0 i , Ŷ 0 i ) ) ] ◆ . ( 7 ) When t = 0 and w = 1 , the definition of the generalization error above reduces to that of vanilla supervised learning . The generalization error gent is used to measure the extent to which the iterative learning algorithm overfits the training data at the t-th iteration . Instead of focusing on the total generalization error induced during the entire process , we are more interested in the following questions . How does gent evolve as the iteration count t increases ? Do the unlabelled data examples in Su help to improve the generalization error ? 3 PRELIMINARIES . 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 generalization error gent for any t 2 [ 0 : ⌧ ] in terms of the mutual information between input data samples ( either labelled or pseudolabelled ) and the output model parameter ✓t , as well as the KL-divergence between the underlying data distributions and the joint distribution of feature vectors and pseudo-labels . We denote an R-sub-Gaussian random variable L 2 R ( Vershynin , 2018 ) as L ⇠ subG ( R ) . Furthermore , let us recall the following non-standard information quantities . Definition 1 . For arbitrary random variables X , Y and U , define the disintegrated mutual information ( Negrea et al. , 2019 ; Haghifam et al. , 2020 ) between X and Y given U as IU ( X ; Y ) : = D ( PX , Y |UkPX|U ⌦ PY |U ) , and the disintegrated KL-divergence between PX and PY given U as DU ( PXkPY ) : = D ( PX|UkPY |U ) . These are ( U ) -measurable random variables . It follows immediately that the conditional mutual information I ( X ; Y |U ) = EU [ IU ( X ; Y ) ] and the conditional KL-divergence D ( PX|UkPY |U |PU ) = EU [ DU ( PXkPY ) ] . Let ✓ ( t ) = ( ✓0 , . . . , ✓t ) for any t 2 [ 0 : ⌧ ] . In iterative SSL , we can upper bound the generalization error as shown in Theorem 1 to follow by applying the law of total expectation . Theorem 1 ( Generalization error upper bound for iterative SSL ) . Suppose l ( ✓ , Z ) ⇠ subG ( R ) under Z ⇠ PZ for all ✓ 2 ⇥ , then for any t 2 [ 0 : ⌧ ] , gent ( PZ , PX , { P✓k|Sl , Su } t k=0 , { f✓k } t 1k=0 ) w n nX i=1 E✓ ( t 1 ) hp 2R2I✓ ( t 1 ) ( ✓t ; Zi ) i + 1 w m tmX i= ( t 1 ) m+1 E✓ ( t 1 ) hq 2R2 I✓ ( t 1 ) ( ✓t ; X 0 i , Ŷ 0 i ) +D✓ ( t 1 ) ( PX0i , Ŷ 0i kPZ ) i . ( 8 ) The proof of Theorem 1 is provided in Appendix B , in which we provide a general version of upper bound not only applicable for sub-Gaussian loss functions . Compared to Bu et al . ( 2020 , Theorem 1 ) and Wu et al . ( 2020 , Theorem 1 ) , this bound focuses on the generalization error at each iteration during the learning process , which depends on the disintegrated mutual information and the disintegrated KL-divergence conditioned on the previous outputs . Intuitively , in the upper bound in Theorem 1 , the mutual information between the individual input data sample Zi and the output model parameter ✓t measures the extent to which the algorithm is sensitive to the input data , and the KL-divergence between the underlying PZ and pseudo-labelled distribution PX0i , Ŷ 0i measures how well the algorithm generalizes to the true data distribution . As n ! 1 and m ! 1 , we show that the disintegrated mutual information I✓ ( t 1 ) ( ✓t ; X 0i , Ŷ 0i ) tends to 0 ( in probability ) , which means that there are sufficient training data such that the algorithm can generalize well . On the other hand , the impact on the generalization error of pseudo-labelling is reflected in the KL-divergence D✓ ( t 1 ) ( PX0i , Ŷ 0i kPZ ) and this term does not necessarily vanish as n , m ! 1 . We quantify this precisely in Remark 1 in Section 4 . In iterative learning algorithms , it is usually difficult to directly calculate the mutual information and KL-divergence between the input and the final output ( Paninski , 2003 ; Nguyen et al. , 2010 ; McAllester & Stratos , 2020 ) . However , by applying the law of total expectation and conditioning the information-theoretic quantities on the output model parameters ✓ ( t 1 ) = { ✓1 , . . . , ✓t 1 } from previous iterations , we are able to calculate the upper bound iteratively . In the next section , we apply the iterated generalization error bound to a classification problem under a specific generative model—the bGMM . This simple model allows us to derive a tractable upper bound on the generalization error as a function of iteration number t that we can compute numerically .

Summary: The paper considers a popular approach to semi-supervised learning based on iterative pseudo-labeling the unlabelled data and refining the model parameters thereafter. The paper is supported theoretically for the case of the Gaussian mixture model establishing a generalization error bound based on the KL divergence between the pseudo-labeled and true data distributions. The paper is also supported empirically for binary classification examples coming from CIFAR10 and MNIST datasets