In this work, we consider the problem of minimizing the sum of Moreau envelopes of given functions, which has previously appeared in the context of meta-learning and personalized federated learning. In contrast to the existing theory that requires running subsolvers until a certain precision is reached, we only assume that a finite number of gradient steps is taken at each iteration. As a special case, our theory allows us to show the convergence of First-Order Model-Agnostic Meta-Learning (FO-MAML) to the vicinity of a solution of Moreau objective. We also study a more general family of first-order algorithms that can be viewed as a generalization of FO-MAML. Our main theoretical achievement is a theoretical improvement upon the inexact SGD framework. In particular, our perturbed-iterate analysis allows for tighter guarantees that improve the dependency on the problem's conditioning. In contrast to the related work on meta-learning, ours does not require any assumptions on the Hessian smoothness, and can leverage smoothness and convexity of the reformulation based on Moreau envelopes. Furthermore, to fill the gaps in the comparison of FO-MAML to the Implicit MAML (iMAML), we show that the objective of iMAML is neither smooth nor convex, implying that it has no convergence guarantees based on the existing theory.

§ INTRODUCTION

Efficient optimization methods for empirical risk minimization have helped the breakthroughs in many areas of machine learning such as computer vision [19] and speech recognition [14]. More recently, elaborate training algorithms have enabled fast progress in the area of meta-learning, also known as learning to learn [29]. At its core lies the idea that one can find a model capable of retraining for a new task with just a few data samples from the task. Algorithmically, this corresponds to solving a bilevel optimization problem [11], where the inner problem corresponds to a single task, and the outer problem is that of minimizing the post-training error on a wide range of tasks.

The success of Model-Agnostic Meta-Learning (MAML) and its first-order version (FO-MAML) [9] in meta-learning applications has propelled the development of new gradient-based meta-learning methods. However, most new algorithms effectively lead to new formulations of meta-learning. For instance, iMAML [27] and proximal meta-learning [35] define two MAML-like objectives with implicit gradients, while Reptile [24] was proposed without defining any objective at all. These dissimilarities cause fragmentation of the field and make it particularly hard to have a clear comparison of meta-learning theory. Nonetheless, having a good theory helps to compare algorithms as well as identify and fix their limitations.

Unfortunately, for most of the existing methods, the theory is either incomplete as is the case with iMAML or even completely missing. In this work, we set out to at least partially mitigate this issue by proposing a new analysis for minimization of Moreau envelopes. We show that a general family of algorithms with multiple gradient steps is stable on this objective and, as a special case, we obtain results even for FO-MAML. Previously, FO-MAML was viewed as a heuristic to approximate MAML [8], but our approach reveals that FO-MAML can be regarded as an algorithm for a the sum of Moreau envelopes. While both perspectives show only approximate convergence, the main justification for the sum of Moreau envelopes is that requires unprecedentedly mild assumptions. In addition, the Moreau formulation of meta-learning does not require Hessian information and is easily implementable by any first-order optimizer, which 35 showed to give good empirical performance.

§.§ Related work

A summary of related work and conceptual differences to our approach. We mark as “N/A” unknown properties that have not been established in prior literature or our work. We say that $F_i$ “Preserves convexity” if for convex $f_i$, $F_i$ is convex as well, which implies that $F_i$ has no extra local minima or saddle points. We say that $F_i$ “Preserves smoothness” if its gradients are Lipschitz whenever the gradients of $f_i$ are, which corresponds to more stable gradients. We refer to [8] for the claims regarding nonconvexity and nonsmoothness of the MAML objective.

Algorithm     $\meta_i$: meta-loss of task $i$     Hessian-
free     Arbitrary
number of steps     No matrix
inversion     Preserves
convexity     Preserves
smoothness     Reference

MAML     $f_i(x-\alpha \nabla f_i(x))$                          [9]

MAML     $f_i(GD(f_i, x))$(1)                          [9]
[16]

iMAML(2)     $f_i(z_i(x))$, where
$z_i(x)=x-\alpha \nabla f_i(z_i(x))$
(Theorem <ref>)
(Theorem <ref>)     [27]

Reptile     N/A(3)                  N/A     N/A     [24]

(original)     $f_i(x-\alpha \nabla f_i(x))$                         [9]

Meta-MinibatchProx     $\min\limits_{x_i}\{f_i(x_i) + \frac{1}{2\alpha}\|x_i-x\|^2 \}$         (4)                   [35]

(extended FO-MAML)     $\min\limits_{x_i}\{f_i(x_i) + \frac{1}{2\alpha}\|x_i-x\|^2 \}$                           This work

(1) Multi-step MAML runs an inner loop with gradient descent applied to task loss $f_i$, so the objective of multi-step MAML is $F_i(x)=f_i(x_s(x))$, where $x_0=x$ and $x_{j+1}=x_j - \alpha \nabla f_i(x_j)$ for $j=0,\dotsc, s-1$.

(2) To the best of our knowledge, iMAML is not guaranteed to work; 27 studied only the approximation error for gradient computation, see the discussion in our special section on iMAML.

(3) Reptile was proposed as an algorithm on its own, without providing any optimization problem. This makes it hard to say how it affects smoothness and convexity. 4 and 17 studied convergence of Reptile on the average loss over the produced iterates, i.e., $F_i(x)=\frac{1}{m}\sum_{j=0}^s f_i(x_j)$, where $x_0=x$ and $x_{j+1}=x_j - \alpha \nabla f_i(x_j)$ for $j=0,\dotsc, s-1$. Analogously to the loss of MAML, this objective seems nonconvex and nonsmooth.

(4) 35 assumed that the subproblems are solved to precision $\varepsilon$, i.e., $x_i$ is found such that $\|\nabla f_i(x_i) + \frac{1}{\alpha}(x_i - x)\|\le \varepsilon$ with an absolute constant $\varepsilon$.

MAML [9] has attracted a lot of attention due to its success in practice. Many improvements have been proposed for MAML, for instance, [36] suggested augmenting each group of tasks with its own global variable, and [2] proposed MAML++ that uses intermediate task losses with weights to improve the stability of MAML. [27] proposed iMAML that makes the objective optimizer-independent by relying on implicit gradients. 35 used a similar implicit objective to that of iMAML with an additional regularization term that, unlike iMAML, does not require inverting matrices. Reptile [24] is an even simpler method that merely runs gradient descent on each sampled task. Based on generalization guarantees, [36] also provided a trade-off between the optimization and statistical errors for a multi-step variant MAML, which shows that it may not improve significantly from increasing the number of gradient steps in the inner loop. We refer to [15] for a recent survey of the literature on meta-learning with neural networks.

On the theoretical side, the most relevant works to ours is that of [35], whose main limitation is that it requires a high-precision solution of the inner problem in Moreau envelope at each iteration. Another relevant work that studied convergence of MAML and FO-MAML on the standard MAML objective is by [8], but they do not provide any guarantees for the sum of Moreau envelopes and their assumptions are more stringent. 8 also study a Hessian-free variant of MAML, but its convergence guarantees still require posing assumptions on the Hessian Lipschitzness and variance.

Some works treat meta-learning as a special case of compositional optimization [31] or bilevel programming [11] and develop theory for the more general problem. Unfortunately, both approaches lead to worse dependence on the conditioning numbers of both inner and outer objective, and provide very pessimistic guarantees. Bilevel programming, even more importantly, requires computation of certain inverse matrices, which is prohibitive in large dimensions. One could also view minimization-based formulations of meta-learning as instances of empirical risk minimization, for which FO-MAML can be seen as instance of inexact (biased) SGD. For example, [1] analyzed SGD with deterministic bias and some of our proofs are inspired by theirs, except in our problem the bias is not deterministic. We will discuss the limitations of their approach in the section on inexact SGD.

Several works have also addressed meta-learning from the statistical perspective, for instance, 33 proposed a Bayesian variant of MAML, and 10 analyzed convergence of MAML in online learning. Another example is the work of 18 who studied the setting of linear regression with task-dependent solutions that are sampled from same normal distribution. These directions are orthogonal to ours, as we want to study the optimization properties of meta-learning.

§ BACKGROUND AND MATHEMATICAL FORMULATION

Before we introduce the considered formulation of meta-learning, let us provide the problem background and define all notions. As the notation in meta-learning varies between papers, we correspond our notation to that of other works in the next subsection.

§.§ Notation

We assume that training is performed over $n$ tasks with task losses $f_1,\dotsc, f_n$ and we will introduce implicit and proximal meta-losses $\{F_i\}$ in the next section.
We denote by $x$ the vector of parameters that we aim to train, which is often called model, meta-model or meta-parameters in the meta-learning literature, and outer variable in the bilevel literature. Similarly, given task $i$, we denote by $z_i$ the task-specific parameters that are also called as ground model, base-model, or inner variable. We will use letters $\alpha, \beta, \gamma$ to denote scalar hyper-parameters such as stepsize or regularization coefficient.

Given a function $\varphi(\cdot)$, we call the following function its Moreau envelope:

\begin{align*}
\Phi(x) = \min_{z\in\R^d}\left\{\varphi(x) + \frac{1}{2\alpha}\|z-x\|^2 \right\},
\end{align*}

where $\alpha>0$ is some parameter. Given the Moreau envelope $F_i$ of a task loss $f_i$, we denote by $z_i(x)$ the solution to the inner objective of $F_i$, i.e., $z_i(x)\eqdef \argmin_{z\in\R^d} \left\{f_i(z) + \frac{1}{2\alpha}\|z - x\|^2\right\}$.

Finally, let us introduce some standard function properties that are commonly used in the optimization literature [23].

We say that a function $\varphi(\cdot)$ is $L$-smooth if its gradient is $L$-Lipschitz, i.e., for any $x, y\in\R^d$,

\begin{align*}
\|\nabla \varphi(x) - \nabla \varphi(y)\|\le L\|x-y\|.
\end{align*}

Given a function $\varphi(\cdot)$, we call it $\mu$-strongly convex if it satisfies for any $x, y\in\R^d$,

\begin{align*}
\varphi(y)\ge \varphi(x) + \<\nabla \varphi(x), y-x> + \frac{\mu}{2}\|y-x\|^2.
\end{align*}

If the property above holds with $\mu=0$, we call $\varphi$ to be convex. If the property does not hold even with $\mu=0$, we say that $\varphi$ is nonconvex.

§.§ MAML objective

Assume that we are given $n$ tasks, and that the performance on task $i$ is evaluated according to some loss function $f_i(x)$. MAML has been proposed as an algorithm for solving the following objective:

\begin{align}
\min_{x\in\R^d} \frac{1}{n}\sumin f_i(x - \alpha \nabla f_i(x)), \label{eq:bad_problem}
\end{align}

where $\alpha>0$ is a stepsize. Ignoring for simplicity minibatching, MAML update computes the gradient of a task meta-loss $\varphi_i(x)=f_i(x - \alpha \nabla f_i(x))$ through backpropagation and can be explicitly written as

\begin{align*}
&= x^k - \beta \left(\mathbf{I}-\alpha\nabla^2 f_i(x^k) \right)\nabla f_i (x^k-\alpha\nabla f_i(x^k)),     \tag{MAML update}
\end{align*}

where $\beta>0$ is a stepsize, $i$ is sampled uniformly from $\{1,\dotsc, n\}$ and $\mathbf{I}\in\R^{d\times d}$ is the identity matrix. Sometimes, MAML update evaluates the gradient of $\varphi_i$ using an additional data sample, but 3 recently showed that this is often unnecessary, and we, thus, skip it.

Unfortunately, objective (<ref>) might be nonsmooth and nonconvex even if the task losses $\{f_i\}$ are convex and smooth [8]. Moreover, if we generalize this objective for more than one gradient step inside $f_i(\cdot)$, its smoothness properties deteriorate further, which complicates the development and analysis of multistep methods.

§.§ iMAML objective

