Original: L EARNING D ISTRIBUTIONALLY R OBUST M ODELS AT S CALE VIA C OMPOSITE O PTIMIZATION
Revised: L EARNING D ISTRIBUTIONALLY R OBUST M ODELS AT S CALE VIA C OMPOSITE O PTIMIZATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: To train machine learning models that are robust to distribution shifts in the data, distributionally robust optimization (DRO) has been proven very effective.
Revised: To train machine learning models that are robust to distribution shifts in the data, distributionally robust optimization (DRO) has been proven very effective.

Original: However, the existing approaches to learning a distributionally robust model either require solving complex optimization problems such as semideﬁnite programming or a ﬁrst-order method whose convergence scales linearly with the number of data samples– which hinders their scalability to large datasets.
Revised: However, the existing approaches to learning a distributionally robust model either require solving complex optimization problems such as semideﬁnite programming or a ﬁrst-order method whose convergence scales linearly with the number of data samples– which hinders their scalability to large datasets.

Original: In this paper, we show how different variants of DRO are simply instances of a ﬁnite-sum composite optimization for which we provide scalable methods.
Revised: In this paper, we show how different variants of DRO are simply instances of a ﬁnite-sum composite optimization for which we provide scalable methods.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Conventional machine learning problem aims at learning a model based on the assumption that training data and test data come from same data distribution.
Revised: Conventional machine learning problem aims at learning a model based on the assumption that training data and test data come from same data distribution.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Deﬁning the empirical distribution of the training data of size m by ˆ P m (cid:44) 1 m (cid:80) mi =1 δ ˆ ξ i where δ is the Dirac delta function, the goal of DRO formulation is to solve the following optimization problem
Revised: Deﬁning the empirical distribution of the training data of size m by ˆ P m (cid:44) 1 m (cid:80) mi =1 δ ˆ ξ i where δ is the Dirac delta function, the goal of DRO formulation is to solve the following optimization problem

Original: [Equation]
Revised: [Equation]

Original: (1) where ξ is a data sample randomly drawn from distribution Q , (cid:96) ( x ; ξ ) is the corresponding loss function and E Q [ (cid:96) ( x , ξ )] is the expected loss over distribution Q which belongs to uncertainty set U m .
Revised: (1) where ξ is a data sample randomly drawn from distribution Q , (cid:96) ( x ; ξ ) is the corresponding loss function and E Q [ (cid:96) ( x , ξ )] is the expected loss over distribution Q which belongs to uncertainty set U m .

Original: The uncertainty set U m is deﬁned as U m (cid:44) { Q : d ( Q, ˆ P m ) ≤ (cid:15) } 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 .
Revised: The uncertainty set U m is deﬁned as U m (cid:44) { Q : d ( Q, ˆ P m ) ≤ (cid:15) } 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 .

Original: We note this uncertainty set captures the distribution shift hence Eq.
Revised: We note this uncertainty set captures the distribution shift hence Eq.

Original: (1) minimizes the worse data distribution.
Revised: (1) minimizes the worse data distribution.

Original: Prior studies (Ben-Tal et al., 2013; Bertsimas et al., 2018; Blanchet et al., 2019; Esfahani & Kuhn, 2018) considered different uncertainty sets (see Deﬁnition 3.1 in Esfahani & Kuhn (2018)) for which they proposed equivalent reformulations of Eq.
Revised: Prior studies (Ben-Tal et al., 2013; Bertsimas et al., 2018; Blanchet et al., 2019; Esfahani & Kuhn, 2018) considered different uncertainty sets (see Deﬁnition 3.1 in Esfahani & Kuhn (2018)) for which they proposed equivalent reformulations of Eq.

Original: (1) based on the speciﬁc choice of U m .
Revised: (1) based on the speciﬁc choice of U m .

Original: To solve the above min-max optimization problems, majority of prior studies heavily rely on either semideﬁnite 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).
Revised: To solve the above min-max optimization problems, majority of prior studies heavily rely on either semideﬁnite 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).

Original: While primal-dual methods can be used as an approach to solve min-max optimization problems, it suffers from a few downsides.
Revised: While primal-dual methods can be used as an approach to solve min-max optimization problems, it suffers from a few downsides.

Original: First and foremost, they need to store a probability distribution of constrained violation of dimension m corresponding to dual variables.
Revised: First and foremost, they need to store a probability distribution of constrained violation of dimension m corresponding to dual variables.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: To overcome these issues, in this paper, we consider three different reformulations of Eq.
Revised: To overcome these issues, in this paper, we consider three different reformulations of Eq.

Original: (1), corresponding to three different choices of uncertainty sets U m 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:
Revised: (1), corresponding to three different choices of uncertainty sets U m 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:

Original: [Equation]
Revised: [Equation]

Original: where we suppose r ( x ) is convex and a relatively simple function, f ( x ) : R p
Revised: where we suppose r ( x ) is convex and a relatively simple function, f ( x ) : R p

Original: →
Revised: →

Original: R
Revised: R

Original: and h
Revised: and h

Original: i ( x ) : R d → R for 1 ≤ i ≤ m are scalar-valued functions, and g i ( x ) : R d → R p for 1 ≤ i ≤ m are vector-valued functions.
Revised: i ( x ) : R d → R for 1 ≤ i ≤ m are scalar-valued functions, and g i ( x ) : R d → R p for 1 ≤ i ≤ m are vector-valued functions.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: We summarize the main contributions of our paper below
Revised: We summarize the main contributions of our paper below

Original: : • We provide a large-scale analysis of DRO with Wasserstein distance and heavily constrained reformulation when the objective function is strongly-convex.
Revised: : • We provide a large-scale analysis of DRO with Wasserstein distance and heavily constrained reformulation when the objective function is strongly-convex.

Original: Our result relies on a novel mini-batch constraint sampling for handling heavily-constrained optimization problems.
Revised: Our result relies on a novel mini-batch constraint sampling for handling heavily-constrained optimization problems.

Original: As summarized in Table.
Revised: As summarized in Table.

Original: 1, our convergence analysis improves the state-of-the-art both in terms of the dependence on the convergence error (cid:15) as well as the number of constraints m .
Revised: 1, our convergence analysis improves the state-of-the-art both in terms of the dependence on the convergence error (cid:15) as well as the number of constraints m .

Original: • We represent a large-scale analysis of DRO with non-convex objectives and χ 2 or KL divergences.
Revised: • We represent a large-scale analysis of DRO with non-convex objectives and χ 2 or KL divergences.

Original: • We verify our theoretical results through various extensive experiments on different datasets.
Revised: • We verify our theoretical results through various extensive experiments on different datasets.

Original: 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.
Revised: 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.

Original: The proofs of all the theorems are provided in the appendix .
Revised: The proofs of all the theorems are provided in the appendix .

Original: 1. R ELATED W ORK
Revised: 1. R ELATED W ORK

Original: DRO and connections to heavily constrained optimization.
Revised: DRO and connections to heavily constrained optimization.

Original: As mentioned earlier, DRO has many different formulations, depending on the divergence metrics used (e.g., Wasserstein , χ 2 or KL).
Revised: As mentioned earlier, DRO has many different formulations, depending on the divergence metrics used (e.g., Wasserstein , χ 2 or KL).

Original: While Namkoong & Duchi (2016); Shapiro (2017); Duchi & Namkoong (2021) consider constrained or penalized DRO formulation, Sinha et al.
Revised: While Namkoong & Duchi (2016); Shapiro (2017); Duchi & Namkoong (2021) consider constrained or penalized DRO formulation, Sinha et al.

Original: ; Levy et al.
Revised: ; Levy et al.

Original: (2020) formulate the underlying optimization problem as unconstrained.
Revised: (2020) formulate the underlying optimization problem as unconstrained.

Original: 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.
Revised: 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.

Original: In particular, when the objective function is convex, Levy et al.
Revised: In particular, when the objective function is convex, Levy et al.

Original: (2020) recently proposed scalable algorithms for different variants of the DRO problems with , e.g., χ 2 or KL divergence metrics.
Revised: (2020) recently proposed scalable algorithms for different variants of the DRO problems with , e.g., χ 2 or KL divergence metrics.

