Adaptive Risk Minimization: Learning to Adapt to Domain Shift

Marvin Zhang; Henrik Marklund; Nikita Dhawan; Abhishek Gupta; Sergey Levine; Chelsea Finn

Source-linked AI summary

Adaptive Risk Minimization: Learning to Adapt to Domain Shift

Marvin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, Chelsea Finn

arXiv:2007.02931v4Jul 6, 2020cs.LGstat.ML

TL;DR

Distribution shift challenges ERM and exposes limits of methods built around invariant relationships. ARM trains models to adapt using unlabeled test batches, and its methods consistently improve average and worst-case performance across evaluated domain-shift problems, though adaptation is not uniformly effective.

Problem
Distribution shift can substantially reduce ERM performance, while prior domain-generalization methods often assume invariant input-output relationships across domains.
Method
ARM meta-learns models on training domains to use unlabeled test batches for adaptation to domain shift.
Results
ARM methods consistently improve average and worst-case metrics over prior domain-shift approaches across the evaluated problems.
Takeaways & Limitations
ARM provides a framework for learning models whose test-time adaptation is directly optimized rather than relying only on a single robust or invariant model.
Takeaways & Limitations
ARM’s adaptation may not suit some benchmark problems, and its formal setting assumes domain structure where p_y|x is determined by p_x.

Abstract

from arXiv · show

A fundamental assumption of most machine learning algorithms is that the training and test data are drawn from the same underlying distribution. However, this assumption is violated in almost all practical applications: machine learning systems are regularly tested under distribution shift, due to changing temporal correlations, atypical end users, or other factors. In this work, we consider the problem setting of domain generalization, where the training data are structured into domains and there may be multiple test time shifts, corresponding to new domains or domain distributions. Most prior methods aim to learn a single robust model or invariant feature space that performs well on all domains. In contrast, we aim to learn models that adapt at test time to domain shift using unlabeled test points. Our primary contribution is to introduce the framework of adaptive risk minimization (ARM), in which models are directly optimized for effective adaptation to shift by learning to adapt on the training domains. Compared to prior methods for robustness, invariance, and adaptation, ARM methods provide performance gains of 1-4% test accuracy on a number of image classification problems exhibiting domain shift.

1 Introduction

Distribution shift can substantially degrade standard ERM, while prior domain-generalization methods often assume invariant input-output relationships. ARM instead learns models that use unlabeled test batches to adapt to new domains.

Distribution shift occurs when test data differ from training data, causing standard ERM performance to deteriorate significantly.
Handwriting classification illustrates shift across new users, where unlabeled batches may help resolve examples whose labels depend on user-specific writing styles.
Domain generalization commonly organizes training data into domains and assumes test distributions represent new domains.
Prior robustness and invariance methods seek a single model or invariant feature space, but can struggle when input-output relationships vary across domains.
ARM optimizes models to leverage an unlabeled adaptation phase, meta-learning adaptability from training domains and evaluating problems with shared domain structure.

2 Related Work

Related work addresses domain shift through invariance, robustness, meta-learning, and test-time adaptation. ARM differs by targeting purely unlabeled adaptation across potentially multiple test distributions and learning adaptation strategies from training domains.

Invariance and robustness to domains: Domain-generalization benchmarks typically assume a consistent input-output relationship across domains, motivating invariant features and worst-case group optimization.
Meta-learning: Prior domain-generalization methods primarily seek non-adaptive generalization, whereas this work extends meta-learning toward unlabeled adaptation for distribution shift.
Meta-learning: Many meta-learning approaches use labeled adaptation data, while ARM targets purely unlabeled adaptation.
Adaptation to shift: Traditional domain adaptation often assumes one test distribution, making some methods difficult to apply when multiple test distributions are present.
Adaptation to shift: ARM builds on adaptive domain-generalization theory and prior test-time methods, adding training-domain meta-learning to improve how models adapt.

3 Preliminaries and Notation

The paper formalizes domains as joint data distributions and adaptive models as predictors that use the marginal input distribution. Under a structural condition linking input and label distributions, adaptive risk provides a principled objective.

A domain is a joint distribution p_xy, with training data organized across S domains and test settings containing unlabeled data from new domains.
An adaptive model f takes an input x and its domain’s marginal distribution p_x, using p_x as a surrogate for information about p_y|x.
The usefulness of adaptation depends on p_x providing information about p_y|x; otherwise, knowing p_x does not benefit prediction.
Adaptive risk E evaluates expected prediction risk across domains sampled from µ, the distribution over joint domain distributions.
If p_y|x is a deterministic function of p_x, minimizing E yields Bayes-optimal predictions for µ-almost all domains.
The required condition excludes settings where domains can share identical input distributions but use different label orderings, as in standard few-shot learning benchmarks.
In practice, the model receives a batch of same-domain unlabeled inputs, whose empirical distribution approximates p_x; this also supports streaming experiments.

4 Adaptive Risk Minimization

