Optimal Algorithms for Stochastic Bilevel Optimization under Relaxed Smoothness Conditions

Introduction

Bilevel optimization is gaining increasing popularity within the machine learning community due to its extensive range of applications, including meta-learning , hyperparameter optimization , data augmentation , and neural architecture search . The objective of bilevel optimization is to minimize a function that is dependent on the solution of another optimization problem:

Stochastic bilevel optimization can be considered as an extension of bilevel empirical risk minimization , allowing for the effective handling of online and streaming data $(\xi,\phi)$.

In many instances, the analytical expression of $y^{*}(x)$ is unknown and can only be approximated using an optimization algorithm. This adds to the complexity of problem (1) compared to its single-level counterpart. Under regular conditions such that $\Phi$ is differentiable, the hypergradient $\nabla\Phi(x)$ derived by the chain rule and the implicit function theorem is given by

Solving (1) using only stochastic oracles poses significant challenges since there is no direct unbiased estimator available for $\left[\nabla_{22}^{2}g(x,y^{*}(x))\right]^{-1}$ and also for $\nabla\Phi(x)$ as a consequence.

To mitigate the estimation bias, many existing methods employ a Hessian Inverse Approximation (HIA) subroutine, which involves drawing a mini-batch of stochastic Hessian matrices and computing a truncated Neumann series . However, this subroutine comes with an increased computational burden and introduces an additional factor of $\log(\epsilon^{-1})$ in the sample complexity. Some alternative methods calculate the explicit inverse of the stochastic Hessian matrix with momentum updates. To circumvent the need for explicit Hessian inversion and the HIA subroutine, recent works propose running Stochastic Gradient Descent (SGD) steps to approximate the solution $z^{*}(x)$ of the linear system (4). In particular, the state-of-the-art Stochastic Bilevel Algorithm (SOBA) only utilizes SGD steps to simultaneously update three variables: the inner variable $y$, the outer variable $x$, and the auxiliary variable $z$. Remarkably, SOBA achieves the same complexity lower bound of its single-level counterpart ($\Phi\in\mathcal{C}_{L}^{1,1}$ $\mathcal{C}_{L}^{p,p}$ denotes $p$-times differentiability with Lipschitz $k$-th order derivatives for $0<k\leq p$.) in the non-convex setting .

Despite the superior computational and sample efficiency of SOBA, its current theoretical framework assumes high-order smoothness for the UL function $f$ and the LL function $g$ such that $z^{*}(x)$ has Lipschitz gradient. Specifically, unlike the typical assumptions in stochastic bilevel optimization that state $f\in\mathcal{C}_{L}^{1,1}$ and $g\in\mathcal{C}_{L}^{2,2}$ (A1), the current theory of SOBA requires $f\in\mathcal{C}_{L}^{2,2}$ and $g\in\mathcal{C}_{L}^{3,3}$ (A2). The necessity of (A2) is counter-intuitive as the partial gradients of $x,y,z$ utilized in constructing SGD steps are already Lipschitz continuous under (A1). Furthermore, assuming $g$ is strongly convex and the partial gradient of the UL function with respect to the inner variable $y$ is bounded for all pairs of $(x,y^{*}(x))$, (i.e., $\left\|\nabla_{2}f(x,y^{*}(x))\right\|\leq L_{f}$ for all $x\in\mathcal{X}$), there exists a subset relation among three function classes as follows (Lemma 2.2 in ):

In light of this, it can be concluded that (A1) is sufficient to ensure the first-order Lipschitzness of $\Phi$, which is the standard assumption in the single-level setting. Therefore, a natural question follows:

Is it possible to develop a fully single-loop and Hessian-inversion-free algorithm for solving stochastic bilevel optimization problems that achieves an optimal sample complexity of $\mathcal{O}(\epsilon^{-2})$ under standard smoothness assumptions $\{f\in\mathcal{C}_{L}^{1,1},g\in\mathcal{C}_{L}^{2,2}\}$9? In this paper, we provide an affirmative answer to the aforementioned question. Our contributions can be summarized as follows:

We propose a class of fully single-loop and Hession-inversion-free algorithm, named Moving-Average SOBA (MA-SOBA), which builds upon the SOBA algorithm by incorporating an additional sequence of average hypergradients. Unlike SOBA, MA-SOBA achieves an optimal sample complexity of $\mathcal{O}(\epsilon^{-2})$ under standard smoothness assumptions, without relying on high-order smoothness. Moreover, the introduced sequence of average hypergradients converges to $\nabla\Phi(x)$, thus offering a reliable termination criterion in the stochastic setting.

We expand the scope of MA-SOBA to tackle a broader class of problems, specifically the min-max multi-objective bilevel optimization problem with significant applications in robust machine learning. We introduce MORMA-SOBA, an algorithm that can find an $\epsilon$-first-order stationary point of the $\mu_{\lambda}$-strongly-concave regularized formulation while achieving a sample complexity of $\mathcal{O}(n^{5}\mu_{\lambda}^{-4}\epsilon^{-2})$, which fills a gap (in terms of the order of $\epsilon$-dependency) in the existing literature.

We conduct experiments on several machine learning problems. Our numerical results show the efficiency and superiority of our algorithms.

