[Paper Review] MEMO: Test Time Robustness via Adaptation and Augmentation

Marvin Zhang, Sergey Levine+1 more|arXiv (Cornell University)|Oct 18, 2021

Advanced Neural Network Applications38 references85 citations

TL;DR

MEMO adapts a pretrained probabilistic model at test time by augmenting a single test input, minimizing the marginal entropy across augmentations to enforce invariances and confidence, achieving improved robustness on ImageNet-C/R/A and CIFAR-10 variants.

ABSTRACT

While deep neural networks can attain good accuracy on in-distribution test points, many applications require robustness even in the face of unexpected perturbations in the input, changes in the domain, or other sources of distribution shift. We study the problem of test time robustification, i.e., using the test input to improve model robustness. Recent prior works have proposed methods for test time adaptation, however, they each introduce additional assumptions, such as access to multiple test points, that prevent widespread adoption. In this work, we aim to study and devise methods that make no assumptions about the model training process and are broadly applicable at test time. We propose a simple approach that can be used in any test setting where the model is probabilistic and adaptable: when presented with a test example, perform different data augmentations on the data point, and then adapt (all of) the model parameters by minimizing the entropy of the model's average, or marginal, output distribution across the augmentations. Intuitively, this objective encourages the model to make the same prediction across different augmentations, thus enforcing the invariances encoded in these augmentations, while also maintaining confidence in its predictions. In our experiments, we evaluate two baseline ResNet models, two robust ResNet-50 models, and a robust vision transformer model, and we demonstrate that this approach achieves accuracy gains of 1-8\% over standard model evaluation and also generally outperforms prior augmentation and adaptation strategies. For the setting in which only one test point is available, we achieve state-of-the-art results on the ImageNet-C, ImageNet-R, and, among ResNet-50 models, ImageNet-A distribution shift benchmarks.

Motivation & Objective

• Motivate and study test-time robustification methods that do not rely on training-time changes or access to large test batches.
• Propose a plug-and-play MEMO approach that adapts all model parameters at test time using single test points.
• Encourage predictions to be invariant across augmentations while remaining confident, via marginal entropy minimization.
• Demonstrate compatibility and additive gains when combined with existing robustness techniques and BN adaptation.

Proposed method

• Given a trained probabilistic model f_theta, present a single test input x and a set of augmentations A.
• Sample B augmentations of x to form a batch of augmented inputs; estimate the marginal output distribution by averaging p_theta(y|a(x)) over augmentations.
• Define the MEMO loss as the entropy of the marginal distribution H(bar{p}_theta(.|x)); update theta via gradient descent to minimize this marginal entropy (one gradient step per test point).
• After adaptation, predict on the original input x using the updated parameters; do not require ground-truth labels at test time.
• Optionally combine MEMO with BN-statistics adaptation or other robustness methods without changing the training procedure.

Experimental results

Research questions

• RQ1 Can test-time adaptation be effective without assumptions about the training process or access to test batches?
• RQ2 Does minimizing the marginal entropy across augmented copies of a single test point improve robustness to distribution shift?
• RQ3 How does MEMO interact with existing robustness techniques (e.g., BN adaptation, AugMix, MoEx) across different model architectures and datasets?
• RQ4 Is adaptation or augmentation the dominant contributor to observed gains, and what is the role of augmentation choice and sample count?

Key findings

• MEMO yields 1–8% accuracy gains over standard evaluation on challenging distribution-shift benchmarks.
• On ImageNet single-point tests, MEMO achieves state-of-the-art results for ResNet-50 and strengthens robustness against ImageNet-C, ImageNet-R, and ImageNet-A.
• MEMO improves robustness of both ResNet and vision transformer models, and outperforms prior augmentation/adaptation strategies on several benchmarks.
• Ablation studies show both the invariance across augmentations and maintaining confidence are important components of MEMO.
• MEMO can be combined with pretrained robust models and BN adaptation to further improve performance, and it shows strong gains even for models trained with heavy data augmentations.