Source-linked AI summary

Exponential Moving Average Normalization for Self-supervised and Semi-supervised Learning

Zhaowei Cai, Avinash Ravichandran, Subhransu Maji, Charless Fowlkes, Zhuowen Tu, Stefano Soatto

arXiv:2101.08482v2cs.LGcs.AIcs.CV

TL;DR

Student-teacher learning can exploit unlabeled data, but standard BN introduces batch-dependent statistics and can mismatch the EMA-updated teacher parameters. EMAN replaces teacher BN statistics with exponential moving averages of student statistics, improving self- and semi-supervised baselines across settings.

  • Problem

    Student-teacher methods use unlabeled data, but standard BN creates cross-sample dependency and can mismatch the teacher’s EMA-averaged parameters.

  • Method

    EMAN replaces standard teacher BN with statistics exponentially moving averaged from the student’s BN statistics, without batch-wise statistics computation in the teacher.

  • Results

    EMAN improves self-supervised MoCo/BYOL by 4-6/1-2 points and semi-supervised FixMatch by about 7/2 points for 1%/10% labels on ImageNet.

  • Takeaways & Limitations

    Improvements are consistent across methods, architectures, training durations, and datasets, supporting EMAN as a generally applicable technique.

Abstract

from arXiv · show

We present a plug-in replacement for batch normalization (BN) called exponential moving average normalization (EMAN), which improves the performance of existing student-teacher based self- and semi-supervised learning techniques. Unlike the standard BN, where the statistics are computed within each batch, EMAN, used in the teacher, updates its statistics by exponential moving average from the BN statistics of the student. This design reduces the intrinsic cross-sample dependency of BN and enhances the generalization of the teacher. EMAN improves strong baselines for self-supervised learning by 4-6/1-2 points and semi-supervised learning by about 7/2 points, when 1%/10% supervised labels are available on ImageNet. These improvements are consistent across methods, network architectures, training duration, and datasets, demonstrating the general effectiveness of this technique. The code is available at https://github.com/amazon-research/exponential-moving-average-normalization.

1. Introduction

Student-teacher methods address learning with limited or no annotation, but standard BN can introduce cross-sample dependency in the teacher. EMA-teacher methods average student parameters over training to improve generalization.

  • Limited annotations motivate semi-supervised and self-supervised representation learning, including student-teacher approaches that use teacher outputs to guide students.
  • EMA-teacher updates the teacher by exponential moving average of student parameters, producing a temporally ensembled teacher.
  • Standard BN makes a sample’s output depend on all other samples in its batch, creating cross-sample information leakage.
  • Existing responses include layer normalization, ShuffleBN, and synchronized BN to reduce or avoid BN-related information leakage.

2. Model parameter mismatch. In the teacher network, its

The paper identifies a mismatch in EMA-teacher networks: parameters are temporally averaged, while BN statistics are collected from the current batch. EMAN aligns both through exponential averaging.

  • Teacher parameters are averaged from previous student iterations, whereas batch-wise BN statistics are collected immediately at the current iteration.
  • EMAN updates teacher means and variances by exponential moving average from student BN statistics, matching the parameter-update mechanism.
  • EMAN improves MoCo and BYOL by 4-6/1-2 points and FixMatch by about 7/2 points for 1%/10% labels on ImageNet.
  • These improvements remain consistent across methods, architectures, training durations, and datasets.

2. Related Work

Related work spans student-teacher and consistency-based semi-supervised learning, self-supervised representation learning, and normalization methods designed for optimization or reduced information leakage.

  • Semi-supervised methods use unlabeled data through pseudo-labeling, temporal ensembling, or consistency under input and model perturbations.
  • Mean Teacher introduced the EMA-teacher framework, while FixMatch enforces consistency between weakly and strongly augmented inputs.
  • Self-supervised vision methods evolved from proxy tasks toward contrastive instance discrimination, with MoCo and SimCLR narrowing supervised-learning gaps in some domains.
  • Student-teacher frameworks vary in whether teachers are pretrained, identical, ensembled, or more complex than students.
  • Normalization alternatives include LN, IN, GN, MABN, SyncBN, and ShuffleBN, addressing convergence, batch-size stability, or BN information leakage.

