On the Convergence of Stochastic Compositional Gradient Descent Ascent Method

1 Introduction

In recent years, minimax optimization has attracted increasing attention in the community of machine learning. This is mainly due to the fact that a broad range of machine learning models can be formulated as the minimax optimization problem, including generative adversarial networks [Goodfellow et al., 2014], adversarial training of deep neural networks [Madry et al., 2017], and distributionally robust machine learning models [Chen et al., 2017]. In the meanwhile, numerous machine learning models can be formulated as the compositional optimization problem, such as policy evaluation in reinforcement learning [Sutton and Barto, 2018], riskaverse portfolio optimization [Rockafellar, 2007], and sparse additive models [Wang et al., 2017a]. Given the growing importance of the compositional minimax problem in machine learning, in this paper, we are interested in optimizing the compositional minimax problem as follows:

where g(·) : X → Rp , f(·, ·) : (Rp , Y) → R, X and Y are convex and compact sets. In particular, f(g(x), y) is a strongly concave function with respect to y for all x ∈ X , and f(g(x), y) is a nonconvex compositional function regarding x for all y ∈ Y. A typical example of this kind of compositional minimax problem is the distributionally robust compositional optimization problem:

∗Corresponding author

where X = Rd and Y = {y ∈ Rn | ∑ni=1 yi = 1, yi ≥ 0, ∀i}. The loss function for each sample fi( 1m ∑mj=1 gj(x)) is a compositional function, and y ∈ Rn weights each sample to handle noisy data to learn a robust model. The goal of this paper is to develop efficient algorithms to solve the compositional minimax optimization problem in Eq. (1).

A lot of efforts [Lin et al., 2020; Luo et al., 2020; Xu et al., 2020; Huang et al., 2020; Chen et al., 2020; Zhang et al., 2020; Tran-Dinh et al., 2020; Yan et al., 2020] have been recently made to develop efficient algorithms to optimize the minimax optimization problem. For instance, [Lin et al., 2020] proposed a stochastic gradient descent ascent method for the nonconvex-strongly-concave minimax problem. [Luo et al., 2020; Xu et al., 2020] developed a variance reduced stochastic gradient descent ascent (SGDA) method to accelerate the convergence speed of SGDA. [Huang et al., 2020] studied an optimization method for the minimax problem on Riemannian manifold. [Chen et al., 2020] presented a projection-free approach for the convex-strongly-concave constrained minimax problem. However, all of these approaches only focus on the non-compositional machine learning model where a stochastic gradient is an unbiased estimation of the full gradient. Thus, these methods are not applicable to optimize the compositional minimax optimization problem in Eq. (1).

Due to the two-level stochasticity in the compositional loss function, a stochastic gradient is a biased estimation of the full gradient. It is much more challenging to optimize the minimax problem in Eq. (1). Although there exist some works studying how to deal with the two-level stochasticity in optimizing the compositional minimization problem, these prior work is not applicable for the minimax problem in Eq. (1). For instance, [Wang et al., 2017a] proposed a stochastic compositional gradient descent (SCGD) method to handle the two-level stochasticity. After that, a string of variance reduced variants [Yuan et al., 2019; Yang and Hu, 2020;

Zhang and Xiao, 2019b] have been proposed to accelerate the convergence speed. However, for these methods, it is not clear how to optimize the maximization subproblem in Eq. (1). Thus, it is necessary and important to design new optimization methods for Eq. (1).

In this study, to address the aforementioned challenges, we develop a novel stochastic compositional gradient descent ascent (SCGDA) method for optimizing Eq. (1). In particular, when optimizing the minimization subproblem, we use a new strategy to estimate the inner function value g(x) and its gradient ∇g(x), which helps to control the estimation variance of the stochastic compositional gradient. Meanwhile, when optimizing the maximization problem, we employ the standard stochastic gradient ascent method to update y. Our theoretical result indicates that SCGDA can achieve the convergence rate of O(κ4 /ϵ4 ), where κ is the condition number of the loss function. To the best of our knowledge, this is the first work achieving such a convergence rate. Additionally, our experimental results confirm the effectiveness of the proposed optimization algorithm. The contributions of this work are summarized as follows:

• We proposed a new optimization algorithm to optimize the compositional minimax optimization problem. This is the first work studying how to optimize this kind of minimax compositional problem.
• We theoretically proved that our proposed algorithm enjoys a convergence rate of O(κ4 /ϵ4 ). This is the first work achieving such a convergence rate.
• We conducted extensive experiments and our experimental results confirm the effectiveness of the proposed algorithm.

2 RelatedWorks

In this section, we briefly revisit related works to motivate our proposed approaches.

