Modeling the Second Player in Distributionally Robust Optimization

1 INTRODUCTION . Machine learning models trained with empirical risk minimization ( ERM ) are able to achieve high aggregate performance on data sampled from their training distribution . However , they often exhibit drops in accuracy when confronted with data from domains that are under-represented in their training data , such as those of different topic ( Gururangan et al. , 2020 ) , sociolect ( Blodgett et al. , 2016 ) , accent ( Amodei et al. , 2016 ) or writer age ( Hovy & Søgaard , 2015 ) in language processing tasks , or skin color ( Grother et al. , 2019 ) or lighting ( Georghiades et al. , 2001 ) in image processing tasks . This is a particularly egregious issue in applications where higher error rates can have far reaching negative implications , such as the silencing of underrepresented minorities in toxicity detection systems ( Dixon et al. , 2018 ) or disparity amplifying feedback loops in credit rating models ( Fuster et al. , 2018 ) . This behaviour often arises from the objective function of ERM , where the parameters θ of the model are learned by minimizing the expectation of a loss function ` under a data distribution p ( or , specifically in practice , an associated empirical data distribution p̂ ) LERM ( θ ) = E ( x , y ) ∼p̂ ` ( x , y , θ ) . ( 1 ) When the model encounters data sampled from a different distribution qtest 6= p , performance can suffer significantly . Distributionally robust optimization ( DRO ) ( Ben-Tal et al. , 2013b ) provides a natural solution to this issue by replacing the expected risk under a single distribution p with the worst expected risk over a pre-determined family of distributions Q ( the “ uncertainty set ” ) LDRO ( θ ) = max q∈Q E ( x , y ) ∼q ` ( x , y , θ ) . ( 2 ) 1Code to reproduce our experiments can be found at https : //github.com/pmichel31415/P-DRO If Q contains qtest , the DRO objective upper bounds the expected risk under qtest . However , a priori knowledge of possible test distributions is not always available or easy to acquire . For example , training a model to be robust to some demographic attributes ( Q = { qdemographic 1 , qdemographic 2 , . . . } ) requires collecting and annotating data with the necessary information , an expensive and ethically fraught endeavour . In the absence of such information , one has to resort to defining the uncertainty set analytically , drawing on one ’ s intuition of what constitutes a possible test distribution given the observed training distribution , such as using moment constraints ( Delage & Ye , 2010 ; Nguyen et al. , 2020 ) , f -divergence ( Ben-Tal et al. , 2013a ; Hu & Hong , 2013 ; Faury et al. , 2020 ) , Wasserstein/IPM ( Sinha et al. , 2018 ; Husain , 2020 ) balls , or coarse-grained mixture models ( Oren et al. , 2019 ; Hu et al. , 2018 ) . However , the need for keeping the inner supremum in Eq . ( 2 ) tractable limits the possible choices . In this paper , we propose that the uncertainty set be instead defined as a family of parametric generative models . The resulting DRO objective ( §2 ) is a differentiable game with two players : the original model ` ( x , y ; θ ) and a model of its worst-case distribution qψ ( x , y ) , the titular “ second player ” which we hereafter refer to as the adversary . Using this formulation — which we call Parametric DRO ( P-DRO ) — allows for more flexibility in the choice of the adversary ’ s architecture ( and so the uncertainty set ) . Unfortunately , finding a solution of this game via direct application of simultaneous gradient descent ( Singh et al. , 2000 ) is difficult ( Balduzzi et al. , 2018 ) . In particular , direct gradient descent on the uncertainty set suffers from instability due to the large variance of the gradients ( Greensmith et al. , 2004 ) , and hyper-parameter selection is not straightforward . To address these challenges , we make two main contributions ( §3 ) : first , we propose a new relaxation of the DRO game ’ s inner maximization problem ( with KL constraints ) . The resulting objective is more amenable to simultaneous gradient update than the original zero-sum game and significantly improves training stability , while still yielding useful adversaries . Second , we develop a principled approach for selecting hyper-parameters : we leverage the learned adversaries to decide which of any two given models trained with P-DRO is more robust than the other . We do an in-depth set of experiments analyzing the effect of our proposed changes on both a toy task as well as a more realistic , yet still synthetic sentiment classification task ( §4 ) . Finally , we show that in the more realistic setting of toxicity detection , P-DRO yields models that are more robust to changes in demographic groups , even though these groups are unknown at training time , opening up applications in combatting dataset bias ( §5 ) . 2 PARAMETERIZING THE UNCERTAINTY SET . Consider a model parameterized by θ ∈ Rdmodel . Minimizing the DRO objective described in Eq . ( 2 ) over the uncertainty set Q turns the optimization problem into the min-max ( or zero-sum ) game min θ∈Rd max q∈Q E ( x , y ) ∼q ` ( x , y , θ ) . ( 3 ) The first player controls the parameters θ , whilst the second player controls the worst-case distribution q . In the absence of explicit information on groups of interest ( such as demographics , domain , etc . ) , an adequate choice of the uncertainty set Q is critical to the success of DRO . This is in fact very much an active area of research ( Sinha et al . ( 2018 ) ; Duchi & Namkoong ( 2018 ) ; Oren et al . ( 2019 ) , see Rahimian & Mehrotra ( 2019 ) for a survey ) . Q must be sufficiently large to contain test distributions of interest , but if it is too large it may contain “ adversarial ” distributions on which no model can perform well . Moreover , the design of Q is also circumscribed by the necessity of keeping the min-max problem tractable , particularly in the context of stochastic optimization . In Hu & Hong ( 2013 ) and Duchi et al . ( 2016 ) for example , the choice of f -divergence balls allows the use of duality arguments to reformulate ( 3 ) as a more manageable min-min problem . Others , like Hu et al . ( 2018 ) or Oren et al . ( 2019 ) , propose using mixture models , the simplicity of which enables them to solve the inner maximization problem efficiently . Instead , we propose to explicitly model the second player in the DRO game as a parametric model qψ of the data . Of course , not all parameterizations ψ ∈ Rdadv of a given generative model represent useful distributions , and we require that the adversary stay “ close ” to the underlying true data distribution p. As a measure of distance between qψ and p , we choose the KL ( Kullback & Leibler , 1951 ) divergence due to its wide acceptance in the machine learning community , as well as its appealing properties in the context of DRO.2 The KL upper bound , κ , is left as a parameter to be decided by the experimenter . We refer to the resulting DRO formulation as Parametric DRO min θ max ψ KL ( qψ‖p ) ≤κ E ( x , y ) ∼qψ ` ( x , y , θ ) ︸ ︷︷ ︸ LP-DRO ( θ , ψ ) . ( 4 ) 3 OPTIMIZING P-DRO . The min-max problem in Eq . ( 4 ) belongs to a class of games called “ differentiable games ” ( another famous representative being generative adversarial networks ( Goodfellow et al. , 2014 ) ) . We can search for a solution of this game with simultaneous gradient descent ( Singh et al. , 2000 ) , i.e . by simultaneously updating θ and ψ with −∇θ LP-DRO and ∇ψ LP-DRO respectively . Unfortunately , in general , there is no theoretical guarantee that simultaneous gradient descent will converge to a Nash equilibrium3 ( Balduzzi et al. , 2018 ) , nor that any such equilibrium even exists if the objective is nonconvex in θ ( or non-concave in ψ ) . The success of GANs and the follow-up literature ( Wang et al. , 2019 ) serves as an encouraging example that gradient based methods can yield useful solutions despite the pessimistic theoretical results . In this section , we discuss difficulties that arise when optimizing θ and ψ jointly , and propose modifications of the objective to address them . 3.1 TRAINING THE MODEL θ We could train the model θ by taking negative gradient steps on E ( x , y ) ∼qψ ` ( x , y ; θ ) . This gradient can be estimated by sampling examples from qψ and averaging the gradient of their losses . Unfortunately , this objective requires that qψ is well-behaved at all iterations , as it is the only source of supervision for θ . If qψ is initialized incorrectly or begins producing unrealistic ( x , y ) , the quality of θ degrades as it begins to learn a predictor on invalid training examples from qψ . As an alternative , we opt to compute the gradients for θ with importance sampling , i.e . rewriting L P-DRO as E ( x , y ) ∼p qψ ( x , y ) p ( x , y ) ` ( x , y ; θ ) , which ensures that all ( x , y ) samples will be derived from the training set itself . Unfortunately , the true density p is unknown to us . As an approximation , we replace qψ ( x , y ) p ( x , y ) with the likelihood ratio between qψ and the maximum likelihood estimate of p , qψ0 : = argmaxqψ E ( x , y ) ∼p log qψ ( x , y ) . This changes the min-max problem to min θ max ψ KL ( qψ‖p ) ≤κ E ( x , y ) ∼p qψ ( x , y ) qψ0 ( x , y ) ` ( x , y , θ ) ︸ ︷︷ ︸ Lmodel . ( 5 ) 2For instance : KL ( q‖p ) < +∞ implies that q stays within the support of p 3Nash equilibria ( Osborne & Rubinstein , 1994 ) can be thought of the game theoretic analog of global minima in optimization . This becomes a simple expected loss objective , which we can estimate by sampling from the empirical distribution p̂ . In experiments , we find that with this formulation we are able to train robust θ even when qψ is only a mediocre generative model ( see Appendix C.2 ) . To further stabilize training at the beginning of the optimization process , we initialize ψ with ψ0 , making the objective exactly the same as ERM for the first gradient step . 3.2 TRAINING THE ADVERSARY ψ According to Eq . ( 5 ) the adversary ψ must maximize E ( x , y ) ∼qψ p ( x , y ) qψ0 ( x , y ) ` ( x , y , θ ) within a KL ball of fixed radius . This is challenging for several reasons : first , enforcing the bound is intractable for complex families of adversaries where e.g . projecting onto the KL ball is another difficult optimization problem of its own . Second , maximizing the expectation with respect to the parameters of the distribution qψ is prone to instability due to large gradient variance ( Greensmith et al. , 2004 ) . Lagrangian Relaxation To address the first difficulty , we loosen the strict KL constraint and instead consider the Lagrangian relaxation L L ( ψ , τ ) = E ( x , y ) ∼qψ p ( x , y ) qψ0 ( x , y ) ` ( x , y , θ ) − τ ( KL ( qψ‖p ) − κ ) . ( 6 ) We fix the Lagrangian multiplier τ > 0 as treat it as a “ temperature ” hyper-parameter . With some reorganization ( which we develop in Appendix A.1 ) , we can show that L ( ψ , τ ) = −τKL ( qψ‖q∗τ , θ ) + C. ( 7 ) Where q∗τ , θ ∝ p ( x , y ) e p ( x , y ) qψ0 ( x , y ) ` ( x , y ; θ ) τ and C is a constant in ψ . In other words , maximizing L in ψ is equivalent to minimizing the KL divergence between qψ and q∗τ , θ . One difficulty with this objective is that q∗τ , θ depends upon the unknown probability density p ( x , y ) . We avoid this problem by treating the density ratio p ( x , y ) qψ0 ( x , y ) as a constant , which is closely related to assumptions that have been used successfully in past formulations of DRO ( Oren et al. , 2019 ) . Empirically , we find that incorporating qψ0 as a surrogate for p is a serviceable approximation , as demonstrated in Section 4 . Reversing the KL Minimizing the KL divergence in this direction is difficult for several reasons . First , it entails optimizing an expectation in qψ over ψ , which is difficult due to the large variance of the gradients ( Greensmith et al. , 2004 ) . Second , computing this KL necessitates access to the true theoretical density p ( x , y ) in order to compute q∗τ , θ ( x , y ) in the argument of the expectation , but this quantity is unknown in practice.4 To sidestep these issues , we elect to minimize the reverse direction KL ( q∗τ , θ‖qψ ) instead . Due to the KL divergence being non-symmetric , this is a rather crude approximation5 , the implications of which are discussed in Norouzi et al . ( 2016 ) . However , we find that this approach dramatically stabilizes the gradient dynamics while still yielding good adversaries , as observed empirically in Section 4.4 . Discarding the entropy term ( constant in ψ ) , the resulting problem is equivalent to minimizing Ladv ( ψ , τ ) : = − 1 Zτ , θ Ep e ` ( x , y ; θ ) τ log qψ ( x , y ) ( 8 ) in ψ , where Zτ , θ = Ep e ` ( x , y ; θ ) τ is the normalizer of q∗ . In this case , we can estimate this expectation by substituting the empirical distribution p̂ for p in the expectation . Computing the Normalizer Approximating the inverse normalizer 1Zτ , θ in a minibatch yields a biased estimator . On the other hand , computing Zτ , θ over the entire training data at each step is prohibitive since it requires computing the loss of every single example . As a middle ground , we keep a running normalizer Z̃k computed from the average of the normalizers over a fixed number 4Note that substituting the empirical distribution p̂ for p poses issues here , because qψ is not absolutely continuous with respect to p̂ . 5For instance , the optimum of the reverse KL doesn ’ t necessarily match that of the forward KL within the parametric confusion set Q k of consecutive minibatches . In other words , if Bi and θi denote the minibatch and adversary parameters at step i respectively , the normalizer at step t will be Z̃k = 1∑t i=t−k |Bi| t∑ i=t−k ∑ x , y∈Bi e ` ( x , y ; θi ) τ . ( 9 ) If k is too low , there is a risk of under-estimating the normalizer , especially if the distribution of weights contains infrequent high weight samples . On the other hand , if k is too high there is a risk of using “ stale ” weights in the normalizer . In experiments , we treat k as a hyper-parameter .

Summary: The paper proposes to define the uncertainty set in the DRO problem as a family of parametric generative models, which is to allow more flexibility in the choice of the uncertainty set architecture. To realize this idea, the paper first proposes a new relaxation of the DRO game's inner maximization problem (with KL constraints) so as to improve the training stability. It then develops a principled approach to select the hyper-parameters of the proposed method.