To avoid differentiating through a graph, 27 proposed an alternative objective to (<ref>) that replaces the gradient step inside each function with an implicit gradient step. In particular, if we define $z_i(x)\eqdef \argmin_{z\in\R^d} \left\{f_i(z) + \frac{1}{2\alpha}\|z - x\|^2\right\}$, then the objective of iMAML is

\[
\min_{x\in\R^d} \avein f_i\left(x-\alpha \nabla f_i(z_i(x))\right).
\]

The idea of iMAML is to optimize this objective during training so that at inference, given a new function $f_{n+1}$ and solution $x_{\mathrm{iMAML}}$ of the problem above, one can find an approximate solution to $\min_{z\in\R^d} \left\{f_{n+1}(z) + \frac{1}{2\alpha}\|z - x_{\mathrm{iMAML}}\|^2\right\}$ and use it as a new model for task $f_{n+1}$.

27 proved, under some mild assumptions, that one can efficiently obtain an estimate of the gradient of $\varphi_i(x)\eqdef f_i\left(x-\alpha \nabla f_i(z_i(x))\right)$ with access only to gradients and Hessian-vector products of $f_i$, which rely on standard backpropagation operations. In particular, 27 showed that

\[
\nabla \varphi_i(x)
= \left(\mathbf{I} + \alpha \nabla^2 f_i(z(x)) \right)^{-1} \nabla f_i(z(x)),
\]

where $\mathbf{I}$ is the identity matrix, and they proposed to run the conjugate gradient method to find $\nabla \varphi_i(x)$.
However, it is not shown in [27] if the objective of iMAML is solvable and what properties it has. Moreover, we are not aware of any result that would show when the problem is convex or smooth. Since SGD is not guaranteed to work unless the objective satisfies at least some properties [34], nothing is known about convergence of SGD when applied to the iMAML objective.

As a sign that the problem is rather ill-designed, we present the following theorem that gives a negative example on the problem's convexity.

There exists a convex function $f$ with Lipschitz gradient and Lipschitz Hessian such that the iMAML meta-objective $\varphi(x)\eqdef f(z(x))$ is nonconvex, where $z(x)=x - \alpha \nabla f(z(x))$.

Similarly, we also show that the objective of iMAML may be harder to solve due to its worse smoothness properties as given by the next theorem.

There exists a convex function $f$ with Lipschitz gradient and Lipschitz Hessian such that the iMAML meta-objective $\varphi(x)\eqdef f(z(x))$ is nonsmooth for any $\alpha>0$, where $z(x)=x - \alpha \nabla f(z(x))$.

FO-MAML: First-Order MAML

[1]
Input: $x^0$, $\alpha, \beta > 0$

Sample a subset of tasks $T_k$
each sampled task $i$ in $T_k$
$z_i^k = x^k - \alpha \nabla f_{i}(x^k)$
$x^{k+1} = x^k - \beta \frac{1}{|T_k|}\sum_{i\in T_k} \nabla f_i(z_i^k)$

§.§ Our main objective: Moreau envelopes

In this work we consider the following formulation of meta-learning

\begin{align}
&\min_{x\in\R^d}  \meta(x) \eqdef \frac{1}{n}\sum_{i=1}^n \meta_i(x),\quad \label{eq:new_pb} \\
\meta_i(x)\eqdef \min_{z\in\R^d} \left\{f_i(z) + \frac{1}{2\alpha}\|z - x\|^2\right\}, \notag
\end{align}

and $\alpha>0$ is a parameter controlling the level of adaptation to the problem. In other words, we seek to find a parameter vector $x$ such that somewhere close to $x$ there exists a vector $z_i$ that verifies that $f_i(z)$ is sufficiently small. This formulation of meta-learning was first introduced by 35 and it has been used by 13 and 32 to study personalization in federated learning.

Throughout the paper we use the following variables for minimizers of meta-problems $\meta_i$:

\begin{align}
z_i(x) \eqdef \argmin_{z\in\R^d} \left\{f_i(z) + \frac{1}{2\alpha}\|z - x\|^2\right\},  i=1,\dotsc, n. \label{eq:z_i}
\end{align}

One can notice that if $\alpha\to 0$, then $\meta_i(x)\approx f_i(x)$, and Problem (<ref>) reduces to the well-known empirical risk minimization:

\begin{align*}
\min_{x\in\R^d} f(x)
\eqdef \frac{1}{n}\sum_{i=1}^n f_i(x).
\end{align*}

If, on the other hand, $\alpha\to +\infty$, the minimization problem in (<ref>) becomes essentially independent of $x$ and it holds $z_i(x)\approx \argmin_{z\in\R^d} f_i(z)$. Thus, one has to treat the parameter $\alpha$ as part of the objective that controls the similarity between the task-specific parameters.

We denote the solution to Problem (<ref>) as

\begin{align}
x^* \eqdef \arg\min_{x\in\R^d} \meta(x).
\end{align}

One can notice that $F(x)$ and $x^*$ depend on $\alpha$. For notational simplicity, we keep $\alpha$ constant throughout the paper and do not explicitly write the dependence of $x^*, F, F_1, z_1, \dotsc, F_n, z_n$ on $\alpha$.

§.§ Formulation properties

We will also use the following quantity to express the difficulty of Problem (<ref>):

\begin{align}
\sigma_*^2\eqdef \frac{1}{n}\sumin \|\nabla \meta_i(\opt)\|^2.  \label{eq:def_sigma}
\end{align}

Because $\nabla \meta(\opt)=0$ by first-order optimality of $\opt$, $\sigma_*^2$ serves as a measure of gradient variance at the optimum. Note that $\sigma_*$ is always finite because it is defined on a single point, in contrast to the maximum gradient variance over all space, which might be infinite.

Now let's discuss properties of our formulation <ref>. Firstly, we state a standard result from [5].

Let $\meta_i$ be defined as in <ref> and $z_i(x)$ be defined as in <ref>.
If $f_i$ is convex, proper and closed, then $\meta_i$ is differentiable and $\frac{1}{\alpha}$-smooth:

\begin{align}
&\nabla \meta_i(x)
= \frac{1}{\alpha}(x-z_i(x)) = \nabla f_i(z_i(x)), \label{eq:implicit} \\
&\|\nabla \meta_i(x) - \nabla \meta_i(y)\|\le \frac{1}{\alpha}\|x-y\|.
\end{align}

The results above only hold for convex functions, while in meta-learning, the tasks are often defined by training a neural network, whose landscape is nonconvex. To address such applications, we also refine <Ref> in the lemma bellow, which also improves the smoothness constant in the convex case. This result is similar to Lemma 2.5 of [6], except their guarantee is a bit weaker because they consider more general assumptions.

Let function $f_i$ be $L$-smooth.

* If $f_i$ is nonconvex and $\alpha<\frac{1}{L}$, then $\meta_i$ is $\frac{L}{1-\alpha L}$-smooth. If $\alpha \le \frac{1}{2L}$, then $\meta_i$ is $2L$-smooth.

* If $f_i$ is convex, then $\meta_i$ is $\frac{L}{1+\alpha L}$-smooth. Moreover, for any $\alpha$, it is $L$-smooth.

* If $f_i$ is $\mu$-strongly convex, then $\meta_i$ is $\frac{\mu}{1+\alpha\mu}$-strongly convex. If $\alpha \le \frac{1}{\mu}$, then $\meta_i$ is $\frac{\mu}{2}$-strongly convex.

Whenever $\meta_i$ is smooth, its gradient is given as in equation (<ref>), i.e., $\nabla \meta_i(x) = \nabla f_i(z_i(x))$.

The takeaway message of <Ref> is that the optimization properties of $\meta_i$ are always at least as good as those of $f_i$ (up to constant factors). Furthermore, the conditioning, i.e., the ratio of smoothness to strong convexity, of $\meta_i$ is upper bounded, up to a constant factor, by that of $f_i$. And even if $f_i$ is convex but nonsmooth ($L\to+\infty)$, $F_i$ is still smooth with constant $\frac{1}{\alpha}$.

Finally, note that computing the exact gradient of $F_i$ requires solving its inner problem as per equation (<ref>). Even if the gradient of task $\nabla f_i(x)$ is easy to compute, we still cannot obtain $ \nabla \meta_i(x)$ through standard differentiation or backpropagation. However, one can approximate $\nabla \meta_i(x)$ in various ways, as we will discuss later.

§ CAN WE ANALYZE FO-MAML AS INEXACT SGD?

As we mentioned before
, the prior literature has viewed FO-MAML as an inexact version of MAML for problem (<ref>). If, instead, we are interested in problem (<ref>), one could still try to take the same perspective of inexact SGD and see what convergence guarantees it gives for (<ref>). The goal of this section, thus, is to refine the existing theory of inexact SGD to make it applicable to FO-MAML. We will see, however, that such approach if fundamentally limited and we will present a better alternative analysis in a future section.

§.§ Why existing theory is not applicable

Let us start with a simple lemma for FO-MAML that shows why it approximates SGD for objective (<ref>).

Let task losses $f_i$ be $L$–smooth and $\alpha>0$. Given $i$ and $x\in\R^d$, we define recursively $z_{i,0} \eqdef x$ and $z_{i,j+1} \eqdef {\color{blue}x} - \alpha \nabla f_i({\color{mydarkred}z_{i,j}})$. Then, it holds for any $s\ge 0$

\begin{align*}
\left\| \nabla f_i(z_{i,s}) - \nabla \meta_i(x) \right\| \le (\alpha L)^{s+1} \|\nabla \meta_i(x)\|.
\end{align*}

In particular, the iterates of FO-MAML (<Ref>) satisfy for any $k$

\begin{align*}
\left\| \nabla f_i(z_i^k) - \nabla \meta_i(x^k) \right\| \le (\alpha L)^2 \|\nabla \meta_i(x^k)\|.
\end{align*}

<Ref> shows that FO-MAML approximates SGD step with error proportional to the stochastic gradient norm. Therefore, we can write

\begin{align*}
\nabla f_i(z_i^k)
= \nabla F(x^k) + \underbrace{\nabla F_i(x^k) - \nabla F(x^k)}_{\eqdef \xi_i^k\ (\mathrm{noise})} + \underbrace{b_i^k}_{\mathrm{bias}},
\end{align*}

where it holds $\mathbb{E}[\xi_i^k]=0$, and $b_i^k$ is a bias vector that also depends on $i$ but does not have zero mean. The best known guarantees for inexact SGD are provided by 1, but they are, unfortunately, not applicable because their proofs use independence of $\xi_i^k$ and $b_i^k$. The analysis of 35 is not applicable either because their inexactness assumption requires the error to be smaller than a predefined constant $\varepsilon$, while the error in <Ref> can be unbounded. To resolve these issues, we provide a refined analysis in the next subsection.

§.§ A new result for inexact SGD

For strongly convex objectives, we give the following result by modifying the analysis of 1.

Let task losses $f_1,\dotsc, f_n$ be $L$-smooth and $\mu$-strongly convex. If $|T_k|=\tau$ for all $k$, $\outers \leq \frac 1 {20L}$ and $\alpha\le \frac{1}{4\sqrt{\kappa}L}$, where $\kappa\eqdef \frac{L}{\mu}$, then for the iterates $x^1, x^2\dots$ of <Ref>, it holds

\begin{align*}
\ec{\|x^k-\opt\|^2}
&\le \left(1 - \frac{\outers\mu}{4}\right)^k\|x^0-\opt\|^2
+ \frac{16}{\mu} \left( \frac {2\alpha^2 L^2} {\mu} + \frac \outers \tau + \outers  \right) \sigma_*^2.
\end{align*}

Let us try to compare this result to that of vanilla SGD as studied by 12. Since the first term decreases exponentially, it requires us $\cO\left(\frac{1}{\beta\mu}\log\frac{1}{\varepsilon} \right)$ iterations to make it smaller than $\varepsilon$. The second term, on the other hand, only decreases if we decrease $\alpha$ and $\beta$. Decreasing $\beta$ corresponds to using decreasing stepsizes in SGD, which is fine, but $\alpha$ is a parameter that defines the objective, so in most cases, we do not want to decrease it. Moreover, the assumptions of <Ref> require $\alpha$ to be smaller than $\frac{1}{\sqrt{\kappa}L}$, which seems quite restrictive. This is the main limitation of this result as it shows that FO-MAML as given in <Ref> may not converge to the problem solution.

