Batch normalization

Overview

Batch normalization normalizes each incoming activation with a mean and a variance, then applies a learned scale and shift 2112.00463 §3.1 2112.00463 eq. 1. During training, the mean and variance are running averages updated with a momentum. At test time the running statistics are fixed 2112.00463 §3.1. That fixed state is what breaks under a distribution shift, and it is the hook that the variant below uses.

Coverage: only one paper card lists this concept so far. The page will grow as more cards are compiled.

Variants

Standard batch normalization (the base layer)

Training uses the statistics of the current batch. Inference uses stored running estimates, with the learned scale and shift applied after 2112.00463 §3.1 2112.00463 eq. 2 2112.00463 eq. 3. Only the statistics change in the variant below; the weights stay fixed.

Test-time statistic adaptation (DUA)

DUA keeps every learned weight fixed and replaces each BN layer's running mean and variance with estimates from an unlabeled target stream. The stored statistics start from the source values as a prior 2112.00463 §3.2. Each target sample updates the estimates once, in order [arxiv:2112.00463 §3.2]. The momentum decays with each sample, with a lower bound zeta 2112.00463 eq. 4 2112.00463 eq. 5. The variance update uses the same momentum 2112.00463 eq. 6 2112.00463 eq. 7. Each sample is expanded into a batch of 64 augmented copies, and the paper states that a batch of 8 keeps most of the benefit 2112.00463 §5.4. Card: The Norm Must Go On: Dynamic Unsupervised Domain Adaptation by Normalization.

Author-reported results on CIFAR-10C at severity 5 show mean top-1 error of 26.8% for ResNet-26, against 49.2% for the source model, using 80 samples 2112.00463 tab. 1. With Wide-ResNet-40-2, the same method reaches 12.1%, against 18.3% for the source model 2112.00463 tab. 1. On ImageNet-C with ResNet-18 the method is slightly worse than NORM (78.2 vs 77.3 mean error) 2112.00463 tab. 3.

When to use which

Gains are not uniform. DUA trails TENT with Wide-ResNet-40-2 on CIFAR-100C (37.6 vs 36.0 mean error) 2112.00463 tab. 2. On the KITTI-Fog and KITTI-Rain object detection tasks, the gains over the source model are large, but the fully supervised model is still ahead 2112.00463 tab. 4.

Implementation sketch

Language-agnostic pseudocode of the DUA loop. Variants differ in the augmentation, the momentum schedule and the reporting rule.

# Adapt BN running statistics one sample at a time [arxiv:2112.00463 §3.2]
# Schedule: rho_0 = 0.1, omega = 0.94, zeta = 0.005 [arxiv:2112.00463 §4.3]
def adapt_stream(model, stream, rho_0, omega, zeta):
    rho = rho_0
    for sample in stream:
        # Momentum decays per sample, floored by zeta [arxiv:2112.00463 eq. 4]
        rho = rho * omega
        momentum = rho + zeta
        # Build a batch of 64 augmented copies of the one sample [arxiv:2112.00463 §5.4]
        batch = augment_copies(sample, n=64)
        # Only BN running statistics change; weights stay fixed [arxiv:2112.00463 §3.2]
        set_bn_momentum(model, momentum)
        model.train()
        model(batch)   # running mean and variance updated [arxiv:2112.00463 eq. 6]
        model.eval()   # discard the output
    return model

Pitfalls, from the card:

• Check the momentum convention when porting. The released code sets the framework's BN momentum to rho_k + zeta, which weights the new batch statistics 2112.00463 eq. 4. The first update index is not stated in the paper.
• The first few samples dominate the final statistics, because momentum starts near 0.1 2112.00463 eq. 5.
• The statistics come from augmented copies of one image, not from real test batches 2112.00463 §3.2. The augmentations are not the test corruptions.
• The method adds forward passes per sample, not backward passes 2112.00463 §1.

Contradictions

• Reported number versus released code. The paper reports results at fixed sample counts 2112.00463 app. A. The released code, as read for the card, reports the minimum error over the whole stream, selected on test labels. The paper does not state a selection rule. See the card's Paper vs code section: The Norm Must Go On.
• Input scale in the adaptation pass. The paper does not describe the adaptation transform. The released code, as read for the card, appears to scale the detection inputs twice. This was not run, so treat it as unverified (see the card's Paper vs code section).

Papers

• The Norm Must Go On: Dynamic Unsupervised Domain Adaptation by Normalization: test-time updates of BN running statistics on an unlabeled stream, with a decaying momentum and no backpropagation.