IAPTT for Bilevel Optimization
Published:
Paper Reading: Towards Gradient-based Bilevel Optimization with Non-convex Followers and Beyond
Introduction and Motivation
文章[1]考虑如下的双层优化问题，
\[\min_{x \in \mathcal{X},y \in Y^*(x)} f(x,y) \quad {\rm s.t.} \quad Y^*(x) = \arg \min_{y \in \mathcal{Y}} g(x,y)\]其中 $\mathcal{X} \subseteq \mathbb{R}^{d_x}, \mathcal{Y} \subseteq \mathbb{R}^{d_y}$ 为紧集。 我们定义如下的hyper-objective
\[\begin{align*} \varphi(x) = \min_{y \in Y^*(x)} f(x,y), \quad {\rm s.t.} \quad Y^*(x) = \arg \min_{y \in \mathcal{Y}} g(x,y) \end{align*}\]则问题可以转化为关于 $\min_x \varphi(x)$ 的单层优化问题。
对于内层问题，可以对于给定的$x$ 以及初始点 $z$ 可以使用梯度算法进行求解
\[\begin{align*} y_{k+1}(x) = {\rm Proj}_{\mathcal{Y}}(y_k(x) - \eta \nabla_y g(x,y_k(x))), \quad y_0 =z. \end{align*}\]文章 [1] 针对内层函数非凸的情况，提出了如下的两个技术，分别为
Initialization Auxiliary (IA), 对于初始点 $z$ 作为一个辅助变量也进行梯度下降
Pessimistic Trajectory Truncation (PTT), 使用悲观的策略选取内层的输出 $y_k(x)$.
文章最终提出的算法结合了上述两个技术，称为IAPTT.
文章称其算法可以适用于非凸优化问题，但其 Assumption 3.1 (5) 的假设似乎对于一般的函数很难验证，一类可验证的假设称为KL不等式，例如 [2] 中的Assumption 2, 假设存在正常数 $\theta,\tau$ 使得下面的不等式永远成立
\[\begin{align*} \tau (g(x,y) -g^*(x))^{\theta} \le {\rm dist}(0, \nabla_y g(x,y) + \partial \delta_{\mathcal{Y}}(y)). \end{align*}\]其中我们定义 $\delta_{\mathcal{Y}}(\,\cdot\,)$ 为集合 $\mathcal{Y}$ 的示性函数。
Theoretical Investigation
本节给出文章的理论保证。如果得到序列
\[\begin{align*} (x_k,z_k) \in \arg \min_{x,z} \{ \max_j f(x,y_j(x,z))\}. \end{align*}\]那么可以证明如果序列 ${ x_k}$ 存在极限点 $\bar x$, 那么成立 $\bar x \in \arg \min_x \varphi(x)$.
其中辅助变量 $z$ 来源于文章所提出的技术IA， 而 $\max$ 操作对应于文章所提出的技术PTT
下面我们首先证明如果 $y_k(x,z) \rightarrow \bar y$, 那么该极限点 $ \bar y$ 一定属于 $Y^*(x)$.
根据梯度下降的结论，以及 $\mathcal{X}, \mathcal{Y}$ 都为紧集的假设，我们知道对于任意的 $x,z$ 都有
\[\begin{align*} \min_{1 \le j \le k} \Vert \mathcal{G}(x,y_j(x,z)) \Vert \rightarrow 0. \end{align*}\]其中 $\mathcal{G}(\,\cdot\,, \cdot\,)$ 为梯度映射，其定义为 $ \mathcal{G}(x,y) = y - {\rm Proj}_{\mathcal{Y}}(y - \eta \nabla_y g(x,y))$.
Proof for a Class of Nonconvex LL Problem
本节我们说明，对于满足如下条件的内层函数：
\[\begin{align*} \mathcal{G}(x,y) = 0 \Leftrightarrow y \in Y^*(x). \end{align*}\]可以证明所需要的收敛结果。
根据上文的结论，我们知道对于任意的 $\epsilon>0$, 都存在 $K_1$ 使得对于任意的 $x,z$ 都存在 $i(k)$ 使得
\[\begin{align*} \Vert \mathcal{G}(x,y_{i(k)}(x,z)) \Vert \le \epsilon \quad \forall k \ge K_1. \end{align*}\]对于序列 ${ y_{i(k)}(x_k,z_k)}$, 我们选取其一个收敛子列，为了方便我们仍然用下标 $k$ 表示该子列。
我们假设该收敛子列的极限点为 $\bar y$, 也即 $ y_{i(k)}(x_k,z_k) \rightarrow \bar y$.
根据梯度映射的连续性，存在 $K_2$ 使得
\[\begin{align*} \Vert \mathcal{G} (\bar x,\bar y ) \Vert \le\Vert \mathcal{G}(x_k,y_{i(k)}(x_k,z_k)) \Vert + \epsilon \le 2 \epsilon, \quad \forall k \ge \max(K_1,K_2). \end{align*}\]那么由于 $\epsilon$ 的任意性，令 $\epsilon \rightarrow 0$ 并且使用平稳点就是最优点的假设就可以得到 $\bar y \in Y^\ast(x)$.
下面我们对我们所关心的hyper-objective 对象 $\varphi(x)$ 进行分析，首先我们有
\[\begin{align*} \varphi(\bar x) = \min_{y \in Y^*(\bar x)} f(\bar x,y) \le f(\bar x, \bar y) \end{align*}\]根据 $f$ 的连续性以及 $(x_k, y_{i(k)}(x_k,z_k)) \rightarrow (\bar x, \bar y)$ ，对于任意的 $\epsilon>0$, 都存在 $K_3$ 使得
\[\begin{align*} \varphi(\bar x) \le f(\bar x, \bar y) \le f(x_k, y_{i(k)} (x_k,z_k)) + \epsilon, \quad \forall k \ge K_3. \end{align*}\]算法并不知道 $i(k)$ 具体取值，但是可以使用悲观策略PTT，那么就有
\[\begin{align*} \varphi(\bar x ) \le \max_{1 \le j \le k} f(x_k , y_j (x_k,z_k)) + \epsilon, \quad \forall k \ge K_3. \end{align*}\]根据辅助初始化策略IA，我们知道对于任意的 $x$ 若 $z \in Y^\ast(x)$, 那么对于任意的 $1 \le j \le k$ 都有 $ y_j (x,z) = z$ .
这告诉我们,
\[\begin{align*} \min_{z} \{ \max_{1 \le j \le k} f(x,y_j(x,z))\} \le \min_{y \in Y^\ast(x)}f(x,y), \quad \forall k. \end{align*}\]根据 $(x_k,z_k) \in \arg \min_{x,z} { \max_j f(x,y_j(x,z))} $, 我们进一步得到对于任意的 $x$ 都成立，
\[\begin{align*} \varphi(\bar x) \le \max_{1 \le j \le k} f(x_k , y_j (x_k,z_k)) + \epsilon \le \min_{y \in Y^*(x)} f(x,y) + \epsilon = \varphi(x) + \epsilon, \quad \forall k \ge K_3. \end{align*}\]最后由于 $\epsilon$ 的任意性，令 $\epsilon \rightarrow 0$ 就可以证明 $\bar x \in \arg \min_x \varphi(x)$.
Proof for Convex LL Problem
本节我们说明，对于内层函数为凸函数的情况，算法并不需要悲观策略PTT.
此时该算法对应于如下的对于序列：
\[\begin{align*} (x_k,z_k) \in \arg \min_{x,z} \{ f(x,y_k(x,z))\}. \end{align*}\]那么可以证明如果序列 ${ x_k}$ 存在极限点 $\bar x$, 那么成立 $\bar x \in \arg \min_x \varphi(x)$.
证明思路与上文类似，但由于凸性证明可以有所简化，这是由于对于凸函数，我们有
\[\begin{align*} g(x,y_k(x)) \rightarrow g^*(x). \end{align*}\]这意味着如果序列 ${y_k(x_k,z_k) }$ 存在极限， 也即 $y_k(x_k,z_k) \rightarrow \bar y$, 那么有 $ \bar y \in Y^\ast(x)$.
这是由于对于任意的 $\epsilon>0$ , 根据 $g,g^\ast$的连续性，存在 $K_1 $ 使得
\[\begin{align*} g(\bar x, \bar y) - g^*(\bar x) \le g(x_k,y_k(x_k,z_k)) - g^*(x_k) + \epsilon, \quad k \ge K_1. \end{align*}\]根据 $g(x,y_k(x,z)) \rightarrow g^\ast(x)$, 存在 $K_2$ 使得
\[\begin{align*} g(\bar x, \bar y) - g^*(\bar x) \le 2 \epsilon, \quad \forall k \ge \max(K_1,K_2). \end{align*}\]最后根据 $\epsilon$ 的任意性并且令 $\epsilon \rightarrow 0$ 就可以证明 $\bar y \in Y^\ast(x)$.
\[\begin{align*} \varphi(\bar x) = \min_{y \in Y^*(\bar x)} f(\bar x,y )\le f(\bar x, \bar y). \end{align*}\]取 ${y_k(x_k,z_k)}$ 的一个收敛子列，为了方便起见我们仍然采用下标 $k$.
根据 $f$ 的连续性以及 $(x_k,y_k(x_k,z_k)) \rightarrow (\bar x, \bar y)$, 对于任意的 $\epsilon>0$ 都存在 $K_3$ 使得
\[\begin{align*} \varphi(\bar x) \le f(x_k, y_k(x_k,z_k)) + \epsilon, \quad \forall k \ge K_3. \end{align*}\]根据辅助初始化策略IA，我们知道对于任意的 $x$ 若 $z \in Y^\ast(x)$, 那么对于任意的 $ k$ 都有 $ y_k (x,z) = z$ .
这告诉我们,
\[\begin{align*} \min_z f(x,y_k(x,y)) \le \min_{y \in Y^*(x)} f(x,y ) = \varphi(x). \end{align*}\]那么根据 $x_k, z_k$ 的定义，我们知道对于任意的 $x$ 都有，
\[\begin{align*} \varphi(\bar x) \le f(x_k,y_k(x_k,z_k)) + \epsilon \le \varphi(x) + \epsilon , \quad \forall k \ge K_3. \end{align*}\]最后由于 $\epsilon$ 的任意性，令 $\epsilon \rightarrow 0$ 就可以证明 $\bar x \in \arg \min_x \varphi(x)$.
当然在内层函数具有凸性的时候，可以使用加速梯度下降（AGD）来替代梯度下降得到更好的性能。
Reference
[1] Towards Gradient-based Bilevel Optimization with Non-convex Followers and Beyond. In NeurIPS, 2021.
[2] Doubly Smoothed GDA for Constrained Nonconvex-Nonconcave Minimax Optimization. arXiv preprint, 2023