Related Work. The concept of bilevel optimization was initially introduced in the work of . Since then, numerous gradient-based bilevel optimization algorithms have been proposed, broadly categorized into two groups: ITerative Differentiation (ITD) based methods and Approximate Implicit Differentiation (AID) based methods . The ITD-based algorithms typically involve approximating the solution of the inner problem using an iterative algorithm and then computing an approximate hypergradient through automatic differentiation. However, a major drawback of this approach is the necessity of storing each iterate of the inner optimization algorithm in memory. The AID-based algorithms leverage the implicit gradient given by (3), which requires the solution of a linear system characterized by (4). Extensive research has been conducted on designing and analyzing deterministic bilevel optimization algorithms with strongly-convex LL functions; see and the references cited therein.

In recent years, there has been a growing interest in stochastic bilevel optimization, especially in the setting of a non-convex UL function and a strongly-convex LL function. To address estimation bias, one set of methods uses SGD iterations for the inner problem and employs truncated stochastic Neumann series to approximate the inverse of the Hessian matrix in $z^{*}(x)$ . The analysis of such methods was refined by to achieve convergence rates similar to those of SGD. However, the Neumann approximation subroutine introduces an additional factor of $\log(\epsilon^{-1})$ in the sample complexity. Some alternative approaches calculate the explicit inverse of the stochastic Hessian matrix with momentum updates. Nevertheless, these methods encounter challenges related to computational complexity in matrix inversion and numerical stability.

To avoid the need for explicit Hessian inversion and the Neumann approximation, recent algorithms propose running SGD steps to approximate the solution $z^{*}(x)$ of the linear system (4). One such algorithm called AmIGO employs a double-loop approach and achieves an optimal sample complexity of $\mathcal{O}(\epsilon^{-2})$ under regular assumptions. However, AmIGO requires a growing batch size inversely proportional to $\epsilon$. On the other hand, the single-loop algorithm SOBA achieves the same complexity lower bound but with constant batch size. Unfortunately, the current analysis of SOBA relies on the assumption of higher-order smoothness for the UL and LL functions. In this work, we introduce a novel algorithm framework that differs slightly from SOBA but can achieve optimal sample complexity in theory without higher-order smoothness assumptions. A summary of our results and comparison to prior work is provided in Table 1.

In addition, there exist several variance reduction-based methods following the line of research by . Some of these methods achieve an improved sample complexity of $\mathcal{O}(\epsilon^{-1.5})$ and match the lower bounds of their single-level counterparts when stochastic functions $F_{\xi}$ and $G_{\phi}$ satisfy mean-squared smoothness assumptions and the algorithm is allowed simultaneous queries at the same random seed . However, since we are specifically considering smoothness assumptions on $f$ and $g$, we will not delve into the comparison with these methods.

The most recent advancements in (stochastic) bilevel optimization focus on several new ideas: (i) addressing constrained lower-level problems , (ii) handling lower-level problems that lack strong convexity , (iii) developing fully first-order (Hesssian-free) algorithms , (iv) establishing convergence to the second-order stationary point , and (v) expanding the framework to encompass multi-objective optimization problems . It is promising to apply some of these advancements to our specific framework. However, in this work, we contribute to multi-objective bilevel problems with a slight modification of our approach. Other directions are left as future work.

Proposed Framework: the MA-SOBA Algorithm

Similar to , our algorithm initiates with inexact hypergradient descent techniques and seeks to offer an alternative in the stochastic setting. To provide a clear illustration, let us initially consider the deterministic setting. The SOBA framework keeps track of three sequences, denoted as $\{x^{k},y^{k},z^{k}\}$, and updates them using $D_{x},D_{y},D_{z}$ as follows:

where (5) is the GD step to minimize $g(x^{k},\cdot)$, (7) is the inexact hyper gradient descent step, and (6) is the GD step to minimize a quadratic minimization problem with $z^{*}(x^{k})$ being the solution, i.e.,

Given that the above update rule, highlighted in blue, does not involve the Hessian matrix inversion, SOBA can directly utilize the stochastic oracles of $\nabla_{1}f,\nabla_{2}f,\nabla_{2}g,\nabla_{22}^{2}g,\nabla_{12}^{2}g$ to obtain unbiased estimators of $D_{x},D_{y},D_{z}$ in Eq.(5), (6), (7). This approach circumvents the requirement for a Neumann approximation subroutine or a direct matrix inversion. However, due to the update rule for $y$, which only utilizes one-step SGD at each iteration $k$, the value of $y^{k}$ does not coincide with $y^{*}(x^{k})$. As a result, a certain bias is introduced in the partial gradient of $z$ in Eq.(6). Similarly, when estimating the hypergradient $\nabla\Phi(x)$, another bias term arises in Eq.(7). Although the bias decreases to zero as $y^{k}\rightarrow y^{*}(x^{k})$ and $z^{k}\rightarrow z^{*}(x^{k})$ under standard smoothness assumptions as indicated by Lemma 3.4 in , the current analysis of SOBA requires more regularity on $f$ and $g$ to carefully handle the bias; it assume that $f$ has Lipschitz Hessian and $g$ has Lipschitz third-order derivative.

The inability to obtain an unbiased gradient estimator is a common characteristic in stochastic optimization involving nested structures; see, for example, stochastic compositional optimization as a specific case of (1). One popular approach is to introduce a sequence of dual variables that approximates the true gradient by aggregating all past biased stochastic gradients using a moving averaging technique . Motivated by this approach, we introduce another sequence of variables, denoted as $\{h^{k}\}$, and update it at $k$-th iteration given the past iterates $\mathscr{F}_{k}$ as follows:

Theoretical Analysis

In this section, we provide convergence rates of MA-SOBA under standard smoothness conditions on $f,g$ and regular assumptions on stochastic oracles. We also present a proof sketch and have detailed discussions about assumptions made in the literature. The complete proofs are deferred in Appendix.

