Machine learning algorithms typically employ the approach of Empirical Risk Minimization (ERM), which minimizes the expected loss under the empirical distribution P 0 of the training dataset and assumes that test samples are generated from the same distribution. However, in practice, there usually exists a mismatch between the training and testing distributions due to various reasons, for example, in domain adaptation tasks domains differ from training to testing (Blitzer, McDonald, and Pereira 2006;Daume III and Marcu 2006); test samples were selected from minority groups which are underrepresented in the training dataset (Grother et al. 2011;Hovy and Søgaard 2015) and there might exist potential adversarial attacks (Goodfellow, Shlens, and Szegedy 2014;Madry et al. 2017). Such a mismatch may lead to a significant performance degradation.
This challenge spurred noteworthy efforts on developing a framework of Distributionally Robust Optimization (DRO) e.g., (Ben-Tal et al. 2013;Shapiro 2017;Rahimian and Mehrotra 2019). Rather than minimizing the expected loss under one fixed distribution, in DRO, one seeks to optimize the expected loss under the worst-case distribution in an uncertainty set of distributions.
Specifically, DRO aims to solve the following problem:
where U(P 0 ) is an uncertainty set of distributions centered at P 0 , P 0 is the empirical distribution of the training dataset, `is the loss function, and x is the optimization variable. For example, the uncertainty set can be defined as
where D is some distance-like metric, e.g., Kullback-Leibler (KL) divergence and 2 divergence, and ⇢ is the uncertainty level. In practice, for ease of implementation and analysis, a relaxed formulation of eq. ( 1), which is referred to as the penalized DRO, is usually solved (Levy et al. 2020;Jin et al. 2021;Qi et al. 2021;Sinha et al. 2017):
where > 0 is a fixed hyperparameter that needs to be chosen manually. In contrast to constrained DRO in eq. ( 1), a regularization term is added to the objective function to keep the distribution Q and the distribution P 0 close, and the hyperparameter is manually chosen beforehand to control the tradeoff with minimizing the loss. Compared with the penalized DRO setting, the constrained DRO problem in eq. ( 1) requires that the distribution Q to be strictly in the uncertainty set, and searches for the optimal solution under the worst-case distribution in the uncertainty set. Therefore, the obtained solution from the constrained DRO is minimax optimal for distributions in the uncertainty set, whereas it is hard to get such a guarantee for the penalized DRO relaxation. In this paper, we focus on the challenging constrained DRO problem in eq. ( 1).
Existing studies on constrained DRO are limited to convex loss functions or require some additional assumptions (Soma and Yoshida 2020;Hashimoto et al. 2018;Levy et al. 2020;Duchi and Namkoong 2018;Duchi, Glynn, and Namkoong 2021;Qi et al. 2022;Wang, Gao, and Xie 2021). Little understanding on the practical non-convex loss functions, e.g., neural network, is known. In this paper, we focus on the constrained DRO problem with non-convex loss.
DRO problems under different uncertainty sets are fundamentally different. As will be discussed later in related works, there is a rich literature on DRO with various uncertainty sets. In this paper, we focus on the general Cressie-Read family divergence defined uncertainty set (Duchi and Namkoong 2018;Jin et al. 2021), which includes, e.g., 2 divergence, as a special case (see Section 2 for more details). We also investigate the smoothed conditional value at risk (CVaR) DRO problem.
More importantly, we focus on the practical yet challenging large-scale scenario, where P 0 is the empirical distribution of N samples and N is very large. In classic stochastic optimization problems, e.g., ERM, it is easy to get an unbiased estimate of the gradient using only a few samples, and therefore the computational complexity at each iteration is independent of the training dataset size. However, in the DRO problems, due to taking the worst-case distributions in the objective, the problem becomes challenging. Many existing DRO algorithms incur a complexity that increases linearly (or even worse) in the training dataset size (Duchi and Namkoong 2018;Namkoong and Duchi 2016;Ghosh, Squillante, and Wollega 2018), which is not feasible for large-scale applications. In this paper, we will design a stochastic algorithm with computational complexity at each iteration being independent of the training dataset size (Qi et al. 2022;Wang, Gao, and Xie 2021;Levy et al. 2020).
The key challenges and main contributions in this paper are summarized as follows.
• For large-scale applications, the number of training samples is large, and therefore directly computing the full gradient is not practical. Nevertheless, as discussed above, it is challenging to obtain an unbiased estimate of the gradient for DRO problems using only a few samples. For '-divergence DRO problem, the distributions in the uncertainty set are continuous w.r.t. the training distribution. Thus the distributions in the uncertainty set can be parameterized by an N -dimensional vector (Namkoong and Duchi 2016). Then the DRO problem becomes a min-max problem and primal-dual algorithms (Rafique et al. 2022;Lin, Jin, and Jordan 2020;Xu et al. 2023) can be used directly. Subsampling methods in DRO were also studied in (Namkoong and Duchi 2016;Ghosh, Squillante, and Wollega 2018). However, all the above studies require a computational complexity linear or even worse in the training dataset size at each iteration and thus is prohibitive in large-scale applications. In (Levy et al. 2020), an efficient subsampling method was proposed, where the batch size is independent of the training dataset size. However, they only showed the sampling bias for 2 and CVaR DRO problems. In this paper, we generalize the analysis of the bias in (Levy et al. 2020) to the general Cressie-Read family. We further develop a Frank-Wolfe update on the dual variables in order to bound the gap between the objective and its optimal value given the optimization variable x and the biased estimate. • The second challenge is due to the non-convex loss function. Existing studies for the Cressie-Read divergence family (Duchi and Namkoong 2018;Levy et al. 2020) are limited to the case with convex loss function, and their approach does not generalize to the non-convex case. The key difficulty lies in that the subgradient of the objective function can not be obtained via subdifferential for nonconvex loss functions. Instead of explicitly calculating the worst-case distribution as in (Duchi and Namkoong 2018;Levy et al. 2020), we propose to design an algorithm for the dual problem which optimizes the objective under a known distribution. Thus the gradient of the objective can be efficiently obtained. • The third challenge is that the dual form of constrained DRO is neither smooth nor Lipschitz, making the convergence analysis difficult. Existing studies, e.g., (Wang, Gao, and Xie 2021), assume that the optimal dual variable is bounded away from zero, i.e., ⇤ > 0 for some 0 > 0, so that it is sufficient to consider 0 . However, this assumption may not necessarily be true as shown in (Wang, Gao, and Xie 2021;Hu and Hong 2013). In this paper, we generalize the idea in (Qi et al. 2022;Levy et al. 2020) to the general Cressie-Read divergence family. We design an approximation of the original problem, and show that it is smooth and Lipschitz. The approximation error can be made arbitrarily small so that the solution to the approximation is still a good solution to the original. We prove the strong duality of the approximated problem. We add a regularizer to the objective and at the same time we keep the hard constraint. In this way, we can guarantee that its dual variable has a positive lower bound. Moreover, our strong duality holds for any '-divergence DRO problem.
• We design a novel algorithm to solve the approximated problem and prove it converges to a stationary point of the constrained DRO problem. The general Proximal Gradient Descent algorithm (Ghadimi, Lan, and Zhang 2016) can be used to solve this approximated problem directly. However, it assumes the objective is non-convex in all parameters and does not provide a tight bound on the bias due to subsampling. We take advantage of the fact that the objective function is convex in and thus the bias due to subsampling can be bounded in a tighter way. Our proposed algorithm converges to a stationary point faster than existing methods.
Various Uncertainty Sets. '-divergence DRO problems (Ali and Silvey 1966;Csiszár 1967) were widely studied, for example, CVaR in (Rockafellar, Uryasev et al. 2000;Soma and Yoshida 2020;Curi et al. 2020;Tamar, Glassner, and Mannor 2015), 2 -divergence in (Hashimoto et al. 2018;Ghosh, Squillante, and Wollega 2018;Levy et al. 2020), KLdivergence in (Qi et al. 2021(Qi et al. , 2022;;Hu and Hong 2013) and Sinkhorn distance (Wang, Gao, and Xie 2021), a variant of Wasserstein distance based on entropic regularization. However, the above studies are for some specific divergence function and can not be extended directly to the general Cressie-Read divergence family.
Penalized DRO. The general '-divergence DRO problem was studied in (Jin et al. 2021) where the proposed algorithm
The Thirty-Eighth AAAI Conference on Artificial Intelligence (AAAI-24)
works for any divergence function with a smooth conjugate.
The authors also designed a smoothed version of the CVaR problem and showed their method converges to a stationary point. However, their method is for the penalized formulation and does not generalize to the constrained DRO. In this paper, we focus on the challenging constrained DRO, the solution of which is minimax optimal over the uncertainty set. Our proposed algorithm can also be applied to solve the smoothed CVaR problem in the constrained setting.
Constrained DRO With Convex Loss. The general '-divergence constrained DRO problem was studied in (Namkoong and Duchi 2016). Instead of optimizing from the dual form, the authors treat the worst-case distribution as a N -dimentional vector and implement a stochastic primaldual method to solve the min-max problem. However, the computational complexity at each iteration is linear in the number of the training samples and can not be used in largescale applications. The same problem was further studied in (Duchi, Glynn, and Namkoong 2021). The authors pointed out that minimizing constrained DRO with '-divergence is equivalent to adding variance regularization for the Empirical Risk Minimization (ERM) objective. The general Cressie-Read divergence family DRO problem was studied in (Duchi and Namkoong 2018), where the basic idea is to calculate the worst-case distribution for the constrained DRO first and then use the subdifferential to get the subgradient. Furthermore, the 2 and CVaR DRO problems were studied in (Levy et al. 2020). Compared with the method in (Duchi and Namkoong 2018), they calculate the worst-case distribution for the penalized DRO and then optimize both the Lagrange multiplier and the loss function. This approach converges to the optimal solution with a reduced complexity. Their method can be extended to the general Cressie-Read divergence family. However, all the above papers are limited to the case with convex loss function. To the best of our knowledge, our work is the first paper on large-scale non-convex constrained DRO with the general Cressie-Read divergence family. We note that the KL DRO was studied in (Qi et al. 2022), which however needs an exponential computational complexity. We achieve a polynomial computational complexity for the Cressie-Read divergence family.
2 Preliminaries and Problem Model
Let s be a sample in S and P 0 be the distribution on the points {s i } N i=1 , where N is the size of the support. Denote by n := {p 2 R n | P n i=1 p i = 1, p i 0} the ndimensional probability simplex. Denote by x 2 R d the optimization variable. We denote by 1 X (x) the indicator function, where 1 X (x) = 0 if x 2 X, otherwise 1 X (x) = 1. Let `: R d ⇥ S ! R be a non-convex loss function. Let k • k be the Euclidean norm and (t) + := max{t, 0} be the positive part of t 2 R. Denote r x by the gradient to x. For a function f :
In this paper, we take the following standard assumptions that are commonly used in the DRO literature (Duchi and Namkoong 2018;Levy et al. 2020;Qi et al. 2021Qi et al. , 2022;;Wang, Gao, and Xie 2021; Soma and Yoshida 2020):
• The non-convex loss function is bounded:
In empirical risk minimization (ERM), the goal is to solve
where the objective function is the expectation of the loss function with respect to the training distribution P 0 . To solve the distributional mismatch between training data and testing data, the formulation of Distributionally Robust Optimization (DRO) (Goodfellow, Shlens, and Szegedy 2014;Madry et al. 2017; Rahimian and Mehrotra 2019) was developed, where the goal is to minimize the expected loss with respect to the worst distribution in an uncertainty set U(P 0 ): ⌘ dP 0 , where ' is a non-negative convex function such that '(1) = 0 and '(t) = +1 for ant t < 0. Then let the uncertainty set U(P 0 ) := {Q : D ' (QkP 0 ) ⇢} where ⇢ > 0 is the radius of the uncertainty set.
In this paper, we study the general Cressie-Read family of '-divergence (Cressie and Read 1984;Van Erven and Harremos 2014), where
This family includes as special cases 2 -divergence (k = 2) and KL divergence (k ! 1). When k > 2, the conjugate function of ' k (t) (which will be introduced later) is not smooth, thus the problem becomes hard to solve even in the penalized formulation (Jin et al. 2021). In this paper, we focus on k 2
Solving ( 6) directly is challenging due to the sup over Q. In (Namkoong and Duchi 2016), a finite-dimensional vector q was used to parameterize the distributions in the uncertainty set since Q ⌧ P 0 for '-divergence. Then the DRO problem becomes a convex concave min-max problem. This method can be extended to the case with non-convex loss function by applying the algorithms for non-convex concave min-max problems (Rafique et al. 2022;Lin, Jin, and Jordan 2020;Xu et al. 2023). However, the dimension of distribution in the uncertainty set is equal to the number of training samples. Thus, the computational complexity at each iteration is linear in the sample size and is prohibitive in largescale applications.
To obtain a complexity independent of the sample size, one alternative is to use its dual. By duality, we can show that the DRO objective (6) can be equivalently written as (Levy et al. 2020;Shapiro 2017)
} is the conjugate function of ' k (t 0 ). In this way, the optimization problem under an unknown distribution is rewritten into one under a known distribution. The subsampling method can then be used, which leads to a complexity independent of the sample size (which will be introduced later). For the Cressie-Read family in (5), the corresponding conjugate function family is
. Therefore, the objective can be written as
Let ⌘ = ⌘ k 1 and the corresponding objective is
Thus the goal is to solve
where F (x; ; ⌘) is defined as
. Therefore, we reformulate the DRO problem as one to minimize an objective function under a known distribution, where subsampling method could be used to reduce the complexity.
In this section, we analyze the constrained DRO problems under Cressie-Read family divergence uncertainty sets with general smooth non-convex loss function. We first discuss the challenges appearing in constrained formulations, then we present how to construct the corresponding approximated problem in order to overcome these challenges.
For 2 [0, +1), ⌘ 2 R, the objective function F (x; ; ⌘) is neither smooth nor Lipschitz. Thus it is difficult to implement gradient-based algorithms. In the following, we will construct an approximation of the original problem so that the objective function F (x; ; ⌘) becomes smooth and Lipschitz by constraining both and ⌘ in some bounded intervals.
Denote by ! = (k(k 1)⇢+1) 1 k⇤ . Since the loss function is bounded such that 0 ` B, we can show that there exists an upper bound ¯ = (k 1)! 1
which only depends on k, ⇢ and B such that the optimal value ⇤ ¯ . In this paper, we do not assume that ⇤ 0 > 0 as in (Wang, Gao, and Xie 2021). Instead, we consider an approximation with 2 [ 0 , ¯ ], and show that the difference between the orignial and the approximation can be bounded. We can show corresponding optimal
k⇤ 1 . The proof can be found in Appendix A. The challenge lies in that the value of ⌘ can be negative. Thus given this ⌘, the optimal value of can be quite large then it is hard to upper bound . In our proof, we change the objective to the function that only depends on ⌘ and find the lower bound on ⌘. Based on this lower bound, we solve this challenge and further get the bound on .
We show that the difference between the original and the approximation can be bounded in the following lemma.
The proof can be found in Appendix B. Note in the proof of this lemma, we provide the strong duality of
where both the hard constraint and regulator are kept. This is different from the approach in Section 3.2 of (Shapiro 2017). Note this strong duality holds for any '-divergence DRO problem.
Lemma 1 demonstrates that the non-smooth objective function can be approximated by a smooth objective function. A smaller 0 makes the gap smaller but the function "less smooth".
The advantage of our approximated problem is that the function is smooth in all x, , and ⌘. Moreover, We find that the objective function is convex in and ⌘ though the loss function is non-convex in x. Lemma 2. Define z = ( , ⌘) 2 M, where M = {( , ⌘) :
The proof can be found in Appendix C. Note the firstorder gradient of the objective function is non-differential at some point when k ⇤ = 2. Therefore, we discuss in cases: k ⇤ > 2 and k ⇤ = 2. In the first case, we can get the Hessian matrix of the objective. In the second case, we show the smoothness and convexity.
Existing constrained stochastic algorithm for general nonconvex functions (Ghadimi, Lan, and Zhang 2016) can be used to solve the approximated problem directly. However, their method optimizes y = (x; ; ⌘) as a whole. It can be seen that the objective function is non-convex in y and the computation complexity to get the ✏-stationary point is O(✏ 3k⇤ 5 ).
In the previous section, we show that
). If we optimize all the parameters together, we need to implement non-convex algorithms to optimize a smooth function with a large smooth constant, which is not computationally efficient. However, if we optimize x and z separately, though L z > L x which requires more resources to optimize z, the convexity in z makes it faster to converge to the optimal value of z.
This motivates us to consider a stronger convergence criterion. Instead of finding the ✏-stationary point for F (y), we can find (x, , ⌘) such that
We then provide our Stochastic gradient and Frank-Wolfe DRO algorithm (SFK-DRO), which optimizes x and z separately (see Algorithm 1). Define D = max z1,z22M kz 1
The convergence rate is then provided in the following theorem.
Theorem 1. With a mini-batch size
✏ 2 ) iterations are needed to guarantee a stationary point (x t 0 +1 ; z t 0 ) in expectation:
The detailed proof can be found in Appendix D and a proof sketch will be provided later. Before that, we introduce a lemma for our subsampling method. Via this lemma, we can show the complexity is independent of the sample size and thus is suitable for our large-scale setting. When we optimize z, an estimator
Though the estimator is unbiased, in our Frank-Wolfe update process (Jaggi 2013;Frank, Wolfe et al. 1956; Lacoste-Julien 2016) we need to estimate min F (x; z) via E min f z (x; z). Obviously, the expectation of minimum is not equal to the minimum of expectation, thus it is a biased estimator. In the following lemma, we show that this gap can be bounded by a decreasing function of the sample batch n z .
Lemma 3. For any bounded loss function
The detailed proof can be found in Appendix E. Note that (Levy et al. 2020) only shows this lemma when k ⇤ = 2, and we extend the results to k ⇤ > 2. This lemma shows that the gap is in the order of O(n z 1 k⇤ ) and is independent of the total number of samples.
The Thirty-Eighth AAAI Conference on Artificial Intelligence (AAAI-24)
We use a stochastic gradient descent method (Moulines and Bach 2011;Gower et al. 2019;Robbins and Monro 1951) to update x. Since the objective function is L x -smooth in x, if ↵ 1 2Lx we have that:
where
and we can show that
Since z 2 M, instead of the stochastic gradient descent method, we employ the Frank-Wolfe method (Frank, Wolfe et al. 1956) to update z. Define e t = arg min e2M he, r z f z (x t+1 ; z t )i and g t = he t z t , r z f z (x t+1 ; z t )i. In addition, we have Jaggi 2013). We can show that gt C ⇠ O( 0 ) thus for small 0 we have gt C 1. Then due to the fact that the objective is L z -smooth in z (( 9) of (Lacoste-Julien 2016)), we have that
By recursively adding ( 9) and ( 11), we have that
Since L z ⇠ O( k⇤ 1 0
) and L z ⇠ O( k⇤ 1 0
), for small
C✏ 2 , and denote = F (x 1 ; z 1 ) min x,z2M F (x; z), for some t 2 [1, T ] we have
We choose (x t+1 , z t ) as our output and we need to bound
By Lemma 3, we pick
By Lemma 1, when 0 = ✏ 8⇢ , we have
Thus we have
which completes the proof.
Our algorithm can also solve other DRO problems efficiently, for example, the Smoothed CVaR proposed in (Jin et al. 2021). The CVaR DRO is an important '-divergence DRO problem, where '(t) = 1 [0, 1 µ ) if 0 t < 1 µ , and 0 < µ < 1 is some constant. The dual expression of CVaR can be written as
The dual of CVaR is non-differentiable, which is undesirable from an optimization viewpoint. To solve this problem, (Jin et al. 2021) proposed a new divergence function, which can be seen as a smoothed version of the CVaR. Their experiment results show the optimization of smoothed CVaR is much easier. However, (Jin et al. 2021)'s method only works for the penalized formulation of DRO. We will show that our method can solve the constrained smoothed CVaR.
Here, the divergence function is
The corresponding conjugate function is
The objective function is then written as
We can show that there exist upper bounds for the optimal values ⇤ and ⌘ ⇤ . There exists a ¯ > 0 only depends on µ, B and ⇢ such that ⇤ 2 [0, ¯ ] and ⌘ ⇤ 2 [0, B]. The proof can be found in Appendix F. This objective function is non-smooth when ! 0. Therefore, we take a similar approach as the one in Section 3.1 to approximate the original problem with 2 [ 0 , ¯ ]. We bound the difference in the following lemma.
The proof is similar to Lemma 1 thus is omitted here. In addition, we can show that F s (x; z) is L 0 z -smooth and convex in z, where L
). Similar to eq. ( 42) and Remark 1 in (Levy et al. 2020), we can prove that min z2M [F s (x t+1 ; z)] E [min z2M f s (x t+1 ; z)] ⇠ O(n 0.5 s ). We then use Algorithm 1 directly and the complexity to get the ✏-stationary point is O(✏ 7 ). The detailed proof can be found in Appendix F.
In this section, we verify our theoretical results in solving an imbalanced classification problem. In the experiment, we consider a non-convex loss function and k is set to be 2 for the Cressie-Read family. We will show that 1) to optimize the same dual objective function, our proposed algorithm converges faster than the general Proximal Gradient Descent(PGD) (Ghadimi, Lan, and Zhang 2016);
2) The performance proposed algorithm for the constrained DRO problem outperforms or is close to the performance of the penalized DRO with respect to the worst classes. Both of them outperform the baseline. Tasks. We conduct experiments on the imbalanced CIFAR-10 dataset, following the experimental setting in (Jin et al. 2021;Chou et al. 2020). The original CIFAR-10 test dataset consists of 10 classes, where each of the classes has 5000 images. We randomly select training samples from the original set for each class with the following sampling ratio: {0.804, 0.543, 0.997, 0.593, 0.390, 0.285, 0.959, 0.806, 0.967, 0.660}. We keep the test dataset unchanged. Models. We learn the standard Alexnet model in (Krizhevsky, Sutskever, and Hinton 2012) with the standard cross-entropy (CE) loss. For the comparison of convergence rate, we optimize the same dual objective with the PGD algorithm in (Ghadimi, Lan, and Zhang 2016). To compare robustness, we optimize the ERM via vanilla SGD. In addition, we propose an algorithm PAN-DRO, which fixes and only optimizes ⌘ and the neural network. Thus it gets the solution for the penalized DRO problem. Training Details. We set 1 = 1, ⌘ 1 = 0, 0 = 0.1, ⌘ = 10, and the upper bounds ¯ = 10, B = 10. To achieve a faster optimization rate, we set the learning rate ↵ = 0.01 before the first 40 epochs and ↵ = 0.001 after. The minibatch size is chosen to be 128. All of the results are moving averaged by a window with size 5. The simulations are repeated by 4 times. Results. In Figure 1, we plot the value of the CE loss using different algorithms through the training process. It can be seen that to optimize the same dual objective function with the same learning rate, the PGD algorithm converges slower than our proposed DRO algorithms, which matches our theoretical results. Moreover, compared with ERM, the DRO algorithms have higher training losses but lower test losses, which demonstrates they are robust.
We also provide the test accuracy of trained models in Table 1. It can be shown that for class 3, 4, 5, the accuracies are the lowest due to the limited samples. For these classes, the performance of our SFK-DRO algorithm for the constrained DRO is better or close to the performance of PAN-DRO for the penalized DRO. Both DRO algorithms outperform the
In this paper, we developed the first stochastic algorithm for large-scale non-convex stochastic constrained DRO problems in the literature with theoretical convergence and complexity guarantee. We developed a smooth and Lipschitz approximation with bounded approximation error to the original problem. Compared with existing algorithms, the proposed algorithm has an improved convergence rate. The computational complexity at each iteration is independent of the size of the training dataset, and thus our algorithm is applicable to large scale applications. Our results hold for a general family of Cressie-Read divergences.
The Thirty-Eighth AAAI Conference on Artificial Intelligence