Section 4.5 Structured Optimization with Compositional Gradient

In this section, we extend the compositional optimization technique to address other structured optimization problems, including min-max optimization, min-min optimization, and bilevel optimization. These problems share a common structure in the form of a compositional gradient, denoted by \(\mathcal{M}(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))\), where \(\mathcal{M}\) is a mapping that is Lipschitz continuous with respect to its second argument, and \(\mathbf{u}^*(\mathbf{w})\) is defined as the solution to a strongly convex optimization problem:

\[\begin{align} \label{eqn:auxu} \mathbf{u}^*(\mathbf{w}) = \arg\min_{\mathbf{u} \in \mathcal{U}} h(\mathbf{w}, \mathbf{u}). \end{align}\]

This structure generalizes the gradient of a compositional function \(f(g(\mathbf{w}))\), whose gradient takes the form \(\mathcal{M}(\mathbf{w}, \mathbf{u}^*(\mathbf{w})) = \nabla g(\mathbf{w})\nabla f(\mathbf{u}^*(\mathbf{w}))\) with

\[\mathbf{u}^*(\mathbf{w}) = \arg\min_{\mathbf{u}} \|\mathbf{u} - g(\mathbf{w})\|_2^2.\]

The high-level idea underlying the algorithms and analysis presented below is summarized as follows. To estimate \(\mathcal{M}(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))\) at \(\mathbf{w}_t\), we use an auxiliary variable \(\mathbf{u}_t\) to track the optimal solution \(\mathbf{u}^*(\mathbf{w}_t)\), which is defined by solving (\(\ref{eqn:auxu}\)) with one step update at \(\mathbf{w}_t\). A key aspect of the analysis is that the error in the approximation of \(\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t)\) is controlled by the estimation error \(\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|_2\), due to the Lipschitz continuity of \(\mathcal{M}\):

\[\begin{align} \|\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t) - \mathcal{M}(\mathbf{w}_t, \mathbf{u}^*(\mathbf{w}_t))\|_2^2 \leq O(\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|_2^2). \end{align}\]

Moreover, since \(\mathbf{u}^*(\mathbf{w})\) is the solution to a strongly convex problem and hence is Lipschitz continuous with respect to \(\mathbf{w}\), we can construct a recursion for \(\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|_2^2\) to effectively bound the cumulative error over iterations.

In cases where \(\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t)\) cannot be computed exactly and is instead approximated by a stochastic estimator \(\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t; \zeta_t)\), where \(\zeta_t\) is a random variable, we employ a moving average (MA) estimator:

\[\begin{align*} \mathbf{v}_{t} = (1 - \beta_t) \mathbf{v}_{t-1} + \beta_t\, \mathcal{M}(\mathbf{w}_t, \mathbf{u}_t; \zeta_t). \end{align*}\]

The model update is then performed using:

\[\begin{align*} \mathbf{w}_{t+1} = \mathbf{w}_t - \eta_t \mathbf{v}_{t}. \end{align*}\]

Alternatively, if \(\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t)\) is directly computable, the update simplifies to:

\[\begin{align*} \mathbf{w}_{t+1} = \mathbf{w}_t - \eta_t\, \mathcal{M}(\mathbf{w}_t, \mathbf{u}_t). \end{align*}\]

4.5.1 Non-convex Min-Max Optimization

We consider a non-convex min-max optimization problem:

\[\begin{align} \label{eqn:mm} \min_{\mathbf{w}\in\mathbb{R}^d}\max_{\mathbf{u}\in\mathcal U} f(\mathbf{w}, \mathbf{u}) := \mathbb{E}_{\xi}[f(\mathbf{w}, \mathbf{u}; \xi)], \end{align}\]

where \(f(\mathbf{w},\mathbf{u})\) is continuous and differentiable and \(\mathcal U\) is a closed convex set. Let \(F(\mathbf{w}) = \max_{\mathbf{u}\in\mathcal U} f(\mathbf{w}, \mathbf{u})\). Denote by \(\nabla_1f(\cdot,\cdot)\) and \(\nabla_2 f(\cdot, \cdot)\) the partial gradients of the first and second variable, respectively.

4.5.1.1 A Double-loop Large mini-batch method

Let us first consider a straightforward approach that updates \(\mathbf{w}_t\) using a large-batch gradient estimator

\[ \mathbf{v}_t = \frac{1}{B} \sum_{i=1}^B \nabla_1 f(\mathbf{w}_t, \mathbf{u}_t; \zeta_{i,t}), \]

and computes \(\mathbf{u}_t\) via an inner-loop SGD with \(K\) updates. It suffices to have \(K = O(L_1^2\sigma_2^2 / (\mu^2 \epsilon^2))\) (by Corollary 3.9) such that

\[\begin{align*} &\mathbb{E}[\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|_2^2] \leq \frac{\epsilon^2}{L_1^2}. \end{align*}\]

If \(B = O(\sigma_1^2/\epsilon^2)\), following Lemma 4.18 below we have

\[\begin{align*} &\mathbb{E}[\|{\mathbf{v}}_t - \nabla F(\mathbf{w}_t)\|^2_2]\leq \mathbb{E}\bigg[\bigg\|\frac{1}{B} \sum_{i=1}^B \nabla_1 f(\mathbf{w}_t, \mathbf{u}_t; \zeta_{i,t}) - \nabla_1 f(\mathbf{w}_t, \mathbf{u}^*(\mathbf{w}_t))\bigg\|_2^2\bigg] \\ &\leq O\left(\frac{\sigma_1^2}{B} + L_1^2\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|_2^2\right)\leq \epsilon^2. \end{align*}\]