We first state some regularity assumptions on the functions $f$ and $g$.

($f\in\mathcal{C}_{L}^{1,1}$ and $g\in\mathcal{C}_{L}^{2,2}$)9 $\nabla f,\nabla g,\nabla^{2}g$ are $L_{\nabla f},L_{\nabla g},L_{\nabla^{2}g}$ Lipschitz continuous respectively.

(SC LL) $g$ is $\mu_{g}$-strongly convex.

$\left\|\nabla_{2}f(x,y^{*}(x))\right\|\leq L_{f}<\infty$ for all $x\in\mathcal{X}$.

The above assumption serves as a sufficient condition for the Lipschitz continuity of $\nabla\Phi$, $y^{*}(x)$, and $z^{*}(x)$, as well as $D_{x}$, $D_{y}$, and $D_{z}$ in Eq. (5), (6), (7). The inclusion of high-order smoothness assumptions ($f\in\mathcal{C}_{L}^{2,2}$ and $g\in\mathcal{C}_{L}^{3,3}$) in the current analysis of SOBA is primarily intended to ensure the Lipschitzness of $\nabla z^{*}(x)$. However, the necessity of such assumptions is subject to doubt, given that $\nabla z^{*}(x)$ is not involved in designing the algorithm. Furthermore, the Lipschitzness of $f$ or uniformly boundedness of $\nabla_{2}f$ made in several previous works is unnecessary. Instead, the boundedness assumption on $\nabla_{2}f$ is only required for all pairs of $(x,y^{*}(x))$ as demonstrated by (c).

Next, we discuss assumptions made on the stochastic oracles.

2 Convergence Results

We have the following theorem characterizing the convergence results of MA-SOBA.

3 Proof Sketch of Theorem 1

Define $V_{k}=\frac{1}{\tau^{2}}\|x_{+}^{k}-x^{k}\|^{2}+\|h^{k}-\nabla\Phi(x^{k})\|^{2}$. To obtain (8), we consider the merit function $W_{k}$:

where $\eta_{\mathcal{X}}(x,h,\tau)=\left\langle h,x_{+}-x\right\rangle+\frac{1}{2\tau}\left\|x_{+}-x\right\|^{2}$. By leveraging the moving average updates of $x^{k}$ (line 2 of Algorithm 1), we can obtain

It is worth noting that requires the existence and Lipschitzness of $\nabla^{2}f$ and $\nabla^{3}g$ to ensure the Lipschitzness of $\nabla z^{*}(x)$ (see (4)) which is used in proving the sufficient decrease of $\left\|z^{k}-z_{*}^{k}\right\|^{2}$. In contrast, based on the moving average updates of $x^{k}$ and $h^{k}$, our refined analysis does not necessitate such high-order smoothness assumptions to obtain that

The proof of Theorem 1 can then be completed by choosing appropriate $\alpha_{k},c_{1},c_{2},c_{3},\tau>0$.

Min-Max Bilevel Optimization

To incorporate robustness in the multi-objective setting where each objective can be expressed as a bilevel optimization problem in (1), the following mini-max bilevel problem formulation was proposed in :

Note that (9) can be reformulated as a general nonconvex-concave min-max optimization problem (with a bilevel substructure):

Instead of solving (10) directly, in this work, we focus on solving the following regularized version,

The proposed algorithm, which we refer as to Multi-Objective Robust MA-SOBA (MORMA-SOBA), for solving (11) is presented in 2. In addition to the basic framework of Algorithm 1, we also maintain a moving average step in the updates of $\lambda^{k}$ for solving the max part of problem 2. It is worth noting that in its single-level counterpart without the inner variable $y$, the proposed MORMA-SOBA algorithm is fundamentally similar to the single-timescale averaged SGDA algorithm proposed in the general nonconvex-strongly-concave setting . Moreover, our algorithm framework can be leveraged to solve the distributionally robust compositional optimization problem discussed in .

2 Convergence Results

We first present additional assumptions required in the analysis of MORMA-SOBA.

For any $k\geq 0$, functions $\Phi(x),\nabla\Phi_{i}(x)$ are bounded, functions $f_{i}$ are $L_{f}$-Lipschitze continuous in the second input, and their stochastic versions are unbiased with bounded variance, i.e., there exists $L_{\Phi},L_{f},\sigma_{f,0}\geq 0$ such that

$\bigcup_{i=1}^{n}\left\{u_{x,i}^{k+1},u_{y,i}^{k+1},v_{i}^{k+1},H_{i}^{k+1},J_{i}^{k+1}\right\}\cup\left\{s^{k+1}\right\}$ are conditionally independent with respect to $\mathscr{F}_{k}$.

We have the following convergence theorem of MORMA-SOBA.

which is commonly used in min-max optimization (see, e.g., ). The construction of the first two terms follows the same idea in Section 3.3, and for the last two terms we have