3. Preliminaries

The EMA-teacher averages student parameters to stabilize consistency-based learning, but standard BN behaves differently during training and inference and depends on batch statistics. Its operation normalizes each sample with batch-wise statistics and learned affine parameters.

  • EMA-teacher: The EMA-teacher updates teacher parameters from student parameters using momentum close to 1, such as 0.999.
  • EMA-teacher: The teacher receives no gradient updates and is discarded after training, while the student learns with standard SGD.
  • EMA-teacher: Temporal averaging smooths student checkpoints and can stabilize training trajectories in consistency-based semi- and self-supervised learning.
  • Batch Normalization: Standard BN is used in both student and teacher despite the teacher’s special parameter-update rule.
  • Batch Normalization: During training, BN computes batch mean and variance, then normalizes every sample using those statistics before applying learned scale and shift.
  • Batch Normalization: At inference, batch-wise statistics are replaced by exponentially averaged proxy statistics so outputs do not depend on other inputs in the batch.

4. Exponential Moving Average Normalization

EMAN replaces the teacher’s batch normalization with exponentially averaged statistics from the student, addressing cross-sample dependence and parameter-statistics mismatch. Applied across FixMatch, MoCo, and BYOL, it provides a general, simple normalization alternative.

  • 4. Exponential Moving Average Normalization: Teacher batch normalization can make pseudo-labels cross-sample dependent and mismatch the teacher’s EMA-averaged parameters.The mismatch may cause non-smoothness in parameter space.
  • 4. Exponential Moving Average Normalization: EMAN uses exponentially moving averaged student normalization statistics for the teacher while the student continues using batch normalization.The teacher’s mean and variance are updated in the same way as other exponentially averaged quantities.
  • 4. Exponential Moving Average Normalization: EMAN removes cross-sample dependence from the teacher and aligns its normalization factors with the EMA-updated teacher parameters.The resulting teacher normalization is no longer dependent on batch statistics.
  • 4.1. Applications: EMAN was applied to FixMatch, MoCo, and BYOL with a few lines of code change.The methods use different student-teacher formulations, including pseudo-labeling, contrastive learning, and regression-based self-supervision.
  • 4.1. Applications: EMAN surpasses standard BN and the ShuffleBN or SyncBN counterparts across FixMatch, MoCo, and BYOL without cross-GPU communication or synchronization.The cited comparison presents EMAN as a simpler and more efficient alternative across the three techniques.

5. Experiments

Experiments evaluate EMAN across normalization choices, self-supervised and semi-supervised methods, label regimes, architectures, training durations, and downstream tasks. EMAN consistently improves strong baselines, especially when labeled data are scarce, while also outperforming alternative normalization choices and state-of-the-art comparisons.

  • Experimental setup: ImageNet experiments use ResNet-50 by default and evaluate EMAN with FixMatch, MoCo, and BYOL under the corresponding baseline settings.The study also varies architectures, training duration, label percentages, and evaluation protocols.
  • Normalization comparison: EMAN outperforms standard BN and technique-specific ShuffleBN or SyncBN across FixMatch, MoCo, and BYOL.ShuffleBN and SyncBN perform well within their original techniques but generalize poorly across methods, whereas EMAN achieves the best results across all three.
  • Self-supervised evaluation: 4-5 points of top-1 accuracy are gained with 1% labels, while 1-2 points are gained with 10% labels in linear evaluation.The pattern is consistent across MoCo and BYOL, ResNet-50 and 2×-width ResNet-50, and 200- versus 800-epoch training.
  • State-of-the-art comparison: 57.4 top-1 accuracy is achieved by MoCo-EMAN with 1% ImageNet labels, while BYOL-EMAN reaches 55.1 with 1% labels after 200 pretraining epochs.The reported comparison notes that BYOL used 1000 epochs, whereas BYOL-EMAN used 200; the comparison is not completely fair because methods differ in training strategies.
  • Downstream evaluation: EMAN consistently improves MoCo and BYOL in kNN classification, image retrieval, and low-shot classification, including 4-6 points for MoCo and 1-3 points for BYOL.In 1000-way 1-shot ImageNet, MoCo-EMAN reaches 35.8% top-1 accuracy, 12.3 points above SwAV.