Combining this with Lemma 4.9, we can set the step size \(\eta_t = O(1/L_F)\) and the number of iterations \(T = O(L_F/\epsilon^2)\), yielding an overall sample complexity of

\[BT + KT = O\left(\frac{L_F\sigma_1^2}{\epsilon^4} + \frac{L_FL_1^2\sigma_2^2}{\mu^2\epsilon^4}\right).\]

4.5.1.2 A Stochastic Momentum Method

We present a solution method in Algorithm 12, referred to as SMDA (Stochastic Momentum Descent-Ascent). The method begins by updating the dual variable using stochastic gradient ascent (Step 4), then computes the moving average gradient estimator \(\mathbf{v}_{t}\) for the primal variable (Step 6), and finally updates the primal variable using this estimator (Step 7). When \(\beta_t = 1\), the method reduces to Algorithm 7. However, setting \(\beta_t < 1\) is crucial for achieving improved complexity. Conceptually, the method shares similarities with SCMA.

Convergence Analysis

We will prove the convergence of the gradient norm of \(F(\mathbf{w})\). We first prove the following lemmas.

Proof

■

Proof

■

Next, we prove two lemmas similar to Lemma 4.8 and Lemma 4.1, regarding the recursion of gradient estimation error and the estimation error of \(\mathbf{u}\), respectively. The descent lemma (Lemma 4.9) still holds.

Proof

■

Proof

■

Finally, we can prove the following theorem regarding the convergence of SMDA.

Proof

■

4.5.2 Non-convex Min-Min Optimization

We can extend SMDA to solving a non-convex strongly-convex min-min problem: \[\begin{align}\label{eqn:minmin} \min_{\mathbf{w}\in\mathbb{R}^d}\min_{\mathbf{u}\in\mathcal U} f(\mathbf{w}, \mathbf{u}): = \mathbb{E}_{\xi}[f(\mathbf{w}, \mathbf{u}; \xi)], \end{align}\] where \(f(\mathbf{w},\mathbf{u})\) is smooth, non-convex in terms of \(\mathbf{w}\) and strongly convex in terms of \(\mathbf{u}\) and \(\mathcal U\) is a closed convex set. If the \(\mathbf{u}^*(\mathbf{w})=\arg\min_{\mathbf{u}\in\mathcal U} f(\mathbf{w}, \mathbf{u})\) exists and unique, then we have \(\nabla F(\mathbf{w}) = \nabla_1 f(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))\). Hence, its gradient also exhibits a compositional structure, where the inner function \(\mathbf{u}^*(\mathbf{w})\) is a solution to a strongly convex problem.

SMDA can be modified by replacing the \(\mathbf{u}\) update with \[\mathbf{u}_{t+1} = \Pi_{\mathcal U}[\mathbf{u}_{t} - \gamma_t \nabla_2 f(\mathbf{w}_t, \mathbf{u}_t; \zeta_t)].\] Then, the same convergence result in the last subsection can be established for min-min problem, which is omitted here.

4.5.2.1 Application to weakly convex minimization

Next, we present an application to solving weakly convex minimization problems: \[\begin{align} \min_{\mathbf{w}} F(\mathbf{w}) := \mathbb{E}[g(\mathbf{w}; \zeta)], \end{align}\] where \(F(\mathbf w) > -\infty\) is \(\rho\)-weakly convex, as discussed in Chapter ch:2.

As argued in Section 3.1.4, an \(\epsilon\)-stationary solution of the Moreau envelope of \(F(\mathbf{w})\) corresponds to a nearly \(\epsilon\)-stationary solution of the original problem. Hence, we consider optimizing the Moreau envelope directly: \[\begin{align}\label{eqn:minmin-weak} \min_{\mathbf{w}} F_{\rho}(\mathbf{w}) := \min_{\mathbf{u}} \mathbb{E}[g(\mathbf{u}; \zeta)] + \rho \|\mathbf{u} - \mathbf{w}\|_2^2. \end{align}\] Define \(f(\mathbf{w}, \mathbf{u}) = \mathbb{E}[g(\mathbf{u}; \zeta)] + \rho \|\mathbf{u} - \mathbf{w}\|_2^2\). Then \(f(\mathbf{w}, \mathbf{u})\) is \(\rho\)-strongly convex with respect to \(\mathbf{u}\) due to the \(\rho\)-weak convexity of \(F\).

For updating \(\mathbf{u}\), we use the standard SGD: \[\begin{align} \mathbf{u}_{t+1} = \mathbf{u}_t - \gamma_t(\mathcal{G}(\mathbf{u}_t; \zeta_t) + 2\rho (\mathbf{u}_t - \mathbf{w}_t)). \end{align}\] where \(\mathcal{G}(\mathbf{u}_t; \zeta_t)\in\partial g(\mathbf{u}_t; \zeta_t)\). For updating \(\mathbf{w}\), then we just apply GD with its gradient given by \(\nabla_1 f(\mathbf{w}_t,\mathbf{u}_t) = 2\rho(\mathbf{w}_t - \mathbf{u}_t)\): \[\begin{align} \mathbf{w}_{t+1} = \mathbf{w}_t - \eta_t 2\rho (\mathbf{w}_t - \mathbf{u}_t) = (1-2\eta_t\rho)\mathbf{w}_t + 2\eta_t\rho\mathbf{u}_t. \end{align}\] We present the updates in Algorithm 13. An interesting observation about this algorithm is that the \(\mathbf{u}\) update is similar to the Momentum update except that the momentum term \(\mathbf{u}_t - \mathbf{u}_{t-1}\) is replaced by \(\mathbf{u}_t - \mathbf{w}_t\), where \(\mathbf{w}_t\) is a MA weight vector.

