Accelerated Proximal Alternating Gradient-Descent-Ascent for Nonconvex Minimax Machine Learning

Minimax optimization is an emerging and important optimization framework that covers a variety of modern machine learning applications. Some popular application examples include generative adversarial networks (GANs) [13], adversarial machine learning [37], game theory [10], reinforcement learning [34], etc. A standard minimax optimization problem is written as $\min_{x\in\mathbb{R}^{m}}\max_{y\in\mathcal{Y}}\leavevmode\nobreak\ f(x,y)$, where $f$ is a differentiable bivariate function. A basic algorithm for solving the above minimax optimization problem is the gradient-descent-ascent (GDA), which simultaneously performs gradient descent update and gradient ascent update on the variables $x$ and $y$, respectively, i.e., $x_{t+1}=x_{t}-\eta_{x}\nabla_{1}f(x_{t},y_{t})$, $y_{t+1}=y_{t}+\eta_{y}\nabla_{2}f({x_{t}},y_{t})$. Here, $\nabla_{1}$ and $\nabla_{2}$ denote the gradient operator with regard to the first and the second variable, respectively. The convergence rate of GDA has been studied under various types of geometries of the minimax problem, including strongly-convex-strongly-concave geometry [9], nonconvex-(strongly)-concave geometry [23] and Lojasiewicz-type geometry [5]. Recently, by leveraging the popular momentum technique [28, 38, 3, 11, 20, 21] for accelerating gradient-based algorithms, accelerated variants of GDA have been proposed for strongly-convex-strongly-concave [45] and nonconvex-strongly-concave [17] minimax optimization. There are other accelerated GDA-type algorithms that achieve a near-optimal complexity [22, 46], but they involve complex nested-loop structures and require function smoothing with many fine-tuned hyper-parameters, which are not used in practical minimax machine learning.

Another important variant of GDA that has been widely used in practical training of minimax machine learning is the alternating-GDA (AltGDA) algorithm, which updates the two variables $x$ and $y$ alternatively via $x_{t+1}=x_{t}-\eta_{x}\nabla_{1}f(x_{t},y_{t})$, $y_{t+1}=y_{t}+\eta_{y}\nabla_{2}f({x_{t+1}},y_{t})$. Compared with the update of GDA, the $y$-update of AltGDA uses the fresh $x_{t+1}$ instead of the previous $x_{t}$, and it is shown to converge faster than the standard GDA algorithm [4, 2, 12, 42]. Despite the popularity of the AltGDA algorithm, it is shown to suffer from a high computation complexity in nonconvex minimax optimization. Therefore, this study aims to improve the complexity of AltGDA by leveraging momentum acceleration techniques. In particular, in the existing literature, the convergence of momentum accelerated AltGDA is only established for convex-concave minimax problems [43] and bilinear minimax problems [12], and it has not been explored in nonconvex minimax optimization that covers modern machine learning applications. Therefore, this study aims to fill in this gap by developing a single-loop proximal-AltGDA algorithm with momentum acceleration for solving a class of regularized nonconvex minimax problems, and analyze its convergence and computation complexity.

In this section, we introduce the problem formulation and technical assumptions. We consider the following class of regularized minimax optimization problems.

where $f:\mathbb{R}^{m}\times\mathcal{Y}\to\mathbb{R}$ is differentiable and nonconvex-strongly-concave, $g,h$ are possibly non-smooth convex regularizers, and $\mathcal{Y}$ is a bounded set with diameter $R$. In particular, define the envelope function $\Phi(x):=\max_{y\in\mathcal{Y}}f(x,y)-h(y)$, then the problem $\mathrm{(P)}$ can be rewritten as the minimization problem $\min_{x\in\mathbb{R}^{m}}\Phi(x)+g(x)$.

Throughout the paper, we adopt the following assumptions on the problem (P). These are standard assumptions that have been considered in the existing literature [23, 5].

In particular, item 3 of the above assumption allows the regularizers $g,h$ to be any convex function that can be possibly non-smooth. Next, consider the mapping $y^{*}(x)=\arg\max_{y\in\mathcal{Y}}f(x,y)-h(y)$, which is uniquely defined due to the strong concavity of $f(x,\cdot)$. The following proposition is proved in [23, 5] that characterizes the Lipschitz continuity of the mapping $y^{*}(x)$ and the smoothness of the function $\Phi(x)$. Throughout the paper, we denote $\kappa=L/\mu>1$ as the condition number, and denote $\nabla_{1}f(x,y),\nabla_{2}f(x,y)$ as the gradients with respect to the first and the second input variables, respectively.

