# Min-Max Bilevel Multi-objective Optimization with Applications in Machine
Learning

Alex Gu<EMAIL_ADDRESS>
Massachusetts Institute of Technology, Cambridge, MA, USA
Songtao Lu<EMAIL_ADDRESS>
IBM Thomas J. Watson Research Center, Yorktown Heights, NY, USA
Parikshit Ram<EMAIL_ADDRESS>
IBM Thomas J. Watson Research Center, Yorktown Heights, NY, USA
Lily Weng<EMAIL_ADDRESS>
University of California San Diego, La Jolla, CA, USA

###### Abstract

This paper is the first to propose a generic min-max bilevel multi-objective
optimization framework, highlighting applications in representation learning
and hyperparameter optimization. In many machine learning applications such as
meta-learning, multi-task learning, and representation learning, a subset of
the parameters are shared by all the tasks, while each specific task has its
own set of additional parameters. By leveraging the recent advances of
nonconvex min-max optimization, we propose a gradient descent-ascent bilevel
optimization (MORBiT) algorithm which is able to extract a set of shared
parameters that is robust over all tasks and further overcomes the
distributional shift between training and testing tasks. Theoretical analyses
show that MORBiT converges to the first-order stationary point at a rate of
$\mathcal{O}(\sqrt{n}K^{-2/5})$ for a class of nonconvex problems, where $K$
denotes the total number of iterations and $n$ denotes the number of tasks.
Overall, we formulate a min-max bilevel multi-objective optimization problem,
provide a single loop two-timescale algorithm with convergence rate
guarantees, and show theoretical bounds on the generalization abilities of the
optimizer. Experimental results on sinusoid regression and representation
learning showcase the superiority of MORBiT over state-of-the-art methods,
validating our convergence and generalization results.

##  1 Introduction

We begin by examining the generic bilevel optimization (BO) problem:

$\min_{x\in\mathcal{X}\subseteq\mathbb{R}^{d_{1}}}f(x,y^{\star}(x))\;\text{s.t.}\;y^{\star}(x)=\operatorname*{arg\,min}_{y\in\mathbb{R}^{d_{2}}}g(x,y).$
(1)

Here, $f$ is the outer-level objective function, $g$ is the inner level
objective function, and $\mathcal{X}\in\mathbb{R}^{d_{1}}$ denotes the outer-
level input space. (1) is called bilevel optimization because the outer
objective function $f$ depends on the solution to the inner objective $g$.
Recently, BO has found a myriad of applications in modern machine learning:
hyperparameter optimization (Franceschi et al., 2018), reinforcement learning
(Hong et al., 2020), personalized federated learning (Fallah et al., 2020),
multi-task learning (MTL), representation learning Arora et al. (2020), etc.
We refer the reader to (Liu et al., 2021) for a comprehensive survey of
applications of BO in machine learning and computer vision.

Our first contribution is in considering a min-max generalized version of (1)
to the multi-objective setting, where we have $n$ different objective function
pairs $(f_{i},g_{i})$. Let, $[n]\triangleq\\{1,2,\cdots,n\\}$ and
$f_{i}(\cdot,\cdot)$, $g_{i}(\cdot,\cdot)$ denote the $i$th outer and inner
objectives, respectively. We consider the following problem:

$\displaystyle\min_{x\in\mathcal{X}\subseteq\mathbb{R}^{d_{1}}}\max_{i\in[n]}$
$\displaystyle\quad f_{i}(x,y_{i}^{\star}(x))$ subject to $\displaystyle\quad
y_{i}^{\star}(x)=\operatorname*{arg\,min}_{y_{i}\in\mathbb{R}^{d_{2}^{i}}}g_{i}(x,y_{i}),i\in[n]$
(2)

To understand the necessity and benefits of our formulation, we present a
motivating example, multi-task representation learning (MTRL). In MTRL, we
have $n$ tasks $\mathcal{T}_{1},\cdots,\mathcal{T}_{n}$, where task
$\mathcal{T}_{i}$ consists of data sampled from a distribution
$\mathcal{D}_{i}$. Let $\mathcal{X}$, $\mathcal{Y}$ be the input and label
spaces respectively. We have a meta-training set $D=\\{D_{i},i\in[n]\\}$,
where each $D_{i}\triangleq D_{i}^{\text{tr}}\cup D_{i}^{\text{test}}$
consists of two data batches of $K$ training and $J$ testing points from
$\mathcal{X}\times\mathcal{Y}$, respectively. In representation learning, we
aim to learn a set of shared representation parameters common to all the
tasks. For each task, we start with the representation parameters and fine-
tune additional task-specific parameters. We denote by $x$ the shared
representation parameters and by $y\triangleq\\{y_{i},i\in[n]\\}$ the task-
specific parameters. Concretely, consider an two-task animal classification
setting, where task $\mathcal{T}_{1}$ aims to distinguish various breeds of
dogs and task $\mathcal{T}_{2}$ aims to distinguish various breeds of cats.
The shared representation parameters $x$ might learn high-level animal
features (e.g., the presence of facial features), while the task-specific
parameters $y_{1}$ might learn features important for distinguishing different
breeds of dogs (e.g., fur patterns). Let $\hat{\ell}_{i}(x,y;D)$ denote the
loss function for task $i$ on data $D$. In this case, we have

$f(x,y)=\sum_{i=1}^{n}\ell_{i}(x,y_{i};D_{i}^{\text{test}}),\;g(x,y)=\sum_{i=1}^{n}\ell_{i}(x,y_{i};D_{i}^{\text{tr}}),$

and solving the optimization problem (1) would give us an optimal shared
representation $x$. A critical issue with this approach is that $f$ and $g$
are measuring the average loss over all tasks, which can lead to a solution
with good average performance but poor worst-case performance. This is
problematic in safety-critical scenarios such as few-shot facial recognition.
If each task represents a different population demographic, a representation
$x$ with a poor worst-case loss but good average-case loss can be very biased
towards a certain demographic. Collins et al. (2020) highlights this problem
in the meta-learning setup, showing that popular algorithms such as MAML (Finn
et al., 2017) indeed show poor worst-case performance. To mitigate this
problem, they propose a task-robust MAML objective to minimize the worst-case
loss rather than the average-case loss.

Inspired by this work, our optimization objective, (1), takes this idea and
generalizes it from the meta-learning setup to a general BO setup. Our
objective allows for the specification of multiple tasks $f_{i}$ and optimizes
over the worst-case objective $\max_{i\in[n]}f_{i}(x,y_{i}^{\star}(x))$. For
example, in the representation learning case, letting
$f_{i}(x,y_{i})=\ell_{i}(x,y_{i};D_{i}^{\text{test}})$ and
$g_{i}(x,y_{i})=\ell_{i}(x,y_{i};D_{i}^{\text{tr}})$, solving (1) allows us to
learn a task-robust representation $x$. Compared to (Collins et al., 2020),
our objective captures a much more general setting of BO. In fact, they do not
analyze their algorithm as a bilevel problem, but rather as a nested
objective, $f(x-\alpha\nabla f(x))$, but their objective can be written as a
special case of ours, with $g_{i}(x,y)=\ell(y)+\frac{1}{2}\|x-y\|^{2}$.

Our contributions are three-fold:

1\. (Section 3) First, we propose a novel min-max formulation of the BO
problem, (1), the first to capture our robust multi-objective desiderata in
BO. When $n=1$, we recover the traditional BO problem (1). We highlight two
key features of our new formulation. First, it allows for multiple objectives,
both inner and outer, giving significantly more flexibility: in a multi-task
setting where inner-loop tasks differ greatly, our formulation allows using
different algorithms $g_{i}$ for each task. Second, in a setting where tasks
are similar, our formulation allows for task-level customization such as
different inner-loop step-sizes for each task. Second, our min-max formulation
ensures that the BO solution is robust, ensuring all $n$ functions achieve a
low value.

2\. (Sections 4, 5, and 6) Second, we propose a simple and intuitive algorithm
to solve (1). Our algorithm, MORBiT (_M_ ulti-_O_ bjective _R_ obust bilevel
optimization with _Bi_ \- _T_ imescale algorithm), is a single-loop, bi-
timescale, stochastic algorithm based on the standard gradient descent-ascent
algorithm. We remark that our algorithm is a generalization of the TTSA
algorithm (Hong et al., 2020) for solving the single-objective bilevel
optimization problem. We also show, theoretically, that MORBiT converges to
the first-order stationary point of a general weakly convex min-max bilevel
optimization problem at a rate of $\mathcal{O}(\sqrt{n}K^{-2/5})$ when the
outer objective is weakly convex and the inner objectiver is strongly convex.
As special cases, we recover the convergence rates of two existing works, TTSA
and TR-MAML, but our algorithm is able to capture a much wider set of bilevel
optimization problems.

3\. (Sections 7 and 8) Third, we extend the generalization result from
(Collins et al., 2020), giving a Rademacher complexity bound on the
generalization error of new tasks in the convex hull of existing meta-training
tasks. In addition, we validate the superiority of our formulation, as well as
convergence and generalization results, on sinusoid regression and
representation learning settings.

##  2 Related Work

Bilevel Optimization Theory:  In the past few years, there have been a flurry
of new algorithms to solve the single-objective bilevel optimization problem
(1). Ghadimi and Wang (2018) and Ji et al. (2020) propose BSA and stocBiO
algorithms respectively that solve (1) in a double-loop manner, meaning that
for each outer step, multiple inner SGD steps are taken to better approximate
$y^{*}(x)$. These works consider the setting where the outer problem is
nonconvex and the inner problem is strongly convex. One drawback to double
loop algorithms is that selecting step-sizes and termination criterion for
inner loop steps is tricky. Motivated by this, Chen et al. (2021a), Hong et
al. (2020), and Chen et al. (2021b) propose STABLE, TTSA, and ALSET,
respectively, as single-loop algorithms that solve (1).

There has been little work analyzing multi-objective bilevel optimization
schemes (Sinha et al., 2015; Deb and Sinha, 2009; Ji et al., 2017). All of
these works analyze the multi-objective from a game theoretic point of view,
using a vector-valued objective with the notion of Pareto optimality. In
contrast, we are the first to analyze the bilevel multi-objective problem from
a traditional optimization perspective and propose a min-max task-robust
objective to ensure fast convergence across all objectives.

Bilevel Optimization in Machine Learning:  Recently, bilevel optimization has
gained popularity in applications such as meta-learning, representation
learning, and hyperparameter optimization (HPO). In meta-learning, Rajeswaran
et al. (2019) proposes iMAML, modeling few-shot learning as a bilevel problem.
The outer level learns parameters that optimally capture the overall structure
of all the tasks, while the inner level optimally adapts to the few-shot
samples. In representation learning, Arora et al. (2020) considers a bilevel
problem where the outer level learns an optimal shared task representation and
the inner level learns optimal task-specific parameters. In HPO, Franceschi et
al. (2018) considers a bilevel problem where the outer level learns the
optimal hyperparameters, while the inner level learns the optimal model for
those hyperparameters. While meta-learning and representation learning are
inherently multi-objective settings, they are treated in practice as a single-
objective problem by summing the objectives over all tasks. Our formulation
gives the flexibity to treat these inherently multi-objective problems as
such, obtain stronger guarantees and ensuring that the loss of each individual
task, rather than just the average loss, converges quickly.

Minimax Optimization:  Minimax optimization is commonly employed to make
learned models more robust, such as in distributionally robust optimization
(DRO) and adversarially robust learning. In DRO, Duchi and Namkoong (2018) and
Shalev-Shwartz and Wexler (2016) showed that minimizing the maximum loss
improves generalization due to a variance regularization effect. Qian et al.
(2019) considers a robust model over multiple domains, showing convergence
rates for both convex and non-convex losses. In MTL, DRO is shown to have
strong generalization guarantees for unseen tasks in (Mehta et al., 2012),
providing a theoretical motivation for robust MTL. In the context of
adversarial learning and robustness, we also consider a minimax problem, where
the minimum is over the loss and the maximum is over the worst-case
perturbation to inputs (Madry et al., 2017; Wang et al., 2019). Cheng et al.
(2020) customizes the perturbation level for each training sample through a
minimax objective. In both DRO and robust learning, the minimax objective is
in the form $\min_{x}\max_{y}f(x,y)$. Observe that this is a special case of
the single-objective BO with independent lower and upper levels. In contrast,
we study a much more general multi-objective BO problem where the upper level
objective is dependent on the solution to the lower level objective. One key
work that we highlight once more is Task-Robust MAML (TR-MAML) (Collins et
al., 2020), which provides a robust meta-learning algorithm for few-shot
tasks. They use a minimax objective (1), posing their optimization objective
as a min-max single level multi-objective problem. In contrast, we focus on
the min-max bilevel multi-objective BO problem.

##  3 Problem Formulation

In this work, we consider the min-max bilevel multi-objective optimization
problem (1). Since (1) cannot be solved efficiently when $f$ is nonconvex, we
reformulate (1) as follows:

$\displaystyle\min_{x\in\mathcal{X}}\max_{\lambda\in\Delta_{n}}\quad$
$\displaystyle\sum_{i=1}^{n}\lambda_{i}f_{i}(x,y_{i}^{\star}(x))$ s.t.
$\displaystyle
y_{i}^{\star}(x)=\operatorname*{arg\,min}_{y_{i}\in\mathbb{R}^{d_{2}^{i}}}g_{i}(x,y_{i}),i\in[n]$
(3)

where $\Delta_{n}$ is the simplex
$\Delta_{n}\subseteq\mathbb{R}^{n}=\\{\lambda\in\mathbb{R}^{n}|\lambda_{1}+\cdots+\lambda_{n}=1,\lambda_{i}\geq
0\\}$. Next, we show one motivating example below as an application of our new
framework.

Multi-Task Representation Learning with Common Embedding:  Consider a
collection of tasks $\mathcal{T}_{1},\mathcal{T}_{2},\cdots,\mathcal{T}_{n}$,
and consider the problem of learning an permutation-invariant representation
feature space among all the tasks. Partition a neural network into two pieces,
an embedding network and a task-specific network. The layers in the embedding
network learn an invariant representation feature space that is common to all
the tasks, while the task-specific layers learn features specific to that
task. During training, the goal is to find a good set of representation
features so that new tasks can be learned quickly. Such a network has been
shown to generalize well for unseen tasks, as everything but the last layer is
learned across all the tasks. This idea has been used in both few-shot
learning, where the task-specific network is just the last layer (Raghu et
al., 2019), and in continual learning (Javed and White, 2019).

In the context of our framework, $f_{i}$ denotes the loss for task
$\mathcal{T}_{i}$, $x$ denotes the weights of the embedding network, and
$y_{i}$ denotes the task-specific weights for task $i$. We remark that the
framework in Collins et al. (2020) is unable to handle this problem because
their framework requires the problem dimension of each of the inner tasks to
be identical, i.e., $d_{2}^{i}=d_{2}^{j}\forall i,j\in[n]$. However, in our
framework, $y_{i}\in\mathbb{R}^{d_{2}^{i}}$ may have different dimensions for
each task $i$.

Notation:  We use $K$ to denote the number of iterates and $k$ as the
iteration count. We also use $n$ to denote the number of tasks and $i$ to be
the task identifier. The dimensions of the variables are
$x\in\mathcal{X}\in\mathbb{R}^{d_{1}}$, $y\in\mathbb{R}^{d_{2}}$, and
$\lambda\in\mathbb{R}^{n}$.

We call $f_{i}(\cdot,\cdot)$’s outer objectives and $g_{i}(\cdot,\cdot)$’s
inner objectives. We denote the inner optimum of the $i$th task as
$y_{i}^{\star}(x)\triangleq\max_{y}g_{i}(x,y)$. We also define a shorthand for
the outer objective $\ell_{i}(x)=f_{i}(x,y_{i}^{\star}(x))$. Finally, we
define the weighted objectives $F(x,\lambda)$ and $F(x,y,\lambda)$ to be

$F(x,\lambda)=\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x),\;F(x,y,\lambda)=\sum_{i=1}^{n}\lambda_{i}f_{i}(x,y).$
(4)

Gradient Estimates:  Now, we define the notation we use for gradients and
gradient approximations. Based on Lemma 2.1 from Ghadimi and Wang (2018)), we
can compute $\nabla\ell_{i}(x)$ exactly:
$\nabla\ell_{i}(x)=\nabla_{x}f_{i}(x,y_{i}^{\star}(x))-\nabla^{2}_{xy}g_{i}(x,y_{i}^{\star}(x))[\nabla^{2}_{yy}g_{i}(x,y_{i}^{\star}(x))]^{-1}\nabla_{y}f_{i}(x,y_{i}^{\star}(x))$.
Let $y$ denote $\\{y_{i}\\}_{i\in[n]}$. In practice, since we cannot compute
$y_{i}^{\star}(x)$ exactly, we instead define an approximation of the gradient
estimate, replacing $\\{y_{i}^{\star}(x)\\}_{i\in[n]}$ with
$\\{y_{i}\\}_{i\in[n]}$:

$\mkern 1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x,y)=\nabla_{x}f_{i}(x,y)-\nabla^{2}_{xy}g_{i}(x,y)[\nabla^{2}_{yy}g_{i}(x,y)]^{-1}\nabla_{y}f_{i}(x,y).$

Similarly, we can also define the exact and approximate gradients of the full
loss $F(x;y;\lambda)$ as follows:

$\displaystyle\nabla_{x}F(x,\lambda)$
$\displaystyle=\sum_{i=1}^{n}\lambda_{i}\nabla\ell_{i}(x),$ (5a)
$\displaystyle\nabla_{\lambda}F(x,\lambda)$
$\displaystyle=[\ell_{1}(x),\cdots,\ell_{n}(x)]^{\top},$ (5b)
$\displaystyle\mkern 1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x;y;\lambda)$ $\displaystyle=\sum_{i=1}^{n}\lambda_{i}\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern 1.5mu_{x}f_{i}(x,y^{i}),$
(5c) $\displaystyle\mkern 1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{\lambda}F(x;y;\lambda)$
$\displaystyle\\!=\\![f_{1}(x,y^{1}),\\!\cdots,\\!f_{n}(x,y^{n})]^{\top}.$
(5d)

Stochastic Estimates:  We will present our formal definitions in Section 5.1.
We use the letter $h$ to denote stochastic estimates of quantities, e.g.
$h_{g}^{k,i},h_{x}^{k},h_{\lambda}^{k}$. We use $B_{k}$ and $B_{k}$ to denote
the absolute error between $h_{x}^{k}$ and its estimate, and $b_{k}$ as an
upper bound of $\|B_{k}\|$. We also use $\sigma$ to denote quantities related
to squared error, e.g. $\sigma_{f}$, $\sigma_{g}$, and $\tilde{\sigma}_{f}$.

##  4 Algorithm Design

###  4.1 Algorithm

We first introduce our algorithm for min-max bilevel optimization and then
provide intuition on our algorithm.

for  _$k=1,2,\cdots,K$_ do

for  _$i=1,2,\cdots,n$_ do

$y_{k+1}^{i}=y_{k}^{i}-\beta h_{g}^{k,i}$

end for

$x_{k+1}=\mathrm{proj}_{\mathcal{X}}(x_{k}-\alpha h_{x}^{k})$;

$\lambda_{k+1}=\mathrm{proj}_{\Delta_{n}}(\lambda_{k}+\gamma
h_{\lambda}^{k})$;

end for

pick  $\tau\in\\{1,2,\cdots,K\\}$ uniformly at random;

return _
$\bar{x}=x_{\tau},\bar{y}_{i}=y_{\tau-1}^{i},\bar{\lambda}=\lambda_{\tau}$_;

Algorithm 1 MORBiT (_M_ ulti-_O_ bjective _R_ obust bilevel optimization with
_Bi_ \- _T_ imescale algorithm)

Our algorithm, MORBiT (Algorithm 1), tracks three sets of iterates:
$x_{k},y_{k}=\\{y_{k}^{1},y_{k}^{2},\cdots,y_{k}^{n}\\}$, and $\lambda_{k}$.
Intuitively, at each step, $x_{k}$ is estimating the outer-level solution and
$y_{k}^{i}$ is estimating an approximation of the inner loop solution. We
expect $y_{k}^{i}\approx y_{i}^{\star}(x_{k})$, and we will show that
theoretically, these quantities are close. Our novelty is the introduction of
the $\lambda_{k}$ terms, representing weights on each of the objectives. Our
algorithm uses the $\lambda$ terms to put more emphasis on tasks that haven’t
been learned well.

Next, we describe the procedure for each iteration $k=1,2,\cdots,K$. First,
$\alpha,\beta,\gamma$ denote step-sizes for updating $x$, $y$, and $\lambda$
iterates, respectively. In each iteration, we perform an inner and outer step.
First, we perform the inner-level updates, taking one SGD step on each of the
inner functions $g_{i}$, $y_{k+1}^{i}=y_{k}^{i}-\beta h_{g}^{k,i}$. Here, we
use the stochastic gradient
$h_{g}^{k,i}\approx\bar{\nabla}_{y}g_{i}(x_{k},y_{k}^{i})$. Next, we perform
the outer-level updates on $x_{k}$ and $\lambda_{k}$. We take a descent step
on $x_{k}$ and an ascent step on $\lambda_{k}$ using the stochastic gradient
approximations of the weighted loss, $h_{x}^{k}\approx\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})$ and $h_{\lambda}^{k}\approx\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{\lambda}F(x_{k},y_{k+1},\lambda_{k})$. Note that we are considering a
constrained setting $x\in\mathcal{X}$ and $\lambda$ is restricted to be in the
simplex, so we also perform projections at each step. These updates are fairly
standard in min-max optimization (eg, Lu et al. (2020); Lin et al. (2020)). We
will formally list our assumptions on $h_{g}^{k,i}$,
$h_{x}^{k},h_{\lambda}^{k}$ in Section 5.1.

We remark that our algorithm is a single-loop, two-timescale algorithm in the
sense that at each outer iteration, we only take one gradient step for each
iterate. Theoretically, we choose our step-sizes as
$\alpha=\mathcal{O}(K^{-3/5})$, $\beta=\mathcal{O}(K^{-2/5})$,
$\gamma=\mathcal{O}(K^{-3/5}n^{-1/2})$. We highlight the two-timescale nature
of the step-sizes, where for a fixed number of tasks $n$, the $\beta$ decays
slower than $\alpha$ and $\gamma$.

###  4.2 Connections to Existing Work

Connection to TTSA (Hong et al., 2020):  When $n=1$, Algorithm 1 reduces to an
algorithm where at each step, $y_{k+1}=y_{k}-\beta h_{g}^{k}$ and
$x_{k+1}=\mathrm{proj}_{\mathcal{X}}(x_{k}-\alpha h_{x}^{k})$. The $\lambda$
updates go away because $\lambda_{k}=1$ for all $k$. This algorithm, then, is
exactly the TTSA algorithm. In the multi-task case where $n>1$, our algorithm
is a natural generalization of TTSA, with $y$ updates behaving identically and
$\lambda$ updates inspired by gradient descent-ascent algorithms such as Lin
et al. (2020). Our framework expands upon TTSA by allowing the optimization of
multiple different functions $g_{i}$ at once.

Connection to TR-MAML (Collins et al., 2020):  Now, consider a setting where
instead of stochastic gradients, we choose a full batch size and set
$g_{i}(x,y)=\ell_{i}(x,y)+\frac{\eta}{2}\|y-x\|^{2}$. This gives us the one-
step update $y_{i}^{\star}(x)=x-\eta\nabla\ell_{i}(x,y)$, which is exactly the
formulation of TR-MAML. Therefore, our algorithm generalizes TR-MAML but adds
flexibility, varying the step-size across tasks or using different update
algorithms for individual tasks.

##  5 Theoretical Guarantees

Our main theoretical result is shown in Theorem 1, which claims a convergence
rate of $\mathcal{O}(\sqrt{n}K^{-2/5})$ for achieving an approximate
stationary and optimal solutions with respect to variables $x$ and
$y,\lambda$. Overall, our analysis follows a similar overall structure as the
analysis in Hong et al. (2020). However, we highlight a few unique differences
due to the task-robust nature of our formulation and algorithm. First, as in
TTSA, the convergence of $x$’s relies on the $y$ being close to $y^{*}(x)$.
However, we are trying to fit one $x$ across multiple functions $f_{i}$. In
our analysis, therefore, convergence is limited by the worst $y_{i}$, i.e.
$\max_{i\in[n]}\|y_{i}-y_{i}^{\star}(x)\|$. Second, since objective values are
weighted by $\lambda$, it may still be small for suboptimal $(x,y)$ if
$\lambda$ does not converge. Therefore, we also require a convergence result
for $\lambda$. To do so, we take a similar approach by leveraging the
theoretical result in Collins et al. (2020).

###  5.1 Assumptions

In this section, we informally list our assumptions. The full list of
assumptions can be found in the Appendix.

Regularity Assumptions:  In this work, we assume our outer functions are
$\mu_{\ell}$-weakly convex, bounded, Lipschitz, and Lipschitz-continuous,
while we assume our inner functions are Lipschitz continuous and
$\mu_{g}$-strongly convex. We use the same assumption here as in Hong et al.
(2020), with the addition of:

$\mathbb{E}[h_{\lambda}^{k}]=\nabla_{\lambda}F(x_{k},\lambda_{k})=\left[f_{1}(x_{k},y_{k+1}^{1}),\cdots,f_{n}(x_{k},y_{k+1}^{n})\right]^{\top}.$
(6)

Notation for Regularity Constants:  We will formally define our regularity
constants in Assumptions 1 & 2 in the appendix and in Lemma 1. In general, we
use the letter $L$ to denote Lipschitz continuity (smoothness) constants, e.g.
$L_{f_{x}}$, $L_{f_{y}}$, $\bar{L}_{f_{y}}$, $L_{g}$, $L_{gxy}$, $L_{gyy}$,
$\bar{L}_{gxy}$, $\bar{L}_{gyy},L,L_{f}$. For functions with multiple
arguments, we distinguish the Lipschitz continuity w.r.t. $y$ and $x$ with $L$
and $\bar{L}$, respectively. We use the letter $G$ to denote Lipschitz
constants, e.g. $G_{f},G_{y}$. We use the letter $B$ to denote bounds on
function values, in the case of $B_{\ell}$. We use the letter $C$ to denote
bounds on gradients and Hessians, e.g. $C_{f_{y}},C_{\ell},C_{gxy}$. Finally,
we use $\mu$ to denote convexity constants, e.g. $\mu_{\ell},\mu_{g}$.

Boundedness of stochastic gradient estimates: For any $k\geq 0$, we assume
that there exists a nonincreasing sequence $\\{b_{k}\\}_{k\geq 0}$ such that

$\displaystyle\mathbb{E}[h_{x}^{k}]$ $\displaystyle=\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+B_{k},\|B_{k}\|\leq b_{k},$ (7)

where $h_{g}^{k,i}$ is an unbiased estimator of
$\nabla_{y}g_{i}(x_{k},y_{k}^{i})$.

In addition, we assume some conditions on the variances, so that there exist
$\sigma_{f},\sigma_{g}$ satisfying

$\displaystyle\mathbb{E}[\|h_{x}^{k}-\mathbb{E}[h_{x}^{k}]\|^{2}]$
$\displaystyle\leq\sigma_{f}^{2},$
$\displaystyle\mathbb{E}[\|h_{g}^{k,i}-\nabla_{y}g_{i}(x_{k},y_{k})\|^{2}]$
$\displaystyle\leq\sigma_{g}^{2}(1+\|\nabla_{y}g_{i}(x_{k},y_{k})\|^{2}).$

###  5.2 Convergence Results

Our main results are the following convergence rates on the output of
Algorithm 1 , $\bar{x},\bar{y_{i}},\bar{\lambda}$, to the stationary points of
problem (3):

$\displaystyle\mathbb{E}\left[\max_{i\in[n]}\|\bar{y}_{i}-y_{i}^{\star}(\bar{x})\|^{2}\right]$
$\displaystyle\leq\mathcal{O}(\sqrt{n}K^{-2/5}),$ (8)
$\displaystyle\max_{\lambda}\mathbb{E}[F(\bar{x},\lambda)]-\mathbb{E}[F(\bar{x},\bar{\lambda})]$
$\displaystyle\leq\mathcal{O}(\sqrt{n}K^{-2/5}),$ (9)
$\displaystyle\mathbb{E}[\|\hat{x}(\bar{x})-\bar{x}\|^{2}]$
$\displaystyle\leq\mathcal{O}(\sqrt{n}K^{-2/5}).$ (10)

First, (8) tells us that in expectation, $\bar{y}_{i}$ converges to
$y_{i}^{\star}(\bar{x})$ for all $i$. Next, (9) shows that $\bar{\lambda}$ is
near-optimal, as using $\bar{\lambda}$ gives an objective value of
$F(\bar{x},\cdot)$ close to optimal. To understand (10), we define the
proximal map as follows:

Proximal Map: For a fixed $\rho$ and some $\lambda\in\mathbb{R}^{n}$, we
define the proximal map $\hat{x}(z)$ as

$\displaystyle\hat{x}(z)\triangleq$
$\displaystyle\operatorname*{arg\,min}_{x}\frac{\rho}{2}\|x-z\|^{2}+\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x).$
(11)

Now, going back to (10), we follow Davis and Drusvyatskiy (2018) and Hong et
al. (2020), quantifying the convergence rate of $x$ by showing the convergence
of $\mathbb{E}[\|\hat{x}(x_{k})-x_{k}\|^{2}]$. For any $\epsilon>0$, we say
that $x_{k}\in\mathcal{X}$ is an $\epsilon$-nearly stationary solution if
$\mathbb{E}[\|\hat{x}(x_{k})-x_{k}\|^{2}]\leq\rho^{2}\epsilon$. In (10), we
hide the $\rho$ factors under the $\mathcal{O}$ notation, so our solution is
$\sqrt{n}K^{-2/5}$-stationary.

For completeness, we also formally state our main result:

######  Theorem 1.

Under appropriate regularity assumptions and step-sizes

$\displaystyle\alpha$
$\displaystyle=\min\Bigg{(}\frac{\mu_{g}}{16G_{y}L}\nu,\frac{K^{-3/5}}{4G_{y}L}\Bigg{)},$
(12a) $\displaystyle\beta$
$\displaystyle=\min\Big{(}\nu,\frac{4K^{-2/5}}{\mu_{g}}\Big{)},$ (12b)
$\displaystyle\gamma$ $\displaystyle=\frac{2K^{-3/5}}{B_{\ell}n^{1/2}},$ (12c)

where
$\nu=\min\left\\{\frac{\mu_{g}}{L_{g}^{2}(1+\sigma_{g}^{2})},\frac{2}{\mu_{g}}\right\\}$,
Algorithm 1 produces $\bar{x},\bar{\lambda},\bar{y}_{i}$ satisfying:

$\displaystyle\mathbb{E}\left[\max_{i\in[n]}\|\bar{y}_{i}-y_{i}^{\star}(\bar{x})\|^{2}\right]\leq
A,$ (13)
$\displaystyle\max_{\lambda}\mathbb{E}[F(\bar{x},\lambda)]-\mathbb{E}[F(\bar{x},\bar{\lambda})]\leq\sqrt{2}B_{\ell}\sqrt{n}K^{-2/5}+G_{f}A,$
(14)
$\displaystyle\mathbb{E}[\|\hat{x}(\bar{x})-\bar{x}\|^{2}]\leq\frac{16\Phi_{1/\rho}(x_{0})G_{y}L}{(-\mu_{\ell}+\rho)\rho
K^{2/5}}+\frac{8(b_{0}^{2}+L^{2}A)}{(-\mu_{\ell}+\rho)^{2}}+\frac{2\alpha(\tilde{\sigma}_{f}^{2}+3b_{0}^{2}+3L^{2}A)}{-\mu_{\ell}+\rho},$

where

$A=\frac{\Delta_{y}^{0,i_{0}}/\mu_{g}}{K^{3/5}}+\frac{16\sigma_{g}^{2}/\mu_{g}^{2}}{K^{7/5}}+\frac{G_{y}/L}{K^{4/5}}+\frac{2\sqrt{n}B_{\ell}G_{y}/L}{K^{2/5}}+\frac{(b_{0}^{2}+\frac{1}{2}\sigma_{f}^{2})/(L^{2})}{K^{2/5}}+\frac{16\sigma_{g}^{2}/\mu_{g}^{2}}{K^{2/5}}.$

###  5.3 Comparison to Existing Convergence Results:

TTSA:  Once again, consider the case of $n=1$. (9) is irrelevant in this case,
because the LHS reduces to $0$, and there is no need for $\lambda$ iterates
anyways when $n=1$. Next, when $n=1$, (10) and (8) exactly recover the
convergence rates of $x$ and $y$ given in Theorem 2 of Hong et al. (2020). We
observe that when $n\neq 1$, the convergence of all three quantities suffers a
factor of $\sqrt{n}$, which is still sub-linear.

TR-MAML:  As described earlier, TR-MAML is a very special instance of our
framework, where each of the $y$ iterates is computed exactly, so (8) is
irrelevant in this case. They have convergence rates for $x$ and $\lambda$ in
Theorem 3, which exactly match our rate of $\mathcal{O}(\sqrt{n}K^{-2/5})$.

##  6 Discussion

###  6.1 Technical Challenges

Although our proof borrows the theorem proving ideas from TTSA (Hong et al.,
2020), we highlight a few key challenges differentiating our proof from
theirs:

1\. Sharpening the dependence on $n$:  One major challenge in this proof was
to achieve a rate of $\mathcal{O}(\sqrt{n}K^{-2/5})$ rather than
$\mathcal{O}(nK^{-2/5})$. Applying the steps in the proof of TTSA’s
convergence leads to this $n$ factor, as there are $n$ inner-loop problems
which potentially additive losses when considering the outer loop. We address
this through careful selection of the step size $\gamma$ and careful bounding
of $\lambda$ terms, highlighted in Lemma 5.