Convergence Analysis

Let us first prove the following lemma.

Proof

The smoothness of \(F_\rho\) has been proved in Proposition 3.1 with \(\lambda=\rho/2\). The Lipschitz continuity of \(\nabla_1 f(\mathbf{w}, \mathbf{u})= 2\rho(\mathbf{w} - \mathbf{u})\) is obvious. Next, let us prove the Lipschitz continuity of \(\mathbf{u}^*(\mathbf{w})\). The proof is similar to that of Lemma 4.17.

■

Since \(\partial_2f(\mathbf{w},\mathbf{u})\) is not Lipschitz continuous with respect to \(\mathbf{u}\), Lemma 4.20 is not directly applicable. We develop a similar one below.

Proof

■

Since \(\nabla F_\rho(\mathbf{w}_t) = 2\rho(\mathbf{w}_t - \mathbf{u}^*(\mathbf{w}_t))\), hence \(\|\nabla F_\rho(\mathbf{w}_t) -\mathbf{z}_t\|_2^2 = 4\rho^2\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|_2^2\), whose recursion has been established in Lemma 4.22. We can combine these two lemmas and establish a complexity of \(O(1/\epsilon^4)\) for Algorithm 13 in order to find an \(\epsilon\)-stationary solution to \(F_\rho(\cdot)\).

4.5.2.2 Application to weakly-convex strongly-concave min-max problems

The same technique can be applied to solving weakly-convex strongly-concave min-max problems \(\min_{\mathbf{w}}\max_{\mathbf{u}\in\mathcal U}f(\mathbf{w}, \mathbf{u})\) with a single loop algorithm. In subsection 4.5.1, we assume the partial gradient \(\nabla_1 f(\mathbf{w}, \mathbf{u})\) is Lipschitz continuous. We replace this assumption by an assumption that \(f(\mathbf{w}, \mathbf{u})\) is \(\rho\)-weakly convex in terms of \(\mathbf{w}\) for any \(\mathbf{u}\in\mathcal U\).

In this case, \(F(\mathbf{w})=\max_{\mathbf{u}\in\mathcal U}f(\mathbf{w}, \mathbf{u})\) is not smooth but weakly convex. Let us consider its Moreau envelope: \[\begin{align*} \min_{\mathbf{w}} F_\rho(\mathbf{w}): = \min_{\mathbf{u}_1} F(\mathbf{u}_1) + \rho \|\mathbf{u}_1 - \mathbf{w}\|_2^2. \end{align*}\] This problem is equivalent to \[\begin{align*} \min_{\mathbf{w}, \mathbf{u}_1} \max_{\mathbf{u}_2\in\mathcal U} f(\mathbf{u}_1, \mathbf{u}_2) + \rho \|\mathbf{u}_1 - \mathbf{w}\|_2^2, \end{align*}\] which is strongly convex in terms of \(\mathbf{u}_1\) and strongly concave in terms of \(\mathbf{u}_2\).

Compared to (\(\ref{eqn:minmin-weak}\)), this problem just adds another layer of inner maximization. However, it can be still mapped to the general framework as discussed at the beginning. The gradient of \(F_\rho(\mathbf{w})\) is given by \(\mathcal{M}(\mathbf{w}, \mathbf{u}_1^*(\mathbf{w}))=\rho (\mathbf{w} - \mathbf{u}_1^*(\mathbf{w}))\). If we track \(\mathbf{u}_1^*(\mathbf{w}_t)\) by \(\mathbf{u}_{1,t}\) and its update relies on the gradient \(\partial_1 f(\mathbf{u}_{1,t}, \mathbf{u}_{2}^*(\mathbf{u}_{1,t}))\). Hence, we just need another variable \(\mathbf{u}_{2,t}\) to track \(\mathbf{u}_2^*(\mathbf{u}_{1,t})\).

We can develop a similar algorithm. First, let us update \(\mathbf{u}_1,\mathbf{u}_2\). Given \(\mathbf{w}_t, \mathbf{u}_{1,t}, \mathbf{u}_{2,t}\), we update \(\mathbf{u}_{1,t+1}, \mathbf{u}_{2,t+1}\) with SGD update by \[\begin{align} \mathbf{u}_{2,t+1} & = \Pi_{\mathcal U}[\mathbf{u}_{2,t} + \gamma_2 \partial_2 f(\mathbf{u}_{1,t}, \mathbf{u}_{2,t}; \zeta_t)]\\ \mathbf{u}_{1,t+1} &= \mathbf{u}_{1,t} - \gamma_1(\partial_1f(\mathbf{u}_{1,t}, \mathbf{u}_{2,t}; \zeta_t)+2\rho (\mathbf{u}_{1,t} - \mathbf{w}_t)). \end{align}\] Then we update \(\mathbf{w}_{t+1}\) with GD update by \[\begin{align} \mathbf{w}_{t+1} = \mathbf{w}_t - \eta 2\rho(\mathbf{w}_t - \mathbf{u}_{1,t}) = (1-2\eta\rho) \mathbf{w}_t + 2\eta\rho\mathbf{u}_{1,t}. \end{align}\] This algorithm also enjoys a complexity of \(O(1/\epsilon^4)\) for finding a nearly \(\epsilon\)-stationary solution of \(F(\mathbf{w})\). We refer the readers to Hu et al. (2024) for a convergence analysis of this algorithm.

4.5.2.3 Application to Compositional Optimization