Note that in Theorem 2 we explicitly characterize the dependency on $n$ and $\mu_{\lambda}$ in the convergence rate and the sample complexity. It is worth noting that two variants of stochastic gradient descent ascent (SGDA) algorithms for solving the nonconvex-strongly-concave min-max optimization problems (without bilevel substructures), have been studied in . While such algorithms are not immediately applicable to solve (11) due to the presence of the additional bilevel substructure, it is instructive to compare to those methods assuming direct access to $y_{i}^{*}(x)$ in (9). Specifically, we observe that the sample complexity of SGDA with batch size $M=\Theta(n^{1.5}\epsilon^{-1})$ in and moving-average SGDA with $\mathcal{O}(1)$ batch size in for solving (11) assuming direct access to $y_{i}^{*}(x)$ will be $\mathcal{O}\left(n^{4}\mu_{\lambda}^{-2}\epsilon^{-2}\right)$ and $\mathcal{O}\left(n^{5}\mu_{\lambda}^{-4}\epsilon^{-2}\right)$Note that $\Phi_{\mu_{\lambda}}(x,\lambda)$ in (10) is quadratic in $\lambda$, and these two sample complexities are obtained under this special case, i.e., $\nabla_{2}^{2}f(x,y)=-\mu\mathbf{I}$ applied to . respectively. Our results in Theorem 2 indicate that the sample complexity of the proposed algorithm MORMA-SOBA for solving min-max bilevel problems has the same dependency on $n$ and $\mu_{\lambda}$ as the sample complexity of the moving-average SGDA introduced in for solving min-max single-level problems, while also computing $y_{i}^{*}(x)$ instead of assuming direct access.

Experiments

Conclusion

In this work, we propose a novel class of algorithms (MA-SOBA) for solving stochastic bilevel optimization problems in (1) by introducing the moving-average step to estimate the hypergradient. We present a refined convergence analysis of our algorithm, achieving the optimal sample complexity without relying on the high-order smoothness assumptions employed in the literature. Furthermore, we extend our algorithm framework to tackle a generic min-max bilevel optimization problem within the multi-objective setting, identifying and addressing the theoretical gap present in the literature.

We thank the authors of for clarifications regarding their paper.

References

Appendix A Experimental Details

All the experiments were conducted using Python. The initial two tasks, involving the comparison of MA-SOBA with other stochastic bilevel optimization algorithms, utilized the Benchopt package and the open-sourced bilevel benchmark https://github.com/benchopt/benchmark_bilevel. The final task, focused on robust multi-objective representation learning, was implemented in PyTorch, following the source code provided by https://github.com/minimario/MORBiT.

Setup. In our experiments, we strictly adhere to the settings provided in bench_bilevel¶ ‣ A, as detailed in Appendix B.1 of . The previous results and setups of have also been available in https://benchopt.github.io/results/benchmark_bilevel.html. For completeness, we provide a summary of the setup below.

To avoid redundant computations, we utilize oracles for the function $F_{\xi},G_{\phi}$, which provide access to quantities such as $\nabla_{1}F_{\xi}(x,y)$, $\nabla_{2}F_{\xi}(x,y)$, $\nabla_{2}G_{\phi}(x,y)$, $\nabla_{22}^{2}G_{\phi}(x,y)v$, and $\nabla_{12}^{2}G_{\phi}(x,y)v$, although this approach may violate the independence assumption in Assumption 2.

In all our experiments, we employ a batch size of 64 for all methods, even for BSA and AmIGO that theoretically require increasing batch sizes.

For methods involving an inner loop (stocBiO, BSA, AmIGO), we perform 10 inner steps per each outer iteration as proposed in those papers.

For methods that involve the Neumann approximation for the Hessian vector product (such as BSA, TTSA, SUSTAIN, and MRBO), we perform 10 steps of the subroutine per outer iteration. For AmIGO, we perform 10 steps of SGD to approximate the inversion of the linear system.

The step sizes and momentum parameters used in all benchmark algorithms are directly adopted from the fine-tuned parameters provided by . From a grid search, we select the best constant step sizes for MO-SOBA.

At present, we have excluded SRBA from the benchmark due to the unavailability of an open-sourced implementation and its limited reported improvement over SABA.

In this experiment, we focus on selecting the regularization parameters for a multi-regularized logistic regression model on the IJCNN1 dataset, where we have one hyperparameter per feature. Specifically, the problem can be formulated as:

To complement the comparison presented in the main paper, we conducted additional experiments that involved comparing all benchmark methods, including the variance reduction based method. In Figure 3, we plot the suboptimality gap ($\Phi(x)-\Phi^{*}$) against runtime and the number of calls to oracles. Unfortunately, the previous results obtained for MRBO and AmIGO on the IJCNN1 dataset are not reproducible at the moment due to some conflicts in the current developer version of Benchopt. As reported in , MRBO exhibits similar performance to SUSTAIN, while the curve of AmIGO initially follows a similar trend as SUSTAIN and eventually reaches a similar level as SABA towards the end. Following a grid search, we have selected the parameters in MA-SOBA as $\alpha_{k}\tau=0.02$, $\beta_{k}=\gamma_{k}=0.01$, and $\theta_{k}=0.1$. As shown in Figure 3, our proposed method MA-SOBA outperforms SOBA significantly, achieving a slightly lower suboptimality gap compared to the state-of-the-art variance reduction-based method SABA.

A.1.2 Data Hyper-Cleaning on MNIST

To supplement the comparison presented in the main paper, we conducted additional experiments that involved comparing all benchmark methods, including the variance reduction-based method. Following a grid search, we have selected the parameters in MA-SOBA as $\alpha_{k}\tau=10^{3}$, $\beta_{k}=\gamma_{k}=10^{-2}$, and $\theta_{k}=10^{-1}$. In Figure 4, we plot the test error against runtime and the number of calls to oracles with different corruption probability $p\in\{0.5,0.7,0.9\}$. We observe that MA-SOBA has comparable performance to the state-of-the-art method SABA. Remarkably, MA-SOBA is the fastest algorithm to reach the best test accuracy when $p=0.5$.

A.1.3 Moving Average vs. Variance Reduction

Through empirical studies, we have demonstrated that our proposed method, MA-SOBA, which utilizes a moving average (MA) technique, achieves comparable performance to the state-of-the-art variance reduction-based approach SABA using SAGA updates . In this context, we would like to highlight the key difference and relationship between these two methods.

