Year: 2021

Abstract:
To train machine learning models that are robust to distribution shifts in the data, distributionally robust optimization (DRO) has been proven very effective. However, the existing approaches to learning a distributionally robust model either require solving complex optimization problems such as semidefinite programming or a first-order method whose convergence scales linearly with the number of data samples– which hinders their scalability to large datasets. In this paper, we show how different variants of DRO are simply instances of a finite-sum composite optimization for which we provide scalable methods. We 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.

heading: 1 INTRODUCTION
text: 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.

heading: 1.1 RELATED WORK
text: 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 cannot 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.

heading: 2 DRO VIA FINITE-SUM COMPOSITE OPTIMIZATION
text: 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) γ )) .

heading: 3 OUR PROPOSED ALGORITHM
text: Having reduced the different notions of DRO to an instance of the composite optimization, in this section we describe our scalable approach for minimizing the objective Ψ(x) = r(x) + Φ(x) where Φ(x) = 1m ∑m i=1 hi(x) + f( 1 m ∑m i=1 gi(x)). We note that the compositional structure in Φ(·) leads to more challenges in optimization compared with the non-compositional finite-sum problem, since the stochastic gradient of the loss function is not an unbiased estimation of the full gradient. To overcome this issue, and by building on incremental variance reduction (Zhang & Xiao, 2019b), we propose a more general algorithm with a new ingredient in which we also employ variance reduction on the extra term h(x) = (1/m) ∑m i=1 hi(x). To handle r(·), similar to Zhang & Xiao (2019b), we assume r is convex and a relatively simple function. We follow the proximal gradient iterates (Nesterov, 2013; Beck, 2017) as follows:
x(t+1) = Πηr
( x(t) − η∇Φ(x(t)) ) (8)
where we apply the proximal operator of r(x) with the learning rate η as Πηr(x) , arg miny [ r(y) + 12η ‖y − x‖ 2 ]. By defining the proximal gradient mapping of Ψ as Gη(x) , 1 η [x−Πηr (x− η∇Φ(x))], the updating rule in Eq. (8) can be equivalently written as x(t+1) = x(t)−ηGη(x(t)). Given any y as an output of randomized algorithm, we say y achieves the stationary point of problem in Eq. (2) in expectation if E [ ‖Gη(y)‖2 ] ≤ holds. Our goal is to achieve an stationary point with the least number of calls to a (mini-batch) stochastic oracle.
Focusing on Φ(·), as detailed in Algorithm 1, we apply three time-scale variance-reduced estimators for gi(x) and its gradient∇gi(x), as well as hi(x). For DRO with Wasserstein divergence metric with
Algorithm 1: Generalized Composite Incremental Variance Reduction (GCIVR (x(0))) Inputs: Number of iterations t = 1, . . . , T , learning rate η, initial global model x(0), the size of
epoch length τt, and mini-batch sizes of Bt and St at time t. for t = 1, . . . , T do
Sample a mini-batch B(t) with size Bt uniformly over [m] and compute
y (t) 0 =
1
Bt
∑ ξ∈B(t) g(x(t)τt ; ξ), z (t) 0 = 1
Bt
∑ ξ∈B(t) ∇g(x(t)τt ; ξ),w (t) 0 = 1
Bt
∑ ξ∈B(t) ∇h(x(t)τt ; ξ)
Compute ∇̃Φ(x(t)0 ) = ( z (t) 0 )> ( f ′(y (t) 0 ) ) +w (t) 0
Update the model as follows: x(t)1 = Π η r
( x
(t) 0 − ∇̃Φ(x (t) 0 )) ) for j = 1, . . . , τt − 1 do in parallel
Sample a mini-batch S(t)j with size St uniformly over [m], and form the estimates
y (t) j = y (t) j−1 +
1
St ∑ ξ∈S(t)j [ g(x (t) j ; ξ)− g(x (t) j−1; ξ) ] (9)
z (t) j = z (t) j−1 +
1
St ∑ ξ∈S(t)j [ ∇g(x(t)j ; ξ)−∇g(x (t) j−1; ξ) ] (10)
w (t) j = w (t) j−1 +
1
St ∑ ξ∈S(t)j [ ∇h(x(t)j ; ξ)−∇h(x (t) j−1; ξ) ] (11)
Compute ∇̃Φ(x(t)j ) = ( z (t) j )> ( f ′(y (t) j ) ) +w (t) j
Update the model as follows: x(t)j+1 = Π η r
( x
(t) j − ∇̃Φ(x (t) j )) ) end
end Output: Return a randomly selected solution from {xtj}t=1,...,Tj=0,...,τt
optimally strongly objective, at the beginning of each epoch t we compute a full-batch gradient over the entire data samples Bt = m, i.e., y (t) 0 = 1 m ∑m i=1 gi(x (t) τt ), z (t) 0 = 1 m ∑m i=1∇gi(x (t) τt ),w (t) 0 = 1 m ∑m i=1∇hi(x (t) τt ). In contrast, for the case of DRO with χ2 or KL divergence metrics and nonconvex objectives we incrementally increase the size of the mini-batch Bt ≤ m in the beginning of each epoch until it reaches the point where we need to compute the full-batch of samples. We should also highlight that a mini-batch in case of Wasserstein DRO indicates a batch of constraints and in DRO with χ2 or KL divergence metrics represents the number of data sample accessed. We denote the length of each epoch and the mini-batch size within each epoch with τt and St, respectively. In each epoch t and iteration j, we estimate mini-batch gradients from Eqs. (9), (10), and (11) in Algorithm 1, with some corrections term applied from the previous iteration. We note that the variance-reduced term corresponding to the correction is inspired by SARAH (Nguyen et al., 2017) and SPIDER (Fang et al., 2018). Finally, the algorithm returns a randomly selected solution from the iterates.
Distributed variant of Algorithm 1. As mentioned before, efficient training of the stochastic composite optimization problem has attracted increasing attention in recent years. Despite much progress, all of existing methods including Algorithm 1 only focus on the single-machine setting. To employ Algorithm 1 in a distributed setting with p machines, in Appendix E we propose a distributed variant of proposed algorithm and establish its convergence rate for convex and non-convex objectives which enjoys a speedup in terms of number of machines.