We can apply a similar strategy to a compositional function \(F(\mathbf{w})=f_0(g(\mathbf{w}))\), where \(f_0\) is smooth convex and \(g\) is weakly convex. With the conjugate of \(f_0\), we can write \[\min_{\mathbf{w}}f_0(g(\mathbf{w})) = \min_{\mathbf{w}}\max_{\mathbf{u}_2\in\mathcal U}f(\mathbf{w},\mathbf{u}_2) := \mathbf{u}_2^{\top} g(\mathbf{w}) - f_0^*(\mathbf{u}_2).\] Since \(f_0\) is smooth, then \(f_0^*\) is strongly convex. Then if \(g\) is weakly convex and \(\mathcal U\) is bounded (i.e., \(f_0\) is Lipschitz), then \(f(\mathbf{w},\mathbf{u})\) is weakly convex and strongly concave. Optimizing the Moreau envelope of \(f_0(g(\mathbf{w}))\) yields: \[\min_{\mathbf{w}, \mathbf{u}_1}\max_{\mathbf{u}_2\in\mathcal U}\mathbf{u}_2^{\top} g(\mathbf{u}_1) - f_0^*(\mathbf{u}_2) + \rho\|\mathbf{u}_1 - \mathbf{w}\|_2^2,\] which is strongly convex in terms of \(\mathbf{u}_1\) and strongly concave in terms of \(\mathbf{u}_2\). We give an update below: \[\begin{align*} \mathbf{u}_{2,t+1} & = \Pi_{\mathcal U}[\mathbf{u}_{2,t} + \gamma_2 g(\mathbf{u}_{1,t}; \zeta_t)]\\ \mathbf{u}_{1,t+1} &= \mathbf{u}_{1,t} - \gamma_1(\partial_1 g(\mathbf{u}_{1,t}; \zeta_t)\mathbf{u}_{2,t}+2\rho (\mathbf{u}_{1,t} - \mathbf{w}_t))\\ \mathbf{w}_{t+1} &= \mathbf{w}_t - \eta 2\rho(\mathbf{w}_t - \mathbf{u}_{1,t}) = (1-2\eta\rho) \mathbf{w}_t + 2\eta\rho\mathbf{u}_{1,t}. \end{align*}\] Then similar convergence analysis can be developed with a complexity of \(O(1/\epsilon^4)\) for finding a nearly \(\epsilon\)-stationary solution to \(F\).

4.5.3 Non-convex Bilevel Optimization

In this section, we discuss the application of the compositional gradient estimation technique to non-convex bilevel optimization defined by

\[\begin{equation}\label{eqn:bi} \begin{aligned} \min_{\mathbf{w}\in\mathbb{R}^d}& f(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))\\ \mathbf{u}^*(\mathbf{w})&=\arg\min_{\mathbf{u}\in\mathbb{R}^{d'}} g(\mathbf{w}, \mathbf{u}), \end{aligned} \end{equation}\]

where \(g\) is twice differentiable and \(\mu_g\)-strongly convex in terms of \(\mathbf{u}\). Let \(F(\mathbf{w}) =f(\mathbf{w}, \mathbf{u}^*(\mathbf{w})).\) The following lemma states the gradient of the objective \(F(\mathbf{w})\).

■

Let us define \[\begin{align*} &\mathcal{M}(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))=\nabla_1 f(\mathbf{w}, \mathbf{u}^*(\mathbf{w})) - \nabla_{21} g(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))^{\top}(\nabla_{22} g(\mathbf{w}, \mathbf{u}^*(\mathbf{w})))^{-1}\nabla_2 f(\mathbf{w}, \mathbf{u}^*(\mathbf{w})). \end{align*}\]

If we can establish the Lipschitz continuity of \(\mathcal{M}(\mathbf{w}, \mathbf{u}^*(\mathbf{w}))\) in terms of the second argument and the Lipschitz continuity of \(\mathbf{u}^*(\mathbf{w})\), then the similar technique can be leveraged. Let \(\mathbf{u}^*(\mathbf{w}_t)\) be tracked by \(\mathbf{u}_t\). It can be updated by SGD: \[\begin{align} \mathbf{u}_{t+1} = \mathbf{u}_t - \gamma_t \nabla_2 g(\mathbf{w}_t, \mathbf{u}_t; \zeta_t). \end{align}\]

With \(\mathbf{u}_t\), the gradient at \(\mathbf{w}_t\) can be estimated by \[\begin{align} \mathcal{M}(\mathbf{w}_t, \mathbf{u}_t)=\nabla_1 f(\mathbf{w}_t, \mathbf{u}_t) +\nabla_{21} g(\mathbf{w}_t, \mathbf{u}_t)^{\top}(\nabla_{22} g(\mathbf{w}_t, \mathbf{u}_t))^{-1}\nabla_2 f(\mathbf{w}_t, \mathbf{u}_t). \end{align}\] However, another challenge is to handle the Hessian inverse \((\nabla_{22} g(\mathbf{w}_t, \mathbf{u}_t)^{-1}\), which itself is a compositional structure. We will discuss three different ways to tackle this challenge. If we have a stochastic estimator of \(\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t)\) denoted by \(\mathbf{v}_{t}\), then we update the model parameter by: \[\begin{align} \mathbf{w}_{t+1} = \mathbf{w}_t - \eta_t \mathbf{v}_{t}. \end{align}\]

4.5.3.1 Approach 1: The MA Estimator