We start with presenting the update rules of the sequence of estimated gradients $\{g^{k}\}$ for the variance reduction techniques SAGA and our moving average method (MA) for the single-level problem.

SAGA (finite-sum:) $\min~{}\frac{1}{n}\sum_{i=1}^{n}f_{i}(x)$ $\displaystyle g^{k}=\nabla f_{i_{k}}(x^{k})-\nabla f_{i_{k}}(\bar{x}_{i_{k}})+\frac{1}{n}\sum_{j=1}^{n}\nabla f_{j}(\bar{x}_{j})$ The SAGA update is designed for finite-sum problems with offline batch data. At each iteration $k$, the algorithm randomly selects an index $i_{k}\in[n]$ and updates the gradient variable $g^{k}$ using a reference point $\bar{x}_{i_{k}}$, which corresponds to the last evaluated point for $\nabla f_{i_{k}}$. However, it should be noted that SAGA requires storing the previously evaluated gradients ${\nabla f_{j}(\bar{x}_{j})}$ in a table, which can be memory-intensive when sample size $n$ or dimension $d$ is large. In the finite-sum setting, there exist several other variance reduction methods, such as SARAH , that can be employed to further enhance the dependence on the number of samples, $n$, for bilevel optimization problems. However, the SARAH-type method requires double gradient evaluations on each iteration of $x^{k}$ and $x^{k-1}$.

However, MSS is a stronger assumption than the general smoothness assumption on $f$:

By Jensen’s inequality, we have that MSS is stronger than the general smoothness assumption on $f$:

In this work, the theoretical results of the proposed methods are only built on the smoothness assumption on the UL and LL functions $f,g$ without further assuming MSS on $F_{\xi}$ and $G_{\phi}$. It is worth noting that a clear distinction in the lower bounds of sample complexity for solving the single-level stochastic optimization has been proven in . Specifically, they establish a separation under the MSS assumption on $F_{\xi}$ and smoothness assumptions on $f$ ($\mathcal{O}(\epsilon^{-1.5})$ vs. $\mathcal{O}(\epsilon^{-2})$). Thus, it is important to emphasize that MA-SOBA achieves the optimal sample complexity $\mathcal{O}(\epsilon^{-2})$ under our weaker assumptions.

Practical Implementation. Variance reduction methods often entail additional space complexity, require double-loop implementation or double oracle computations per iteration. These requirements can be unfavorable for large-scale problems with limited computing resources. For instance, in the second task, the runtime improvement achieved by using SABA is limited. This limitation can be attributed to the dimensionality of the variables $\nu$ (with a dimension of $20,000$) and $W$ (with a dimension of $10\times 784$). The benefit of using variance reduction methods is expected to be less significant for more complex problems involving computationally expensive oracle evaluations.

A.2 Experimental Details for MORMA-SOBA

We adopt the same setup as described in , which can be summarized as follows.

Setup. We consider binary classification tasks generated from the FashionMNIST data set where we select 8 “easy” tasks (lowest loss $\sim 0.3$ from independent training) and 2 “hard” tasks (lowest loss $\sim 0.45$ from independent training) for multi-objective robust representation learning:

“easy” tasks: (0, 9), (1, 7), (2, 7), (2, 9), (4, 7), (4, 9), (3, 7), (3, 9)

For each task $i\in$ above, we partition its dataset into the training set $\mathcal{D}_{i}^{\text{train}}$, validation set $\mathcal{D}_{i}^{\text{val}}$, and test set $\mathcal{D}_{i}^{\text{test}}$. We also generate 7 (unseen) binary classification tasks for testing:

“easy” tasks: (1, 9), (2, 5), (4, 5), (5, 6)

We train a shared representation network that maps the 784-dimensional (vectorized 28x28 images) input to a 100-dimensional space. Subsequently, each task learns a binary classifier based on this shared representation. To learn a shared representation and per-task models that generalize well on each task, we aim to solve the following min-max bilevel optimization problem:

In the experiment, the regularization parameter in the LL function $\rho=5\times 10^{-4}$. The implementation of MORBiT follows the same manner described in . Specifically, the code of MORBiT uses vanilla SGD with a learning rate scheduler and incorporates momentum and weight decay techniques to optimize each variable:

Outer variable: learning rate = $0.01$, momentum = $0.9$, weight_decay = $10^{-4}$

Inner variable: learning rate = $0.01$, momentum = $0.9$, weight_decay = $10^{-4}$

Simplex variable: learning rate = $0.3$, momentum = $0.9$, weight_decay = $10^{-4}$

In addition, MORBiT adopts a straightforward iterative auto-differentiation to calculate the hypergradient without using the Neumann approximation of the Hession inversion.

For the implementation of MORMA-SOBA, the regularization parameter $\mu_{\lambda}$ in 11 is set to be $0.01$. All remaining parameters are chosen as constant values, as listed below:

Outer variable: $\tau_{x}=1,\alpha_{k}=0.02,$

Simplex variable: $\tau_{\lambda}=1,\alpha_{k}=0.02$

Both evaluated methods use batch sizes of 8 and 128 to compute $g_{i}$ for each inner step and $f_{i}$ for each outer iteration, respectively. In addition to Figure 2, which showcases the performance on 10 seen tasks used for representation learning, we present Figure 5. This figure displays the maximum/average loss values against the number of iterations on test sets consisting of 10 seen tasks and 7 unseen tasks. Our proposed approach, MORMA-SOBA, demonstrates superior performance in terms of faster reduction of both the maximum and average loss.

