Section 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 3.

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\).

← Go Back