Adaptive Risk Minimization: A Meta-Learning Approach for Tackling Group Shift

1 INTRODUCTION . The standard assumption in empirical risk minimization ( ERM ) is that the data distribution at test time will match the distribution at training time . When this assumption does not hold , the performance of standard ERM methods typically deteriorates rapidly , and this setting is commonly referred to as distribution or dataset shift ( Quiñonero Candela et al. , 2009 ; Lazer et al. , 2014 ) . For instance , we can imagine a handwriting classification system that , after training on a large database of past images , is deployed to specific end users . Some new users have peculiarities in their handwriting style , leading to shift in the input distribution . This test scenario must be carefully considered when building machine learning systems for real world applications . Algorithms for handling distribution shift have been studied under a number of frameworks ( Quiñonero Candela et al. , 2009 ) . Many of these frameworks aim for zero shot generalization to shift , which requires more restrictive but realistic assumptions . For example , one popular assumption is that the training data are provided in groups and that distributions at test time will represent either new group distributions or new groups altogether . This assumption is used by , e.g. , group distributionally robust optimization ( DRO ) ( Hu et al. , 2018 ; Sagawa et al. , 2020 ) , robust federated learning ( Mohri et al. , 2019 ; Li et al. , 2020 ) , and domain generalization ( Blanchard et al. , 2011 ; Gulrajani & Lopez-Paz , 2020 ) . Constructing training groups or tasks in practice is generally accomplished by using meta-data , which exists for most commonly used datasets . This assumption allows for more tractable optimization and still permits a wide range of realistic distribution shifts . However , achieving strong zero shot generalization in this setting is still a hard problem . For example , DRO methods , which focus on achieving maximal worst case performance , can often be overly pessimistic and learn models that do not perform well on the actual test distributions ( Hu et al. , 2018 ) . In this work , we take a different approach to combating group distribution shift by learning models that are able to deal with shift by adapting to the test time distribution . To do so , we assume that we can access a batch of unlabeled data points at test time – as opposed to individual isolated inputs – which can be used to implicitly infer the test distribution . This assumption is reasonable in many standard supervised learning setups . For example , we do not access single handwritten characters from an end user , but rather collections of characters such as sentences or paragraphs . When combined with the group assumption above , we arrive at a problem setting that is similar to the standard meta-learning setting ( Vinyals et al. , 2016 ) . This allows us to extend well established tools and techniques from meta-learning to address distribution shift problems . Meta-learning typically assumes that training data are grouped into tasks and new tasks are encountered at meta-test time , however these new tasks still include labeled examples for adaptation . As illustrated in Figure 1 , we instead aim to train a model that uses unlabeled data to adapt to the test distribution , thereby not requiring the model to generalize zero shot to all test distributions as in prior approaches . The main contribution of this paper is to introduce the framework of adaptive risk minimization ( ARM ) , in which models have the opportunity to adapt to the data distribution at test time based on unlabeled data points . This contribution provides a principled approach for designing meta-learning methods to tackle distribution shift . We introduce an algorithm and instantiate a set of methods for solving ARM that , given a set of candidate distribution shifts , meta-learns a model that is adaptable to these shifts . One such method is based on meta-training a model such that simply updating batch normalization statistics ( Ioffe & Szegedy , 2015 ) provides effective adaptation at test time , and we demonstrate that this simple approach can produce surprisingly strong results . Our experiments demonstrate that the proposed methods , by leveraging the meta-training phase , are able to outperform prior methods for handling distribution shift in image classification settings exhibiting group shift , including benchmarks for federated learning ( Caldas et al. , 2019 ) and testing image classifier robustness ( Hendrycks & Dietterich , 2019 ) . 2 RELATED WORK . A number of prior works have studied distributional shift in various forms ( Quiñonero Candela et al. , 2009 ) . In this section , we review prior work in robust optimization , meta-learning , and adaptation . Robust optimization . DRO methods optimize machine learning systems to be robust to adversarial data distributions , thus optimizing for worst case performance against distribution shift ( Globerson & Roweis , 2006 ; Ben-Tal et al. , 2013 ; Liu & Ziebart , 2014 ; Esfahani & Kuhn , 2015 ; Miyato et al. , 2015 ; Duchi et al. , 2016 ; Blanchet et al. , 2016 ) . Recent work has shown that these algorithms can be utilized with deep neural networks , with additional care taken for regularization and model capacity ( Sagawa et al. , 2020 ) . Unlike DRO methods , ARM methods do not require the model to generalize zero shot to all test time distribution shifts , but instead trains it to adapt to these shifts . Also of particular interest are methods for robustness or adaptation to different users ( Horiguchi et al. , 2018 ; Chen et al. , 2018 ; Jiang et al. , 2019 ; Fallah et al. , 2020 ; Lin et al. , 2020 ) , a setting commonly referred to as robust or fair federated learning ( McMahan et al. , 2017 ; Mohri et al. , 2019 ; Li et al. , 2020 ) . Unlike these works , we consider the federated learning problem setting in which we do not assume access to any labels from any test users , as we partition users into disjoint train and test sets . We argue that this is a realistic setting for many practical machine learning systems – oftentimes , the only available information from the end user is an unlabeled batch of data . Meta-learning . Meta-learning ( Schmidhuber , 1987 ; Bengio et al. , 1992 ; Thrun & Pratt , 1998 ; Hochreiter et al. , 2001 ) has been most extensively studied in the context of few shot supervised learning methods ( Santoro et al. , 2016 ; Vinyals et al. , 2016 ; Ravi & Larochelle , 2017 ; Finn et al. , 2017 ; Snell et al. , 2017 ) , i.e. , labeled adaptation . The aim of this work is to extend meta-learning paradigms to problems requiring unlabeled adaptation , with the goal of tackling distribution shift . We demonstrate in the next section how paradigms such as contextual meta-learning ( Garnelo et al. , 2018 ; Requeima et al. , 2019 ) are 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 ( Ren et al. , 2018 ; Zhang et al. , 2018 ; Li et al. , 2019 ) or the transductive learning setting ( Liu et al. , 2019 ; Antoniou & Storkey , 2019 ; Hu et al. , 2020 ) . These works do not focus on the same setting of distribution shift and all assume access to labeled data for adaptation . Prior works in meta-learning for unlabeled adaptation include Yu et al . ( 2018 ) , which adapts a policy to imitate human demonstrations in the context of robotic learning , and Metz et al . ( 2019 ) , which meta-learns an update rule for unsupervised representation learning , though they still require labels to learn a predictive model . Unlike these prior works , the ARM framework facilitates the development of meta-learning methods for quickly adapting a predictive model using unlabeled examples . 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 , similar to transductive learning ( Vapnik , 1998 ) . Some of these methods , such as importance weighting approaches ( Shimodaira , 2000 ) , only handle a single predefined shift and do not constitute test time adaptation ( Csurka , 2017 ; Wilson & Cook , 2020 ) . Certain domain adaptation methods , however , are applicable in the setting with training groups , such as methods for learning invariant features ( Ganin & Lempitsky , 2015 ; Li et al. , 2018 ) , and we compare to these methods in Section 4 . Several methods for adaptation at test time have been developed specifically for dealing with label shift ( Royer & Lampert , 2015 ; Lipton et al. , 2018 ; Sulc & Matas , 2019 ) . Other methods adapt using statistics of the test inputs ( Li et al. , 2017 ) or optimize self-supervised surrogate losses ( Sun et al. , 2020 ) , and these methods have been shown to perform well across a number of image classification domains . We also compare against these prior methods in Section 4 . 3 ADAPTIVE RISK MINIMIZATION . In this section , we first formally describe the ARM problem setting , which builds on the settings used in prior work for tackling distribution shift . The novel aspect of the ARM setting is that it is amenable to meta-learning solutions to shift , and we demonstrate this by proposing an objective for the ARM setting that resembles typical meta-learning objectives . The problem setting and objective together constitute the ARM problem formulation . We subsequently propose a general algorithm as well as specific meta-learning approaches for solving the ARM problem . 3.1 THE ARM PROBLEM SETTING . A key goal in machine learning is to develop methods that can go beyond the standard ERM setting and generalize in the face of distribution shift . Accomplishing this goal necessitates the use of additional assumptions beyond ERM , and we wish to carefully craft these assumptions such that they fulfill two properties : they are realistic and applicable to real world problems , and they allow for powerful and tractable methods . In this work , we choose two assumptions that are well established in the literature on distribution shift , in order to fulfill the first property , and we develop a novel meta-learning framework using these assumptions , thus fulfilling the second . The first assumption is that the training data are provided in groups , which , as discussed above , mirrors analogous assumptions made in group DRO ( Hu et al. , 2018 ) , federated learning ( McMahan et al. , 2017 ) , and meta-learning ( Vinyals et al. , 2016 ) , among other settings . The second assumption is that we observe batches of test points all together , rather than one point at a time . Assuming access to multiple test points has been standard in domain adaptation ( Csurka , 2017 ; Wilson & Cook , 2020 ) , which makes this assumption for training , as well as recent works studying test time adaptation ( Li et al. , 2017 ; Sun et al. , 2020 ; Wang et al. , 2020 ) . To our knowledge , these assumptions have not been considered simultaneously in prior work . However , as we detail in this section , it is their conjunction that allows us to develop meta-learning solutions to shift . In the ARM problem setting , we assume access to a training dataset that consists of N labeled data points ( x ( i ) , y ( i ) , z ( i ) ) sampled i.i.d . from the training distribution p. As noted , this differs from standard supervised learning in that we additionally observe the group z ( i ) associated with each point , which is a discrete value z ∈ { 1 , . . . , S } that can represent tasks , users , or other types of meta-data . The goal is to learn a model g ( · ; θ ) : X → Y that is parameterized by θ ∈ Θ and predicts the output y ∈ Y given the input x ∈ X . At test time , we are given batches of K unlabeled data points , where each batch is drawn from a distribution that may differ from both p and the other batch distributions , and we do not observe either y or z . For example , we can imagine a test scenario that separately considers each user ’ s images , as discussed in Section 1 .

Summary: This paper studies domain adaptation under the assumption that only unlabeled target data is available in training and the domain shift follows a special group shift. The main idea for the proposed method is having an adaptation model that takes only the unlabeled data in and output updated parameters. The proposed method also involves the test-time training, which means the adaptation model takes in unlabeled training data in training but takes in unlabeled target data in the adaptation phase. The method is called adaptive risk minimization and there are two meta-learning approaches provided, contextual and gradient-based, in the paper. In the experiment, the proposed method outperforms a limited set of baselines. The paper also discusses a few cases when the assumption is violated like the group indicators are unknown.