Original: 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.
Revised: 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.

Original: (2018) studied DRO with χ 2 -divergence but did not provide any convergence guarantee).
Revised: (2018) studied DRO with χ 2 -divergence but did not provide any convergence guarantee).

Original: Similarly, Kuhn et al.
Revised: Similarly, Kuhn et al.

Original: ; Esfahani & Kuhn (2018) formulated DRO with Wasserstein distance as an instance of constrained optimization.
Revised: ; Esfahani & Kuhn (2018) formulated DRO with Wasserstein distance as an instance of constrained optimization.

Original: Notably, they require ti impose one constraint per training data point and to solve such a constrained problem they proposed a semi-deﬁnite program.
Revised: Notably, they require ti impose one constraint per training data point and to solve such a constrained problem they proposed a semi-deﬁnite program.

Original: Even though the formulation is very novel, it cannot scale.
Revised: Even though the formulation is very novel, it cannot scale.

Original: We, in contrast, consider such a heavily constrained optimization as an instance of a composite optimization for which we provide a scalable solution.
Revised: We, in contrast, consider such a heavily constrained optimization as an instance of a composite optimization for which we provide a scalable solution.

Original: What is rather surprising about our method is that it only checks a batch of constraints per iteration, inspired by Cotter et al.
Revised: What is rather surprising about our method is that it only checks a batch of constraints per iteration, inspired by Cotter et al.

Original: (2016), and performs a single projection at the ﬁnal stage of the algorithm in order to provide an (cid:15) -optimal solution in the case of strongly convex objectives.
Revised: (2016), and performs a single projection at the ﬁnal stage of the algorithm in order to provide an (cid:15) -optimal solution in the case of strongly convex objectives.

Original: Moreover, in contrast to Cotter et al.
Revised: Moreover, in contrast to Cotter et al.

Original: (2016), we do not keep a probability distribution over the set of constraints.
Revised: (2016), we do not keep a probability distribution over the set of constraints.

Original: We should also remark that our convergence guarantees achieve the known lower bounds in terms of accuracy (cid:15) and the number of constraints m .
Revised: We should also remark that our convergence guarantees achieve the known lower bounds in terms of accuracy (cid:15) and the number of constraints m .

Original: Finally, we should highlight the difference of our algorithm and Frank-Wolfe (FW) (Frank et al., 1956;
Revised: Finally, we should highlight the difference of our algorithm and Frank-Wolfe (FW) (Frank et al., 1956;

Original: Jaggi, 2013; Zhang et al., 2020b).
Revised: Jaggi, 2013; Zhang et al., 2020b).

Original: While FW does not require a projection oracle, it performs a linear program over the set of constraints at each iteration.
Revised: While FW does not require a projection oracle, it performs a linear program over the set of constraints at each iteration.

Original: In contrast, our heavily- constrained optimization solution performs a single projection without the overhead of running a linear program at each iteration.
Revised: In contrast, our heavily- constrained optimization solution performs a single projection without the overhead of running a linear program at each iteration.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Stochastic composite optimization.
Revised: Stochastic composite optimization.

Original: The general stochastic composite optimization min x (cid:2) Ψ( x ) (cid:44) r ( x ) + f ( E ξ [ g ξ ( x ))] (cid:3) has recently received a lot of attentions (Qi et al., 2020b;a; Wang et al., 2017).
Revised: The general stochastic composite optimization min x (cid:2) Ψ( x ) (cid:44) r ( x ) + f ( E ξ [ g ξ ( x ))] (cid:3) has recently received a lot of attentions (Qi et al., 2020b;a; Wang et al., 2017).

Original: Our reformulation of DRO variants is a ﬁnite-sum instance of this general problem.
Revised: Our reformulation of DRO variants is a ﬁnite-sum instance of this general problem.

Original: More concretely, Lian et al.
Revised: More concretely, Lian et al.

Original: ; Huo et al.
Revised: ; Huo et al.

Original: ; Zhang & Xiao (2019a) aimed to solve the following ﬁnite-sum problem min x
Revised: ; Zhang & Xiao (2019a) aimed to solve the following ﬁnite-sum problem min x

Original: [Equation]
Revised: [Equation]

Original: , using SVRG or SAGA (Defazio et al., 2014).
Revised: , using SVRG or SAGA (Defazio et al., 2014).

Original: In contrast, our proposed algorithm is inspired by
Revised: In contrast, our proposed algorithm is inspired by

Original: Zhang & Xiao (2019a) and generalizes their method to the case where the extra terms
Revised: Zhang & Xiao (2019a) and generalizes their method to the case where the extra terms

Original: h i ( x ) in
Revised: h i ( x ) in

Original: Eq.
Revised: Eq.

Original: (2) are non-zero.
Revised: (2) are non-zero.

Original: We should also note that Qi et al.
Revised: We should also note that Qi et al.

Original: (2020a) proposed a similar idea in the context of online learning for DRO problems with KL divergence.
Revised: (2020a) proposed a similar idea in the context of online learning for DRO problems with KL divergence.

Original: Our work in contrast provides guarantees for DRO with both constraints or penalty terms.
Revised: Our work in contrast provides guarantees for DRO with both constraints or penalty terms.

Original:  DRO VIA F INITE -S UM C OMPOSITE O PTIMIZATION
Revised:  DRO VIA F INITE -S UM C OMPOSITE O PTIMIZATION

Original: In this section, we discuss in detail how a ﬁnite-sum composite optimization (2) can unify various notions of distributionally robust learning, where some of which rely on heavily constrained optimization subroutines.
Revised: In this section, we discuss in detail how a ﬁnite-sum composite optimization (2) can unify various notions of distributionally robust learning, where some of which rely on heavily constrained optimization subroutines.

Original: 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.
Revised: 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.

Original: DRO
Revised: DRO

Original: with Wasserstein distance.
Revised: with Wasserstein distance.

Original: Kuhn et al.
Revised: Kuhn et al.

Original: [Equation]
Revised: [Equation]

Original: ); Esfahani & Kuhn (2018) provide the equivalent and tractable reformulation of Eq.
Revised: ); Esfahani & Kuhn (2018) provide the equivalent and tractable reformulation of Eq.

Original: (1) which can be regarded as a heavily constrained optimization problem as follows:
Revised: (1) which can be regarded as a heavily constrained optimization problem as follows:

Original: [Equation]
Revised: [Equation]

Original: where ˜ g
Revised: where ˜ g

Original: i ( x ) are functions related to loss function as well as slack variables (please see Appendix A for more details).
Revised: i ( x ) are functions related to loss function as well as slack variables (please see Appendix A for more details).

Original: Naively solving optimization problem (3) suffers from the computational complexity due to the large number of constraints m .
Revised: Naively solving optimization problem (3) suffers from the computational complexity due to the large number of constraints m .

Original: To efﬁciently solve the optimization problem (3), inspired by
Revised: To efﬁciently solve the optimization problem (3), inspired by

Original: Mahdavi et al.
Revised: Mahdavi et al.

Original: (2012) and Cotter et al.
Revised: (2012) and Cotter et al.

Original: , we pursue the smoothed constrained reduction approach and introduce the augmented optimization problem (see Appendix B) of the form min x Ψ( x ) (cid:44) [ r ( x ) + γ ln ( g ( x ))] where g i ( x ) (cid:44) exp (cid:16) α
Revised: , we pursue the smoothed constrained reduction approach and introduce the augmented optimization problem (see Appendix B) of the form min x Ψ( x ) (cid:44) [ r ( x ) + γ ln ( g ( x ))] where g i ( x ) (cid:44) exp (cid:16) α

Original: ˜ g
Revised: ˜ g

Original: i ( x ) γ (cid:17) and g ( x ) = 1 m +

Revised: i ( x ) γ (cid:17) and g ( x ) = 1 m +

Original: [1 + (cid:80) mi
Revised: [1 + (cid:80) mi

Original: =

Revised: =

Original: g i ( x )] .
Revised: g i ( x )] .

Original: We can see that  optimization problem(3) is a special case of the optimization problem Eq.
Revised: We can see that this optimization problem is a special case of the optimization problem Eq.

