Diffusion Stochastic Optimization for Min-Max Problems

Minimax optimization problems have been largely studied under the single node setting [6, 24, 5, 9, 25, 26, 27, 28, 29, 30]. The fundamental setting assumed within these works is the strong convexity-strong concavity/strong monotonicity. For this setting, the work [6] established the tight complexity bound of the first-order algorithm in finding a near-optimal solution. The work [5] devised the first optimal algorithm that matches this bound. The work [24] showed that the alternating gradient descent-ascent (AGDA) method converges faster than its simultaneous counterpart. The work [25] proposed a stochastic variance-reduced extragradient (EG) method, which showed improved performance in training GANs. The quasi-Newton methods have also been proposed to achieve a faster convergence rate in finding the saddle point [26]. Dropping “strongly” setting from both sides of strong convexity-strong concavity brings about convergence issues where simultaneous GDA will suffer from divergent rotation while AGDA will get stuck in a limit cycle rather than converge to a limit point [24, 31]. For the convex-concave setting, the work [29] established the sublinear rate of $\mathcal{O}(\frac{1}{T})$ for the best-iterate of the stochastic OG method, namely the best iterate produced over running $T$ iterations. Achieving last-iterate convergence is more difficult; this aspect of OG was not addressed until the work [28]. Furthermore, the work [27] showed that the average-iterate of EG, namely the average of all iterates produced over iterations, has a convergence rate of $\mathcal{O}(\frac{1}{T^{2}})$ which is faster than its last-iterate. The work [30] also showed anchoring EG has a convergence rate of $\mathcal{O}(\frac{1}{T^{2}})$ without considering stochasticity. There is another line of research that investigates the structured nonmonotone problem, particularly under the premise of the (weak) Minty variational inequality— see [16, 17, 15].

In recent years, there has been an increasing interest towards the stochastic nonconvex (one-sided) minimax optimization problem. The work [9] established the convergence of stochastic GDA (SGDA) in finding the primal solution under a nonconvex-SC formulation, with the use of batch size $\mathcal{O}(\varepsilon^{-2})$ to reach $\varepsilon$ error accuracy. The work [32] proposed to incorporate a concave maximizer and variance reduction into SGDA, which is computationally costly since it performs multi-loop steps in updating the dual variable. The work [10] studied the two-sided PL setting and established the linear convergence of variance-reduced stochastic AGDA (SAGDA). Additionally, the work [8] showed that employing two iid samples for SAGDA obtains $\mathcal{O}(\frac{1}{T^{1/2}})$ convergence rate for the primal solution in nonconvex-PL setting. However, vanilla AGDA does not converge in bilinear scenarios[31], which casts doubt on its ability to solve more complex settings. Exploiting the convergence of the backbone minimax algorithms across various settings, say OG/EG, is always of independent interest. Unfortunately, the recent work [7] showed a less than ideal result under the nonconvex-SC formulation as the large batch issue persists. This fact strongly motivates us to pursue better results that theoretically address this limitation.

Over the past decade, distributed optimization has become an increasingly important field due to its flexibility in processing large volume and physically dispersed data [22, 33]. The distributed minimax optimization scenario remains an open field to be exploited since there are many diverse settings. Some distributed nonconvex minimax works appear in [34, 35, 36, 37, 14, 20]. The work [34] showed the convergence rate of $\mathcal{O}(\frac{1}{T^{2/3}})$ for the primal solution using STORM momentum-based [38] SGDA, under a nonconvex-SC setting. The work [37] extended the analysis to the nonconvex-PL setting. Other works [35, 36] introduced variance-reduced SGDA for the nonconvex-SC setting. However, variance reduction requires an excessively large batch size to estimate the gradient at the checkpoint, making it hard to tune. On the other hand, the works [34, 38, 37, 36, 35] rely on the Lipschitz continuity of the stochastic loss gradient to carry out the analysis. Relaxing this condition leads to a limited applicability of the results [39]. The work [14] proposed a distributed version of stochastic EG for solving stochastic MVI in a time-varying network. The work [20] proposed a single-call EG [40] for distributed parallel training of the GANs. However, both results face the same large batch size issue.

Our main contributions in this work are as follows:

• •

  We introduce and analyze a new variant of stochastic OG that achieves a convergence rate of $\mathcal{O}(\frac{1}{T^{1/2}})$ in finding the primal solution under the nonconvex-PL setting. Our analysis improves the existing performance bound by addressing the issue of large batch size [7]. We also provide insights into theoretical limitations that prevent the conventional stochastic OG from achieving better results.

• •

  We extend the applicability of stochastic OG to a fully distributed setting. We consider left-stochastic combination matrices, which are more general than the doubly-stochastic setting considered in [34, 35, 36, 37, 14, 20]. We also carry out convergence analysis without assuming the smoothness of the stochastic loss gradient.

• •

  We establish convergence results under more relaxed conditions, ensuring convergence for both primal and dual solutions. Previous works primarily focused on the convergence of the primal objective [34, 35, 36, 37]. We show that the dual solution converges at a rate of $\mathcal{O}(\frac{1}{T})$ for our method. Please refer to Table I for the comparison of our results with existing works.

We start with a centralized scenario where the data processing and model training are carried out at a fusion center. Specifically, we consider initially the minimax problem with $J(x,y)=\mathbb{E}Q(x,y,\bm{\xi})$ formulated at a stand-alone agent or fusion center. Several variants of the stochastic OG method will be revisited here to motivate our new formulation. It is noted that these variants were originally proposed in [19, 40] in the deterministic setting, whereas we consider the stochastic version.

The extragradient (EG) method traces back to the seminal work [19], which is based on a heuristic idea, namely “extrapolation”. It proposed to utilize the gradient at the extrapolated point, rather than the gradient at the current point for carrying out the update. Doing so leads to increased stability by alleviating the rotational force. Such a scheme was also applied to other optimization problems[41]. The stochastic EG (S-EG) was studied in the works [7, 14]. This algorithm starts from iteration $i=0$, say, with randomly initialized variables $\overline{\bm{x}}_{0},\overline{\bm{y}}_{0}$ and recursively performs the following steps:

Here, $(\bm{x}_{i},\bm{y}_{i})$ denotes the extrapolated point, and $(\overline{\bm{x}}_{i+1},\overline{\bm{y}}_{i+1})$ is the output. The variables $(\overline{\bm{\xi}}_{x,i},\overline{\bm{\xi}}_{y,i},\bm{\xi}_{x,i},\bm{\xi}_{y,i})$ are iid samples drawn during different steps and for updating different variables. We use the subscript and bar notation in these iid samples to highlight at which point and for which variables they are drawn. For instance, $\overline{\bm{\xi}}_{x,i}$ denotes an iid sample drawn at iteration $i$ for constructing the stochastic gradient at point $(\overline{\bm{x}}_{i},\overline{\bm{y}}_{i})$ and for updating the primal variable $\bm{x}$. Additionally, the scalars $\mu_{x},\mu_{y}>0$ are the stepsizes or learning rates. The S-EG needs to query two stochastic gradient oracles to update each variable. It is also executed in a sequential manner, i.e., determining the extrapolated point first before performing the update.

To reduce the computational complexity, a single-call version of EG known as past EG (PEG), was introduced in the work [40]. It proposed to use the past gradient at the point $(\bm{x}_{i-1},\bm{y}_{i-1})$. The stochastic form of PEG is characterized by the following steps [42, 20]:

PEG can also be derived from the forward-reflected-backward framework [43]. It has attracted wide interest in theoretical aspects and obtains empirical success in practical applications [28, 21, 20].

The standard form of S-PEG, also referred to as the stochastic OG (S-OG), is often utilized for theoretical analysis [7]. The updates in Eqs. (5a)–(5d) can be simplified by rewriting them in terms of $(\bm{x}_{i+1},\bm{y}_{i+1})$ at each iteration $i$ [7, 42]:

S-OG involves two different samples, namely $\bm{\xi}_{x,i}$ and $\bm{\xi}_{x,i-1}$, for updating the primal variable in different steps (similar to the dual variable). We observe that $\bm{\xi}_{x,i}$ and $\bm{\xi}_{x,i-1}$ are actually approximating different losses whose landscape differs. The gradient $\nabla_{x}Q(\bm{x}_{i},\bm{y}_{i};\bm{\xi}_{x,i})$ constructed by $\bm{\xi}_{x,i}$ can deviate from that at the same point $(\bm{x}_{i},\bm{y}_{i})$ when it is evaluated at $\bm{\xi}_{x,i-1}$. In other words, S-OG is not making the “optimistic” direction under the same landscape in a strict sense. When the sample variance is high, the deviation between the loss landscapes constructed from $\bm{\xi}_{x,i}$ and $\bm{\xi}_{x,i-1}$ could be high, leading to a poor approximation of the update direction [44]. On the other hand, the stochastic gradients involved in (6a)–(6b) are not martingale processes, which cause difficulty for analysis (see Section III-E for discussion).

We instead introduce a new formulation using the same-sample approximation. For each variable, it utilizes the gradients from the same loss landscape to perform “optimistic” scheme; these losses are also unbiased for the true risk. This can be verified by taking expectation over the distribution of the current samples $(\bm{\xi}_{x,i}$, $\bm{\xi}_{y,i})$. This approach involves recomputing the stochastic gradient at the point $(\bm{x}_{i-1},\bm{y}_{i-1})$ using the current samples $(\bm{\xi}_{x,i},\bm{\xi}_{y,i})$:

SS-OG requires double call of stochastic gradient oracles for each variable. However, it can be executed in parallel with some extra memory overhead, making it preferred in comparison to the sequential approaches, e.g., S-EG and S-PEG.

In this section, we extend SS-OG to the distributed setting (2a)–(2b). Under this scenario, each agent $k$ utilizes a local sample $\{\bm{\xi}_{k}\}$ to approximate the true gradient by using its loss value. Several distributed learning strategies for minimization problems have been investigated in [22, 33] and the references therein. Here, we integrate (6a)-(6b) with the adapt-then-combine (ATC) diffusion learning strategy from [22, 33]. This strategy has been shown in these references to outperform the consensus strategy, which has been applied to minimax problems as well in works such as [20, 35]. For this reason, we focus on diffusion learning and refer to the proposed algorithm as the Diffusion Stochastic Same-Sample Optimistic Gradient (DSS-OG). It is listed in Algorithm 1. In this algorithm, $a_{\ell k}$ is the scalar weight used by agent $k$ to scale information received from agent $\ell\in\mathcal{N}_{k}$ and $a_{\ell k}=0$ for $\ell\notin\mathcal{N}_{k}$. This algorithm starts from $i=0$ with randomly initialized local weights $\bm{x}_{k,i},\bm{y}_{k,i}(i=-1,-2)$. Each agent first collect two iid samples, namely $\bm{\xi}^{x}_{k,i}$ and $\bm{\xi}^{y}_{k,i}$, to approximate the OG direction for updating the local primal and dual variables, respectively. After then the local models of each agent are sent to its immediate neighbors for aggregation using the weights $\{a_{\ell k}\}$.

For convenience of the analysis, we introduce the following network quantities:

and similarly for ${\bm{y}}_{c,i},\bm{g}_{y,c,i}\in\mathbb{R}^{M_{2}\times 1},\bm{\mathcal{Y}}_{i},\bm{\mathcal{Y}}_{c,i},\bm{\mathcal{G}}_{y,i}\in\mathbb{R}^{M_{2}K\times 1}$. Additionally, we use ${{g}}_{x,c,i},{{g}}_{y,c,i},{\mathcal{G}}_{x,i},{\mathcal{G}}_{y,i}$ refer to deterministic realizations of the random quantities $\bm{g}_{x,c,i}$, $\bm{g}_{y,c,i}$, $\bm{\mathcal{G}}_{x,i}$, $\bm{\mathcal{G}}_{y,i}$. If we further introduce the network combination matrices:

then, the DSS-OG recursions be expressed globally the following recursion:

Since the primal problem is to minimize over $x$ in (2a), we introduce the primal objective[34, 35, 36, 37]:

The above assumptions can be relaxed by adopting more sophisticated distributed learning strategies, e.g., exact diffusion and gradient tracking [45, 46, 47]; we leave these extensions for future works.