FO-MuML: First-Order Multistep Meta-Learning (general formulation)

[1]
Input:$x^0$, $\beta>0$, accuracy $\delta\geq0$ or $\varepsilon\ge 0$.
Sample a subset of tasks $T_k$
each sampled task $i$ in $T_k$
Find $z_i^k$ s.t. $\norm{\frac 1 \inners \left(x^k -z_i^k \right) - \nabla \meta_i(x^k)} \leq \delta \norm{\nabla \meta_i(x^k)}$
$x^{k+1} = x^k - \beta\frac{1}{|T_k|}\sum_{i\in T_k} \nabla f_i(z_i^k)$

To fix the nonconvergence of FO-MAML, let us turn our attention to <Ref>, which may perform multiple first-order steps.

Let task losses $f_1,\dotsc, f_n$ be $L$-smooth and $\mu$-strongly convex. If $|T_k|=\tau$ for all $k$, $\alpha\le \frac{1}{L}, \outers \leq \frac 1 {20L}$, and $\delta \leq \frac 1 {4 \sqrt{ \kappa}}$, where $\kappa\eqdef \frac{L}{\mu}$, then the iterates of <Ref> satisfy

\begin{align*}
\ec{\|x^k-\opt\|^2}
&\le \left(1 - \frac{\outers\mu}{4}\right)^k\|x^0-\opt\|^2
+ \frac{16}{\mu} \left( \frac {2\delta^2} {\mu} + \frac \outers \tau + \outers \delta^2 \right) \sigma_*^2.
\end{align*}

The result of <Ref> is better than that of <Ref> since it only requires the inexactness parameter $\delta$ to go to 0 rather than $\alpha$, so we can solve the meta-learning problem (<ref>) for any $\alpha\le \frac{1}{L}$. The rate itself, however, is not optimal, as we show in the next section with a more elaborate approach.

§ IMPROVED THEORY

FO-MuML (example of implementation)

[1]
Input: $x^0$, number of steps $s$, $\alpha > 0$, $\beta>0$
Sample a subset of tasks $T_k$
each sampled task $i$ in $T_k$
$z_{i, 0}^k = x^k$
$l = 0, \dotsc, s-1$
$z^k_{i, l+1} = {\color{blue} x^k} - \alpha \nabla f_{i}({\color{red}z_{i, l}^k})$
$z_i^k = z_{i,s}^k$
$x^{k+1} = x^k - \beta\frac{1}{|T_k|}\sum_{i\in T_k} \nabla f_i(z_i^k)$

In this section, we provide improved convergence theory of FO-MAML and FO-MuML based on a sequence of virtual iterates that appear only in the analysis. Surprisingly, even though the sequence never appears in the algorithm, it allows us to obtain tighter convergence bounds.

§.§ Perturbed iterate is better than inexact gradient

Before we introduce the sequence, let us make some observations from prior literature on inexact and biased variants of SGD. For instance, the literature on asynchronous optimization has established that getting gradient at a wrong point does not significantly worsen its rate of convergence [21]. A similar analysis with additional virtual sequence was used in the so-called error-feedback for compression [30], where the goal of the sequence is to follow the path of exact gradients even if compressed gradients are used by the algorithm itself. Motivated by these observations, we set out to find a virtual sequence that could help us analyze FO-MAML.

§.§ On what vector do we evaluate the gradients?

The main difficulty that we face is that we never get access to the gradients of $\{F_i\}$ and have to use the gradients of $\{f_i\}$. However, we would still like to write

\begin{align*}
=x^k - \frac \inners \tau \sum_{i \in T_k} \nabla f_i(z_i^k)
= x^k - \frac \inners \tau \sum_{i \in T_k} \nabla \meta_i(y_i^k)
\end{align*}

for some point $y_i^k$. If this is possible, using point $y_i^k$ would allow us to avoid working with functions $f_i$ in some of our recursion.

Why exactly would this sequence help? As mentioned before, FO-MAML is a biased method, so we cannot evaluate expectation of $\E{\nabla f_i(z_i^k)}$. However, if we had access to $\nabla F_i(x^k)$, its expectation would be exactly $\nabla F(x^k)$. This suggests that if we find $y_i^k$ that satisfies
$\nabla \meta_i(y_i^k) \approx \nabla \meta_i(x^k)$, then

\[
x^{k+1} = x^k - \frac \inners \tau \sum_{i \in T_k} \nabla \meta_i(y_i^k) \approx x^k - \frac \inners \tau \sum_{i \in T_k} \nabla \meta_i(x^k),
\]

which would allow us to put the bias inside the gradient.

Fortunately, objective (<ref>) allows us to find such point easily. In particular, for Moreau Envelopes, the following proposition holds.

For any points $z, y \in \R^d$ it holds $y= z + \inners \nabla f_i(z)$ if and only if $z = y - \inners \nabla \meta_i(y)$. Therefore, given $z$, we can define $y=z+\alpha \nabla f_i(z)$ and obtain $\nabla f_i(z)=\nabla F_i(y)$.

The result follows immediately from the last statement of Lemma <ref>.

The second part of <Ref> is exactly what we need. Indeed, we can choose $y_i^k \eqdef z_{i}^k + \inners \nabla  f_i(z_{i}^k)$ so that $z_{i}^k = y_i^k - \inners \nabla \meta_i(y_i^k)$ and $\nabla f_i(z_i^k)=\nabla F_i(y_i^k)$. As we have explained, this can help us to tackle the bias of FO-MAML.

§.§ Main results

We have established the existence of variables $y_i^k$ such that $\nabla f_i(z_i^k)=\nabla F_i(y_i^k)$. This allows us to write

\begin{align*}
\nabla f_i(z_i^k)
= \nabla F_i(y_i^k)
= \nabla F(x^k) + \underbrace{\nabla F_i(x^k) - \nabla F(x^k)}_{\mathrm{noise}}
+ \underbrace{\nabla F_i(y_i^k) - \nabla F_i(x^k)}_{\textrm{reduced bias}}&.
\end{align*}

As the next theorem shows, we can use this to obtain convergence guarantee to a neighborhood even with a small number of steps in the inner loop.

Consider the iterates of <Ref> (with general $\delta$) or <Ref> (for which $\delta=\alpha L$).
Let task losses be $L$–smooth and $\mu$–strongly convex and let objective parameter satisfy $\inners \leq \frac {1}{\sqrt 6 L}$. Choose stepsize $ \beta \leq \frac \tau {4 L}$, where $\tau = |T_k|$ is the batch size. Then we have

\begin{align*}
\E{\norm{x^k-x^*}^2} &\leq \left(1 - \frac {\beta \mu}{12}  \right)^k \norm{x^0 - x^*}^2
+ \frac { 6\left( \frac \beta \tau + 3 \delta^2 \inners^2 L \right) \varopt} {\mu}.
\end{align*}

Similarly to <Ref>, the theorem above guarantees convergence to a neighborhood only. However, the radius of convergence is now $\cO\left(\frac{\frac{\beta}{\tau} + \alpha^2L}{\mu} \right)$ in contrast to $\cO\left(\frac{\beta + \kappa\alpha^2L}{\mu} \right)$. If the first term is dominating, then it implies an improvement proportional to the batch size $\tau$. If, in contrast, the second term is larger, then the improvement is even more significant and the guarantee is $\cO(\kappa)$ times better, which is often a very large constant.

The proof technique for this theorem also uses recent advances on the analysis of biased SGD methods by 22. In particular, we show that the three-point identity (provided in the Appendix) is useful for getting a tighter recursion.

Next, we extend this result to the nonconvex convergence as given under the following assumption on bounded variance.

We assume that the variance of meta-loss gradients is uniformly bounded by some $\sigma^2$, i.e.,

\begin{align}
\E{\|\nabla F_i(x) - \nabla F(x)\|^2}
\le \sigma^2. \label{eq:bounded_var}
\end{align}

The new assumption on bounded variance is different from the one we used previously of variance being finite at the optimum, which was given in equation (<ref>). At the same time, it is very common in literature on stochastic optimization when studying convergence on nonconvex functions.

Let <Ref> hold, functions $f_1,\dotsc, f_n$ be $L$–smooth and $F$ be lower bounded by $F^*>-\infty$. Assume $\alpha\le \frac{1}{4L}, \beta\le \frac{1}{16L}$. If we consider the iterates of <Ref> (with $\delta=\alpha L$) or <Ref> (with general $\delta$), then

\begin{align*}
\min_{t\le k}\E{\|\nabla F(x^t)\|^2}
&\le \frac{4}{\beta k}\E{F(x^0)-F^*} + 4(\alpha L)^2\delta^2 \sigma^2
+ 32 \beta(\alpha L)^2 \left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right) \sigma^2.
\end{align*}

Notice that this convergence is also only until some neighborhood of first-order stationarity, since the second term does not decrease with $k$. This size of the upper bound depends on the product $\cO((\alpha L)^2 \delta^2)$, so to obtain better convergence one can simply increase approximation accuracy to make $\delta$ smaller. However, the standard FO-MAML corresponds to $\delta=\alpha L$, so its convergence guarantees directly depend on the problem parameter $\alpha$.

For Algorithm <ref>, we have $\delta=\cO((\alpha L)^s)$ as per Lemma <ref>, and we recover convergence guarantee up to a neighborhood of size $\cO((\alpha L)^2\delta^2)=\cO((\alpha L)^{2s+2})$. Therefore, to make this smaller than some given target accuracy $\varepsilon>0$, we need at most $s=\cO(\log\frac{1}{\varepsilon})$ inner-loop iterations. If we can plug-in $s=1$, we also get that FO-MAML converges to a neighborhood of size $\cO((\alpha L)^4)$.

Our <Ref> is very similar to the one obtained by 8, except their convergence neighborhood depends on $\alpha$ as $\cO(\alpha^2)$, whereas ours is of size $\cO(\alpha^4)$, which goes to 0 much faster when $\alpha\to 0$. Moreover, in contrast to their theory, ours does not require any assumptions on the Hessian smoothness. Note, in addition, that the main difference comes from the kind of objectives that we study, as 8 considered minimization of problems not involving Moreau envelopes.

§ CONCLUSION

In this paper, we presented a new analysis of first-order meta-learning algorithms for minimization of Moreau envelopes. Our theory covers both nonconvex and strongly convex smooth losses and guarantees convergence of the family of methods covered by Algorithm <ref>. As a special case, all convergence bounds apply to Algorithm <ref> with an arbitrary number of inner-loop steps. Compared to other results available in the literature, ours are more general as they hold with an arbitrary number of inner steps and do not require Hessian smoothness. The main theoretical difficulty we faced was the limitation of the inexact SGD framework, which we overcame by presenting a refined analysis using virtual iterates. As a minor contribution, we also pointed out that standard algorithms, such as SGD, are not immediately guaranteed to work on the iMAML objective, which might be nonconvex and nonsmooth even for convex and smooth losses. To show this, we presented examples of losses whose convexity and smoothness cease when the iMAML objective is constructed.

[1]
Ajalloeian, A. and Stich, S. U. (2020).
Analysis of SGD with biased gradient estimators.
arXiv preprint arXiv:2008.00051.

[2]
Antoniou, A., Edwards, H., and Storkey, A. J. (2018).
How to train your MAML.
In International Conference on Learning Representations.

[3]
Bai, Y., Chen, M., Zhou, P., Zhao, T., Lee, J., Kakade, S., Wang, H., and
Xiong, C. (2021).
How important is the train-validation split in meta-learning?
In International Conference on Machine Learning, pages
543–553. PMLR.