Lastly, recall that the minimax problem (P) is equivalent to the minimization problem $\min_{x\in\mathbb{R}^{m}}\Phi(x)+g(x)$. Therefore, the optimization goal of the minimax problem (P) is to find a critical point $x^{*}$ of the nonconvex function $\Phi(x)+g(x)$ that satisfies the optimality condition $\mathbf{0}\in\partial(\Phi+g)(x^{*})$, where $\partial$ denotes the subdifferential operator.

In this section, we propose a single-loop proximal alternating-GDA with momentum (proximal-AltGDAm) algorithm to solve the regularized minimax problem (P).

We first recall the update rules of the basic proximal alternating-GDA (proximal-AltGDA) algorithm [4] for solving the problem (P). Specifically, proximal-AltGDA alternates between the following two proximal-gradient updates (a.k.a. forward-backward splitting updates [24]).

(Proximal-AltGDA):

To elaborate, the first update is a proximal gradient descent update that aims to minimize the nonconvex function $f(x,y_{t})+g(x)$ from the current point $x_{t}$, and the second update is a proximal gradient ascent update that aims to maximize the strongly-concave function $f({x_{t+1}},y)-h(y)$ from the current point $y_{t}$. More specifically, the two proximal gradient mappings are formally defined as

Compared with the standard (proximal) GDA algorithm [23, 5], the proximal ascent step of proximal-AltGDA evaluates the gradient at the freshest point $x_{t+1}$ instead of $x_{t}$. Such an alternative update is widely used in practice and usually leads to better convergence properties [4, 2, 12, 42].

Next, we introduce momentum schemes to accelerate the convergence of proximal-AltGDA. As the two proximal gradient steps of proximal-AltGDA are used to solve two different types of optimization problems, namely, the nonconvex problem $f(x,y_{t})+g(x)$ and the strongly-concave problem $f({x_{t+1}},y)-h(y)$, we consider applying different momentum schemes to accelerate these proximal gradient updates. Specifically, the proximal gradient descent step minimizes a composite nonconvex function, and we apply the heavy-ball momentum scheme [33] that was originally designed for accelerating nonconvex optimization. Therefore, we propose the following proximal gradient descent with heavy-ball momentum update for minimizing the nonconvex part of the problem (P).

(Heavy-ball momentum):

To explain, the first step is an extrapolation step that applies the momentum term $\beta(x_{t}-x_{t-1})$ (with momentum coefficient $\beta>0$), and the second proximal gradient step updates the extrapolation point $\widetilde{x}_{t}$ using the original gradient $\nabla_{1}f(x_{t},y_{t})$. In conventional gradient-based optimization, gradient descent with such a heavy-ball momentum is guaranteed to find a critical point of smooth nonconvex functions [32, 31].

On the other hand, as the proximal gradient ascent step of proximal-AltGDA maximizes a composite strongly-concave function, we are motivated to apply the popular Nesterov’s momentum scheme [29], which was originally designed for accelerating strongly-concave (convex) optimization. Specifically, we propose the following proximal gradient ascent with Nesterov’s momentum update for maximizing the strongly-concave part of the problem (P).

(Nesterov’s momentum):

To elaborate, the first step is a regular extrapolation step that involves momentum, which is the same as the first step of the heavy-ball scheme. The only difference from the heavy-ball scheme is that the starting point of the second proximal gradient step changes from $y_{t}$ to its extrapolated point $\widetilde{y}_{t}$.

We refer to the above algorithm design as proximal-AltGDA with momentum (proximal-AltGDAm), and the algorithm updates are formally presented in Algorithm 1. It can be seen that proximal-AltGDAm is a simple algorithm that has a single loop structure, and adopts alternating updates with momentum acceleration. More importantly, it involves only standard hyper-parameters such as the learning rates and momentum parameters and therefore is easy to implement in practice. Such a practical algorithm is much simper than the other accelerated GDA-type algorithms that involve complex nested-loop structure and require fine-tuned function smoothing [22, 46].

Although the proposed proximal-AltGDAm algorithm applies the popular heavy-ball momentum and Nesterov’s momentum to the GDA updates in a straightforward way, its convergence analysis cannot directly follow from the existing studies of momentum accelerated gradient descent algorithms. To explain more specifically, notice that in the $x$-proximal gradient update of Algorithm 1, it involves the partial gradient $\nabla_{1}f(x_{t},y_{t})$, which corresponds to the gradient of the time-varying function $f(\cdot,y_{t})$ (since $y_{t}$ changes over time $t$). Similarly, the $y$-proximal gradient update involves the gradient of another time-varying function $f(x_{t+1},\cdot)$. Therefore, both momentum accelerated proximal gradient updates are actually applied to time-varying functions due to the nature of GDA updates. In this sense, the existing analysis of momentum accelerated gradient descent algorithms cannot be directly applied to analyze this algorithm.