Adaptive risk minimization trains models to use unlabeled test inputs for domain adaptation, rather than only seeking invariant features or a single robust model. The framework meta-learns adaptation across training domains and instantiates contextual, batch-normalization, and gradient-based methods.

Adaptive Risk Minimization: ARM optimizes models to leverage an unlabeled adaptation phase for handling domain shift.Training domains provide the meta-learning episodes used to produce models that adapt at test time.
Evaluation setting: The framework is evaluated for settings where unlabeled test batches or streams can be leveraged for better performance, including test domains not sampled identically to training domains.The paper also considers streaming inputs, where adaptation occurs incrementally.
Adaptive model: The adaptive model combines a prediction network g with an adaptation model h that maps model parameters and K unlabeled points to adapted parameters θ′.The prediction model produces outputs using θ′, while h performs adaptation using only unlabeled data.
Meta-learning procedure: Algorithm 1 simulates test-time adaptation during meta-training, using unlabeled batches to adapt and labels only to optimize the training update.At test time, the procedure applies h to an unlabeled batch and predicts with the resulting parameters.
Methods: ARM supports contextual, batch-normalization, and learned-loss methods as extensions of existing meta-learning or unlabeled-adaptation approaches.ARM-CML summarizes batches into context; ARM-BN adapts normalization statistics; ARM-LL updates parameters with an unlabeled loss.
ARM-BN: ARM-BN trains on batches sampled from a single domain and recomputes normalization statistics at test time instead of using training running averages.The single-domain batch-training difference is identified as novel to ARM-BN, and meta-training improves performance where BN adaptation has favorable inductive bias.

5 Experiments

The experiments evaluate ARM across image-classification benchmarks with domain shifts, compare it against robustness, invariance, and adaptation methods, and test both batch and streaming adaptation. ARM generally achieves the strongest worst-case and average accuracy, while results on WILDS show that method suitability varies by problem.

Experimental goals: The experiments address whether ARM can leverage unlabeled adaptation, outperform prior methods, and adapt when test points arrive as a stream.These questions are evaluated across four proposed image-classification benchmarks and WILDS datasets.
Evaluation domains and protocol: The proposed benchmarks include rotated MNIST, FEMNIST, and corrupted CIFAR-10-C and Tiny ImageNet-C, with domains defined by rotations, users, or corruptions.The corrupted-image protocol uses 56 training corruptions and 22 disjoint test corruptions, measuring worst-case and average accuracy.
Comparisons and ablations: Context and learned-loss ablations test whether sampling from individual domains and meta-training on domain shifts contribute to improved performance.The ablations instead sample unlabeled training batches uniformly from the entire training set.
Comparisons and ablations: ARM methods consistently improve both worst-case and average accuracy over prior methods, with ARM-CML performing broadly well and ARM-BN performing best overall on corrupted-image testbeds.BN adaptation and TTT are the strongest prior methods, but ARM methods are less reliant on favorable image-classification inductive biases.
Streaming evaluation: ARM-CML and ARM-BN reach near their original worst-case and average accuracy within 50 streaming test points on Tiny ImageNet-C, despite training batch sizes of 100.The methods update average context or normalization statistics after each new test point.
WILDS results: On WILDS, different methods perform best on different problems: ARM-BN improves RxRx1 and performs well on PovertyMap but struggles on FMoW.Adaptation broadly boosts Camelyon17 performance, while CORAL performs best on iWildCam and no method significantly outperforms ERM on FMoW or PovertyMap.

6 Discussion and Future Work

The paper presents ARM as a framework for adapting to domain shift using unlabeled test batches and reports consistent gains over prior approaches. It also identifies unresolved theoretical and broader-impact concerns.

ARM meta-learns models that adapt to new domains using only unlabeled test examples.The framework optimizes models for effective test-time adaptation based on training domains.
ARM methods consistently improve both average and worst case performance compared with prior domain-shift approaches.
Future work includes extending ARM beyond contextual meta-learning and learning domains when domains are unavailable during training.The paper suggests unsupervised domain construction and diverse domain generation as possible directions.
Distribution shift can make deployed machine-learning systems unreliable, harmful, or dangerous.
The work does not address failure modes and theoretical limits that could create false confidence in deployed systems.

B More Details on the ARM Methods

The ARM methods form adapted prediction models from context computed over test inputs, with batch and streaming procedures that avoid retaining observed test points. Training optimizes post-adaptation predictions.

ARM-CML represents post-adaptation parameters as [θ, c] and uses the prediction model with context c at training and test time.
ARM-CML averages context vectors produced for batch inputs and feeds the resulting context into the prediction network.The context network processes each image, while training backpropagates the post-adaptation loss to update θ and φ.
In streaming ARM-CML, the running average context is updated after each test point before predicting on that point.After K observations, this produces the same context as the batch setting without storing past test points.
Streaming ARM-BN tracks the mean and variance of batch-normalization activations as its context.Its update is more complex because it requires computing second moments.
ARM-LL used α = 0.1 and one gradient step for both meta-training and meta-testing.The paper did not evaluate ARM-LL in the streaming setting, though it describes a possible extension.