[4]
Balcan, M.-F., Khodak, M., and Talwalkar, A. (2019).
Provable guarantees for gradient-based meta-learning.
In International Conference on Machine Learning, pages
424–433. PMLR.

[5]
Beck, A. (2017).
First order methods in optimization.
MOS-SIAM Series on Optimization.

[6]
Davis, D. and Drusvyatskiy, D. (2021).
Proximal methods avoid active strict saddles of weakly convex
Foundations of Computational Mathematics, pages 1–46.

[7]
Deleu, T., Würfl, T., Samiei, M., Cohen, J. P., and Bengio, Y. (2019).
Torchmeta: A Meta-Learning library for PyTorch.

[8]
Fallah, A., Mokhtari, A., and Ozdaglar, A. (2020).
On the convergence theory of gradient-based model-agnostic
meta-learning algorithms.
In International Conference on Artificial Intelligence and
Statistics, pages 1082–1092. PMLR.

[9]
Finn, C., Abbeel, P., and Levine, S. (2017).
Model-agnostic meta-learning for fast adaptation of deep networks.
In International Conference on Machine Learning, pages
1126–1135. PMLR.

[10]
Finn, C., Rajeswaran, A., Kakade, S., and Levine, S. (2019).
Online meta-learning.
In International Conference on Machine Learning, pages
1920–1930. PMLR.

[11]
Franceschi, L., Frasconi, P., Salzo, S., Grazzi, R., and Pontil, M. (2018).
Bilevel programming for hyperparameter optimization and
In International Conference on Machine Learning, pages
1568–1577. PMLR.

[12]
Gower, R. M., Loizou, N., Qian, X., Sailanbayev, A., Shulgin, E., and
Richtárik, P. (2019).
SGD: general analysis and improved rates.
In Proceedings of the 36th International Conference on Machine
Learning, volume 97, pages 5200–5209. PMLR.

[13]
Hanzely, F., Hanzely, S., Horváth, S., and Richtárik, P. (2020).
Lower bounds and optimal algorithms for personalized federated
arXiv preprint.

[14]
Hinton, G., Deng, L., Yu, D., Dahl, G. E., Mohamed, A.-r., Jaitly, N., Senior,
A., Vanhoucke, V., Nguyen, P., Sainath, T. N., and Kingsbury, B. (2012).
Deep neural networks for acoustic modeling in speech recognition: The
shared views of four research groups.
IEEE Signal processing magazine, 29(6):82–97.

[15]
Hospedales, T. M., Antoniou, A., Micaelli, P., and Storkey, A. J. (2021).
Meta-learning in neural networks: A survey.
IEEE transactions on pattern analysis and machine intelligence.

[16]
Ji, K., Yang, J., and Liang, Y. (2020).
Theoretical convergence of multi-step model-agnostic meta-learning.
arXiv e-prints, pages arXiv–2002.

[17]
Khodak, M., Balcan, M.-F. F., and Talwalkar, A. S. (2019).
Adaptive gradient-based meta-learning methods.
Advances in Neural Information Processing Systems, 32.

[18]
Konobeev, M., Kuzborskij, I., and Szepesvári, C. (2021).
A distribution-dependent analysis of meta-learning.
In International Conference on Machine Learning. PMLR.

[19]
Krizhevsky, A., Sutskever, I., and Hinton, G. E. (2012).
Imagenet classification with deep convolutional neural networks.
Advances in neural information processing systems,

[20]
Long, L. (2018).
MAML-pytorch implementation.

[21]
Mania, H., Pan, X., Papailiopoulos, D., Recht, B., Ramchandran, K., and Jordan,
M. I. (2017).
Perturbed iterate analysis for asynchronous stochastic optimization.
SIAM Journal on Optimization, 27(4):2202–2229.

[22]
Mishchenko, K., Khaled, A., and Richtárik, P. (2020).
Random reshuffling: Simple analysis with vast improvements.
Advances in Neural Information Processing Systems,

[23]
Nesterov, Y. (2013).
Introductory lectures on convex optimization: a basic course,
volume 87.

[24]
Nichol, A., Achiam, J., and Schulman, J. (2018).
On first-order meta-learning algorithms.
arXiv preprint arXiv:1803.02999.

[25]
Planiden, C. and Wang, X. (2016).
Strongly convex functions, Moreau envelopes, and the generic nature
of convex functions with strong minimizers.
SIAM Journal on Optimization, 26(2):1341–1364.

[26]
Poliquin, R. and Rockafellar, R. T. (1996).
Prox-regular functions in variational analysis.
Transactions of the American Mathematical Society,

[27]
Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. (2019).
Meta-learning with implicit gradients.
In Advances in Neural Information Processing Systems, pages

[28]
Ravi, S. and Larochelle, H. (2017).
Optimization as a model for few-shot learning.
In ICLR.

[29]
Schmidhuber, J. (1987).
Evolutionary principles in self-referential learning.
Diploma thesis, Technische Universität München.

[30]
Stich, S. U., Cordonnier, J.-B., and Jaggi, M. (2018).
Sparsified SGD with memory.
Advances in Neural Information Processing Systems,

[31]
Sun, Y., Chen, T., and Yin, W. (2021).
An optimal stochastic compositional optimization method with
applications to meta learning.
In ICASSP 2021-2021 IEEE International Conference on Acoustics,
Speech and Signal Processing (ICASSP), pages 3665–3669. IEEE.

[32]
T. Dinh, C., Tran, N., and Nguyen, T. D. (2020).
Personalized federated learning with Moreau envelopes.
Advances in Neural Information Processing Systems, 33.

[33]
Yoon, J., Kim, T., Dia, O., Kim, S., Bengio, Y., and Ahn, S. (2018).
Bayesian model-agnostic meta-learning.
In Proceedings of the 32nd International Conference on Neural
Information Processing Systems, pages 7343–7353.

[34]
Zhang, J., Lin, H., Jegelka, S., Sra, S., and Jadbabaie, A. (2020).
Complexity of finding stationary points of nonsmooth nonconvex
In Proceedings of the 37th International Conference on Machine
Learning, pages 11173–11182.

[35]
Zhou, P., Yuan, X., Xu, H., Yan, S., and Feng, J. (2019).
Efficient meta learning via minibatch proximal update.
Advances in Neural Information Processing Systems,

[36]
Zhou, P., Zou, Y., Yuan, X., Feng, J., Xiong, C., and Hoi, S. C. (2020).
Task similarity aware meta learning: Theory-inspired improvement on
In 4th Workshop on Meta-Learning at NeurIPS.

§ CONTENT LEFT OUT

§.§ Table of frequently used notation

For clarity, we provide a table of frequently used notation.

Notation     Meaning

$f_i$     The loss of task $i$

$F_i(x) = \min_z \{f_i(z) + \frac{1}{2\alpha}\|z-x\|^2\}$     Meta-loss

$F(x)=\frac{1}{n}\sum_{i=1}^n F_i(x)$     Full meta loss

$z_i(x)=\argmin_z \{f_i(z) + \frac{1}{2\alpha}\|z-x\|^2\}$      The minimizer of regularized loss

$L, \mu$     Smoothness and strong convexity constants of $f_i$

$\Lmeta$     Smoothness constant of $\meta$

$\alpha$     Objective parameter

$\outers$     Stepsize of the outer loop

$\locstep, s$     Stepsize and number of steps in the inner loop

$\delta$     Precision of the proximal oracle

§.§ Parametrization of the inner loop of <Ref>

Note that <Ref> depends on only one parameter – $\outers$. We need to keep in mind that parameter $\inners$ is fixed by the objective (<ref>) and changing $\inners$ shifts convergence neighborhood. Nevertheless, we can still investigate the case wehn $\inners$ from (<ref>) and $\inners$ from Line $6$ of <Ref> are different, as we can see in the following remark.

If we replace line $6$ of <Ref> by $z_{l+1}^k = x^k - \locstep \nabla f_i(z_{i,l}^k)$, we will have freedom to choose $\locstep$. However, if we choose stepsize $\locstep \neq \inners$, then similar analysis to the proof of <Ref> yields

\begin{align} \label{eq:locstep_inexact}
\frac 1 \locstep \|z_{i,s}^k - (x^k - \locstep\nabla \meta_i(x^k))\|
&\leq  \left( (\locstep L)^s + |\alpha - \locstep| L \right)  \| \nabla \meta_i(x^k)\|.
\end{align}

Note that in case $\locstep \neq \inners$, we cannot set number of steps $s$ to make the right-hand side of (<ref>) smaller than $\delta \| \nabla \meta_i(x^k)\|$ when $\delta$ is small. In particular, increasing the number of local steps $s$ will help only as long as $\delta > |\inners - \locstep|L$.

This is no surprise, for the modified algorithm (using inner loop stepsize $\locstep$) will no longer be approximating $\nabla \meta_i(x^k)$. It will be exactly approximating $\nabla \tilde \meta_i(x^k)$, where  $\tilde \meta_i(x)\eqdef \min_{z\in\R^d} \left\{f_i(z) + \frac{1}{2\locstep}\|z - x\|^2\right\}$ (see <Ref>).
Thus, choice of stepsize in the inner loop affects what implicit gradients do we approximate and also what objective we are minimizing.s

§ PROOFS

§.§ Basic facts

For any vectors $a, b\in\R^d$ and scalar $\nu>0$, Young's inequality states that

\begin{align}
2\<a, b>\le \nu\|a\|^2+ \frac{1}{\nu}\|b\|^2. \label{eq:young}
\end{align}

Moreover, we have

\begin{align}
\|a+b\|^2
\le 2\|a\|^2 + 2\|b\|^2.\label{eq:a_plus_b}
\end{align}

More generally, for a set of $m$ vectors $a_1,\dotsc, a_m$ with arbitrary $m$, it holds

\begin{align}
\Bigl\|\frac{1}{m}\sum_{i=1}^m a_i\Bigr\|^2
\le \frac{1}{m}\sum_{i=1}^m\|a_i\|^2. \label{eq:jensen_for_sq_norms}
\end{align}

For any random vector $X$ we have

\begin{align}
\ec{\|X\|^2}=\|\ec{X}\|^2 + \ec{\|X-\ec{X}\|^2}. \label{eq:rand_vec_sq_norm}
\end{align}

If $f$ is $L_f$-smooth, then for any $x,y\in\R^d$, it is satisfied

\begin{equation}
\le f(x) + \<\nabla f(x), y-x> + \frac{L_f}{2}\|y-x\|^2. \label{eq:smooth_func_vals}
\end{equation}

Finally, for $L_f$-smooth and convex function $f$, it holds

\begin{equation}
f(x) \leq f(y)+\langle\nabla f(x), x-y\rangle-\frac{1}{2 L_f}\|\nabla f(x)-\nabla f(y)\|^{2} \label{eq:smooth_conv}.
\end{equation}

[Three-point identity]
For any $u,v,w \in \R^d$, any $f$ with its Bregman divergence $D_f(x,y) = f(x)-f(y) - \langle \nabla f(y), x-y\rangle $, it holds

\begin{align*}
\langle \nabla f(u) - \nabla f(v), w-v \rangle
&= D_f (v,u) + D_f(w,v) - D_f(w,u).
\end{align*}

§.§ Proof of <Ref>

The counterexample that we are going to use is given below:

\begin{align*}
f(x) &= \min\left\{\frac{1}{4}x^4 - \frac{1}{3}|x|^3 + \frac{1}{6}x^2, \frac{2}{3}x^2 - |x| + \frac{5}{12} \right\} \\
\begin{cases}
\frac{1}{4}x^4 - \frac{1}{3}|x|^3 + \frac{1}{6}x^2, &\textrm{if } |x|\le 1,\\
\frac{2}{3}x^2 - |x| + \frac{5}{12}, & \textrm{otherwise}.
\end{cases}
\end{align*}

See also Figure <ref> for its numerical visualization.

Values of functions $f$ and $\varphi$.

Illustration of nonconvexity: the value of $\varphi$ goes below its tangent line from $x_0$, which means that $\varphi$ is nonconvex at $x_0$.

