[Paper Review] ReMixMatch: Semi-Supervised Learning with Distribution Matching and Augmentation Anchoring

David Berthelot, Nicholas Carlini+5 more|arXiv (Cornell University)|Apr 30, 2020

Domain Adaptation and Few-Shot LearningComputer Science26 references338 citations

TL;DR

ReMixMatch improves semi-supervised learning by introducing distribution alignment and augmentation anchoring, which align unlabeled predictions with true label distributions and encourage consistency across strong and weak augmentations. The method achieves state-of-the-art accuracy with significantly fewer labeled examples—reaching 93.73% on CIFAR-10 using only 250 labeled examples, outperforming MixMatch with 4,000 labels.

ABSTRACT

We improve the recently-proposed ``MixMatch semi-supervised learning algorithm by introducing two new techniques: distribution alignment and augmentation anchoring. - Distribution alignment encourages the marginal distribution of predictions on unlabeled data to be close to the marginal distribution of ground-truth labels. - Augmentation anchoring} feeds multiple strongly augmented versions of an input into the model and encourages each output to be close to the prediction for a weakly-augmented version of the same input. To produce strong augmentations, we propose a variant of AutoAugment which learns the augmentation policy while the model is being trained. Our new algorithm, dubbed ReMixMatch, is significantly more data-efficient than prior work, requiring between 5 times and 16 times less data to reach the same accuracy. For example, on CIFAR-10 with 250 labeled examples we reach 93.73% accuracy (compared to MixMatch's accuracy of 93.58% with 4000 examples) and a median accuracy of 84.92% with just four labels per class.

Motivation & Objective

• To improve data efficiency in semi-supervised learning by reducing reliance on large amounts of labeled data.
• To address distribution shift in unlabeled data predictions by aligning them with true label distributions.
• To enhance model robustness by enforcing consistency between weakly and strongly augmented views of the same input.
• To develop a training-time augmentation policy learning method that improves generalization.

Proposed method

• Distribution alignment minimizes the divergence between the marginal distribution of model predictions on unlabeled data and the true label distribution.
• Augmentation anchoring uses multiple strong augmentations per input and encourages each to predict similarly to the weakly-augmented version of the same input.
• A variant of AutoAugment is trained jointly with the model to learn optimal augmentation policies during training.
• The model is trained using a combination of cross-entropy loss on labeled data and consistency regularization on unlabeled data with both weak and strong augmentations.
• The method leverages entropy minimization and distribution alignment to stabilize and improve predictions on unlabeled data.

Experimental results

Research questions

• RQ1 Can distribution alignment improve generalization in semi-supervised learning by aligning model predictions with true label distributions?
• RQ2 Does enforcing consistency between weak and strong augmentations improve model robustness and accuracy?
• RQ3 Can a training-time augmentation policy search method outperform fixed augmentation policies in semi-supervised settings?
• RQ4 To what extent can ReMixMatch reduce the labeled data requirement while maintaining or improving accuracy?

Key findings

• ReMixMatch achieves 93.73% accuracy on CIFAR-10 using only 250 labeled examples, surpassing MixMatch's 93.58% accuracy with 4,000 labeled examples.
• With just four labeled examples per class, ReMixMatch attains a median accuracy of 84.92% on CIFAR-10, demonstrating high data efficiency.
• The method reduces the labeled data requirement by 5 to 16 times compared to prior state-of-the-art methods to achieve comparable accuracy.
• Distribution alignment significantly improves prediction consistency and reduces distribution shift in unlabeled predictions.