This assumption ensures that $A$ has a single maximum eigenvalue at $1$ with Perron eigenvector $p$ with positive entries such that $Ap=p,\mathbbm{1}^{\top}p=1$. Furthermore, the combination matrix $A$ admits the Jordan decomposition $A=VJV^{-1}$ where [22]:

Here, the submatrix $J_{\gamma}$ consists of Jordan blocks with arbitrarily small values on the lower diagonal, and it holds that $\mathbbm{1}^{\top}V_{R}=0$, $V^{\top}_{L}V_{R}=I_{K-1}$, $\|J_{\gamma}\|<1$. For convenience of the analysis, we let $\mathcal{J}_{\gamma_{1}}\triangleq J_{\gamma}\otimes I_{M_{1}},\mathcal{J}_{\gamma_{2}}\triangleq J_{\gamma}\otimes I_{M_{2}},\mathcal{V}_{R_{1}}\triangleq V_{R}\otimes I_{M_{1}},\mathcal{V}_{R_{2}}\triangleq V_{R}\otimes I_{M_{2}},\mathcal{V}_{L_{1}}\triangleq V_{L}\otimes I_{M_{1}},\mathcal{V}_{L_{2}}\triangleq V_{L}\otimes I_{M_{2}}$. It is not difficult to verify $\|\mathcal{J}_{\gamma_{1}}\|=\|\mathcal{J}_{\gamma_{2}}\|,\|\mathcal{V}_{R_{1}}\|=\|\mathcal{V}_{R_{2}}\|,\|V_{L_{1}}\|=\|V_{L_{2}}\|$.

In Theorem 1, we prove the convergence rate of the best-iterate of the primal objective, which is the standard criterion for convergence analysis involving the nonconvex problem. In Theorem 2, we show the last-itearate convergence of the dual optimality gap. Given the conditions from these theorems, we establish the complexity of the DSS-OG in finding an $\varepsilon$-stationary point.

Theorem 1 states that by adopting a two-time-scale scheme for the step sizes $\mu_{x}$ and $\mu_{y}$, DSS-OG is able to output an $\epsilon$-stationary point for the primal objective with a sublinear rate dominated by $\mathcal{O}\Big{(}\frac{1}{T^{1/2}}\Big{)}$. The two-time-scale condition reflects the unsymmetric assumption of the risk function over the primal and dual variables. Similar conditions were also established in the works [7, 9, 8] when asymmetric conditions are assumed over the primal and dual variables.

The above theorem implies that finding a dual solution that shrinks the dual optimality gap has a faster sublinear convergence rate of $\mathcal{O}\Big{(}\frac{1}{T}\Big{)}$ than the primal convergence rate. Eqs. (18) and (22) also show the fundamental difference of the complexity in finding the solution for a nonconvex problem and for a PL problem. From this theoretical aspect, it is hard to guarantee a pair of primal and dual solutions with the same accuracy level (24a)–(24b) simultaneously. Fortunately, we can prove that the same accuracy level of the primal and dual solutions can be obtained in an asynchronous manner (See Theorem 3).

We point out the key differences between SS-OG and S-OG:

• •

  The stochastic gradients involved in S-OG are not martingale process. To see this, we observe that the backward gradients $[\nabla_{x}Q(\bm{x}_{i-1};\bm{y}_{i-1};\bm{\xi}_{x,i-1}),\nabla_{y}Q(\bm{x}_{i-1},\bm{y}_{i-1};\bm{\xi}_{y,i-1})]$ used at iteration $i$ are constructed by the past samples $(\bm{\xi}_{x,i-1},\bm{\xi}_{y,i-1})$. They are biased estimators when taking expectation over the current samples $(\bm{\xi}_{x,i},\bm{\xi}_{y,i})$. In contrast, it is easy to verify the stochastic gradients involved in SS-OG are unbiased.

• •

  SS-OG strictly performs the “optimistic” scheme under the same loss landscape. The losses utilized there are unbiased estimators for the local risk. In contrast, S-OG involves gradients from two different loss landscape. The gradient direction at the same point in two different loss landscapes can be notably different when data variance is high.