If the lower level problem is low-dimensional such that the inverse of the Hessian matrix can be efficiently computed, we can estimate \(\nabla_{22} g(\mathbf{w}_t, \mathbf{u}_t)\) by a MA estimator: \[\begin{align*} H_{22,t} = S_{\mu_g}[(1-\beta) H_{22,t-1} + \beta \nabla_{22} g(\mathbf{w}_t, \mathbf{u}_t; \zeta_{2,t})]. \end{align*}\] where \(S_{\mu_g}[\cdot]\) is a projection operator that projects a matrix into a matrix whose minimum eigen-value is lower bounded by \(\mu_g\), where \(\mu_g\) is the lower bound of eigen-values of \(\nabla_{22} g(\mathbf{w}, \mathbf{u})\). The projection ensures that \([H_{22,t}]^{-1}\) is Lipschitz continuous with respect to \(H_{22,t}\).

The a vanilla stochastic gradient estimator of \(\mathbf{w}_t\) and its MA estimator are computed by

\[\begin{equation}\label{eqn:bio-hma} \begin{aligned} &\mathbf{z}_{t} = \nabla_1 f(\mathbf{w}_t, \mathbf{u}_t; \xi_{t}) +\nabla_{21} g(\mathbf{w}_t, \mathbf{u}_t; \zeta'_{2,t})^{\top}(H_{22,t})^{-1}\nabla_2 f(\mathbf{w}_t, \mathbf{u}_t; \xi_{t})\\ &\mathbf{v}_{t} =(1-\beta)\mathbf{v}_{t-1} + \beta \mathbf{z}_t. \end{aligned} \end{equation}\]

Convergence Analysis

The proof is largely similar to that of Theorem 4.3. We provide a sketch of proof below. Recall that \[\begin{align*} &\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t)=\nabla_1 f(\mathbf{w}_t, \mathbf{u}_t) +\nabla_{21} g(\mathbf{w}_t, \mathbf{u}_t)^{\top}(\nabla_{22} g(\mathbf{w}_t, \mathbf{u}_t))^{-1}\nabla_2 f(\mathbf{w}_t, \mathbf{u}_t). \end{align*}\] Define: \[\begin{align*} \hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t)=\nabla_1 f(\mathbf{w}_t, \mathbf{u}_t) +\nabla_{21} g(\mathbf{w}_t, \mathbf{u}_t)^{\top} H_{22,t}^{-1}\nabla_2 f(\mathbf{w}_t, \mathbf{u}_t). \end{align*}\]

First, similar to Lemma 4.9, we have the following: \[\begin{align} F(\mathbf{w}_{t+1})& \leq F(\mathbf{w}_t) + \frac{\eta_t}{2} \|\mathbf{v}_t - \nabla F(\mathbf{w}_t)\|_2^2-\frac{\eta_t}{2}\|\nabla F(\mathbf{w}_t)\|_2^2- \frac{1}{4\eta_t} \|\mathbf{w}_{t+1} - \mathbf{w}_t\|_2^2. \end{align}\]

We establish a recursion of the error \(\|\mathbf{v}_t - \nabla F(\mathbf{w}_t)\|_2^2\) similar to Lemma 4.7 by noting that \(\mathbb E_{\xi_t, \zeta'_{2,t}}[\mathbf{z}_t]=\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t)\) and there exists \(\sigma>0\) such that \(\mathbb E_{\xi_t, \zeta'_{2,t}}[\|\mathbf{z}_t - \hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t)\|_2^2]\leq \sigma^2\). Thus, Lemma 4.7 implies that

\[\begin{align}\label{eqn:bio-v-err} &\mathbb E_{\xi_t, \zeta'_{2,t}}\left[\|\mathbf{v}_{t} - \nabla F(\mathbf{w}_t)\|_2^2\right] \leq (1-\beta_t)\|\mathbf{v}_{t-1} - \nabla F(\mathbf{w}_{t-1})\|_2^2 \\ & \quad + \frac{2L_F^2}{\beta_t}\|\mathbf{w}_{t-1}-\mathbf{w}_t\|_2^2 + 4\beta_t\|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t) - \nabla F(\mathbf{w}_t)\|_2^2 + \beta_t^2\sigma^2.\notag \end{align}\]

Then, we bound \(\|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t)-\nabla F(\mathbf{w}_t)\|_2^2\) by \[\begin{align*} \|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t)-\nabla F(\mathbf{w}_t)\|_2^2&\leq 2\|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t) -\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t) \|^2_2\\ & +2\|\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t) - \nabla F(\mathbf{w}_t)\|^2_2\\ & \leq O(\|H_{22,t} - \nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\|_2^2) + O(\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|^2_2). \end{align*}\]

As a result, we have \[\begin{align*} \mathbb E[\|\mathbf{v}_t - &\nabla F(\mathbf{w}_t)\|_2^2] \leq (1-\beta_t)\|\mathbf{v}_{t-1}- \nabla F(\mathbf{w}_{t-1})\|^2 + \frac{2L_F^2}{\beta_t}\|\mathbf{w}_t - \mathbf{w}_{t-1}\|_2^2\\ &+ \beta_t (O(\|H_{22,t} - \nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\|_2^2) + O(\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|^2_2)) + \beta_t^2O(\sigma^2). \end{align*}\] This result is similar to that in Lemma 4.8.

We can further build the error recursion of \(\|H_{22,t} - \nabla_{22}g(\mathbf{w}_t,\mathbf{u}_t)\|_2^2\) similar to Lemma 4.1, and the error recursion of \(\|\mathbf{u}_t - \mathbf{u}^*(\mathbf{w}_t)\|^2_2\) similar to Lemma 4.20. Combining these results, we can establish a complexity of \(O(1/\epsilon^4)\) for finding an \(\epsilon\)-stationary solution of \(F(\cdot)\) in expectation.

4.5.3.2 Approach 2: The Neumann Series (Matrix Taylor Approximation)