C Contrasting with Prior Benchmarks

The paper contrasts ARM’s adaptation-oriented benchmarks with prior domain-generalization and DRO testbeds centered on invariance or fixed groups. Results indicate that adaptation is problem-dependent.

The paper targets problems where unlabeled adaptation is feasible, helpful, and potentially crucial, leading its benchmarks to differ from prior domain-shift work.
ARM-CML showed middling performance across most DomainBed testbeds, suggesting adaptation may not suit their standard formulations.
Adaptation and ARM methods do not improve performance on some WILDS problems, potentially because marginal p(x) carries limited information about p(y|x) or domains lack shared structure.
Meta-learning motivates benchmarks with many domains to reduce task overfitting, whereas DomainBed benchmarks use fewer domains.
Group DRO testbeds usually reuse the same groups for training and testing, unlike settings with disjoint held-out test domains.

D Additional Experimental Details

The experiments use controlled data splits, repeated-seed evaluation, and specified neural-network and optimization settings. Additional details define rotation construction, validation, and model architectures.

Code for the Table 1 results is publicly available in the authors’ GitHub repository.
Each method was run across three seeds, with results reported as the mean and standard error across seeds.
Models were checkpointed every epoch and evaluated using the checkpoint with the best worst case validation accuracy.Hyperparameters were selected in preliminary experiments using validation accuracy only.
Dataset splits were designed without considering any method’s train, validation, or test accuracy.The choices were intuitive, random, or made for benign alternate purposes.
The rotation experiment used 32292 training points with rotation frequencies decreasing at larger angles.Rotations from 0 through 20 degrees had 7560 points each, while 120 to 130 degrees had 108 each.
Rotation models trained for 200 epochs with batch size 300, Adam learning rate 0.0001, and ARM meta-batches arranged as 6 × 50.
Validation sampled 300 held-out points per rotation, while testing sampled 3000 test points per rotation.
Non-context models used four convolutional layers, average pooling, a size-200 fully connected layer, and a linear output layer.The context network used two convolutional layers followed by a 12-filter convolutional layer.

D.2 FEMNIST details

FEMNIST experiments use user-structured splits and evaluate methods under variable user-level data sizes, while additional tables compare ARM with adaptation baselines and probabilistic contexts.

D.2 FEMNIST details: FEMNIST is more challenging than MNIST because it has 62 classes, severe label imbalance, and inherent ambiguities.Filtering users with fewer than 100 examples leaves 262 training, 50 validation, and 35 test users.
D.2 FEMNIST details: Validation accuracy is computed once per validation user, matching the test-evaluation procedure.User data sizes are not always multiples of 50, so methods can receive small batches.
D.2 FEMNIST details: CIFAR-10-C and Tiny ImageNet-C use disjoint corruption domains across training, validation, and test sets.The supplied passage specifies 56, 17, and 22 domains for each split, respectively.
D.2 FEMNIST details: The experiments use support size 100 and meta batch size 3, with architecture changes for RGB images and pretrained ResNet-50 for Tiny ImageNet-C.Context and learned-loss ablations use small auxiliary networks.
D.2 FEMNIST details: ARM-CML remains more performant than DANN despite DANN retraining separately for each known test domain.The comparison attributes ARM-CML’s advantage to leveraging meta-training and adaptation.
D.2 FEMNIST details: Standard ARM-CML performs better than its probabilistic-context modification on rotated MNIST and FEMNIST.The passage suggests this may reflect the standard objective’s direct emphasis on predictive accuracy.

E Additional Experiments

Additional experiments compare ARM with unsupervised and zero-shot domain adaptation methods, including a probabilistic-context variant of ARM-CML.

E Additional Experiments: UDA methods are evaluated against ARM using DANN, which assumes the particular test domain is known during training.With multiple test shifts, UDA retrains separately for each shift.
E Additional Experiments: A probabilistic-context variant of ARM-CML generally performs worse than the original formulation.The variant places independent unit Gaussian priors on context dimensions and optimizes an evidence lower bound.

E.2 Additional results with loosened assumptions

Additional results show that ARM can adapt rapidly in streaming evaluation and can operate when domains are learned rather than explicitly specified.

E.2 Additional results with loosened assumptions: Fewer than 10 test inputs suffice for ARM methods to reach their batch-adaptation performance on rotated MNIST streaming evaluation.This occurs despite meta-training with batches of 50 points.
E.2 Additional results with loosened assumptions: ARM-CML is evaluated for unknown domains by training a VAE with discrete latent variables on training images and labels.The experiment tests unsupervised discovery of domain structure in the training data.
E.2 Additional results with loosened assumptions: Learned domains let ARM-CML outperform ERM and match TTT on rotated MNIST.The reported result may improve with techniques that learn more diverse domains.
E.2 Additional results with loosened assumptions: The VAE is conditioned on y and c, and learns to use c to represent rotations.The figure organizes generated samples by y on the x axis and c on the y axis.