An Online Method for A Class of Distributionally Robust Optimization with Non-Convex Objectives

Abstract:
In this paper, we propose a practical online method for solving a class of distributionally robust optimization (DRO) with non-convex objectives, which has important applications in machine learning for improving the robustness of neural networks. In the literature, most methods for solving DRO are based on stochastic primal-dual methods. However, primal-dual methods for DRO suffer from several drawbacks: (1) manipulating a high-dimensional dual variable corresponding to the size of data is time expensive; (2) they are not friendly to online learning where data is coming sequentially. To address these issues, we consider a class of DRO with an KL divergence regularization on the dual variables, transform the minmax problem into a compositional minimization problem, and propose practical duality-free online stochastic methods without requiring a large mini-batch size. We establish the state-of-the-art complexities of the proposed methods with and without a Polyak-Łojasiewicz (PL) condition of the objective. Empirical studies on large-scale deep learning tasks (i) demonstrate that our method can speed up the training by more than 2 times than baseline methods and save days of training time on a large-scale dataset with ∼ 265K images, and (ii) verify the supreme performance of DRO over Empirical Risk Minimization (ERM) on imbalanced datasets. Of independent interest, the proposed method can be also used for solving a family of stochastic compositional problems with state-of-the-art complexities.

heading: 1 Introduction
text: Distributionally robust optimization (DRO) has received tremendous attention in machine learning due to its capability to handle noisy data, adversarial data and imbalanced classification data [42, 33, 4]. Given a set of observed data {z1, . . . , zn}, where zi = (xi, yi), a DRO formulation can be written as:
min w∈Rd max p∈∆n Fp(w) = n∑ i=1 pi`(w; zi)− h(p,1/n) + r(w), (1)
where w denotes the model parameter, ∆n = {p ∈ Rn : ∑ i pi = 1, pi ≥ 0} denotes a ndimensional simplex, `(w; z) denotes a loss function on data z, h(p,1/n) is a divergence measure between p and uniform probabilities 1/n, and r(w) is convex regularizer of w. When `(w; z) is a convex function (e.g., for learning a linear model), many stochastic primal-dual methods can be employed for solving the above min-max problem [35, 20, 49, 48, 32]. When `(w; z) is a non-convex function (e.g., for learning a deep neural network), some recent studies also proposed stochastic primal-dual methods [41, 48]. ∗The first two authors make equal contributions. Correspondence to tianbao-yang@uiowa.edu.
35th Conference on Neural Information Processing Systems (NeurIPS 2021), virtual.
However, stochastic primal-dual methods for solving DRO problems with a non-convex `(w; z) loss function (e.g., the predictive model is a deep neural network) suffer from several drawbacks. First, primal-dual methods need to maintain and update a high-dimensional dual variable p ∈ Rn for large-scale data, whose memory cost is as high as O(n) per-iteration. Second, existing primal-dual methods usually need to sample data according to probabilities p in order to update w, which brings additional costs than random sampling. Although random sampling can be used for computing the stochastic gradient in terms of w, the resulting stochastic gradient could have n-times larger variance than using non-uniform sampling according to p (please refer to the supplement for a simple illustration). Third, due to the constraint on p ∈ ∆n, the min-max formulation (1) is not friendly to online learning in which the data is received sequentially and n is rarely known in prior.
Can we design an efficient online algorithm to address the DRO formulation (1) without dealingwith p ∈ Rn for a non-convex objective that is applicable to deep learning?
To address this question, we restrict our attention to a family of DRO problems, in which the KL divergence h(p,1/n) = λ ∑ i pi log(npi) is used for regularizing the dual variables p, where λ > 0 is a regularization parameter. We note that this consideration does not impose strong restriction to the modeling capability. It has been shown that for a family of divergence functions h(p,1/n), different DRO formulations are statistically equivalent to a certain degree [12]. The proposed method is based on an equivalent minimization formulation for h(p,1/n) = λ ∑ i pi log(npi). In particular, by maximizing over p exactly, (1) is equivalent to
min w∈Rd
{ Fdro(w) = λ log ( 1/n
n∑ i=1 exp(`(w; zi)/λ))
) + r(w) } . (2)
In an online learning setting, we can consider a more general formulation:
min w∈Rd
{ Fdro(w) = λ log (Ez exp (`(w; z)/λ)) + r(w) } . (3)
The above problem is an instance of stochastic compositional problems of the following form:
min w∈Rd F (w) := f(Ez[gz(w)]) + r(w), (4)
by setting f(s) = λ log(s), s ≥ 1 and gz(w) = exp(`(w; z)/λ). Stochastic algorithms have been developed for solving the above compositional problems. [44] proposed the first stochastic algorithms for solving (4), which are easy to implement. However, their sample complexities are sub-optimal for solving (4). Recently, a series of works have tried to improve the convergence rate by using advanced variance reduction techniques (e.g., SVRG [19], SPIDER [13], SARAH [36]). However, most of them require using a mega mini-batch size in the order of O(n) or O(1/ )2 at every iteration or many iterations for updating w, which hinders their applications on large-scale problems. In addition, these algorithms usually use a constant step size, which may harm the generalization performance.
This paper aims to develop more practical stochastic algorithms for solving (3) without suffering from the above issues in order to enable practitioners to explore the capability of DRO for deep learning with irregular data (e.g., imbalanced data, noisy data). To this end, we proposed an online stochastic method (COVER) and its restarted variant (RECOVER). We establish a state-of-the-art complexity of COVER for finding an -stationary solution and a state-of-the-art complexity of RECOVER under a Polyak-Łojasiewicz (PL) condition of the problem. PL condition has been widely explored for developing practical optimization algorithms for deep learning [52]. Compared with other stochastic algorithms, the practical advantages of RECOVER are:
1. RECOVER is an online duality-free algorithm for addressing large-scale KL regularized DRO problem that is independent of the high dimensional dual variable p ∈ Rn, which makes it suitable for deep learning applications.
2. RECOVER also enjoys the benefits of stagewise training similar to existing stochastic methods for deep learning [52], i.e., the step size is decreased geometrically in a stagewise manner.
2 is either the objective gap accuracy F (w) − minF (w) ≤ or the gradient norm square bound ‖∇F (w)‖2 ≤
In addition, this paper also makes several important theoretical contributions for stochastic non-convex optimization, including
1. We establish a nearly optimal complexity for finding -stationary point, i.e., ‖∇F (w)‖2 ≤ , for a class of two-level compositional problems in the order of Õ(1/ 3/2) without a large mini-batch size, which is better than existing results [44, 45, 14, 5].
2. We etablish an optimal complexity for finding -optimal solution under an µ−PL condition for a class of two-level compositional problems in the order of O(1/(µ )) without a large mini-batch size, which is better than existing results [53].
A theoretical comparison between our results and existing results is shown in Table 1. Empirical studies vividly demonstrate the effectiveness of RECOVER for deep learning on imbalanced data.

heading: 2 Related Work
text: DRO has been extensively studied in machine learning [31, 11, 40], statistics, and operations research [42]. In [33], the authors proved that minimizing the DRO formulation with a quadratic regularization in a constraint form is equivalent to minimizing the sum of the empirical loss and a variance regularization defined on itself. Variance regularization can enjoy better generalization error compared with the empirical loss minimization [33], and was also observed to be effective for imbalanced data [33, 56]. Recently, [12] also establishes this equivalence for a broader family of regularization function h(p,1/n) including the KL divergence.
Several recent studies have developed stochastic primal-dual methods for solving DRO with a nonconvex loss function `(w; z) assuming it is smooth or weakly convex [41, 26, 30, 48]. [41] proposed the first primal-dual methods for solving weakly convex concave problems. For online problems, their algorithms for finding an -stationary solution whose gradient norm square is less than have a sample complexity of O(1/ 2) or O(1/ 3) with or without leveraging the strong concavity of h(p,1/n) for finding an -stationary point. Recently Liu et al. [28] proposed to leverage the PL condition of the objective function to improve the convergence for a non-convex min-max formulation of AUC maximization. Then, a PES-SGDA algorithm is proposed to solve a more general class of non-convex min-max problems by leveraging the PL condition [15]. Both [28] and [15] have used geometrically decreasing step sizes in a stagewise manner. However, their algorithms’ complexity is in the order of O(1/µ2 ), which is worse than O(1/µ ) achieved in this paper. Similarly, [50] also leveraged PL conditions to solve non-convex min-max problems and has a sample complexity of O(1/µ2 ). Nevertheless, the step size of their algorithm is decreased polynomially in the order of O(1/t), which usually yields poor performance for deep learning.
All the methods reviewed above require maintaining and updating both the primal variable w and a high dimensional dual variable p ∈ Rn. Recently, Levy et al. [23] considered different formulations of DRO, which includes our considered KL-regularized DRO formulation as a special case. Their assumed that the loss function is convex and proposed a stochastic method with a sample complexity O(1/ 3) for sovling the KL-regularized DRO formulation. In contrast, we provide a better sample
complexity in the order of O(1/ ) under a PL condition without convexity assumption. Additionally, their method requires a large batch size in the order of O(1/ ), while our method only requires a constant batch size which is more practical. We also notice that a recent work [24] and its extended version [25] have considered a formulation similar to (2) and proposed a stochastic algorithm. However, their algorithm has a slower convergence rate with an O(1/ 2) complexity for finding an - stationary point and anO(1/(µ2 )) complexity for finding an -optimal solution under a PL condition. Our work is a concurrent work appearing online earlier than [24]. To the best of our knowledge, this is the first work trying to solve the non-convex DRO problem with a duality-free stochastic method by formulating the min-max formulation into an equivalent stochastic compositional problem.
There are extensive studies for solving stochastic compositional problems. [44] considered a more general family of stochastic compositional problems and proposed two algorithms. When the objective function is non-convex, their algorithm’s complexity isO(1/ 7/2) for finding an -stationary solution. This complexity was improved in their later works [14], reducing to O(1/ 2). When the objective is smooth, several papers proposed to use variance reduction techniques (e.g., SPIDER, SARAH) to improve the complexity for finding a stationary point [53, 18, 55, 27]. The best sample complexity achieved for online problems is O(1/ 3/2) [53, 55]. [53] also considered the PL condition for developing a faster algorithm called restarted CIVR, whose sample complexity is O(1/µ ). However, these variance reduction-based methods require using a very large mini-batch size at many iterations, which has detrimental influence on training deep neural networks [43]. To address this issue, [9] proposed a new technique called STORM that integrates momentum and the recursive variance reduction technique for solving stochastic smooth non-convex optimization. Their algorithm does not require a large mini-batch size at every iterations and enjoys a sample complexity of O(log2/3(1/ )/ 3/2) for finding an -stationary point. However, their algorithm uses a polynomially decreasing step size, which is not practical for deep learning, and is not directly applicable to stochastic composite problems.

heading: 3 Preliminaries
text: In this section, we provide some definitions and assumptions for next section. For more generality, we consider the stochastic compositional problem (4):
min w∈Rd F (w) := f(Ez[gz(w)]) + r(w) (5)
where gz(w) : Rd → Rp. Define g(w) = Ez[gz(w)]. Let ‖ · ‖ denote the Euclidean norm of a vector or the Frobenius norm of a matrix. We make the following standard assumptions regarding the problem (5). Assumption 1. Let Cf , Lf , Cg and Lg be positive constants. Assume that
(a) f : Rp → R is a Cf -Lipschitz function and its gradient∇f is Lf -Lipschitz. (b) gz : Rd → Rp satisfies E‖gz(w1)− gz(w2)‖2] ≤ C2g‖w1 −w2‖2 for any w1,w2 and its
Jacobian ∇gz satisfies E[‖∇gz(w1)−∇gz(w2)‖2] ≤ L2g‖w1 −w2‖2.
(c) r : Rd → R ∪ {∞} is a convex and lower-semicontinuous function. (d) F∗ = infw F (w) ≥ −∞ and F (w1)− F∗ ≤ ∆F for the initial solution w1.
Remark: When f(s) = s is a linear function, the assumption E‖gz(w1)− gz(w2)‖2] ≤ C2g‖w1 − w2‖2 is not needed. To upper bound continuity and smoothness of function F , we denote L = 2 max{LgCgLf , CfCgLf , C2f , LgCf , C2gLf , Cf , CgLf , C2f , C2g , C2gL2f} for simple derivation in the appendix. Assumption 2. Let σg and σg′ be positive constants and σ2 = σ2g + σ2g′ . Assume that
Ez[‖gz(w)− g(w)‖2] ≤ σ2g , Ez[‖∇gz(w)−∇g(w)‖2] ≤ σ2g′ . Remark: We remark how the minimization formulation of DRO problem (3) can satisfy Assumption 1, in particular Assumption 1(a) and (b). In order to satisfy Assumption 1(b), we can define a bounded loss function `(w, z) ∈ [0, `max] and then use a shifted loss `(w; z) − `max in (3). Then gz(w) = exp((`(w; z) − `max)/λ) is Lipchitz continuous and smooth if `(w; z) is Lipchitz and smooth. f(s) = λ log(s) is Lipschitz continuous and smooth since s ≥ exp(−`max/λ).
For more generality, we allow for a non-smooth regularizer r(·) in this section. To handle nonsmoothness of r, we can use the proximal operator of r: proxηr(w̄) = arg minw 12‖w−w̄‖
2+ηr(w). When r = 0, the above operator reduces to the standard Euclidean projection. Correspondingly, we define the proximal gradient measure for the compositional problem (5):
Gη(w) = 1
η (w − proxηr(w − η∇g(w)>∇f(g(w)))).
When r = 0, the proximal gradient reduces to the standard gradient measure, i.e., Gη(w) = ∇F (w). To facilitate our discussion, we define sample complexity below. Definition 1. The sample complexity is defined as the number of samples z in order to achieve E[‖Gη(w)‖2] ≤ for a certain η > 0 or E[F (w)− F∗] ≤ .

heading: 4 Basic Algorithm: COVER
text: We present our Algorithm 1,which serves as the foundation for proving the the convergence of the objective gap under a PL condition in next section. The convergence results in this section might be of independent interest to those who are interested in convergence analysis without a PL condition. The motivation is to develop a stochastic algorithm with fast convergence in terms of gradient norm. We refer to the algorithm as Compositional Optimal VariancE Reduction (COVER). It will be clear shortly why it is called optimal variance reduction. Note that in order to compute a stochastic estimator of the gradient f(g(w)) given by ∇g(w)>∇f(g(w)), we maintain and update two estimators denoted by {u}Tt=1 and {v}Tt=1 sequence, respectively. The {ut}Tt=1 sequence maintains an estimation of {g(wt)}Tt=1 and the {vt}Tt=1 sequence maintains an estimation of {∇g(wt)}Tt=1. The strategy of maintaining and updating two individual sequences was first proposed in [44] and has been widely used for solving compositional problems [55, 53]. However, the key difference from previous algorithms lies in the method for updating the two sequences. COVER is inspired by the STORM technique [9]. To understand the update, let us consider update that applied to the DRO problem (3) by let f(·) = λ log(·), gz(·) = exp( `(·;z)λ ) and ignoring r for the moment. Plugging the gradient of f(·) and gz(·), we have
wt+1 = wt − ηt 1
ut ṽt,
ṽt+1 = exp( `(wt+1; zt+1)
λ )∇`(wt+1; zt+1) + (1− at+1)(ṽt − exp(
`(wt; zt+1)
λ )∇`(wt; zt+1)),
where ut becomes a scalar, which is an online variance-reduced estimator of Ez[exp(`(wt; z)/λ)], and ṽt is a scaled version of vt, which is an online variance-reduced estimator of Ez[exp(`(wt; z)/λ)∇`(wt; z)]. Finally, we notice that a similar method for updating the ut sequence for estimating g(wt) has been adopted in a recent work [6]. However, different from the present work they just use an unbiased stochastic gradient to estimate∇g(wt), which yields a worse convergence rate.

heading: 4.1 Convergence of Proximal Gradient
text: In this section, we present the convergence result of COVER.
Theorem 1. Assume the Assumption 1 and 2, for any C > 0, k = Cσ 2/3
L , c = 128L+ σ 2/(7Lk3),
w = max((16Lk3), 2σ2, ( ck4L ) 3), and ηt = k/(w + σ2t)1/3. The output of COVER satisfies
E[‖Gηt∗ (wt∗)‖ 2] ≤ Õ ( ∆F T 2/3 + σ2 T 2/3 ) . (6)
where t∗ is sampled from {1, . . . , T}.
Remark: Theorem 1 implies that with a polynomially decreasing step size, COVER is able to find an -stationary point, i.e., E[‖Gηt∗ (wt∗)‖2] ≤ for a regularized objective and E[‖∇F (w)‖2] ≤ for a non-regularized objective, with a near-optimal sample complexity Õ( 1
3/2 ). Note that the complexity
Õ(1/ 3/2) is optimal up to a logarithmic factor for making the (proximal) gradient’s norm smaller than in expectation for solving non-convex smooth optimization problems [2].
Algorithm 1: COVER (w1,u1,v1, {ηt}, T,PL = False) 1: Let at = cη2t 2: if not PL then 3: Draw a samples z and construct the estimates: u1 = gz(w1), v1 = ∇gz(w1) 4: end if 5: for t = 1, . . . , T − 1 do 6: wt+1 ← proxηtr (wt − ηtv>t ∇f(ut)) 7: Draw a samples zt+1, and update
ut+1 = gzt+1(wt+1) + (1− at+1)(ut − gzt+1(wt))
vt+1 = ∇gzt+1(wt+1) + (1− at+1)(vt −∇gzt+1(wt)) 8: end for 9: Return: (wτ ,uτ ,vτ ) for randomly selected τ ∈ {1, . . . , T}.

heading: 5 A Practical Variant (RECOVER) under a PL condition
text: The issue of COVER is that the polynomially decreasing step size is not practical for deep learning applications and obstacles its generalization performance [52]. A stagewise step size is widely and commonly used [17, 22, 52] for deep learning optimization. To this end, we develop a multi-stage REstarted version of COVER, called RECOVER, which uses a geometrically decreasing step size in a stagewise manner. In oder to analyze RECOVER, we assume the following PL condition of the objective with a smooth regularization r term [52]. Assumption 3. F (w) satisfies the µ-PL condition if there exists µ > 0 such that
2µ(F (w)− min w∈Rd F (w)) ≤ ‖∇F (w)‖2. (7)
In the following, we simply consider the objective F (w) = f(Ez[gz(w)]), where r(·) is absorbed into f(Ez[gz(w)]). As a result, Gη(w) = ∇F (w). Although the PL condition has been considered in various papers for developing stagewise algorithms and improving the convergence rate of non-convex optimization [53, 52, 28, 15]. In order to establish the improved rate, we have innovations in twofold (i) at the algorithmic level, we utilize the variance reduction techniques at the inner and outer level without using mega large mini-batch size at any iterations; (ii) at the analysis level, we innovatively prove that the estimation error of the two sequences, u and v, are decreasing geometrically after a stage (Lemma 3). These innovations at two levels yield the optimal convergence rate in the order of O(1/(µ )).

heading: 5.1 Theoretical Verification of PL Assumption for KL-regularized DRO
text: Before presenting the proposed algorithm and its convergence, we discuss how the Fdro can satisfy Assumption 3. First, we note that a PL condition of the weighted loss implies that of the primal objective. Lemma 1. Let Fp(w) = ∑n i=1 pi`(w; zi). If for any p ∈ ∆n, Fp(w) satisfies a µ-PL condition,
then Fdro(w) = λ log( 1n ∑ i exp(`(w; zi)/λ)) satisfies the µ-PL condition.
Remark: The assumption that the weighted loss satisfies a PL condition can be proven for a simple square loss `(w; zi) = (w>xi − yi)2, where zi = (xi, yi) consists of a feature vector xi and a label yi. In order to see this, we can write Fp(w) = ∑n i=1(w >xi √ pi − yi √ pi)
2 = ‖Aw − b‖2, where A = (x1 √ p1, . . . ,xn √ pn) >,b = (y1 √ p1, . . . , yn √ pn) >. It has been shown in many previous studies that such Fp(w) satisfies a PL condition [47, 51, 34]. Hence, the above lemma indicates Fdro(w) satisfies a PL condition.
We can also justify that Fdro(w) satisfies a PL condition for deep learning with ReLU activation function in a neighborhood around a random initialized point following the result in [1]. Lemma 2. Assume that input {(x1, y1), . . . , (xn, yn)} satisfies ‖xi‖ = 1 and ‖xi−xj‖ ≥ δ, where xn ∈ Rd1 , yi ∈ Rd0 and ‖yi‖ ≤ O(1). Consider a deep neural network with hi,0 = φ(Axi), hi,l =
Algorithm 2: RECOVER(w0, 0, c) 1: Initialization: Draw a sample z0 and construct the estimates u0 = gz0(w0), v0 = ∇gz0(w0) 2: for k = 1, . . . ,K do 3: (wk,uk,vk) = COVER(wk−1,uk−1,vk−1, ηk, Tk,True) 4: change ηk, Tk according to Theorem 2 5: end for 6: Return: wK
φ(Wlhi,l−1), l = 1, . . . , L̃, ŷi = Bhi,L̃ where A ∈ Rd2×d1 Wl ∈ Rd2×d2 , B ∈ Rd0×d2 , φ is the ReLU activation function, and `(W ; zi) = (ŷi − yi)2 is a square loss. Suppose that for any W , p∗i = exp(`(W ; zi)/λ)/ ∑n i=1 exp(`(W ; zi)/λ) ≥ p0 > 0, then with a high probability over randomness of W0, A,B for every W with ‖W −W0‖ ≤ O(1/poly(n, L̃, p−10 , δ−1), there exists a small µ > 0 such that ‖∇Fdro(W )‖2F +O( ) ≥ µ(Fdro(W )−minW Fdro(W )).
Remark: The O( ) term in the left side of the PL condition is caused by using the covering net argument for proving the high probability result. Nevertheless, it does not affect the final convergence rate.

heading: 5.2 Theoretical Analysis of RECOVER
text: Now, we are ready to present the proposed algorithm under the PL condition and its convergence result. The algorithm is described in Algorithm 2.
The first key feature of RECOVER is equipped with the practical geometrical decreases step size between stages. At each stage, we adopt a constant step size ηk rather than the polynomial decreases step size used by COVER as in Theorem 1. Another key feature of RECOVER is that it uses not only wk for restarting but also uk,vk, the corresponding online estimator of g(wk) and ∇g(wk), for restarting the next stage. It is this feature that allows us to avoid the large batch size required in other variance reduction methods to achieve the optimal sample complexity. With this feature, we can show that the variance of uk,vk is decreased by a constant factor stagewisely as shown in the following lemma.
Lemma 3. Define constants 1 = c 2σ2 64µL3 and k = 1/2 k−1, with ηk = min{
√ µ kL
2cσ , 1 16L}, Tk = O(max{ 96cσ
µ3/2 √ kL
, 16c 2σ2 µL2 k , ∆Fσ2 }), the variance of the stochastic estimator of Algorithm 2 at wk
satisfies: E[‖uk − g(wk))‖2 + ‖vk −∇g(wk))‖2] ≤ µ k. (8)
With the above lemma and the convergence bound for E[‖∇F (wk)‖2] at the k-th stage, we can show that the objective gap E[F (wk) − F∗] is decreased by a factor of 2 after each stage under the PL condition. Hence, we have the following convergence for RECOVER.
Theorem 2. Assume that assumption 1,2,3 hold. Define constants 1 = c 2σ2 64µL4 and k = 1/2 k−1. By setting ηk = min{ √ µ kL 2cσ , 1 16L}, Tk = O(max{ 96cσ µ3/2 √ kL , 2c 2σ2 µL2 k , ∆Fσ2 }), c = 104L
2, then after K = O(log( 1/ )) stages, the output of RECOVER satisfies E[F (wK)− F∗] ≤ .
Remark: It is not difficult to derive the sample complexity of RECOVER isO(max{ 1 µ3/2 √ , 1µ )) = O( 1µ ) for ≤ µ. It is notable this complexity is optimal for the considered general stochastic compositional problem, which includes stochastic strongly convex optimization as a special case, whose lower bound is O(1/(µ )) [16].
In addition, it is notable that the proposed multi-stage algorithm is very different from many other multi-stage algorithms for non-convex optimization that are based on the proximal point framework [7, 48, 41, 15]. In particular, in these previous studies, a quadratic function γ/2‖w − wk−1‖2 with an appropriate regularization parameter γ is added into the objective function at the k-th stage in order to convextify the objective function. In RECOVER, no such regularization is manually added. Nevertheless, we can still obtain strong convergence guarantee.

heading: 6 Experimental Results
text: We focus on the task of classification with imbalanced data in our experiments. Firstly, we compare RECOVER with five State-Of-The-Art (SOTA) baselines from two categories: (i) primal-dual algorithms for solving the primal-dual formulation of DRO (1), and (ii) algorithms that are designed for the stochastic compositional formulation of DRO (3). Secondly, we verify the advantages of DRO over Emperical Risk Minimization (ERM) for imbalanced data problems by comparing the test accuracy learned by optimizing DRO using RECOVER and optimizing ERM using SGD on the imbalanced datasets. Then we show the RECOVER is also an effective fine-tuning algorithm for large-scale imbalanced data training. The code for reproducing the results is released here [39].

heading: 6.1 Comparison with SOTA DRO Baselines
text: We compare RECOVER with five baselines: Restarted CIVR [53] (RCIVR), ASC-PG [45], StocAGDA [50], PG-SMD2 [41] and PES-SGDA [15]. RCIVR and ASC-PG are the state-of-the-art algorithms for solving stochastic compositional problems. RCVIR uses variance reduction techniques and leverages the PL condition, while ASC-PG does neither. Stoc-AGDA and PG-SMD2 are the primal-dual algorithms with and without leveraging the PL condition explicitly, respectively. PESSGDA is a variant of PG-SMD2 and was proposed by leveraging the PL condition for achieving faster convergence. Please note that ASC-PG and Stoc-AGDA use polynomially decreasing step sizes, RECOVER, PG-SMD2 and PES-SGDA use stagewise decreasing step size, and RCIVR uses a constant step size. The parameters of each algorithm are appropriately tuned for the best performance. All the algorithms are implemented using Pytorch and run on GeForce GTX 1080 Ti GPU.
We conduct experiments on four datasets, namely STL10 [8], CIFAR10, CIFAR100 [21], and iNaturalist2019 [46]. The original STL10, and CIFAR10, CIFAR100 are balanced data, where STL10 has 10 classes and each class has 500 training images, CIFAR10 (resp. CIFAR100) has 10 (resp. 100) classes and each class has 5K (resp. 500) training images. For STL10, CIFAR10 and CIFAR100, we artificially construct imbalanced training data, where we only keep the last 100 images of each class for the first half classes. iNaturallist2019 itself is an imbalanced dataset that contains 265,213 images with 1010 classes. We train ResNet-20 on STL10, CIFAR10, CIFAR100, and Inception-V3 on iNaturalist2019.
For fair comparison, we use the same constant batch size b for all methods except for RCIVR in which the inner loop batch size b′ and outer loop batch size Bk are hyperparameters that relate to convergence. We use the constant batch size b = 128 on CIFAR10, CIFAR100, and b = 64 on iNaturalist2019, and b = 32 on STL. For RCIVR, both the fixed inner loop batch size b′ and the initial outer loop batch size B0 are tuned in {32, 64, 128}. The outer loop mini-batch size Bk is also increased by a factor of 10 per-stage according to the theory.
For RECOVER, the initial step size η0 and the momentum parameter a0 at the first stage are tuned in {0.1, 0.2, ..., 1}, and ηk is divided by 10 after each stage and ak is updated accordingly. For RCIVR, the constant step size is tuned η ∈ {0.1, 0.2, · · · , 1}. For the ASC-PG, the step size is set to be η = c0/ta, and the momentum parameter is set to be β = 2c0/tb, where c0 is tuned from 0.01 ∼ 1 and a, b are tuned ranging from 0.1 to 0.9 by grid search, t is the number of iterations. For Stoc-AGDA, the step size for primal variable is set to be β1/(τ1 + t) and the step size for dual variable p is set to be β2/(τ2 + t). β1, β2 are tuned in [10−1, 1, 10, 102, 500, 103] and τ1, τ2 are tuned in [1, 10, 102, 500, 103]. For PES-SGDA and PG-SMD2, the algorithm have multiple stages, where each stage solves a strongly-convex strongly-concave subproblem, and step size decrease after each stage. For PES-SGDA, the number of iteration per-stage is increased by a factor of 10 and step sizes for the primal and the dual variables are decreased by 10 times after each stage, with their initial values tuned. In particular, η1 (for primal variable) is tuned in {0.1, 0.2, · · · 1} and η2 (for the dual variable) is tuned in {10−5, 10−4, 10−3}, T0 (the number of iterations for the first stage) is tuned in {5, 10, 30, 60}nb , where n is the number of training examples. As we aim to compare the optimization for the same objective in this section, λ is set to 5 both in the compositional objective (3) and min-max formulation of (1) with regularizer h(p,1/n) = λ ∑ i pi log(npi). Following the standard training strategy, we run all algorithms 120 epochs and set the time threshold 150 hours for early stopping on iNaturalist data.
We compare testing accuracy vs running time and vs the number of processed training examples separately. We present the convergence of testing accuracy in terms of running time in Figure 1 and in terms of processed training examples in Figure 2. From the results, we can observe that: (i) in terms of running time RECOVER converges faster than all baselines on all data except on the smallest data STL10, on which PES-SGDA has similar running time performance as RECOVER. The reason is that STL10 is the smallest data, which only has 3000 imbalanced training data samples and hence PES-SGDA has marginal overhead per-iteration; (ii) when the training data size is moderately large, the primal dual methods (PES-SGDA, PG-SMD2, Stoc-AGDA) have significant overhead, which makes them converge much slower than RECOVER in terms of running time. On the large iNaturalist2019 data, RECOVER can save days of training time; (iii) RECOVER is much faster than RCIVR on all datasets; (iv) ASC-PG performs reasonably well but is still not as good as RECOVER in terms of both running time and sample complexity. The convergence instability of ASC-PG verifies the robustness of RECOVER for addressing the compositional problems.

heading: 6.2 Comparison between SGD and DRO.
text: We compare the generalization performance of DRO optimized by RECOVER with traditional ERM optimized by SGD for imbalance multi-classification tasks on STL10, CIFAR10, CIFAR100. The IMbalance RATIO (IMRATIO) is defined as the number of samples in the minority classes over
the number of samples in the majority classes. We mannually construct different training sets with different IMRATIO, i.e., we only keep the last IMRATIO portion of images in the first half of classes.
Different from previous experiments, we tune λ in a certain of range {1, 5, 10, 20, 100} by a crossvalidation approach and report the best testing results. Other parameters of RECOVER is tuned according to the setting in previous experiments. We use ResNet-32 for CIFAR10, CIFAR100, and ResNet-20 for STL10. For SGD, the step size is set as η0 in the first 60 epochs, and is decreased by a factor of 10 at 60, and 90 epochs following the practical strategy [17], where η0 is tuned in {0.1, 0.5, 1} and 1 epoch means one pass of training data. We report averaged test accuracy over 5 runs with mean (variance) in Table 2. We can see that DRO with RECOVER achieves higher test accuracy with smaller variance over multiple runs on all datasets than ERM with SGD. In addition, we report the results over 5 runs of different λ on CIFAR10 with different IMRATIO in Figure 3 (left). It is obvious to see that an appropriate regularization on the dual variable can improve the performance.

heading: 6.3 Effectiveness of RECOVER as a Fine-tuning Method
text: Fine-tuning high level layers from a pertained model is widely used for transfer learning and is also an effective method to update the models without increasing the computational cost too much when receiving new samples. For this purpose, we demonstrate that DRO is a better objective than the Cross Entropy (CE) loss and focal loss for fine-tuning on imbalanced datasets.
ImageNet-LT [29] and Places-LT [29] are two popular imbalanced data sets and are the Long-Tailed (LT) version of ImageNet-2012 [10] and Places-2 [54] by sampling a subset following the Pareto distribution [3] with the power value 6. ImageNet-LT has 115.8K images from 1000 categories, and Places-LT contains 62.5K training images from 365 classes. The head class is 4980 images and the tail class contains 5 images in both datasets.
To verify that DRO is a better objective and that RECOVER is an efficient optimization algorithm, we compare the test accuracy of the model trained with different objectives: DRO, CE loss and focal loss, where DRO is optimized by RECOVER and the other two losses are optimized by SGD. All methods start from the same pretrained model. We apply the ImangeNet pretrained ResNet152 as the pre-trained model for Places-LT. For ImageNet-LT, we train ResNet50 using CE loss for 90 epochs following the standard training strategy proposed in [17] as the pre-trained model. We then fine tune the last block of the convolutions layer and the classifier layer for 30 epochs by using RECOVER for optimizing DRO and using SGD for optimizing ERM, respectively. The initial step size for RECOVER and SGD are both tuned in η0 ∈ {0.1, 0.5, 1}. For DRO, λ is tunes in {1, 5, 10}. The test accuracy over 3 runs with mean (variance) is reported in Figure 3 (right). It is clear to see that DRO optimized by RECOVER outperforms ERM with the CE loss and focal loss optimized by SGD more than 1(%) on both datasets. This vividly verifies the effectiveness of RECOVER as a fine-tuning method on imbalanced data.

heading: 7 Conclusion
text: In this paper, we proposed a duality-free online method for solving a class of distributionally robust optimization problems. We used a KL divergence regularization on the dual variable and transformed the problem into a two-level stochastic compositional problem. By leveraging a practical PL condition, we developed a practical method RECOVER based on recursive variance-reduced estimators and established an optimal sample complexity. Experiments verify the effectiveness of the proposed algorithm in terms of both running time and prediction performance on large-scale imbalanced data. An open question remains is how to solve the DRO problem with a KL constraint on the dual variable by a pratical stochatic algorithm without maintaining and updating the high dimensional dual variable. We plan to address this challenge in the future work.

heading: Acknowledgments
text: The authors thank anonymous reviewers for constructive comments. This work was supported by NSF Career Award #1844403, NSF Award #2110545 and NSF Award #1933212.

Review summaries:
This work proposes a practical online optimization method COVER and RECOVER for solving DRO with f-divergence, see eq(1).  They propose an online algorithm (RECOVER) for solving a more general formulation, by recasting the formulation into a composition minimization problem via fully exploiting the hidden structure. Moreover, they present the convergence result for the proposed method under the PL condition. Finally, the authors conduct experiments for several deep learning problems.

This paper propose an online algorithm to solve a class of Distributionally Robust Optimization problem with non-convex objective functions. The standard primal-dual method typically suffer from optimization over high dimensional variable, which is impractical in many real world applications. The sota complexity of proposed method is provided with and without the PL condition. Empirical experiments demonstrate the superiority of the proposed algorithm.

This paper studies the family of distributionally robust optimization (DRO) problems with non-convex objectives and KL-divergence. With the motivation of improving the computational efficiency on existing methods and applying to online settings, the authors proposed a new duality-free online stochastic methods for solving the considered DRO problem. Theoretical analyses on complexities are provided for the proposed algorithm with and without a PL condition. Comparisons with state-of-the-art DRO methods are presented over ERM on imbalanced datasets.