Minimax optimization problem. The minimax optimization problem is an important type of models in machine learning, which motivates numerous efforts to develop efficient algorithms to solve the problem. Currently, the basic idea for minimax optimization is to alternatively optimize the minimization and maximization subproblems. For instance, based on the regular SGD method, [Lin et al., 2020] developed the stochastic gradient descent ascent (SGDA) method, which uses stochastic gradient to alternatively update the two subproblems. [Luo et al., 2020; Xu et al., 2020] applied a recursive variance reduction technique [Fang et al., 2018] for SGDA to improve the convergence rate of SGDA. [Yan et al., 2020] provided an epoch-wise stochastic gradient descent ascent method for strongly-convex-stronglyconcave problems. Moreover, [Huang et al., 2020] proposed the Riemannian stochastic gradient descent ascent method and some variants for the Riemannian minimax optimization problem. [Qiu et al., 2020] reformulated nonlinear TemporalDifference (TD) learning as a minimax optimization problem and proposed the single-timescale stochastic gradient descent ascent method. [Chen et al., 2020] developed a stochastic Frank-Wolfe method to optimize the constrained minimax problem. All of these methods fail to handle the compositional structure in the compositional minimax optimization problem shown in Eq. (1).

Compositional minimization problem. The compositional optimization problem is defined as follows:

A typical challenge to optimize the compositional minimization problem is that the stochastic gradient is not an unbiased estimation of the full gradient. To address this issue, various methods have been proposed in recent years. For instance, [Wang et al., 2017a] uses the stochastic gradient for the inner function and uses a momentum-like method to compute the inner function value when computing the stochastic gradient. However, its convergence rate is as slow as O(1/ϵ8 ) for nonconvex problems, which is much worse than the regular SGD method. Consequently, some accelerated variants [Wang et al., 2017a; Wang et al., 2017b] have been proposed to improve the convergence rate. However, their convergence rates are still inferior to that of the regular SGD method. Recently, based on the variance reduction technique such as SPIDER [Fang et al., 2018]and STORM [Cutkosky and Orabona, 2019], a series of methods [Zhang and Xiao, 2019b; Zhang and Xiao, 2019a; Yuan et al., 2019; Yang and Hu, 2020] have been proposed to further improve the convergence rate. Nonetheless, these methods can only be applied to the compositional minimization problem rather than the minimax problem.

To sum up, the aforementioned minimax optimization algorithms and compositional minimization algorithms cannot be applied to optimize the compositional minimax problem defined in Eq. (1). Therefore, in this study, we will develop new optimization algorithms to optimize this challenging problem.

Algorithm 1 The Stochastic Compositional Gradient Descent Ascent Method (SCGDA)

Initialization: x1, y1 = y∗ (x1), u′1 = ∇g(x1), u1 = g(x1), w1 = ∇yf(u1, y1), v1 = (u′t )T w1, α > 0, γ > 0,

• λ > 0, ηt > 0 1: for t = 1, · · · , T do 2: Update x: xt+1 = xt - γηtvt 3: Update y: ˜ yt+1 = P(yt + λwt) yt+1 = yt + ηt(˜ yt+1 - yt) 4: Randomly sample a subset Bξ,t+1, compute ∇g(xt+1; Bξ,t+1) and g(xt+1; Bξ,t+1), ut+1 = (1 - αηt)ut + αηtg(xt+1; Bξ,t+1) u′t+1 = (1 - αηt)u′t + αηt∇g(xt+1; Bξ,t+1) 5: Randomly sample a subset Bζ,t+1, compute ∇gf (ut+1, yt+1; Bζ,t+1) and ∇yf (ut+1, yt+1; Bζ,t+1), vt+1 = (u′t+1 )T ∇gf(ut+1, yt+1; Bζ,t+1) wt+1 = ∇yf (ut+1, yt+1; Bζ,t+1) 6: end for

3 Stochastic Compositional Gradient Descent Ascent Method

In this section, we present the details of our proposed algorithm and provide its theoretical convergence rate.

In Algorithm 1, we propose a stochastic compositional gradient descent ascent (SCGDA) method to optimize the compositional minimax problem defined in Eq. (1). Here, a key challenge is that the stochastic gradient regarding x is NOT an unbiased estimation of the full compositional gradient shown as follows:

To alleviate this issue, when optimizing the compositional minimization problem, [Wang et al., 2017a] employs the following strategy 1 to estimate the compositional gradient:

where g(xt;Bξ,t) = 1 |Bξ,t | ∑j∈Bξ,t gj(xt), ∇g(xt;Bξ,t) = 1 |Bξ,t | ∑j∈Bξ,t ∇gj(xt) , α > 0, and ηt > 0. Here, ut estimates g(x) and vt is the estimation of ∇g(x)T ∇gf(g(x)). With the strategy in Eq. (5), ut is supposed to have a smaller estimation variance than g(xt;Bξ,t). In fact, this strategy of controlling variance is also used in non-compositional optimization methods [Mokhtari et al., 2020; Qiu et al., 2020]. However, theoretical results in [Wang et al., 2017a] indicate that its convergence rate is as slow as O(1/ϵ8 ).

In Algorithm 1, when t > 1, we use the following strategy to compute ut to estimate the inner function g(xt) and u′t for its gradient ∇g(xt) :

In this way, both ut and u′t will have a small estimation variance, which will be shown in the next section. Based on ut and u′t , we compute the stochastic compositional gradient regarding x as follows:

where ∇gf(ut, yt;Bζ,t) = 1 |Bζ,t | ∑i∈Bζ,t ∇gfi(ut, yt). Then, we use the following strategy to update x:

where γ > 0. This strategy can facilitate tightly bounding the estimation variance E[‖ut -g(xt)‖2 ] and E[‖u′t -∇g(xt)‖2 ], which will be demonstrated in the next section.

To optimize the maximization subproblem regarding y, we directly compute the stochastic gradient wt = ∇yf(ut, yt;Bζ,t) = 1 |Bζ,t | ∑i∈Bζ,t ∇yfi(ut, yt) and update y as follows:

1 Note that this strategy is designed for the compositional minimization problem, not the compositional minimax problem. Here, we just use it to motivate our method.

where λ > 0, 0 < ηt < 1, and PY(·) is a projection operator to make sure that ˜ yt+1 satisfy the constraint Y. Since Y is a convex set, the combination between ˜ yt+1 and yt in Eq. (9) still lies in Y.

In the following, we will investigate the convergence rate of Algorithm 1. We first introduce some commonlyused assumptions for compositional optimization. To make these assumptions easy to follow, we denote ∇f(a, b) = (∇af(a, b),∇bf(a, b)) for (a, b) ∈ A × B where A = {g(x)|x ∈ X} and B = Y.

Assumption 1. (Smoothness) For ∀(a1, b1), (a2, b2) ∈ A× B, there exists L > 0, such that

For ∀x1, x2 ∈ X , there exists Lg > 0, such that

Assumption 2. (Bounded gradient) There exist Cg > 0 and Cf > 0, such that

Assumption 3. (Bounded variance) There exist σf > 0, σg > 0, and σ′g > 0, such that

Assumption 4. (Strong concavity) For ∀a ∈ Aand ∀b1, b2 ∈ B, there exists µ > 0, such that

Additionally, we introduce two auxiliary functions: Φ(x) = maxy∈Y f(g(x), y) and y∗ (x) = argmaxy∈Y f(g(x), y). Here, Φ(x) is LΦ-smooth where LΦ = 2C 2 g L2 /µ + CfLg. We use Φ∗ to represent the minimum value of Φ(x). Based on these assumptions and definitions, we have the following convergence result.

From Theorem 1, it can be seen that ηγ = O( 1 κ2 ) where κ = L/µ. Then, by setting B = T, we can get the convergence rate O( κ √T ). Hence, to achieve ϵ-accuracy solution, the total sample complexity is B×T = O(κ4 /ϵ4 ). Note that, although it is incomparable with SCGD [Wang et al., 2017a] for the compositional minimization problem, we observe that our sample complexity has a much better dependence on ϵ.

4 Convergence Analysis

In this section, we present the high-level idea of the convergence analysis for Theorem 1.

As discussed in the last section, our Algorithm 1 results in a tighter bound for E[‖ut -g(xt)‖2 ] and E[‖u′t -∇g(xt)‖2 ], which makes our algorithm converge faster comparing with existing work. In particular, we have the following bound for these two variances.

Lemma 1. Given Assumptions 1-3 for Algorithm 1, by setting ηt ≤min{ 1 8α , 1} and |Bζt | = |Bξt | = B, we can get

Lemma 2. Given Assumptions 1-3 for Algorithm 1, by setting ηt ≤min{ 1 8α , 1} and |Bζt | = |Bξt | = B, we can get

Note that, [Wang et al., 2017a] employs a similar way to estimate g(xt) and obtains a similar bound for the estimation variance. However, our bound is tighter. In particular, for our bound in Lemma 1, the learning rate ηt < 1 lies in the nominator of the second term on RHS while it resides in the denominator in the bound of [Wang et al., 2017a]. Thus, our bound for the estimation variance is much tighter.

Furthermore, the following two lemmas demonstrate that controlling the estimation variance of g(x) and ∇g(x) can benefit controlling the estimation variance of the compositional gradient ∇xf(g(x), y) and that of the gradient ∇yf (g(x), y).

