Adaptive Risk Minimization: Learning to Adapt to Domain Shift

1 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 [ 54 , 38 ] . 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 setting involves dealing with shift . In Figure 1 , we visualize a batch of 50 examples from a test user , and we highlight an ambiguous example which may be either a “ 2 ” ( written with a loop ) or an “ a ” ( in the double-storey style ) depending 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-storey ) from this user . Can we somehow leverage this unlabeled data to better handle test shifts caused by new users ? ⇤equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 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 . [ 35 ] ) . However , prior benchmarks for domain generalization and similar settings typically center around invariances – i.e. , in these benchmarks , there is a consistent input-output 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 [ 23 ] ) . Thus , prior methods aim for generalization to shifts by discovering this relationship , through techniques such as robust optimization and learning an invariant feature space [ 41 , 3 , 60 ] . 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 , e.g. , the handwriting classification example above . 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 a 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 as sentences or paragraphs . Unlabeled adaptation has been shown empirically to be useful for distribution shift problems [ 69 , 63 , 75 ] , such as for dealing with image corruptions [ 25 ] . Taking inspiration from these findings , we propose and evaluate on a number of problems , detailed in Section 5 , for which adaptation is beneficial in dealing with domain shift . 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 a set of methods that , given a set of training domains , meta-learns a model that is 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 [ 9 ] 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 designed with invariance in mind [ 3 , 53 , 23 ] . 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 state-of-the-art methods by 1-4 % test accuracy . 2 Related Work . A number of prior works have studied distribution shift in various forms [ 54 ] . In this section , we review prior work in domain generalization , group robustness , meta-learning , and adaptation . Invariance and robustness to domains . As discussed above , a 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 as “ 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 to new datasets [ 17 ] , new image types [ 39 , 53 ] , and underrepresented demographics [ 60 ] . Prior benchmarks typically assume the existence of a 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 , 3 ] or optimizing for worst case group performance [ 30 , 60 ] . In particular , methods for domain generalization – sometimes referred to as multi-source domain adaptation [ 68 ] or zero shot domain adaptation [ 78 ] – have largely focused on learning invariant features [ 19 , 67 , 41 , 44 , 53 ] . Gulrajani and Lopez-Paz [ 23 ] provide a 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 Section 5 , 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 . Meta-learning . Meta-learning [ 62 , 6 , 71 , 27 ] has been most extensively studied in the context of few shot labeled adaptation [ 61 , 74 , 55 , 18 , 65 ] . Our aim is not to address few-shot recognition problems , nor to propose a novel meta-learning algorithm , but rather to extend 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 [ 40 , 15 ] , which seek to metatrain models for non adaptive generalization performance . We discuss in Section 4 how paradigms such as contextual meta-learning [ 20 , 57 ] can be readily extended using the ARM framework . 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 , 46 , 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-learning for unlabeled adaptation include Yu et al . [ 80 ] , who adapt a policy to imitate human demonstrations in the context of robotic learning ; Metz et al . [ 48 ] , who meta-learn an update rule for unsupervised representation learning , though they still require labels to learn a predictive model ; and Alet et al . [ 1 ] , 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 a prominent framework that assumes access to test examples at training time [ 13 , 76 ] , similar to transductive learning [ 73 ] . As such , most domain adaptation methods consider the problem setting where there is a single test distribution [ 64 , 14 , 22 , 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 as methods for learning invariant features [ 19 , 67 , 41 ] , and we compare to these methods in Section 5 . Adaptive methods for domain generalization include Muandet et al . [ 49 ] and Kumagai and Iwata [ 37 ] , who propose a method similar to one of the ARM methods described below . We compare to a version of this method in Appendix E. Blanchard et al . [ 7 ] and Blanchard et al . [ 8 ] 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-learning , allowing us to devise new methods in a 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 [ 59 , 45 , 66 ] and crafting favorable inductive biases for the domain of interest . For image classification , techniques such as normalizing via the test inputs [ 43 ] and optimizing self-supervised surrogate losses [ 69 ] have proven effective for adapting to image corruptions [ 25 ] . We compare to these prior methods in Section 5 and empirically demonstrate the advantage of using training domains to learn how to adapt .

Summary: This paper proposes a framework of adaptive risk minimization (ARM), which adapts to unlabeled instances at the test phase  by pretraining with labeled instances obtained in many groups (domains or users).  To adapt to unlabeled instances, the authors proposed three approaches (ARM-CML, ARM-BN, and ARM-LL). The effectiveness of the proposed method was shown using real-world image datasets.