It is straightforward to observe that this function is smooth and convex because its Hessian is

\[
= \begin{cases}
3x^2 - 2|x| + \frac{1}{3}, &\textrm{if } |x|\le 1,\\
\frac{4}{3}, & \textrm{otherwise}.
\end{cases},
\]

which is always nonnegative and bounded. However, the function $\varphi(x) = f(z(x))$ is not convex at point $x_0=0.4 + \alpha\nabla f(0.4)$, because its Hessian is negative, i.e., $\varphi''(x_0)<0$, which we shall prove below. First of all, by definition of $x_0$, it holds that $0.4=x_0 - \alpha\nabla f(0.4)$, which is equivalent to the definition of $z(x)$, implying $z(x_0)=0.4$. Next, let us obtain the expression for the Hessian of $\varphi$. As shown in [27], it holds in general that

\[
\nabla \varphi(x)
= \frac{dz(x)}{dx}\nabla f(z(x)),
\]

where $\frac{dz(x)}{dx}$ is the Jacobian matrix of the mapping $z(x)$. Differentiating this equation again, we obtain

\[
\nabla^2 \varphi(x)
= \frac{d^2z(x)}{dx^2}\nabla f(z(x)) + \nabla^2 f(z(x))\frac{dz(x)}{dx}\Bigl(\frac{dz(x)}{dx}\Bigr)^\top.
\]

Moreover, we can compute $\frac{d^2z(x)}{dx^2}$ by differentiating two times the equation $z(x) = x - \alpha \nabla f(z(x))$, which gives

\[
\frac{dz(x)}{dx} = \mathbf{I} - \alpha \nabla^2 f(z(x))\frac{dz(x)}{dx},
\]

where $\mathbf{I}$ is the identity matrix. Rearranging the terms in this equation yields

\[
\frac{dz(x)}{dx} = (\mathbf{I} + \alpha \nabla^2 f(z(x)))^{-1}.
\]

At the same time, if we do not rearrange and instead differentiate the equation again, we get

\[
\frac{d^2z(x)}{dx^2}
= -\alpha \nabla^2 f(z(x))\frac{d^2z(x)}{dx^2} - \alpha \nabla^3 f(z(x)) \left[\frac{dz(x)}{dx}, \frac{dz(x)}{dx}\right],
\]

where $\nabla^3 f(z(x))[\frac{dz(x)}{dx},\frac{dz(x)}{dx}]$ denotes tensor-matrix-matrix product, whose result is a tensor too. Thus,

\[
\frac{d^2z(x)}{dx^2}
= -\alpha (\mathbf{I} + \alpha\nabla^2 f(z(x)))^{-1}\nabla^3 f(z(x))\left[\frac{dz(x)}{dx}, \frac{dz(x)}{dx}\right],
\]

and, moreover,

\[
\nabla^2 \varphi(x)
= -\alpha(\mathbf{I} + \alpha\nabla^2 f(z(x)))^{-1}\nabla^3 f(z(x))\left[\frac{dz(x)}{dx}, \frac{dz(x)}{dx}\right] + \nabla^2 f(z(x))\frac{dz(x)}{dx}\Bigl(\frac{dz(x)}{dx}\Bigr)^\top.
\]

For any $x\in(0, 1]$, our counterexample function satisfies $f''(x)=3x^2 - 2x + \frac{1}{3}$ and $f'''(x)=6x-2$. Moreover, since $z(x_0)=0.4$, we have $f''(z(x_0))=\frac{1}{75}$,  $f'''(z(x_0))=\frac{2}{5}$, $\frac{dz(x)}{dx}= \frac{1}{1+ \alpha/75}$, and

\[
\varphi''(x)=- \frac{2\alpha}{5(1+ \alpha/75)^3} + \frac{1}{75(1+ \alpha/75)^2}.
\]

It can be verified numerically that $\varphi''(x)$ is negative at $x_0$ for any $\alpha>\frac{75}{2249}$. Notice that this value of $\alpha$ is much smaller than the value of $\frac{1}{L}=\frac{3}{4}$, which can be obtained by observing that our counterexample satisfies $f''(x)\le \frac{4}{3}$.

Let us also note that obtaining nonconvexity of this objective for a fixed function and arbitrary $\alpha$ is somewhat challenging. Indeed, in the limit case $\alpha\to 0$, it holds that $\varphi(x)''\to f''(x)$ for any $x$. If $f''(x)>0$ then for a sufficiently small $\alpha$ it would also hold $\varphi''(x)>0$. Finding an example that works for any $\alpha$, thus, would require $f''(x_0)=0$.

§.§ Proof of <Ref>

Consider the following simple function

\[
f(x) = \frac{1}{2}x^2 + \cos(x).
\]

The Hessian of $f$ is $f''(x)=1 - \cos(x)\ge 0$, so it is convex. Moreover, it is apparent that the gradient and the Hessian of $f$ are Lipschitz. However, we will show that the Hessian of $\varphi$ is unbounded for any fixed $\alpha>0$. To establish this, let us first derive some properties of $z(x)$. First of all, by definition $z(x) $ is the solution of $\alpha f'(z(x))+(z(x)-x)=0$, where by definition of $f$, it holds $f'(z(x))=z(x) - \sin(z(x))$. Plugging it back, we get

\[
(\alpha + 1) z(x) - \alpha \sin (z(x)) = x.
\]

Differentiating both sides with respect to $x$, we get $(\alpha+1)\frac{d z(x)}{dx} - \alpha \cos(z(x))\frac{d z(x)}{dx} = 1$ and

\[
\frac{dz(x)}{dx} = \frac{1}{1 + \alpha - \alpha \cos(z(x))}.
\]

Thus, using the fact that $\varphi(x)=\varphi(z(x))$, we get

\[
\varphi'(x)
= \frac{d\varphi(x)}{dx}
= \frac{df(z)}{dz}\frac{d z(x)}{dx}
= \frac{z(x) - \sin(z(x)) }{1 + \alpha - \alpha \cos(z(x))}.
\]

Denoting, for brevity, $z(x)$ as $z$, we differentiate this identity with respect to $z$ and derive $\frac{d\varphi'(x)}{dz} = \frac{1 + 2\alpha - \alpha z \sin(z) - (1+2\alpha)\cos(z)}{(1+\alpha - \alpha\cos(z))^2}$.
Therefore, for the Hessian of $\varphi$, we can produce an implicit identity,

\[
\varphi''(x)
= \frac{d^2\varphi(x)}{dx^2}
= \frac{d\varphi'(x)}{dz}\frac{dz(x)}{dx}
= \frac{1 + 2\alpha - \alpha z \sin(z) - (1+2\alpha)\cos(z)}{(1+\alpha - \alpha\cos(z))^3}.
\]

The denominator of $\varphi''(x)$ satisfies $|1+\alpha - \alpha\cos(z)|^3\le (1 + 2\alpha)^3$, so it is bounded for any $x$. The numerator, on the other hand, is unbounded in terms of $z(x)$ since $|1 + 2\alpha - \alpha z \sin(z) - (1+2\alpha)\cos(z)|\ge \alpha |z\sin(z)| - 2(1+2\alpha)$. Therefore, $|\varphi''(x)|$ is unbounded. Moreover, $z(x)$ is itself unbounded, since the previously established identity for $z(x)$ can be rewritten as $|z(x)|=\left|\frac{1}{1+\alpha}x - \frac{\alpha}{1+\alpha}\sin(z(x))\right|\ge \frac{1}{1+\alpha}|x| - 1$. Therefore, $z(x)$ is unbounded, and since $\varphi''(x)$ grows with $z$, it is unbounded too. The unboundedness of $\varphi''(x)$ implies that $\varphi$ is not $L$-smooth for any finite $L$.

§.§ Proof of <Ref>

The statement that $\meta_i$ is $\frac{\mu}{1+\alpha\mu}$-strongly convex is proven as Lemma 2.19 in [25], so we skip this part.

For nonconvex $\meta_i$ and any $x\in\R^d$, we have by first-order stationarity of the inner problem that $\nabla \meta_i(x) = \nabla f_i(z_i(x))$, where $z_i(x)=\arg\min_z \{f_i(z) + \frac{1}{2\alpha}\|z-x\|^2\} = x - \alpha \nabla \meta_i(x)$. Therefore,

\begin{align*}
\|\nabla \meta_i(x) - \nabla \meta_i(y)\|
&= \|\nabla f_i(z_i(x)) - \nabla f_i(z_i(y))\|
\le L \|z_i(x) - z_i(y)\| \\
&= L\|x-y-\alpha(\nabla \meta_i(x) - \nabla \meta_i(y))\| \\
&\le L\|x-y\| + \alpha L \|\nabla \meta_i(x) - \nabla \meta_i(y)\|.
\end{align*}

Rearranging the terms, we get the desired bound:

\[
\|\nabla \meta_i(x) - \nabla \meta_i(y)\|
\le \frac{L}{1-\alpha L}\|x-y\|.
\]

For convex functions, our proof of smoothness of $\meta_i$ follows the exact same steps as the proof of Lemma 2.19 in [25]. Let $f_i^*$ be the convex-conjugate of $f_i$. Then, it holds that $\meta_i = (f_i^* + \frac{\alpha}{2}\|\cdot\|^2)^*$, see Theorem 6.60 in [5]. Therefore, $\meta_i^* = f_i^* + \frac{\alpha}{2}\|\cdot\|^2$. Since $f_i$ is $L$-smooth, $f_i^*$ is $\frac{1}{L}$-strongly convex. Therefore, $\meta_i^*$ is $(\frac{1}{L} + \alpha)$-strongly convex, which, finally, implies that $\meta_i$ is $\frac{1}{\frac{1}{L}+\alpha}$-smooth.

The statement $\frac{L}{1+\alpha L}\le L$ holds trivially since $\alpha>0$. In case $\alpha\le \frac{1}{\mu}$, we get the constants from the other statements by mentioning that $\frac{\mu}{1+\alpha\mu}\ge \frac{\mu}{2}$.

The differentiability of $\meta_i$ follows from Theorem 4.4 of 26, who show differentiability assuming $f_i$ is prox-regular, which is a strictly weaker property than $L$-smoothness, so it automatically holds under the assumptions of <Ref>.

§.§ Proof of <Ref>

Let task losses $f_i$ be $L$–smooth and $\alpha>0$. Given $i$ and $x\in\R^d$, we define recursively $z_{i,0}=x$ and $z_{i,j+1} = {\color{blue}x} - \alpha \nabla f_i({\color{mydarkred}z_{i,j}})$. Then, it holds for any $s\ge 0$

\begin{align*}
\left\| \nabla f_i(z_{i,s}) - \nabla \meta_i(x) \right\|
\le (\alpha L)^{s+1} \|\nabla \meta_i(x)\|.
\end{align*}

In particular, the iterates of FO-MAML (<Ref>) satisfy for any $k$

\begin{align*}
\left\| \nabla f_i(z_i^k) - \nabla \meta_i(x^k) \right\|
\le (\alpha L)^2 \|\nabla \meta_i(x^k)\|.
\end{align*}

First, observe that by <ref> it holds

\begin{align*}
z_i(x)=x - \alpha \nabla\meta_i(x)=x-\alpha \nabla f_i(z_i(x)).
\end{align*}

For $s=0$, the lemma's claim then follows from initialization, $z_{i, 0}=x$, since

\[
\|\nabla f_i(z_{i,s}) - \nabla \meta_i(x)\| = \|\nabla f_i(x) - \nabla f_i(z_i(x))\|
\le L \|x - z_i(x)\|
= \alpha L \|\nabla \meta_i(x)\|.
\]

For $s>0$, we shall prove the bound by induction.
We have for any $l\ge 0$

\begin{align*}
\|z_{i,l+1} - (x - \alpha\nabla \meta_i(x))\|
&= \alpha \|\nabla f_i (z_{i,l}) - \nabla \meta_i(x)\|
= \alpha \|\nabla f_i (z_{i,l}) - \nabla f_i(z_i(x))\|
\le \alpha L\|z_{i, l} - z_i(x)\| \\
&= \alpha L\|z_{i, l} - (x - \alpha\nabla \meta_i(x))\|.
\end{align*}

This proves the induction step as well as the lemma itself.

If task losses $f_1,\dotsc, f_n$ are $L$-smooth and $\outers \le \frac{1}{L}$, then it holds

\begin{align}
\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} g_i^k\Bigr\|^2
&\le \left(1 + 2 (\inners L)^{2s} + \frac 2 {|T|} \right)4L(\meta(x^k) - \meta(\opt)) + 4\left(\frac 1 {|T_k|} + (\inners L)^{2s}\right)\sigma_*^2\\
&\le 20L(\meta(x^k) - \meta(\opt)) + 4\left(\frac 1 {|T_k|} + \delta^2\right)\sigma_*^2. \label{eq:ave_grad_norm}
\end{align}

First, let us replace $g_i^k$ with $\nabla \meta_i(x^k)$, which $g_i^k$ approximates:

\begin{align*}
\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} g_i^k\Bigr\|^2
&= \Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k) + \frac{1}{|T_k|}\sum_{i\in T_k} (g_i^k - \nabla\meta_i(x^k))\Bigr\|^2 \\
&\overset{\eqref{eq:a_plus_b}}{\le} 2\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\Bigr\|^2 + 2\Bigl\| \frac{1}{|T_k|}\sum_{i\in T_k} (g_i^k - \nabla\meta_i(x^k))\Bigr\|^2 \\
&\overset{\eqref{eq:jensen_for_sq_norms}}{\le} 2\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\Bigr\|^2 + \frac{2}{|T_k|}\sum_{i\in T_k}\| g_i^k - \nabla\meta_i(x^k)\|^2 \\
&\le 2\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\Bigr\|^2 + \frac{2}{|T_k|}\sum_{i\in T_k} \delta^2 \| \nabla\meta_i(x^k)\|^2.
\end{align*}

