ReMixMatch: Semi-Supervised Learning with Distribution Matching and Augmentation Anchoring

1 INTRODUCTION . Semi-supervised learning ( SSL ) provides a means of leveraging unlabeled data to improve a model ’ s performance when only limited labeled data is available . This can enable the use of large , powerful models when labeling data is expensive or inconvenient . Research on SSL has produced a diverse collection of approaches , including consistency regularization ( Sajjadi et al. , 2016 ; Laine & Aila , 2017 ) which encourages a model to produce the same prediction when the input is perturbed and entropy minimization ( Grandvalet & Bengio , 2005 ) which encourages the model to output highconfidence predictions . The recently proposed “ MixMatch ” algorithm ( Berthelot et al. , 2019 ) combines these techniques in a unified loss function and achieves strong performance on a variety of image classification benchmarks . In this paper , we propose two improvements which can be readily integrated into MixMatch ’ s framework . First , we introduce “ distribution alignment ” , which encourages the distribution of a model ’ s aggregated class predictions to match the marginal distribution of ground-truth class labels . This concept was introduced as a “ fair ” objective by Bridle et al . ( 1992 ) , where a related loss term was shown to arise from the maximization of mutual information between model inputs and outputs . After reviewing this theoretical framework , we show how distribution alignment can be straightforwardly added to MixMatch by modifying the “ guessed labels ” using a running average of model predictions . Second , we introduce “ augmentation anchoring ” , which replaces the consistency regularization component of MixMatch . For each given unlabeled input , augmentation anchoring first generates a weakly augmented version ( e.g . using only a flip and a crop ) and then generates multiple strongly augmented versions . The model ’ s prediction for the weakly-augmented input is treated as the basis of the guessed label for all of the strongly augmented versions . To generate strong augmentations , we introduce a variant of AutoAugment ( Cubuk et al. , 2018 ) based on control theory which we dub “ CTAugment ” . Unlike AutoAugment , CTAugment learns an augmentation policy alongside model training , making it particularly convenient in SSL settings . We call our improved algorithm “ ReMixMatch ” and experimentally validate it on a suite of standard SSL image benchmarks . ReMixMatch achieves state-of-the-art accuracy across all labeled data amounts , for example achieving an accuracy of 93.73 % with 250 labels on CIFAR-10 compared to the previous state-of-the-art of 88.92 % ( and compared to 96.09 % for fully-supervised classification with 50,000 labels ) . We also push the limited-data setting further than ever before , ultimately achieving a median of 84.92 % accuracy with only 40 labels ( just 4 labels per class ) on CIFAR-10 . To quantify the impact of our proposed improvements , we carry out an extensive ablation study to measure the impact of our improvements to MixMatch . Finally , we release all of our models and code to facilitate future work on semi-supervised learning . 2 BACKGROUND . The goal of a semi-supervised learning algorithm is to learn from unlabeled data in a way that improves performance on labeled data . Typical ways of achieving this include training against “ guessed ” labels for unlabeled data or optimizing a heuristically-motivated objective that does not rely on labels . This section reviews the semi-supervised learning methods relevant to ReMixMatch , with a particular focus on the components of the MixMatch algorithm upon which we base our work . Consistency Regularization Many SSL methods rely on consistency regularization to enforce that the model output remains unchanged when the input is perturbed . First proposed in ( Bachman et al. , 2014 ) , ( Sajjadi et al. , 2016 ) and ( Laine & Aila , 2017 ) , this approach was referred to as “ Regularization With Stochastic Transformations and Perturbations ” and the “ Π-Model ” respectively . While some work perturbs adversarially ( Miyato et al. , 2018 ) or using dropout ( Laine & Aila , 2017 ; Tarvainen & Valpola , 2017 ) , the most common perturbation is to apply domain-specific data augmentation ( Laine & Aila , 2017 ; Sajjadi et al. , 2016 ; Berthelot et al. , 2019 ; Xie et al. , 2019 ) . The loss function used to measure consistency is typically either the mean-squared error ( Laine & Aila , 2017 ; Tarvainen & Valpola , 2017 ; Sajjadi et al. , 2016 ) or cross-entropy ( Miyato et al. , 2018 ; Xie et al. , 2019 ) between the model ’ s output for a perturbed and non-perturbed input . Entropy Minimization Grandvalet & Bengio ( 2005 ) argues that unlabeled data should be used to ensure that classes are well-separated . This can be achieved by encouraging the model ’ s output distribution to have low entropy ( i.e. , to make “ high-confidence ” predictions ) on unlabeled data . For example , one can explicitly add a loss term to minimize the entropy of the model ’ s predicted class distribution on unlabeled data ( Grandvalet & Bengio , 2005 ; Miyato et al. , 2018 ) . Related to this idea are “ self-training ” methods ( McLachlan , 1975 ; Rosenberg et al. , 2005 ) such as Pseudo-Label ( Lee , 2013 ) that use the predicted class on an unlabeled input as a hard target for the same input , which implicitly minimizes the entropy of the prediction . Standard Regularization Outside of the setting of SSL , it is often useful to regularize models in the over-parameterized regime . This regularization can often be applied both when training on labeled and unlabeled data . For example , standard “ weight decay ” ( Hinton & van Camp , 1993 ) where the L2 norm of parameters is minimized is often used alongside SSL techniques . Similarly , powerful MixUp regularization ( Zhang et al. , 2017 ) which trains a model on linear interpolants of inputs and labels has recently been applied to SSL ( Berthelot et al. , 2019 ; Verma et al. , 2019 ) . Other Approaches The three aforementioned categories of SSL techniques does not cover the full literature on semi-supervised learning . For example , there is a significant body of research on “ transductive ” or graph-based semi-supervised learning techniques which leverage the idea that unlabeled datapoints should be assigned the label of a labeled datapoint if they are sufficiently similar ( Gammerman et al. , 1998 ; Joachims , 2003 ; 1999 ; Bengio et al. , 2006 ; Liu et al. , 2018 ) . Since our work does not involve these ( or other ) approaches to SSL , we will not discuss them further . A more substantial overview of SSL methods is available in ( Chapelle et al. , 2006 ) . 2.1 MIXMATCH . MixMatch ( Berthelot et al. , 2019 ) unifies several of the previously mentioned SSL techniques . The algorithm works by generating “ guessed labels ” for each unlabeled example , and then using fullysupervised techniques to train on the original labeled data along with the guessed labels for the unlabeled data . This section reviews the necessary details of MixMatch ; see ( Berthelot et al. , 2019 ) for a full definition . Let X = { ( xb , pb ) : b ∈ ( 1 , . . . , B ) } be a batch of labeled data and their corresponding one-hot labels representing one of L classes and let x̂b be augmented versions of these labeled examples . Similarly , let U = { ub : b ∈ ( 1 , . . . , B ) } be a batch of unlabeled examples . Finally , let pmodel ( y | x ; θ ) be the predicted class distribution produced by the model for input x. MixMatch first produces K weakly augmented versions of each unlabeled datapoint ûb , k for k ∈ { 1 , . . . , K } . Then , it generates a “ guessed label ” qb for each ub by computing the average prediction q̄b across the K augmented versions : q̄b = 1K ∑ k pmodel ( y | ûb , k ; θ ) . The guessed label distribution is then sharpened by adjusting its temperature ( i.e . raising all probabilities to a power of 1⁄T and renormalizing ) . Finally , pairs of examples ( x1 , p1 ) , ( x2 , p2 ) from the combined set of labeled examples and unlabeled examples with label guesses are fed into the MixUp ( Zhang et al. , 2017 ) algorithm to compute examples ( x′ , p′ ) where x′ = λx1 + ( 1 − λ ) x2 for λ ∼ Beta ( α , α ) , and similarly for p′ . Given these mixed-up examples , MixMatch performs standard fully-supervised training with minor modifications . A standard cross-entropy loss is used for labeled data , whereas the loss for unlabeled data is computed using a mean square error ( i.e . the Brier score ( Brier , 1950 ) ) and is weighted with a hyperparameter λU . The terms K ( number of augmentations ) , T ( sharpening temperature ) , α ( MixUp Beta parameter ) , and λU ( unlabeled loss weight ) are MixMatch ’ s hyperparameters . For augmentation , shifting and flipping was used for the CIFAR-10 , CIFAR-100 , and STL-10 datasets , and shifting alone was used for SVHN . 3 REMIXMATCH . Having introduced MixMatch , we now turn to the two improvements we propose in this paper : Distribution alignment and augmentation anchoring . For clarity , we describe how we integrate them into the base MixMatch algorithm ; the full algorithm for ReMixMatch is shown in algorithm 1 . 3.1 DISTRIBUTION ALIGNMENT . Our first contribution is distribution alignment , which enforces that the aggregate of predictions on unlabeled data matches the distribution of the provided labeled data . This general idea was first introduced over 25 years ago ( Bridle et al. , 1992 ) , but to the best of our knowledge is not used in modern SSL techniques . A schematic of distribution alignment can be seen in fig . 1 . After reviewing and extending the theory , we describe how it can be straightforwardly included in ReMixMatch . Algorithm 1 ReMixMatch algorithm for producing a collection of processed labeled examples and processed unlabeled examples with label guesses ( cf . Berthelot et al . ( 2019 ) Algorithm 1 . ) 1 : Input : Batch of labeled examples and their one-hot labels X = { ( xb , pb ) : b ∈ ( 1 , . . . , B ) } , batch of unlabeled examples U = { ub : b ∈ ( 1 , . . . , B ) } , sharpening temperature T , number of augmentations K , Beta distribution parameter α for MixUp . 2 : for b = 1 to B do 3 : x̂b = StrongAugment ( xb ) // Apply strong data augmentation to xb 4 : ûb , k = StrongAugment ( ub ) ; k ∈ { 1 , . . . , K } // Apply strong data augmentation K times to ub 5 : ũb = WeakAugment ( ub ) // Apply weak data augmentation to ub 6 : qb = pmodel ( y | ũb ; θ ) // Compute prediction for weak augmentation of ub 7 : qb = Normalize ( qb × p ( y ) / p̃ ( y ) ) // Apply distribution alignment 8 : qb = Normalize ( q 1/T b ) // Apply temperature sharpening to label guess 9 : end for 10 : X̂ = ( ( x̂b , pb ) ; b ∈ ( 1 , . . . , B ) ) // Augmented labeled examples and their labels 11 : Û1 = ( ( ûb,1 , qb ) ; b ∈ ( 1 , . . . , B ) ) // First strongly augmented unlabeled example and guessed label 12 : Û = ( ( ûb , k , qb ) ; b ∈ ( 1 , . . . , B ) , k ∈ ( 1 , . . . , K ) ) // All strongly augmented unlabeled examples 13 : Û = Û ∪ ( ( ũb , qb ) ; b ∈ ( 1 , . . . , B ) ) // Add weakly augmented unlabeled examples 14 : W = Shuffle ( Concat ( X̂ , Û ) ) // Combine and shuffle labeled and unlabeled data 15 : X ′ = ( MixUp ( X̂i , Wi ) ; i ∈ ( 1 , . . . , |X̂ | ) ) // Apply MixUp to labeled data and entries fromW 16 : U ′ = ( MixUp ( Ûi , Wi+|X̂ | ) ; i ∈ ( 1 , . . . , |Û | ) ) // Apply MixUp to unlabeled data and the rest ofW 17 : return X ′ , U ′ , Û1

Summary: This paper presents ReMixMatch an improved version of MixMatch. The main contributions are the distribution alignment and the augmentation anchoring. Distribution alignment rescales the predictions based on the difference between the model marginals and the ground truth running average estimation. Augmentation anchoring instead of computing the guessed probabilities on unlabelled data as the average probabilities on transformed samples (as in MixMatch), it considers as guessed labels the average probabilities obtained from weak transformations (flip+crop) even when using stronger transformations (Autoaugment like).