2\. Bounding the maximum of $\|y_{k}^{i}-y_{i}^{\star}(x_{k})\|$ rather than
the sum over $i$:  When considering the proof of TTSA’s convergence, they
consider a descent equation on $\mathbb{E}[\|y_{k}-y_{*}(x_{k})\|]$. This can
be extended in our case to a descent equation on
$\mathbb{E}[\|y_{k}^{i}-y_{i}^{\star}(x_{k})\|]$ for each individual $i$.
However, bounding $\mathbb{E}[\max_{i}\|y_{k}^{i}-y_{i}^{\star}(x_{k})\|]$ is
trickier, as the maximum resides inside the expectation, and just because the
expectation of each individual term is bounded doesn’t directly imply that the
expectation of the maximum will be bounded. We address this in Lemma 3 and
Lemma 4 below, where we take special care to avoid the loose bound
$\sum_{i=1}^{n}\lambda_{i}a_{i}\leq\sum_{i=1}^{n}a_{i}$ and instead use
$\sum_{i=1}^{n}\lambda_{i}a_{i}\leq\max_{i}a_{i}$. Since the index $i$
maximizing $\|y_{k}^{i}-y_{i}^{\star}(x_{k})\|^{2}$ is different for each $k$,
in Lemma 4 we also need to adjust the telescoping sum accordingly.

3\. Using a weighted potential function:  In our proof, we use the potential
function
$\mathcal{L}^{k}=\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\ell_{i}(x_{k})]$
rather than just $\ell(x_{k})$, as the TTSA paper does. Telescoping this
weighted potential is trickier, as smoothness cannot directly be used to bound
$\mathcal{L}^{k+1}-\mathcal{L}^{k}$. This challenge is highlighted in the
proof of Lemma 3.

4\. Convergence of $\lambda$ terms:  When $n>1$, convergence of $x$ and $y$ is
not sufficient. We also require convergence of $\lambda$ terms, which TTSA
does not consider at all. This challenge, like challenge 2, is highlighted in
Lemma 5.

###  6.2 Proof Sketch

Overall Roadmap:  In Lemma 1, we first bound the averaged squared norm of the
stochastic gradient estimate, $\mathbb{E}[\|h_{x}^{k}\|^{2}]$. In Lemma 2, we
give a descent equation involving
$\mathbb{E}\|y_{k+1}^{i}-y_{i}^{\star}(x_{k})\|^{2}$. In Lemma 3, we give a
descent equation involving the objective
$\mathcal{L}^{k}=F(x_{k},\lambda_{k})$. Lemmas 2 and 3 provide recursively
coupled descent inequalities for the convergence of both $x$ and
$\mathcal{L}$. Lemma 4 shows how to solve the coupled equations from Lemmas 2
and 3, which directly leads to the result on the convergence of $y$. Lemma 5
then leverages the convergence of $y$ to bound the convergence of $\lambda$,
and Lemma 6 concludes the trifecta of convergence lemmas by showing the bound
on $x$. By plugging in our step-sizes into Lemmas 4, 5, and 6, Theorem 1
directly follows. In what follows, we give an abbreviated and more readable
proof of these lemmas, with constant terms abstracted away with $\mathcal{O}$
notation.

######  Lemma 1.

Under our regularity assumptions, the average squared norm of $h_{x}^{k}$ can
be bounded as follows:

$\mathbb{E}[\|h_{x}^{k}\|^{2}]\leq\mathcal{O}\left(\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|y_{i}^{\star}(x_{k})-y_{k+1}^{i}\|^{2}]\right).$

Next, in Lemma 2, we begin to bound the iterates, showing that
$\Delta_{y}^{k+1,i}=\mathbb{E}[\|y_{k+1}^{i}-y_{i}^{\star}(x_{k})\|^{2}]$
satisfies a descent equation:

######  Lemma 2.

The distance between the algorithm’s iterates $y_{k}^{i}$ and the true inner
optimum $y_{i}^{\star}(x_{k})$ satisfy the following descent equation:

$\mathbb{E}[\|y_{k+1}^{i}-y_{i}^{\star}(x_{k})\|^{2}]\leq\mathcal{O}((1-c_{1})\mathbb{E}[\|y_{k}^{i}-y_{i}^{\star}(x_{k-1})\|^{2}]+c_{2}\mathbb{E}[\|x_{k-1}-x_{k}\|^{2}]).$
(15)

In Lemma 3, we show that
$\mathcal{L}^{k}=\mathbb{E}[F(x_{k},\lambda_{k})]=\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\ell_{i}(x_{k})]$
satisfies a descent equation by bounding $\mathcal{L}^{k+1}-\mathcal{L}^{k}$:

######  Lemma 3.

Let
$\mathcal{L}^{k}\triangleq\mathbb{E}[F(x_{k},\lambda_{k})]=\sum\limits_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\ell_{i}(x_{k})]$.
Then, $\mathcal{L}$ satisfies the descent equation

$\mathcal{L}^{k+1}-\mathcal{L}^{k}\leq\mathcal{O}(c_{1}\mathbb{E}[\|x_{k+1}-x_{k}\|^{2}]+c_{2}\max_{i\in[n]}\mathbb{E}[\|y_{k+1}^{i}-y_{i}^{\star}(x_{k})\|^{2}]+\sqrt{n}\gamma+\alpha).$

At this point, observe that Lemma 2 and Lemma 3 give us a set of coupled
recurrent inequalities involving
$\mathbb{E}[\|y_{k+1}^{i}-y_{i}^{\star}(x_{k})\|^{2}]$, $\mathcal{L}^{k}$, and
$\mathbb{E}[\|x_{k+1}-x_{k}\|^{2}]$. Lemma 4 solves this set of coupled
inequalities and directly leads to (8).

######  Lemma 4.

Assume that
$\Omega^{i},\Theta^{i},\Upsilon^{i},\lambda_{i},c_{0},c_{1},c_{2}$,
$d_{0},d_{1},d_{2}$ are real numbers such that for all $0\leq k\leq K-1$,
$\Omega^{k+1}\leq\Omega^{k}-c_{0}\Theta^{k+1}+c_{1}\max_{i\in[n]}\Upsilon^{k+1,i}+c_{2}$
and also for all $1\leq k\leq K,1\leq i\leq N$,
$\Upsilon^{k+1,i}\leq(1-d_{0})\Upsilon^{k,i}+d_{1}\Theta^{k}+d_{2}$ . Then, if
$i_{0}=\operatorname*{arg\,max}_{i\in[n]}\Upsilon^{0,i}$, we have

$\frac{1}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Upsilon^{k,i}\leq(d_{0}-d_{1}c_{0}^{-1}c_{1})^{-1}\left(\frac{\Upsilon^{0,i_{0}}+d_{1}\Theta^{0}+d_{2}+d_{1}c_{0}^{-1}\Omega^{0}}{K}\\!+\\!d_{1}c_{0}^{-1}c_{2}\\!+\\!d_{2}\right).$

In Lemma 5, we bound the optimality gap of $\bar{\lambda}$, and we conclude by
showing the optimality of $\bar{x}$ in Lemma 6. Theorem 1 then follows from
Lemmas 4, 5, and 6.

######  Lemma 5.

For any $\lambda\in\Delta_{n}$, the iterates of Algorithm 1 satisfy

$\frac{1}{K}\mathbb{E}\left[\sum_{k=1}^{K}F(x_{k},\lambda)-F(x_{k},\lambda_{k})\right]\leq\mathcal{O}(\sqrt{n}K^{-2/5}).$

######  Lemma 6.

The iterates of Algorithm 1 satisfy

$\displaystyle\frac{1}{K}\sum_{k=1}^{K}\|\hat{x}(x_{k})-x_{k}\|^{2}\leq\mathcal{O}(\sqrt{n}K^{-2/5}).$

##  7 Generalization Bounds

In addition to convergence, we also present a result showing the
generalization abilities of the bilevel optimizer. The proofs and theorems in
this section closely follow those in Collins et al. (2020), Section 5.
However, since our setting is a generalization of their setting, our results
hold in more general bilevel optimization problems beyond the meta-learning
setup in Collins et al. (2020).

Assume that for task $i$, we observe $m_{i}$ batches of train/test data,
$\mathcal{D}_{i,j}^{\text{train}}$ and $\mathcal{D}_{i,j}^{\text{test}}$.
Assume that each train and test batch has $K$ and $J$ input-output pairs,
respectively, so that the $\mathcal{D}_{i,j}$ sets are drawn from a common
distribution $\mathcal{D}_{i}$ over
$\mathcal{Z}\triangleq(\mathcal{X}\times\mathcal{Y})^{K+J}$. Also, let
$\hat{y}_{i}^{*}(x;D_{i}^{\text{tr}})=\operatorname*{arg\,min}_{y_{i}}\frac{1}{m_{i}}\sum_{j^{\prime}=1}^{m_{i}}\hat{g}_{i}(x,y_{i},D_{i,j^{\prime}}^{\text{train}})$
be the value of $y_{i}$ that minimizes the empirical inner loss on some
dataset $D_{i}^{\text{tr}}$. For outer objectives $f_{i}$ and inner objectives
$g_{i}$, we consider the following function class $\mathcal{F}$, where
$x\in\mathcal{X}$ is the set of optimization parameters introduced in Eq. 1
and $(D_{i}^{\text{tr}},D_{i}^{\text{test}})$ are any train and test datasets
sampled from $\mathcal{Z}$:

$\mathcal{F}_{i}=\left\\{\hat{f}_{i}\left(x,\hat{y}_{i}^{*}(x;D_{i}^{\text{tr}}),D_{i}^{\text{test}}\right),x\in\mathcal{X}\right\\}$

Here, we use $\hat{f}$ and $\hat{g}$ to denote the empirical function values
evaluated at the points in $D_{i}^{\text{tr}}$ and $D_{i}^{\text{test}}$.
Therefore, the empirical Rademacher complexity of $\mathcal{F}_{i}$ on $m_{i}$
samples
$\mathbf{D}_{i}\triangleq\\{(D_{i,j}^{\text{train}},D_{i,j}^{\text{test}})\\}_{j=1}^{m_{i}}\sim(\mathcal{D}_{i})^{m_{i}}$
is

$\mathcal{R}_{m_{i}}^{i}(\mathcal{F}_{i})=E_{\mathbf{D}_{i}}\mathbb{E}_{\epsilon_{j}}\Bigg{[}\sup_{x\in\mathcal{X}}\frac{1}{m_{i}}\sum_{j=1}^{m_{i}}\epsilon_{j}\hat{f}_{i}\Bigg{(}x,\hat{y}_{i}^{*}(x;D_{i}^{\text{tr}}),D_{i,j}^{\text{test}}\Bigg{)}\Bigg{]}.$
(16)

Now, let’s define, for fixed samples $D_{i,j}^{\text{train}}$ and
$D_{i,j}^{\text{test}}$,

$\hat{F}_{i}(x)\triangleq\frac{1}{m_{i}}\sum_{j=1}^{m_{i}}\hat{f}_{i}\left(x,\hat{y}_{i}^{*}(x;D_{i}^{\text{tr}}),D_{i,j}^{\text{test}}\right).$

Similarly, define $F_{i}(x)\triangleq\mathbb{E}_{\mathbf{D}}[\hat{F}_{i}(x)]$.
First, from classical generalization results such as Theorem 26.5 of Shalev-
Shwartz and Ben-David (2014) or Theorem 3.3 of Mohri et al. (2018), we
directly conclude the following proposition, which bounds the true loss of the
classifier as a function of the empirical loss.

Proposition 7.1:  Assume the regularity assumptions given in Section 5,
specifically that the function $\ell$ is $B_{\ell}$-bounded. Then, with
probability at least $1-\delta$, we have

$F_{i}(x)\leq\hat{F}_{i}(x)+2\mathcal{R}_{m_{i}}^{i}(\mathcal{F})+B_{\ell}\sqrt{\frac{\log
1/\delta}{2m_{i}}}.$ (17)

Next, we proceed to bound the generalization on unseen tasks. Consider a new
task with distribution $\mathcal{D}_{n+1}=\sum_{i=1}^{n}a_{i}\mathcal{D}_{i}$,
for some $a\in\Delta_{n}$, meaning that the distribution of the new task is
anywhere in the convex hull of the distribution of the old tasks. We make this
assumption because if the new task is very dissimilar to the existing tasks,
there is no reason to expect good generalization in the first place. We then
show the following proposition:

Proposition 7.2:  Let $x^{\star}$ be the optimal solution to the bilevel
optimization objective in (3). Then, we have

$F_{n+1}(x^{\star})\leq\min_{x\in\mathcal{X}}\max_{p\in\Delta_{n}}p_{i}\hat{F}_{i}(x)+2a_{i}\mathcal{R}_{m_{i}}^{i}(\mathcal{F}_{i})+a_{i}B_{\ell}\sqrt{\frac{\log(n/\delta)}{2m_{i}}}.$
(18)

This proposition shows the generalization abilities of our function class for
new tasks in the convex hull of the seen tasks, giving an upper bound for the
loss on new tasks.

##  8 Experimental Results

While the focus of our work is mostly theoretical, we present an empirical
evaluation on a synthetic sinusoid regression setup and a nonlinear
representation learning setting, validating the effectiveness of our proposed
formulation and theoretical results.

###  8.1 Sinusoid Regression

Through the sinusoid regression example, we highlight the ability of min-max
training to converge quickly on the worst-case seen task loss (as in Theorem
1) and generalize well on unseen tasks (as in Proposition 7.2). We defer
experiment details to B and focus on highlighting the significance of our
results here. We trained using two different methods, the first using the min-
max multi-objective we propose and the second using a standard averaged
single-objective. We label these methods min-max and min-avg, respectively.
During training, at each epoch, we plotted the worst-task MSE losses for both
methods on seen data in Figure 1. During testing, we train the inner-level
weights on the new unseen tasks while fixing the outer-weights, again plotting
the MSE losses on unseen data in Figure 1. In the figure, we can see that min-
max training both outperforms min-avg training on seen tasks and generalizes
better to unseen tasks.

Figure 1: Comparison of the standard min-avg single-objective training with
our proposed min-max multi-objective training for sinusoid regression. Observe
that min-max training leads to better performance on the worst-case loss for
both seen and unseen tasks.

###  8.2 Representation Learning

Through the representation learning example, we again start by highlighting
the ability of min-max training to converge quickly on both average-case and
worst-case seen task losses (as in Theorem 1). In addition, we showcase the
generalization ability of min-max training on unseen instances of seen tasks
(as in Proposition 7.1). During training, we keep track of two values at each
epoch: the loss of our algorithm on meta-training data (seen instances of seen
tasks) and the generalization of our algorithm to meta-testing data (unseen
instances of seen tasks). We keep track of both the average loss and the
worst-case loss. As Figure 2 highlights, our min-max training procedure
outperforms min-avg training in both cases, showing both faster convergence
and better generalization on seen tasks.

Figure 2: The aggregated (mean/max) objectives during the optimization
trajectory of the min-average and the min-max schemes for the nonlinear
representation learning with binary classification tasks from FashionMNIST.
The plot reports the outer level max objectives $\max_{i}f_{i}$.

Convergence rate:  Next, we highlight the empirical convergence rate of the
gradient norm $\|\nabla f(x)\|_{2}$, which is a tractable proxy for the
theoretical convergence of $\bar{x}$ derived (10). As shown in Figure 3, the
convergence of the gradient norm very closely matches the theoretical
convergence rate of $O(K^{-2/5})$.

Figure 3: The trajectory of the gradient norm $\|\nabla_{x}\|_{2}$ through the
optimization trajectory of the nonlinear representation learning with binary
classification problems from FashionMNIST. We compare this trajectory to
$K^{-\nicefrac{{2}}{{5}}}$.

##  9 Concluding Remarks

In this work, inspired by the desiderata of robust multi-task and
representation learning, we first propose a new framework for task-robust
bilevel optimization. Our framework builds off of the existing bilevel
optimization framework and adds an extra maximization step for robustness. Our
new framework provides full flexibility and generality, allowing each task to
have a different loss function and inner optimization procedure. Next, we
proposed MORBiT, a single-loop gradient descent-ascent based algorithm for
finding an optimal solution to our proposed framework. Finally, we show
theoretically, that MORBiT converges at a rate of $O(K^{-2/5})$. As special
cases, MORBiT reduces to TR-MAML (Collins et al., 2020) and TTSA (Hong et al.,
2020), and our theory recovers the convergence results in those two works. We
hope that our framework inspires future algorithms for robust bilevel
optimization and that others find novel applications of our ideas.

Limitations: Since our focus has been on the algorithm development and
analysis of a general robust bilevel optimization framework, the main
limitation of our work is a thorough empirical evaluation of the scheme across
different multi-task learning problems and other applications of min-max
bilevel optimization such as hyperparameter optimization. In addition, our
bounds still scale as $n$ grows, which could potentially pose a problem with a
large number of tasks.

