Distributionally Robust Models with Parametric Likelihood Ratios

1 INTRODUCTION . It is well acknowledged that modern neural network based machine learning models tend to underperform when they are evaluated on data distributions that differ from the one they were trained on . For example , machine learning model performance has been observed to degrade under train-test mismatch in topics ( Gururangan et al. , 2020 ) , demographics ( Blodgett et al. , 2016 ; Amodei et al. , 2016 ; Hovy & Søgaard , 2015 ; Grother et al. , 2019 ) , geographic regions ( Koh et al. , 2020 ) , and even data collection processes ( Beery et al. , 2018 ; Zech et al. , 2018 ; Michel & Neubig , 2018 ) . In particular , these models often perform poorly when evaluated on subpopulations , domains that are present but underrepresented in their training data ( Sagawa et al. , 2020 ) , and they can latch on to spurious correlations ( McCoy et al. , 2019 ) . This has problematic real-world consequences : when such models are deployed at large , this representation disparity can , for example , unfairly affect minority groups ( Buolamwini & Gebru , 2018 ; Hashimoto et al. , 2018 ) . This behaviour can largely be attributed to the empirical risk minimization ( ERM ) principle which underlies the majority of learning algorithms used in practice . In ERM , models are trained to minimize the average loss over a finite sample from a fixed training distribution ( Vapnik , 1992 ) , as a proxy for the expected loss on a random example drawn from the fixed , albeit unknown , data distribution . This favors models which perform well on average on a fixed training set , as opposed to models which would perform equally well on a variety of subpopulations that better reflects the diverse set of distributions that can be encountered at test time . On the other hand , Distributionally robust optimization ( DRO ) proposes an appealing alternative to ERM . In DRO , models are trained to minimize their worst case risk ( or an empirical estimate thereof computed on a finite sample , via e.g . importance weighting ) under a pre-determined family of distributions Q , called the “ uncertainty set ” ( or “ ambiguity set ” ) : LDRO ( θ ) = max q∈Q E ( x , y ) ∼q ` θ ( x , y ) . ( 1 ) In the absence of explicit information about the subpopulations of interest ( which would naturally define Q ) , it is up to the practitioner to carefully define this uncertainty set . This has been the subject of much work in the literature ( see Rahimian & Mehrotra ( 2019 ) for a survey ) . Recently , Michel et al . ( 2021 ) proposed P-DRO , a promising approach where the uncertainty set is defined by a parametric family of generative models , which allows for more flexibility in defining the uncertainty set . P-DRO shows significant improvement over comparable baselines , but it suffers from several drawbacks . First , it presupposes the availability of generative models capable of outputting exact densities , which limits its field of application to modalities where such models are readily available ( such as language models in NLP ) . Second , it is challenging to use in practice due to its reliance on a number of hyper-parameters and approximations to the objective function . In this paper , we propose a new approach for DRO , called RP-DRO , based on a key modification of the P-DRO algorithm : instead of modeling the worst-case distributions directly , we parametrize the likelihood ratio between the training distribution and the worst-case distribution . This removes the dependency on an unwieldy generative model , making the method useful for more applications . While likelihood ratio formulations of DRO have been tried in prior work ( Sagawa et al. , 2020 ) , we show that they are particularly effective for parametric , neural network based adversaries . Our approach relies on three simple ideas : a mini-batch level normalization strategy to enforce likelihood ratio constraints , a penalty-form of the KL divergence uncertainty set , and the use of simultaneous gradient updates for training . RP-DRO consistently achieves equal or better robust subpopulation accuracy compared to P-DRO and other baselines on a variety of standard benchmarks in image and text classification . In addition , we find it is both faster than P-DRO and depends on fewer hyperparameters . Additional ablation experiments demonstrate that both our mini-batch normalization strategy and simultaneous gradient updates are necessary for high performance . Finally , we perform experimental analyses to shed light on the advantages brought by parametric adversaries compared to their nonparametric counterparts . 2 BACKGROUND . In the following , we consider a model parametrized by θ ∈ Rdmodel . Our goal is to find a model which minimizes the loss function ` θ ( x , y ) on pairs of inputs and outputs ( x , y ) ∈ X × Y . For instance , x might represent images and y a categorical label . Parameters θ are estimated on a training dataset Dtrain = { ( xi , yi ) } i=1 ... Ntrain which we assume to be drawn from a training distribution p. The DRO optimization problem with uncertainty set Q is min θ max q∈Q Eq ` θ ( x , y ) . ( 2 ) Note that the DRO loss in Eq . 1 is the inner maximum of the DRO problem , and it provides an upper bound on the expected loss of the model under any distribution in the uncertainty set Q . This motivates the use the minimizer of the min-max game in Eq . 2 as a robust model . We refer to the solution of the inner maximum as the “ adversary ” from now on However this objective is only useful insofar that ( 1 ) Q covers test distributions of interest ( corresponding to different domains , demographics , etc . ) and ( 2 ) Q is not overly pessimistic . To fulfil this second condition , there should exist some model θ∗ that achieves low loss simultaneously on the test distribution as well as Q . This often requires that Q only contain distributions that are covariate shifts of the test distribution , i.e . that are such that the conditional distribution q ( y | x ) coincides with that of training distribution p ( y | x ) . 2.1 NONPARAMETRIC DRO . There is substantial existing work on nonparametric formulations of DRO , where Q is expressed as a divergence ball centered at the training distribution . This includes f -divergences ( Ben-Tal et al. , 2013 ; Hu & Hong , 2013 ; Faury et al. , 2020 ) , Wasserstein/IPM ( Sinha et al. , 2018 ; Husain , 2020 ) , moment constraints ( Delage & Ye , 2010 ; Nguyen et al. , 2020 ) , and CVaR ( Fan et al. , 2017 ; Curi et al. , 2020 ; Levy et al. , 2020 ) based uncertainty sets . These nonparametric approaches are appealing as they require very little domain-specific knowledge , have well-understood theory ( Duchi & Namkoong , 2018 ) , and optimization procedures ( e.g . Hu & Hong ( 2013 ) for KL constraints and Levy et al . ( 2020 ) for χ2 and CVaR constraints ) . Unfortunately , nonparametric DRO algorithms suffer from being overly pessimistic . Their uncertainty sets tend to include distributions that are exceedingly difficult to learn , or not representative of realworld distribution shifts . Furthermore , they often can not enforce even basic constraints such as covariate shift structures ( Duchi et al. , 2020 ; Hu et al. , 2018 ) . Group-structured DRO uncertainty sets ( Sagawa et al. , 2020 ) overcome some of these challenges , but require significant domain expertise to pre-specify target subpopulations that a model should be robust to . 2.2 PARAMETRIC DRO . Parametric DRO ( Michel et al. , 2021 ) is a method for DRO in which the uncertainty set Q is defined as a family of parametric generative models , which avoids the extreme pessimism of nonparametric DRO without the explicit specification of subpopulations . Specifically , given a generative model qψ parameterized by ψ ∈ Rdadv , the KL-constrained parametric DRO objective can be written as follows : min θ max ψ KL ( qψ‖p ) ≤κ E ( x , y ) ∼qψ ` ( x , y , θ ) . ( 3 ) As demonstrated by Michel et al . ( 2021 ) , P-DRO yields significant improvements over its nonparametric counterpart . However , the difficulty of optimizing Eq . 3 directly results in a number of approximations and additional hyper-parameters that are hard to tune . In addition , a central drawback of P-DRO is that it necessitates training an auxiliary generative model of the data . This can be difficult for several reasons . First , this limits the applicability of the method to domains with generative models that allow for exact probability computations . Moreover , even when such generative models are available , they are often more computationally demanding than their discriminative counterparts . In language models for instance , probabilities for sequences of text are obtained by iteratively producing conditional probabilities over all tokens in the vocabulary . This additional step results in considerable computational overhead compared to discriminative models . 3 PARAMETRIC LIKELIHOOD RATIO . 3.1 DRO AS A LIKELIHOOD RATIO OPTIMIZATION PROBLEM . In the situation that all distributions in Q are absolutely continuous with respect to p ( i.e . for all measurable subset A ⊂ X × Y , all q ∈ Q , q ( A ) > 0 only if p ( A ) > 0 ) the inner maximum in Eq . 2 can be rewritten purely as a function of the likelihood ratio qp E ( x , y ) ∼q ` θ ( x , y ) = E ( x , y ) ∼p q ( x , y ) p ( x , y ) ` θ ( x , y ) . ( 4 ) Such absolute continuity assumptions are standard in f -divergence and group DRO methods , which both rely upon re-weighting the training distributions . In fact , the KL divergence constraint in P-DRO presupposes absolute continuity . This suggests that the inner maximum can be re-written as an optimization problem on functions r : X × Y −→ R+ within the uncertainty setR ∈ { r | pr ∈ Q } min θ max r∈R E ( x , y ) ∼p r ( x , y ) ` θ ( x , y ) . ( 5 ) This reparametrization of the problem will allow us to replace a parametric family of generative models with a parametric family over probability ratios . 3.2 RATIO-BASED P-DRO . The likelihood ratio formulation described above is appealing for P-DRO because it enables the use of discriminative style neural architectures for parametrizing the ratio r , which opens up many more options for defining the parametric uncertainty set . Specifically , we can set the adversary to be any parametric function rψ : X × Y −→ R+ verifying Ex , y∼p rψ ( x , y ) = 1 . The key insight that we use to realize our proposed method is that we do not need to restrict the choice of adversaries to those that implicitly satisfy this normalization condition ( i.e . generative models ) . Instead , we can pick any adversary and treat normalization as an additional constraint ( the “ normalization constraint ” ) . Note that in this case , the KL constraint takes the simple form KL ( prψ‖p ) = Eprψ log prψ p = Ep rψ log rψ . The final min-max problem , which we dub ratio-based P-DRO ( RP-DRO ) , is : min θ max ψ Ep rψ log rψ≤κ Eprψ=1 E ( x , y ) ∼prψ ( x , y ) ` θ ( x , y ) ︸ ︷︷ ︸ LRP-DRO . ( 6 ) As in P-DRO , we can look for equilibria of this differentiable min-max game by performing simultaneous gradient updates ( Singh et al. , 2000 ) to θ and ψ in directions −∇θLRP-DRO and +∇ψLRP-DRO respectively . Although finding global equilibria is not guaranteed in high dimensional non-convex settings ( Balduzzi et al. , 2018 ) , empirical evidence suggests that models trained in this manner still reach useful solutions ( Michel et al. , 2021 ) . In experiments , we adopt an exponential parametrization rψ ( x , y ) ∝ efψ ( x , y ) where fψ is the output of any parametric model with values in R. Similarly to P-DRO , we do not explicitly enforce the KL constraint ( due to the difficulty of projecting onto the KL ball ) , and instead we relax it in the form of a term τ Ep rψ log rψ added to the loss function . The regularization strength τ is treated as a hyper-parameter .

Summary: As far as I can see, this work is about practical methods for distributionally robust optimization (DRO) with a special focus on trying to overcome the limitations of previous methods that have been proposed for this kind of problem. This work builds on the idea of instance-reweighting of the loss function via a re-weighting function that plays the role of the likelihood ratios (between the distributions in the class used for enforcing robustness and the training distribution). The re-weighting functions considered here are the well-known "exponential weights" and this work proposes and explores (i) using mini-batch level renormalizations and (ii) ensuring the KL-constraint by adding a KL penalty term to the training objective. The framework and corresponding objectives are described, and they are tested in experiments on benchmark image and text classification data sets, and compared to other existing DRO methods and to plain ERM.