PaperlayerView exact revision on arXiv

Source-linked AI summary

MEMO: Test Time Robustness via Adaptation and Augmentation

Marvin Zhang, Sergey Levine, Chelsea Finn

arXiv:2110.09506v3cs.LGcs.CV

TL;DR

Distribution shift exposes brittleness in deep models, while prior test-time adaptation can require multiple test points or specialized assumptions. MEMO adapts all parameters using augmented versions of a single test input and marginal-entropy minimization. Across CIFAR-10 and ImageNet distribution-shift benchmarks, it consistently improves performance and achieves several state-of-the-art results in the single-test-point setting.

  • Problem

    Deep models are brittle under distribution shift, and existing test-time adaptation methods often require additional assumptions such as access to multiple test points.

  • Method

    MEMO augments one test input and adapts all model parameters by minimizing the entropy of the marginal predictive distribution across augmentations.

  • Results

    MEMO consistently improves performance across CIFAR-10 and ImageNet distribution-shift benchmarks and achieves several state-of-the-art results when only one test point is available.

  • Takeaways & Limitations

    MEMO is broadly applicable to pretrained models and can be composed with robust training and model-statistics adaptation.

  • Takeaways & Limitations

    MEMO inference is more computationally expensive than standard inference because it performs augmentation and adaptation for each test point.

Abstract

from arXiv · show

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.

1 Introduction

Deep models can be brittle under distribution shift, motivating test-time methods that use individual test inputs without changing training. MEMO adapts all model parameters using augmented versions of one test point and marginal-entropy minimization, improving robustness across challenging benchmarks.

  • Deep neural networks may degrade sharply under input perturbations and natural domain shifts despite strong in-distribution performance.
  • Prior robustification methods often require modifying training, specialized procedures, or access to batches or sets of test inputs.
  • MEMO adapts a pretrained model on one test point by augmenting it and minimizing the marginal entropy of predictions across augmentations.
  • 1-10% improvement over non adaptive marginal distribution predictions is reported on ImageNet-C and ImageNet-R.
  • MEMO uses pretrained models without training-procedure or architecture assumptions and reports new state-of-the-art results when only one test point is available.

2 Related work

Distribution-shift research spans frameworks that add training or test assumptions, while empirical robustness methods modify models or training data. MEMO is positioned as complementary, applying single-point test-time adaptation to pretrained models, including robust ones.

  • Distribution-shift frameworks commonly add training or test assumptions to make the problem more tractable.
  • Training-time robustness methods increase model or dataset size or use heavy training augmentations, whereas MEMO operates at test time.
  • Prior test-time adaptation methods often use batches or entire test datasets for statistics, prototypes, or entropy minimization.
  • Test-time augmentation averages outputs over augmented test copies, while MEMO combines augmentation with adaptation using single test points.

3 Augmenting and Adapting at Test Time

MEMO adapts pretrained models on a single test input by enforcing consistent, confident predictions across sampled augmentations. It is broadly applicable and can be composed with robust training and batch-normalization adaptation.

  • 3 Augmenting and Adapting at Test Time: Test-time robustness assumes an adaptable model with differentiable probabilistic outputs, but no special training procedure, architecture, or ground-truth labels.
  • 3.2 Composing MEMO with Prior Methods: Single-point BN adaptation is not required by MEMO, which applies beyond models containing batch-normalization layers.
  • 3.1 Marginal Entropy Minimization with One test point: MEMO samples augmentations of one test point, estimates their marginal output distribution, minimizes its entropy, and predicts on the original input.
  • 3.1 Marginal Entropy Minimization with One test point: Marginal-entropy minimization encourages both invariant predictions across augmentations and confident predictions, unlike average conditional-entropy minimization.
  • 3.1 Marginal Entropy Minimization with One test point: MEMO adapts all model parameters with gradient-based optimization and uses one gradient step per test point before the final prediction.
  • 3.2 Composing MEMO with Prior Methods: MEMO can be combined with robust pretrained models and adapted batch-normalization statistics, with the latter generally improving test-time adaptation.

4 Experiments

The experiments evaluate MEMO across five distribution-shift benchmarks, model architectures, and pretraining procedures, comparing it with prior adaptation and augmentation strategies. MEMO generally improves robustness, reaches state-of-the-art results in several single-test-point settings, and benefits from combining augmentation with adaptation.

  • Experimental setup: MEMO is evaluated on five distribution-shift benchmarks spanning CIFAR-10 and ImageNet test sets.The benchmarks are CIFAR-10-C, CIFAR-10.1, ImageNet-C, ImageNet-R, and ImageNet-A.
  • Experimental setup: The study compares MEMO with TTT, Tent, BN adaptation, and non-adaptive test-time augmentation across multiple model architectures and training procedures.ImageNet experiments include robust ResNet-50 models and RVT∗-small, while CIFAR-10 experiments use ResNet-26 models.
  • CIFAR results: MEMO consistently improves test error over baselines, non-adaptive TTA, and TTT across CIFAR-10-C corruption types.It also provides a larger performance gain on CIFAR-10.1 than TTT, while non-adaptive TTA remains competitive on the simpler test sets.
  • ImageNet results: MEMO combined with robust training achieves new state-of-the-art performance on ImageNet-C and ImageNet-R for ResNet-50 and RVT∗-small models using one test point.For ResNet-50 models, MEMO achieves new state-of-the-art performance on each ImageNet benchmark in the single-test-point setting.
  • ImageNet results: MEMO substantially improves ImageNet-A and achieves state-of-the-art performance among ResNet-50 models, although TTA performs better for RVT∗-small.Unlike prior test-time adaptation methods, MEMO reports successful ImageNet-A results using a learning signal from single test points.
  • Ablative study: Ablations show that marginal entropy minimization generally outperforms pairwise cross-entropy and conditional-entropy adaptation because both invariance and confidence provide useful learning signals.Small augmentation counts such as B = 4 and 8 can already provide significant performance gains, creating an efficiency–accuracy tradeoff.