If the lower level problem is high-dimensional such that it is prohibited to compute the Hessian, one approach is to leverage the Neuman series: \[\begin{align} A^{-1}= \sum_{i=0}^\infty (I -A)^i, \quad \text{ if }\quad \|A\|\leq 1. \end{align}\] Hence, if \(\|\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\|\leq L_{22}\), we estimate the inverse of \(\frac{1}{L_{22}}\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\), yielding \[\begin{align} \left(\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\right)^{-1}\approx \frac{1}{L_{22}}\sum_{i=0}^{K-1} \left(I - \frac{1}{L_{22}}\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\right)^i. \end{align}\] This can be further estimated by a stochastic route, by sampling \(k\) from \(\{0,\ldots,K- 1\}\) randomly, then estimate the Hessian inverse by \[\begin{align} Q_{22,t} = \left\{\begin{array}{ll}\frac{K}{L_{22}}\prod_{i=1}^{k} \left(I - \frac{1}{L_{22}}\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t; \zeta_{i})\right)& \text{ if } k\geq 1\\ \frac{K}{L_{22}}I & \text{ if } k=0\end{array}\right.. \end{align}\] This is can be justified by \[\begin{align*} \mathbb{E}[Q_{22,t}] &= \frac{1}{K}\frac{K}{L_{22}}I + \frac{K-1}{K}\mathbb{E}_{k\sim\{1,\ldots, K-1\}}\left[\frac{K}{L_{22}}\prod_{i=1}^{k} \left(I - \frac{1}{L_{22}}\mathbb{E}[\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t; \zeta_{i})]\right)\right]\\ & = \mathbb{E}_{k}\frac{K}{L_{22}}\left(I - \frac{1}{L_{2}}\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\right)^k = \sum_{k=0}^{K-1}\frac{1}{L_{22}}\left(I -\frac{1}{L_{22}}\nabla_{22}g(\mathbf{w}_t, \mathbf{u}_t)\right)^k. \end{align*}\]

Then the vanilla gradient estimator of \(\mathbf{w}_t\) and its MA estimator are computed by