Lemma 3. Given Assumptions 1-3, for Algorithm 1, by setting |Bζt | = |Bξt | = B, we can get

Lemma 4. Given Assumptions 1-3, for Algorithm 1, by setting |Bζt | = |Bξt | = B, we can get

Moreover, we need an additional lemma to prove Theorem 1.

Lemma 5. Given Assumptions 1-3, for Algorithm 1, by setting λ ≤ 1 6L and ηt < 1, we can get

Based on aforementioned lemmas, we will provide the main proof for Theorem 1.

Proof. Due to the smoothness of Φ(x), we can get

where the last inequality follows from ηt ≤ 1 2γLΦ and the following inequality:

where the last inequality follows from Assumptions 1 and 2. Furthermore, according to Lemma 3 and taking expectation for both sides, we can get

Furthermore, define the Lyapunov function

By summing t over 1, · · · , T and setting ηt = η, we get

From the initialization condition, it is easy to get

Finally, we can get

which completes the proof.

5 Experiments

In this section, we conduct experiments to verify the convergence of our proposed SCGDA in Algorithm 1.

In our experiments, we use SCGDA to optimize the distributionally robust linear value function approximation in reinforcement learning [Zhang and Xiao, 2019a]. The value function in reinforcement learning is an important component to compute the reward. In detail, given a Markov decision process (MDP) {S, P π , R, r} where S = {1, 2, · · · , S} represents the state space, P π s,s′ denotes the transition probability from state s to state s′ for a given policy π, Rs,s′ is the reward when state s goes to state s′ , and r is the discount factor, then the value function at state s is defined as V (s) = ∑Ss′=1 P π s,s′ (Rs,s′ + rV (s′ )). To estimate the value function, a typical choice is to parameterize it with a linear function: ˜ Vw(s) = zT s w where zs ∈ Rd is fixed and w ∈ Rd is the model parameter which needs to be optimized. To obtain the model parameter w, we need to optimize the following problem:

Obviously, the loss function is a compositional function [Yuan et al., 2019; Zhang and Xiao, 2019c]. Here, we are interested in its distributionally robust variant, which is defined as follows:

where w = [wi] ∈ Rd is the model parameter, Y = {y = [ys] ∈ RS | ∑Ss=1 ys = 1, ys ≥ 0, ∀s}, β > 0. Obviously, it is nonconvex regarding w and strongly concave regarding y.

Following [Yuan et al., 2019], we generate an MDP which has 400 states and each state is associated with 10 actions. Regarding the transition probability, P π s,s′ is drawn from [0, 1] uniformly. Additionally, to guarantee the ergodicity, we add 10-5 to P π s,s′ . Then, we use SCGDA to optimize Eq. (31) on this dataset. Note that, our method is the first stochastic compositional gradient descent ascent method where there are no baseline methods to compare with. Thus, we will conduct experiments to show the convergence of our proposed SCGDA under different conditions.

In our experiments, we set the batch size to 20, α = 3, β = 10-5 . Then, we verify the convergence performance of SCGDA with different learning rates η. Specifically, in Figure 1, we fix γ = λ = 0.1 and change η to show the value function gap 1S ∑Ss=1 ( ˜ Vw(s) - ∑Ss′=1 P π s,s′(Rs,s′ + r ˜ Vw(s′ )))2 regarding the number of iterations. It can be seen that the value function gap decreases with the training going

on for all cases, which confirms that our algorithm is effective to optimize Eq. (31). Additionally, with a larger learning rate η = 0.1, our SCGDA method converges much faster than the other two cases.

Furthermore, in Figure 2, we fix the learning rate η and change λ, as well as γ. Here, we set λ = γ to make the minimization subproblem and maximization subproblem update in the single-timescale manner. Similarly, from Figure 2, it can be seen that the value function gap decreases with the optimization going on for all cases, and larger parameters λ, γ lead to a much faster convergence rate.

Finally, we plot the learned parameter y in Figure 3. Here, we set η to 0.1 and γ = λ = 0.1. Then, we plot y - 1S in Figure 3 for all sates. We observe that our method can learn different weights for different states, which is consistent with the idea of the distributionally robust optimization.

6 Conclusions

In this paper, we proposed a novel stochastic compositional gradient descent ascent method for optimizing the compositional minimax problem. We provided the theoretical convergence rate of the proposed method for the nonconvexstrongly-concave compositional minimax problem. We then conducted extensive experimental results that confirmed the effectiveness of our proposed algorithm.

Acknowledgments

This work is partially supported by NSF (#1955890 to X.W. and #1750632 to X.S.).