There Are Many Consistent Explanations of Unlabeled Data: Why You Should Average

1 INTRODUCTION . Recent advances in deep unsupervised learning , such as generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) , have led to an explosion of interest in semi-supervised learning . Semisupervised methods make use of both unlabeled and labeled training data to improve performance over purely supervised methods . Semi-supervised learning is particularly valuable in applications such as medical imaging , where labeled data may be scarce and expensive ( Oliver et al. , 2018 ) . Currently the best semi-supervised results are obtained by consistency-enforcing approaches ( Bachman et al. , 2014 ; Laine and Aila , 2017 ; Tarvainen and Valpola , 2017 ; Miyato et al. , 2017 ; Park et al. , 2017 ) . These methods use unlabeled data to stabilize their predictions under input or weight perturbations . Consistency-enforcing methods can be used at scale with state-of-the-art architectures . For example , the recent Mean Teacher ( Tarvainen and Valpola , 2017 ) model has been used with the Shake-Shake ( Gastaldi , 2017 ) architecture and has achieved the best semi-supervised performance on the consequential CIFAR benchmarks . This paper is about conceptually understanding and improving consistency-based semi-supervised learning methods . Our approach can be used as a guide for exploring how loss geometry interacts with training procedures in general . We provide several novel observations about the training objective and optimization trajectories of the popular ⇧ ( Laine and Aila , 2017 ) and Mean Teacher ( Tarvainen and Valpola , 2017 ) consistency-based models . Inspired by these findings , we propose to improve SGD solutions via stochastic weight averaging ( SWA ) ( Izmailov et al. , 2018 ) , a recent method that averages weights of the networks corresponding to different training epochs to obtain a single model with improved generalization . On a thorough empirical study we show that this procedure achieves the best known semi-supervised results on consequential benchmarks . In particular : • We show in Section 3.1 that a simplified ⇧ model implicitly regularizes the norm of the Jacobian of the network outputs with respect to both its inputs and its weights , which in turn encourages flatter solutions . Both the reduced Jacobian norm and flatness of solutions have been related to generalization in the literature ( Sokolić et al. , 2017 ; Novak et al. , 2018 ; Chaudhari et al. , 2016 ; Schmidhuber and Hochreiter , 1997 ; Keskar et al. , 2017 ; Izmailov et al. , 2018 ) . Interpolating between the weights corresponding to different epochs of training we demonstrate that the solutions of ⇧ and Mean Teacher models are indeed flatter along these directions ( Figure 1b ) . • In Section 3.2 , we compare the training trajectories of the ⇧ , Mean Teacher , and supervised models and find that the distances between the weights corresponding to different epochs are much larger for the consistency based models . The error curves of consistency models are also wider ( Figure 1b ) , which can be explained by the flatness of the solutions discussed in section 3.1 . Further we observe that the predictions of the SGD iterates can differ significantly between different iterations of SGD . • We observe that for consistency-based methods , SGD does not converge to a single point but continues to explore many solutions with high distances apart . Inspired by this observation , we propose to average the weights corresponding to SGD iterates , or ensemble the predictions of the models corresponding to these weights . Averaging weights of SGD iterates compensates for larger steps , stabilizes SGD trajectories and obtains a solution that is centered in a flat region of the loss ( as a function of weights ) . Further , we show that the SGD iterates correspond to models with diverse predictions – using weight averaging or ensembling allows us to make use of the improved diversity and obtain a better solution compared to the SGD iterates . In Section 3.3 we demonstrate that both ensembling predictions and averaging weights of the networks corresponding to different training epochs significantly improve generalization performance and find that the improvement is much larger for the ⇧ and Mean Teacher models compared to supervised training . We find that averaging weights provides similar or improved accuracy compared to ensembling , while offering the computational benefits and convenience of working with a single model . Thus , we focus on weight averaging for the remainder of the paper . • Motivated by our observations in Section 3 we propose to apply Stochastic Weight Averaging ( SWA ) ( Izmailov et al. , 2018 ) to the ⇧ and Mean Teacher models . Based on our results in Section 3.3 we propose several modifications to SWA in Section 4 . In particular , we propose fast-SWA , which ( 1 ) uses a learning rate schedule with longer cycles to increase the distance between the weights that are averaged and the diversity of the corresponding predictions ; and ( 2 ) averages weights of multiple networks within each cycle ( while SWA only averages weights corresponding to the lowest values of the learning rate within each cycle ) . In Section 5 , we show that fast-SWA converges to a good solution much faster than SWA . • Applying weight averaging to the ⇧ and Mean Teacher models we improve the best reported results on CIFAR-10 for 1k , 2k , 4k and 10k labeled examples , as well as on CIFAR-100 with 10k labeled examples . For example , we obtain 5.0 % error on CIFAR-10 with only 4k labels , improving the best result reported in the literature ( Tarvainen and Valpola , 2017 ) by 1.3 % . We also apply weight averaging to a state-of-the-art domain adaptation technique ( French et al. , 2018 ) closely related to the Mean Teacher model and improve the best reported results on domain adaptation from CIFAR-10 to STL from 19.9 % to 16.8 % error . • We release our code at https : //github.com/benathi/fastswa-semi-sup 2 BACKGROUND . 2.1 CONSISTENCY BASED MODELS . We briefly review semi-supervised learning with consistency-based models . This class of models encourages predictions to stay similar under small perturbations of inputs or network parameters . For instance , two different translations of the same image should result in similar predicted probabilities . The consistency of a model ( student ) can be measured against its own predictions ( e.g . ⇧ model ) or predictions of a different teacher network ( e.g . Mean Teacher model ) . In both cases we will say a student network measures consistency against a teacher network . Consistency Loss In the semi-supervised setting , we have access to labeled data DL = { ( xLi , yLi ) } NL i=1 , and unlabeled data DU = { xUi } NU i=1 . Given two perturbed inputs x0 , x00 of x and the perturbed weights w0f and w 0 g , the consistency loss penalizes the difference between the student ’ s predicted probablities f ( x0 ; w0f ) and the teacher ’ s g ( x00 ; w0g ) . This loss is typically the Mean Squared Error or KL divergence : ` MSEcons ( wf , x ) = kf ( x0 ; w0f ) g ( x00 , w0g ) k2 or ` KLcons ( wf , x ) = KL ( f ( x0 ; w0f ) ||g ( x00 , w0g ) ) . ( 1 ) The total loss used to train the model can be written as L ( wf ) = X ( x , y ) 2DL ` CE ( wf , x , y ) | { z } LCE + X x2DL [ DU ` cons ( wf , x ) | { z } Lcons , ( 2 ) where for classification LCE is the cross entropy between the model predictions and supervised training labels . The parameter > 0 controls the relative importance of the consistency term in the overall loss . ⇧ Model The ⇧ model , introduced in Laine and Aila ( 2017 ) and Sajjadi et al . ( 2016 ) , uses the student model f as its own teacher . The data ( input ) perturbations include random translations , crops , flips and additive Gaussian noise . Binary dropout ( Srivastava et al. , 2014 ) is used for weight perturbation . Mean Teacher Model The Mean Teacher model ( MT ) proposed in Tarvainen and Valpola ( 2017 ) uses the same data and weight perturbations as the ⇧ model ; however , the teacher weights wg are the exponential moving average ( EMA ) of the student weights wf : wkg = ↵ · wk 1g + ( 1 ↵ ) · wkf . The decay rate ↵ is usually set between 0.9 and 0.999 . The Mean Teacher model has the best known results on the CIFAR-10 semi-supervised learning benchmark ( Tarvainen and Valpola , 2017 ) . Other Consistency-Based Models Temporal Ensembling ( TE ) ( Laine and Aila , 2017 ) uses an exponential moving average of the student outputs as the teacher outputs in the consistency term for training . Another approach , Virtual Adversarial Training ( VAT ) ( Miyato et al. , 2017 ) , enforces the consistency between predictions on the original data inputs and the data perturbed in an adversarial direction x0 = x + ✏radv , where radv = arg maxr : krk=1 KL [ f ( x , w ) kf ( x + ⇠r , w ) ] . 3 UNDERSTANDING CONSISTENCY-ENFORCING MODELS . In Section 3.1 , we study a simplified version of the ⇧ model theoretically and show that it penalizes the norm of the Jacobian of the outputs with respect to inputs , as well as the eigenvalues of the Hessian , both of which have been related to generalization ( Sokolić et al. , 2017 ; Novak et al. , 2018 ; Dinh et al. , 2017a ; Chaudhari et al. , 2016 ) . In Section 3.2 we empirically study the training trajectories of the ⇧ and MT models and compare them to the training trajectories in supervised learning . We show that even late in training consistency-based methods make large training steps , leading to significant changes in predictions on test . In Section 3.3 we show that averaging weights or ensembling predictions of the models proposed by SGD at different training epochs can lead to substantial gains in accuracy and that these gains are much larger for ⇧ and MT than for supervised training . 3.1 SIMPLIFIED ⇧ MODEL PENALIZES LOCAL SHARPNESS . Penalization of the input-output Jacobian norm . Consider a simple version of the ⇧ model , where we only apply small additive perturbations to the student inputs : x0 = x + ✏z , z ⇠ N ( 0 , I ) with ✏ ⌧ 1 , and the teacher input is unchanged : x00 = x.1 Then the consistency loss ` cons ( Eq . 1 ) becomes ` cons ( w , x , ✏ ) = kf ( w , x + ✏z ) f ( w , x ) k2 . Consider the estimator Q̂ = lim✏ ! 0 1 ✏2 1 m Pm i=1 ` cons ( w , xi , ✏ ) . We show in Section A.5 that E [ Q̂ ] = Ex [ kJxk2F ] and Var [ Q̂ ] = 1 m ✓ Var [ kJxk2F ] + 2E [ kJTx Jxk2F ] ◆ , where Jx is the Jacobian of the network ’ s outputs with respect to its inputs evaluated at x , k · kF represents Frobenius norm , and the expectation Ex is taken over the distribution of labeled and 1This assumption can be relaxed to x00 = x+ ✏z2 z2 ⇠ N ( 0 , I ) without changing the results of the analysis , since ✏ ( z1 z2 ) = 2✏z̄ with z̄ ⇠ N ( 0 , I ) . unlabeled data . That is , Q̂ is an unbiased estimator of Ex [ kJxk2F ] with variance controlled by the minibatch size m. Therefore , the consistency loss implicitly penalizes Ex [ kJxk2F ] . The quantity ||Jx||F has been related to generalization both theoretically ( Sokolić et al. , 2017 ) and empirically ( Novak et al. , 2018 ) . For linear models f ( x ) = Wx , penalizing ||Jx||F exactly corresponds to weight decay , also known as L2 regularization , since for linear models Jx = W , and kWk2F = kvec ( W ) k22 . Penalizing Ex [ kJxk2F ] is also closely related to the graph based ( manifold ) regularization in Zhu et al . ( 2003 ) which uses the graph Laplacian to approximate Ex [ krMfk2 ] for nonlinear models , making use of the manifold structure of unlabeled data . Isotropic perturbations investigated in this simplified ⇧ model will not in general lie along the data manifold , and it would be more pertinent to enforce consistency to perturbations sampled from the space of natural images . In fact , we can interpret consistency with respect to standard data augmentations ( which are used in practice ) as penalizing the manifold Jacobian norm in the same manner as above . See Section A.5 for more details . Penalization of the Hessian ’ s eigenvalues . Now , instead of the input perturbation , consider the weight perturbation w0 = w + ✏z . Similarly , the consistency loss is an unbiased estimator for Ex [ kJwk2F ] , where Jw is the Jacobian of the network outputs with respect to the weights w. In Section A.6 we show that for the MSE loss , the expected trace of the Hessian of the loss Ex [ tr ( H ) ] can be decomposed into two terms , one of which is Ex [ kJwk2F ] . As minimizing the consistency loss of a simplified ⇧ model penalizes Ex [ kJwk2F ] , it also penalizes Ex [ tr ( H ) ] . As pointed out in Dinh et al . ( 2017a ) and Chaudhari et al . ( 2016 ) , the eigenvalues of H encode the local information about sharpness of the loss for a given solution w. Consequently , the quantity tr ( H ) which is the sum of the Hessian eigenvalues is related to the notion of sharp and flat optima , which has recently gained attention as a proxy for generalization performance ( see e.g . Schmidhuber and Hochreiter , 1997 ; Keskar et al. , 2017 ; Izmailov et al. , 2018 ) . Thus , based on our analysis , the consistency loss in the simplified ⇧ model encourages flatter solutions .

Summary: This paper proposes to apply Stochastic Weight Averaging to the semi-supervised learning context. It makes an interesting argument that the semi-supervised MT/Pi models are especially amenable to SWA since they are empirically observed to traverse a large flat region of the weight space during the later stages of training. To speed up training, the authors propose fast-SWA.