Original: (2) where r ( x ) = f
Revised: (2) where r ( x ) = f

Original: ( x ) , f ( 1 m (cid:80) mi =

Revised: ( x ) , f ( 1 m (cid:80) mi =

Original: g i ( x ))
Revised: g i ( x ))

Original: =
Revised: =

Original: γ ln g ( x ) , and h ( x ) = 0 .
Revised: γ ln g ( x ) , and h ( x ) = 0 .

Original: 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 (cid:15) , we propose an algorithm that simply checks a batch of constraints and achieves the optimum dependency in terms of m and (cid:15) .
Revised: 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 (cid:15) , we propose an algorithm that simply checks a batch of constraints and achieves the optimum dependency in terms of m and (cid:15) .

Original: DRO with χ 2 -divergence.
Revised: DRO with χ 2 -divergence.

Original: The second type of DRO problem we consider utilizes the χ 2divergence metric as follows:
Revised: The second type of DRO problem we consider utilizes the χ 2divergence metric as follows:

Original: [Equation]
Revised: [Equation]

Original: where the χ 2 divergence is deﬁned as the distance between the uniform distribution and an arbitrary probability distribution p , i.e., D χ 2 ( p ) (cid:44) m 2 (cid:80) mi =1 (cid:0) p i − 1 m (cid:1) 2 .
Revised: where the χ 2 divergence is deﬁned as the distance between the uniform distribution and an arbitrary probability distribution p , i.e., D χ 2 ( p ) (cid:44) m 2 (cid:80) mi =1 (cid:0) p i − 1 m (cid:1) 2 .

Original: Levy et al.
Revised: Levy et al.

Original: (2020) studied this problem only for the case of convex objectives.
Revised: (2020) studied this problem only for the case of convex objectives.

Original: In this paper, we allow objective functions
Revised: In this paper, we allow objective functions

Original: f i for 1 ≤ i ≤ m to be both non-convex or strongly-convex.
Revised: f i for 1 ≤ i ≤ m to be both non-convex or strongly-convex.

Original: The following claim derives the equivalent ﬁnite-sum composite optimization.
Revised: The following claim derives the equivalent ﬁnite-sum composite optimization.

Original: Claim 2.1.
Revised: Claim 2.1.

Original: The optimization problem (4) is equivalent to the following problem:
Revised: The optimization problem (4) is equivalent to the following problem:

Original: [Equation]
Revised: [Equation]

Original: We note that the optimization problem (5) ﬁts into the formulation of ﬁnite-sum composite optimization (2) by choosing r ( x )
Revised: We note that the optimization problem (5) ﬁts into the formulation of ﬁnite-sum composite optimization (2) by choosing r ( x )

Original: = 1 , h ( x ) = 1 m (cid:80) mi
Revised: = 1 , h ( x ) = 1 m (cid:80) mi

Original: [Equation]
Revised: [Equation]

Original: 1 − ( f i ( x
Revised: 1 − ( f i ( x

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: with h i ( x ) = − f 2 i ( x ) 2 γ , g i ( x ) =
Revised: with h i ( x ) = − f 2 i ( x ) 2 γ , g i ( x ) =

Original: f i ( x ) and f ( x ) = x
Revised: f i ( x ) and f ( x ) = x

Original: 2 2 γ .
Revised: 2 2 γ .

Original: DRO with KL divergence.
Revised: DRO with KL divergence.

Original: Finally, for DRO with KL-divergence, usually considered in online settings (Qi et al., 2020a), we consider solving the following optimization problem:
Revised: Finally, for DRO with KL-divergence, usually considered in online settings (Qi et al., 2020a), we consider solving the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where H ( p 1 , . . . , p m ) = − (cid:80) mi =1 p
Revised: where H ( p 1 , . . . , p m ) = − (cid:80) mi =1 p

Original: i log p i is the entropy function.
Revised: i log p i is the entropy function.

Original: To solve the min-max problem (6), it is straightforward to convert it to the equivalent stochastic composite optimization problem of the following form:
Revised: To solve the min-max problem (6), it is straightforward to convert it to the equivalent stochastic composite optimization problem of the following form:

Original: [Equation]
Revised: [Equation]

Original: As it can be seen, the optimization problem (7) ﬁts into the composite optimization (2) by choosing r ( x ) = h ( x ) = 

Revised: As it can be seen, the optimization problem (7) ﬁts into the composite optimization (2) by choosing r ( x ) = h ( x ) =

Original: and f ( g ( x ))
Revised: and f ( g ( x ))

Original: =
Revised: =

Original: ln (cid:16) 1 m (cid:80) mi =1 exp (cid:16) f i (
Revised: ln (cid:16) 1 m (cid:80) mi =1 exp (cid:16) f i (

Original: [Equation]
Revised: [Equation]

Original: γ (cid:17)(cid:17) .
Revised: γ (cid:17)(cid:17) .

Original:  O UR P ROPOSED A LGORITHM
Revised:  O UR P ROPOSED A LGORITHM

Original: 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 )
Revised: 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 )

Original: =
Revised: =

Original: r ( x ) + Φ( x ) where Φ( x ) = 1 m (cid:80) mi =1 h i ( x ) + f ( 1 m (cid:80) mi
Revised: r ( x ) + Φ( x ) where Φ( x ) = 1 m (cid:80) mi =1 h i ( x ) + f ( 1 m (cid:80) mi

Original: =

Revised: =

Original: g i ( x )) .
Revised: g i ( x )) .

Original: We note that the compositional structure in Φ( · ) leads to more challenges in optimization compared with the non-compositional ﬁnite-sum problem, since the stochastic gradient of the loss function is not an unbiased estimation of the full gradient.
Revised: We note that the compositional structure in Φ( · ) leads to more challenges in optimization compared with the non-compositional ﬁnite-sum problem, since the stochastic gradient of the loss function is not an unbiased estimation of the full gradient.

Original: 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 ) =
Revised: 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 ) =

Original: (1 /m ) (cid:80) mi =1 h i ( x ) .
Revised: (1 /m ) (cid:80) mi =1 h i ( x ) .

Original: To handle r ( · ) , similar to Zhang & Xiao (2019b), we assume r is convex and a relatively simple function.
Revised: To handle r ( · ) , similar to Zhang & Xiao (2019b), we assume r is convex and a relatively simple function.

Original: We follow the proximal gradient iterates (Nesterov, 2013; Beck, 2017) as follows:
Revised: We follow the proximal gradient iterates (Nesterov, 2013; Beck, 2017) as follows:

Original: [Equation]
Revised: [Equation]

Original: where we apply the proximal operator of r ( x ) with the learning rate η as Π ηr (
Revised: where we apply the proximal operator of r ( x ) with the learning rate η as Π ηr (

Original: x ) (cid:44) arg min y (cid:2) r ( y )
Revised: x ) (cid:44) arg min y (cid:2) r ( y )

Original: + 12 η (cid:107) y − x (cid:107) 2 (cid:3) .
Revised: + 12 η (cid:107) y − x (cid:107) 2 (cid:3) .

Original: By deﬁning the proximal gradient mapping of Ψ as G η ( x ) (cid:44) 1 η [ x − Π η r ( x − η ∇ Φ( x ))] , the updating rule in Eq.
Revised: By deﬁning the proximal gradient mapping of Ψ as G η ( x ) (cid:44) 1 η [ x − Π η r ( x − η ∇ Φ( x ))] , the updating rule in Eq.

Original: (8) can be equivalently written as x ( t +1) =
Revised: (8) can be equivalently written as x ( t +1) =

Original: x ( t ) − η G η ( x ( t ) ) .
Revised: x ( t ) − η G η ( x ( t ) ) .

Original: Given any y as an output of randomized algorithm, we say y achieves the stationary point of problem in Eq.
Revised: Given any y as an output of randomized algorithm, we say y achieves the stationary point of problem in Eq.

Original: (2) in expectation if E (cid:2) (cid:107)G η ( y ) (cid:107) 2 (cid:3) ≤ (cid:15) holds.
Revised: (2) in expectation if E (cid:2) (cid:107)G η ( y ) (cid:107) 2 (cid:3) ≤ (cid:15) holds.

Original: Our goal is to achieve an (cid:15) stationary point with the least number of calls to a (mini-batch) stochastic oracle.
Revised: Our goal is to achieve an (cid:15) stationary point with the least number of calls to a (mini-batch) stochastic oracle.

Original: Focusing on Φ( · ) , as detailed in Algorithm 1, we apply three time-scale variance-reduced estimators for g i ( x ) and its gradient ∇ g i ( x ) , as well as h i ( x ) .
Revised: Focusing on Φ( · ) , as detailed in Algorithm 1, we apply three time-scale variance-reduced estimators for g i ( x ) and its gradient ∇ g i ( x ) , as well as h i ( x ) .

Original: For DRO with Wasserstein divergence metric with for t = 1 , . . .
Revised: For DRO with Wasserstein divergence metric with for t = 1 , . . .

Original: , T do
Revised: , T do

Original: Sample a mini-batch B ( t ) with size B t uniformly over [ m ] and compute
Revised: Sample a mini-batch B ( t ) with size B t uniformly over [ m ] and compute

Original: [Equation]
Revised: [Equation]

Original: Compute ˜ ∇ Φ( x ( t ) 0 ) =
Revised: Compute ˜ ∇ Φ( x ( t ) 0 ) =

Original: (cid:16) z ( t ) 0 (cid:17) (cid:62) (cid:16) f (cid:48)
Revised: (cid:16) z ( t ) 0 (cid:17) (cid:62) (cid:16) f (cid:48)

Original: ( y ( t ) 0 ) (cid:17) + w ( t )
Revised: ( y ( t ) 0 ) (cid:17) + w ( t )

Original: [Equation]
Revised: [Equation]

Original: Update the model as follows:
Revised: Update the model as follows:

Original: x ( t ) 1 =
Revised: x ( t ) 1 =

Original: Π ηr (cid:16) x ( t ) 

Revised: Π ηr (cid:16) x ( t )

Original: − ˜ ∇ Φ( x ( t ) 0 ))
Revised: − ˜ ∇ Φ( x ( t ) 0 ))