• •

  We will bound the terms involving $\mathbb{E}\|\nabla_{x}J(\bm{x}_{c,i},\bm{y}_{c,i})-{g}_{x,c,i}\|^{2}$ for DSS-OG rather than $\mathbb{E}\|\nabla_{x}J(\bm{x}_{c,i},\bm{y}_{c,i})-\bm{g}_{x,c,i}\|^{2}$ for distributed S-OG. To see this, we can repeat the arguments in (I) and use the fact that the two terms involved in the inner product $\langle\nabla{{P}}(\bm{x}_{c,i}),\bm{g}_{x,c,i}\rangle$ are correlated for distributed S-OG. Therefore, we observe that the distributed S-OG needs to bound the deviation between the deterministic and stochastic quantities. The consequence of such deviation is that a constant on the order of $\mathcal{O}(\sigma^{2})$ will appear in the final bound because the Lipschitz continuity to be used is only assumed for determistic quantities. In constrast, we only need to bound the deviation between deterministic quantities for DSS-OG. Therefore, employing a large batch size will become necessary for S-OG in order to control the magnitude of $\mathcal{O}(\sigma^{2})$.

In this example, we consider the application of an $\ell_{2}$-regularized Wasserstein GAN (WGAN) for learning the mean and variance of a one-dimensional Gaussian [8]. The objective is to optimize the following global risk function through the cooperation of agents:

Here, the dual variable $y$ is $\ell_{2}$-regularized by a constant $\lambda$ to ensure the PL condition. We consider a setting where the generator has a single hidden layer with 5 neurons, and the discriminator is parameterized by $y\in\mathbb{R}^{2}$ and is of the form $D(y;\bm{u})=y_{1}\bm{u}_{k}+y_{2}\bm{u}_{k}^{2}$. As for the true samples, $\bm{u}_{k}$, they are drawn from the distribution $\mathcal{N}(\pi_{k},\sigma^{2}_{k})$.

We generate a strongly connected network with $K=8$ agents and use the averaging rule [48, 33]. In addition to the fully distributed setting, we also implement centralized SS-OG (CSS-OG). We compare DSS-OG and CSS-OG with the distributed version of SAGDA, smoothed SAGDA and the adaptive algorithms, such as RMSprop and Adam which are also evaluated in this same application [8]. For these algorithms, we follow the optimal settings reported in [8]. For DSS-OG and CSS-OG, we tune the step sizes to $\mu_{x}=\mu_{y}=0.1$. We set the same random seed before running each algorithm to ensure the same generated training set. The code is available at the attached link¹¹1https://github.com/MulGarFed/Minimax_DSSOG.git.

The simulation results are demonstrated in Figs.LABEL:fig:example1_variance0001_grad-LABEL:fig:example1_variance01_mse. We observe both DSS-OG and CSS-OG outperform the other algorithms in terms of gradient norm and mean square error (MSE) when $\sigma^{2}_{k}=0.001$. From Fig. LABEL:fig:example1_variance01_mse, the MSE of DSS-OG and CSS-OG is able to match the results of RMSprop when $\sigma^{2}_{k}=0.1$. We find that the adaptive algorithm RMSprop converges more rapidly than constant step size algorithms. Another interesting observation is that the gradient norm does not necessarily imply the quality of a solution. For instance, despite RMSprop not having the lowest gradient norm in Fig. LABEL:fig:example1_variance0001_grad, it exhibits rapid convergence and good performance in terms of MSE, as seen in Fig. LABEL:fig:example1_variance0001_mse. This discrepancy suggests a promising research direction in identifying more practical metrics for evaluating minimax solutions.

In this example, we train a deep convolutional GANs (DCGANs) using the MNIST dataset. We consider a similar neural network structure described in [49].

We use the same network topology from the previous example. To enhance performance, we implement DSS-OG with a Adam optimizer through replacing its original stochastic gradient term with the SS-OG gradients from (7a)–(7b). For this experiment, we set the momentum factors to $\beta_{1}=0.2,\beta_{2}=0.999$ and vary the step sizes. We plot the results of the Fréchet Inception Distance (FID) score [50] in Fig. 2(a) The sample images generated by a single node are demonstrated in Fig. 2(b). It is seen that the simulated algorithms can attain good results of FID score with the appropriately chosen hyperparameters. Moreover, the single agent is capable of generating all classes of digits, demonstrating the generative diversity of a single node in the distributed network.