Taking the expectation on both sides, we get

\begin{align*}
\ec{\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} g_i^k\Bigr\|^2}
\overset{\eqref{eq:rand_vec_sq_norm}}{\le} 2\|\nabla\meta(x^k)\|^2 + 2\ec{\Bigl\| \frac{1}{|T_k|}\sum_{i\in T_k}\nabla\meta_i(x^k) - \nabla\meta(x^k)\Bigr\|^2} + \frac{2}{n}\sum_{i=1}^n \delta^2\| \nabla\meta_i(x^k)\|^2.
\end{align*}

Moreover, each summand in the last term can be decomposed as

\begin{align*}
\| \nabla\meta_i(x^k)\|^2
\overset{\eqref{eq:a_plus_b}}{\le} 2\| \nabla\meta_i(\opt)\|^2 + 2\|\nabla \meta_i(x^k) - \nabla \meta_i(\opt)\|^2
\overset{\eqref{eq:def_sigma}}{=} 2\sigma_*^2 + 2\|\nabla \meta_i(x^k) - \nabla \meta_i(\opt)\|^2.
\end{align*}

Since $\meta_i$ is convex and $L$-smooth, we have for any $i$

\begin{align*}
\|\nabla \meta_i(x^k) - \nabla \meta_i(\opt)\|^2 \le 2L(\meta_i(x^k) - \meta_i(\opt) - \<\nabla \meta_i(\opt), x^k-\opt>).
\end{align*}

Averaging and using $\frac{1}{n}\sumin \nabla \meta_i(\opt)=0$, we obtain

\begin{align*}
\frac{1}{n}\sum_{i=1}^n\| \nabla\meta_i(x^k)-\nabla\meta_i(\opt)\|^2
\le 2L(\meta(x^k) - \meta(\opt)).
\end{align*}

\begin{align}
\frac{2}{n}\sum_{i=1}^n \delta^2 \| \nabla\meta_i(x^k)\|^2
&\le 4 \delta^2 \sigma_*^2 + 8L \delta^2 (\meta(x^k)-\meta(\opt)) \label{eq:meta_grad_norm} \\
&\le 4 \delta^2 \sigma_*^2 + 8L(\meta(x^k)-\meta(\opt)). \notag
\end{align}

Proceeding to another term in our initial bound, by independence of sampling $i\in T_k$ we have

\begin{align*}
\ec{\Bigl\| \frac{1}{|T_k|}\sum_{i\in T_k}\nabla\meta_i(x^k) - \nabla\meta(x^k)\Bigr\|^2}
&= \frac{1}{|T_k|} \avein \ec{\|\nabla\meta_i(x^k)\|^2} \\
&\overset{\eqref{eq:a_plus_b}}{\le} \frac{2}{|T_k|} \avein \left(\ec{\|\nabla\meta_i(x^k) - \nabla\meta_i(\opt)\|^2} + \ec{\|\nabla \meta_i(\opt)\|^2} \right) \\
&\overset{\eqref{eq:smooth_conv}}{\leq} \frac 2 {|T_k|} \left(2L (\meta(x^k) - \meta(\opt)) + \sigma_*^2 \right)\\
&\le \frac {4L} {|T_k|} (\meta(x^k) - \meta(\opt)) + \frac 2 {|T_k|} \sigma_*^2.
\end{align*}

Finally, we also have $\|\nabla \meta(x^k)\|^2\le 2L(\meta(x^k) - \meta(\opt))$. Combining all produced bounds, we get the claim

\begin{align}
\Bigl\|\frac{1}{|T_k|}\sum_{i\in T_k} g_i^k\Bigr\|^2
\le \left(1 + 2  \delta^2 + \frac 2 {|T|} \right)4L(\meta(x^k) - \meta(\opt)) + 4\left(\frac 1 {|T_k|} + \delta^2 \right)\sigma_*^2.
\end{align}

§.§ Proof of <Ref>

Let task losses $f_1,\dotsc, f_n$ be $L$-smooth and $\mu$-strongly convex. If $|T_k|=\tau$ for all $k$, $\alpha\le\frac{1}{L}, \outers \leq \frac 1 {20L}$ and $\delta \leq \frac 1 {4 \sqrt \kappa}$, where $\kappa\eqdef \frac{L}{\mu}$, then the iterates of <Ref> satisfy

\begin{align*}
\ec{\|x^k-\opt\|^2}
\le \left(1 - \frac{\outers\mu}{4}\right)^k\|x^0-\opt\|^2 + \frac{16}{\mu} 	\left( \frac {2\delta^2} \mu + \frac \outers {\tau} + \outers  \delta^2  \right) \sigma_*^2.
\end{align*}

For the iterates of <Ref>, we can write

\begin{align*}
= x^k - \frac{\outers}{\tau}\sum_{i\in T_k} g_i^k.
\end{align*}

We also have by <Ref> that

\[
\|g_i^k - \nabla \meta_i(x^k)\|^2 \le (\alpha L)^2\delta^2 \|\nabla \meta_i (x^k)\|^2 \le \delta^2 \|\nabla \meta_i (x^k)\|^2,
\]

so let us decompose $g_i^k$ into $\nabla \meta_i(x^k)$ and the approximation error:

\begin{align*}
\|x^{k+1} - \opt\|^2
&= \|x^k - \opt\|^2 - \frac{2\outers}{\tau}\sum_{i\in T_k} \<g_i^k, x^k - \opt> + \outers^2\Bigl\|\frac{1}{\tau}\sum_{i\in T_k} g_i^k\Bigr\|^2 \\
&= \|x^k - \opt\|^2 - \frac{2\outers}{\tau}\sum_{i\in T_k} \<\nabla\meta_i(x^k), x^k - \opt> + \frac{2\outers}{\tau}\sum_{i\in T_k} \<\nabla\meta_i(x^k) - g_i^k, x^k - \opt> + \outers^2\Bigl\|\frac{1}{\tau}\sum_{i\in T_k} g_i^k\Bigr\|^2.
\end{align*}

First two terms can be upperbounded using strong convexity (recall that by <Ref>, $\meta_i$ is $\frac{\mu}{2}$-strongly convex):

\begin{align*}
\|x^k - \opt\|^2 - \frac{2\outers}{\tau}\sum_{i\in T_k} \<\nabla\meta_i(x^k), x^k - \opt>
&\le \left(1 - \frac{\outers\mu}{2}\right)\|x^k - \opt\|^2 - \frac{2\outers}{\tau}\sum_{i\in T_k}(\meta_i(x^k) - \meta_i(\opt)).
\end{align*}

For the third term, we will need Young's inequality:

\begin{align*}
2\<\nabla\meta_i(x^k) - g_i^k, x^k - \opt>
&\overset{\eqref{eq:young}}{\le}\frac{4}{\mu}\|\nabla\meta_i(x^k) - g_i^k \|^2 + \frac{\mu}{4}\|x^k - \opt\|^2 \le\frac{4}{\mu} \delta^2 \|\nabla\meta_i(x^k)\|^2 + \frac{\mu}{4}\|x^k - \opt\|^2,
\end{align*}

which we can scale by $\beta$ and average over $i\in T_k$ to obtain

\[
\frac{2\outers}{\tau}\sum_{i\in T_k} \<\nabla\meta_i(x^k) - g_i^k, x^k - \opt>
\leq  \frac{4 \outers \delta^2}{\mu} \frac 1 \tau \sum_{i \in T_k} \|\nabla \meta_i(x^k)\|^2 + \frac{\beta \mu}{4}\|x^k - \opt\|^2.
\]

Plugging in upper bounds and taking expectation yields

\begin{align*}
\ec{\|x^{k+1}-\opt\|^2}
&\le \left(1 - \frac{\outers\mu}{4}\right)\|x^k-\opt\|^2 - 2\outers (\meta(x^k) - \meta(\opt)) + \frac{4}{\mu}\outers \delta^2 \avein \|\nabla \meta_i(x^k)\|^2 + \outers^2\Bigl\| \frac{1}{\tau}\sum_{i\in T_k} g_i^k\Bigr\|^2  \\
&\overset{\eqref{eq:ave_grad_norm}}{\le} \left(1 - \frac{\outers\mu}{4}\right)\|x^k-\opt\|^2 - 2\outers(1 - 10\beta L)  (\meta(x^k) - \meta(\opt)) + \frac{4}{\mu}\outers  \delta^2 \avein\|\nabla \meta_i(x^k)\|^2 \\
&\quad + 4\outers^2 \left(\frac 1 {\tau} + \delta^2 \right)\sigma_*^2 \\
&\overset{\eqref{eq:meta_grad_norm}}{\le} \left(1 - \frac{\outers\mu}{4}\right)\|x^k-\opt\|^2 - 2\outers(1 - 10\outers L)  (\meta(x^k) - \meta(\opt))\\
&\quad + \frac{8}{\mu}\outers \delta^2\left(\sigma_*^2 + 2L(\meta(x^k) - \meta(\opt))\right) + 4\outers^2 \left(\frac 1 {\tau} + \delta^2\right)\sigma_*^2 \\
&= \left(1 - \frac{\outers\mu}{4}\right)\|x^k-\opt\|^2 - 2\outers \left(1 - 10\outers L - \frac {8L} \mu  \delta^2 \right)  (\meta(x^k) - \meta(\opt)) + \frac{8}{\mu}\outers \delta^2 \sigma_*^2 + 4\outers^2 \left(\frac 1 {\tau} + \delta^2 \right)\sigma_*^2.
\end{align*}

By assumption $\outers \leq \frac 1 {20L}, \delta \leq \frac 1 {4\sqrt{ \kappa}}$, we have $10\outers L \leq \frac 1 2$ and $8 \frac{L}{\mu} \delta^2\leq  \frac 1 2$, so $1 - 10\outers L - \frac {8L} \mu  \delta^2  \geq 0$, hence