heading: 4 CONVERGENCE ANALYSIS
text: In this section we establish the convergence of proposed algorithm for different DRO notions discussed in Section 2. We start by stating the general assumptions and then discuss the obtained rates. Due to lack of space we only include the rates on the convergence of DRO with Wasserstein metric for strongly convex, and χ2 and KL divergence metrics for non-convex objectives. We defer
the analysis of χ2 and KL divergence metrics with strongly objectives to Appendix D. To establish the convergence rates, we first introduce some standard assumptions. Assumption 1. We make the following assumptions on the components of objective Ψ(x) = r(x) + 1 m ∑m i=1 hi(x) + f( 1 m ∑m i=1 gi(x)):
1) ‖∇h(x1)−∇h(x2)‖ ≤ Lh ‖x1 − x2‖ where x1,x2 ∈ Rd. We also assume that ‖h(x1)− h(x2)‖ ≤ `h ‖x1 − x2‖. 2) ‖f ′(x1)− f ′(x2)‖ ≤ Lf ‖x1 − x2‖ where x1,x2 ∈ R. We also assume that ‖f(x1)− f(x2)‖ ≤ `f ‖x1 − x2‖. 3) ‖∇g(x1; ξ)−∇g(x2; ξ)‖ ≤ Lg ‖x1 − x2‖ and ‖g(x1; ξ)− g(x2; ξ)‖ ≤ `g ‖x1 − x2‖ for ∀x1,x2 ∈ Rd. 4) We suppose that Ψ(x) in Eq. (2) is bounded from below that Ψ∗ = infxΨ(x) > −∞. 5) We assume r(x) ∈ R ∪ {∞} is a convex and lower-semicontinues function.
An immediate implication of above assumption is that f(g(x)) , f( 1m ∑m i=1 gi(x)) is smooth with
module ( `2gLf + `fLg ) , hence Ψ(x) is smooth with module LΦ , [( `2gLf + `fLg ) + Lh ] (Wang et al., 2016; Zhang & Xiao, 2019b). Assumption 2 (Unbiased estimation). We assume that the mini-batch of samples ξ over functions gi for i = 1, 2, . . . ,m is unbiased, that is E [g(x; ξ)] = g(x) and E [∇g(x; ξ)] = ∇g(x).

heading: 4.1 OPTIMALLY STRONGLY CONVEX OBJECTIVES
text: We here establish the convergence for optimally strongly convex objectives under Wasserstein metric as defined below. Definition 1 (Optimally strongly convex). We say that the objective function Ψ(x) is optimally strongly convex objective with module µ if Ψ(x)−Ψ(Px∗ (x)) ≥ µ ‖x− Px∗ (x)‖2.
Following Cotter et al. (2016) and Mahdavi et al. (2012), we make the following assumption. Assumption 3. There exists a constant ρ such that if we define g(x) = max1≤i≤m g̃i(x) we have min g(x)=0 ‖∇g(x)‖2 ≥ ρ.
Assumption 4. Function r(x) is smooth with module G. Theorem 4.1. Let us assume that Ψ is µ-optimally strongly convex. Let us set τt = St = √ Bt = √ m, T = 5√ mµη
, η < 2 LΦ+ √ L2Φ+36G0 where G0 , 3(`4gL 2 f + ` 2 fL 2 g + ` 2 h), α > G ρ , and γ = exp(−K) ln(m+1) . Let us denote x(k+1) = GCIVR ( x(k) ) for k = 0, . . . ,K − 1 (using Algorithm 1). Under assumptions
1-4 and by letting K = ln (1/ ), the solution of DRO with Wasserstein divergence is obtained by projectingx(K) onto the constraint setK = {x|gi(x) ≤ 0, i = 1, 2, . . . ,m}, i.e., x̄(K) = ΠK(x(K)). In order to achieve r(x̄(K))− r(x(∗)) ≤ we require O ( (m+ κ √ m) ln 1 ) number of calls to the stochastic oracle.
Comparison with previous results. Compared to the MidTouch approach by Cotter et al. (2016), our obtained rate improves both on the dependency on the number of constraints m, as well as convergence error . To better understand the intuition behind achieving such a double folded improvement in terms of m and , we note that unlike Cotter et al. (2016) which utilizes a primal-dual approach in order to obtain an approximate to the optimal distribution over the constraints, we directly find the optimal distribution exactly. Second, while Cotter et al. (2016) does not apply any variance reduction technique to primal variable x, our algorithm benefits from variance reduction over x too. Furthermore, Theorem 4.1 entails tighter rate in terms of final accuracy compared to Yu et al. (2021). These results are summarized in Table. 1.
In Algorithm 1 for optimally strongly convex objectives, we need to do a single projection at the end. The following corollary bounds the error between the solution before and after the projection. Corollary 4.2. Under the assumptions made in Theorem 4.1, the error between the solution with and without projection is bounded by
r(x̄(K))− r(x(K)) ≤ G [ γ ln(m+ 1)
αρ−G + 1 αρ−GO (exp(−K)) ] + γ ln(m+ 1). (12)
Therefore, with a proper choice of γ we can make the right hand side matches the convergence error of Theorem 4.1.
Remark 1. It is worthy to highlight that through reduction of heavily-constrained optimization to composite optimization, GCIVR algorithm can be considered as an alternative method to solving constrained optimization problems via mini-batch sampling of constraints. In particular, Theorem 4.1 shows that under certain conditions, we can solve any heavily-constrained optimization problem with sampling a mini-batch of constraints and achieve similar guarantees compared to projection-based counterparts while avoiding the heavy dependency on the number of constraints. Another implication of Theorem 4.1 is that we can solve heavily constrained optimization with the sample complexity similar to that of an unconstrained optimization problem.
Remark 2. To understand the tightness of our obtained rate, consider minx [ Ψ(x) , r(x) + 1m ∑m i=1 gi(x) ] which is an instance of our general optimization problem (2). Clearly, any lower bound to solve this instance also holds for the original optimization problem (2). Xie et al. (2019) provides a lower-bound of O ((m+ √ κm) ln (1/ )) for above instance, matching our upper bound in terms of the dependency on the number of data samples, while the dependency on κ can still be improved. Furthermore, for the DRO problem with Wasserstein divergence metric, rather than solving constrained optimization problem we solve unconstrained compositional optimization problem. Thus, since solving constrained optimization is more complex than unconstrained optimization problem, we do not expect to obtain any better bound regarding m even for DRO with Wasserstein metric.
Remark 3. In distributed setting with p devices, we are able to improve the sample complexity to O (( m+ mp + κ √ m ) ln 1 ) . The proof can be found in Appendix F.3.