Original: (cid:17) for j = 1 , . . .
Revised: (cid:17) for j = 1 , . . .

Original: , τ t − 1 do in parallel
Revised: , τ t − 1 do in parallel

Original: Sample a mini-batch S ( t ) j with size S t uniformly over [ m ] , and form the estimates
Revised: Sample a mini-batch S ( t ) j with size S t uniformly over [ m ] , and form the estimates

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Update the model as follows: x
Revised: Update the model as follows: x

Original: [Equation]
Revised: [Equation]

Original: end end
Revised: end end

Original: Output: Return a randomly selected solution from { x tj } t =1 ,...,T j =0 ,...,τ t optimally strongly objective, at the beginning of each epoch t we compute a full-batch gradient over the entire data samples B t = m , i.e., y ( t ) 

Revised: Output: Return a randomly selected solution from { x tj } t =1 ,...,T j =0 ,...,τ t optimally strongly objective, at the beginning of each epoch t we compute a full-batch gradient over the entire data samples B t = m , i.e., y ( t )

Original: = 1 m (cid:80) mi
Revised: = 1 m (cid:80) mi

Original: =

Revised: =

Original: g i ( x ( t ) τ t ) , z ( t ) 

Revised: g i ( x ( t ) τ t ) , z ( t )

Original: = 1 m (cid:80) mi =1 ∇ g i ( x ( t ) τ t ) , w ( t ) 

Revised: = 1 m (cid:80) mi =1 ∇ g i ( x ( t ) τ t ) , w ( t )

Original: = 1 m (cid:80) mi =1 ∇ h
Revised: = 1 m (cid:80) mi =1 ∇ h

Original: i ( x ( t ) τ t ) .
Revised: i ( x ( t ) τ t ) .

Original: 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 B t ≤ m in the beginning of each epoch until it reaches the point where we need to compute the full-batch of samples.
Revised: 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 B t ≤ m in the beginning of each epoch until it reaches the point where we need to compute the full-batch of samples.

Original: We should also highlight that a mini-batch in case of Wasserstein DRO indicates a batch of constraints and in
Revised: We should also highlight that a mini-batch in case of Wasserstein DRO indicates a batch of constraints and in

Original: DRO with χ 2 or KL divergence metrics represents the number of data sample accessed.
Revised: DRO with χ 2 or KL divergence metrics represents the number of data sample accessed.

Original: We denote the length of each epoch and the mini-batch size within each epoch with τ t and S t , respectively.
Revised: We denote the length of each epoch and the mini-batch size within each epoch with τ t and S t , respectively.

Original: In each epoch t and iteration j , we estimate mini-batch gradients from Eqs.
Revised: In each epoch t and iteration j , we estimate mini-batch gradients from Eqs.

Original: (9), (10), and (11) in Algorithm 1, with some corrections term applied from the previous iteration.
Revised: (9), (10), and (11) in Algorithm 1, with some corrections term applied from the previous iteration.

Original: 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).
Revised: 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).

Original: Finally, the algorithm returns a randomly selected solution from the iterates.
Revised: Finally, the algorithm returns a randomly selected solution from the iterates.

Original: Distributed variant of Algorithm 1.
Revised: Distributed variant of Algorithm 1.

Original: As mentioned before, efﬁcient training of the stochastic composite optimization problem has attracted increasing attention in recent years.
Revised: As mentioned before, efﬁcient training of the stochastic composite optimization problem has attracted increasing attention in recent years.

Original: Despite much progress, all of existing methods including Algorithm 1 only focus on the single-machine setting.
Revised: Despite much progress, all of existing methods including Algorithm 1 only focus on the single-machine setting.

Original: 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.
Revised: 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.

Original:  C ONVERGENCE A NALYSIS
Revised:  C ONVERGENCE A NALYSIS

Original: In this section we establish the convergence of proposed algorithm for different DRO notions discussed in Section 2.
Revised: In this section we establish the convergence of proposed algorithm for different DRO notions discussed in Section 2.

Original: We start by stating the general assumptions and then discuss the obtained rates.
Revised: We start by stating the general assumptions and then discuss the obtained rates.

Original: 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.
Revised: 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.

Original: We defer the analysis of χ 2 and KL divergence metrics with strongly objectives to Appendix D. To establish the convergence rates, we ﬁrst introduce some standard assumptions.
Revised: We defer the analysis of χ 2 and KL divergence metrics with strongly objectives to Appendix D. To establish the convergence rates, we ﬁrst introduce some standard assumptions.

Original: Assumption 1.
Revised: Assumption 1.

Original: We make the following assumptions on the components of objective Ψ( x ) =
Revised: We make the following assumptions on the components of objective Ψ( x ) =

Original: r ( x ) +m (cid:80) mi =1 h i ( x ) + f ( 1 m (cid:80) mi =

Revised: r ( x ) +m (cid:80) mi =1 h i ( x ) + f ( 1 m (cid:80) mi =

Original: g i ( x )) :
Revised: g i ( x )) :

Original: 1) (cid:107)∇ h ( x 1 ) − ∇ h ( x 2 ) (cid:107) ≤ L h (cid:107) x 1 − x 2 (cid:107) where x 1 , x 2 ∈ R d .
Revised: 1) (cid:107)∇ h ( x 1 ) − ∇ h ( x 2 ) (cid:107) ≤ L h (cid:107) x 1 − x 2 (cid:107) where x 1 , x 2 ∈ R d .

Original: We also assume that (cid:107) h ( x 1 ) − h ( x 2 ) (cid:107) ≤ (cid:96) h (cid:107) x 1 − x 2 (cid:107) .
Revised: We also assume that (cid:107) h ( x 1 ) − h ( x 2 ) (cid:107) ≤ (cid:96) h (cid:107) x 1 − x 2 (cid:107) .

Original: 2) (cid:107) f
Revised: 2) (cid:107) f

Original: (cid:48) ( x 1 ) − f (cid:48) ( x 2 ) (cid:107) ≤
Revised: (cid:48) ( x 1 ) − f (cid:48) ( x 2 ) (cid:107) ≤