\begin{align*}
\ec{\|x^{k+1}-\opt\|^2}
&\le \left(1 - \frac{\outers\mu}{4}\right)\|x^k-\opt\|^2 + \frac{8}{\mu}\outers \delta^2 \sigma_*^2 + 4\outers^2 \left(\frac 1 {\tau} +  \delta^2 \right)\sigma_*^2.
%	&\le \left(1 - \frac{\outers\mu}{4}\right)\|x^k-\opt\|^2 + \frac{8}{\mu}\outers\delta\sigma_*^2 + 2\outers^2\sigma_*^2.
\end{align*}

Recurring this bound, which is a standard argument, we obtain the theorem's claim.

\begin{align*}
\ec{\|x^{k}-\opt\|^2}
&\le \left(1 - \frac{\outers\mu}{4}\right)^k\|x^0-\opt\|^2 + \left( \frac{8}{\mu}\outers \delta^2 \sigma_*^2 + 4\outers^2 \left(\frac 1 {\tau} + \delta^2\right)\sigma_*^2 \right) \frac {1-\left( 1- \frac {\outers \mu} 4 \right)^k}{\frac {\outers \mu} 4}\\
&\le \left(1 - \frac{\outers\mu}{4}\right)^k\|x^0-\opt\|^2 + \frac{32}{\mu^2} \delta^2 \sigma_*^2  + \frac{16}{\mu \tau} \outers \sigma_*^2 + \frac{16}{\mu} \outers  \delta^2 \sigma_*^2\\
&\le \left(1 - \frac{\outers\mu}{4}\right)^k\|x^0-\opt\|^2 + \frac{16}{\mu} \left( \frac {2\delta^2} \mu + \frac \outers {\tau} + \outers  \delta^2  \right) \sigma_*^2.
\end{align*}

§.§ Proof of <Ref>

Consider the iterates of <Ref> (with general $\delta$) or <Ref> (for which $\delta=\alpha L$).
Let task losses be $L$–smooth and $\mu$–strongly convex and let objective parameter satisfy $\inners \leq \frac {1}{\sqrt 6 L}$. Choose stepsize $ \beta \leq \frac \tau {4 L}$, where $\tau = |T_k|$ is the batch size. Then we have

\begin{align*}
\E{\norm{x^k-x^*}^2} \leq \left(1 - \frac {\beta \mu}{12}  \right)^k \norm{x^0 - x^*}^2 + \frac { 6\left( \frac \beta \tau + 3 \delta^2 \inners^2 L \right) \varopt} {\mu}.
\end{align*}

Denote $L_\meta$, $\mu_\meta, \kappa_\meta = \frac \Lmeta {\mu_\meta}$ smoothness constant, strong convexity constant, condition number of Meta-Loss functions $\meta_1, \dots, \meta_n$, respectively.
We have

\begin{align*}
\norm{x^{k+1}-x^*}^2
&= \biggl\|x^k - x^* - \frac \beta \tau \sum_{i \in T_k}  \nabla \meta_i (y_i^k)\biggr\|^2\\
&=\norm{x^k - x^*}^2 - \frac {2 \beta} \tau \sum_{i \in T_k} \langle \nabla \meta_i (y_i^k), x^k-x^* \rangle + \beta^2 \norm{ \frac 1 \tau \sum_{i \in T_k} \nabla  \meta_i(y_i^k)}^2\\
&\leq \norm{x^k - x^*}^2 - \frac {2 \beta} \tau \sum_{i \in T_k} \langle \nabla \meta_i (y_i^k) - \nabla \meta_i(x^*), x^k-x^* \rangle + 2\beta^2 \norm{ \frac 1 \tau \sum_{i \in T_k} (\nabla  \meta_i(y_i^k) - \nabla  \meta_i(x^*) ) }^2\\
& \quad  - \frac {2 \beta} \tau \sum_{i \in T_k} \langle \nabla \meta_i (x^*), x^k-x^* \rangle + 2 \beta^2 \biggl\| \frac 1 \tau \sum_{i \in T_k} \nabla  \meta_i(x^*)\biggr\|^2.
\end{align*}

Using Proposition <ref>, we rewrite the scalar product as
\langle \nabla \meta_i (y_i^k) - \nabla \meta_i(x^*), x^k-x^* \rangle
= D_{\meta_i}(x^*, y_i^k) + D_{\meta_i}(x^k, x^*) - D_{\meta_i}( x^k, y_i^k) %\label{eq:three_point_used}
%	& \geq D_{\meta_i}(x^k, y_i^k) + D_{\meta_i}(x^*, x^k) - \frac 1 {2 \mu_F} \norm{\nabla \meta_i (x^*) - \nabla \meta_i (y_i^k)}^2
which gives

\begin{align*}
\norm{x^{k+1}-x^*}^2
& \leq \norm{x^k - x^*}^2 - \frac {2 \beta} \tau \sum_{i \in T_k} \left[ D_{\meta_i}(x^*, y_i^k) + D_{\meta_i}(x^k, x^*) - D_{\meta_i}( x^k, y_i^k) \right]\\
&+ 2\beta^2 \biggl\| \frac 1 \tau \sum_{i \in T_k} (\nabla  \meta_i(y_i^k) - \nabla  \meta_i(x^*) ) \biggr\|^2  - \frac {2 \beta} \tau \sum_{i \in T_k} \langle \nabla \meta_i (x^*), x^k-x^* \rangle + 2 \beta^2 \biggl\| \frac 1 \tau \sum_{i \in T_k} \nabla  \meta_i(x^*)\biggr\|^2.
\end{align*}

Since we sample $T_k$ uniformly and $\{\nabla \meta_i (x^*)\}_{i\in T_k}$ are independent random vectors, we obtain

\begin{align*}
\E{\norm{x^{k+1}-x^*}^2}
&\leq \norm{x^k - x^*}^2 + \frac {2 \beta} \tau \E{ \sum_{i \in T_k} \left[ - D_{\meta_i}(x^*, y_i^k) - D_{\meta_i}(x^k, x^*) + D_{\meta_i}( x^k, y_i^k)\right] }\\
& \quad + \frac {2\beta^2} {\tau^2} \E{ \sum_{i \in T_k}  \norm{ \nabla  \meta_i(y_i^k) - \nabla  \meta_i(x^*) }^2 } + \frac {2 \beta^2} \tau \varopt.
\end{align*}

Next, we are going to use the following three properties of Bregman divergence:

\begin{align}
-D_{\meta_i}( x^*, y_i^k)
& \overset{\eqref{eq:smooth_conv}}{\le} -\frac 1 {2 L_\meta} \norm{ \nabla  \meta_i(y_i^k) - \nabla  \meta_i(x^*) }^2  \notag \\%\label{eq:div1}\\
-D_{\meta_i}(x^k, x^*)
&\le  -\frac {\mu_\meta} 2\norm {x^k-x^*}^2\label{eq:div2} \\
D_{\meta_i}(x^k, y_i^k)
& \le \frac \Lmeta 2 \norm{x^k-y_i^k}^2. \notag %\label{eq:div3}
\end{align}

Moreover, using identity $y_i^k = z_i^k + \alpha \nabla F_i(y_i^k)$, we can bound the last divergence as

\begin{align*}
D_{\meta_i}(x^k, y_i^k)
&\le \frac \Lmeta 2 \norm{x^k-z_i^k - \inners \nabla \meta_i (y_i^k)}^2\\
&= \frac 1 2 \inners^2 \Lmeta \Bigl\|\frac 1 \inners (x^k-z_i^k) - \nabla \meta_i(y_i^k)\Bigr\|^2\\
& \le \frac 3 2 \inners^2 \Lmeta \left( \Bigl\|\frac 1 \inners (x^k-z_i^k) - \nabla \meta_i(x^k)\Bigr\|^2 + \norm{\nabla \meta_i (x^k)- \nabla \meta_i(x^*)}^2 + \norm{\nabla \meta_i (x^*)- \nabla \meta_i(y_i^k)}^2  \right)\\
& \le \frac 3 2 \inners^2 \Lmeta \left( \delta^2 \norm{\nabla \meta_i(x^k)}^2 + \norm{\nabla \meta_i (x^k)- \nabla \meta_i(x^*)}^2 + \norm{\nabla \meta_i (x^*)- \nabla \meta_i(y_i^k)}^2  \right),
\end{align*}

where the last step used the condition in <Ref>. Using inequality (<ref>) on $\nabla \meta_i(x^k) = \nabla \meta_i(x^*) + (\nabla \meta_i(x^k) - \nabla \meta_i(x^*))$, we further derive

\begin{align*}
D_{\meta_i}(x^k, y_i^k)
& \le \frac 3 2 \inners^2 \Lmeta \left( 2\delta^2 \norm{\nabla \meta_i(x^*)}^2 + (1+2\delta^2) \norm{\nabla \meta_i (x^k)- \nabla \meta_i(x^*)}^2 + \norm{\nabla \meta_i (x^*)- \nabla \meta_i(y_i^k)}^2  \right)\\
& \overset{\eqref{eq:smooth_conv}}{\le} \frac 3 2 \inners^2 \Lmeta \left( 2\delta^2 \norm{\nabla \meta_i(x^*)}^2 + (1+2\delta^2)\Lmeta D_{\meta_i}(x^k,x^*) + \norm{\nabla \meta_i (x^*)- \nabla \meta_i(y_i^k)}^2  \right).	\end{align*}

Assuming $\inners\leq \sqrt{\frac 2 3 (1+2\delta^2)} \frac 1 {L_\meta}$ so that $ 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) >0$, we get

\begin{align*}
-D_{\meta_i}(x^k, x^*) + D_{\meta_i}(x^k, y_i^k)
&\le - \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right) D_{\meta_i}(x^k,x^*) \\
&\quad + \frac 3 2 \inners^2 \Lmeta \left( 2\delta^2 \norm{\nabla \meta_i(x^*)}^2 + \norm{\nabla \meta_i (x^*)- \nabla \meta_i(y_i^k)}^2  \right)\\
&\stackrel{\eqref{eq:div2}}\le - \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right) \frac {\mu_\meta} 2\norm {x^k-x^*}^2\\
&\quad +  \frac 3 2 \inners^2 \Lmeta \left( 2\delta^2 \norm{\nabla \meta_i(x^*)}^2 + \norm{\nabla \meta_i (x^*)- \nabla \meta_i(y_i^k)}^2  \right).
\end{align*}

Plugging these inequalities yields

\begin{align*}
\E{\norm{x^{k+1}-x^*}^2}
&\leq \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right) \norm{x^k - x^*}^2 \\
& \quad+  \frac \beta \tau \left(3 \inners^2 \Lmeta +  \frac {2\beta} \tau - \frac 1 {L_\meta}  \right) \E{ \sum_{i \in T_k}  \norm{ \nabla  \meta_i(y_i^k) - \nabla  \meta_i(x^*) }^2 }\\
& \quad + 2 \beta\left( \frac \beta \tau + 3 \inners^2 \delta^2 \Lmeta \right)\varopt.
\end{align*}

Now, if $\inners \leq \frac {1}{\sqrt 6 L_\meta}$ and $ \beta \leq \frac \tau {4 L_\meta}$, then $3 \inners^2 \Lmeta +  \frac {2\beta} \tau  - \frac 1 {L_\meta} \leq 0$, and consequently

\begin{align*}
\E{\norm{x^{k+1}-x^*}^2}
&\leq \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right) \norm{x^k - x^*}^2 + 2 \beta \left( \frac \beta \tau + 3 \inners^2 \delta^2 \Lmeta \right)\varopt.
\end{align*}

We can unroll the recurrence to obtain the rate