heading: 4.2 NON-CONVEX OBJECTIVES
text: In this section, we estbalish the convergence of DRO problems with χ2 and KL divergence metrics for non-convex objectives. To this end, we make an additional assumption on the variance of stochastic mini-batches. Assumption 5 (Bounded variance). The mini-batch sampling has bounded variance that is E [ ‖∇g(x; ξ)−∇g(x)‖2 ] ≤ σ2B where B indicates the batch size.
Theorem 4.3. Under Assumptions 1,2, and 5 using Algorithm 1 for some positive constants β and 0 ≤ ζ < √m, denote T0 , √ m−ζ β = O ( √ m). For t ≤ T0, let us set τt = St = √ Bt = βt+ ζ and for t > T0 set τt = St = √ m. Then, if η < 4
LΦ+ √ L2Φ+12G0
, after T = Õ (min(1/ √ , 1/ √ m ))
iterations, with Õ ( min{√m/ , 1/ 1.5} ) number of calls to the stochastic oracle it holds that
E [ ∥∥Gη(x(T ))∥∥2 ] ≤ , where we used Õ(.) notation to hide terms with logarithmic dependency.
We also remark that Assumption 5 is only needed for convergence until T0 = √ m−ζ β and after t > T0 we will not need this assumption. We also note that this assumption is not required in optimally strongly convex setting as we utilize a full batch at the beginning of each epoch.
Discussion on lower bound. As we discussed in Remark 2, any impossibility result for optimizing minx [ Ψ(x) , r(x) + 1m ∑m i=1 gi(x) ] also holds for general compositional optimization problem in Eq. (2). For former problem, Fang et al. (2018) shows that to achieve the stationary point with error we require at least Õ ( min{ √ m , 1 1.5 } ) (almost optimal) gradient oracle calls for strongly
optimal objectives. As a consequence, for the DRO with χ2 or KL divergence metrics we do not expect less number of constraint checks or sample complexity, respectively.
Remark 4. In distributed setting with p devices, we can improve the convergence rate of DRO with χ2 and KL metrics to Õ ( min{ √ m p + √ m , 1 p 1.5 + 1 1.5 } ) . For details, please refer to Appendix F.3.