Original: L f (cid:107) x 1 − x 2 (cid:107) where x 1 , x 2 ∈ R .
Revised: L f (cid:107) x 1 − x 2 (cid:107) where x 1 , x 2 ∈ R .

Original: We also assume that (cid:107) f
Revised: We also assume that (cid:107) f

Original: ( x 1 ) − f ( x 2 ) (cid:107) ≤ (cid:96) f
Revised: ( x 1 ) − f ( x 2 ) (cid:107) ≤ (cid:96) f

Original: (cid:107) x 1 − x 2 (cid:107) .
Revised: (cid:107) x 1 − x 2 (cid:107) .

Original: (cid:107)∇ g ( x 1 ; ξ ) − ∇ g ( x 2 ; ξ ) (cid:107) ≤ L g (cid:107) x
Revised: (cid:107)∇ g ( x 1 ; ξ ) − ∇ g ( x 2 ; ξ ) (cid:107) ≤ L g (cid:107) x

Original: 1 − x 2 (cid:107) and (cid:107) g
Revised: 1 − x 2 (cid:107) and (cid:107) g

Original: ( x 1 ; ξ ) − g ( x 2 ; ξ ) (cid:107) ≤ (cid:96) g (cid:107) x 1 − x 2 (cid:107) for ∀ x 1 , x 2 ∈ R d . 4)
Revised: ( x 1 ; ξ ) − g ( x 2 ; ξ ) (cid:107) ≤ (cid:96) g (cid:107) x 1 − x 2 (cid:107) for ∀ x 1 , x 2 ∈ R d . 4)

Original: We suppose that Ψ( x ) in Eq.
Revised: We suppose that Ψ( x ) in Eq.

Original: (2) is bounded from below that Ψ ∗ =
Revised: (2) is bounded from below that Ψ ∗ =

Original: inf x Ψ( x ) >
Revised: inf x Ψ( x ) >

Original: −∞ . 5)
Revised: −∞ . 5)

Original: We assume r ( x ) ∈ R ∪ {∞} is a convex and lower-semicontinues function.
Revised: We assume r ( x ) ∈ R ∪ {∞} is a convex and lower-semicontinues function.

Original: An immediate implication of above assumption is that f ( g ( x ))
Revised: An immediate implication of above assumption is that f ( g ( x ))

Original: (cid:44) f
Revised: (cid:44) f

Original: ( 1 m (cid:80) mi =

Revised: ( 1 m (cid:80) mi =

Original: g i ( x )) is smooth with module (cid:0) (cid:96) 2 g L f
Revised: g i ( x )) is smooth with module (cid:0) (cid:96) 2 g L f

Original: + (cid:96) f L g (cid:1) , hence Ψ( x ) is smooth with module L Φ (cid:44) (cid:2)(cid:0) (cid:96) 2 g L f
Revised: + (cid:96) f L g (cid:1) , hence Ψ( x ) is smooth with module L Φ (cid:44) (cid:2)(cid:0) (cid:96) 2 g L f

Original: + (cid:96) f L g (cid:1)
Revised: + (cid:96) f L g (cid:1)

Original: + L h (cid:3) (Wang et al., 2016; Zhang & Xiao, 2019b).
Revised: + L h (cid:3) (Wang et al., 2016; Zhang & Xiao, 2019b).

Original: Assumption 2 (Unbiased estimation) .
Revised: Assumption 2 (Unbiased estimation) .

Original: We assume that the mini-batch of samples ξ over functions
Revised: We assume that the mini-batch of samples ξ over functions

Original: g i for i = 1 , 2 , . . .
Revised: g i for i = 1 , 2 , . . .

Original: , m is unbiased, that is E
Revised: , m is unbiased, that is E

Original: [ g ( x ; ξ )]
Revised: [ g ( x ; ξ )]

Original: = g ( x ) and E [ ∇ g ( x ; ξ )]
Revised: = g ( x ) and E [ ∇ g ( x ; ξ )]

Original: = ∇ g ( x ) .
Revised: = ∇ g ( x ) .

Original: 4. O PTIMALLY S TRONGLY C ONVEX O BJECTIVES
Revised: 4. O PTIMALLY S TRONGLY C ONVEX O BJECTIVES

Original: We here establish the convergence for optimally strongly convex objectives under Wasserstein metric as deﬁned below.
Revised: We here establish the convergence for optimally strongly convex objectives under Wasserstein metric as deﬁned below.

Original: Deﬁnition 1 (Optimally strongly convex) .
Revised: Deﬁnition 1 (Optimally strongly convex) .

Original: We say that the objective function Ψ( x ) is optimally strongly convex objective with module µ if Ψ( x ) − Ψ
Revised: We say that the objective function Ψ( x ) is optimally strongly convex objective with module µ if Ψ( x ) − Ψ

Original: ( P x ∗ ( x )) ≥ µ (cid:107) x − P x ∗ ( x ) (cid:107) 2 .
Revised: ( P x ∗ ( x )) ≥ µ (cid:107) x − P x ∗ ( x ) (cid:107) 2 .

Original: Following Cotter et al.
Revised: Following Cotter et al.

Original: (2016) and Mahdavi et al.
Revised: (2016) and Mahdavi et al.

Original: (2012), we make the following assumption.
Revised: (2012), we make the following assumption.

Original: Assumption 3.
Revised: Assumption 3.

Original: There exists a constant ρ such that if we deﬁne g ( x ) = max 1 ≤ i ≤ m ˜ g
Revised: There exists a constant ρ such that if we deﬁne g ( x ) = max 1 ≤ i ≤ m ˜ g

Original: i ( x ) we have min g ( x )=

Revised: i ( x ) we have min g ( x )=

Original: (cid:107)∇ g ( x ) (cid:107) 2 ≥ ρ .
Revised: (cid:107)∇ g ( x ) (cid:107) 2 ≥ ρ .

Original: Assumption 4.
Revised: Assumption 4.

Original: Function r ( x ) is smooth with module G .
Revised: Function r ( x ) is smooth with module G .

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Let us assume that Ψ is µ -optimally strongly convex.
Revised: Let us assume that Ψ is µ -optimally strongly convex.

Original: Let us set τ t = S t = √ B t = √ m ,
Revised: Let us set τ t = S t = √ B t = √ m ,

Original: T = 5 √ mµη , η < 2 L Φ + √ L 2Φ +36 G 0 where G 0 (cid:44) 3( (cid:96) 4 g L 2 f + (cid:96) 2 f L 2 g + (cid:96) 2 h ) , α > Gρ , and γ = exp( − K ) ln( m +1) .
Revised: T = 5 √ mµη , η < 2 L Φ + √ L 2Φ +36 G 0 where G 0 (cid:44) 3( (cid:96) 4 g L 2 f + (cid:96) 2 f L 2 g + (cid:96) 2 h ) , α > Gρ , and γ = exp( − K ) ln( m +1) .

Original: Let us denote x
Revised: Let us denote x

Original: ( k +1) =
Revised: ( k +1) =

Original: GCIVR (cid:0) x ( k ) (cid:1) for k = 0 , . .
Revised: GCIVR (cid:0) x ( k ) (cid:1) for k = 0 , . .

Original: , K − 1 (using Algorithm 1).
Revised: , K − 1 (using Algorithm 1).

Original: Under assumptions 1-4 and by letting K = ln (1 /(cid:15) ) , the solution of DRO with Wasserstein divergence is obtained by projecting x ( K ) onto the constraint set K = { x | g i ( x ) ≤ 0 , i = 1 , 2 , . . .
Revised: Under assumptions 1-4 and by letting K = ln (1 /(cid:15) ) , the solution of DRO with Wasserstein divergence is obtained by projecting x ( K ) onto the constraint set K = { x | g i ( x ) ≤ 0 , i = 1 , 2 , . . .

Original: , m } , i.e., ¯ x ( K ) =
Revised: , m } , i.e., ¯ x ( K ) =

Original: Π K ( x ( K ) ) .
Revised: Π K ( x ( K ) ) .

Original: In order to achieve r (¯ x ( K ) ) − r ( x ( ∗ ) ) ≤ (cid:15) we require O (cid:0) ( m + κ √ m ) ln 1 (cid:15) (cid:1) number of calls to the stochastic oracle.
Revised: In order to achieve r (¯ x ( K ) ) − r ( x ( ∗ ) ) ≤ (cid:15) we require O (cid:0) ( m + κ √ m ) ln 1 (cid:15) (cid:1) number of calls to the stochastic oracle.