\[\begin{equation}\label{eqn:bio-neu} \begin{aligned} &\mathbf{z}_{t} = \nabla_1 f(\mathbf{w}_t, \mathbf{u}_t; \zeta_{1,t}) +\nabla_{21} g(\mathbf{w}_t, \mathbf{u}_t; \zeta'_{2,t})^{\top}Q_{22,t}\nabla_2 f(\mathbf{w}_t, \mathbf{u}_t; \zeta_{1,t})\\ &\mathbf{v}_{t} =(1-\beta)\mathbf{v}_{t-1} + \beta \mathbf{z}_t. \end{aligned} \end{equation}\]

Convergence Analysis

We provide a proof sketch below. We can understand that \(\mathbf{z}_t\) is an unbiased stochastic estimator of \[\begin{align*} \hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t)=\nabla_1 f(\mathbf{w}_t, \mathbf{u}_t) +\nabla_{21} g(\mathbf{w}_t, \mathbf{u}_t)^{\top}\mathbb{E}[Q_{22,t}]\nabla_2 f(\mathbf{w}_t, \mathbf{u}_t). \end{align*}\] We decompose the estimation error of \(\mathbf{v}_t\) similarly as in (\(\ref{eqn:bio-v-err}\)) and bound \(\|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t) - \nabla F(\mathbf{w}_t)\|_2^2\) by \[\begin{align*} \|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t) - \nabla F(\mathbf{w}_t)\|_2^2 &\leq 2\|\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t) -\nabla F(\mathbf{w}_t)\|_2^2 \\ & + 2\|\hat{\mathcal{M}}(\mathbf{w}_t, \mathbf{u}_t) - \mathcal{M}(\mathbf{w}_t,\mathbf{u}_t)\|_2^2 . \end{align*}\] The error recursion of the first term on the right hand side can be similarly bounded as before. To bound the last error, since \[\begin{equation*} \left[\nabla_{22}^2 g(\mathbf{w}, \mathbf{u})\right]^{-1} = \mathbb{E}[Q_{22}] + \frac{1}{L_{22}} \sum_{i=K}^{\infty} \left[I - \frac{1}{L_{22}} \nabla_{22}^2 g(\mathbf{w}, \mathbf{u})\right]^i, \end{equation*}\] we have \[\begin{align*} &\|\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t) - \hat{\mathcal{M}}(\mathbf{w}_t,\mathbf{u}_t)\|_2^2\leq O(\|\left[\nabla_{22}^2 g(\mathbf{w}, \mathbf{u})\right]^{-1} - \mathbb{E}[Q_{22}]\|_2^2),\\ &\left\| \left[\nabla_{22}^2 g(\mathbf{w}, \mathbf{u})\right]^{-1} - \mathbb{E}[Q_{22}] \right\|_2 \le \frac{1}{L_{22}} \sum_{i=K}^{\infty} \left\|I - \frac{1}{L_{22}} \nabla_{22}^2 g(\mathbf{w}, \mathbf{u})\right\|_2^i \le \frac{1}{\mu_g} \left(1 - \frac{\mu_g}{L_{22}}\right)^K. \end{align*}\] As a result, if \(K=O(\frac{L_{22}}{\mu_g}\log(1/(\mu_g\beta_t\sigma^2)))\), then \(\|\mathcal{M}(\mathbf{w}_t, \mathbf{u}_t) - \mathcal{M}'(\mathbf{w}_t,\mathbf{u}_t)\|_2^2\leq O(\beta_t\sigma^2)\). Then similar to the analysis of Approach 1, we can establish a complexity of \(O(1/\epsilon^4)\) for finding an \(\epsilon\)-stationary solution of \(F(\cdot)\) in expectation.

4.5.3.3 Approach 3: The penalty method

An alternative approach to avoid computing the Hessian inverse and Jacobian matrices is to reformulate the problem as a constrained optimization problem: \[\begin{align*} \min_{\mathbf{w},\mathbf{u}} &\quad f(\mathbf{w}, \mathbf{u}) \\ \text{s.t.} &\quad g(\mathbf{w}, \mathbf{u}) \leq \min_{\mathbf{u}} g(\mathbf{w}, \mathbf{u}). \end{align*}\] This constrained problem can be addressed using a penalty method (see Section 6.7): \[\begin{align*} \min_{\mathbf{w},\mathbf{u}} \; f(\mathbf{w}, \mathbf{u}) + \lambda \big( g(\mathbf{w}, \mathbf{u}) - \min_{\mathbf{y}} g(\mathbf{w}, \mathbf{y}) \big)_+, \end{align*}\] where \(\lambda > 0\) is a penalty parameter and \((\cdot)_+\) denotes the positive part. Since \(g(\mathbf{w}, \mathbf{u}) \geq \min_{\mathbf{y}} g(\mathbf{w}, \mathbf{y})\), the formulation simplifies to:

\[\begin{align}\label{eqn:bio-minmax} &\min_{\mathbf{w},\mathbf{u}} \; f(\mathbf{w}, \mathbf{u}) + \lambda \left( g(\mathbf{w}, \mathbf{u}) - \min_{\mathbf{y}} g(\mathbf{w}, \mathbf{y}) \right)\\ &= \min_{\mathbf{w},\mathbf{u}} \max_{\mathbf{y}} f(\mathbf{w}, \mathbf{u}) + \lambda \left( g(\mathbf{w}, \mathbf{u}) - g(\mathbf{w}, \mathbf{y}) \right). \end{align}\]

If both \(f\) and \(g\) are smooth and \(g\) is strongly convex in its second argument, the resulting formulation becomes a non-convex strongly-concave min-max problem, which can be effectively addressed using the SMDA algorithm with the following update for \(t\geq 1\):

\[\begin{equation}\label{eqn:bi-smda} \begin{aligned} &\mathbf{y}_{t+1} = \mathbf{y}_{t} + \gamma_t\lambda\nabla_2 g(\mathbf{w}_t, \mathbf{y}_{t}; \xi_t),\\ &\mathbf{z}_t = \nabla f(\mathbf{w}_t, \mathbf{u}_{t}; \zeta_t) + \lambda \left(\nabla g(\mathbf{w}_t, \mathbf{u}_{t}; \xi_t) - \left[\nabla_1 g(\mathbf{w}_t, \mathbf{y}_{t}; \xi_t)\atop 0\right]\right),\\ &\mathbf{v}_{t} = (1-\beta_t)\mathbf{v}_{t-1} + \beta_t \mathbf{z}_t, \\ &\left[\mathbf{w}_{t+1}\atop \mathbf{u}_{t+1}\right] = \left[\mathbf{w}_t\atop \mathbf{u}_{t}\right] - \eta_t \mathbf{v}_{t}. \end{aligned} \end{equation}\]

Convergence Analysis

The convergence analysis of (\(\ref{eqn:bi-smda}\)) for the min–max problem (\(\ref{eqn:bio-minmax}\)) follows a similar approach to that of Theorem 4.5 for SMDA. However, a remaining challenge lies in converting the convergence result for the min–max formulation into that of the original problem. To address this, we provide the detailed convergence analysis below. We begin by stating the following assumption.

Let us define \(\bar{\mathbf{w}} : = (\mathbf{w}, \mathbf{u})\) and

\[\begin{align} &\bar f(\bar{\mathbf{w}}, \mathbf{y}): = f(\mathbf{w}, \mathbf{u}) + \lambda \left( g(\mathbf{w}, \mathbf{u}) - g(\mathbf{w}, \mathbf{y}) \right),\\ &\bar F(\bar{\mathbf{w}}) : = \max_{\mathbf{y}}\bar f(\bar{\mathbf{w}}, \mathbf{y}). \end{align}\]

Then \[\begin{align*} &\nabla_1\bar f(\bar{\mathbf{w}}, \mathbf{y})=\nabla f(\mathbf{w}, \mathbf{u}) + \lambda \left(\nabla g(\mathbf{w}, \mathbf{u}) - \left[\nabla_1 g(\mathbf{w}, \mathbf{y})\atop 0\right]\right),\\ &\nabla_2\bar f(\bar{\mathbf{w}}, \mathbf{y})= -\lambda \nabla_2 g(\mathbf{w}, \mathbf{y}),\\ &\nabla_1\bar f(\bar{\mathbf{w}}, \mathbf{y};\varepsilon)=\nabla f(\mathbf{w}, \mathbf{u}; \zeta) + \lambda \left(\nabla g(\mathbf{w}, \mathbf{u};\xi) - \left[\nabla_1 g(\mathbf{w}, \mathbf{y}; \xi)\atop 0\right]\right),\\ &\nabla_2\bar f(\bar{\mathbf{w}}, \mathbf{y};\xi)= -\lambda \nabla_2 g(\mathbf{w}, \mathbf{y};\xi). \end{align*}\] where \(\varepsilon=(\zeta, \xi)\). We first show \(\bar f(\bar{\mathbf{w}}, \mathbf{y})\) satisfies the conditions in Assumption 4.8).

■

■

Convergence of the original function

Next, we derive the convergence of the original function in terms of \(\|\nabla F(\mathbf{w})\|_2\). We need the following additional assumption.