heading: 5 EXPERIMENTS
text: In this section, we empirically examine the efficacy of the proposed algorithm in different use cases. The main algorithm that we compare against is Heavily-Constrained algorithm that uses a parametric multiplier model proposed in Narasimhan et al. (2020b). We call this algorithm Heavily-constrained, where they learn the Lagrange multiplier values using a parametric model such as a neural network. The tasks we apply the algorithms are mainly focused on fairness constraints since they fit greatly to the problem definition in this paper due to the large number of constraints they add to the main learning problem. In addition, similar to distribution robust optimization approaches, the main goal of fairness constraint is to find a solution that is agnostic to the distribution of protected groups. The code for the experiments is available at this link.
Distributionally robust optimization for fairness constraints. The first experiment setup is based on Wang et al. (2020); Narasimhan et al. (2020b), where we want to enforce equality of opportunity (Hardt et al., 2016) constraints for different groups while the group membership is noisy and changing during the training. Hence, the problem is to make the solution distributionally robust among different protected groups in the problem. Based on the setup in Narasimhan et al. (2020b), we assume we have access to the marginal probability of the true groups (P (gi = j|ĝi = k), where gi is the true group membership and ĝi is the noisy group membership). Hence, to enforce fairness constraints, we consider all possible proxy groups using this marginal probability, which can increase the number of constraints greatly. The goal in this case for equal opportunity is to have the true positive rate of each group in the vicinity of the true positive rate (tpr) of the overall data, that is: tpr(g = j) ≥ tpr(ALL)− for every proxy group we define. In this experiment, we use the Adult dataset (Dua & Graff, 2017), and consider race groups of “white”, “black”, and “other” as protected groups. We train a linear classifier with logistic regression, and report the overall error rate of the classifier, as well as the maximum violation of the fairness constraints (equal opportunity) over true group memberships. We set = 0.05 and the noise level to 0.3. Again, we compare with unconstrained optimization and heavily-constrained algorithm with a linear model as its multiplier model.First row of Figure 1 shows the results for this experiment, where the proposed GCIVR algorithm can achieve the same level of constraint violation on true group memberships as the heavily constraint while outperforms it in terms of overall error rate. Hence, the solution found by the GCIVR dominates heavily-constrained solution. In terms of runtime speed, the first row of Figure 1(c) clearly shows the advantage of GCIVR over heavily-constrained with much lower overhead to the unconstrained optimization.
Fairness Constraints on intersectional groups. In this task we ought to learn a linear classifier with logistic regression to predict the crime rate for a community on Communities and Crime dataset (Dua & Graff, 2017). This dataset contains 1994 instances of communities each with 128 features to predict the per capita crime rate for each community in the US. The labels are high and low crime rates to represent if a community is above the 70th percentile of the data or not. To add fairness constraints, we first determine different communities based on the percentages of the Black, Hispanic and Asian population in each community, as discussed in Narasimhan et al. (2020b); Cotter et al. (2019). Similar to Narasimhan et al. (2020b), we generate 1000 thresholds of the form (τ1, τ2, τ3) ∈ [0, 1]3 to define each group. Consider the population of each race as (p1, p2, p3), then a community belongs to group gτ1,τ2,τ3 if pi ≥ τi ∀i ∈ [3]. Then the fairness constraints enforce that the error rate of each group should be in the neighborhood of the overall error by margin of . That is err (gτ1,τ2,τ3) ≤ err (ALL) + . Similar to Narasimhan et al. (2020b), we set = 0.01 and not consider groups with less than 1% of data.
We compare with the unconstrained optimization and Heavily-constrained algorithm with a 2-layer neural network, each with 100 nodes as their Lagrange model, as described in their paper. We compare the trade-off between the error rate and maximum violation of fairness constraints among groups. The second row of Figure 1 portrays this comparison for the test dataset. As it can be inferred, the proposed GCIVR algorithm achieves the same error rate as Heavily-constrained, but both higher than the unconstrained optimization. However, comparing the maximum violation of constraints, it is clear that GCIVR outperforms both optimization methods. Also, the third figure shows the runtime of different algorithms. It is clear that GCIVR adding a minimal overhead to the unconstrained optimization, while heavily-constrained increases the runtime more than 100 times.
Per-query fairness constraints in ranking. In the third problem we evaluate our algorithm on the fairness in the ranking problem. In this case, we intend to impose per-query constraints on the
learning to rank problem as defined in Narasimhan et al. (2020a;b). Similar to those setups, we divide document-query pairs into two groups based on the 40th percentile of the their QualityScore features. Hence, in this problem, we want to learn a ranking function f : D ×Q → R, which maps a pair of document-query features to a real number as the score. Consider the groups of g0 and g1 as mentioned before, we want the error rate for these groups to be close to each other. In other terms, if the pairwise error rate is defined as: erri,j(q) = E [1 {f(d, q) < f (d′, q)} | y > y′, d ∈ gi, d′ ∈ gj ], where y and y′ are the respective binary labels. Then, the fairness constraints can be satisfied as |err0,1(q)− err1,0(q)| ≤ ∀q ∈ Q. For this experiment, we use Microsoft Learning to Rank Dataset (MSLR-WEB10K) (Qin & Liu, 2013), which contains 10K queries and 136 features. For this experiment we use a non-convex objective, where the model is a two-layer neural network each with 128 nodes and cross-entropy as the loss function. We compare against the unconstrained optimization and the heavily-constrained algorithm with an one-layer neural network with 64 nodes as its multiplier model. We use 1000 queries in the training and 100 queries in the test datasets. The third row in Figure 1 shows the error rate and maximum violation of groups constraints. As it is clear GCIVR outperforms heavily-constrained in both error rate and maximum violations of groups by a large margin. In terms of runtime, GCIVR is very close to the unconstrained optimization, and about 2× faster than heavily-constrained.

heading: 6 CONCLUSION
text: In this paper, we showed that many DRO problems or heavily constrained optimization problems can be cast into a general framework based on finite-sum composite optimization. To solve this composite finite-sum optimization, we introduced centralized and distributed algorithms. We theoretically illustrated that our algorithm converges with an almost optimal number of constraint checks (for Wasserstein distance) or gradient calls (for χ2 or KL divergence metrics). Finally, we validated our theory with a number of experiments.

Review summaries:
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. 

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.