Original: Comparison with previous results.
Revised: Comparison with previous results.

Original: Compared to the MidTouch approach by Cotter et al.
Revised: Compared to the MidTouch approach by Cotter et al.

Original: , our obtained rate improves both on the dependency on the number of constraints m , as well as convergence error (cid:15) .
Revised: , our obtained rate improves both on the dependency on the number of constraints m , as well as convergence error (cid:15) .

Original: To better understand the intuition behind achieving such a double folded improvement in terms of m and (cid:15) , we note that unlike Cotter et al.
Revised: To better understand the intuition behind achieving such a double folded improvement in terms of m and (cid:15) , we note that unlike Cotter et al.

Original: (2016) which utilizes a primal-dual approach in order to obtain an approximate to the optimal distribution over the constraints, we directly ﬁnd the optimal distribution exactly.
Revised: (2016) which utilizes a primal-dual approach in order to obtain an approximate to the optimal distribution over the constraints, we directly ﬁnd the optimal distribution exactly.

Original: Second, while Cotter et al.
Revised: Second, while Cotter et al.

Original: (2016) does not apply any variance reduction technique to primal variable x , our algorithm beneﬁts from variance reduction over x too.
Revised: (2016) does not apply any variance reduction technique to primal variable x , our algorithm beneﬁts from variance reduction over x too.

Original: Furthermore, Theorem 4.1 entails tighter rate in terms of ﬁnal accuracy compared to Yu et al.
Revised: Furthermore, Theorem 4.1 entails tighter rate in terms of ﬁnal accuracy compared to Yu et al.

Original: These results are summarized in Table.
Revised: These results are summarized in Table.

Original: In Algorithm 1 for optimally strongly convex objectives, we need to do a single projection at the end.
Revised: In Algorithm 1 for optimally strongly convex objectives, we need to do a single projection at the end.

Original: The following corollary bounds the error between the solution before and after the projection.
Revised: The following corollary bounds the error between the solution before and after the projection.

Original: Corollary 4.2.
Revised: Corollary 4.2.

Original: Under the assumptions made in Theorem 4.1, the error between the solution with and without projection is bounded by
Revised: Under the assumptions made in Theorem 4.1, the error between the solution with and without projection is bounded by

Original: [Equation]
Revised: [Equation]

Original: Therefore, with a proper choice of γ we can make the right hand side matches the convergence error of Theorem 4.1.
Revised: Therefore, with a proper choice of γ we can make the right hand side matches the convergence error of Theorem 4.1.

Original: Remark 1.
Revised: Remark 1.

Original: 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.
Revised: 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.

Original: In particular, Theorem 4. 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.
Revised: In particular, Theorem 4. 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.

Original: 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.
Revised: 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.

Original: Remark 2.
Revised: Remark 2.

Original: To understand the tightness of our obtained rate, consider min x (cid:104) Ψ( x ) (cid:44)
Revised: To understand the tightness of our obtained rate, consider min x (cid:104) Ψ( x ) (cid:44)

Original: r
Revised: r

Original: ( x ) + 1 m (cid:80) mi
Revised: ( x ) + 1 m (cid:80) mi

Original: =

Revised: =

Original: g i ( x ) (cid:105) which is an instance of our general optimization problem (2).
Revised: g i ( x ) (cid:105) which is an instance of our general optimization problem (2).

Original: Clearly, any lower bound to solve this instance also holds for the original optimization problem (2).
Revised: Clearly, any lower bound to solve this instance also holds for the original optimization problem (2).

Original: Xie et al.
Revised: Xie et al.

Original: (2019) provides a lower-bound of O (( m + √ κm ) ln
Revised: (2019) provides a lower-bound of O (( m + √ κm ) ln

Original: (1 /(cid:15) ))
Revised: (1 /(cid:15) ))

Original: 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.
Revised: 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.

Original: Furthermore, for the DRO problem with Wasserstein divergence metric, rather than solving constrained optimization problem we solve unconstrained compositional optimization problem.
Revised: Furthermore, for the DRO problem with Wasserstein divergence metric, rather than solving constrained optimization problem we solve unconstrained compositional optimization problem.

Original: 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.
Revised: 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.

Original: Remark 3.
Revised: Remark 3.

Original: In distributed setting with p devices, we are able to improve the sample complexity to
Revised: In distributed setting with p devices, we are able to improve the sample complexity to

Original: O (cid:16)(cid:16) m
Revised: O (cid:16)(cid:16) m

Original: + mp
Revised: + mp

Original: + κ √ m (cid:17) ln 1 (cid:15) (cid:17) .
Revised: + κ √ m (cid:17) ln 1 (cid:15) (cid:17) .

Original: The proof can be found in Appendix F.3.
Revised: The proof can be found in Appendix F.3.

Original: 4. N ON -C ONVEX O BJECTIVES
Revised: 4. N ON -C ONVEX O BJECTIVES

Original: In this section, we estbalish the convergence of DRO problems with χ 2 and KL divergence metrics for non-convex objectives.
Revised: In this section, we estbalish the convergence of DRO problems with χ 2 and KL divergence metrics for non-convex objectives.

Original: To this end, we make an additional assumption on the variance of stochastic mini-batches.
Revised: To this end, we make an additional assumption on the variance of stochastic mini-batches.

Original: Assumption 5 (Bounded variance) .
Revised: Assumption 5 (Bounded variance) .

Original: The mini-batch sampling has bounded variance that is
Revised: The mini-batch sampling has bounded variance that is

Original: E (cid:2) (cid:107)∇ g ( x ; ξ ) − ∇ g ( x ) (cid:107) 2 (cid:3) ≤ σ 2 B where B indicates the batch size.
Revised: E (cid:2) (cid:107)∇ g ( x ; ξ ) − ∇ g ( x ) (cid:107) 2 (cid:3) ≤ σ 2 B where B indicates the batch size.

Original: Theorem 4.3.
Revised: Theorem 4.3.

Original: Under Assumptions 1,2, and 5 using Algorithm 1 for some positive constants β and 0 ≤ ζ < √ m , denote T 0 (cid:44) √ m − ζ β
Revised: Under Assumptions 1,2, and 5 using Algorithm 1 for some positive constants β and 0 ≤ ζ < √ m , denote T 0 (cid:44) √ m − ζ β

Original: = O ( √ m ) .
Revised: = O ( √ m ) .

Original: For t ≤ T 0 , let us set τ t =
Revised: For t ≤ T 0 , let us set τ t =

Original: S t = √ B t = βt
Revised: S t = √ B t = βt

Original: + ζ and for t > T 0 set τ t =
Revised: + ζ and for t > T 0 set τ t =

Original: S t = √ m .
Revised: S t = √ m .

Original: Then, if η < 4 L Φ + √ L 2Φ +12 G 0 , after T =
Revised: Then, if η < 4 L Φ + √ L 2Φ +12 G 0 , after T =

Original: ˜
Revised: ˜

Original: O
Revised: O

Original: (min(1 / √ (cid:15), 1 / √ m(cid:15) )) iterations, with ˜ O (cid:0) min {√ m/(cid:15),
Revised: (min(1 / √ (cid:15), 1 / √ m(cid:15) )) iterations, with ˜ O (cid:0) min {√ m/(cid:15),

Original: 1 /(cid:15) 1 . 5 } (cid:1) number of calls to the stochastic oracle it holds that
Revised: 1 /(cid:15) 1 . 5 } (cid:1) number of calls to the stochastic oracle it holds that

Original: E (cid:2) (cid:13)(cid:13)
Revised: E (cid:2) (cid:13)(cid:13)

Original: G η ( x ( T ) )
Revised: G η ( x ( T ) )

Original: (cid:13)(cid:13) 2 (cid:3) ≤ (cid:15) , where we used ˜ O ( . ) notation to hide terms with logarithmic dependency.
Revised: (cid:13)(cid:13) 2 (cid:3) ≤ (cid:15) , where we used ˜ O ( . ) notation to hide terms with logarithmic dependency.

