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 proposes two modifications for the MixMatch method [1] and achieves improved accuracy on a range of semi-supervised benchmarks. The first modification enforces the distribution of predicted labels to match the distribution of labeled data. The second modification is adding a learned data augmentation strategy, and adapting the method to work with strong data augmentation. The final method is titled ReMixMatch, and improves significantly over MixMatch, especially in low-data regime.