To analyze the convergence of Algorithm 1, we first study the $x$-proximal gradient update with heavy-ball momentum and obtain the following characterization of per-iteration progress on the objective function value.

The above proposition tracks the per-iteration optimization progress made by the $x$-proximal gradient update with heavy-ball momentum. To elaborate, the increment terms $\|x_{t+1}-x_{t}\|^{2},\|x_{t}-x_{t-1}\|^{2}$ are induced by the heavy-ball momentum scheme. Moreover, since the $x$-update uses the partial gradient $\nabla_{1}f(x_{t},y_{t})$ to approximate the exact gradient $\nabla\Phi(x_{t})=\nabla_{1}f(x_{t},y^{*}(x_{t}))$, it naturally induces a tracking error term $\|y_{t}-y^{*}(x_{t})\|^{2}$ that tracks the optimization gap of the $y$-update. Hence, we need to further bound this tracking error term by analyzing the $y$-proximal gradient update with Nesterov’s momentum, which we explore next.

As explained earlier, the momentum accelerated $y$-updates in proximal-AltGDAm are applied to a time-varying strongly-concave function $f(x_{t+1},\cdot)$. Hence, the tracking error term $\|y_{t}-y^{*}(x_{t})\|^{2}$ cannot be directly bounded using the standard convergence result of Nesterov’s accelerated proximal gradient algorithm [28]. Instead, we can view these $y$-updates as inexact accelerated proximal gradient updates. To elaborate, note that in the $t$-th iteration, the $y$-proximal gradient update is applied to the function $f(x_{t+1},\cdot)$. Then, we can rewrite the $y$-updates in all the previous iterations $k=0,1,...,t-1$ as follows.

That is, we can view all the previous $y$-updates as applied to the fixed function $f(x_{t+1},\cdot)$ with an inexactness $\mathbf{e}_{k+1}$ involved in the computation of the partial gradient $\nabla_{2}f(x_{t+1},\widetilde{y}_{k})$. In summary, the $y$-updates of proximal-AltGDAm can be understood as inexact accelerated gradient updates applied to the function $f(x_{t+1},\cdot)$ at time $t$. In particular, under the smoothness condition in 1, the inexactness is bounded as $\|\mathbf{e}_{k+1}\|\leq L\|x_{k+1}-x_{t+1}\|$. Consequently, we can leverage the existing convergence result of inexact accelerated gradient algorithm [36] to bound the optimality gap $\|y_{t}-y^{*}(x_{t})\|^{2}$ as follows.

Intuitively, in the above bound, the first term on the right hand side corresponds to the normal accelerated convergence rate, and the other term is induced by the inexactness $\mathbf{e}_{k}$. As both terms are scaled by the accelerated convergence factor $(1-\kappa^{-0.5})$, we expect that the above bound converges fast and further facilitates the convergence of eq. 2. Next, substituting eq. 4 into eq. 2 and telescoping, we obtain the following asymptotic stability properties of proximal-AltGDAm.

Remark 1. In [5], the proximal-GDA algorithm (without alternating update and momentum) uses a small learning rate $\eta_{x}\leq\mathcal{O}({L^{-2}}\kappa^{-3})$ to establish convergence. As a comparison, proximal-AltGDAm allows to choose a much larger learning rate $\eta_{x}\leq{\mathcal{O}(L^{-1}\kappa^{-\frac{11}{6}})}$ to guarantee a faster convergence (proved later), thanks to the momentum acceleration schemes.

Therefore, if we can show that $\{x_{t}\}_{t}$ converges to a desired critical point $x^{*}$, then the above stability properties of proximal-AltGDAm implies that $\{y_{t}\}_{t}$ converges to the corresponding best response point $y^{*}(x^{*})$.

To further characterize the global convergence property of proximal-AltGDAm, we define the following proximal gradient mapping associated with the objective function $\Phi(x)+g(x)$.

The proximal gradient mapping is a standard metric for evaluating the optimality of nonconvex composite optimization problems [28]. It can be shown that $x$ is a critical point of $(\Phi+g)(x)$ if and only if $G(x)=\mathbf{0}$, and it reduces to the normal gradient when there is no regularizer $g$. Hence, our convergence criterion is to find a near-critical point $x$ that satisfies $\|G(x)\|\leq\epsilon$ for some pre-determined accuracy $\epsilon>0$.

Next, by leveraging Proposition 2 and Proposition 3, we obtain the following global convergence result of proximal-AltGDAm and characterize its computational complexity (number of gradient and proximal mapping evaluations).