## References

  * Franceschi et al. (2018) Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil.  Bilevel programming for hyperparameter optimization and meta-learning.  In _International Conference on Machine Learning_ , pages 1568–1577. PMLR, 2018.
  * Hong et al. (2020) Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang.  A two-timescale framework for bilevel optimization: Complexity analysis and application to actor-critic.  _arXiv preprint arXiv:2007.05170_ , 2020.
  * Fallah et al. (2020) Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar.  Personalized federated learning: A meta-learning approach.  _arXiv preprint arXiv:2002.07948_ , 2020.
  * Arora et al. (2020) Sanjeev Arora, Simon Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi.  Provable representation learning for imitation learning via bi-level optimization.  In _International Conference on Machine Learning_ , pages 367–376. PMLR, 2020.
  * Liu et al. (2021) Risheng Liu, Jiaxin Gao, Jin Zhang, Deyu Meng, and Zhouchen Lin.  Investigating bi-level optimization for learning and vision from a unified perspective: A survey and beyond.  _arXiv preprint arXiv:2101.11517_ , 2021.
  * Collins et al. (2020) Liam Collins, Aryan Mokhtari, and Sanjay Shakkottai.  Task-robust model-agnostic meta-learning.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks.  _arXiv preprint arXiv:1703.03400_ , 2017.
  * Ghadimi and Wang (2018) Saeed Ghadimi and Mengdi Wang.  Approximation methods for bilevel programming.  _arXiv preprint arXiv:1802.02246_ , 2018.
  * Ji et al. (2020) Kaiyi Ji, Junjie Yang, and Yingbin Liang.  Provably faster algorithms for bilevel optimization and applications to meta-learning.  _arXiv preprint arXiv:2010.07962_ , 2020.
  * Chen et al. (2021a) Tianyi Chen, Yuejiao Sun, and Wotao Yin.  A single-timescale stochastic bilevel optimization method.  _arXiv preprint arXiv:2102.04671_ , 2021a.
  * Chen et al. (2021b) Tianyi Chen, Yuejiao Sun, and Wotao Yin.  Tighter analysis of alternating stochastic gradient method for stochastic nested problems.  _arXiv preprint arXiv:2106.13781_ , 2021b.
  * Sinha et al. (2015) Ankur Sinha, Pekka Malo, and Kalyanmoy Deb.  Towards understanding bilevel multi-objective optimization with deterministic lower level decisions.  In _International Conference on Evolutionary Multi-Criterion Optimization_ , pages 426–443. Springer, 2015.
  * Deb and Sinha (2009) Kalyanmoy Deb and Ankur Sinha.  Solving bilevel multi-objective optimization problems using evolutionary algorithms.  In _International conference on evolutionary multi-criterion optimization_ , pages 110–124. Springer, 2009.
  * Ji et al. (2017) Ying Ji, Shaojian Qu, and Zhensheng Yu.  A new method for solving multiobjective bilevel programs.  _Discrete Dynamics in Nature and Society_ , 2017, 2017.
  * Rajeswaran et al. (2019) Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine.  Meta-learning with implicit gradients.  In _Advances in Neural Information Processing Systems_ , pages 113–124, 2019.
  * Duchi and Namkoong (2018) John Duchi and Hongseok Namkoong.  Learning models with uniform performance via distributionally robust optimization.  _arXiv preprint arXiv:1810.08750_ , 2018.
  * Shalev-Shwartz and Wexler (2016) Shai Shalev-Shwartz and Yonatan Wexler.  Minimizing the maximal loss: How and why.  In _International Conference on Machine Learning_ , pages 793–801. PMLR, 2016.
  * Qian et al. (2019) Qi Qian, Shenghuo Zhu, Jiasheng Tang, Rong Jin, Baigui Sun, and Hao Li.  Robust optimization over multiple domains.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 4739–4746, 2019.
  * Mehta et al. (2012) Nishant A. Mehta, Dongryeol Lee, and Alexander G. Gray.  Minimax multi-task learning and a generalized loss-compositional paradigm for mtl.  In _Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 2_ , page 2150–2158, 2012.
  * Madry et al. (2017) Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu.  Towards deep learning models resistant to adversarial attacks.  _arXiv preprint arXiv:1706.06083_ , 2017.
  * Wang et al. (2019) Jingkang Wang, Tianyun Zhang, Sijia Liu, Pin-Yu Chen, Jiacen Xu, Makan Fardad, and Bo Li.  Towards a unified min-max framework for adversarial exploration and robustness.  _arXiv preprint arXiv:1906.03563_ , 2019.
  * Cheng et al. (2020) Minhao Cheng, Qi Lei, Pin-Yu Chen, Inderjit Dhillon, and Cho-Jui Hsieh.  Cat: Customized adversarial training for improved robustness.  _arXiv preprint arXiv:2002.06789_ , 2020.
  * Raghu et al. (2019) Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals.  Rapid learning or feature reuse? towards understanding the effectiveness of maml.  _arXiv preprint arXiv:1909.09157_ , 2019.
  * Javed and White (2019) Khurram Javed and Martha White.  Meta-learning representations for continual learning.  _arXiv preprint arXiv:1905.12588_ , 2019.
  * Lu et al. (2020) Songtao Lu, Ioannis Tsaknakis, Mingyi Hong, and Yongxin Chen.  Hybrid block successive approximation for one-sided non-convex min-max problems: algorithms and applications.  _IEEE Transactions on Signal Processing_ , 68:3676–3691, 2020.
  * Lin et al. (2020) Tianyi Lin, Chi Jin, and Michael Jordan.  On gradient descent ascent for nonconvex-concave minimax problems.  In _International Conference on Machine Learning_ , pages 6083–6093. PMLR, 2020.
  * Davis and Drusvyatskiy (2018) Damek Davis and Dmitriy Drusvyatskiy.  Stochastic subgradient method converges at the rate $\mathcal{O}(k^{-1/4})$ on weakly convex functions.  _arXiv preprint arXiv:1802.02988_ , 2018.
  * Shalev-Shwartz and Ben-David (2014) Shai Shalev-Shwartz and Shai Ben-David.  _Understanding machine learning: From theory to algorithms_.  Cambridge university press, 2014.
  * Mohri et al. (2018) Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar.  _Foundations of machine learning_.  MIT press, 2018.

##  Appendix A Proofs

###  A.1 Assumptions

First, we begin by listing the assumptions we make:

Assumption 1 (Regularity of the outer functions):  For all $i\in[n]$, assume
that outer functions $f_{i}(x,y)$ and $\ell_{i}(x)=f_{i}(x,y_{i}^{*}(x))$
satisfy the following properties:

  * •

For any $x\in\mathcal{X}$, $f_{i}(x,\cdot)$ is Lipschitz (w.r.t. $y$) with
constant $G_{f}>0$.

  * •

For any $x\in\mathcal{X}$, $\nabla_{x}f_{i}(x,\cdot)$ and
$\nabla_{y}f_{i}(x,\cdot)$ are Lipschitz continuous (w.r.t. $y$) with
constants $L_{f_{x}}>0$ and $L_{f_{y}}>0$.

  * •

For any $y\in\mathbb{R}^{d_{2}}$, $\nabla_{y}f_{i}(\cdot,y)$ is Lipschitz
continuous (w.r.t. $x$) with constant $\bar{L}_{f_{y}}>0$.

  * •

For any $x\in\mathcal{X},y\in\mathbb{R}^{d_{2}}$, we have
$\|\nabla_{y}f_{i}(x,y)\|\leq C_{fy}$ for $C_{fy}>0$.

  * •

For all $i$, the function $\ell_{i}(\cdot)$ is $\mu_{\ell}$ weakly convex (in
$x$), so that for all $v,w\in\mathcal{X}$,

$\ell(w)\geq\ell(v)+\langle\nabla\ell(v),w-v\rangle-\mu_{\ell}\|w-v\|^{2}.$
(19)

  * •

For all $x\in\mathcal{X}$, $\|\ell_{i}(x)\|\leq B_{\ell}$ for $B_{\ell}>0$.

  * •

For all $x\in\mathcal{X}$, $\|\nabla\ell_{i}(x)\|\leq C_{\ell}$, for
$C_{\ell}>0$.

Assumption 2 (Regularity of the inner functions):  Assume that inner functions
$g_{i}(x,y),\forall i$ satisfy:

  * •

For any $x\in\mathcal{X}$ and $y\in\mathbb{R}^{d_{2}}$, $g_{i}(x,y)$ is twice
continuously differentiable in $(x,y)$.

  * •

For any $x\in\mathcal{X}$, $\nabla_{y}g_{i}(x,\cdot)$ is Lipschitz continuous
(w.r.t. $y$) with constant $L_{g}$.

  * •

For any $x\in\mathcal{X}$, $g_{i}(x,\cdot)$ is $\mu_{g}$-strongly convex in
$y$, so that for all $v,w\in\mathbb{R}^{d_{2}}$,

$g_{i}(x,w)\geq g_{i}(x,v)+\langle\nabla
g_{i}(x,v),w-v\rangle+\mu_{g}\|w-v\|^{2}.$ (20)

  * •

For any $x\in\mathcal{X}$, $\nabla_{xy}^{2}g_{i}(x,\cdot)$ and
$\nabla_{yy}^{2}g_{i}(x,\cdot)$ are Lipschitz continuous (w.r.t. $y$) with
constants $L_{gxy}>0$ and $L_{gyy}>0$, respectively.

  * •

For any $x\in\mathcal{X}$ and $y\in\mathbb{R}^{d_{2}}$, we have
$\|\nabla_{xy}^{2}g_{i}(x,y)\|\leq C_{gxy}$ for some $C_{gxy}>0$.

  * •

For any $y\in\mathbb{R}^{d_{2}}$, $\nabla_{xy}^{2}g_{i}(\cdot,y)$ and
$\nabla_{yy}^{2}g_{i}(\cdot,y)$ are Lipschitz continuous (w.r.t. $x$) with
constants $\bar{L}_{gxy}>0$ and $\bar{L}_{gyy}>0$, respectively.

Assumption 3 (Boundedness of stochastic gradient estimates): For any $k\geq
0$, we assume that there exists a nonincreasing sequence $\\{b_{k}\\}_{k\geq
0}$ such that

$\displaystyle\mathbb{E}[h_{g}^{k,i}]$
$\displaystyle=\nabla_{y}g_{i}(x_{k},y_{k}^{i}),$ (21)
$\displaystyle\mathbb{E}[h_{x}^{k}]$ $\displaystyle=\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+B_{k},\|B_{k}\|\leq b_{k},$ (22)
$\displaystyle h_{\lambda}^{k}$
$\displaystyle=\nabla_{\lambda}F(x_{k},\lambda_{k}),$ (23)
$\displaystyle=\left[f_{1}(x_{k},y_{k+1}^{1}),\cdots,f_{n}(x_{k},y_{k+1}^{n})\right]^{\top}.$
(24)

In addition, we assume some conditions on the variances, so that there exist
$\sigma_{f},\sigma_{g}$ satisfying

$\displaystyle\mathbb{E}[\|h_{x}^{k}-\mathbb{E}[h_{x}^{k}]\|^{2}]$
$\displaystyle\leq\sigma_{f}^{2},$ (25)
$\displaystyle\mathbb{E}[\|h_{g}^{k,i}-\nabla_{y}g_{i}(x_{k},y_{k})\|^{2}]$
$\displaystyle\leq\sigma_{g}^{2}(1+\|\nabla_{y}g_{i}(x_{k},y_{k})\|^{2}).$
(26)

From these assumptions, we can show a few additional regularity-type
conditions. Since these conditions can also be found as Lemma 2.2 in [Ghadimi
and Wang, 2018] and/or Lemma 1 in [Hong et al., 2020], we state these results
without proof.

Lemma 0 (Corollary of Assumptions): ([Ghadimi and Wang, 2018] Lemma 2.2, [Hong
et al., 2020] Lemma 1) Under assumptions $1$ and $2$ stated above, for all
$x,x_{1},x_{2}\in\mathcal{X},y\in\mathbb{R}^{d_{2}},i\in[n]$, we have
$\|\overline{\nabla}_{x}f_{i}(x,y)-\nabla\ell_{i}(x)\|\leq
L\|y_{i}^{*}(x)-y\|$, $\|y_{i}^{*}(x_{1})-y_{i}^{*}(x_{2})\|\leq
G_{y}\|x_{1}-x_{2}\|$, and
$\|\nabla\ell_{i}(x_{1})-\nabla\ell_{i}(x_{2})\|\leq L_{f}\|x_{1}-x_{2}\|$.
Here, we define

$\displaystyle L$ $\displaystyle\triangleq
L_{f_{x}}+\frac{L_{f_{y}}C_{gxy}}{\mu_{g}}+C_{f_{y}}\left(\frac{L_{gxy}}{\mu_{g}}+\frac{L_{gyy}C_{gxy}}{\mu_{g}^{2}}\right),$
(27) $\displaystyle L_{f}$ $\displaystyle\triangleq
L_{f_{x}}+\frac{(\bar{L}_{f_{y}}+L)C_{gxy}}{\mu_{g}}+C_{f_{y}}\left(\frac{\bar{L}_{gxy}}{\mu_{g}}+\frac{\bar{L}_{gyy}C_{gxy}}{\mu_{g}^{2}}\right),G_{y}\triangleq\frac{C_{g}}{\mu_{g}}.$
(28)

###  A.2 Main Theorem and Remarks

Theorem 1 (Convergence of MORBiT): Under Assumptions $1,2,3$, when step sizes
are chosen as

$\displaystyle\alpha=\min\Bigg{(}\frac{\mu_{g}}{16G_{y}L}\nu,\frac{K^{-3/5}}{4G_{y}L}\Bigg{)},\beta=\min\Big{(}\nu,\frac{4K^{-2/5}}{\mu_{g}}\Big{)},\gamma=\frac{2K^{-3/5}}{B_{\ell}n^{1/2}},$
(29)

where
$\nu=\min(\frac{\mu_{g}}{L_{g}^{2}(1+\sigma_{g}^{2})},\frac{2}{\mu_{g}})$,
Algorithm 1 produces $\bar{x},\bar{\lambda},\bar{y}_{i}$ satisfying:

$\displaystyle\mathbb{E}\left[\max_{i\in[n]}\|\bar{y}_{i}-y_{i}^{*}(\bar{x})\|^{2}\right]$
$\displaystyle\leq A,$ (30)
$\displaystyle\max_{\lambda}\mathbb{E}[F(\bar{x},\lambda)]-\mathbb{E}[F(\bar{x},\bar{\lambda})]$
$\displaystyle\leq\sqrt{2}B_{\ell}\sqrt{n}K^{-2/5}+G_{f}A,$ (31)
$\displaystyle\mathbb{E}[\|\hat{x}(\bar{x})-\bar{x}\|^{2}]$
$\displaystyle\leq\frac{16\Phi_{1/\rho}(x_{0})G_{y}L}{(-\mu_{\ell}+\rho)\rho
K^{2/5}}+\frac{8(b_{0}^{2}+L^{2}A)}{(-\mu_{\ell}+\rho)^{2}}+\frac{2\alpha(\tilde{\sigma}_{f}^{2}+3b_{0}^{2}+3L^{2}A)}{-\mu_{\ell}+\rho},$
(32)

where

$A=\frac{\Delta_{y}^{0,i_{0}}/\mu_{g}}{K^{3/5}}+\frac{16\sigma_{g}^{2}/\mu_{g}^{2}}{K^{7/5}}+\frac{G_{y}/L}{K^{4/5}}+\frac{2\sqrt{n}B_{\ell}G_{y}/L}{K^{2/5}}+\frac{(b_{0}^{2}+\frac{1}{2}\sigma_{f}^{2})/(L^{2})}{K^{2/5}}+\frac{16\sigma_{g}^{2}/\mu_{g}^{2}}{K^{2/5}}.$
(33)

We remark that as we generalize [Hong et al., 2020], our proof follows a
similar structure. In particular, Lemma 1 is a generalization of (10) in [Hong
et al., 2020], and our Lemma 2 combines Lemmas 2 and 3 from [Hong et al.,
2020]. Lemmas 3, 4, and 6 in our work parallel Lemmas 5, 4, and 6 from [Hong
et al., 2020], respectively. Lemma 5 in our work deals with the maximization
problem w.r.t. $\lambda$, so there is no analogue in [Hong et al., 2020].
However, it borrows techniques from Theorem 1 in [Collins et al., 2020].

We also discuss the convergence rate. Note that $A$ is dominated by the fourth
term, $\sqrt{n}/K^{2/5}$, so it is clear that (30) and (31) converge at a rate
of $\mathcal{O}(\sqrt{n}K^{-2/5})$. We give special attention to (32). Apart
from the $\frac{8b_{0}^{2}}{(-\mu_{\ell}+\rho)^{2}}$ term, we see that the RHS
of (32) converges at a rate of $\mathcal{O}(\sqrt{n}K^{-2/5})$. To understand
the convergence of this term, we turn to (7) from Assumption 3. As discussed
in Section E.4 of the Appendix in [Hong et al., 2020], $b_{k}$ can be made
arbitrarily small by running more iterations of the subroutine for estimating
$h_{x}^{k}$. Therefore, as long as we run enough iterations such that
$b_{k}^{2}\leq\alpha$, (32) will also converge at a rate of
$\mathcal{O}(\sqrt{n}K^{-2/5})$.

