Adaptive Risk Minimization: Learning to Adapt to Domain Shift
Authors: Marvin Mengxin Zhang
Henrik Marklund
Nikita Dhawan
Abhishek Gupta
Sergey Levine
Chelsea Finn
Venue: NeurIPS 2021 Poster (2021)
Keywords: distribution shift
domain generalization
test time adaptation

Abstract: 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.

Adaptive Risk Minimization: Learning to Adapt to Domain Shift
Marvin Zhang Henrik Marklund Nikita Dhawan" Abhishek Gupta', Sergey Levine', Chelsea Finn
UC Berkeley, Stanford University
Abstract
A fundamental assumption of most machine learning algorithms is that the training and test data are drawn from the same underlying distribution. However, this as- sumption 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 Iearn a single robust model or invariant feature space that performs well all domains. In contrast, 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 adapt on the training domains. Compared to prior methods for robustness, in- variance, and adaptation, ARM methods provide performance gains of 1-4% test accuracy on a number of image classification problems exhibiting domain shift.
aim to learn models
Introduction
The standard assumption in empirical risk minimization (ERM) is that the data distribution at test time will match the training distribution. When this assumption does not hold, i.e., when there is distribution shift, the performance of standard ERM methods can deteriorate significantly [z 3X]
As an example which we study quantitatively in Section 5 consider a handwriting classification model that, after training on data from past users, is deployed to new end users. Each new user represents a new test distribution that differs from the training distribution. Thus, each test set- ting involves dealing with shift. In Figure 1] we visualize a batch of 50 examples from test user, and we highlight an ambiguous example which may be either "2" (written with loop) or an "a" (in the double-storey style) depend- ing on the user's handwriting. Due to the biases in the training data, an ERM trained model incorrectly classifies this example as "2". However, we can see that the batch of images from this test user contains other examples of "2" (written without loops) and "a" (also double-store storey) from this user. Can we somehow leverage this unlabeled data to better handle test shifts caused by users?
S
5
ERM
ARM
Figure 1: An example of ambiguous data points in handwriting classification, eval- uated quantitatively Section 5
equal contribution
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Any framework that aims to address this question must use additional assumptions beyond the ERM setting, and many such frameworks have been proposed [54 One commonly used assumption within several frameworks, such as domain generalization 7 23 is that the training data are provided in domains and distributions at test time will represent new domains. The example above neatly fits this description if we equate users with domains- we would be assuming that the training data are organized by users and that the model will be tested separately on new users, and these are reasonable assumptions. Constructing training domains in practice is generally accomplished by using meta-data, which exists for many commonly used datasets. Thus, this domain assumption is applicable for a wide range of realistic distribution shift problems (see, e.g., Koh et al. (51
However, prior benchmarks for domain generalization and similar settings typically center around invariances i.e., in these benchmarks, there is a consistent input- input-ouiput relationship across all domains, and the goal is to learn this relationship while ignoring the spurious correlations within the domains (see, e.g., Gulrajani and Lopez-Paz. (231). Thus, prior methods aim for generalization to shifts by discovering this relationship, through techniques such as robust optimization and learning an invariant feature space 4
These methods are appealing in that they make minimal assumptions about the information provided at test time - in particular, they do not require test labels, and the learned model can be immediately applied to predict on a single point. Nevertheless, these methods also have limitations, such as in dealing with problems where the input-output relationship varies across domains, c.g., the handwriting classification example above.
3
60
In this paper, we instead focus on methods that aim to adapt at test time to domain shift. To do so, we study problems in which it is both feasible and helpful (and perhaps even necessary) to assume access to batch or stream of inputs at test time. Leveraging this test assumption does not require labels for any test data and is feasible in many practical setups. For example, for handwriting classification, we do not access only single handwritten characters from an end user, but rather collections of characters such sentences or paragraphs. Unlabeled adaptation has been shown empirically to be useful for distribution shift problems [69 [3 [5), such as for dealing with image corruptions [25]. number of problems, detailed in Section 5] for which adaptation is beneficial in dealing with domain shift.
Taking inspiration from these findings, we propose and evaluate
Our main contribution is to introduce the framework of adaptive risk minimization (ARM), which proposes the following objective: optimize the model such that it can maximally leverage the unlabeled adaptation phase to handle domain shift. To do so, we instantiate set of methods that, given a set of training domains, meta-learns model that adaptable to these domains. These methods are straightforward extensions of existing meta-learning approaches, thereby demonstrating that tools from the meta-learning toolkit can be readily adapted to tackle domain shift. Our experiments in [Section 5 test on several image classification problems, derived from benchmarks for federated learning g and image classifier robustness 25 in which training and test domains share structure that can be leveraged for improved performance. These testbeds are also a contribution of our work, as we believe these problems can supplement existing benchmarks which, almost exclusively, are We also evaluate on the WILDS suite of distribution shift problems [35, which have been curated to faithfully represent important real world problems. Empirically, we demonstrate that the proposed ARM methods, by leveraging meta-training and test time adaptation, are often able to outperform prior -ftte--rrrt methods by 1-4% test accuracy.
designed with invariance in mind 3
[53 23
2 Related Work
A number of prior works have studied distribution shift in various forms [54] review prior work in domain generalization, group robustness, meta-learning. and adaptation.
In this section, we
Invariance and robustness to domains. As discussed above, number of frameworks leverage training domains to address test time shift. The terminology in prior work is scattered and, depending on the application, includes terms such "groups", "datasets", "subpopulations"" and "users": in this work, we adopt the term "domains" which we believe is an appropriate unifying term. A number of testbeds for this problem setting have been proposed for image classification, including generalizing and underrepresented demographics [60).
to new datasets 11 new image types 19 [
Prior benchmarks typically assume the existence of consistent input-output relationship across domains that is learnable by the specified model, thus motivating methods such as learning an invariant feature space 41
44 i
or optimizing for worst case group performance 330 60 In

---

particular, methods for domain generalization - sometimes referred to as multi domain adaptation 68 or zero shot domain adaptation 778 have largely focused on learning invariant
features 19 67
53
Gulrajani and Lopez-Paz
2 provide comprehensive survey of
domain generalization benchmarks and find that, surprisingly, ERM is competitive with the state of the art across all the benchmarks considered. In Appendix C we discuss this finding as well as the performance of an ARM method on this benchmark suite. In Secction S] we identify different problems for which adaptation is helpful, and we find that, on these problems, ARM methods consistently outperform ERM and other non adaptive methods for robustness and invariance.
27 has been most extensively studied in the context of few Our aim is not to address few-shot recognition problems, nor to propose novel meta- meeaallerning algorithm, but rather to extend meta- meta-learning paradigms to problems requiring unlabeled adaptation, with the primary goal of tackling distribution shift. This aim differs from previous work in meta-learning for domain generalization 44 15 which seek to meta- train models for non adaptive generalization performance. We discuss Section 4 how paradigms can be readily extended using the ARM framework.
Meta-learning. Meta-learning 62
6
shot labeled adaptation 61
74
55
18
such contextual meta-learning 20 57
Some other meta-learning methods adapt using both labeled and unlabeled data, either in the semi supervised learning setting (56 81
42] or the transductive learning setting [51] 4d, 2 29] These works all assume access to labeled data for adaptation, whereas we propose methods and problems for purely unlabeled adaptation. Prior works in meta meta-learning for unlabeled adaptation include Yu et al. |.. who adapt policy to imitate human demonstrations in the context of robotic learning: Metz et al. 4S, who meta-learn an update rule for unsupervised representation learning, though they still require labels to learn a predictive model; and Alet al. M. who meta-learn adaptive models based on task specific unsupervised objectives. Unlike these prior works, we propose a general framework for tackling distribution shift problems by meta-learning unsupervised adaptation strategies. This framework simplifies the extension of meta-learning paradigms to these problems, encapsulates previous approaches such as the gradient based meta-learning approach of Yu et al. [80], and sheds light on how to improve existing strategies such as adaptation via batch normalization
43]
Adaptation to shift. Unlabeled adaptation has primarily been studied separately from meta-learning.
Domain adaptation is prominent framework that assumes access to test examples at training
time 13 7d. similar to transductive learning 7 As such, most domain adaptation methods
consider the problem setting where there is single test distribution [64 44 19
72 10] and
some of these methods are difficult to apply to problems where there are multiple test distributions.
Certain domain adaptation methods have also been applied in the domain generalization setting, such and we compare to these methods in Section 5)
as methods for learning invariant features 19 31
Adaptive methods for domain generalization include Muandet et al. 40 and Kumagai and Iwata ky7 who propose a method similar to one of the ARM methods described below. We compare to a version of this method in Appenddi El Blanchard et al. (7] and Blanchard et al. [ provide a theoretic study of domain generalization and establish favorable generalization bounds for models that can adapt to domain shift at test time. We summarize some of these results in /Section 3 In comparison, our work establishes a framework that makes explicit the connection between adaptation to domain shift and meta-leanan learning, allowing us to devise new methods in straightforward and principled manner. These methods are amenable to expressive models such as deep neural networks, which enables us to propose and evaluate on problems with raw image observations.
Test time adaptation has also been studied for dealing with label shift 5 favorable inductive biases for the domain of interest. For image classification, techniques such as normalizing via the test inputs 4 and optimizing self-supervised surrogate losses [69] have proven
35 66
and crafting
effective for adapting to image corruptions [25]. We compare to these prior methods in Secciiion 5) and empirically demonstrate the advantage of using training domains to learn how to adapt.
3 Preliminaries and Notation
In this section, we discuss the domain generalization problem setting and formally describe adaptive models. In ISSection 4 we discuss how adaptive models can be meta-trained via the ARM objective and approach, and we instantiate ARM methods which we empirically evaluate in Section 5
3

---

L E X' and y E Y represent the input and output, respectively. We can formalize the domain generalization problem setting using the following data generation process [] first, a joint data distribution Pxy is sampled from set of distributions Pxy. and then some data points are sampled from Pxy We refer to each Pxy domain, e.g., particular dataset or user, thus Pxy represents the set of all possible domains. We assume that the training dataset is composed of data from S runs of this generative process, organized by domain. An equivalent characterization which we will use for clarity is that, within the training set, there are S domains, and each data point (x(*) y()) is annotated with a domain label 20(). Each 22() is an integer that takes on a value between 1 and S, indicating which Pxy generated the i-th training point (though, of course, we do not have access to, or knowledge of, Pxy itself). At test time, there may be multiple evaluation settings, where each setting is considered separately and contains only unlabeled data sampled via a new run of the same generative process. This data may represent, e.g., new dataset or user, and the test domains are likely to be distinct from the training domains when |Pxy] is large or infinite.
Our formal goal is to optimize for expected performance, e.g., classification accuracy, at test time. To do so, let first consider predictive models of the form f X' x Px + Y', where the model f takes in not just an input x but also the marginal input distribution Px E Px that x was sampled from. We refer to an adaptive model, as it has the opportunity to use px to adapt its predictions on x. The underlying assumption is that px provides information about Py|x. i.c., Px is used as a surrogate input in place of Pxy. In the worst case, if Px and Pylx are sampled independently, then the model does not benefit at all from knowing Px- In many problems, however, we expect knowledge about px to be useful, e.g., for resolving ambiguity as in the handwriting classification example in Section
Theoretically, when Px provides information about Pylx- and when training and test domains are drawn from the same distribution over Pxy. we can establish favorable generalization bounds for the expected performance of f in adapting to domain shift at test time. We can formalize this as follows. First, define prediction model to be non adaptive model of the form g X Y, and define the risk for prediction model and loss function E, under a data distribution Pxy. as
R(g. Pxy) A Epn [((g(x), v)]
Further, define the Bayes optimal risk for E under Pxy as
R¹ (Pxx) A min Ríg. Pxy)
Let yl denote the distribution on Pay from which training and test domains Pxy sampled. To avoid overlapping terms, define the adaptive risk for an adaptive model ʃ and E, under p., to be
{(f, 1²) A E [Ep,, (L((x,x)).v)
(I)
We state the following result from Blanchard al. R. which details a condition on μ under which is a strongly principled objective for learning adaptive models.
Lemma 9 from Blanchard et al. |. Let f denote a minimizer of & for the given p. If j is a
distribution on Pxy such that p-almost surely it holds that pyix = M (px) for some deterministic mapping M, then for p-almost all Pxy. we have
R( (. Pxy) Pxy) = R² (Pxx)
E(J", (1) = Ep [R (pxy))
In other words, an adaptive model which minimizes the adaptive risk E coincides with a Bayes optimal decision function for Pxy. for p-almost all domains Pxy-
Remark. The required condition on μ that Pylx is determined by Px - holds if, and only if, an expert (or oracle) is able to correctly label inputs from a given domain provided only information about the input distribution. This condition holds for the testbeds proposed in this paper, those in Gulrajani and Lopez-Paz [3 and those in WILDS 35 The condition does not hold for, e.g., standard few shot learning testbeds, where is possible for two domains with identical input distributions to shuffle their label orderings differently 7Z4 Thus, these problems are outside the scope of this work.
This result provides strong justification for learning adaptive models f by minimizing the adaptive risk E. However, a practical instantiation of this approach requires some approximations. First, we
Formally, the number of points sampled is another random variable with support over the positive integers.

---

do not know and cannot input Px to f in most cases. Instead, we instantiate f such that it takes in batch of inputs X!...., Xx. all from the same domain, where K can vary. f makes predictions on the whole batch, which also serves as empirical approximation (i.e., histogram) f>x of Px [8] In our exposition, we will assume that a batch of unlabeled points is available at test time for adaptation. However, we also experiment in |Section 5] with the streaming setting where the test inputs are observed one at a time and adaptation occurs incrementally.
Notice that, if we instead passed in an approximation Pxy of Pxy to the model, such as a batch of labeled data (x1, 81 yi)..... (XK. yK), then this setup would resemble the standard few shot meta- learning problem [74]. Formally, a meta-learning model takes in both an input x and pxy, which approximates the distribution that x was sampled from and thus can be used adapt the prediction on X. Compared to our problem setting, the meta- meta-learning formalism can tackle wider range of problems but also requires more restrictive assumptions, specifically, labels at test time via Pxy- Transductive meta-learning methods further assume that, in addition to Pxy. a full batch of inputs XK passed into the model, which allows for better estimation of the input distribution The model then makes predictions on this entire batch. In meta-learning terminology, Pxy and XI...., XK are often referred to as the support and query, respectively. Therefore, another interpretation of the adaptive models that we study in this work is that they resemble transductive meta-learning models, but they are given only the unlabeled query and not the labeled support set.
X1,
Px
[5144 29
In the next section, we expand on this connection to develop the ARM framework, which then allows us bring forward tools from meta- reta-learning to tackle domain shift problems.
Adaptive Risk Minimization
In this section, formally describe the ARM framework, which defines objective for training adaptive models to tackle domain shift. Furthermore, we propose a general meta-learning algorithm as well as specific methods for optimizing the ARM objective. In [Section 5] we test these ARM methods on problems for which unlabeled adaptation can be leveraged for better test performance.
4.1 Devising the ARM objective
We wish learn adaptive model f X
K - > to tackle domain shift. As noted, meta-learning methods for labeled adaptation study a similar form of model, and a common approach in many of these methods is to define f such that it is composed of two parts: first, a learner which ingests the data and produces parameters, and second, a prediction model which uses these parameters to make predictions 74 [8.
We will follow a similar strategy which, as we will discuss in subsection 4.2 allows us to easily extend and design meta-learning methods towards our goal.
In particular, we will decompose the model f into two modules: standard prediction model g(::0) : X - y, that is parameterized by 6 E θ and predicts y given x. and an adaptation model h(.,- :)) : θ x x² + θ, which is parameterized by 6. h takes in the prediction model parameters 0 and K unlabeled data points and uses the K points to produce adapted parameters l'. This is analogous the learner in meta-learning. however, h adapts the model parameters using only unlabeled data. We defer the discussion of how to instantiate h to subsection 4.2
The ARM objective is to optimize & and 0 such that h can adapt g using unlabeled data sampled according to a particular domain z. This can be expressed as the optimization problem
K
min
E(i,E(0,0 E(, = Ey. Exxuis is 0.6
1 K
(g((x;00), mk)
where 0' PP=ho =1(0,xx, (100xx...... xxiΦ). (2)
1=
Note that & is the empirical form of the adaptive risk in Equation 1 for the form of f we have defined. Mimicking the generative process from Section that we assume generated the training data, p- is
a categorical distribution over (1,..., ... ST which places uniform probability mass on each training domain, and Pxyl= assigns uniform probability to only the training points within a particular domain. As we have established theoretically, we expect the trained models to perform well at test time if the
2 For some meta- learning methods, the learner does not take as input the unadapted model parameters 74] and we also devise methods of this form. In the formalism above, these methods simply ignore the input 6.
5

---

test domains are sampled independently and identically - i.e. from the same distribution over Pxy as the training domains. In practice, similar to how meta-learned few shot classification models are evaluated on new and unseen meta-test classes 74 18 we empirically show in Seectiion that the trained models can generalize to test domains that are not sampled identically to the training domains.
4.2 Optimizing the ARM objective
Algorithm presents a general meta-learning approach for optimizing the ARM objective. As described above, outputs updated pa- rameters 0' using an unlabeled batch of data (line 5). This mimics the adaptation proce- dure test time, where we do not assume
Algorithm Meta Meta-Learning for ARM
// Training procedure
Require: training steps T, batch size K, learning rate 7) 1: Initialize: 0,6
2: fo do
access to labels (lines 7-8). However, the training update itself does rely on the labels
3: Sample z uniformly from training domains
4: Sample (xx, yk) ~ p(-,- (z) for f 1,.... K
(line 6). We assume that h is differentiable with respect to its input 0 and 6, thus we use gradient updates on both 0 and é to optimize for post adaptation performance on mini batch of data sampled according to a partic- ular domain 2. In practice, we also sample mini batches of domains, rather than just one domain (as written in line 3), to provide a better gradient signal for optimizing & and 0.
5: 8' + h(8, Xi...... xxi
6:
() < (0. - (0.0) Ex=r ((g(xe; Mb)
// Test time adaptation
procedure
Require: 0, o. test batch XI,..,, .K
7: <- h(00, x1...., ........... xk: 6)
8: Ux + g(xx:00) 0') for I....,
Together, Equation 2] and Algorithm I shed light on a number of ways to devise methods for solving the ARM problem. First, we can extend meta-learning paradigms to the ARM problem setting, and any paradigm in which the adaptation model h can be augmented to operate on unlabeled data is readily applicable. As an example, we propose the ARM-CML method, which is inspired by recent works in contextual meta-learning (CML) 20 57. Second, we can enhance prior unlabeled adaptation methods by incorporating meta meta-training phase that allows the model to better leverage the adaptation. To this end, we propose the ARM- -BN method, based on the general approach of Third, we can
adapting using batch normalization (BN) statistics of the test inputs 4 53 33
50.
incorporate existing methods for meta-learning unlabeled adaptation to solve domain shift problems.
We demonstrate this by proposing the ARM-LL method, which is based on the robotic imitation learning method from Yu et al. [0 which adapts via a learned loss (LL). All of these methods are straightforward extensions of existing meta-lcarning and adaptation methods, and this is intentional we aim to show how existing tools can be readily adapted to tackle domain generalization problems. We summarize the methods here and refer the reader to Appendix B for complete details.
ARM-CML, In ARM-CML, the parameters & of h define the weights of a context network (- Ø): X - IR,, parameterized by the adaptation model parameters d. We also instantiate the model with a prediction network (. :) : X' x TR. -> Y/, parameterized by When given mini batch of inputs, foom processes each example X in the mini batch separately and outputs
C₁ E R for 1 1, .......... K, which averaged together into context K Σk Ck- D is hyperparameter, and in our experiments, we choose D to be the dimensionality of x, such that we concatenate each image X2 and the context along the channel dimension to produce the input to freed In other words, free processes each X separately to produce an estimate of the output ÿt. but it additionally receives as input. In this way, from can provide information about the entire batch of K unlabeled data points to fred for predicting the correct outputs.
Note that the difference between ARM-CML and prior contextual meta-learning approaches is that, in prior approaches, the context network processes both inputs and outputs to produce each Ck- ARM-CML is designed for the domain generalization setting in which we do not assume access to labels at test time, thus we meta-train for unlabeled adaptation performance at training time.
ARM-BN. ARM-BN is particularly simple method that is applicable for any model that has BN Practically, training g via ARM-BN follows the same protocol as loffe and Szegedy except for two key differences: first, the training batches are sampled from single domain, rather than from the entire dataset, and second, the normalization statisties are recomputed at test time rather than using training running average. As noted, this second difference has been explored
layers 31
31

---

by several works method for test time adaptation, but the first difference is novel to ARM-BN. Following Algorithm I ARM-BN defines meta-training procedure which g learns to adapt i.e., compute normalization statistics using batches of training points sampled from the same domain. We empirically show in Section 3] that, for problems where BN adaptation already has a favorable inductive bias, such as for image classification, further introducing meta-training boosts its performance. We believe that other test time adaptation methods, such as those based on optimizing surrogate losses [69 75, may similarly benefit from their corresponding meta-training procedures.
At high level, ARM-BN operates in a simi- lar fashion to ARM-CML, thus we group these methods together into the umbrella of contextual approaches, shown in Figure 2 (top). The inter-
X;
pretation of ARM-BN through the contextual approach is that h replaces the running statistics used by standard BN with statistics computed on the batch of inputs, which then serves as the context c. Thus, for ARM-BN, there is no con- text network, and h has no parameters beyond the model parameters 8 involved in computing BN statisties. The model 9 is again specified via a prediction network fpeed: which must have BN layers. BN typically tracks running average of the first and second moments of the activations in these layers, which are then used at test time. ARM-BN defines h such that it swaps out these moments for the moments computed via the acti- vations on the test batch, thus giving us adapted parameters 0' if we view the moments as part of the model parameters. This method is remark- ably simple, and in deep learning libraries such as PyTorch [52], implementing ARM-BN in- volves changing single line of code. However, as shown in Seection 5] this method also performs very well empirically, and the adaptation effec- tiveness is further boosted by meta-training.
X
grackin updata
FaCCa (1at
04
Figure 2: In the contextual approach (top), x1,...; XK are summarized into context c, and we propose two methods for this summarization, either through a separate context network or us- ing batch normalization activations in the model itself. can then be used by the model to infer additional information about the input distribution. In the gradient based approach (bottom), an unla- beled loss function C is used for gradient updates to the model parameters, in order to produce pa- rameters that are specialized to the test inputs and can produce more accurate predictions.
ARM-LL. ARM-LL, depicted in Figure 2 (bottom), follows the gradient based meta-learning paradigm 18] and learns parameters 0 that are amenable to gradient updates on loss function in order to quickly adapt to new problem. other words, h produces /' = 0 a V CoL( XI...., : d). where a is hyperparameter. Note that the loss function L used in the gradient updates is different from the original supervised loss function in that operates on only the inputs x, rather than the input output pairs that l receives. We follow the general implementation of this approach proposed in Yu et al. [ We define to produce output features o E IR yll that are used as logits when making predictions. We then define the unlabeled loss function L to be the composition of g and loss network (:))(::)) RDI + R, which takes in the output features from g and produces scalar. We use the f2-norm of these scalars across the batch of inputs as the loss for updating 0. In other words, h(0 X1..... Xx:0 Ø) 0 - a where v [foss (g(x1; 0); Φ)..... ..... floss (g(Xx; f); Φ)].
5 Experiments
Our experiments are designed to answer the following questions:
Do ARM methods learn models that can leverage unlabeled adaptation to tackle domain shift?
How do ARM methods compare to prior methods for robustness, invariance, and adaptation?
Can models trained via ARM methods adapt successfully in the streaming test setting?
5.1 Evaluation domains and protocol
We propose four image classification problems, which we present below and describe in full detail in We also present results on datasets from the WILDS benchmark [35] in subsection 5.4
Appendix D
7

---

We believe that the problems we propose in this paper can supplement existing benchmarks for domain shift, which, as discussed above, are designed to test invariances. A key characteristic of the problems presented here is the potential for adaptation to improve test performance, and this differs from prior benchmarks such as the problems compiled by DomainBed (2). In Appendix C we compare our testbeds to DomainBed and group robustness benchmarks, and we briefly discuss the results in Gulrajani and Lopez-Paz 23
which also evaluate ARM-CML.
Rotated MNIST. We study a modified version of MNIST where images are rotated in 10 degree increments, from 0 130 degrees. We treat each rotation as separate domain, i.e., different value of z. We use only 108 training data points for each of the 2 smallest domains (120 and 130 degrees), and 324 points each for rotations 90 to 110, whereas the overall training set contains 32292 points. In this setting, we hypothesize that adaptation can specialize the model to specific domains, in particular the rare domains the training set. For each test evaluation, we generate images from the MNIST test set with certain rotation. We measure both worst case and average accuracy across domains.
Federated Extended MNIST (FEMNIST). The extended MNIST (EMNIST) dataset consists of images of handwritten uppercase and lowercase letters, in addition to digits [2]. FEMNIST is the same dataset, but it also provides the meta-data of which user generated each data point + We treat each user as a domain. We measure each method's worst case and average accuracy across 35 test users, which are held out and thus disjoint from the training users. As discussed in Section adaptation may help for this problem for specializing the model and resolving ambiguous data points.
Corrupted image datasets. CIFAR-IO-C and Tiny ImageNet-C 25 augment the CIFAR-10 [36] and Tiny ImageNet test sets with common image corruptions that vary in type and severity. The original goal of these augmented test sets to benchmark how well methods could handle these corruptions without access to any corruptions during training 225 Thus, successful methods for these problems typically have relied on domain knowledge and heuristics designed specifically for image classification. For example, prior work has shown that carefully designed test time adaptation
procedures are effective for these problems (69 63 7/) One possible reason for this phenomenon is
that convolutional networks are biased toward texture 21, which is distorted by corruptions, thus adaptation can help the model recover its performance for each corruption type.
We study whether meta-training for adaptation performance can improve upon these results. To do so, we modify the protocol from Hendrycks and Dietterich (25) to fit into the ARM problem setting by applying a set of 56 corruptions to the training data, and we define each corruption to be a domain. We use disjoint set of 22 corruptions for the test data, which mostly of different types from the training corruptions (thus, not sampled identically), and we measure worst case and average accuracy across the test corruptions. This modification allows us study, for both ARM and prior methods, whether seeing corruptions at training time can help the model deal with new corruptions at test time.
5.2 Comparisons and ablations
We compare the ARM methods against several prior methods designed for robustness, invariance, and adaptation. We describe the comparisons here and provide additional details in Appendix D
Test time adaptation. We evaluate the general approach of using test batches to compute BN statis- tics 43 (3] [33 (50). which we term BN adaptation. We also compare to test time training (TTT) 39 which adapts the model at test time using a self-supervised rotation prediction loss. These meth- ods have previously achieved strong results for image classification, likely because they constitute favorable inductive biases for improving on the true classification task 669
Ablations. We also include ablations of the ARM-CML and ARM-LL methods, which sample training batches of unlabeled examples uniformly from the entire training set, rather than sampling from a single domain B These "context ablation" and "learned loss ablation" are similar to test time adaptation methods in that they do not require training domains, thus they allow us to evaluate whether or not meta- vetaatraiinng on domain shifts is important for improved performance.
Group robustness and invariance. Sagawa et al. 600 recently proposed a state-of-the-arr method for group robustness, and we refer to this approach as distributionally robust neural networks (DRNN). Their work also evaluates strong upweighting (UW) baseline that samples uniformly from each group, and so we also evaluate this approach in our experiments. Additionally, we compare to
Note that the corresponding ablation of ARM-BN BN is simply the BN adaptation method.
8

---

Table 1: Worst case (WC) and average (Avg) top I accuracy on all testbeds, where means and standard errors are reported across three separate runs of each method. Horizontal lines separate methods that make use of (from top to bottom)): neither, training domains, test batches, or both. ARM methods consistently achieve greater robustness, measured by WC, and Avg performance compared to prior methods. *UW is identical to ERM for CIFAR-IO-C and Tiny ImageNet-C.
These results have been updated from an earlier version of the paper, primarily for CIFAR- 10-C, due to significant refactoring of the code, additional hyperparameter tuning for both the ARM methods and the prior methods, and efforts to standardize results across the authors' different computing environments and library versions. These results are reproducible from the publicly available code: https://githii www/comm////ee http://nniiimakkkkkm kmarkl arm
MNIST
FEMNIST
CIEAR-IOC
Tiny ImageNNe--C
Method
WC
Avg
WC
Avg
WC
Avg
WC
Avg
ERM
71.5 141
93.6
62.1
7.3
51.1 + 0.3
70.1 0.1
20.3 0.5
11.9 0.1
UW
1.2
95.1 0.1
65.7
80.3
0.6
DRNN
7.7
94.9+ 0.1
57.5 + 1.7
76.5 + 1.2
49.3 + 0.9
65.7 0.5
14.2 +
31.6 1.0
DANN
7.8
941
65.4
8
53.9
= 2.2
69.8
0.3
20 0.7
40.9 + 0.2
MMD
0.9
0.3
6..1=0.7
7.1
52.2
0.3
69.5 0.1
19.7
0.1
BN adaptation
7..0
9..
65.7
1:5 1.5
80.0 1: 0.5
60.6 0.3
70.9
2.03
42.8 1 0.0
TTT
81.1+003
85.4++00
68.6 0.4
84.2 + 0.1
61.5 + 0.3
71.7 0.5
27.6 + 0.5
37.7 0.3
CML ablation
18
0.2
6.8
81.6
58.8
0.2
0.6
42.5
LL. ablation
1.1
95.0 0.3
64.1 1.6
80.8 + 0.2
60.9 + 0.4
71.3 + 0.0
21.6
32.6
3.2
ARM CML
88.0:1.08
96.3
0.4
1..9
86.4 0.3
61.2 0.4
70.3 0.2
2..1
43.3 0.1
ARM-BN
8.5
9.1
64.5 3.2
83.2 2
61.7 + 0.3
72.4
0.3
0.3
43.3
0.T
ARM-LL
88.9:10.8
96.9 0.2
67.0 0.9
84.3 0.7
61.2 1 0.7
72.5 1: 0.4
0.04
35.7
0.4
domain adversarial neural networks (DANN) 19 and maximum mean discrepancy (MMD) feature learning 4II. two state-of-the-an methods for adversarial learning of invariant predictive features. For the WILDS datasets, we include the numbers reported in Koh et al. (35) for DRNN and two other invariance methods, correlation alignment (CORAL) 67 and invariant risk minimization (IRM) 3
Robustness and invariance methods assume access to training domains but not test batches, whereas adaptation methods assume the opposite. Thus, at a high level, can view the comparisons to these methods as evaluating the importance of each of these assumptions for the specified problems.
5.3 Quantitative evaluation and comparisons
The results for the four proposed benchmarks are presented in Table The best results, stratified by classes of methods. are bolded, with the single best result across all methods underlined. Across all of these problems, ARM methods increase both worst case and average accuracy compared to all other methods. ARM-CML performs well across all tasks, and despite its simplicity, ARM-BN achieves the best performance overall on the corrupted image testbeds, demonstrating the effectiveness of meta- training on top of an already strong adaptation procedure. BN adaptation and TTT are the strongest prior methods, as these adaptation procedures constitute inductive biases that are generally well suited for image classification. However, ARM methods are comparatively less reliant on favorable inductive biases and consistently attain better results. In general, we observe poor performance from robustness methods, varying performance from invariance methods, strong performance from adaptation methods, and the strongest performance from ARM methods.
When we cannot access batch of test points all once, and instead the points are observed in streaming fashion, we can augment the proposed ARM methods to perform sequential model updates. For example, ARM-CML and ARM-BN can update their average context and normalization statistics, respectively, after observing each new test point. In Figure 3] we study this test setting for the Tiny ImageNet-CC problem. We see that both models trained with ARM-CML and ARM-BN are able to achieve near their original worst case and average accuracy within observing 50 data points, well before the training batch size of 100. This result demonstrates that ARM methods are applicable for problems where test points must be observed one at a time, provided that the model is permitted to adapt using each point. We describe in detail how each ARM method can be applied to the streaming setting in Appendix B and we provide streaming results rotated MNIST in Appendix E

---

Worst Case Accuracy Tiny ImageNet-C
Average Accuracy Tiny ImageNet-C
0.35
0.30
04-
Accuracy
0.3
0.20
ARM CML (Stresming)
0.-
ARM-CML [Streamingi
0.10
ARM BN (Streaming)
ARM--N (Streaming)
ERM
ERM
AR06-BN
ARM
0.00
examples
# examples
Figure 3: In the streaming setting, ARM methods reach strong performance on Tiny ImageNet-C after fewer than 50 data points, despite using training batch sizes of 100. This highlights that the trained models are able to adapt successfully in the standard streaming evaluation setting.
Table 2: Results on the WILDS image testbeds. Different methods are best suited for different problems, motivating the need for wide range of methods. ARM-BN struggles on FMoW but performs well on the other datasets, in particular RxRxI.
iWildCam
Camelyon17
RaRxl
FMoW
Poverty Mup
Method
Acc
Macro Fl
Acc
Acc
WC
Ace
Avg Acc
WC Pearson
Peurson
ERM
71.6 + 2.5
310
1.3
70.3
+ 6.4
29.9 0.4
32.3
1.25
53.0 0.55
0.46 0.06
0.78 + 0.04
DRNN
72.7 + 2.0
23.9 +
68.4 + 7.3
23.0
30.8
0.81
52.1
0.5
0.30 +-006
0.75 + 0.07
CORAL
73.3
4.3
32.8
0.1
59.5 1 7.7
28.4 28.4
31.7
1.24
50.5
1 0.36
0.44 1.4.1.0.66
0.78 1 0.05
IRM
59.8
15.1 +
61.2
8.2
30.0
1.37
50.8
4
0.43
0.17 + 0.05
adaptation
46.4
1.0
13.8
-
88.6 14
20.0
30.2 0.26
5L.+
0.39
0.82 +
ARM-BN
70.3 + 2.4
23.2
+ 2.7
87.2 + 0.9
31.2+01
24.6 + 0.04
42.0
0.49++0.21
0.84 + 0.05
5.4 WILDS results
Finally, we present results on the WILDS benchmark ARM-BN on these testbeds. We see that, these real world distribution shift problems, different methods perform well for different problems. CORAL, a method for invariance (67), performs best on the iWildCam animal classification problem [5. whereas methods outperform ERM by significant margin on the FMoW [ or PovertyMap (79 satellite imagery problems. ARM-BN performs particularly poorly the FMoW problem. However, it performs well Poverty Map and significantly improves performance on the RxRx1 (70 problem of treatment classification from medical images. On the other medical imagery problem of Camelyon17 4 tumor identification, adaptation in general boosts performance dramatically. These results indicate the need to consider wide range of tools, including meta-learning and adaptation, for combating distribution shift.
35 in Table 2
We evaluate BN adaptation and
Discussion and Future Work
We presented adaptive risk minimization (ARM), a framework and problem formulation for learning models that can adapt in the face of domain shift at test time using only a batch of unlabeled test examples. We devised an algorithm and instantiated a set of methods for optimizing the ARM objective that meta-learns models that are adaptable different domains of training data. Empirically, we observed that ARM methods consistently improve performance terms both average and worst case metrics, as compared to number of prior approaches for handling domain shift.
Though we provided contextual meta-learning as a concrete example, number of other meta-learning paradigms would also be interesting to extend to the ARM setting. For example, few shot generative modeling objectives would be natural fit for unlabeled adaptation [6 26 Another exciting direction for future work is to explore the problem setting where domains are not provided at training time. As discussed in Appendix in this setting, we can instead construct domains via unsupervised learning techniques. Similar to Hsu et al. 28 one promising approach is to generate a diverse set of domains in order to learn generally effective adaptation strategies. Robustness and invariance methods cannot be used easily with multiple different groupings, learned or otherwise, as techniques such as group weighted loss functions 600 and domain classifiers [9 are not immediately extendable to this setup. Thus. ARM methods may be uniquely suited to be paired with domain learning.
10

---

Acknowledgments and Disclosure of Funding
MZ. thanks Matt Johnson and Sharad Vikram for helpful discussions and was supported by an NDSEG fellowship. HM is funded by scholarship from the Dr. Tech. Marcus Wallenberg Foundation for Education in International Industrial Entrepreneurship. AG was supported by an NSF graduate research fellowship. CF is a CIFAR Fellow in the Learning in Machines and Brains program. This research was supported by the DARPA Assured Autonomy and Learning with Less Labels programs.
References
[1] Alet, M. Bauza, K. Kawaguchi, N. Kuru, T. Lozano-Péree and L. Kaelbling. Tailoring: Encoding inductive biases by optimizing unsupervised objectives prediction time. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[2] A. Antoniou and A. Storkey. Learning to learn via self-critique. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[3] M. Arjovsky. L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 02.9999,,,,, 2019.
[4] P. Bandi, O. Geessink, Q Manson, M. Van Dijk, M. Balkenhol, M. Hermsen, B. Bejnordi, B. Lee, K. Paeng, A. Zhong, et al. From detection of individual metastases to classification of lymph node status at the patient level: the CAMELYON17 challenge. IEEE Transactions on Medical Imaging, 2018
[5] S. Beery, E. Cole, and A. Gjoka. The iWildCam 2020 competition dataset. arXiv preprint arXiv:2004. 10340, 2020.
[6] S. Bengio, Y. Bengio, J. Cloutier, and Gecsei. On the optimization a synaptic learning rule. In Optimality in Artificial and Biological Neural Networks, 1992.
[7] G. Blanchard, G. Lee, and C. Scott. Generalizing from several related classification tasks to a new unlabeled sample. In Advances in Neural Information Processing Systems (NIPS), 2011.
[8] G. Blanchard, A. Deshmukh, U. Dogan, G. Lee, and Scott. Domain generalization by marginal transfer learning. Journal of Machine Learning Research (JMLR), 2021.
[9] S. Caldas, S. Duddu, P. Wu, T. J. Konečný, H. McMahan, V. Smith, and A. Talwalkar. LEAF:
A benchmark for federated settings. In Workshop on Federated Learning for Data Privacy and Confidentiality, 2019.
[10] F. Carlucci, L. Porzi, B. Caputo, E. Ricci, and S. Rota Bulò. AutoDIAL: Automatic domaln alignment layers. In IEEE International Conference on Computer Vision (ICCV), 2017.
[11] G. Christie, N. Fendley, J. Wilson, and R. Mukherjee. Functional map of the world. In Conference Computer Vision and Pattern Recognition (CVPR), 2018.
[12] G. Cohen, Afshar, Tapson, and A. van Schaik. EMNIST: An extension of MNIST to handwritten letters. arXiv preprint
2017.
[13] G. Csurka, Domain adaptation for visual applications: comprehensive survey. arXiv preprint arXiv:1702.055: 2017.
[14] H. Daumé III. Frustratingly easy domain adaptation. In Conference of the Association for Computational Linguistics (ACL.), 2007.
[15] Q. Dou, D. Castro, K. Kamnitsas, and B. Glocker. Domain generalization via model-agnostic learning of semantic features. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[16] H. Edwards and Storkey. Towards neural statistician. In International Conference on Learning Representations (ICLR), 2017.

---

[17] C. Fang, Y. Xu, and D. Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In IEEE International Conference on Computer Vision (ICCV), 2013.
[18] C. Finn, P. Abbeel, and S. Levine. Model Model-agnostic meta-learning for fast adaptation of deep networks, International Conference on Machine Learning (ICML), 2017.
[19] Y. Ganin and Lempitsky. Unsupervised domain adaptation by backpropagation. In Interna- tional Conference on Machine Learning (ICML), 2015.
[20] M. Garnelo, D. Rosenbaum, C. Maddison, T. Ramalho, D. Saxton, M. Shanahan, Y. Teh,
D. Rezende, and S. Eslami. Conditional neural processes. International Conference on Machine Learning (ICML), 2018.
[21] R. Geirhos, P. Rubisch, C. Michaelis, M. Bethge, F. Wichmann, and W. Brendel. Imagenet- trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations (ICLR), 2019.
[22] B. Gong, Y. Shi, F. Sha, and K. Grauman. Geodesic flow kernel for unsupervised domain adaptation. In Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
[23] Gulrajani and Lopez-Paz. search of lost domain generalization. International Conference on Learning Representations (ICLR), 2021.
[24] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
[25] D. Hendrycks and T. Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In International Conference on Learning Representations (ICLR), 2019.
[26] L. Hewitt, M. Nye, A. Gane, Jaakkola, and Tenenbaum. The variational homoencoder: Learning to learn high capacity generative models from few examples. In Conference on Uncertainty in Artificial Intelligence (UAI), 2018.
[27] S. Hochreiter, A. Younger, and Conwell. Learning to learn using gradient descent. In International Conference on Artificial Neural Networks (ICANN), 2001.
[28] K. Hsu, S. Levine, and C. Fino. Unsupervised learning via meta-learning. In International Conference on Learning Representations (ICLR), 2019.
[29] Hu, Moreno, Y. Xiao, X. Shen, G. Obozinski, N. Lawrence, and A. Damianou. Empirical Bayes transductive meta-learning with synthetic gradients. In International Conference Learning Representations (ICLR), 2020.
[30] W. Hu, Niu, I. Sato, and M. Sugiyama. Does distributionally robust supervised learning give robust classifiers? In International Conference on Machine Learning (ICML), 2018.
[31] S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International Conference on Machine Learning (ICMI.), 2015.
[32] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with Gumbel-softmax. In International Conference on Learning Representutions (ICLR), 2017.
[33] A. Kaku, S. Mohan, A. Parnandi, H. Schambra, and C. Fernandez-Granda. Be like wa- ter: Robustness extraneous variables via adaptive feature normalization. arXiv preprint arXiv:2002.04001 2020.
[34] D. Kingma and M. Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations (ICLR), 2014.
[35] Koh, S. Sagawa, H. Marklund, S. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. Phillips, I. Gao, T. Lee, E. David, I. Stavness, W. Guo, B. Earnshaw, I. Haque, S. Beery, J. Leskovee, A. Kundaje, E. Pierson, Levine, C. Finn, and P. Liang. WILDS: A benchmark of in ti-the-wild distribution shifts. In International Conference on Machine Learning (ICML), 2021.
12

---

[36] A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[37] A. Kumagai and T. Iwata. Zero-shot domain adaptation with domain semantic descriptors. arXiv preprint arXiv: (.07.00977 2018.
[38] D. Lazer, R. Kennedy, G. King, and A. Vespignani. The parable of Google flu: Traps in big data analysis. Science, 2014.
[39] D. Li, Y. Yang. Y. Song, and T. Hospedales. Deeper, broader and artier domain generalization. In IEEE International Conference on Computer Vision (ICCV), 2017.
[40] D. Li, Y. Yang, Y. Song, and T. Hospedales. Learning to generalize: Meta-learning for domain generalization. In AAA! Conference on Artificial Intelligence (AAAI), 2018.
[41] H. Li, S. Pan, Wang, and A. Kot. Domain generalization with adversarial feature learning. In Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
[42] X. Li, Q. Sun, Y. Liu, Q. Zhou, S. Zheng. T. Chua, and B. Schiele. Learning to self- telf-train for semi-supervised few-shot classification. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[43] Y. Li, N. Wang, J. Shi, J. Liu, and X. Hou. Revisiting batch normalization for practical domain adaptation. In International Conference on Learning Representations Workshop (ICLRW), 2017.
[44] Y. Li, X. Tian, M. Gong, Y. Liu, T. Liu, K. Zhang, and D. Tao. Deep domain generalization via conditional invariant adversarial networks. In European Conference on Computer Vision (ECCV), 2018
[45] Z. Lipton. Y. Wang, and A. Smola. Detecting and correcting for label shift with black box predictors. In International Conference on Machine Learning (ICML), 2018.
[46] Y. Liu, J. Lee, M. Park, S. Kim, E. Yang, S. Hwang, and Y. Yang. Learning to propagate labels: Transductive propagation network for few-shot learning. International Conference on Learning Representations (ICLR), 2019.
[47] C. Maddison, A. Mnih, and Y. Teh. The Concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations (ICLR), 2017.
[48] L. Metz, N. Maheswaranathan, B. Cheung, and J. Sohl-Dickstein. Meta-learning update rules for unsupervised representation learning. In International Conference on Learning Representations (ICLR), 2019.
[49] K. Muandet, D. Balduzzi, and B. Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning (ICMI., 2013.
[50] Z. Nado, S. Padhy, D. Sculley, A. D' A'muu, B. Lakshminarayanan, and J. Snoek. Evaluat- ing prediction-time batch normalization for robustness under covariate shift. arXiv preprint arXiv:2006. 10963, 2020.
[51] A. Nichol, Achiam, and J. Schulman. On first-order meta-learning algorithms. arXiv preprint arXiv:1803.0220 2018.
[52] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, Desmaison, A. Kopf, Yang, De Vito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[53] X. Peng. Q. Bai, X. Xia, Z. Huang, K. Saenko, and B. Wang. Moment matching for multi-source domain adaptation. In IEEE International Conference on Computer Vision (ICCV), 2019.
[54] Quiñonero Candela, M. Sugiyama, Schwaighofer, and N. Lawrence. Dataset Shift in Machine Learning. The MIT Press, 2009.
13

---

[55] S. Ravi and H. Larochelle. Optimization as model for few-shot learning. In International Conference on Learning Representations (ICLR). 2017.
[56] M. Ren, E. Triantafillou, S. Ravi, Snell, K. Swersky, J. Tenenbaum, H. Larochelle, and R. Zemel. Meta-learning for semi-supervised few-shot classification. In International Confer- ence on Learning Representations (ICLR), 2018.
[57] Requeima, J. Gordon, Bronskill, S. Nowozin, and R. Turner. Fast and flexible multi-task classification using conditional neural adaptive processes. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[58] D. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning (ICML), 2014.
[59] A. Royer and C. Lampert. Classifier adaptation at prediction time. In Conference on Computer Vision and Pattern Recognition (CVPR), 2015.
[60] S. Sagawa, P. Koh, T. Hashimoto, and P. Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference Learning Representations (ICLR), 2020.
[61] A. Santoro, S. Bartunov, M. Botvinick, D. Wierstra, and T. Lillicrap. Meta-learning with memory-augmented neural networks. In International Conference on Machine Learning (ICML), 2016.
[62] Schmidhuber. Evolutionary principles in self-referential learning. Diploma thesis, Institut f. Informatik, Tech. Univ. Munich, 1987.
[63] S. Schneider, F. Rusak, L. Eck, O. Bringmann, W. Brendel, and M. Bethge. Improving robustness against common corruptions by covariate shift adaptation. Advances Neural Information Processing Systems (NeurIPS), 2020.
[64] H. Shimodaira. Improving predictive inference under covariate shift by weighting the log- likelihood function. Journal of Statistical Planning and Inference (JSPI), 2000.
[65] Snell, K. Swersky, and R. Zemel. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems (NIPS), 2017.
[66] M. Sule and J. Matas. Improving CNN classifiers by estimating test-time priors. In IEEE International Conference on Computer Vision (ICCV), 2019.
[67] B. Sun, J. Feng, and K. Saenko. Return of frustratingly easy domain adaptation. In AAAI Conference Artificial Intelligence (AAAI), 2016.
[68] S. Sun, H. Shi, and Y. Wu. A survey of multi-source domain adaptation. Information Fusion, 2015.
[69] Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt. Test-time training with self- supervision for generalization under distribution shifts, In International Conference on Machine Learning (ICML), 2020.
[70] J. Taylor, B. Earnshaw, B. Mabey, M. Victors, and J. Yosinski. RxRxI: An image set for cellular morphological variation across many experimental batches. In International Conference on Learning Representations (ICLR), 2019.
[71] S. Thrun and L. Pratt. Learning to Learn. Springer Science & Business Media, 1998.
[72] E. Tzeng, J. Hoffman, K. Saenko, and T. Darrell. Adversarial discriminative domain adaptation. In Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
[73] V. Vapnik. Statistical Learning Theory. Wiley New York, 1998.
[74] O. Vinyals, C. Blundell, T. Lillicrap, K. Kavukcuoglu, and D. Wierstra. Matching networks for one shot learning. In Advances in Neural Information Processing Systems (NIPS), 2016.
14

---

[75] D. Wang, E. Shelhamer, S. Liu, B. Olshausen, and T. Darrell. Tent: Fully test-time adaptation by entropy minimization. In International Conference on Learning Representations (ICLR), 2021.
[76] G. Wilson and D. Cook. A survey of unsupervised deep domain adaptation. ACM Transactions on Intelligent Systems and Technology (TIST), 2020.
[77] M. Wu, K. Choi, N. Goodman, and S. Ermon. Meta-amortized variational inference and learning.
In AAAI Conference Artificial Intelligence (AAAI), 2019.
[78] Y. Yang and T. Hospedales. A unified perspective on
and multi-task learning. In International Conference on Learning Representations (ICLR), 2015.
[79] C. Yeh, A. Perez, A. Driscoll, G. Azzari, Z. Tang, D. Lobell, S. Ermon, and M. Burke. Using publicly available satellite imagery and deep learning to understand economic well-being in Africa. Nature Communications, 2020.
[80] T. Yu, C. Finn, A. Xie, S. Dasari, P. Abbeel, and S. Levine. One-shot imitation from observing humans via domain comain-adaptive meta-learning, Robotics: Science and Systems (RSS), 2018.
[81] R. Zhang, T. Che, Z. Ghahramani, Y. Bengio, and Y. Song. MetaGAN: An adversarial approach few-shot learning. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
15


Reviews:
[{"review_id": "zDqs2oevZR4", "reviewer": "Reviewer_dgsw", "summary": "This work proposes to modify MAML where the inner optimization problem is adapting to a different distribution in an unsupervised manner rather than finetuning on a new task. That is, the inner step must reduce loss but without the labels.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I find this work to be quite limited in novelty, and it seems like a very shallow application. Given the similarities between domain generalization, domain adaptation, and meta-learning/multi-task learning, it's not insightful to just apply the MAML objective, and so the contribution of this work seems to be trying out a few types of adaptation to find which ones work reasonably well.\n\nIn fact, I think this work is better described as \"MAML when the test task samples are unlabeled\" rather than a method for distribution shift. The proposed algorithm is _literally exactly MAML_, where the inner gradient steps are replaced with alternative adaptation techniques for when the labels are unavailable. **The fact that this work is basically just \"how can we slightly modify MAML\" is I think the primary factor for me, but below I include some additional discussion which is also relevant and I think essential to properly address.**\n\nI feel the discussion of general \"distribution shift\" is poorly framed. While Blanchard et al. did indeed consider a prior over tasks, the current formulation of domain generalization is a much more difficult problem of generalizing in a _minimax_ sense, when the statistics across domains may have very little in common (further, note that that work provides actual theoretical guarantees, where this work provides none).\n\nThis paper heavily conflates \"group shift\" with general distribution shift, such as in the problem of domain generalization. In the title and abstract, the authors indicate that this method is meant for dealing with group shift---this is not the same thing as the typical domain generalization setup or distribution shift in general. This is a _very_ specific type of change at test time where the mixture components of a set of fixed sub-populations (all of which are observed at train time) changes. That's fine, but I think the introduction discussing domain generalization and general distribution shift is quite misleading because this paper addresses a very niche problem that's nowhere near as difficult as the problems discussed in the introduction. The experiments show that the algorithm is successful at unsupervised meta-learning, or adapting to new domains which share a _great deal of structure_ with the training data. The experiments do _not_ show that the algorithm is successful at generalization under general distribution shift.\n\nAs further evidence, and particularly damning, I note that in Appendix C the authors discuss the performance of ARM on the DomainBed benchmark. See the results here:\nhttps://github.com/facebookresearch/DomainBed/blob/master/domainbed/results/2020_10_06_7df6f06/results.png\nI find it quite disingenuous that the authors write: \"ARM-CML handily outperforms ERM and all prior methods on the colored MNIST benchmark\", as CMNIST is an extremely contrived and unrealistic dataset. Aside from the two MNIST-based datasets (the two easiest ones), on _all other datasets in the DomainBed benchmark_ ARM is nowhere near the top, and on average it is quite in the middle of the pack, failing to outperform ERM. **Cherry-picking results on a benchmark suite defeats the purpose, and gives me the impression that this comparison was made in bad faith.**\n\nI think \"MAML when the test task labels are unavailable\" is an ok idea, lacking in novelty but not totally without merit. I also personally feel that group shift is not a realistic assumption for what to expect at test time without some a priori structural knowledge of the task we are trying to solve. Regardless, the experimental results demonstrate its ability to perform well when there is heavily shared structure among the domains (as might be expected, given the empirical success of MAML). But the algorithm does not work for more difficult distribution shift (and there is *no theoretical intuition or justification whatsoever*), so really this is just a slight tweak to an existing approach to meta-learning.\n\nThis paper would greatly benefit from an extensive reframing, and the results---positive _and_ negative---need to be honestly and frankly reported. \n\n--------------------\n**Update after rebuttal**. The authors have indicated that they will fix the issues that I surfaced, an after reconsideration I consider this to be marginally above the publication threshold.\n\nI think it is **essential** that these changes are actually implemented. This method works well, but only in *very* specific settings, with *much* stronger assumptions than previously made. As an example, the comparison in table 1 stratifies by assumption, but then still only bolds the top performing algorithm. This gives the casual reader the impression that ARM achieves SOTA, when in reality it is simply not being compared to any other algorithms with the same access to information at test time. This fact should be conveyed **much more clearly**, in addition to the other issues I found with the experimental reporting and general discussion of domain generalization and group shift.\n\n**In fact, I think the algorithm suggested here is more akin to unsupervised domain adaptation than domain generalization**... maybe the comparisons should be updated to reflect this.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Sure", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "sqdP0TmLEBT", "reviewer": "Reviewer_sJVX", "summary": "This paper proposes a framework of adaptive risk minimization (ARM), which adapts to unlabeled instances at the test phase \nby pretraining with labeled instances obtained in many groups (domains or users). \nTo adapt to unlabeled instances, the authors proposed three approaches (ARM-CML, ARM-BN, and ARM-LL).\nThe effectiveness of the proposed method was shown using real-world image datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- Problem formulation is interesting.\n- Experiments with sequential arriving data are interesting.\n\nCons:\n- Some related works are missing.\n- Experiments can be improved.\n\nDetails:\n\nThis paper proposed an interesting and useful problem setting.\nUnlike the standard meta-learning methods that require labeled instances as support set at the test phase, \nthis approach requires unlabeled instances at the test phase to adapt to new groups.\nAlthough the proposed method for adaptation is straightforward and simple, I like this research direction.\nBut, I think that there are some improvements that can be made.\n\nFirst, about the experiment, how does the proposed method compare to the unsupervised domain adaptation, \nwhich require both labeled training data and unlabeled test data during training?\nI think that this comparison gives us a better insight into the proposed approach.\nIn addition, this paper uses only image datasets in the experiments.\nIf the authors can show the effectiveness of the proposed method on datasets other than image datasets,\nit will strengthen the contribution of the paper (although it is not always necessary).\n\nSecond, the problem setting of this paper is not completely new.\nFor example, the paper [1] proposes almost the same problem setting and a method similar to the ARM-CML.\nThis method uses only unlabeled instances to adapt to new target domains at the test time by learning labeled instances in multiple source domains.\nFor adaptation, the deepset network, which outputs a latent context vector from the set of unlabeled instances, was used.\nI think it is necessary to discuss in more depth the differences between these methods and the advantages of the proposed method.\n\n[1] Kumagai, Atsutoshi, and Tomoharu Iwata. \"Zero-shot domain adaptation without domain semantic descriptors.\" arXiv preprint arXiv:1807.02927 (2018).\n\n===after rebuttal===\nThank you for your answer. \nI think the additional experimental comparison will be useful. I would like to raise the score from 6 to 7, provided that the authors reflect these results in the revised version and discuss them in more depth.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concerns", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "o2JiyqNUzC6", "reviewer": "Reviewer_dMsf", "summary": "The authors propose a method to adapt to group distribution shifts using meta-learning. Their proposed method is general for adapting to different test distributions. Empirical results show that their method is effective in improving the classification accuracies of common datasets by 1-4%. \n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The key observation of the authors is that changes in the underlying group distributions in input x can tell us something about the distribution of labels, and this is a natural problem that occurs for example with the different handwriting styles of different users in handwriting recognition. \n\n- I think the authors can provide more details on the different adaptation approaches in ARM-CML, ARM-BN, and ARM-LL, especially ARM-LL. This is an important part of the algorithm and crucial for understanding why it works, but only half a page is used for these methods. \n\n- When I look at Algorithm 1, apart from the sampling of z from different groups, there is nothing specific in the algorithm that is tailored for group distribution shifts. It can be applied to other types of test distribution shifts as well, as long as we have a sample of the test data. This is both a pro and a con: it is general enough to apply to other distribution shifts, but on the other hand there does not seem to be a lot of new insights and contributions to the base meta-learning approach. \n\n- For the empirical evaluations, the results are in general good and contain small improvements over previous adaptation approaches. \n\n- Overall this is a borderline paper to me because although the problem of adapting to group distribution shifts is interesting, the current paper does not add a lot on top of a direct application of the meta-learning method.  \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "3haqV6lyjJJ", "reviewer": "Reviewer_6LKK", "summary": "This paper proposes a framework for unsupervised multi-source domain adaptation in which training samples are grouped according to which domain they come from. (The paper uses \"groups\" to refer to the domains, but I will simply use \"domains\".) The authors assume that predictions can be adapted using the information of the input distribution of the target domain, and the learner is allowed to use a batch of unlabeled test samples at test time in order to adapt predictions.\n\nThe authors approach this problem by learning an adaptation model h that inputs a batch of unlabeled samples and transforms the parameters of a prediction model g (Algorithm 1). They present three concrete methods based on this approach and compare them with previous methods using a few benchmark datasets, which confirms the effectiveness of the approach.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Overall evaluation\nThe clarity of the paper could be improved, but the core idea of the proposed approach is interesting and the experimental results look promising. I recommend accepting the paper provided that the authors address clarity issues.\n\n# Originality\n* As far as I understand, the problem setup can be viewed as multi-source unsupervised domain adaptation, and this itself is not new.\n* The proposed approach of learning to adapt using a batch of unlabeled samples seems novel as far as I can see after seeing several papers.\n\n# Quality\n* The paper provides good discussions on related work.\n* The technical part is clear except for the concrete algorithms.\n* The general algorithm is natural and reasonable although the paper does not provide any theoretical guarantees.\n* On the other hand, the concrete algorithms look quite heuristic and it would be nice if there were theoretical analysis.\n* The experiments are well-designed and the results look convincing.\n\n# Clarity\n* Section 1 and Section 2 are somewhat too abstract and unclear to me. In fact, I could not understand the problem setup until I read Section 4. A reason may be the authors try to avoid technical terminology and give high-level arguments in the introduction part, which had unfortunately negative impact on my reading experience.\n* For example, the authors use \"groups\", which is too general and ambiguous, for describing what I would call \"domains\".\n* Also, the introduction does not seem to explain well what the proposed method really does.\n* I suggest defining the problem setup a little bit more formally in early part of the paper so that readers can have clear understandings what the paper is talking about.\n* The descriptions of the concrete methods are hard to understand without looking at the supplementary material.\n\n# Significance\n* The problem setup is slightly more restrictive than the standard unsupervised domain adaptation in that we need samples from multiple source-domains (groups).\n* However, there may be many real-world applications indeed satisfying the requirement such as the examples that the authors present in the paper.\n* The idea of incorporating information of the input distribution into adaptation is interesting and seems reasonable. It does put strong assumptions like covariate shift but seems to be solving more general problems.\n* The experimental results look convincing and indeed show superiority of the proposed method.\n\n# Other comments\n* It would be nice to add discussions about previous work on multi-source unsupervised domain adaptation and zero-shot domain adaptation.\n\n\n-----\n\n# Update after the authors' response\nI have read the other review comments and the authors' response. My opinion has not changed, and I would like to keep my score at 7. I hope the authors improve the clarity of the paper and the experiment part, as discussed in the rebuttal.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper does not mention societal impact although I have no suggestions on that.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to apply a meta-learning method (similar to MAML) to solve the domain generalization problem. After spending some time reading the paper and trying to understand the main idea, I agree with Reviewer dMsf, sJVX, and dgsw that there is a limited novelty in terms of both methodology and theory. Firstly, from the domain generalization perspective, the idea of learning models that can adapt to a collection of unlabeled data at test time has appeared as early as in [1, 2] and subsequent works in this direction. In these works, the models also take as input the empirical marginal distributions through the kernel function on them. Hence, the key distinction here is the use of deep learning trained in a meta-learning style as a way to equip the models with adaptation capability. Secondly, from the meta-learning perspective, the proposed idea is very similar to the classic MAML, as also pointed out by Reviewer dgsw. The key difference from the standard setting is the use of unlabeled data alone. Lastly, the authors seem to miss several important works in both domain generalization and meta-learning when discussing the related works in Section 3. I encourage the authors to properly conduct a literature review as this is an increasingly important area of research. In the final version, I encourage the authors to improve the clarity and to better highlight the key contributions of this work.

Nevertheless, there is a positive consensus among the expert reviewers that the experimental results are promising, and that the rebuttal has adequately addressed their concerns. The authors are willing to incorporate the reviewers' suggestions in improving the final version of this work. Hence, I recommend acceptance for publication at NeurIPS2021 as a poster.

Remark: Line 75-76: "Invariance methods would fare no better, as there is no feature space that can separate the ambiguous points and make them classifiable". This is indeed misleading. One can in fact separate the ambiguous points if the group/domain information is incorporated into the feature representation as has been commonly done in the literature. 

- [1] G. Blanchard, G. Lee, and C. Scott. Generalizing from several related classification tasks to a new unlabeled sample. In Advances in Neural Information Processing Systems 24, pages 2178–2186, 2011.

- [2] Muandet, K., Balduzzi, D., Schölkopf, B., Domain generalization via invariant feature representation. In: Proceedings of the 30th International Conference on Machine Learning (ICML13). pages 10–18, 2013. 
Decision: Accept (Poster)