To elaborate, the first statement of 1 shows that the sequence generated by proximal-AltGDAm converges to critical points of the minimax problem. The second statement proves that the computation complexity of proximal-AltGDAm for finding a near-critical point is of the order $\mathcal{O}\big{(}\kappa^{\frac{11}{6}}\epsilon^{-2}\big{)}$, which strictly improves the complexity $\mathcal{O}\big{(}\kappa^{2}\epsilon^{-2}\big{)}$ of both GDA [23] and proximal-AltGDA [4] in nonconvex-strongly concave minimax optimization. We note that the improvement is substantial when the problem condition number $\kappa=L/\mu$ is large, while the dependence on $\epsilon^{-2}$ is generally unimprovable in nonconvex optimization. In addition, thanks to the momentum acceleration schemes, our choice of learning rate $\eta_{x}=\mathcal{O}(L^{-1}\kappa^{-\frac{11}{6}})$ is more flexible than that $\eta_{x}=\mathcal{O}(L^{-1}\kappa^{-2})$ adopted by these GDA-type algorithms. These improvements are not only attributed to momentum acceleration, but also to the elaborate selection of the coefficients in the proof that aims to minimize the dependence of the complexity on $\kappa$.

In this section, we compare the performance of proximal-AltGDAm with that of other GDA-type algorithms via numerical experiments. Specifically, we compare proximal-AltGDAm with the standard proximal-GDA/AltGDA algorithm [5] and the single-loop accelerated AltGDA algorithm (APDA) [47]. All these algorithms have a single-loop structure.

We consider solving the following regularized Wasserstein robustness model (WRM) [37] using the MNIST dataset [19].

where $n=60$k is the number of training samples, $\theta$ is the model parameter, $(x_{i},y_{i})$ corresponds to the $i$-th data sample and label, respectively, and $\xi_{i}$ is the adversarial sample corresponding to $x_{i}$. We choose the cross-entropy loss function $\ell$. We add an $\ell_{1}$ regularization to impose sparsity on the adversarial examples, and add an $\ell_{2}^{2}$ regularization to prevent divergence of the model parameters.

We set $\lambda=1.0$ that suffices to make the maximization part be strongly-concave, and set $\lambda_{1}=\lambda_{2}=10^{-4}$. We use a convolution network that consists of two convolution blocks followed by two fully connected layers. Specifically, each convolution block contains a convolution layer, a max-pooling layer with stride step $2$, and a ReLU activation layer. The convolution layers in the two blocks have $1,10$ input channels and $10,20$ output channels, respectively, and both of them have kernel size $5$, stride step $1$ and no padding. The two fully connected layers have input dimensions $320,50$ and output dimensions $50,10$, respectively.

We implement all three algorithms using full gradients with the whole training set of $60$k images. We choose the same learning rates $\eta_{x}=\eta_{y}=10^{-3}$ for all algorithms. For proximal-AltGDAm, we choose momentum $\beta=0.25$ and $\gamma=0.75$. For APDA, we choose the fine-tuned $\eta=2\eta_{x}$. As the function $\Phi(x)$ cannot be exactly evaluated, we run $100$ steps of stochastic gradient ascent updates with learning rate $0.1$ to maximize $f(x_{t},y)-h(y)$ and obtain an approximated $y^{*}(x_{t})$, which is used to estimate $\Phi(x)+g(x)$.

Figure 1 (Left) compares the estimated objective function value achieved by all the three algorithms. It can be seen that proximal-AltGDAm achieves the fastest convergence among these algorithms and is significantly faster than the proximal-GDA and the proximal-AltGDA. This demonstrates the effectiveness of our simple momentum scheme. Figure 1 (Right) further demonstrates the advantage of proximal-AltGDAm in the test accuracy. It can be seen that the robust model trained by proximal-AltGDAm achieves a much higher test accuracy.

We develop a single-loop and fast AltGDA algorithm that leverages proximal gradient updates and momentum acceleration to solve general regularized nonconvex-strongly-concave minimax optimization problems. By viewing the GDA updates of the algorithm as inexact accelerated gradient updates, we prove that the algorithm converges to a $\epsilon$-critical point with a computational complexity $\mathcal{O}(\kappa^{\frac{11}{6}}\epsilon^{-2})$, which substantially improves the state-of-the-art result. In the future work, it is interesting to develop a stochastic variant of this algorithm to further improve the sample complexity.

The work of Ziyi Chen, Shaocong Ma and Yi Zhou was supported in part by U.S. National Science Foundation under the Grants CCF-2106216 and DMS-2134223.