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