Original: We also remark that Assumption 5 is only needed for convergence until T 0 = √ m − ζ β and after t > Twe will not need this assumption.
Revised: We also remark that Assumption 5 is only needed for convergence until T 0 = √ m − ζ β and after t > Twe will not need this assumption.

Original: 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.
Revised: 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.

Original: Discussion on lower bound.
Revised: Discussion on lower bound.

Original: As we discussed in Remark 2, any impossibility result for optimizing min x (cid:104) Ψ
Revised: As we discussed in Remark 2, any impossibility result for optimizing min x (cid:104) Ψ

Original: ( x ) (cid:44) r ( x ) + 1 m (cid:80) mi =1 g i ( x ) (cid:105) also holds for general compositional optimization problem in Eq.
Revised: ( x ) (cid:44) r ( x ) + 1 m (cid:80) mi =1 g i ( x ) (cid:105) also holds for general compositional optimization problem in Eq.

Original: For former problem, Fang et al.
Revised: For former problem, Fang et al.

Original: (2018) shows that to achieve the stationary point with error (cid:15) we require at least ˜ O (cid:16) min { √ m (cid:15) , 1 (cid:15) 1 . 5 } (cid:17) (almost optimal) gradient oracle calls for strongly optimal objectives.
Revised: (2018) shows that to achieve the stationary point with error (cid:15) we require at least ˜ O (cid:16) min { √ m (cid:15) , 1 (cid:15) 1 . 5 } (cid:17) (almost optimal) gradient oracle calls for strongly optimal objectives.

Original: 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.
Revised: 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.

Original: Remark 4.
Revised: Remark 4.

Original: In distributed setting with p devices, we can improve the convergence rate of DRO with χ 2 and KL metrics to ˜ O (cid:16) min { √ mp(cid:15) + √ m(cid:15) , 1 p(cid:15) 1 .
Revised: In distributed setting with p devices, we can improve the convergence rate of DRO with χ 2 and KL metrics to ˜ O (cid:16) min { √ mp(cid:15) + √ m(cid:15) , 1 p(cid:15) 1 .

Original: 5 + 1 (cid:15) 1 . 5 } (cid:17) .
Revised: 5 + 1 (cid:15) 1 . 5 } (cid:17) .

Original: For details, please refer to Appendix F.3.
Revised: For details, please refer to Appendix F.3.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we empirically examine the efﬁcacy of the proposed algorithm in different use cases.
Revised: In this section, we empirically examine the efﬁcacy of the proposed algorithm in different use cases.

Original: The main algorithm that we compare against is Heavily-Constrained algorithm that uses a parametric multiplier model proposed in Narasimhan et al.
Revised: The main algorithm that we compare against is Heavily-Constrained algorithm that uses a parametric multiplier model proposed in Narasimhan et al.

Original: (2020b).
Revised: (2020b).

Original: We call this algorithm Heavily-constrained, where they learn the Lagrange multiplier values using a parametric model such as a neural network.
Revised: We call this algorithm Heavily-constrained, where they learn the Lagrange multiplier values using a parametric model such as a neural network.

Original: The tasks we apply the algorithms are mainly focused on fairness constraints since they ﬁt greatly to the problem deﬁnition in this paper due to the large number of constraints they add to the main learning problem.
Revised: The tasks we apply the algorithms are mainly focused on fairness constraints since they ﬁt greatly to the problem deﬁnition in this paper due to the large number of constraints they add to the main learning problem.

Original: In addition, similar to distribution robust optimization approaches, the main goal of fairness constraint is to ﬁnd a solution that is agnostic to the distribution of protected groups.
Revised: In addition, similar to distribution robust optimization approaches, the main goal of fairness constraint is to ﬁnd a solution that is agnostic to the distribution of protected groups.

Original: The code for the experiments is available at this link.
Revised: The code for the experiments is available at this link.

Original: Distributionally robust optimization for fairness constraints.
Revised: Distributionally robust optimization for fairness constraints.

Original: The ﬁrst experiment setup is based on Wang et al.
Revised: The ﬁrst experiment setup is based on Wang et al.

Original: Narasimhan et al.
Revised: Narasimhan et al.

Original: (2020b), where we want to enforce equality of opportunity
Revised: (2020b), where we want to enforce equality of opportunity

Original: (Hardt et al., 2016) constraints for different groups while the group membership is noisy and changing during the training.
Revised: (Hardt et al., 2016) constraints for different groups while the group membership is noisy and changing during the training.

Original: Hence, the problem is to make the solution distributionally robust among different protected groups in the problem.
Revised: Hence, the problem is to make the solution distributionally robust among different protected groups in the problem.

Original: Based on the setup in Narasimhan et al.
Revised: Based on the setup in Narasimhan et al.

Original: (2020b), we assume we have access to the marginal probability of the true groups ( P ( g i =
Revised: (2020b), we assume we have access to the marginal probability of the true groups ( P ( g i =

Original: j | ˆ g i = k ) , where g i is the true group membership and ˆ g i is the noisy group membership).
Revised: j | ˆ g i = k ) , where g i is the true group membership and ˆ g i is the noisy group membership).

Original: Hence, to enforce fairness constraints, we consider all possible proxy groups using this marginal probability, which can increase the number of constraints greatly.
Revised: Hence, to enforce fairness constraints, we consider all possible proxy groups using this marginal probability, which can increase the number of constraints greatly.

Original: 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 ) − (cid:15) for every proxy group we deﬁne.
Revised: 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 ) − (cid:15) for every proxy group we deﬁne.

Original: In this experiment, we use the Adult dataset (Dua & Graff, 2017), and consider race groups of “white”, “black”, and “other” as protected groups.
Revised: In this experiment, we use the Adult dataset (Dua & Graff, 2017), and consider race groups of “white”, “black”, and “other” as protected groups.

Original: We train a linear classiﬁer with logistic regression, and report the overall error rate of the classiﬁer, as well as the maximum violation of the fairness constraints (equal opportunity) over true group memberships.
Revised: We train a linear classiﬁer with logistic regression, and report the overall error rate of the classiﬁer, as well as the maximum violation of the fairness constraints (equal opportunity) over true group memberships.

Original: We set (cid:15) = 0 .
Revised: We set (cid:15) = 0 .

Original: 05 and the noise level to 0 .
Revised: 05 and the noise level to 0 .

Original: Again, we compare with unconstrained optimization and heavily-constrained algorithm with a linear model as its multiplier model.
Revised: Again, we compare with unconstrained optimization and heavily-constrained algorithm with a linear model as its multiplier model.

Original: 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.
Revised: 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.

Original: Hence, the solution found by the GCIVR dominates heavily-constrained solution.
Revised: Hence, the solution found by the GCIVR dominates heavily-constrained solution.

Original: In terms of runtime speed, the ﬁrst row of Figure 1(c) clearly shows the advantage of GCIVR over heavily-constrained with much lower overhead to the unconstrained optimization.
Revised: In terms of runtime speed, the ﬁrst row of Figure 1(c) clearly shows the advantage of GCIVR over heavily-constrained with much lower overhead to the unconstrained optimization.

Original: Fairness Constraints on intersectional groups.
Revised: Fairness Constraints on intersectional groups.

Original: In this task we ought to learn a linear classiﬁer with logistic regression to predict the crime rate for a community on Communities and Crime dataset (Dua & Graff, 2017).
Revised: In this task we ought to learn a linear classiﬁer with logistic regression to predict the crime rate for a community on Communities and Crime dataset (Dua & Graff, 2017).

Original: This dataset contains 1994 instances of communities each withfeatures to predict the per capita crime rate for each community in the US.
Revised: This dataset contains 1994 instances of communities each withfeatures to predict the per capita crime rate for each community in the US.

Original: The labels are high and low crime rates to represent if a community is above the 70 th percentile of the data or not.
Revised: The labels are high and low crime rates to represent if a community is above the 70 th percentile of the data or not.

Original: To add fairness constraints, we ﬁrst determine different communities based on the percentages of the
Revised: To add fairness constraints, we ﬁrst determine different communities based on the percentages of the