Appendix B Proofs

In addition, they are conditionally independent conditioned on $\mathscr{F}_{k}$.

Next we state some technical lemmas that will be used in both sections.

Suppose $f(x)$ is $\mu$-strongly convex and $L$-smooth. For any $x$ and $\gamma<\frac{2}{\mu+L}$, define $x^{+}=x-\gamma\nabla f(x),\ x^{*}=\operatorname*{arg\,min\,}f(x)$. Then we have

Suppose Assumption 1 holds. Then $\Phi(x)$ is differentiable and $\nabla\Phi(x)$ is given by Then $\Phi(x),y^{*}(x),z^{*}(x)$ are differentiable and $\nabla\Phi(x),y^{*}(x),z^{*}(x)$ are $L_{\nabla\Phi},L_{y^{*}},L_{z^{*}}$-Lipschitz continuous respectively, with their expressions as

where the first inequality uses triangle inequality, the second and third inequalities use Assumption 1, and the fourth inequality uses the Lipschitz continuity of $y^{*}(x)$. The inequality in (16) holds since $g(x,\cdot)$ is $\mu_{g}$-strongly convex and $\left\|\nabla_{2}f(x,y^{*}(x))\right\|\leq L_{f}$ (see Assumption 1). ∎

For any convex compact set $\mathcal{X}$, function $\eta_{\mathcal{X}}(x,h,\tau)$ defined in Section 3.3 is differentiable and $\nabla\eta_{\mathcal{X}}$ is $L_{\nabla\eta_{\mathcal{X}}}$-Lipschitz continuous, with the closed form exression and constant given by

For simplicity, we summarize the notations that will be used in Section B.1 as follows.

In this section we suppose Assumptions 1 and 2 hold. We assume stepsizes in Algorithm 1 satisfy

where $c_{1},c_{2},c_{3}>0$ are constants to be determined. We will utilize the following merit function in our analysis:

By definition of $\eta_{\mathcal{X}}$, we can verify that $W_{k,1}\geq 0$. Moreover, as discussed in Section 3.3, we consider the following optimality measure:

The following Lemma characterizes the relation between $V_{k}$ and gradient mapping of problem 1.

Suppose Assumptions 1 and 2 hold. In Algorithm 1 we have

where the first inequality uses Cauchy-Schwarz inequality and the second inequality uses the non-expansiveness of projection onto a convex compact set. This completes the proof. ∎

Next we present a technical lemma about the variance of $w^{k+1}$ and the bound for $\left\|h^{k+1}-h^{k}\right\|$.

Suppose Assumptions 1 and 2 hold. In Algorithm 1 we have

where the first equality uses independence, the first inequality uses Cauchy-Schwarz inequality, and the second inequality uses (16). This proves (22). Next for $\left\|h^{k+1}-h^{k}\right\|$ we have

which proves of (23) by taking expectation on both sides. ∎

Note that $D_{x}^{t}$ in is the same as our $w^{k+1}$ (see (7), line 5 of Algorithm 1 and definition of $w^{k+1}$ in (18)). The second moment bound can directly imply the variance bound, i.e.,

This implies that some stronger assumptions are needed to guarantee Assumption 3.7 in , as also pointed out by the authors (see discussions right below it). Instead, our refined analysis does not require that.

Note that Assumptions 3.1 and 3.2 in state that the upper-level function $f$ is twice differentiable, the lower-level function $g$ is three times differentiable and $\nabla^{2}f,\nabla^{3}g$ are Lipschitz continuous so that $z_{*}^{k}$, as a function of $x^{k}$ (see (18)), is smooth, which is a crucial condition for (63) - (67) in , which follows the analysis in Equation (49) in . In this section we show that, by incorporating the moving average technique recently introduced to decentralized bilevel optimization , we can remove this additional assumption. We have the following lemma characterizing the error induced by $y^{k}$ and $z^{k}$.

Suppose Assumptions 1 and 2 hold. If the stepsizes satisfy

We first consider the error induced by $y^{k}$. We have

where the first inequality uses Cauchy-Schwarz inequality:

Thanks to the moving average step of $x^{k}$, our analysis of $\left\|y_{*}^{k+1}-y_{*}^{k}\right\|$ is simplified comparing to that in . We also have

where the first inequality uses Assumption (2) and Lemma B.1, and the second inequality uses Lemma B.1 (which requires strong convexity of $g$, Lipschtiz continuity of $\nabla_{2}g$, and the first inequality in (24)). Combining (26) and (27), we know

where the second inequality uses $\beta_{k}<\frac{2}{\mu_{g}+L_{\nabla g}}\leq\frac{1}{\mu_{g}}$. Taking summation ($k$ from to $K$) on both sides and taking expectation, we know

which proves the first inequality in (25) by dividing $c_{1}\mu_{g}$ on both sides. Next we analyze the error induced by $z^{k}$. Our analysis is substantially different from . We first notice that

where we use Cauchy-Schwarz inequality in the first and second inequality, we use the facts that $\nabla y^{*}$ is Lipschitz continuous. For $\left\|z^{k+1}-z_{*}^{k}\right\|$, we may follow the analysis of SGD under the strongly convex setting:

where the first inequality uses Assumption 2, the second inequality uses Cauchy-Schwarz inequality and the definition of $z_{*}^{k}$, the third inequality uses Cauchy-Schwarz inequality and the fact that $g$ is $\mu_{g}$-strongly convex, and the fourth inequality uses Cauchy-Schwarz inequality, (16) and

