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: The authors try to tackle the *distribution shift* problem with a meta learning approach. The algorithm, namely ARM, is proposed. Following regular meta learning regime, ARM uses an updated version of parameter $\theta'$ to calculate the loss for back propagation. Several specific implementations are put forward, i.e. the contextual / gradient-based methods. Experiments are performed in small and large scale datasets to demonstrate the effectiveness of the proposed algorithm. Detailed ablation study and qualitative analysis are also conducted.