Original: Black, Hispanic and Asian population in each community, as discussed in Narasimhan et al.
Revised: Black, Hispanic and Asian population in each community, as discussed in Narasimhan et al.

Original: (2020b);
Revised: (2020b);

Original: Cotter et al.
Revised: Cotter et al.

Original: Similar to Narasimhan et al.
Revised: Similar to Narasimhan et al.

Original: (2020b)
Revised: (2020b)

Original: , we generate 1000 thresholds of the form ( τ 1 , τ 2 , τ 3 ) ∈ [0 , 1] 3 to deﬁne each group.
Revised: , we generate 1000 thresholds of the form ( τ 1 , τ 2 , τ 3 ) ∈ [0 , 1] 3 to deﬁne each group.

Original: Consider the population of each race as ( p 1 , p 2 , p 3 ) , then a community belongs to group g τ 1 ,τ 2 ,τ 3 if p i ≥ τ
Revised: Consider the population of each race as ( p 1 , p 2 , p 3 ) , then a community belongs to group g τ 1 ,τ 2 ,τ 3 if p i ≥ τ

Original: i ∀ i ∈ [3] .
Revised: i ∀ i ∈ [3] .

Original: Then the fairness constraints enforce that the error rate of each group should be in the neighborhood of the overall error by margin of (cid:15) .
Revised: Then the fairness constraints enforce that the error rate of each group should be in the neighborhood of the overall error by margin of (cid:15) .

Original: That is err ( g τ
Revised: That is err ( g τ

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ) ≤ err
Revised: ) ≤ err

Original: [Equation]
Revised: [Equation]

Original: ALL )
Revised: ALL )

Original: [Equation]
Revised: [Equation]

Original: Similar to Narasimhan et al.
Revised: Similar to Narasimhan et al.

Original: (2020b), we set (cid:15) = 0 .
Revised: (2020b), we set (cid:15) = 0 .

Original: 01 and not consider groups with less than 1% of data.
Revised: 01 and not consider groups with less than 1% of data.

Original: 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.
Revised: 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.

Original: We compare the trade-off between the error rate and maximum violation of fairness constraints among groups.
Revised: We compare the trade-off between the error rate and maximum violation of fairness constraints among groups.

Original: The second row of Figure 1 portrays this comparison for the test dataset.
Revised: The second row of Figure 1 portrays this comparison for the test dataset.

Original: As it can be inferred, the proposed GCIVR algorithm achieves the same error rate as Heavily-constrained, but both higher than the unconstrained optimization.
Revised: As it can be inferred, the proposed GCIVR algorithm achieves the same error rate as Heavily-constrained, but both higher than the unconstrained optimization.

Original: However, comparing the maximum violation of constraints, it is clear that GCIVR outperforms both optimization methods.
Revised: However, comparing the maximum violation of constraints, it is clear that GCIVR outperforms both optimization methods.

Original: Also, the third ﬁgure shows the runtime of different algorithms.
Revised: Also, the third ﬁgure shows the runtime of different algorithms.

Original: It is clear that GCIVR adding a minimal overhead to the unconstrained optimization, while heavily-constrained increases the runtime more than 100 times.
Revised: It is clear that GCIVR adding a minimal overhead to the unconstrained optimization, while heavily-constrained increases the runtime more than 100 times.

Original: Per-query fairness constraints in ranking.
Revised: Per-query fairness constraints in ranking.

Original: In the third problem we evaluate our algorithm on the fairness in the ranking problem.
Revised: In the third problem we evaluate our algorithm on the fairness in the ranking problem.

Original: In this case, we intend to impose per-query constraints on the
Revised: In this case, we intend to impose per-query constraints on the

Original: [Figure]
Revised: [Figure]

Original: I n t e r s ec ti on a l F a i r n e ss
Revised: I n t e r s ec ti on a l F a i r n e ss

Original: [Figure]
Revised: [Figure]

Original: i ng
Revised: i ng

Original: [Figure]
Revised: [Figure]

Original: a i r n
Revised: a i r n

Original: [Figure]
Revised: [Figure]

Original: ss
Revised: ss

Original: [Figure]
Revised: [Figure]

Original: learning to rank problem as deﬁned in Narasimhan et al.
Revised: learning to rank problem as deﬁned in Narasimhan et al.

Original: (2020a;b).
Revised: (2020a;b).

Original: Similar to those setups, we divide document-query pairs into two groups based on the 40 th percentile of the their QualityScore features.
Revised: Similar to those setups, we divide document-query pairs into two groups based on the 40 th percentile of the their QualityScore features.

Original: 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.
Revised: 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.

Original: Consider the groups of g 0 and g 1 as mentioned before, we want the error rate for these groups to be close to each other.
Revised: Consider the groups of g 0 and g 1 as mentioned before, we want the error rate for these groups to be close to each other.

Original: In other terms, if the pairwise error rate is deﬁned as: err i,j ( q ) =
Revised: In other terms, if the pairwise error rate is deﬁned as: err i,j ( q ) =

Original: E [ 1 { f ( d, q ) < f ( d (cid:48) , q ) } | y > y (cid:48) , d ∈ g i , d (cid:48)
Revised: E [ 1 { f ( d, q ) < f ( d (cid:48) , q ) } | y > y (cid:48) , d ∈ g i , d (cid:48)

Original: ∈
Revised: ∈

Original: g j ] , where y and y (cid:48) are the respective binary labels.
Revised: g j ] , where y and y (cid:48) are the respective binary labels.

Original: Then, the fairness constraints can be satisﬁed as | err 0 , 1 ( q ) − err 1 , 0 ( q ) | ≤ (cid:15) ∀ q ∈ Q .
Revised: Then, the fairness constraints can be satisﬁed as | err 0 , 1 ( q ) − err 1 , 0 ( q ) | ≤ (cid:15) ∀ q ∈ Q .

Original: For this experiment, we use Microsoft Learning to Rank Dataset (MSLR-WEB10K) (Qin & Liu, 2013), which contains 10 K queries and 136 features.
Revised: For this experiment, we use Microsoft Learning to Rank Dataset (MSLR-WEB10K) (Qin & Liu, 2013), which contains 10 K queries and 136 features.

Original: 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.
Revised: 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.

Original: We compare against the unconstrained optimization and the heavily-constrained algorithm with an one-layer neural network with 64 nodes as its multiplier model.
Revised: We compare against the unconstrained optimization and the heavily-constrained algorithm with an one-layer neural network with 64 nodes as its multiplier model.

Original: We use 1000 queries in the training and 100 queries in the test datasets.
Revised: We use 1000 queries in the training and 100 queries in the test datasets.

Original: The third row in Figure 1 shows the error rate and maximum violation of groups constraints.
Revised: The third row in Figure 1 shows the error rate and maximum violation of groups constraints.

Original: As it is clear GCIVR outperforms heavily-constrained in both error rate and maximum violations of groups by a large margin.
Revised: As it is clear GCIVR outperforms heavily-constrained in both error rate and maximum violations of groups by a large margin.

Original: In terms of runtime, GCIVR is very close to the unconstrained optimization, and about 2 × faster than heavily-constrained.
Revised: In terms of runtime, GCIVR is very close to the unconstrained optimization, and about 2 × faster than heavily-constrained.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we showed that many DRO problems or heavily constrained optimization problems can be cast into a general framework based on ﬁnite-sum composite optimization.
Revised: In this paper, we showed that many DRO problems or heavily constrained optimization problems can be cast into a general framework based on ﬁnite-sum composite optimization.

Original: To solve this composite ﬁnite-sum optimization, we introduced centralized and distributed algorithms.
Revised: To solve this composite ﬁnite-sum optimization, we introduced centralized and distributed algorithms.

Original: We theoretically illustrated that our algorithm converges with an almost optimal number of constraint checks (for
Revised: We theoretically illustrated that our algorithm converges with an almost optimal number of constraint checks (for

Original: Wasserstein distance) or gradient calls (for χ 2 or KL divergence metrics).
Revised: Wasserstein distance) or gradient calls (for χ 2 or KL divergence metrics).

Original: Finally, we validated our theory with a number of experiments.
Revised: Finally, we validated our theory with a number of experiments.