which is a direct result from the bound of $\gamma_{k}$ in (24). It is worth noting that our estimation can be viewed as a refined version of (72) - (75) in Combining (29) and (30) we may obtain

where the equality uses the definition of $\sigma_{w}^{2}$ in (22) and the third inequality uses $\gamma_{k}\mu_{g}\leq\frac{1}{4}$. Taking summation ($k$ from to $K$) and expectation, we know

Suppose Assumptions 1 and 2 hold. We have

Note that we have the following decomposition:

which, together with Cauchy-Schwarz inequality, implies

B.1.2 Primal Convergence

The $L_{\nabla\Phi}$-smoothness of $\Phi(x)$ and $L_{\nabla\eta_{\mathcal{X}}}$-smoothness of $\eta_{\mathcal{X}}$ in Lemma B.2 and B.3 imply

where the first inequality uses $L_{\nabla\eta_{\mathcal{X}}}$-smoothness of $\nabla\eta_{\mathcal{X}}$, and the second inequality uses the optimality condition (17) (with $d=x^{k}$). Hence by computing \eqref{ineq: phi_decrease_simple}+\eqref{ineq: eta_x_decrease_simple}/c_{3} and taking conditional expectation with respect to $\mathscr{F}_{k}$ we know

where the second inequality uses Young’s inequality and the following inequalities:

where the second inequality uses (32). Taking summation and expectation on both sides of (36) and using (37), we obtain (33)

B.1.3 Dual Convergence

Suppose Assumptions 1 and 2 hold. In Algorithm 1 we have

Note that by moving average update of $h^{k}$, we have

where the first equality uses the fact that $x^{k},h^{k},x^{k+1},$ are all $\mathscr{F}_{k}$-measurable and are independent of $w^{k+1}$ given $\mathscr{F}_{k}$, the first inequality uses the convexity of $\left\|\cdot\right\|^{2}$ and (22), the second inequality uses Cauchy-Schwarz inequality, the third inequality uses the Lipschitz continuity of $\nabla\Phi$ in Lemma B.10, and the update rules of $x^{k+1}$. Taking summation, expectation on both sides of (40), dividing $c_{3}$ and using (22), we know (38) holds. ∎

B.1.4 Proof of Theorem 1

Now we are ready to prove Theorem 1. From Lemma B.4 we know it suffices to bound $V_{k}$. By definition of $V_{k}$ in (21), (33) and (38) we have

in the second inequality. The constants are defined as

Using constants defined in Lemma B.6, we know

so that the conditions ((24), (32) and (41)) in previous lemmas hold. Then we have

which, together with Lemma B.4, proves Theorem 1.

B.2 Proof of Theorem 2

In this section we present our proof of Theorem 2. For simplicity, we summarize the notations that will be used in our proof as follows.

In this subsection we suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. We suppose stepsizes in Algorithm 2 satisfy

where $c_{1},c_{2},c_{3}>0$ are constants to be determined. We will utilize the following merit function in our analysis:

The following lemma provides some smoothness of functions that we will use in our proof.

Functions $\nabla\Psi(\cdot),\nabla_{1}\Phi_{\mu_{\lambda}}(\cdot,\lambda),\nabla_{1}\Phi(\cdot,\lambda),\nabla_{1}\Phi_{\mu_{\lambda}}(x,\cdot),\nabla_{1}\Phi(x,\cdot),\nabla_{2}\Phi_{\mu_{\lambda}}(\cdot,\lambda),\\ \nabla_{2}\Phi_{\mu_{\lambda}}(x,\cdot)$ are $L_{\nabla\Psi},L_{\nabla\Phi},L_{\nabla\Phi},L_{\nabla_{1}\Phi_{\mu_{\lambda}}},L_{\nabla_{1}\Phi_{\mu_{\lambda}}},L_{\nabla_{2}\Phi_{\mu_{\lambda}}},\mu_{\lambda}$-Lipschitz continuous respectively, with the constants given by

For $\nabla\Psi$ we first notice that the nonconvex-strongly-concave problem in (11) can be reformulated as a bilevel problem:

from which we know $\nabla\Psi(\cdot)$ is $L_{\nabla\Psi}$-Lipschitz continuous since

(47), (48) and (49) imply $\nabla_{1}\Phi_{\mu_{\lambda}}(\cdot,\lambda),\nabla_{1}\Phi(\cdot,\lambda)$ are $L_{\nabla\Phi}$-Lipschitz continuous and $\nabla_{1}\Phi_{\mu_{\lambda}}(x,\cdot),\nabla_{1}\Phi(x,\cdot)$ are $L_{\nabla_{1}\Phi_{\mu_{\lambda}}}$-Lipschitz continuous. Finally, for $\nabla_{2}\Phi_{\mu_{\lambda}}(x,\lambda)$ we have

and thus $\nabla_{2}\Phi_{\mu_{\lambda}}(\cdot,\lambda),\nabla_{2}\Phi_{\mu_{\lambda}}(x,\cdot)$ are $\sqrt{n}L_{\Phi},\mu_{\lambda}$-Lipschitz continuous respectively. ∎

Next we present a technical lemma that will be used in analyzing the strongly convex function over a convex compact set.

Suppose $f(x)$ is $\mu$-strongly convex and $L$-smooth over a convex compact set $\mathcal{X}$. For any $\tau\leq\frac{1}{L}$ define $x_{+}=\Pi_{\mathcal{X}}(x-\tau\nabla f(x))$ and $x_{*}=\operatorname*{arg\,min\,}_{x\in\mathcal{X}}f(x)$, we have