6. Conclusion

EMAN is proposed as a normalization technique for EMA-teacher-based semi- and self-supervised learning, addressing limitations of standard BN and improving performance across settings.

  • EMAN is a normalization technique for EMA-teacher-based semi- and self-supervised learning.
  • EMAN addresses cross-sample dependency and parameter mismatch associated with standard BN in EMA-teacher frameworks.
  • EMAN improvements are consistent across techniques, network architectures, training durations, and datasets.

A. Appendix

The appendix reports FixMatch accuracy with 10% labels in Figure 4.

  • Figure 4 reports FixMatch accuracy.
  • The evaluation uses 10% labels.
  • The reported results are presented as a figure.

A.1. FixMatch

The FixMatch appendix uses an EMA-teacher formulation and finds that the EMA teacher provides higher and more stable validation accuracy than the baseline during training.

  • FixMatch uses batches containing 64 labeled and 320 unlabeled images.The supervised and unsupervised loss weights are 1.0 and 10.0, respectively.
  • The EMA teacher achieves higher and more stable ImageNet validation accuracy than baseline FixMatch, especially during early epochs.
  • The observed EMA-teacher reliability motivates reformulating baseline FixMatch within an EMA-teacher framework.

A.2. MoCo

The MoCo appendix reports training instability and late-epoch accuracy decline with EMAN, while showing that EMAN changes training behavior and may prevent cheating.

  • MoCo-EMAN training is less stable initially because the whole network, including normalization statistics, is updated slowly with momentum m = 0.999.
  • MoCo-EMAN accuracy slightly decreases in the final training epochs because the learning rate becomes too small.
  • Evaluating the 90%-th epoch checkpoint or using a minimum learning rate of 0.001 produces closely matched performances.
  • EMAN makes MoCo more difficult to train, producing higher loss and lower instance discrimination accuracy than MoCo-ShuffleBN.
  • The altered training behavior suggests that EMAN may better prevent MoCo from cheating and potentially improve representation power.

A.3. BYOL

In BYOL experiments, EMAN makes the self-supervised task more difficult, reflected by higher training loss. Evaluation uses both linear and finetuning protocols with specified optimization settings and standard image preprocessing.

  • BYOL experiments use batch size 512 across eight GPUs, with weight decay 0.000001 and initial EMA momentum 0.98.The initial learning rate is 0.9 for 50 epochs and 1.8 for 200 epochs.
  • EMAN produces higher BYOL training loss, indicating a more difficult self-supervised learning task.The authors interpret this as evidence that EMAN may better prevent the model from cheating and potentially improve representation power.
  • ImageNet semi-supervised experiments sample 1% or 10% of images per class, totaling 12,820 or 128,118 images.
  • Linear evaluation uses zero weight decay, while finetuning uses weight decay 0.0001; both use momentum 0.9.Training uses random resized crops and random horizontal flips, while inference uses a centered 224×224 crop.
  • Figure 5 reports MoCo training loss and instance discrimination top-1 accuracy alongside BYOL training loss.
  • kNN evaluation uses an unweighted classifier with top k = 20, assigning each query the majority class among its neighbors.The center 224×224 crop is also used for image retrieval and low-shot classification.
Loading 2101.08482v2…