5 Discussion

MEMO improves test-time robustness across several model families and distribution-shift benchmarks without modifying training or requiring multiple test inputs. Its main costs are higher inference computation and potential degeneracy during continual adaptation.

  • Discussion: MEMO consistently improves performance for ResNet, vision transformer, and, to an extent, ResNext models in the single-test-point setting.The method achieves several new state-of-the-art results across CIFAR-10 and ImageNet distribution-shift benchmarks.
  • Discussion: MEMO requires neither access to nor changes in model training procedures and adapts using single test inputs.This broadens applicability across pretrained architectures and training methods.
  • Discussion: MEMO inference is more computationally expensive than standard inference because it performs augmentation and adaptation.Using fewer augmentations per test point can improve the efficiency–accuracy tradeoff.
  • Discussion: Continual adaptation as more test data arrives produced degenerate solutions in preliminary experiments, including constant-label predictions with maximal confidence.The authors suggest restricting adapted parameters or adding regularization as possible remedies.

Checklist

The checklist reports that the paper describes its contributions, limitations, training details, and reproducibility materials, while several ethics and compute disclosures are marked not applicable or absent.

  • Checklist: The paper reports that it describes its contributions and limitations, with limitations discussed in Section 5.
  • Checklist: The authors report providing code, data, and reproduction instructions in supplementary material.
  • Checklist: Training details and hyperparameter-selection procedures are reported in Appendix A.
  • Checklist: The checklist marks compute-resource reporting and asset-license discussion as absent, while noting that all assets are publicly and freely available.
  • Checklist: Participant-related disclosures and personally identifiable or offensive-content discussion are marked not applicable.

A Experimental Protocol

The experimental protocol selects hyperparameters on disjoint CIFAR-10-C and ImageNet-C validation corruptions, reuses them on benchmarks without validation sets, and specifies comparison settings for models and baselines.

  • Hyperparameter selection: Hyperparameters are selected using four disjoint CIFAR-10-C and ImageNet-C validation corruptions.The same settings are reused on benchmarks without validation sets, with no additional tuning.
  • Hyperparameter selection: The grid search varies learning rate, gradient steps, an entropy threshold, and prior strength N.The tested learning rates range from 10^-3 to 10^-6 with additional multiplicative adjustments.
  • Hyperparameter selection: Threshold-based adaptation slightly worsens validation performance, while prior strength N = 16 performs best for single-point BN adaptation.
  • Model-specific settings: For RVT∗-small, the search additionally varies the update rule, weight decay, and prior strength.
  • Model-specific settings: AdamW with learning rate 0.00001, weight decay 0.01, and B = 64 is used for RVT∗-small and, with B = 32, for ResNext-101.
  • Baseline comparisons: The CIFAR comparison uses a ResNet-26 architecture aligned with TTT's group-normalization modifications, where single-point BN adaptation is inapplicable.
  • Baseline comparisons: TTA uses the same AugMix augmentations as MEMO, while BN adaptation uses N = 16 for single points and N = 256 with batches of 256 test inputs.
  • Model sources: Baseline and robust model parameters are obtained from torchvision, released links, or specified pretrained checkpoints.

B Analysis on Augmentations

The augmentation analysis finds that augmentation choice matters most under severe, systematic shifts, while adaptation without augmentations provides only marginal gains beyond single-point BN adaptation.

  • Need for augmentations: Single-point entropy minimization generally provides only marginal gains beyond single-point BN adaptation.The comparison uses episodic Tent with test batch size 1, equivalent to entropy minimization without augmented copies.
  • Need for augmentations: Augmentations are important for achieving MEMO's reported results in the single-test-point setting.
  • Augmentation choice: Heavier AugMix augmentations yield greater gains than standard augmentations under the severe and systematic shifts of CIFAR-10-C.The difference is smaller for unchanged or subtly shifted test sets.
  • Augmentation choice: AugMix may offer greater test-time advantages for robust models trained with heavy augmentations than for the standard-augmentation ResNet-26 model.

C Full CIFAR-10-C and ImageNet-C Results

The supplementary results report CIFAR-10-C test errors by corruption and severity level, while ImageNet-C results are supplied separately in CSV files.

  • CIFAR-10-C results are broken down by corruption type and severity level.The listed tables cover levels 1 through 5.
  • Joint training and TTT are omitted because their results are available from Sun et al..
  • ImageNet-C test results are provided in supplementary CSV files.
  • Tables 7–11 report test error (%) for CIFAR-10-C corruption levels 1–5.Table 7 covers level 5, Table 8 level 4, Table 9 level 3, Table 10 level 2, and Table 11 level 1.
Loading 2110.09506v3…