where $r=\frac{\mu}{\frac{1}{\tau}+\mu}\leq\frac{1}{2}$. Applying Young’s inequality to the left hand side of the above inequality, we know

Suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. If $\tau_{\lambda}\mu_{\lambda}=1$, then in Algorithm 2 we have

where the first inequality uses Cauchy-Schwarz inequality and the second inequality uses the non-expansiveness of projection onto a convex compact set. Recall that

which is a minimizer (over the probability simplex) of a $\mu_{\lambda}$-smooth and $\mu_{\lambda}$-strongly convex function $\Phi_{\mu_{\lambda}}(x^{k},\cdot)$. Hence we know from Lemma B.11 that

where the second inequality uses Cauchy-Schwarz inequality and the third inequality uses non-expansiveness of the projection onto a convex compact set. Setting $\tau_{\lambda}\mu_{\lambda}=1$ completes the proof. ∎

Suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. In Algorithm 2 we have

which proves (50). Next for $\left\|h_{x}^{k+1}-h_{x}^{k}\right\|$ we have

which proves the first inequality of (51). Similarly we have

which proves the second inequality of (51). ∎

Suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. In Algorithm 2 if the stepsizes satisfy

where constants $C_{yx},C_{y,1},C_{zx},C_{z,1}$ are defined the same as those in Lemma B.6. $C_{y_{i},0},C_{z_{i},0}$ are defined as

Note that the proof follows almost the same reasoning in Lemma B.6. Since Assumptions 1 and 2 hold for all $f_{i},g_{i}$, by replacing $y^{k},y_{*}^{k},z^{k},z_{*}^{k}$ with $y_{i}^{k},y_{*,i}^{k},z_{i}^{k},z_{*,i}^{k}$ respectively, we have similar results hold for each $1\leq i\leq n$

Taking summation on both sides of (55), we complete the proof.

Suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. We have

Note that we have the following decomposition:

which, together with Cauchy-Schwarz inequality, implies

Applying Cauchy-Schwarz inequality, Assumption 1 and Lemma B.10 to the above equation and (56), we know

which together with Lemma B.12 completes the proof. ∎

B.2.2 Primal Convergence

Suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. If

where the second inequality uses (57). Taking summation and expectation on both sides of (59) and using (60), we obtain the first inequality in (58). For the second inequality in (58), the $L_{\nabla\Psi}$-smoothness of $\Psi(x)$ and $L_{\nabla\eta_{\mathcal{X}}}$-smoothness of $\eta_{\mathcal{X}}$ in Lemma B.10 imply

where the inequality uses Lemma B.10 and $(c)$ in Assumption 1 to obtain

Taking conditional expectation with respect to $\mathscr{F}_{k}$ on $\eqref{ineq: phi_decrease}+\eqref{ineq: eta_lam_decrease_exp}/c_{3}+\eqref{eq: F_decrease}$, we know

where the second inequality uses Lemma B.10, and the third inequality uses Young’s inequality and the conditions on $\alpha_{k}$ (see (57)):

in (57). Combining (65), (66), and (67), we have

which implies the second inequality in (58) by taking summation. ∎

B.2.3 Dual Convergence

Suppose Assumptions 1, 2 hold for all $f_{i},g_{i}$ and Assumption 3 holds. In Algorithm 2 we have

where the first equality uses the fact that $x^{k},\lambda^{k},h_{x}^{k},x^{k+1},\lambda^{k+1}$ are all $\mathscr{F}_{k}$-measurable and are independent of $w^{k+1}$ given $\mathscr{F}_{k}$, the first inequality uses the convexity of $\left\|\cdot\right\|^{2}$ and (50), the second inequality uses Cauchy-Schwarz inequality, the third inequality uses the Lipschitz continuity of $\nabla_{1}\Phi$ in Lemma B.10, and the update rules of $x^{k+1}$ and $\lambda^{k+1}$. Taking summation, expectation on both sides of (71), dividing $c_{3}$, and applying (50), we know the first inequality in (69) holds.

where the second equality uses $\nabla_{2}\Phi_{\mu_{\lambda}}(x^{k},\lambda^{k})=\nabla_{2}\Phi(x^{k},\lambda^{k})-\mu_{\lambda}\left(\lambda^{k}-\frac{\mathbf{1}_{n}}{n}\right)$. Hence we know

where the third inequality uses Lemma B.10 and the fact that

Taking summation, expectation on both sides of (73), and dividing $c_{3}$, we know the second inequality in (69) holds. ∎

B.2.4 Proof of Theorem 2

According to the definition of the constants in Lemmas B.6 and B.14, we could obtain (for simplicity we omit the dependency on $\kappa$ here)

Hence we can pick $\alpha_{k},c_{1},c_{2},c_{3},\tau_{x},\tau_{\lambda}$ such that

and the conditions ((53), (57), and (76)) in previous lemmas hold. Moreover, using the above conditions in (77) and (79), we can get

Combining the above two inequalities, we have

which completes the proof of Theorem 2 since we have

where the second inequality uses non-expansiveness of projection operator and $\sqrt{n}L_{\Phi}$-Lipschitz continuity of $\nabla_{1}\Phi_{\mu_{\lambda}}(x,\cdot)$ in Lemma B.10. Note that we have $n^{2}$ in the numerator since we explicitly write out the Lipschitz constant $L_{\nabla_{1}\Phi_{\mu_{\lambda}}}$.

Appendix C Discussions on the Prior Work [29]

In this section, we discuss several issues in the current form of , which introduces a Multi-Objective Robust Bilevel Two-timescale optimization algorithm (MORBiT).