Learning Distributionally Robust Models at Scale via Composite Optimization

1 INTRODUCTION . Conventional machine learning problem aims at learning a model based on the assumption that training data and test data come from same data distribution . However , this assumption may not hold in various practical learning problems where there is label shift ( Zhang et al. , 2020a ) , distribution shift ( Sagawa et al. , 2019 ) , fairness constraints ( Hashimoto et al. , 2018 ) , and adversarial examples ( Sinha et al. , 2017 ) , to name a few . Distributionally robust optimization ( DRO ) , which has recently attracted remarkable attention from the machine learning community , is a common approach to deal with the aforementioned uncertainties ( Duchi & Namkoong , 2016 ; Chen et al. , 2017 ; Rahimian & Mehrotra , 2019 ) . Defining the empirical distribution of the training data of size m by P̂m , 1m ∑m i=1 δξ̂i where δ is the Dirac delta function , the goal of DRO formulation is to solve the following optimization problem inf x [ Ψ ( x ) , sup ξ∈Q EQ [ ` ( x ; ξ ) ] ] , ( 1 ) where ξ is a data sample randomly drawn from distribution Q , ` ( x ; ξ ) is the corresponding loss function and EQ [ ` ( x , ξ ) ] is the expected loss over distributionQ which belongs to uncertainty set Um . The uncertainty set Um is defined as Um , { Q : d ( Q , P̂m ) ≤ } indicates the ball of a distribution with center P̂m and also d ( P , Q ) is a distance measure between probability distribution P and Q . We note this uncertainty set captures the distribution shift hence Eq . ( 1 ) minimizes the worse data distribution . Prior studies ( Ben-Tal et al. , 2013 ; Bertsimas et al. , 2018 ; Blanchet et al. , 2019 ; Esfahani & Kuhn , 2018 ) considered different uncertainty sets ( see Definition 3.1 in Esfahani & Kuhn ( 2018 ) ) for which they proposed equivalent reformulations of Eq . ( 1 ) based on the specific choice of Um . To solve the above min-max optimization problems , majority of prior studies heavily rely on either semidefinite programming ( Esfahani & Kuhn , 2018 ) or stochastic primal-dual methods both for convex ( Nemirovski et al. , 2009 ; Juditsky et al. , 2011 ; Yan et al. , 2019 ; 2020 ; Namkoong & Duchi , 2016 ) and non-convex ( deep learning ) objectives ( Yan et al. , 2020 ) . While primal-dual methods can be used as an approach to solve min-max optimization problems , it suffers from a few downsides . First and foremost , they need to store a probability distribution of constrained violation of dimension m corresponding to dual variables . Additionally , available primal-dual methods often demand data sampling that corresponds to the probability distribution over m data samples which introduces additional cost over uniform sampling . Finally , while majority of prior studies are limited to DRO problems with convex objectives , establishing tight convergence rate for DRO problems with penalty with non-convex objectives is still lacking . To overcome these issues , in this paper , we consider three different reformulations of Eq . ( 1 ) , corresponding to three different choices of uncertainty sets Um namely , ( 1 ) DRO with Wasserstein metrics , ( 2 ) DRO with χ2 divergence metrics , and ( 3 ) DRO with regularized entropy metrics ( also known as KL ) and show in Section 2 that all aforementioned DRO notions are indeed different instances of a deterministic composite optimization and can be solved by reducing to an instances of the following problem : min x [ Ψ ( x ) , r ( x ) + 1 m ∑m i=1 hi ( x ) + f ( 1 m ∑m i=1 gi ( x ) ) ] , ( 2 ) where we suppose r ( x ) is convex and a relatively simple function , f ( x ) : Rp → R and hi ( x ) : Rd → R for 1 ≤ i ≤ m are scalar-valued functions , and gi ( x ) : Rd → Rp for 1 ≤ i ≤ m are vector-valued functions . On the road to solve problem ( 2 ) at scale , we also develop a novel algorithm for heavily constrained optimization problems ( Wang & Bertsekas , 2015 ; 2016 ; Narasimhan et al. , 2020b ) that rather surprisingly invokes a single projection through the course of optimization . This algorithm is of independent interest and addresses the scalability issues raised in applications such as fairness ( Donini et al. , 2018 ; Zafar et al. , 2019 ) . We summarize the main contributions of our paper below : • We provide a large-scale analysis of DRO with Wasserstein distance and heavily constrained reformulation when the objective function is strongly-convex . Our result relies on a novel mini-batch constraint sampling for handling heavily-constrained optimization problems . As summarized in Table . 1 , our convergence analysis improves the state-of-the-art both in terms of the dependence on the convergence error as well as the number of constraints m. • We represent a large-scale analysis of DRO with non-convex objectives and χ2 or KL divergences . • We verify our theoretical results through various extensive experiments on different datasets . In particular , we show that our proposed method outperforms recent methods in DRO for heavily constrained problems with a great reduction in time complexity over them . The proofs of all the theorems are provided in the appendix . 1.1 RELATED WORK . DRO and connections to heavily constrained optimization . As mentioned earlier , DRO has many different formulations , depending on the divergence metrics used ( e.g. , Wasserstein , χ2 or KL ) . While Namkoong & Duchi ( 2016 ) ; Shapiro ( 2017 ) ; Duchi & Namkoong ( 2021 ) consider constrained or penalized DRO formulation , Sinha et al . ( 2017 ) ; Levy et al . ( 2020 ) formulate the underlying optimization problem as unconstrained . One of the contributions of our paper is to provide a unifying framework through the language of composite optimization and treat all these variants similarly . In particular , when the objective function is convex , Levy et al . ( 2020 ) recently proposed scalable algorithms for different variants of the DRO problems with , e.g. , χ2 or KL divergence metrics . Our unifying approach readily extends those results to the more challenging non-convex setting for which we are unaware of any prior work with convergence guarantees ( for instance , Hashimoto et al . ( 2018 ) studied DRO with χ2-divergence but did not provide any convergence guarantee ) . Similarly , Kuhn et al . ( 2019 ) ; Esfahani & Kuhn ( 2018 ) formulated DRO with Wasserstein distance as an instance of constrained optimization . Notably , they require ti impose one constraint per training data point and to solve such a constrained problem they proposed a semi-definite program . Even though the formulation is very novel , it can not scale . We , in contrast , consider such a heavily constrained optimization as an instance of a composite optimization for which we provide a scalable solution . What is rather surprising about our method is that it only checks a batch of constraints per iteration , inspired by Cotter et al . ( 2016 ) , and performs a single projection at the final stage of the algorithm in order to provide an -optimal solution in the case of strongly convex objectives . Moreover , in contrast to Cotter et al . ( 2016 ) , we do not keep a probability distribution over the set of constraints . We should also remark that our convergence guarantees achieve the known lower bounds in terms of accuracy and the number of constraints m. Finally , we should highlight the difference of our algorithm and Frank-Wolfe ( FW ) ( Frank et al. , 1956 ; Jaggi , 2013 ; Zhang et al. , 2020b ) . While FW does not require a projection oracle , it performs a linear program over the set of constraints at each iteration . In contrast , our heavily- constrained optimization solution performs a single projection without the overhead of running a linear program at each iteration . Stochastic composite optimization . The general stochastic composite optimization minx [ Ψ ( x ) , r ( x ) + f ( Eξ [ gξ ( x ) ) ] ] has recently received a lot of attentions ( Qi et al. , 2020b ; a ; Wang et al. , 2017 ) . Our reformulation of DRO variants is a finite-sum instance of this general problem . More concretely , Lian et al . ( 2017 ) ; Huo et al . ( 2018 ) ; Zhang & Xiao ( 2019a ) aimed to solve the following finite-sum problem minx [ Ψ ( x ) , r ( x ) + 1n ∑n j=1 fj ( 1 m ∑m i=1 gi ( x ) ) ] , using SVRG or SAGA ( Defazio et al. , 2014 ) . In contrast , our proposed algorithm is inspired by Zhang & Xiao ( 2019a ) and generalizes their method to the case where the extra terms hi ( x ) in Eq . ( 2 ) are non-zero . We should also note that Qi et al . ( 2020a ) proposed a similar idea in the context of online learning for DRO problems with KL divergence . Our work in contrast provides guarantees for DRO with both constraints or penalty terms . 2 DRO VIA FINITE-SUM COMPOSITE OPTIMIZATION . In this section , we discuss in detail how a finite-sum composite optimization ( 2 ) can unify various notions of distributionally robust learning , where some of which rely on heavily constrained optimization subroutines . While much research effort has been devoted to develop a specialized algorithm for each notion , our reduction paves the way to developing a scalable algorithm , discussed in Section 3 . DRO with Wasserstein distance . Kuhn et al . ( 2019 ) ; Esfahani & Kuhn ( 2018 ) provide the equivalent and tractable reformulation of Eq . ( 1 ) which can be regarded as a heavily constrained optimization problem as follows : min x r ( x ) , 1 m m∑ i=1 fi ( x ) subject to g̃i ( x ) ≤ 0 , ∀i ∈ [ m ] . ( 3 ) where g̃i ( x ) are functions related to loss function as well as slack variables ( please see Appendix A for more details ) . Naively solving optimization problem ( 3 ) suffers from the computational complexity due to the large number of constraintsm . To efficiently solve the optimization problem ( 3 ) , inspired by Mahdavi et al . ( 2012 ) and Cotter et al . ( 2016 ) , we pursue the smoothed constrained reduction approach and introduce the augmented optimization problem ( see Appendix B ) of the form min x Ψ ( x ) , [ r ( x ) + γ ln ( g ( x ) ) ] where gi ( x ) , exp ( αg̃i ( x ) γ ) and g ( x ) = 1m+1 [ 1 + ∑m i=1 gi ( x ) ] . We can see that this optimization problem is a special case of the optimization problem Eq . ( 2 ) where r ( x ) = f ( x ) , f ( 1m ∑m i=1 gi ( x ) ) = γ ln g ( x ) , and h ( x ) = 0 . In contrast to Cotter et al . ( 2016 ) that requires an extra storage cost of probability distribution of dimension m , and relatively poor convergence rate in terms of m and accuracy , we propose an algorithm that simply checks a batch of constraints and achieves the optimum dependency in terms of m and . DRO with χ2-divergence . The second type of DRO problem we consider utilizes the χ2divergence metric as follows : min x max 0≤pi≤1 , ∑m i=1 pi=1 m∑ i=1 pifi ( xi ) − γDχ2 ( p ) . ( 4 ) where the χ2 divergence is defined as the distance between the uniform distribution and an arbitrary probability distribution p , i.e. , Dχ2 ( p ) , m2 ∑m i=1 ( pi − 1m ) 2 . Levy et al . ( 2020 ) studied this problem only for the case of convex objectives . In this paper , we allow objective functions fi for 1 ≤ i ≤ m to be both non-convex or strongly-convex . The following claim derives the equivalent finite-sum composite optimization . Claim 2.1 . The optimization problem ( 4 ) is equivalent to the following problem : min x Ψ ( x ) , 1− 1 2γm m∑ i=1 [ fi ( x ) ] 2 + 1 2γ [ 1 m m∑ i=1 fi ( x ) ] 2 ( 5 ) We note that the optimization problem ( 5 ) fits into the formulation of finite-sum composite optimization ( 2 ) by choosing r ( x ) = 1 , h ( x ) = 1m ∑m i=1− ( fi ( x ) ) 2 2γ and f ( g ( x ) ) = 1 2γ [ 1 m ∑m i=1 fi ( x ) ] 2 with hi ( x ) = − f 2 i ( x ) 2γ , gi ( x ) = fi ( x ) and f ( x ) = x2 2γ . DRO with KL divergence . Finally , for DRO with KL-divergence , usually considered in online settings ( Qi et al. , 2020a ) , we consider solving the following optimization problem : min x max 0≤pi≤1 , ∑m i=1 pi=1 [ m∑ i=1 pifi ( xi ) + γH ( p1 , . . . , pm ) ] , ( 6 ) whereH ( p1 , . . . , pm ) = − ∑m i=1 pi log pi is the entropy function . To solve the min-max problem ( 6 ) , it is straightforward to convert it to the equivalent stochastic composite optimization problem of the following form : min x [ Ψ ( x ) , ln ( 1 m m∑ i=1 exp ( fi ( x ) γ ) ) ] . ( 7 ) As it can be seen , the optimization problem ( 7 ) fits into the composite optimization ( 2 ) by choosing r ( x ) = h ( x ) = 0 and f ( g ( x ) ) = ln ( 1 m ∑m i=1 exp ( fi ( x ) γ ) ) .

Summary: The paper views different variants of DRO are simply instances of a finite-sum composite optimization, from which efficient optimization algorithms were proposed. The convergence analysis was established for strongly-convex and non-convex settings.  The effectiveness of the proposed algorithm are well demonstrated in experiments.