###  A.3 Proof of Lemma 1

Lemma 1:  Under Assumptions 1, 2, and 3, the average squared norm of
$h_{x}^{k}$ can be bounded as follows, where
$\tilde{\sigma}_{f}^{2}=\sigma_{f}^{2}+3C_{\ell}^{2}$:

$\mathbb{E}[\|h_{x}^{k}\|^{2}]\leq\tilde{\sigma}_{f}^{2}+3b_{k}^{2}+3L^{2}\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}].$
(34)

Proof:  We can derive the following:

$\displaystyle\mathbb{E}[\|h_{x}^{k}\|^{2}]$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\mathbb{E}[\|h_{x}^{k}-\mathbb{E}[h_{x}^{k}]\|^{2}]+\|\mathbb{E}[h_{x}^{k}]\|^{2}$
(35)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\mathbb{E}[\|h_{x}^{k}-\mathbb{E}[h_{x}^{k}]\|^{2}]+\mathbb{E}[\|h_{x}^{k}\|]^{2}$
(36)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\sigma_{f}^{2}+\mathbb{E}[\|h_{x}^{k}\|]^{2}$
(37)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{=}}\sigma_{f}^{2}+\mathbb{E}\left[\left\|\sum_{i=1}^{n}(\lambda_{k})_{i}\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i})+B_{k}\right\|^{2}\right]$ (38)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\leq}}\sigma_{f}^{2}+3b_{k}^{2}+\frac{3}{2}\mathbb{E}\left[\left\|\sum_{i=1}^{n}(\lambda_{k})_{i}\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i})\right\|^{2}\right].$ (39)

(1) is true because

$\displaystyle\mathbb{E}[\|h_{x}^{k}-\mathbb{E}[h_{x}^{k}]\|^{2}]+\|\mathbb{E}[h_{x}^{k}]\|^{2}$
$\displaystyle=\mathbb{E}[\|h_{x}^{k}\|^{2}]+\|\mathbb{E}[h_{x}^{k}]\|^{2}-2\mathbb{E}\langle
h_{x}^{k},\mathbb{E}[h_{x}^{k}]\rangle+\|\mathbb{E}[h_{x}^{k}]\|^{2}$ (40)
$\displaystyle=\mathbb{E}[\|h_{x}^{k}\|^{2}]+\|\mathbb{E}[h_{x}^{k}]\|^{2}-\langle\mathbb{E}[h_{x}^{k}],\mathbb{E}[h_{x}^{k}]\rangle+\|\mathbb{E}[h_{x}^{k}]\|^{2}$
(41) $\displaystyle=\mathbb{E}[\|h_{x}^{k}\|^{2}].$ (42)

(2) follows from Jensen’s inequality and the convexity of norms, (3) follows
from (25), (4) follows from definition of $h_{x}^{k}$, and (5) follows from
Young’s inequality, $\|a+b\|^{2}\leq 3\|a\|^{2}+\frac{3}{2}\|b\|^{2}$. Next,
we bound the last term in (39). We start by using the fact that

$\displaystyle\mathbb{E}\biggl{[}\biggl{\|}\sum_{i=1}^{n}(\lambda_{k})_{i}\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i})\biggr{\|}^{2}\biggr{]}$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}2\mathbb{E}\biggl{[}\biggl{\|}\sum_{i=1}^{n}(\lambda_{k})_{i}(\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i})-\nabla\ell_{i}(x_{k}))\biggr{\|}^{2}\biggr{]}+2\mathbb{E}\biggl{[}\biggl{\|}\sum_{i=1}^{n}(\lambda_{k})_{i}\nabla\ell_{i}(x_{k})\biggr{\|}^{2}\biggr{]}$
(43)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}2\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i})-\nabla\ell_{i}(x_{k})\|^{2}]+2\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|\nabla\ell_{i}(x_{k})\|^{2}],$
(44)

where (1) follows from $\|a+b\|^{2}\leq 2\|a\|^{2}+2\|b\|^{2}$, and (2)
follows from
$\Big{\|}\sum\limits_{i=1}^{n}p_{i}a_{i}\Big{\|}^{2}\leq\sum\limits_{i=1}^{n}p_{i}\|a_{i}\|^{2}$.
Next, we bound the first term in (44). From Lemma 0 (A.1), we have

$\mathbb{E}[\|\mkern 1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i})-\nabla\ell_{i}(x_{k})\|^{2}]\leq
L\mathbb{E}[\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}].$ (45)

Therefore, we can obtain

$\displaystyle\mathbb{E}[\|h_{x}^{k}\|^{2}]$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}\sigma_{f}^{2}+3b_{k}^{2}+3L^{2}\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}]+3\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|\nabla\ell_{i}(x_{k})\|^{2}]$
(46)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\tilde{\sigma}_{f}^{2}+3b_{k}^{2}+3L^{2}\sum_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}]$
(47)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\tilde{\sigma}_{f}^{2}+3b_{k}^{2}+3L^{2}\max_{i\in[n]}\mathbb{E}[\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}],$
(48)

where (1) comes from plugging (45) into (44) and (44) into (39), (2) comes
from the definition of $\tilde{\sigma}_{f}^{2}$ using
$\|\nabla\ell_{i}(x_{k})\|^{2}\leq C_{\ell}^{2}$ and
$\lambda_{k}\in\Delta_{n}$, and (3) comes from the fact that if
$\lambda\in\Delta_{n}$, then
$\sum\limits_{i=1}^{n}\lambda_{i}a_{i}\leq\max_{i\in[n]}a_{i}$.

###  A.4 Proof of Lemma 2

Lemma 2:  Under Assumptions 1, 2, and 3, when
$\beta^{2}(1+\sigma_{g})^{2}L_{g}^{2}\leq\beta\mu_{g}$, the iterates
$y_{k}^{i}$ satisfy the descent equation:

$\displaystyle\mathbb{E}[\|y_{k+1}^{i}-y_{i}^{*}(x_{k})\|^{2}]\leq\Big{(}1-\frac{\mu_{g}\beta}{2}\Big{)}\|y_{k}^{i}-y_{i}^{*}(x_{k-1})\|^{2}+\Big{(}\frac{2}{\mu_{g}\beta}-1\Big{)}G_{y}^{2}\|x_{k-1}-x_{k}\|^{2}+\beta^{2}\sigma_{g}^{2}.$
(49)

Proof:  First, observe that for a particular (fixed) realization of the
iterates $x_{1},\cdots,x_{k},y_{1},\cdots,y_{k}$, we have

$\displaystyle\mathbb{E}[\|h_{g}^{k,i}\|^{2}]$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\mathbb{E}[\|h_{g}^{k,i}-\mathbb{E}[h_{g}^{k,i}]\|^{2}]+\|\mathbb{E}[h_{g}^{k,i}]\|^{2}$
(50)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\mathbb{E}[\|h_{g}^{k,i}-\nabla
g_{i}(x_{k},y_{k})\|^{2}]+\|\nabla g_{i}(x_{k},y_{k})\|^{2}$ (51)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\sigma_{g}^{2}+(1+\sigma_{g}^{2})\|\nabla
g_{i}(x_{k},y_{k})\|^{2}$ (52)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\leq}}\sigma_{g}^{2}+(1+\sigma_{g}^{2})\|\nabla
g_{i}(x_{k},y_{k})-\nabla g(x_{k},y_{i}^{*}(x_{k}))\|^{2}$ (53)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\leq}}\sigma_{g}^{2}+(1+\sigma_{g}^{2})L_{g}^{2}\|y_{k}-y^{*}(x_{k})\|^{2},$
(54)

where (1) follows from algebra, (2) follows from
$\mathbb{E}[h_{g}^{k,i}]=\nabla g_{i}(x_{k},y_{k})$ in Assumption 3, (3) is
from (26) in Assumption 3, (4) is from $\nabla
g_{i}(x_{k},y_{i}^{*}(x_{k}))=0$ due to the optimality of $y_{i}^{*}(x_{k})$,
and (5) is due to the $L_{g}$-Lipschitz continuity of
$\nabla_{y}g_{i}(x,\cdot)$.

Next, we can bound the difference between $y_{k+1}^{i}$ and $y_{i}^{*}(x_{k})$
as the following, where again we assume that
$x_{1},\cdots,x_{k},y_{1},\cdots,y_{k}$ is fixed and the expectation is over
the stochasticity of the gradient estimates:

$\displaystyle\mathbb{E}[\|y_{k+1}^{i}-y_{i}^{*}(x_{k})\|^{2}]$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\mathbb{E}[\|y_{k}^{i}-\beta
h_{g}^{k,i}-y_{i}^{*}(x_{k})\|^{2}]$ (55)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\|y_{k}^{i}-y_{i}^{*}(x_{k})\|^{2}+\beta^{2}\mathbb{E}[\|h_{g}^{k,i}\|]^{2}-2\beta\langle
y_{k}^{i}-y_{i}^{*}(x_{k}),\nabla g_{i}(x_{k},y_{k})\rangle$ (56)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}(1-2\beta\mu_{g})\|y_{k}^{i}-y_{i}^{*}(x_{k})\|^{2}+\beta^{2}\mathbb{E}[\|h_{g}^{k,i}\|^{2}]$
(57)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\leq}}(1-2\beta\mu_{g})\|y_{k}^{i}-y_{i}^{*}(x_{k})\|^{2}+\beta^{2}\sigma_{g}^{2}+\beta^{2}(1+\sigma_{g}^{2})L_{g}^{2}\|y_{k}-y^{*}(x_{k})\|^{2}$
(58)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\leq}}(1-\beta\mu_{g})\|y_{k}^{i}-y_{i}^{*}(x_{k})\|^{2}+\beta^{2}\sigma_{g}^{2}$
(59)
$\displaystyle\stackrel{{\scriptstyle(6)}}{{\leq}}(1-\beta\mu_{g})\left[(1+c)\|y_{k}^{i}-y_{i}^{*}(x_{k-1})\|^{2}+\left(1+\frac{1}{c}\right)\|y_{i}^{*}(x_{k-1})-y_{i}^{*}(x_{k})\|^{2}\right]+\beta^{2}\sigma_{g}^{2}$
(60)
$\displaystyle\stackrel{{\scriptstyle(7)}}{{\leq}}(1-\beta\mu_{g})\left[(1+c)\|y_{k}^{i}-y_{i}^{*}(x_{k-1})\|^{2}+\left(1+\frac{1}{c}\right)G_{y}^{2}\|x_{k-1}-x_{k}\|^{2}\right]+\beta^{2}\sigma_{g}^{2},$
(61)

where (1) is by definition, and (2) is by direct algebra and the unbiasedness
assumption $\mathbb{E}[h_{g}^{k,i}]=\nabla g_{i}(x_{k},y_{k})$. (3) is from
strong convexity, $\beta\langle\nabla
g_{i}(x_{k},y_{k}),y_{k}^{i}-y_{i}^{*}(x_{k})\rangle\geq\beta\mu_{g}\|y_{k}^{i}-y_{i}^{*}(x_{k})\|^{2}$.
(4) is from (54), (5) is from the assumption
$\beta^{2}(1+\sigma_{g}^{2})L_{g}^{2}\leq\beta\mu_{g}$, (6) is from the
inequality $\|a+b\|^{2}\leq(1+1/c)\|a\|^{2}+(1+c)\|b\|^{2}$, and (7) is from
the $G_{y}$-lipschitzness of $y_{i}^{*}$ in Lemma 0 (A.1).

Then, we choose $c=\frac{\mu_{g}\beta}{2(1-\mu_{g}\beta)}$, so that
$(1-\beta\mu_{g})(1+c)=1-\beta\mu_{g}/2$ and $1+1/c=\frac{2}{\mu_{g}\beta}-1$.
Plugging these expressions into (61), we get

$\displaystyle\|y_{k+1}^{i}-y_{i}^{*}(x_{k})\|^{2}\leq\left(1-\frac{\mu_{g}\beta}{2}\right)\|y_{k}^{i}-y_{i}^{*}(x_{k-1})\|^{2}+\left(\frac{2}{\mu_{g}\beta}-1\right)G_{y}^{2}\|x_{k-1}-x_{k}\|^{2}+\beta^{2}\sigma_{g}^{2},$
(62)

which completes the proof.

###  A.5 Proof of Lemma 3

Lemma 3:  Let
$\mathcal{L}^{k}\triangleq\mathbb{E}[F(x_{k},\lambda_{k})]=\sum\limits_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\ell_{i}(x_{k})]$.
Under Assumptions 1, 2, and 3, assume that the iterates
$\\{x_{k},y_{k},\lambda_{k},\forall k\\}$ generated by MORBiT, then,
$\mathcal{L}$ satisfies the descent equation

$\mathcal{L}^{k+1}-\mathcal{L}^{k}\leq 4\alpha
L^{2}\max_{i\in[n]}\Delta_{y}^{k+1,i}+\left(\frac{L_{f}}{2}-\frac{1}{4\alpha}\right)\mathbb{E}[\|x_{k+1}-x_{k}\|^{2}]+\gamma\|h_{\lambda}^{k}\|\sqrt{n}B_{\ell}+4\alpha
b_{0}^{2}+2\alpha\sigma_{f}^{2}.$ (63)

Proof:  First, since $\ell_{i}$ is $L_{f}$-smooth, we know that for all
$i\in[n]$,

$\ell_{i}(x_{k+1})\leq\ell_{i}(x_{k})+\langle
x_{k+1}-x_{k},\nabla\ell_{i}(x_{k})\rangle+\frac{L_{f}}{2}\|x_{k+1}-x_{k}\|^{2}.$
(64)

Taking $(\lambda_{k})_{i}$ times the equation for $i$ in (64) and summing, we
can get

$\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k+1})\leq\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k})+\langle
x_{k+1}-x_{k},\sum_{i=1}^{n}(\lambda_{k})_{i}\nabla\ell_{i}(x_{k})\rangle+\frac{L_{f}}{2}\|x_{k+1}-x_{k}\|^{2}.$
(65)

Therefore, we have

$\displaystyle\sum_{i=1}^{n}(\lambda_{k+1})_{i}\ell_{i}(x_{k+1})-\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k})$
$\displaystyle\leq\underbrace{\sum_{i=1}^{n}(\lambda_{k+1})_{i}\ell_{i}(x_{k+1})-\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k+1})}_{(A)}+\underbrace{\langle
x_{k+1}-x_{k},\sum_{i=1}^{n}(\lambda_{k})_{i}\nabla\ell_{i}(x_{k})\rangle}_{(B)}+\frac{L_{f}}{2}\|x_{k+1}-x_{k}\|^{2}.$
(66)

Next, we bound $(A)$ and $(B)$ respectively as follows. First, we upper bound
term $(A)$. First, from the non-expansiveness of projections and
$\lambda_{k+1}=\mathrm{proj}_{\Delta_{n}}(\lambda_{k}+\gamma
h_{\lambda}^{k})$, we have $\|\lambda_{k+1}-\lambda_{k}\|\leq\|\gamma
h_{\lambda}^{k}\|$. Since $\lambda_{k+1},\lambda_{k}\in\Delta_{n}$,
$\|\lambda_{k+1}-\lambda_{k}\|\leq\sqrt{2}$. Therefore, we know that
$\|\lambda_{k+1}-\lambda_{k}\|\leq\Lambda\triangleq\min\\{\sqrt{2},\gamma\|h_{\lambda}^{k}\|$}.
Based on these facts, we can have

$\displaystyle(A)$
$\displaystyle=\sum_{i=1}^{n}(\lambda_{k+1})_{i}\ell_{i}(x_{k+1})-\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k+1})$
(67)
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\sum_{i=1}^{n}(\lambda_{k+1}-\lambda_{k})_{i}\ell_{i}(x_{k+1})$
(68)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\|\lambda_{k+1}-\lambda_{k}\|\|\langle\ell_{1}(x_{k+1}),\ell_{2}(x_{k+1}),\cdots,\ell_{n}(x_{k+1})\rangle\|$
(69) $\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\Lambda
B_{\ell}\sqrt{n}$ (70)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\leq}}\sqrt{n}\min\\{\sqrt{2},\|\gamma
h_{\lambda}^{k}\|\\}B_{\ell}$ (71)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\leq}}\sqrt{n}\gamma\|h_{\lambda}^{k}\|B_{\ell},$
(72)

where (1) is straightforward, (2) follows from Cauchy-Schwarz, (3) follows
from the update rule for $\lambda$ and the fact that $|\ell_{i}(\cdot)|\leq
B_{\ell}$, (4) follows from Assumption 1, $|\ell_{i}(x_{k+1})|\leq B_{\ell}$,
(4) is from plugging in the definition of $\Lambda$, and (5) follows from
$\gamma\|h^{k}_{\lambda}\|\leq\sqrt{2}$.