Proof.
Let \(\mathbf{u}^* = \mathbf{u}^*(\mathbf{w})\). Then, \[\begin{align*} \nabla_{1} \bar{F}(\mathbf{w}, \mathbf{u}) &= \nabla_{1} f(\mathbf{w}, \mathbf{u}) + \lambda (\nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*)) \\ \nabla_{2} \bar{F}(\mathbf{w}, \mathbf{u}) &= \nabla_{2} f(\mathbf{w}, \mathbf{u}) + \lambda \nabla_{2} g(\mathbf{w}, \mathbf{u}). \end{align*}\] Due to Lemma 4.24, we have

\[\begin{equation} \label{eq:31_new} \begin{split} &\nabla F(\mathbf{w}) - \nabla_{1} \bar{F}(\mathbf{w}, \mathbf{u}) = \nabla_{1} f(\mathbf{w}, \mathbf{u}^*) - \nabla_{1} f(\mathbf{w}, \mathbf{u}) \\ & -\nabla_{12} g(\mathbf{w}, \mathbf{u}^*) \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} \nabla_{2} f(\mathbf{w}, \mathbf{u}^*) - \lambda (\nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*)). \end{split} \end{equation}\]

We can rearrange terms for \((\nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*))\) as the following:

\[\begin{equation} \label{eq:32_new} \begin{split} \nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*) &= \nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} (\mathbf{u} - \mathbf{u}^*) \\ & + \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} (\mathbf{u} - \mathbf{u}^*). \end{split} \end{equation}\]

To continue, we leverage the following equality: \[\begin{equation*} \begin{split} \mathbf{u} - \mathbf{u}^* =& -\nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} g(\mathbf{w}, \mathbf{u}) - \nabla_{2} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{22} g(\mathbf{w}, \mathbf{u}^*) (\mathbf{u} - \mathbf{u}^*)) \\ & +\nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} g(\mathbf{w}, \mathbf{u}) - \nabla_{2} g(\mathbf{w}, \mathbf{u}^*)). \end{split} \end{equation*}\] By the optimality condition for \(\mathbf{u}^*\), \(\nabla_{2} g(\mathbf{w}, \mathbf{u}^*) = 0\), and \(\nabla_{2} \bar{F}(\mathbf{w}, \mathbf{u})=\nabla_{2} f(\mathbf{w}, \mathbf{u}) + \lambda \nabla_{2} g(\mathbf{w}, \mathbf{u})\), we can express \(\mathbf{u} - \mathbf{u}^*\) as

\[\begin{equation}\label{eq:33_new} \begin{split} \mathbf{u} - \mathbf{u}^* &= -\nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} g(\mathbf{w}, \mathbf{u}) - \nabla_{2} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{22} g(\mathbf{w}, \mathbf{u}^*) (\mathbf{u} - \mathbf{u}^*)) \\ & + \frac{1}{\lambda} \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} \bar{F}(\mathbf{w}, \mathbf{u}) - \nabla_{2} f(\mathbf{w}, \mathbf{u})). \end{split} \end{equation}\]

Combining (\(\ref{eq:33_new}\)) and (\(\ref{eq:32_new}\)) and then plugging the result back to (\(\ref{eq:31_new}\)), we have \[\begin{equation*} \begin{split} &\nabla F(\mathbf{w}) - \nabla_{1} \bar{F}(\mathbf{w}, \mathbf{u}) = \nabla_{1} f(\mathbf{w}, \mathbf{u}^*) - \nabla_{1} f(\mathbf{w}, \mathbf{u}) \\ & -\nabla_{12} g(\mathbf{w}, \mathbf{u}^*) \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} \nabla_{2} f(\mathbf{w}, \mathbf{u}^*) \\ &- \lambda (\nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} (\mathbf{u} - \mathbf{u}^*))\\ & +\lambda \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} g(\mathbf{w}, \mathbf{u}) - \nabla_{2} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{22} g(\mathbf{w}, \mathbf{u}^*) (\mathbf{u} - \mathbf{u}^*)) \\ & - \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} \bar{F}(\mathbf{w}, \mathbf{u}) - \nabla_{2} f(\mathbf{w}, \mathbf{u})). \end{split} \end{equation*}\]

As a result, we have \[\begin{equation*} \begin{split} &\nabla F(\mathbf{w}) - \nabla_{1} \bar{F}(\mathbf{w}, \mathbf{u})+ \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} \nabla_{2} \bar{F}(\mathbf{w}, \mathbf{u}) \\ & = \nabla_{1} f(\mathbf{w}, \mathbf{u}^*) - \nabla_{1} f(\mathbf{w}, \mathbf{u}) \\ & -\nabla_{12} g(\mathbf{w}, \mathbf{u}^*) \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} f(\mathbf{w}, \mathbf{u}^*) - \nabla_{2} f(\mathbf{w}, \mathbf{u})) \\ &- \lambda (\nabla_{1} g(\mathbf{w}, \mathbf{u}) - \nabla_{1} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} (\mathbf{u} - \mathbf{u}^*))\\ & +\lambda \nabla_{12} g(\mathbf{w}, \mathbf{u}^*)^{\top} \nabla_{22} g(\mathbf{w}, \mathbf{u}^*)^{-1} (\nabla_{2} g(\mathbf{w}, \mathbf{u}) - \nabla_{2} g(\mathbf{w}, \mathbf{u}^*) - \nabla_{22} g(\mathbf{w}, \mathbf{u}^*) (\mathbf{u} - \mathbf{u}^*)) . \end{split} \end{equation*}\]

■

Next, we bound \(\|\mathbf{u}^*_\lambda(\mathbf{w}) - \mathbf{u}^*(\mathbf{w})\|_2\).

■

■

← Go Back