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: This paper targets on solving distributionally robust optimization (DRO) that considering   distribution shifts in the data. In this paper, they show that  how different variants of DRO are simply instances of a finite-sum composite optimization for which they provide scalable methods by utilizing variance reduction algorithm. They also provide empirical results that demonstrate the effectiveness of our proposed algorithm with respect to the prior art in order to learn robust models from very large datasets.