Then, we upper bound $(B)$. First, from the non-expansiveness of projection
and the update rule $x_{k+1}=\mathrm{proj}_{\mathcal{X}}(x_{k}-\alpha
h_{x}^{k})$, we know that

$\displaystyle\|x_{k+1}-x_{k}+\alpha h_{x}^{k}\|^{2}\leq\|-\alpha
h_{x}^{k}\|^{2},$ (73)
$\displaystyle\Rightarrow\|x_{k+1}-x_{k}\|^{2}+2\alpha\langle
x_{k+1}-x_{k},h_{x}^{k}\rangle\leq 0,$ (74)
$\displaystyle\Rightarrow\frac{1}{2\alpha}\|x_{k+1}-x_{k}\|^{2}+\langle
x_{k+1}-x_{k},h_{x}^{k}\rangle\leq 0.$ (75)

Therefore, we can have

$\displaystyle(B)$
$\displaystyle=\Big{\langle}\sum_{i=1}^{n}(\lambda_{k})_{i}\nabla\ell_{i}(x_{k}),x_{k+1}-x_{k}\Big{\rangle}$
(76)
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\langle\nabla_{x}F(x_{k},\lambda_{k}),x_{k+1}-x_{k}\rangle$
(77)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\langle\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k},x_{k+1}-x_{k}\rangle$
$\displaystyle+\langle\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+B_{k},x_{k+1}-x_{k}\rangle$ (78)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\langle\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k},x_{k+1}-x_{k}\rangle$
$\displaystyle+\langle\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+B_{k}-h_{x}^{k},x_{k+1}-x_{k}\rangle-\frac{1}{2\alpha}\|x_{k+1}-x_{k}\|^{2}$
(79)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\leq}}\frac{1}{2c}\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k}\|^{2}+\frac{c}{2}\|x_{k+1}-x_{k}\|^{2}$
$\displaystyle+\frac{1}{2d}\|\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+B_{k}-h_{x}^{k}\|^{2}+\frac{d}{2}\|x_{k+1}-x_{k}\|^{2}-\frac{1}{2\alpha}\|x_{k+1}-x_{k}\|^{2}$
(80)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\leq}}\frac{1}{2c}\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k}\|^{2}+\frac{c}{2}\|x_{k+1}-x_{k}\|^{2}$
$\displaystyle+\frac{\sigma_{f}^{2}}{2d}+\frac{d}{2}\|x_{k+1}-x_{k}\|^{2}-\frac{1}{2\alpha}\|x_{k+1}-x_{k}\|^{2}$
(81)
$\displaystyle\stackrel{{\scriptstyle(6)}}{{=}}\frac{1}{2c}\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k}\|^{2}+\left(\frac{c+d}{2}-\frac{1}{2\alpha}\right)\|x_{k+1}-x_{k}\|^{2}+\frac{\sigma_{f}^{2}}{2d},$
(82)

where (1) is by definition of $F(x_{k},\lambda_{k})$, (2) is from adding and
subtracting $\langle\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+B_{k},x_{k+1}-x_{k}\rangle$, (3) is from
adding (75) to the previous inequality, (4) is from applying the inequality
$\langle a,b\rangle\leq\frac{1}{2c}\|a\|^{2}+\frac{c}{2}\|b\|^{2}$ to both
inner product terms, (5) is from (25), and (6) is from algebra.

Plugging in our expressions for (A) and (B) into (66), we get

$\displaystyle\sum_{i=1}^{n}(\lambda_{k+1})_{i}\ell_{i}(x_{k+1})-\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k})$
$\displaystyle\leq\gamma\|h_{\lambda}^{k}\|\sqrt{n}B_{\ell}+\frac{1}{2c}\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k}\|^{2}$ (83)
$\displaystyle+\left(\frac{c+d+L_{f}}{2}-\frac{1}{2\alpha}\right)\|x_{k+1}-x_{k}\|^{2}+\frac{\sigma_{f}^{2}}{2d}.$
(84)

Next, we work on bounding $\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k}\|^{2}$ in (82). Observe that

$\displaystyle\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-B_{k}\|^{2}$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}2\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})\|^{2}+2\|B_{k}\|^{2}$ (85)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}2\|\nabla_{x}F(x_{k},\lambda_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})\|^{2}+2b_{k}^{2}$ (86)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}2\Bigg{\|}\sum_{i=1}^{n}[(\lambda_{k})_{i}(\nabla\ell_{i}(x_{k})-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}f_{i}(x_{k},y_{k+1}^{i}))]\Bigg{\|}^{2}+2b_{k}^{2}$ (87)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\leq}}2\left(\sum_{i=1}^{n}(\lambda_{k})_{i}L\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|\right)^{2}+2b_{k}^{2}$
(88)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\leq}}2L^{2}\max_{i\in[n]}\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}+2b_{k}^{2},$
(89)

where (1) comes from $\|a+b\|^{2}\leq 2\|a\|^{2}+2\|b\|^{2}$, (2) comes from
$\|B_{k}\|\leq b_{k}$ in Assumption 3, (3) is from expanding the definitions
of $\nabla_{x}F(\cdot,\cdot)$ and $\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(\cdot,\cdot,\cdot)$, (4) is from Lemma 0 (A.1). Therefore, plugging
in (89) into (84), using $c=d=\frac{1}{4\alpha}$, and taking expectation over
the stochasticity of the gradient estimates, we get:

$\displaystyle\sum_{i=1}^{n}(\lambda_{k+1})_{i}\ell_{i}(x_{k+1})-\sum_{i=1}^{n}(\lambda_{k})_{i}\ell_{i}(x_{k})$
$\displaystyle\quad\leq\sqrt{n}\gamma\|h_{\lambda}^{k}\|B_{\ell}+4\alpha
L^{2}\max_{i\in[n]}\mathbb{E}\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}+4\alpha
b_{k}^{2}+\left(\frac{L_{f}}{2}-\frac{1}{4\alpha}\right)\mathbb{E}[\|x_{k+1}-x_{k}\|^{2}]+2\alpha\sigma_{f}^{2}.$
(90)

Now, observe that the LHS looks like a telescoping sum. To make this more
apparant, define
$\mathcal{L}^{k}=\sum\limits_{i=1}^{n}(\lambda_{k})_{i}\mathbb{E}[\ell_{i}(x_{k})]$
and $\Delta_{y}^{k,i}=\mathbb{E}[\|y_{k}^{i}-y_{i}^{*}(x_{k-1})\|^{2}]$.
Therefore, with the assumption that $b_{k}\leq b_{0}$, we have

$\mathcal{L}^{k+1}-\mathcal{L}^{k}\leq 4\alpha
L^{2}\max_{i\in[n]}\Delta_{y}^{k+1,i}+\left(\frac{L_{f}}{2}-\frac{1}{4\alpha}\right)\mathbb{E}[\|x_{k+1}-x_{k}\|^{2}]+\sqrt{n}\gamma\|h_{\lambda}^{k}\|B_{\ell}+4\alpha
b_{0}^{2}+2\alpha\sigma_{f}^{2}.$ (91)

###  A.6 Proof of Lemma 4

Lemma 4:  Assume that
$\Omega^{i},\Theta^{i},\Upsilon^{i},\lambda_{i},c_{0},c_{1},c_{2},d_{0},d_{1},d_{2}$
are real numbers such that for all $0\leq k\leq K-1$,

$\Omega^{k+1}\leq\Omega^{k}-c_{0}\Theta^{k+1}+c_{1}\max_{i\in[n]}\Upsilon^{k+1,i}+c_{2}$
(92)

and also for all $1\leq k\leq K,1\leq i\leq N$,

$\Upsilon^{k+1,i}\leq(1-d_{0})\Upsilon^{k,i}+d_{1}\Theta^{k}+d_{2}.$ (93)

In addition, assume that $1-d_{0}>0$, $d_{0}-d_{1}c_{1}c_{0}^{-1}>0$ and
$c_{0}-c_{1}d_{1}d_{0}^{-1}$, and that $\Upsilon^{k,i},\Omega^{k}\geq 0$ for
all $k,i$. Then, if $i_{0}=\operatorname*{arg\,max}_{i\in[n]}\Upsilon^{0,i}$,
we have

$\displaystyle\frac{1}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Upsilon^{k,i}\leq(d_{0}-d_{1}c_{0}^{-1}c_{1})^{-1}\left(\frac{\Upsilon^{0,i_{0}}+d_{1}\Theta^{0}+d_{2}+d_{1}c_{0}^{-1}\Omega^{0}}{K}+d_{1}c_{0}^{-1}c_{2}+d_{2}\right).$
(94)

Proof:  First, let
$i_{k}\triangleq\operatorname*{arg\,max}\limits_{i\in[n]}\Upsilon^{k,i}$, so
that $\Upsilon^{k,i_{k}}=\max\limits_{i\in[n]}\Upsilon^{k,i}$. Summing (92)
from $k=0,1,\cdots,K-1$, we get:

$c_{0}\sum_{k=1}^{K}\Theta^{k}\leq\Omega^{0}-\Omega^{K}+c_{1}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}+c_{2}K.$
(95)

Next, we apply (93) for $i=i_{k+1}$. Noting that $1-d_{0}>0$ and
$\Upsilon^{k,i_{k+1}}\leq\Upsilon^{k,i_{k}}$ by definition of $i_{k}$, we have

$\displaystyle\Upsilon^{k+1,i_{k+1}}$
$\displaystyle\leq(1-d_{0})\Upsilon^{k,i_{k+1}}+d_{1}\Theta^{k}+d_{2}$ (96)
$\displaystyle\leq(1-d_{0})\Upsilon^{k,i_{k}}+d_{1}\Theta^{k}+d_{2}.$ (97)

Then summing for $k=1$ to $K$, we get

$d_{0}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}\leq\Upsilon^{1,i_{1}}-\Upsilon^{K+1,i_{K+1}}+d_{1}\sum_{k=1}^{K}\Theta^{k}+d_{2}K.$
(98)

Now, we have

$\displaystyle d_{0}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}\Upsilon^{1,i_{1}}-\Upsilon^{K+1,i_{K+1}}+d_{1}c_{0}^{-1}\left(\Omega^{0}-\Omega^{K}+c_{1}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}+c_{2}K\right)+d_{2}K$
(99)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\Upsilon^{1,i_{1}}+d_{1}c_{0}^{-1}\left(\Omega^{0}+c_{1}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}+c_{2}K\right)+d_{2}K$
(100)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\Upsilon^{1,i_{1}}+d_{1}c_{0}^{-1}\Omega^{0}+d_{1}c_{0}^{-1}c_{1}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}+d_{1}c_{0}^{-1}c_{2}K+d_{2}K,$
(101)

where (1) holds from plugging (95) into (98), (2) is true because
$\Upsilon,\Omega\geq 0$, and (3) follows from the distributive property. We
can rewrite this equation as

$\displaystyle(d_{0}-d_{1}c_{0}^{-1}c_{1})\sum_{k=1}^{K}\Upsilon^{k,i_{k}}\leq\Upsilon^{1,i_{1}}+d_{1}c_{0}^{-1}\Omega^{0}+d_{1}c_{0}^{-1}c_{2}K+d_{2}K$
(102)
$\displaystyle\Rightarrow\frac{1}{K}\sum_{k=1}^{K}\Upsilon^{k,i_{k}}\leq(d_{0}-d_{1}c_{0}^{-1}c_{1})^{-1}\left(\frac{\Upsilon^{1,i_{1}}+d_{1}c_{0}^{-1}\Omega^{0}}{K}+d_{1}c_{0}^{-1}c_{2}+d_{2}\right).$
(103)

By plugging in $\Upsilon^{1,i_{1}}=\Upsilon^{0,i_{0}}+d_{1}\Theta^{0}+d_{2}$
into (103), we get the statement of the lemma.

###  A.7 Proof of Theorem 1, $y$ bound

Now, we show how get from Lemmas 2, 3, and 4 to the convergence of the $y$’s,
(8). First, we define the following quantities, which will be plugged into
Lemma 4:

$\displaystyle\Omega^{k}$ $\displaystyle=\mathcal{L}^{k},$
$\displaystyle\Theta^{k}$ $\displaystyle=\mathbb{E}[\|x_{k}-x_{k-1}\|^{2}],$
$\displaystyle\Upsilon^{k,i}$ $\displaystyle=\Delta_{y}^{k,i},$ $\displaystyle
c_{0}$ $\displaystyle=\frac{1}{4\alpha}-\frac{L_{f}}{2},$ $\displaystyle
c_{1}$ $\displaystyle=4\alpha L^{2},$ $\displaystyle c_{2}$
$\displaystyle=\sqrt{n}\gamma\|h_{\lambda}^{k}\|B_{\ell}+4\alpha
b_{0}^{2}+2\alpha\sigma_{f}^{2},$ $\displaystyle d_{0}$
$\displaystyle=\mu_{g}\beta/2,$ $\displaystyle d_{1}$
$\displaystyle=\left(\frac{2}{\mu_{g}\beta}-1\right)G_{y}^{2},$ $\displaystyle
d_{2}$ $\displaystyle=\beta^{2}\sigma_{g}^{2}.$ (104)

Next, recall that our step sizes were

$\displaystyle\alpha=\min\Bigg{(}\frac{\mu_{g}}{16G_{y}L}\nu,\frac{K^{-3/5}}{4G_{y}L}\Bigg{)},\beta=\min\Big{(}\nu,\frac{4K^{-2/5}}{\mu_{g}}\Big{)},\gamma=\frac{2K^{-3/5}}{B_{\ell}n^{1/2}},$
(105)

where
$\nu=\min(\frac{\mu_{g}}{L_{g}^{2}(1+\sigma_{g}^{2})},\frac{2}{\mu_{g}})$.
Note that the choice of $\nu$ was motivated by the conditions of Lemma 5.
First, observe that $1-d_{0}>0$ is true because we chose $\beta<2/\mu_{g}$.
Now, observe that $\frac{\alpha}{\beta}\leq\frac{\mu_{g}}{16G_{y}L}$. Finally,
will next show that $d_{0}-d_{1}c_{1}(c_{0})^{-1}>0$ and
$c_{0}-c_{1}d_{1}(d_{0})^{-1}>0$, completing the set of conditions in Lemma 5.
By direct algebraic manipulation, we have

$\displaystyle d_{0}-d_{1}c_{1}(c_{0})^{-1}$
$\displaystyle=\frac{\mu_{g}\beta}{2}-\frac{\left(\frac{2}{\mu_{g}\beta}-1\right)G_{y}^{2}\cdot
4L^{2}\alpha}{\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)}\geq\frac{\mu_{g}\beta}{2}-\frac{\frac{2}{\mu_{g}\beta}\cdot
G_{y}^{2}\cdot 4L^{2}\alpha}{\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)}$
$\displaystyle\geq\frac{\mu_{g}\beta}{2}-\frac{8L^{2}G_{y}^{2}\alpha}{\beta\mu_{g}\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)}\geq\frac{\mu_{g}\beta}{2}-\frac{64L^{2}G_{y}^{2}}{\mu_{g}^{2}}\cdot\frac{\alpha^{2}}{\beta^{2}}\cdot\mu_{g}\beta$
$\displaystyle\geq\frac{\mu_{g}\beta}{2}-\frac{64L^{2}G_{y}^{2}}{\mu_{g}^{2}}\cdot\frac{\mu_{g}^{2}}{256G_{y}^{2}L^{2}}\cdot\mu_{g}\beta=\frac{\mu_{g}\beta}{2}-\frac{\mu_{g}\beta}{4}=\frac{\mu_{g}\beta}{4}.$
(106)

Similarly, we also have

$\displaystyle c_{0}-c_{1}d_{1}(d_{0})^{-1}$
$\displaystyle=\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)-\frac{\cdot
4L^{2}\alpha\cdot\left(\frac{2}{\mu_{g}\beta}-1\right)G_{y}^{2}}{\frac{\mu_{g}\beta}{2}}=\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)-\frac{8L^{2}\alpha\left(\frac{2}{\mu_{g}\beta}-1\right)G_{y}^{2}}{\mu_{g}\beta}$
$\displaystyle\geq\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)-\frac{8L^{2}\alpha\left(\frac{2}{\mu_{g}\beta}\right)G_{y}^{2}}{\mu_{g}\beta}\geq\left(\frac{1}{4\alpha}-\frac{L_{f}}{2}\right)-\frac{16L^{2}\alpha
G_{y}^{2}}{\mu_{g}^{2}}\cdot\frac{\alpha^{2}}{\beta^{2}}\cdot\frac{1}{\alpha}$
$\displaystyle\geq\frac{1}{8\alpha}-\frac{16L^{2}\alpha
G_{y}^{2}}{\mu_{g}^{2}}\cdot\frac{\mu_{g}^{2}}{256G_{y}^{2}L^{2}}\cdot\frac{1}{\alpha}=\frac{1}{8\alpha}-\frac{1}{16\alpha}=\frac{1}{16\alpha}.$
(107)

Now, we can bound the optimality of $y$ by bounding the maximum difference
$\Delta_{y}^{k,i}$:

$\displaystyle\frac{1}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}(d_{0}-d_{1}c_{0}^{-1}c_{1})^{-1}\left(\frac{\Upsilon^{0,i_{0}}+d_{1}\Theta^{0}+d_{2}+d_{1}c_{0}^{-1}\Omega^{0}}{K}+d_{1}c_{0}^{-1}c_{2}+d_{2}\right)$
(108)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\frac{4}{\mu_{g}\beta}\left(\frac{\Upsilon^{0,i_{0}}+d_{1}\Theta^{0}+d_{2}+d_{1}c_{0}^{-1}\Omega^{0}}{K}+d_{1}c_{0}^{-1}c_{2}+d_{2}\right)$
(109)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\frac{4}{\mu_{g}\beta}\left(\frac{\Delta_{y}^{0,i_{0}}+\beta^{2}\sigma_{g}^{2}+\frac{2G_{y}^{2}}{\mu_{g}\beta}(8\alpha)\mathcal{L}^{0}}{K}+\frac{2G_{y}^{2}}{\mu_{g}\beta}(8\alpha)(\sqrt{n}\gamma\|h_{\lambda}^{k}\|B_{\ell}+4\alpha
b_{0}^{2}+2\alpha\sigma_{f}^{2})+\beta^{2}\sigma_{g}^{2}\right)$ (110)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{=}}\frac{\frac{4\Delta_{y}^{0,i_{0}}}{\mu_{g}\beta}+4\beta\frac{\sigma_{g}^{2}}{\mu_{g}}+\frac{64G_{y}^{2}\alpha}{\mu_{g}^{2}\beta^{2}}}{K}+\frac{64G_{y}^{2}\alpha}{\mu_{g}^{2}\beta^{2}}\left(\sqrt{n}\gamma\|h_{\lambda}^{k}\|B_{\ell}+4\alpha
b_{0}^{2}+2\alpha\sigma_{f}^{2}\right)+\frac{4\beta\sigma_{g}^{2}}{\mu_{g}}$
(111)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{=}}\frac{4\Delta_{y}^{0,i_{0}}}{\mu_{g}}\frac{1}{\beta
K}+\frac{4\sigma_{g}^{2}}{\mu_{g}}\frac{\beta}{K}+\frac{64G_{y}^{2}}{\mu_{g}^{2}}\frac{\alpha}{\beta^{2}K}+\frac{64G_{y}^{2}}{\mu_{g}^{2}}\gamma||h_{\lambda}^{k}||B_{\ell}\frac{\sqrt{n}\alpha}{\beta^{2}}+\frac{64G_{y}^{2}(4b_{0}^{2}+2\sigma_{f}^{2})}{\mu_{g}^{2}}\frac{\alpha^{2}}{\beta^{2}}+\frac{4\sigma_{g}^{2}}{\mu_{g}}\beta$
(112)
$\displaystyle\stackrel{{\scriptstyle(6)}}{{\leq}}\frac{\Delta_{y}^{0,i_{0}}/\mu_{g}}{K^{3/5}}+\frac{16\sigma_{g}^{2}/\mu_{g}^{2}}{K^{7/5}}+\frac{G_{y}/L}{K^{4/5}}+\frac{2\sqrt{n}B_{\ell}G_{y}/L}{K^{2/5}}+\frac{(b_{0}^{2}+\frac{1}{2}\sigma_{f}^{2})/(L^{2})}{K^{2/5}}+\frac{16\sigma_{g}^{2}/\mu_{g}^{2}}{K^{2/5}}.$
(113)

Here, (1) follows directly from plugging $\Upsilon^{k,i}$ from (104) into
Lemma 4, (2) follows from (106), (3) comes from plugging in the rest of (104),
(4) is direct algebraic manipulation, (5) separates the step sizes and $n,K$
factors from the rest of the constants, and (6) applies the definition of the
step sizes.

###  A.8 Proof of Lemma 5

Lemma 5:  For any $\lambda\in\Delta_{n}$, under Assumptions 1, 2, and 3,
assume that the iterates $\\{x_{k},y_{k},\lambda_{k},\forall k\\}$ generated
by MORBiT, then we have

$\frac{1}{K}\mathbb{E}\left[\sum_{k=1}^{K}F(x_{k},\lambda)-F(x_{k},\lambda_{k})\right]\leq\frac{\sqrt{2}}{2}\left(B_{\ell}\sqrt{n}K^{-3/5}+B_{\ell}\sqrt{n}K^{-2/5}\right)+\frac{G_{f}}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}.$
(114)

Proof:  Recall that we defined

$F(x,\lambda)\coloneqq\sum_{i=1}^{n}\lambda_{i}f_{i}(x,y_{i}^{*}(x))=\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x).$
(115)

For a fixed realization of $x_{1},\cdots,x_{k},y_{1},\cdots,y_{k}$, we have

$\displaystyle F(x_{k},\lambda)-F(x_{k},\lambda_{k})$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}f_{i}(x_{k},y_{i}^{*}(x_{k}))$
(116)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}(f_{i}(x_{k},y_{i}^{*}(x_{k}))-f_{i}(x_{k},y_{k+1}^{i})+f_{i}(x_{k},y_{k+1}^{i})))$
(117)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{=}}\langle\lambda-\lambda_{k},\langle
f_{1}(x_{k},y_{k+1}^{1}),\cdots,f_{n}(x_{k},y_{k+1}^{n})\rangle+\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}(f_{i}(x_{k},y_{i}^{*}(x_{k}))-f_{i}(x_{k},y_{k+1}^{i}))$
(118)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{=}}\langle\lambda-\lambda_{k},h_{\lambda}^{k}\rangle+\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}(f_{i}(x_{k},y_{i}^{*}(x_{k}))-f_{i}(x_{k},y_{k+1}^{i}))$
(119)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{=}}\frac{\|\lambda-\lambda_{k}\|^{2}+\gamma^{2}\|h_{\lambda}^{k}\|^{2}-\|\lambda-\lambda_{k}-\gamma
h_{\lambda}^{k}\|^{2}}{2\gamma}+\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}(f_{i}(x_{k},y_{i}^{*}(x_{k}))-f_{i}(x_{k},y_{k+1}^{i}))$
(120)
$\displaystyle\stackrel{{\scriptstyle(6)}}{{\leq}}\frac{\|\lambda-\lambda_{k}\|^{2}+\gamma^{2}\|h_{\lambda}^{k}\|^{2}-\|\lambda-\lambda_{k+1}\|^{2}}{2\gamma}+\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}(f_{i}(x_{k},y_{i}^{*}(x_{k}))-f_{i}(x_{k},y_{k+1}^{i}))$
(121)
$\displaystyle\stackrel{{\scriptstyle(7)}}{{\leq}}\frac{\|\lambda-\lambda_{k}\|^{2}+\gamma^{2}\|h_{\lambda}^{k}\|^{2}-\|\lambda-\lambda_{k+1}\|^{2}}{2\gamma}+G_{f}\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}(y_{i}^{*}(x_{k})-y_{k+1}^{i})$
(122)

where (1) comes from the definition of $F$, (2) follows from adding and
subtracting $(\lambda-\lambda_{k})_{i}f_{i}(x_{k},y_{k+1}^{i})$ terms, (3)
follows from splitting the preceding sum and writing the first term as a dot
product, (4) follows from definition of $h_{\lambda}^{k}$, (5) uses $\langle
a,b\rangle=\frac{\|a\|^{2}+\|b\|^{2}-\|a-b\|^{2}}{2}$, (6) follows from the
update $\lambda_{k+1}=\mathrm{proj}_{\Delta_{n}}(\lambda_{k}-\gamma
h_{\lambda}^{k})$ and the projection property, and (7) follows from
Lipschitzness of $f$.

Therefore, applying the telescoping sum by adding the preceding inequality
over $k=1,2,\ldots,K$, and taking expectation, we get:

$\displaystyle\mathbb{E}\left[\sum_{k=1}^{K}F(x_{k},\lambda)-F(x_{k},\lambda_{k})\right]$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}\frac{\gamma}{2}\sum_{k=1}^{K}\mathbb{E}\|h_{\lambda}^{k}\|^{2}+\frac{\mathbb{E}[\|\lambda-\lambda_{1}\|^{2}]}{2\gamma}+G_{f}\sum_{k=1}^{K}\sum_{i=1}^{n}(\lambda-\lambda_{k})_{i}\Delta_{y}^{k,i}$
(123)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\frac{\gamma}{2}\sum_{k=1}^{K}\mathbb{E}\|h_{\lambda}^{k}\|^{2}+\frac{1}{\gamma}+G_{f}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}$
(124)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\leq}}\frac{nKB^{2}\gamma}{2}+\frac{1}{\gamma}+G_{f}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}$
(125)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\leq}}\frac{\sqrt{2}}{2}\left(B_{\ell}\sqrt{n}K^{2/5}+B_{\ell}\sqrt{n}K^{3/5}\right)+G_{f}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}$
(126)

where (1) follows directly from (122) and the telescoping sum, (2) follows
from $\lambda\in\Delta_{n}$ and $\|\lambda-\lambda_{1}\|^{2}\leq 2$, (3)
follows from $\|h_{\lambda}^{k}\|^{2}\leq nB^{2}$, and (4) follows from
selecting $\gamma=\frac{\sqrt{2}}{B_{\ell}\sqrt{n}K^{3/5}}$.

Therefore, we obtain

$\frac{1}{K}\mathbb{E}\left[\sum_{k=1}^{K}F(x_{k},\lambda)-F(x_{k},\lambda_{k})\right]\leq\frac{\sqrt{2}}{2}\left(B_{\ell}\sqrt{n}K^{-3/5}+B_{\ell}\sqrt{n}K^{-2/5}\right)+G_{f}O(\sqrt{n}K^{-2/5}).$
(127)

###  A.9 Proof of Lemma 6

Lemma 6:  Under Assumptions 1, 2, and 3, assume that the iterates
$\\{x_{k},y_{k},\lambda_{k},\forall k\\}$ generated by MORBiT, then we have

$\displaystyle\frac{1}{K}\sum_{k=1}^{K}\|\hat{x}(x_{k})-x_{k}\|^{2}$ (128)
$\displaystyle\leq\frac{4}{-\mu_{\ell}+\rho}\Bigg{(}\frac{\Phi_{1/\rho}(x_{0})}{\rho\alpha
K}+\frac{2}{-\mu_{\ell}+\rho}b_{0}^{2}+\left(\frac{2L^{2}}{-\mu_{\ell}+\rho}+\frac{3L^{2}\alpha}{2}\right)\Bigg{(}\frac{1}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}\Bigg{)}+\frac{\alpha}{2}\left(\tilde{\sigma}_{f}^{2}+3b_{0}^{2}\right)\Bigg{)}.$

Proof:  Recall that we defined the Moreau envelope and proximal map as
follows:

$\Phi_{1/\rho}(z)\triangleq\min_{x}\frac{\rho}{2}\|x-z\|^{2}+\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x),\qquad\hat{x}(z)\triangleq\operatorname*{arg\,min}_{x}\frac{\rho}{2}\|x-z\|^{2}+\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x).$
(129)

Therefore, we have

$\displaystyle\Phi_{1/\rho}(x_{k+1})$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\sum_{i}\lambda_{i}\ell_{i}(\hat{x}(x_{k+1}))+\frac{\rho}{2}\|x_{k+1}-\hat{x}(x_{k+1})\|^{2}$
(130)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{\leq}}\sum_{i}\lambda_{i}\ell_{i}(\hat{x}(x_{k}))+\frac{\rho}{2}\|x_{k+1}-\hat{x}(x_{k})\|^{2}$
(131)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{=}}\sum_{i}\lambda_{i}\ell_{i}(\hat{x}(x_{k}))+\frac{\rho}{2}\|x_{k+1}-x_{k}+x_{k}-\hat{x}(x_{k})\|^{2}$
(132)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{=}}\sum_{i}\lambda_{i}\ell_{i}(\hat{x}(x_{k}))+\frac{\rho}{2}\|x_{k+1}-x_{k}\|^{2}+\frac{\rho}{2}\|x_{k}-\hat{x}(x_{k})\|^{2}+\rho\langle
x_{k+1}-x_{k},x_{k}-\hat{x}(x_{k})\rangle$ (133)

where (1) is by definition of the proximal map, (2) comes from the optimality
of the Moreau envelope, (3) is by adding and subtracting $x_{k}$, and (4) is
from expanding out $\|a+b\|^{2}$ into $\|a\|^{2}+\|b\|^{2}+2\langle
a,b\rangle$. Next, from the optimality condition of the update
$x_{k+1}=\mathrm{proj}_{\Delta_{n}}(x_{k}-\alpha h_{x}^{k})$, we have

$\displaystyle\langle x_{k+1}-\hat{x}(x_{k}),x_{k+1}-x_{k}+\alpha
h_{x}^{k}\rangle\leq 0$ (134)
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\Rightarrow}}\langle
x_{k+1}-x_{k}+x_{k}-\hat{x}(x_{k}),x_{k+1}-x_{k}+\alpha h_{x}^{k}\rangle\leq
0$ (135) $\displaystyle\stackrel{{\scriptstyle(2)}}{{\Rightarrow}}\langle
x_{k+1}-x_{k},x_{k}-\hat{x}(x_{k})\rangle\leq-\|x_{k+1}-x_{k}\|^{2}-\alpha\langle
h_{x}^{k},x_{k+1}-x_{k}\rangle-\alpha\langle
h_{x}^{k},x_{k}-\hat{x}(x_{k})\rangle$ (136)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{\Rightarrow}}\rho\langle
x_{k+1}-x_{k},x_{k}-\hat{x}(x_{k})\rangle\leq-\rho\|x_{k+1}-x_{k}\|^{2}-\rho\alpha\langle
h_{x}^{k},x_{k+1}-x_{k}\rangle-\rho\alpha\langle
h_{x}^{k},x_{k}-\hat{x}(x_{k})\rangle$ (137)
$\displaystyle\stackrel{{\scriptstyle(4)}}{{\Rightarrow}}\rho\langle
x_{k+1}-x_{k},x_{k}-\hat{x}(x_{k})\rangle\leq-\rho\|x_{k+1}-x_{k}\|^{2}+\rho\langle\alpha
h_{x}^{k},x_{k}-x_{k+1}\rangle+\rho\alpha\langle
h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$ (138)
$\displaystyle\stackrel{{\scriptstyle(5)}}{{\Rightarrow}}\rho\langle
x_{k+1}-x_{k},x_{k}-\hat{x}(x_{k})\rangle\leq-\rho\|x_{k+1}-x_{k}\|^{2}$
$\displaystyle\qquad+\frac{\rho}{2}\left(\|\alpha
h_{x}^{k}\|^{2}+\|x_{k}-x_{k+1}\|^{2}\right)+\rho\alpha\langle
h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$ (139)
$\displaystyle\stackrel{{\scriptstyle(6)}}{{\Rightarrow}}\rho\langle
x_{k+1}-x_{k},x_{k}-\hat{x}(x_{k})\rangle\leq-\frac{\rho}{2}\|x_{k+1}-x_{k}\|^{2}+\frac{\rho\alpha^{2}}{2}\|h_{x}^{k}\|^{2}+\rho\alpha\langle
h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$ (140)

where (1) is by adding and subtracting $x_{k}$, (2) is from distributing the
inner product
$\langle(x_{k+1}-x_{k})+(x_{k}-\hat{x}(x_{k})),(x_{k+1}-x_{k})+\alpha
h_{x}^{k}\rangle$, (3) is from multiplying both sides by $\rho$, (4) is from
simple algebra, (5) is from rewriting $\langle
a,b\rangle\leq\frac{\|a\|^{2}+\|b\|^{2}}{2}$, and (6) is from combining terms.
Therefore, substituting (140) into (133), we get

$\displaystyle\Phi_{1/\rho}(x_{k+1})$
$\displaystyle\leq\sum_{i}\lambda_{i}\ell_{i}(\hat{x}(x_{k}))+\frac{\rho}{2}\|x_{k}-\hat{x}(x_{k})\|^{2}+\frac{\rho\alpha^{2}}{2}\|h_{x}^{k}\|^{2}+\rho\alpha\langle
h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$ (141)
$\displaystyle=\Phi_{1/\rho}(x_{k})+\frac{\rho\alpha^{2}}{2}\|h_{x}^{k}\|^{2}+\rho\alpha\langle
h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$ (142)

where the second equality is by definition of the Moreau envelope. Now, we
bound the last term in (142):

$\displaystyle\langle$ $\displaystyle h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$
(143)
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\langle\hat{x}(x_{k})-x_{k},h_{x}^{k}-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})+\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-\nabla_{x}F(x_{k},\lambda_{k})+\nabla_{x}F(x_{k},\lambda_{k})\rangle$
(144)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\underbrace{\langle\hat{x}(x_{k})-x_{k},B_{k}\rangle}_{(A)}+\underbrace{\langle\hat{x}(x_{k})-x_{k},\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-\nabla_{x}F(x_{k},\lambda_{k})\rangle}_{(B)}$
$\displaystyle\quad+\underbrace{\langle\hat{x}(x_{k})-x_{k},\nabla_{x}F(x_{k},\lambda_{k})\rangle}_{(C)}$
(145)