\begin{align*}
\E{\norm{x^k-x^*}^2}
& \le \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right)^k \norm{x^0 - x^*}^2 \\
&\quad + \left( \sum_{i=0}^{k-1} \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right)^i \right) 2 \beta \left( \frac \beta \tau + 3 \inners^2 \delta^2 \Lmeta \right)\varopt\\
&= \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right)^k \norm{x^0 - x^*}^2 \\
&\quad + \left( \frac {1- \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right)^k }{1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2)} \right) \frac 2 {\mu_\meta} \left( \frac \beta \tau + 3 \inners^2 \delta^2 \Lmeta \right)\varopt \\
&\leq \left(1 - \beta \mu_\meta \left( 1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2) \right)  \right)^k \norm{x^0 - x^*}^2 + \frac {2 \left( \frac \beta \tau + 3 \inners^2 \delta^2 \Lmeta \right) \varopt} {\mu_\meta (1- \frac 3 2 \inners^2 \Lmeta^2 (1+2\delta^2))}.
\end{align*}

Choice of $\delta$ implies $0 \leq \delta \leq 1$; <Ref> yields $\frac \mu 2 \leq \frac \mu {1+\inners \mu} \leq \mu_\meta$ and $L_\meta \leq \frac L {1 + \inners L} \leq L$, so we can simplify

\begin{align*}
\E{\norm{x^k-x^*}^2} \leq \left(1 - \frac {\beta \mu}2 \left( 1- 5 \inners^2 L^2 \right)  \right)^k \norm{x^0 - x^*}^2 + \frac { 4\left( \frac \beta \tau + 3 \inners^2 L \delta^2 \right) \varopt} {\mu (1 - 2 \inners^2 L^2 )}.
\end{align*}

§.§ Proof of <Ref>

Theorem <ref>
Let <Ref> hold, functions $f_1,\dotsc, f_n$ be $L$–smooth and $F$ be lower bounded by $F^*>-\infty$. Assume $\alpha\le \frac{1}{4L}, \beta\le \frac{1}{16L}$. If we consider the iterates of <Ref> (with $\delta=\alpha L$) or <Ref> (with general $\delta$), then

\begin{align*}
\min_{t\le k}\E{\|\nabla F(x^t)\|^2}
\le \frac{4}{\beta k}\E{F(x^0)-F^*}+ 16 \beta(\alpha L)^2 \left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right) \sigma^2.
\end{align*}

We would like to remind the reader that for our choice of $z_i^k$ and $y_i^k$, the following three identities hold. Firstly, by definition $y_i^k = z_i^k + \alpha \nabla f_i(z_i^k)$. Secondly, as shown in Lemma <ref>, $ z_i^k = y_i^k - \alpha\nabla \meta_i(y_i^k)$. And finally, $\nabla f_i(z_i^k) = \nabla \meta_i(y_i^k)$. We will frequently  use these identities in the proof.

Since functions $f_1,\dotsc, f_n$ are $L$-smooth and $\alpha \le \frac{1}{4L}$, functions $F_1,\dotsc, F_n$ are $(2L)$-smooth as per <Ref>. Therefore, by smoothness of $F$, we have for the iterates of <Ref>

\begin{align*}
\E{F(x^{k+1})}
&\overset{\eqref{eq:smooth_func_vals}}{\le} \E{F(x^k) + \<\nabla F(x^k), x^{k+1}-x^k> + L\|x^{k+1}-x^k\|^2 }\\
&= \E{F(x^k) - \beta  \biggl\langle\nabla F(x^k), \aveis i {T_k} \nabla f_i(z_i^k) \biggr\rangle + \beta^2 L \biggl\|\aveis i {T_k} \nabla f_i(z_i^k)}\biggr\|^2\\
%		&= F(x^k) - \beta \|\nabla F(x^k)\|^2 + \mathbb E\Bigg [\beta\<\nabla F(x^k), \nabla F(x^k) - \aveis i {T_k} \nabla f_i(z_i^k)>  \\
%			& \hspace{140pt} + \beta^2 L \norms{\aveis i {T_k} \nabla f_i(z_i^k)} \Bigg] \\
&= F(x^k) - \beta \|\nabla F(x^k)\|^2 + \beta\mathbb{E} \Biggl[\biggl\langle\nabla F(x^k), \nabla F(x^k) - \avein \nabla f_i(z_i^k)\biggr\rangle\Biggr]  \\
& \qquad + \beta^2 L \E{\biggl\| \aveis i{T_k} \nabla f_i(z_i^k)\biggr\|^2 } \\
%		&\le F(x^k) - \beta \|\nabla F(x^k)\|^2 + \frac{\beta}{2}\norms{\nabla F(x^k)} + \frac{\beta}{2} \norms{\nabla F(x^k) - \avein \nabla f_i(z_i^k)}  \\
%			& \qquad + \frac{\beta^2 L}{2} \E{ \norms{ \aveis i {T_k} \nabla f_i(z_i^k)}}  \\
&\overset{\eqref{eq:a_plus_b}}{\le} F(x^k) - \frac{\beta}{2} \|\nabla F(x^k)\|^2 + \frac{\beta}{2} \avein \norms{\nabla F_i(x^k) - \nabla f_i(z_i^k)}  + \beta^2 L \E{\biggl\|\aveis i {T_k}\nabla f_i(z_i^k)\biggr\|^2}.
\end{align*}

Next, let us observe, similarly to the proof of Lemma <ref>, that the gradient approximation error satisfies

\begin{align*}
\norm{\nabla F_i(x^k) - \nabla f_i(z_i^k)}
&= \norm{\nabla F_i(x^k) - \nabla \meta_i(y_i^k)}
\le L\norm{x^k - y_i^k}
= L \norm{x^k - z_i^k - \alpha \nabla\meta_i(y_i^k)} \\
&\le \alpha L\norm{ \nabla \meta(x^k) -\nabla\meta_i(y_i^k)} + \alpha L\Bigl\|\frac{1}{\alpha}(x^k-z_i^k) - \nabla \meta_i(x^k)\Bigr\|\\
&= \alpha L\norm{\nabla \meta(x^k) -\nabla f_i(z_i^k) } + \alpha L\Bigl\|\frac{1}{\alpha}(x^k-z_i^k) - \nabla \meta_i(x^k)\Bigr\|.
\end{align*}

By rearranging and using our assumption on error $\delta$ as formulated in <Ref>, we have

\[
\norm{\nabla \meta_i(x^k) -\nabla f_i(z_i^k)}
\le \frac{\alpha L}{1-\alpha L}\Bigl\|\frac{1}{\alpha}(x^k-z_i^k) - \nabla \meta_i(x^k)\Bigr\|
\le \frac{\alpha L}{1-\alpha L}\delta\|\nabla \meta_i(x^k)\|
\overset{\alpha\le \frac{1}{4L}}{\le} \frac{4}{3}\alpha L\delta \|\nabla \meta_i(x^k)\|.
\]

Squaring this bound and averaging over $i$, we obtain

\begin{align*}
\avein\norms{\nabla \meta_i(x^k) -\nabla f_i(z_i^k)}
&\le \frac{16}{9}(\alpha L)^2\delta^2 \avein\|\nabla \meta_i(x^k)\|^2 \\
&= \frac{16}{9}(\alpha L)^2\delta^2 \|\nabla \meta(x^k)\|^2 + \frac{16}{9}(\alpha L)^2\delta^2 \avein\|\nabla \meta_i(x^k) - \nabla \meta (x^k)\|^2 \\
&\overset{\eqref{eq:bounded_var}}{\le} \frac{16}{9}(\alpha L)^2\delta^2 \|\nabla \meta(x^k)\|^2 + \frac{16}{9}(\alpha L)^2\delta^2 \sigma^2 \\
&\le \frac{1}{9}\|\nabla \meta(x^k)\|^2 + 2 (\alpha L)^2\delta^2 \sigma^2.
\end{align*}

For the other term in the smoothness upper bound, we can write

\begin{align*}
\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla f_i(z_i^k)\biggr\|^2}
&= \E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k) + \frac{1}{|T_k|}\sum_{i\in T_k} (\nabla f_i(z_i^k) - \nabla\meta_i(x^k))\biggr\|^2} \\
&\overset{\eqref{eq:a_plus_b}}{\le} 2\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\biggr\|^2} + 2\E{\biggl\| \frac{1}{|T_k|}\sum_{i\in T_k} (\nabla f_i(z_i^k) - \nabla\meta_i(x^k))\biggr\|^2} \\
&\overset{\eqref{eq:jensen_for_sq_norms}}{\le} 2\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\biggr\|^2} + \frac{2}{|T_k|}\E{\sum_{i\in T_k}\| \nabla f_i(z_i^k) - \nabla\meta_i(x^k)\|^2} \\
&\le 2\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\biggr\|^2} + \E{\frac{32}{9}\frac{1}{|T_k|}\sum_{i\in T_k} (\alpha L)^2\delta^2 \| \nabla\meta_i(x^k)\|^2}.
\end{align*}

Using bias-variance decomposition, we get for the first term in the right-hand side

\begin{align*}
2\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k)\biggr\|^2}
&\overset{\eqref{eq:rand_vec_sq_norm}}{=} 2\|\nabla \meta (x^k)\|^2 + 2\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla\meta_i(x^k) - \nabla\meta(x^k)\biggr\|^2} \\
&= 2\|\nabla \meta (x^k)\|^2 + \frac{2}{|T_k|}\frac{1}{n}\sum_{i=1}^n\| \nabla\meta_i(x^k) - \nabla\meta(x^k)\|^2.
\end{align*}

Similarly,  we simplify the second term using $\frac{32}{9}< 4$ and then obtain

\begin{align*}
\frac{32}{9}\E{\frac{1}{|T_k|}\sum_{i\in T_k} (\alpha L)^2\delta^2 \| \nabla\meta_i(x^k)\|^2}
\overset{\eqref{eq:rand_vec_sq_norm}}{\le} 4(\alpha L)^2\delta^2\|\nabla\meta(x^k)\|^2 + \frac{4(\alpha L)^2\delta^2}{n}\sum_{i=1}^n\| \nabla\meta_i(x^k) - \nabla\meta(x^k)\|^2.
\end{align*}

Thus, using $\alpha\le\frac{1}{4L}$ and $\delta\le 1$, we get

\begin{align*}
\E{\biggl\|\frac{1}{|T_k|}\sum_{i\in T_k} \nabla f_i(z_i^k)\biggr\|^2}
&\le 3\|\nabla \meta (x^k)\|^2 + \left(\frac{2}{|T_k|} + 4(\alpha L)^2\delta^2\right)\sum_{i=1}^n\| \nabla\meta_i(x^k) - \nabla\meta(x^k)\|^2 \\
&\overset{\eqref{eq:bounded_var}}{\le} 3\|\nabla \meta (x^k)\|^2 + 4\left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right)\sigma^2.
\end{align*}

Now we plug these inequalities back and continue:

\begin{align*}
\E{F(x^{k+1})} - F(x^k)
& \leq - \frac{\beta}{2} \|\nabla F(x^k)\|^2 + \frac{\beta}{18}\|\nabla\meta (x^k)\|^2 + \beta(\alpha L)^2\delta^2 \sigma^2 \\
&\quad  + 3\beta^2 L\|\nabla \meta (x^k)\|^2 + 4\beta^2 L\sigma^2\left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right)\sigma^2 \\
&\overset{\beta\le \frac{1}{16L}}{\le }   - \frac{\beta}{4} \|\nabla F(x^k)\|^2 + 4\beta^2 L\sigma^2\left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right)\sigma^2 + \beta(\alpha L)^2\delta^2 \sigma^2.
\end{align*}

Rearranging the terms and telescoping this bound, we derive

\begin{align*}
\min_{t\le k}\E{\|\nabla F(x^t)\|^2}
&\le \frac{4}{\beta k}\E{F(x^0)-F(x^{k+1})}+ 16 \beta \left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right) \sigma^2 + 4(\alpha L)^2\delta^2 \sigma^2\\
&\le \frac{4}{\beta k}\E{F(x^0)-F^*}+ 16 \beta \left(\frac{1}{|T_k|} + (\alpha L)^2\delta^2\right) \sigma^2 + 4(\alpha L)^2\delta^2 \sigma^2.
\end{align*}

The result for <Ref> is obtained as a special case with $\delta=\alpha L$.