where (1) follows from adding and subtracting $\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k}),\nabla_{x}F(x_{k},\lambda_{k})$ terms
and (2) is from splitting the inner product and applying $h_{x}^{k}-\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1}=B_{k}$ from Assumption 3. To bound $(A)$ and $(B)$,
we simply apply $\langle
a,b\rangle\leq\frac{c}{4}\|a\|^{2}+\frac{1}{c}\|b\|^{2}$ to both inner
products:

$\displaystyle(A)$
$\displaystyle=\langle\hat{x}(x_{k})-x_{k},B_{k}\rangle\leq\frac{c}{4}\|\hat{x}(x_{k})-x_{k}\|^{2}+\frac{1}{c}b_{k}^{2}$
(146) $\displaystyle(B)$ $\displaystyle=\langle\hat{x}(x_{k})-x_{k},\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-\nabla_{x}F(x_{k},\lambda_{k})\rangle$
$\displaystyle\leq\frac{1}{c}\|\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-\nabla_{x}F(x_{k},\lambda_{k})\|^{2}+\frac{c}{4}\|\hat{x}(x_{k})-x_{k}\|^{2}.$
(147)

We proceed to bound $(C)$. First, from weak convexity of $\mu_{\ell}$, we have
that for all $i\in[n]$,

$\ell_{i}(\hat{x}(x_{k}))\geq\ell_{i}(x_{k})+\langle\nabla\ell(x_{k}),\hat{x}(x_{k})-x_{k}\rangle-\frac{\mu_{\ell}}{2}\|\hat{x}(x_{k})-x_{k}\|^{2}.$
(148)

Taking $\lambda_{i}$ times the $i$th of these equations, we get

$\sum_{i=1}^{n}\lambda_{i}\ell_{i}(\hat{x}(x_{k}))\geq\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x_{k})+\big{\langle}\sum_{i=1}^{n}\nabla\ell_{i}(x_{k}),\hat{x}(x_{k})-x_{k}\big{\rangle}-\frac{\mu_{\ell}}{2}\|\hat{x}(x_{k})-x_{k}\|^{2}.$
(149)

By definition of the Moreau envelope, we also have

$\sum_{i=1}^{n}\lambda_{i}\ell_{i}(x_{k})\geq\sum_{i=1}^{n}\lambda_{i}\ell_{i}(\hat{x}(x_{k}))+\frac{\rho}{2}\|\hat{x}(x_{k})-x_{k}\|^{2}.$
(150)

Adding (149) and (150), we have

$\frac{\mu_{\ell}-\rho}{2}\|\hat{x}(x_{k})-x_{k}\|^{2}\geq\langle\nabla_{x}F(x_{k},\lambda_{k}),\hat{x}(x_{k})-x_{k}\rangle.$
(151)

Taking $c=\frac{-\mu_{\ell}+\rho}{2}$ in (146) and (147), we can rewrite (145)
as

$\displaystyle\langle$ $\displaystyle h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$
$\displaystyle\leq\frac{c}{2}\|\hat{x}(x_{k})-x_{k}\|^{2}+\frac{1}{c}b_{k}^{2}+\frac{1}{c}\|\mkern
1.5mu\overline{\mkern-1.5mu\nabla\mkern-1.5mu}\mkern
1.5mu_{x}F(x_{k},y_{k+1},\lambda_{k})-\nabla_{x}F(x_{k},\lambda_{k})\|^{2}-\frac{-\mu_{\ell}+\rho}{2}\|\hat{x}(x_{k})-x_{k}\|^{2}$
(152)
$\displaystyle\leq\frac{2}{-\mu_{\ell}+\rho}b_{k}^{2}+\frac{2L^{2}}{-\mu_{\ell}+\rho}\left(\max_{i\in[n]}\mathbb{E}\|y_{i}^{*}(x_{k})-y_{k+1}^{i}\|^{2}\right)-\frac{-\mu_{\ell}+\rho}{4}\|\hat{x}(x_{k})-x_{k}\|^{2}.$
(153)

Therefore, rewriting everything into (149), we have

$\displaystyle\Phi_{1/\rho}(x_{k+1})$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{\leq}}\Phi_{1/\rho}(x_{k})+\frac{\rho\alpha^{2}}{2}\|h_{x}^{k}\|^{2}+\rho\alpha\langle
h_{x}^{k},\hat{x}(x_{k})-x_{k}\rangle$ (154)
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\Phi_{1/\rho}(x_{k})+\frac{2\rho\alpha}{-\mu_{\ell}+\rho}b_{k}^{2}+\frac{2L^{2}\rho\alpha}{-\mu_{\ell}+\rho}\left(\max_{i\in[n]}\Delta_{y}^{k,i}\right)-\frac{\rho\alpha(-\mu_{\ell}+\rho)}{4}\|\hat{x}(x_{k})-x_{k}\|^{2}$
$\displaystyle\quad+\frac{\rho\alpha^{2}}{2}\left(\tilde{\sigma}_{f}^{2}+3b_{k}^{2}+3L^{2}\max_{i\in[n]}\mathbb{E}\Delta_{y}^{k,i}\right)$
(155)
$\displaystyle\stackrel{{\scriptstyle(3)}}{{=}}\Phi_{1/\rho}(x_{k})+\frac{2\rho\alpha}{-\mu_{\ell}+\rho}b_{k}^{2}+\left(\frac{2L^{2}\rho\alpha}{-\mu_{\ell}+\rho}+\frac{3L^{2}\rho\alpha^{2}}{2}\right)\left(\max_{i\in[n]}\Delta_{y}^{k,i}\right)$
$\displaystyle\quad+\frac{\rho\alpha(\mu_{\ell}-\rho)}{4}\|\hat{x}(x_{k})-x_{k}\|^{2}+\frac{\rho\alpha^{2}}{2}(\tilde{\sigma}_{f}^{2}+3b_{k}^{2}),$
(156)

where (1) is a copy of (142), (2) is from plugging in $\|h_{x}^{k}\|^{2}$ from
Lemma 2 and (156), and (3) is combining terms via algebra. Finally, summing up
from $k=0,1,\cdots,K-1$, we get the following bound:

$\displaystyle\frac{1}{K}\sum_{k=1}^{K}\|\hat{x}(x_{k})-x_{k}\|^{2}$
$\displaystyle\leq\frac{4}{-\mu_{\ell}+\rho}\left(\frac{\Phi_{1/\rho}(x_{0})}{\rho\alpha
K}+\frac{2}{-\mu_{\ell}+\rho}b_{0}^{2}+\left(\frac{2L^{2}}{-\mu_{\ell}+\rho}+\frac{3L^{2}\alpha}{2}\right)\left(\frac{1}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}\right)+\frac{\alpha}{2}\left(\tilde{\sigma}_{f}^{2}+3b_{0}^{2}\right)\right)$
(157)
$\displaystyle=\frac{4}{-\mu_{\ell}+\rho}\left(\underbrace{\frac{\Phi_{1/\rho}(x_{0})}{\rho\alpha
K}}_{K^{-2/5}}+\frac{2}{-\mu_{\ell}+\rho}b_{0}^{2}+\underbrace{\frac{2L^{2}}{-\mu_{\ell}+\rho}\left(\frac{1}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}\right)}_{\sqrt{n}K^{-2/5}}+\underbrace{\frac{\alpha}{2}\left(\tilde{\sigma}_{f}^{2}+3b_{0}^{2}+\frac{3L^{2}}{K}\sum_{k=1}^{K}\max_{i\in[n]}\Delta_{y}^{k,i}\right)}_{K^{-3/5}+\sqrt{n}K^{-1}}\right).$
(158)

###  A.10 Proof of Proposition 7.2

Proposition 7.2: Let $x^{\star}$ be the optimal solution to the bilevel
optimization objective in Section 3. Then, we have

$F_{n+1}(x^{\star})\leq\min_{x\in\mathcal{X}}\max_{p\in\Delta_{n}}p_{i}\hat{F}_{i}(x)+2a_{i}R_{m_{i}}^{i}(\mathcal{F}_{i})+a_{i}B_{\ell}\sqrt{\frac{\log(n/\delta)}{2m_{i}}}.$
(159)

Proof:  First, by definition, for all $x$, we have that

$\displaystyle F_{n+1}(x)$
$\displaystyle=\mathbb{E}_{(\mathcal{D}_{n+1,j}^{\text{train}},\mathcal{D}_{n+1,j}^{\text{test}})\sim\mathcal{D}_{n+1}}\left[\hat{f}_{i}\left(x,\operatorname*{arg\,min}_{y_{i}}\sum_{j^{\prime}=1}^{m_{i}}\hat{g}_{i}(x,y_{i},D_{i,j^{\prime}}^{\text{train}}),D_{i,j}^{\text{test}}\right)\right]$
(160)
$\displaystyle=\sum_{i=1}^{n}a_{i}\mathbb{E}_{(\mathcal{D}_{n+1,j}^{\text{train}},\mathcal{D}_{n+1,j}^{\text{test}})\sim\mathcal{D}_{i}}\left[\hat{f}_{i}\left(x,\operatorname*{arg\,min}_{y_{i}}\sum_{j^{\prime}=1}^{m_{i}}\hat{g}_{i}(x,y_{i},D_{i,j^{\prime}}^{\text{train}}),D_{i,j}^{\text{test}}\right)\right]$
(161) $\displaystyle=\sum_{i=1}^{n}a_{i}F_{i}(x).$ (162)

Therefore, plugging in $x^{\star}$ and using a union bound over $1\leq i\leq
n$, we get that with probability at least $1-n\delta^{\prime}$, we have

$F_{n+1}(x^{*})=\sum_{i=1}^{n}a_{i}F_{i}(x^{\star})\leq\sum_{i=1}^{n}a_{i}\hat{F}_{i}(x^{\star})+2a_{i}R_{m_{i}}^{i}(\mathcal{F})+a_{i}B\sqrt{\frac{\log
1/\delta^{\prime}}{2m_{i}}}.$ (163)

Letting $\delta=n\delta^{\prime}$ in (163), we have

$\displaystyle F_{n+1}(x^{*})$
$\displaystyle\leq\max_{p\in\Delta_{n}}p_{i}\hat{F}_{i}(x^{\star})+2a_{i}R_{m_{i}}^{i}(\mathcal{F})+a_{i}B_{\ell}\sqrt{\frac{\log(n/\delta)}{2m_{i}}}$
(164)
$\displaystyle=\min_{x\in\mathcal{X}}\max_{p\in\Delta_{n}}p_{i}\hat{F}_{i}(x)+2a_{i}R_{m_{i}}^{i}(\mathcal{F})+a_{i}B_{\ell}\sqrt{\frac{\log(n/\delta)}{2m_{i}}}.$
(165)

##  Appendix B Implementation details and Compute Resources

We perform our experiments in Python 3.7.10 and PyTorch 1.8.1 with Intel(R)
Core(TM) i5-8265U CPU @ 1.60GHz.

###  B.1 Sinusoid Regression Task

We consider the sinusoid regression experiment [Finn et al., 2017], a multi-
task representation learning problem where each task $\mathcal{T}_{i}$ is a
regression problem $y=t_{i}(x)=a_{i}\sin(x-\phi_{i})$. We uniformly sample the
amplitude $a_{i}\in[0.1,5]$, frequency and phase $\phi_{i}\in[0,\pi]$ for each
task. We use $n=3$ training tasks and $3$ testing tasks, with $2$ "easy tasks"
$(a_{i}\in[0.1,1.05])$ and one "hard tasks" $(a_{i}\in[4.95,5])$ for each set.
We use easy and hard tasks following the setup in [Collins et al., 2020].
During training, for each task $i$, the learner is given samples $(x,y)$,
$x\in[-5,5]$. The goal is to learn a function approximating $t_{i}$ as best as
possible in the mean squared error sense.

As described in Section 2, we use a neural network divided into two pieces,
i.e., an embedding network and a task-specific network. The embedding network
$f:\mathbb{R}\to\mathbb{R}^{10}$ consists of two hidden ReLU layers of size 80
and a final fully connected layer of size 10. Each task-specific network
$g_{i}:\mathbb{R}^{10}\to\mathbb{R},i\in[n]$ is a one-layer linear layer.
Therefore, the loss on an input $x\in\mathbb{R}$ and $y=t_{i}(x)$ for task $i$
is $(g_{i}(f(x))-y)^{2}$, and the true loss of the network with parameters
$f,g_{i}$ are $\ell_{i}(f;g_{i})=\mathbb{E}_{(x,y)}[(g_{i}(f(x))-y)^{2}]$. The
embedding network is as follows:

Input $(\mathbb{R})$ |  |  |
---|---|---|---
Linear FC Layer (output in $\mathbb{R}^{80}$) |  |  |
ReLU |  |  |
Linear FC Layer (output in $\mathbb{R}^{80}$) |  |  |
ReLU |  |  |
Linear FC Layer (output in $\mathbb{R}^{10}$) |  |  |

Training:  At each iteration, we first perform the inner loop optimization
step (meta-training) by sampling $10$ shots from each of the tasks in order to
update each of the task-specific network weights. We use just 1 inner loop
step. Pseudocode for the inner loop is shown below in PyTorch-style:

⬇

for task_id in task_list:

xs, ys = sample_batch(task_id, n_shots)

embedding = embedding_network(xs)

for _ in range(n_inner):

head_optimizers[task_id].zero_grad()

total_loss = get_loss(task_id, xs, ys)

total_loss.backward()

head_optimizers[task_id].step()

For each outer loop optimization, we run a meta-validation batch again
containing $10$ shots from each of the tasks. We then take an outer-loop step,
optimizing the embedding weights using the results of the meta-validation
batch. The meta-validation batch is sampled in the exact same way as the meta-
training batch shown above.

Regularization:  First, as in [Ji et al., 2020], we add weight regualarization
during inner loop training of the form
$\epsilon_{w}\sum_{w\in\mathcal{W}}\|w\|$, where $\mathcal{W}$ denotes the set
of weight parameters, where $\epsilon_{w}=0.01$. In PyTorch, this is expressed
as

⬇

l2_reg_constant = 0.01

for p in heads[task].parameters():

l2_reg += p.norm(2)

total_loss += l2_reg * l2_reg_constant

Next, for the inner $\lambda$ updates, we add a regularization term to the
overall loss,
$-\epsilon_{\lambda}\sum_{i=1}^{n}(\lambda_{i}-\frac{1}{n})^{2}$, where
$\epsilon_{\lambda}=3$, which pulls the $\lambda$s closer to uniform. In
PyTorch, the $\lambda$ update is expressed as

⬇

task_gradient = task_losses[i]

reg_gradient = -mu_lambda * (lambdas[task] - 1/n)

lambdas[task] += (task_update + reg_update) * gamma

Parameters:  For the Task-Robust version of the algorithm, we use
$\alpha=0.007,\beta=0.005,\gamma=0.003$. For the standard version of the
algorithm, we use $\alpha=0.007,\beta=0.011,\gamma=0.003$.

Loss curves:  To approximate the true loss for measurement purposes, we use
$100$ equally-spaced samples from $[-5,5]$. After each iteration, we
calculated the maximum loss among all the tasks. In 1, we show the minimum of
these maximum losses up until each epoch.

Results with more tasks:  Finally, we show another figure similar to Figure 1,
but with 20 training tasks and 20 test tasks. It can be observed that both the
task-robust training loss and the task-robust testing loss greatly outperform
their respective standard losses.

Figure 4: Comparison of standard training and robust training using 20 tasks

###  B.2 Nonlinear Representation Learning

We consider binary classification tasks generated from the FashionMNIST data
set where we select 8 “easy” tasks (lowest log loss $\sim 0.3$ from
independent training) and 2 “hard” tasks (lowest loss $\sim 0.45$ from
independent training). We learn a shared representation network that maps the
784 dimensional (vectorized 28$\times$28 images) to a 100 dimensional space.
Each tasks then learns a binary classifier on top of this representation. The
task specific objective $g_{i}$ for task $i$ corresponds to the cross-entropy
loss on the training set, while the upper level objective $f_{i}$ corresponds
to the loss of the $y_{i}^{\star}(x)$ with the learned representation $x$ on a
validation set. We also maintain a heldout test set which we use to evaluate
the generalization of the learned representation and per-task models.

For our data, we had $x\in\mathbb{R}^{784\times 100}$ and
$y\in\mathbb{R}^{100\times 2}$. We used step sizes $\alpha=0.01,\beta=0.01$,
and $\gamma=0.3$, using 4 inner steps per outer iteration. We used batch sizes
of 8 and 128 to compute $g_{i}$ for each inner step and $f_{i}$ for each outer
iteration, respectively. In addition, we included $\ell_{2}$-regularization of
$y$ with regularization penalty 0.0005. We used a learning rate scheduler
(ReduceLROnPlateau), invoked every 100 outer iterations, with patience of 10.
The results are generated by aggregation over runs with 10 different seeds