# Memory-based Optimization Methods
for Model-Agnostic Meta-Learning

Bokun Wang<EMAIL_ADDRESS>
Department of Computer Science
The University of Iowa
Iowa City, IA 52242, USA Zhuoning Yuan<EMAIL_ADDRESS>
Department of Computer Science
The University of Iowa
Iowa City, IA 52242, USA Yiming Ying<EMAIL_ADDRESS>
Department of Mathematics and Statistics
University at Albany
Albany, NY 12222, USA Tianbao Yang<EMAIL_ADDRESS>
Department of Computer Science
The University of Iowa
Iowa City, IA 52242, USA

###### Abstract

Recently, model-agnostic meta-learning (MAML) has garnered tremendous
attention. However, the stochastic optimization of MAML is still immature.
Existing algorithms for MAML exploit the “episode” idea by sampling some tasks
and some data points for each sampled task at each iteration for updating the
meta-model. However, they either do not necessarily guarantee convergence with
a constant mini-batch size or require processing a large number of tasks at
every iteration, which is not viable for continual learning or cross-device
federated learning where only a small number of tasks are available per-
iteration or per-round. This paper addresses these issues by (i) proposing
efficient memory-based stochastic algorithms for MAML with the vanishing
convergence error, which only requires sampling a constant number of tasks and
a constant number of data samples per-iteration; (ii) proposing communication-
efficient distributed memory-based MAML algorithms for personalized federated
learning in both the cross-device (with client sampling) and the cross-silo
(without client sampling) settings. The theoretical results significantly
improve the optimization theory for MAML and the empirical results also
corroborate the theory.

Keywords: Meta-Learning, Federated Learning, Model-Agnostic Meta-Learning,
Personalized Federated Learning, Memory-based Algorithms

##  1 Introduction

Although modern deep learning approaches have achieved great success in many
areas, they draw criticism for the reliance on large volumes of data (Marcus,
2018). In contrast, humans can learn with a small amount of data due to their
capabilities of continual learning from many tasks. Recently, meta-learning
has attracted tremendous attention due to its ability to perform well on new
tasks with a small amount of data using the adaptation and the prior knowledge
learned from previous tasks. In particular, the gradient-based model-agnostic
meta-learning (MAML) approach (Finn et al., 2017) has made it viable for a
broad range of problems whose models can be trained with a gradient descent
procedure. Mathematically, MAML (Finn et al., 2017) proposes to solve the
following optimization problem.

$\displaystyle\min_{\mathbf{w}\in\mathbb{R}^{d}}F(\mathbf{w})=\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}_{i}(\mathbf{w}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w})).$
(1)

where $\mathbf{w}\in\mathbb{R}^{d}$ represents the meta model, $n$ denotes the
number of tasks, and $\mathcal{L}_{i}$ denotes the risk function for the
$i$-th task. In particular,
$\mathcal{L}_{i}(\mathbf{w})=\mathbb{E}_{\mathbf{z}\sim\mathbf{D}_{i}}[\ell_{i}(\mathbf{w},\mathbf{z})]$,
where $\mathbf{D}_{i}$ denotes the distribution of data for th $i$-th task and
$\ell_{i}(\cdot)$ denotes a loss function. The inner gradient step
$\mathbf{w}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w})$ denotes a quick
adaptation from the meta model $\mathbf{w}$ to the individual task. Due to
this adaptation and personalization feature, MAML has also been applied to
personalized federated learning (Jiang et al., 2019), where each client has
access to its corresponding task.

Although many studies have been devoted to MAML (Nichol et al., 2018; Antoniou
et al., 2018; Behl et al., 2019; Yoon et al., 2018; Raghu et al., 2019; Nichol
et al., 2018; Li et al., 2017; Grant et al., 2018), the stochastic
optimization algorithms for solving the problem (1) are still immature. In the
optimization algorithms for MAML, there are two key quantities in each
“episode" (in other words, every iteration), namely the number of sampled
tasks in each iteration denoted by $B$ and the number of sampled data points
per task in each iteration denoted by $K$. It is notable that the episode is
designed to mimic the few-shot task by sub-sampling tasks as well as data
points (Vinyals et al., 2016; Snell et al., 2017; Ravi and Larochelle, 2017).
Unfortunately, the original MAML method based on the episode does not
necessarily converge to a stationary solution of the objective function $F$ in
(1) unless $K$ is large enough. In particular, Fallah et al. (2020a) provide
the first convergence analysis of the original MAML approach. Their theory
suggests that in order to find an $\epsilon$-stationary solution $\mathbf{w}$
satisfying $\mathbb{E}[\|\nabla F(\mathbf{w})\|]\leq\epsilon$, one needs to
run MAML for $T=\mathcal{O}(1/\epsilon^{2})$ iterations and sample
$K=\mathcal{O}(1/\epsilon^{2})$ data points for all $n$ tasks in each
iteration, where the settings of batch sizes are not practical for driving the
error level $\epsilon$ to be sufficiently small. Recently, Hu et al. (2020)
propose two biased stochastic methods namely BSGD and BSpiderBoost, where BSGD
is essentially the original MAML algorithm (Finn et al., 2017) when applied to
(1). Their convergence results guarantee that in order to find an
$\epsilon$-stationary solution one needs to set
$K=\mathcal{O}(1/\epsilon^{2})$, $B=\mathcal{O}(1)$ and
$T=\mathcal{O}(1/\epsilon^{4})$ for BSGD or
$K=\mathcal{O}(1/\epsilon^{2}),B=\sqrt{n}$ and $T=\mathcal{O}(1/\epsilon^{2})$
for BSpiderBoost. Neither setting is practical for driving the error level
$\epsilon$ to be sufficiently small, not to mention the imposed additional
assumptions (see Table 1 for a more thorough comparison). These studies imply
that if a small batch size $K$ is used, the original MAML approach might not
converge to an accurate solution. There are also some studies tackling the
optimization of (1) by considering it as an instance of two-level
compositional functions (Chen et al., 2020; Tutunov et al., 2020) or bilevel
optimization problems (Franceschi et al., 2018; Ji et al., 2020b; Chen et al.,
2021), and designing stochastic algorithms for solving these optimization
problems. However, these existing algorithms either require $K$ to be very
large or require passing through all $n$ tasks at each iteration, which is
prohibitive for continual learning and cross-device federated learning.

This paper aims to improve MAML optimization by addressing the following
question:

Can we design efficient stochastic optimization algorithms for MAML, which can
converge to a stationary solution with only $K=\mathcal{O}(1)$ and
$B=\mathcal{O}(1)$ to update the model?

###  1.1 Contributions

We summarize our contributions below.

  * •

We consider the stochastic optimization for MAML in both a centralized setting
and a federated learning setting, where in the former setting the server has
access to all tasks and their data, and in the latter setting the central
sever has no access to the individual tasks and their data on distributed
clients. We propose two memory-based stochastic algorithms, namely MOML and
LocalMOML. MOML is tailored to the centralized setting and LocalMOML can be
implemented in both centralized and federated learning settings. The key
feature of the proposed algorithms is to maintain and update individualized
models (i.e., memories) for individual tasks by a Momentum update, which is
essentially a moving average of historical stochastic updates of individual
models.

  * •

We establish the convergence guarantees of MOML and LocalMOML for finding a
stationary solution to the non-convex objective with only $K=\mathcal{O}(1)$
data samples per task and $B=\mathcal{O}(1)$ tasks per iteration. We compare
our theoretical results and key features of the proposed algorithms with other
existing results in Table 1. It is notable that LocalMOML consistently
improves the sample complexity of the existing Per-FedAvg algorithm (Fallah et
al., 2020b) in both cross-silo (without client sampling) and cross-device
(with client sampling) settings.

[b]  Centralized Learning Algorithm Task Sampling Sample Complexity #Data
points ($K$) Per Iteration Strict Assumptions Bounded Gradient Stochastic
Lipschitz (1) MAML (Fallah et al., 2020a) ✗
$\mathcal{O}\left(n\epsilon^{-4}\right)$ $\mathcal{O}(\epsilon^{-2})$ ✗ ✗ NASA
(Ghadimi et al., 2020) ✗ $\mathcal{O}(n\epsilon^{-4})$ $\mathcal{O}(1)$ ✓ ✗
BSGD (Hu et al., 2020) ✓ $\mathcal{O}(\epsilon^{-6})$
$\mathcal{O}(\epsilon^{-2})$ ✓(2) ✓ BSpiderBoost (Hu et al., 2020) ✓
$\mathcal{O}\left(n\epsilon^{-2}+\sqrt{n}\epsilon^{-4}\right)$
$\mathcal{O}(\epsilon^{-2})$ ✓ ✓ $\text{MOML}^{\text{v1}}$ (This work) ✓
$\mathcal{O}\left(n\epsilon^{-5}\right)$ $\bm{\mathcal{O}(1)}$ ✓ ✗
$\text{MOML}^{\text{v2}}$ (This work) ✓
$\mathcal{O}\left(n\epsilon^{-5}\right)$ $\bm{\mathcal{O}(1)}$ ✗ ✗
Personalized Federated Learning Algorithm Client Sampling Sample Complexity
Communication Complexity Avg. #Data points ($K$) Per Iteration Per-FedAvg
(Fallah et al., 2020b) ✗ $\mathcal{O}\left(n\epsilon^{-6}\right)$
$\mathcal{O}(\epsilon^{-3})$ $\mathcal{O}(\epsilon^{-2})$ Per-FedAvg (This
work) ✓(3) $\mathcal{O}(\epsilon^{-7})$ $\mathcal{O}(\epsilon^{-4})$
$\mathcal{O}(\epsilon^{-2})$ LocalMOML (This work) ✗
$\bm{\mathcal{O}\left(n\epsilon^{-5}\right)}$ $\mathcal{O}(\epsilon^{-3})$
$\bm{\mathcal{O}(1)}$ LocalMOML (This work) ✓
$\bm{\mathcal{O}\left(\epsilon^{-6}\right)}$ $\mathcal{O}(\epsilon^{-4})$
$\bm{\mathcal{O}(1)}$

  * (1)

Stochastic Lipschitz: $\nabla\ell(\cdot;\mathbf{z})$ is Lipschitz continuous
for each $\mathbf{z}$, which is stronger than the Lipschitzness of
$\mathcal{L}_{i}(\cdot)$ in Assumption 1.

  * (2)

BSGD can obtain the same rate without the bounded gradient assumption if the
weak convexity of $F(\cdot)$ is additionally assumed.

  * (3)

The analysis of Per-FedAvg with client sampling in Fallah et al. (2020b) seems
to be problematic. See Section E.3 for details.

Table 1: Comparison of our proposed algorithms with existing approaches. The
$\epsilon$ denotes the accuracy level for an $\epsilon$-stationary point
$\mathbb{E}\left[\|\nabla F(\mathbf{w})\|\right]\leq\epsilon$.

##  2 Related Work

The gradient-based MAML was introduced in Finn et al. (2017) and has become a
popular algorithm to learn from prior experience, which has many applications
in supervised learning, reinforcement learning, etc. Later on, some works dive
deeper into MAML to better understand its practical performance and provide
some tricks of the trade for further improving its practicability (Antoniou et
al., 2018; Raghu et al., 2019; Behl et al., 2019). The vanilla MAML has been
generalized from different perspectives. For example, probabilistic MAML is
introduced in Finn et al. (2018) to model a distribution over prior model
parameters. Algorithms with multi-step gradient descent (Ji et al., 2020c) and
with adapting partial parameters (Ji et al., 2020a) have been considered.
Rajeswaran et al. (2019) propose meta-learning with implicit gradients by
formulating the problem as bilevel optimization.

In the past few years, the computational/optimization issues of MAML and its
variants have been studied. In the convex online learning scheme, provable
guarantees of a generalized framework of gradient-based MAML is provided in
Balcan et al. (2019). When $\mathcal{L}_{i}$ is convex, the global convergence
of MAML is established for meta-supervised learning and meta-reinforcement
learning in Wang et al. (2020). When $\mathcal{L}_{i}$ is nonconvex, the
convergence to stationary points of MAML and its first-order and Hessian-free
variants is proved in Fallah et al. (2020a). The BSGD algorithm (Hu et al.,
2020) is shown to have theoretical and practical advantages compared to the
result in (Fallah et al., 2020a) at the cost of stronger asssumptions.
However, these results are still not satisfactory for MAML as discussed
before. Rajeswaran et al. (2019) propose the implicit MAML algorithm (iMAML)
for solving the formulated bilevel optimization problem. The convergence to
stationary points of iMAML in the nonconvex setting is also provided. However,
their theory requires processing all tasks at each iteration. The Hessian-free
alternatives of MAML are also proposed to improve the computational efficiency
(Finn et al., 2017; Nichol et al., 2018; Zhou et al., 2019; Song et al., 2020;
Fallah et al., 2020a). Although we do not incorporate the Hessian-free updates
into the proposed algorithms, it can be done by the techniques in Fallah et
al. (2020a).

Recently, the connection between meta-learning and personalized federated
learning (FL) has been noticed, since both tasks in meta-learning and clients
in federated learning are heterogeneous. The classical federated averaging
(FedAvg) algorithm in FL is connected to MAML in Jiang et al. (2019).
Moreover, they derive a heuristic-based algorithm that runs FedAvg for several
iterations and switch to a meta-learning approach for fine-tuning, which leads
to a solid initial model for any client and the personalized performance on a
client is also improved in practice even if the local data is limited. The
first convergence theory of federated variant of MAML is built in Fallah et
al. (2020b). They also show that FedAvg leads to superior personalized
federated learning performance compared to FedAvg on some numerical
experiments. Personalized federated learning similar to but not exactly the
same as MAML has also considered in several recent works (Hanzely and
Richtárik, 2020; T. Dinh et al., 2020). However, this work will focus on the
federated learning for the vanilla MAML formulation similar to Fallah et al.
(2020b).

Finally, we notice that using a memory buffer for tracking each task has been
also explored in other continual learning paradigms to overcome the
catastrophic forgetting of a learning agent, for example, memory-based
lifelong learning (Lopez-Paz and Ranzato, 2017; Kirkpatrick et al., 2016; Guo
et al., 2020). However, we emphasize that the memory used in our algorithms
for MAML and that used in lifelong learning are naturally different, where the
former is used for tracking individual models of different tasks and the
latter is used for storing some training data of different tasks.

##  3 Preliminaries

In this section, we present the notation, assumptions, and the key challenges
in solving (1).

###  3.1 Notation

Let $\|\cdot\|$ denote the Euclidean norm of a vector and the spectral norm of
a matrix. We use calligraphic font and capital characters, for example,
$\mathcal{B},\mathcal{S}$, to denote a set. For a data distribution
$\mathbf{D}$, we let $\mathcal{S}\sim\mathbf{D}$ denote a set of i.i.d.
samples following the distribution $\mathbf{D}$. $\mathbb{I}()$ denotes an
indicator function. We denote the unbiased stochastic gradient and stochastic
Hessian of the risk function $\mathcal{L}_{i}$ based on a random set
$\mathcal{S}\sim\mathbf{D}_{i}$ of size $K$ by
$\widehat{\nabla}_{\mathcal{S}}\mathcal{L}_{i}(\mathbf{w})=\frac{1}{K}\sum_{\mathbf{z}_{i}\in\mathcal{S}}\nabla\ell_{i}(\mathbf{w};\mathbf{z}_{i})$,
and
$\widehat{\nabla}_{\mathcal{S}}^{2}\mathcal{L}_{i}(\mathbf{w})=\frac{1}{K}\sum_{\mathbf{z}_{i}\in\mathcal{S}}\nabla^{2}\ell_{i}(\mathbf{w};\mathbf{z}_{i})$,
respectively. Refer to Table 4 for a complete list of notations we use.

###  3.2 Assumptions

Throughout the paper, we assume that 1, 2, and 3 are satisfied, which are
standard in the literature (Fallah et al., 2020a; Ji et al., 2020c; Rajeswaran
et al., 2019).

###### Assumption 1

$\mathcal{L}_{i}(\cdot)$ has $L$-Lipschitz continuous gradient and
$\rho$-Lipschitz continuous Hessian, that is,
$\|\nabla\mathcal{L}_{i}(\mathbf{w})-\nabla\mathcal{L}_{i}(\mathbf{w}^{\prime})\|\leq
L\|\mathbf{w}-\mathbf{w}^{\prime}\|$, and
$\|\nabla^{2}\mathcal{L}_{i}(\mathbf{w})-\nabla^{2}\mathcal{L}_{i}(\mathbf{w}^{\prime})\|\leq\rho\|\mathbf{w}-\mathbf{w}^{\prime}\|$
for any $\mathbf{w},\mathbf{w}^{\prime}\in\mathbb{R}^{d}$.

###### Assumption 2

The variance of stochastic gradient $\nabla\ell(\mathbf{w},\mathbf{z})$ and
stochastic Hessian $\nabla^{2}\ell(\mathbf{w},\mathbf{z})$ are upper bounded:

$\displaystyle\mathbb{E}_{\mathbf{z}\sim\mathbf{D}_{i}}\bigg{[}{\left\|\nabla\ell(\mathbf{w};\mathbf{z})-\nabla\mathcal{L}_{i}(\mathbf{w})\right\|^{2}}\bigg{]}\leq\sigma_{G}^{2},\quad\mathbb{E}_{\mathbf{z}\sim\mathbf{D}_{i}}\bigg{[}{\left\|\nabla^{2}\ell(\mathbf{w};\mathbf{z})-\nabla^{2}\mathcal{L}_{i}(\mathbf{w})\right\|^{2}}\bigg{]}\leq\sigma_{H}^{2}.$

###### Assumption 3

$F$ is bounded below,
$\min_{\mathbf{w}\in\mathbb{R}^{d}}F(\mathbf{w})>-\infty$.

Most of the existing results to solve (1) in the literature (Finn et al.,
2017; Rajeswaran et al., 2019; Wang et al., 2017; Ghadimi et al., 2020; Chen
et al., 2020; Fallah et al., 2020b) are under the bounded gradient assumption
described below.

###### Assumption 4

There exists $G>0$, $\left\|\nabla\mathcal{L}_{i}(\mathbf{w})\right\|\leq G$
for any $\mathbf{w}\in\mathbb{R}^{d}$.

Instead of 4, Fallah et al. (2020a) establish the convergence theory of an
MAML variant based on 5, where the gradients are not necessarily bounded.

###### Assumption 5

There exists $\gamma_{G}\geq 0$,
$\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla\mathcal{L}_{i}(\mathbf{w})-\nabla\mathcal{L}(\mathbf{w})\right\|^{2}\leq\gamma_{G}^{2}$
for all $\mathbf{w}\in\mathbb{R}^{d}$ and
$\mathcal{L}(\mathbf{w})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}_{i}(\mathbf{w})$.

###  3.3 Main Challenges

A key to the design of stochastic optimization of (1) is to estimate the
gradient of the objective $\nabla
F(\mathbf{w})=\frac{1}{n}\sum_{i=1}^{n}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}))\nabla\mathcal{L}_{i}(\mathbf{w}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}))$
based on random samples. A straightforward approach adopted by existing
algorithms (Fallah et al., 2020a; Hu et al., 2020) is to use mini-batch
averaging

$\displaystyle\widehat{\Delta}_{\mathcal{B}}=\frac{1}{B}\sum_{i\in\mathcal{B}}(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{w}-\alpha\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w})),$
(2)

where $\mathcal{B}$ denotes the set of $B$ sampled tasks,
$\mathcal{S}_{1}^{i},\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}$ denote three
independent sample sets of size $K$ for each sampled task $\mathcal{T}_{i}$.
However, this naïve approach could suffer from a large optimization error when
$K$ is not large enough.

Another popular approach to tackle the problem (1) is by solving it as an
instance of stochastic two-level compositional problems in the form of
$\mathbb{E}_{\xi}[f_{\xi}(\mathbb{E}_{\xi^{\prime}}[\mathbf{g}(\mathbf{w};\xi^{\prime})];\xi)]$,
where
$\xi^{\prime}=(\mathbf{z}^{\prime}_{1},\ldots,\mathbf{z}^{\prime}_{N})\sim\mathbf{D}_{1}\times\ldots\times\mathbf{D}_{n}$,
$\xi=(\mathbf{z}_{1},\ldots,\mathbf{z}_{n})\sim\mathbf{D}_{1}\times\ldots\times\mathbf{D}_{n}$,
$\mathbf{g}(\mathbf{w};\xi^{\prime})=[\mathbf{w}-\alpha\nabla\ell_{1}(\mathbf{w};\mathbf{z}^{\prime}_{1});\cdots;\mathbf{w}-\alpha\nabla\ell_{n}(\mathbf{w};\mathbf{z}^{\prime}_{n})]\in\mathbb{R}^{nd}$,
and
$f_{\xi}(\mathbf{g})=\frac{1}{n}\sum_{i=1}^{n}\ell_{i}([\mathbf{g}]_{i},\mathbf{z}_{i})$.
However, the existing approaches for two-level compositional problems applied
to solving MAML would require passing through all $n$ tasks at each iteration
(Wang et al., 2017; Ghadimi et al., 2020), which is not suitable for continual
learning with only a small proportion of tasks available every iteration.
Moreover, these works assume both $\nabla f(\cdot)$ and
$\nabla\mathbf{g}(\cdot)$ are bounded.

Thus, the first challenge is how to design an algorithm that provably
converges with subsampled tasks and a constant number of data points. To
tackle this challenge, we borrow the idea in Wang et al. (2017) that keeps
track of the sequence
$\mathbf{v}_{i}(\mathbf{w}_{t})=\mathbf{w}_{t}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}_{t})$
with a variance-reduced estimator $\mathbf{u}^{i}$ for each task
$\mathcal{T}_{i}$ (also known as the personalized model). Compared to the
prior work (Wang et al., 2017), the first novelty of our work lies in the
support of task sampling, that is to say, the algorithm only needs to sample
data and compute the stochastic gradients for those sampled $B$ tasks, instead
of all $n$ tasks.

Moreover, it is even more challenging to establish similar convergence
guarantees without the bounded gradient assumption (4). As shown in Lemma 12,
the gradient-Lipschitz parameter of the meta-objective $F(\mathbf{w})$ is
$L(\mathbf{w})\coloneqq
4L+\frac{2\rho\alpha}{n}\sum_{i=1}^{n}\left\|\nabla\mathcal{L}_{i}(\mathbf{w})\right\|$.
To handle the unbounded gradients, Fallah et al. (2020a) estimate the
gradient-Lipschitz parameter $L(\mathbf{w}_{t})$ by a stochastic estimator
$\widehat{L}(\mathbf{w}_{t})\coloneqq
4L+\frac{2\rho\alpha}{|\mathcal{B}_{L_{t}}|}\sum_{i\in\mathcal{B}_{L_{t}}}\left\|\widehat{\nabla}_{\mathcal{S}_{L_{t}}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t})\right\|$
and set the stepsize $\eta_{t}$ to be inversely proportional to
$\widehat{L}(\mathbf{w}_{t})$ . However, MAML with that stepsize $\eta_{t}$
still requires $K=\mathcal{O}(1/\epsilon^{2})$ data samples per task in each
iteration to ensure the convergence. Thus, it was still an open problem
whether the variance reduction technique can be extended to this setting for
getting rid of the unrealistic requirement of large batch size.

##  4 Memory-based MAML (MOML) in the Centralized Learning Setting

We tackle the challenges mentioned in Section 3.3 by proposing the MOML
algorithm.

###  4.1 Algorithm Outline

The proposed $\text{MOML}^{\text{v1}}$ (Algorithm 1) updates the personalized
model of a sampled task $\mathcal{T}_{i}$, $i\in\mathcal{B}_{t}$ by a momentum
step while those of the other tasks $i\not\in\mathcal{B}_{t}$ are untouched,
that is,

$\displaystyle\mathbf{u}_{t+1}^{i}=\begin{cases}(1-\beta_{t})\mathbf{u}_{t}^{i}+\beta_{t}\widehat{\bf
v}_{t}^{i}&i\in\mathcal{B}_{t}\\\
\mathbf{u}_{t}^{i}&i\not\in\mathcal{B}_{t}\end{cases},$ (v1)

where $\beta\in(0,1]$ is the momentum factor. It is worth noting that
$\text{MOML}^{\text{v1}}$ with $\beta=1$ recovers the original MAML algorithm.
Based on the updated personalized models $\mathbf{u}_{t+1}^{i}$, we can
compute the stochastic gradient by

$\displaystyle\widehat{\Delta}_{\mathcal{B}_{t}}=\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i}).$
(3)

Algorithm 1 $\text{MOML}^{\text{v1}}$

1: for $t=0,1,\ldots,T-1$ do

2: Select a batch of $B$ tasks $\mathcal{B}_{t}$ from $n$ tasks

3: for each task $\mathcal{T}_{i}$, $i\in\mathcal{B}_{t}$ do

4: Select $K$ samples $\mathcal{S}_{1}^{i}\sim\mathbf{D}_{i}$ and compute
$\widehat{\bf
v}_{t}^{i}=\mathbf{w}_{t}-\alpha\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t})$

5: Update the personalized model by
$\mathbf{u}_{t+1}^{i}=(1-\beta_{t})\mathbf{u}_{t}^{i}+\beta_{t}\widehat{\bf
v}_{t}^{i}$.

6: end for

7: Select $K$ samples $\mathcal{S}_{2}^{i}$ and $\mathcal{S}_{3}^{i}$ from
$\mathbf{D}_{i}$ of task $\mathcal{T}_{i}$, $i\in\mathcal{B}_{t}$ and compute
$\widehat{\Delta}_{\mathcal{B}_{t}}$ by (3)

8: Update the meta-model by
$\mathbf{w}_{t+1}=\mathbf{w}_{t}-\eta_{t}\widehat{\Delta}_{\mathcal{B}_{t}}$

9: end for

###  4.2 Convergence Analysis

We establish the convergence guarantees of $\text{MOML}^{\text{v1}}$ based on
Assumption 1, 2, 3 and 4. With these assumptions, the meta-objective
$F(\cdot)$ is $L_{F}$-smooth (see Lemma 12). Based on this fact, we can derive
the lemma below.

###### Lemma 1

If $\alpha\in(0,1/L]$, the iterates $\\{\mathbf{w}_{t}\\}_{t=0}^{T-1}$ of
_$\text{MOML}^{\text{v1}}$_ satisfy that

$\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]$
$\displaystyle\leq\frac{2F(\mathbf{w}_{0})}{\eta T}+\eta
L_{F}C_{\Delta}+\frac{8L^{2}}{BT}\mathbb{E}\left[\sum_{t=0}^{T-1}\sum_{i\in\mathcal{B}_{t}}\left\|\mathbf{v}_{i}(\mathbf{w}_{t})-\mathbf{u}_{t+1}^{i}\right\|^{2}\right].$
(4)

Next, we need to show the error of tracking $\mathbf{v}_{i}(\mathbf{w}_{t})$
(the last term in Equation 4) is vanishing when $\beta\in(0,1)$. Different
from existing analysis of stochastic compositional optimization (Wang et al.,
2017), the estimators $\mathbf{u}^{i}$ for tracking the inner functions
$\mathbf{v}(\mathbf{w})\coloneqq([\mathbf{v}_{i}(\mathbf{w})]^{\top},\dotsc,[\mathbf{v}_{n}(\mathbf{w})]^{\top})^{\top}$
are only partially updated due to the task sampling. Hence, we need a
different technique to bound the error. Given the total number of iterations
$T$, we define $n$ totally ordered sets
$\mathbb{T}_{1},\dotsc,\mathbb{T}_{n}$, where $\mathbb{T}_{i}\subseteq[T]$
contains the iteration indices that the $i$-th task is sampled, that is to
say, $\mathbb{T}_{i}=\\{t^{i}_{1},\ldots,t^{i}_{k},\ldots,\\}$. If task
$\mathcal{T}_{i}$ is sampled in iteration $t$ and the index of $t$ in
$\mathbb{T}_{i}$ is $k$, then $t^{i}_{k}=t$. Hence, we can define a mapping
from $t$ to $k$ for $i\in\mathcal{B}_{t}$. Based on this definition, we can
obtain that

$\displaystyle\sum_{t=0}^{T-1}\sum_{i\in\mathcal{B}_{t}}\left\|\mathbf{v}_{i}(\mathbf{w}_{t})-\mathbf{u}_{t+1}^{i}\right\|^{2}=\sum_{i=1}^{n}\sum_{k=0}^{T_{i}-1}\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{u}_{t_{k}^{i}+1}^{i}\right\|^{2}.$
(5)

Here $T_{i}$ is the cardinality of set $\mathbb{T}_{i}$, which is random and
depends on the task sampling $\\{\mathcal{B}_{t}\\}_{t=0}^{T-1}$. Lemma 2
upper bounds the right hand side of (5).

###### Lemma 2

Suppose that the batch of tasks $\mathcal{B}_{t}$ is sampled uniformly at
random. The error of _$\text{MOML}^{\text{v1}}$_ with
$|\mathcal{S}_{1}^{i}|=K$ to keep track of $\mathbf{v}_{i}(\mathbf{w}_{t})$
can be upper bounded as

$\displaystyle\mathbb{E}\left[\frac{1}{T}\sum_{i=1}^{n}\sum_{k=0}^{T_{i}-1}\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{u}_{t_{k}^{i}+1}^{i}\right\|^{2}\right]\leq\left(\frac{n\sigma_{G}^{2}}{\beta
KT}+\frac{16\eta^{2}n^{2}C_{\Delta}}{\beta^{2}B^{2}}\right)\mathbb{I}[\beta\in(0,1)]+\frac{\beta\alpha^{2}\sigma_{G}^{2}}{K}.$
(6)

###### Remark 3

Note that $\beta<1$ is the key to remove the unrealistic requirement of batch
size $K$. It can be seen that the vanilla _MAML_ algorithm ($\beta=1$)
requires extremely large batch size $K$ for a sufficiently small error of
tracking $\mathbf{v}_{i}(\mathbf{w})$. On the contrary, a constant batch size
$K=\mathcal{O}(1)$ and proper values of $\eta,~{}\beta$ are enough to
eliminate the error for _$\text{MOML}^{\text{v1}}$_.

Next, we are ready to present the main convergence theorem of
$\text{MOML}^{\text{v1}}$.

######  Theorem 4 (Informal)

Under Assumptions 1, 2 3, 4, _$\text{MOML}^{\text{v1}}$_ with stepsizes
$\eta_{t}=\eta=\mathcal{O}(\epsilon^{-3})$,
$\beta_{t}=\beta=\mathcal{O}(\epsilon^{-2})$ and constant batch sizes
$|\mathcal{S}_{1}^{i}|=|\mathcal{S}_{2}^{i}|=|\mathcal{S}_{3}^{i}|=K=\mathcal{O}(1)$,
$|\mathcal{B}_{t}|=B=\mathcal{O}(1)$ can find a stationary point
$\mathbf{w}_{\tau}$ in $T=\mathcal{O}(n\epsilon^{-5})$ iterations.

###### Remark 5

Compared to previous works on solving (1) under the same assumptions, our
_$\text{MOML}^{\text{v1}}$_ can guarantee the convergence without the
requirement of an extremely large batch $K=\mathcal{O}(1/\epsilon^{2})$ (Hu et
al., 2020) or processing all $n$ tasks (Ghadimi et al., 2020).

###  4.3 Handling Unbounded Gradients

We propose another variant of MOML — $\text{MOML}^{\text{v2}}$ in Algorithm 2
that is provably convergent with possibly unbounded gradients and only require
$K=\mathcal{O}(1)$ data samples for each sampled task in one iteration. In
$\text{MOML}^{\text{v2}}$, the personalized model is updated as
$\mathbf{u}^{i}$

$\displaystyle\mathbf{u}_{t+1}^{i}=\begin{cases}(1-\beta_{t})\mathbf{u}_{t}^{i}+\beta_{t}\mathbf{w}_{t}+\frac{\beta_{t}}{p_{i}}(\widehat{\bf
v}_{t}^{i}-\mathbf{w}_{t})&i\in\mathcal{B}^{\prime}_{t}\\\
(1-\beta_{t})\mathbf{u}_{t}^{i}+\beta_{t}\mathbf{w}_{t}&i\not\in\mathcal{B}^{\prime}_{t}\end{cases},$
(v2)

where $\mathcal{B}^{\prime}_{t}$ is independent of $\mathcal{B}_{t}$ and
$p_{i}$ is the probability of selecting task $i$, that is,
$p_{i}=\mathrm{Prob}(i\in\mathcal{B}^{\prime}_{t})$. Besides, we set
$\eta_{t}=\frac{\eta_{0}}{\widehat{L}(\mathbf{w}_{t})}$ and
$\beta_{t}=6L^{2}\eta_{0}^{-1/3}\eta_{t-1}$ for $\text{MOML}^{\text{v2}}$.
Under the same set of assumptions as Fallah et al. (2020a), we can show the
error of tracking the inner function $\mathbf{v}_{i}(\mathbf{w}_{t})$ is
diminishing for $\text{MOML}^{\text{v2}}$ when $\eta_{0}$ is properly chosen.

###### Lemma 6

If $\eta_{0}\leq\min\\{(1/3L)^{3/2},(3L^{2}/C_{3}C_{6})^{3/2}\\}$, we have

$\displaystyle\mathbb{E}\left[\Upsilon_{t+1}\mid\mathcal{F}_{t}\right]\leq\left(1-3L^{2}\eta_{0}^{-\frac{1}{3}}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\right)\mathbb{E}\left[\Upsilon_{t}\mid\mathcal{F}_{t}\right]+\eta_{0}^{\frac{1}{3}}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]C_{9}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta_{0}^{\frac{4}{3}}C_{10},$

where
$\Upsilon_{t}\coloneqq\frac{1}{n}\sum_{i=1}^{n}\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}$
and the constants are defined as $C_{1}\coloneqq\frac{1}{1-2\alpha
L-\alpha^{2}L^{2}}$, $C_{2}\coloneqq\frac{2\alpha L+\alpha^{2}L^{2}}{1-2\alpha
L-\alpha^{2}L^{2}}$, $C_{3}\coloneqq 2L^{2}((1+\alpha
L)^{2}+\alpha^{2}\sigma_{H}^{2}/|\mathcal{S}_{2}^{i}|)$, $C_{4}\coloneqq
4((1+\alpha L)^{2}+\alpha^{2}\sigma_{H}^{2}/|\mathcal{S}_{2}^{i}|)(1+\alpha
L)^{2}C_{1}$, $C_{5}\coloneqq((1+\alpha
L)^{2}+\alpha^{2}\sigma_{H}^{2}/|\mathcal{S}_{2}^{i}|)(\sigma_{G}^{2}/|\mathcal{S}_{3}^{i}|+2(1+\alpha
L)^{2}(2C_{2}^{2}+1)\gamma_{G}^{2})$, $C_{6}\coloneqq(4(1+\alpha
L)^{2}+\alpha^{2}LC_{p})/(3L^{2})$, $C_{7}\coloneqq
18L^{3}\alpha^{2}C_{p}C_{1}^{2}/L^{2}$, $C_{8}\coloneqq
3L^{2}\left(\alpha^{2}C_{p}(2C_{2}^{2}+1)\gamma_{G}^{2}+2\alpha^{2}C_{p}\sigma_{G}^{2}/|\mathcal{S}_{1}^{i}|\right)$,
$C_{9}\coloneqq C_{7}+C_{4}C_{6}$, $C_{10}\coloneqq C_{8}+C_{5}C_{6}/(4L)$,
and $C_{p}\coloneqq\max_{i}1/p_{i}-1$.

######  Theorem 7 (Informal)

Under Assumptions 1, 2, 3, and 5, it is guaranteed that
_$\text{MOML}^{\text{v2}}$_ with
$\eta_{t}=\frac{\eta_{0}}{\widehat{L}(\mathbf{w}_{t})}$,
$\beta_{t}=6L^{2}\eta_{0}^{-1/3}\eta_{t-1}$,
$\eta_{0}=\mathcal{O}(\epsilon^{-3})$ and constant batch sizes
$|\mathcal{S}_{1}^{i}|=|\mathcal{S}_{2}^{i}|=|\mathcal{S}_{3}^{i}|=K=\mathcal{O}(1)$,
$|\mathcal{B}_{t}|=|\mathcal{B}_{t}^{\prime}|=B=\mathcal{O}(1)$ can find an
$\epsilon$-stationary point $\mathbf{w}_{\tau}$ in $T=O(C_{p}\epsilon^{-5})$
iterations, where $C_{p}=\max_{i}1/p_{i}-1$.

Algorithm 2 $\text{MOML}^{\text{v2}}$

1: for $t=0,1,\ldots,T-1$ do

2: Select two mutually independent batches of tasks $\mathcal{B}_{t}$,
$\mathcal{B}_{t}^{\prime}$ from $n$ tasks

3: for each task $\mathcal{T}_{i}$, $i\in\mathcal{B}_{t}^{\prime}$ do

4: Select $K$ samples $\mathcal{S}_{1}^{i}\sim\mathbf{D}_{i}$ and compute
$\widehat{\bf
v}_{t}^{i}=\mathbf{w}_{t}-\alpha\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t})$

5: end for

6: Update the personalized model $\mathbf{u}^{i}$ by
$\displaystyle\mathbf{u}_{t+1}^{i}=\begin{cases}(1-\beta_{t})\mathbf{u}_{t}^{i}+\beta_{t}\mathbf{w}_{t}+\frac{\beta_{t}}{p_{i}}(\widehat{\bf
v}_{t}^{i}-\mathbf{w}_{t})&i\in\mathcal{B}_{t}^{\prime}\\\
(1-\beta_{t})\mathbf{u}_{t}^{i}+\beta_{t}\mathbf{w}_{t}&i\not\in\mathcal{B}_{t}^{\prime}\end{cases}.$

7: Select $K$ samples $\mathcal{S}_{2}^{i}$ and $\mathcal{S}_{3}^{i}$ from
$\mathbf{D}_{i}$ of task $\mathcal{T}_{i}$, $i\in\mathcal{B}_{t}$ and compute
$\widehat{\Delta}_{\mathcal{B}_{t}}$ by (3)

8: Update the meta-model by
$\mathbf{w}_{t+1}=\mathbf{w}_{t}-\eta_{t}\widehat{\Delta}_{\mathcal{B}_{t}}$

9: end for

###### Remark 8

Theorem 7 shows that _$\text{MOML}^{\text{v2}}$_ improves the theory of MAML
(Fallah et al., 2020a) by removing the requirement of sampling
$K=\mathcal{O}(1/\epsilon^{2})$ data samples in each iteration. However, both
the _MAML_ variant in Fallah et al. (2020a) and _$\text{MOML}^{\text{v2}}$_
are more of theoretical interests because they need some possibly unknown
parameters $L,\rho,\alpha$ and extra tasks/samples
$\mathcal{B}_{\mathcal{L}_{t}}$, $\mathcal{S}_{L_{t}}$ to estimate the
gradient-Lipschitz parameter and then compute the stepsize $\eta_{t}$ in every
iteration. As shown in previous work (Ji et al., 2020c) and our experiments,
the gradients $\nabla\mathcal{L}_{i}(\mathbf{w}_{t})$ are well bounded during
the meta-training process such that _$\text{MOML}^{\text{v1}}$_ with a
constant step size is more practical.

##  5 LocalMOML for Personalized Federated Learning

In this section, we present a stochastic algorithm for solving MAML in the
federated learning setting. In particular, we assume that there are $n$
clients and the $i$-th task and its corresponding data is only accessible at
the $i$-th client 111This can be relaxed to a setting where the $n$ tasks are
allocated to $m<n$ clients with non-overlapping.. These clients are only
allowed to aggregate the models instead of exchanging any data. A naïve
implementation of MOML in the federated learning setting would require
aggregating the local gradient estimator $\widehat{\Delta}_{\mathcal{B}}^{i}$
at every iteration or equivalently aggregating the local copies of the meta
model at every iteration. Hence, the communication complexity is as high as
the iteration complexity $T=\mathcal{O}(n/\epsilon^{5})$. Our goal here is to
reduce the communication complexity by proposing communication-efficient
federated learning algorithms.

###  5.1 Algorithm Outline

Our second algorithm named LocalMOML is presented in Algorithm 3. This
algorithm is partially motivated by the voluminous algorithms in federated
learning that use local computations to trade off the communications (Yang,
2013; McMahan et al., 2017; Deng et al., 2020; Karimireddy et al., 2019;
Stich, 2019; Lin et al., 2020; Woodworth et al., 2020; Khaled et al., 2020).
Each client not only maintains and updates its personalized model
$\mathbf{u}^{i}$ but also maintains and updates its local copy of the meta
model denoted by $\mathbf{w}^{i}$. A key feature of LocalMOML is running $H$
local steps on each sampled client before these clients communicate to
aggregate the local meta models.

Algorithm 3 LocalMOML

1: for $r=1,\ldots,R$ do

2: Select a batch of $B$ clients $\mathcal{B}_{r}$.

3: for each client $\mathcal{C}_{i}$, $i\in\mathcal{B}_{r}$ do

4: if client sampling then

5: Select $K_{0}$ samples $\mathcal{S}_{0}^{i}\sim\mathbf{D}_{i}$, reset
$\mathbf{u}_{r,1}^{i}=\mathbf{w}_{r,1}^{i}-\alpha\widehat{\nabla}_{\mathcal{S}_{0}^{i}}\mathcal{L}_{i}(\mathbf{w}_{r,1}^{i})$

6: else

7: Set $\mathbf{u}_{r,1}^{i}=\mathbf{u}_{r-1,H}^{i}$

8: end if

9: for $h=1,\dotsc,H$ do

10: Sample $\mathcal{S}_{1}^{i}$ to update the personalized model
$\mathbf{u}_{r,h}^{i}$ by (7)

11: Select two sets $\mathcal{S}_{2}^{i}$ and $\mathcal{S}_{3}^{i}$ of from
$\mathbf{D}_{i}$ to compute $\widehat{\Delta}_{r,h}^{i}$ by (8), and update
the local model by
$\mathbf{w}_{r,h+1}^{i}=\mathbf{w}_{r,h}^{i}-\eta\widehat{\Delta}_{r,h}^{i}$

12: end for

13: Client $\mathcal{C}_{i}$ sends $\mathbf{w}_{r,H+1}^{i}$ to the server.

14: end for

15: The server aggregates and broadcasts
$\mathbf{w}_{r+1}=\frac{1}{B}\sum_{i\in\mathcal{B}_{r}}\mathbf{w}_{r,H+1}^{i}$

16: end for

We consider both the cross-silo setting and the cross-device setting in the
literature of federated learning. In the cross-device setting only a partial
set of $B<n$ clients are sampled to update their local models at each round,
while in the cross-silo setting all $n$ clients will participate in updating
the model at each round, that is, $B=n$. In these two settings, LocalMOML
needs different ways to initialize the personalized model
$\mathbf{u}_{r,h}^{i}$ at the beginning of each round. In the cross-silo
setting ($B=n$), personalized models are directly copied from the end of
previous round, that is, $\mathbf{u}_{r,1}^{i}=\mathbf{u}_{r-1,H}^{i}$. In the
cross-device setting, ($B<n$), the personalized models for the sampled tasks
are restarted as
$\mathbf{u}_{r,1}^{i}=\mathbf{w}_{r,1}^{i}-\alpha\widehat{\nabla}_{\mathcal{S}_{0}^{i}}\mathcal{L}_{i}(\mathbf{w}_{r,1}^{i})$
for $i\in\mathcal{B}_{r}$. Then, the personalized model $\mathbf{u}_{r,h}^{i}$
for a sampled client $i$ is updated as:

$\displaystyle\mathbf{u}_{r,h}^{i}=(1-\beta)\mathbf{u}_{r,h-1}^{i}+\beta\left(\mathbf{w}_{r,h}^{i}-\alpha\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right).$
(7)

Based on the local personalized model $\mathbf{u}^{i}$ and meta model
$\mathbf{w}^{i}$, the stochastic gradient estimator is computed as:

$\displaystyle\widehat{\Delta}_{r,h}^{i}=(I-\alpha\widehat{\nabla}^{2}_{\mathcal{S}_{2}^{i}}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i}).$
(8)

After all $H$ iterations in each round, the meta model is averaged for
synchronization. It is notable that the Per-FedAvg algorithm (Fallah et al.,
2020b) is a special case of LocalMOML (Algorithm 3) with $\beta=1$ and the
original MAML algorithm is a special case of LocalMOML with $\beta=1$, $H=1$.

Before presenting the convergence results of LocalMOML, it is worth mentioning
that LocalMOML can be also implemented in the centralized learning setting.
Unlike MOML that only works with finite $n$, LocalMOML can also solve the
problem (1) with infinite number of tasks 222In Table E.2 of the appendix, we
present a table that compares the convergence rates of LocalMOML with the
baselines MAML, BSGD, BSpiderBoost in the infinite $n$ (online tasks)
setting.. Instead of maintaining personalized models $\mathbf{u}^{i}$ for all
tasks $i\in[n]$ at all iterations, the personalized model
$\mathbf{u}^{i}_{r,h}$ in LocalMOML for the sampled task $i\in[n]$ can be
erased from the memory after round $r$ and will not be used in any later round
$r^{\prime}>r$, no matter whether that task will be sampled again or not.

###  5.2 Convergence results of LocalMOML

We analyze LocalMOML under the same assumptions as the Per-FedAvg (Fallah et
al., 2020b), that is, Assumptions 1, 2, 3, 4, 5, and the assumption below.

###### Assumption 6

There exists $\gamma_{H}\geq 0$ that:
$\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla^{2}\mathcal{L}_{i}(\mathbf{w})-\frac{1}{n}\sum_{i=1}^{n}\nabla^{2}\mathcal{L}_{i}(\mathbf{w})\right\|^{2}\leq\gamma_{H}^{2}$.

Note that Assumptions 1 and 4 implies Assumptions 5 and 6. However, directly
utilizing Assumptions 5 and 6 in the analysis can explicitly show the impact
of the dissimilarity of data distribution on the final performance (Fallah et
al., 2020b).

###### Theorem 9

$R$ rounds of _LocalMOML_ with the stepsize
$\eta\leq\min\left\\{\frac{C_{4}}{H},\frac{C_{5}\beta}{\mathbb{I}[\beta\in(0,1)]}\right\\}$
leads to

$\displaystyle\frac{1}{R}\sum_{r=1}^{R}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}\right]$
$\displaystyle\leq\frac{4F(\mathbf{w}_{1})}{\eta
T}+32\eta^{2}H(H-1)C_{1}(\hat{\sigma}^{2}+2\gamma_{F}^{2})+\frac{4\eta}{B}\left(\hat{\sigma}^{2}+\frac{4(n-B)}{(n-1)}H\gamma_{F}^{2}\right)+\frac{64C_{3}\beta\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}$
$\displaystyle+\frac{64C_{3}\sigma_{G}^{2}}{\beta\left(\mathbb{I}[B<n]HS_{0}+\mathbb{I}[B=n]HR\right)}\mathbb{I}[\beta\in(0,1)]+\frac{3072C_{3}\eta^{2}(\hat{\sigma}^{2}+2\gamma_{F}^{2})}{\beta^{2}}\mathbb{I}[\beta\in(0,1)],$

where $T=RH$, $C_{1}\coloneqq C_{\rho,G,L}+2L_{F}+\frac{2(n-B)L_{F}}{B(n-1)}$,
$C_{2}\coloneqq 1+\frac{n-B}{B(n-1)}+\frac{C_{1}}{4L_{F}^{2}}$, and
$C_{3}\coloneqq 1+\frac{L^{2}}{4L_{F}}+\frac{L^{2}C_{1}}{8L_{F}^{2}}$,
$C_{4}\coloneqq\frac{1}{4\max\\{L_{F},8C_{2}\\}}$,
$C_{5}\coloneqq\frac{1}{16\max\\{L,\sqrt{96C_{3}}\\}}$, $C_{\rho,G,L}\coloneqq
2G^{2}\rho^{2}/L^{2}+16L^{2}$,
$\hat{\sigma}^{2}\coloneqq\frac{2\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+\frac{2\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}\left(\frac{\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}+L^{2}\right)+\frac{\alpha^{2}G^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}$.

We consider the special case $\alpha=0$, $H=1$, $\beta=1$, $B=n$ where
LocalMOML becomes SGD and Theorem 9 recovers the standard result of SGD:
$\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]\leq\mathcal{O}\left(\frac{F(\mathbf{w}_{1})}{\eta
T}+\eta\sigma_{G}^{2}\right)$. Besides, ther term $\gamma_{F}^{2}\coloneqq
3G^{2}\alpha^{2}\gamma_{H}^{2}+192\gamma_{G}^{2}$ in Theorem 9 explicitly show
the impact of the dissimilarity of data distribution on the final performance.

To better understand the complexities of LocalMOML, we present two corollaries
below corresponding to the cross-silo $B=n$ and cross-device
$B=\mathcal{O}(1)$ settings .

######  Corollary 1 (Cross-device Setting)

In this setting, we set
$\beta=O(\epsilon^{2}),\eta=\mathcal{O}(\epsilon^{4}),H=\mathcal{O}(1/\epsilon^{2})$,
$|\mathcal{S}_{0}^{i}|=K_{0},|\mathcal{S}_{1}^{i}|=|\mathcal{S}_{2}^{i}|=|\mathcal{S}_{3}^{i}|=K$,
$K=O(1),K_{0}=H$. Then, we can conclude that $T=\mathcal{O}(1/\epsilon^{6})$
and $R=\mathcal{O}(1/\epsilon^{4})$. Hence, the average number of data points
per-iteration is $(K_{0}+3HK)/H=\mathcal{O}(1)$, the total sample complexity
is $KBT+RK_{0}=\mathcal{O}(1/\epsilon^{6})$ and the communication complexity
is $R=\mathcal{O}(1/\epsilon^{4})$.

######  Corollary 2 (Cross-silo Setting)

In this setting, we set
$\beta=\mathcal{O}(\epsilon^{2}),\eta=\mathcal{O}(\epsilon^{3}),H=\mathcal{O}(1/\epsilon^{2})$,
$|\mathcal{S}_{0}^{i}|=0,|\mathcal{S}_{1}^{i}|=|\mathcal{S}_{2}^{i}|=|\mathcal{S}_{3}^{i}|=K$,
and $K=\mathcal{O}(1)$. Then, we can conclude that
$T=\mathcal{O}(1/\epsilon^{5})$ and $R=\mathcal{O}(1/\epsilon^{3})$. Hence,
the total sample complexity is $KNT=\mathcal{O}(n/\epsilon^{5})$ and the
communication complexity is $R=\mathcal{O}(1/\epsilon^{3})$.

Thus, in both cross-silo and cross-device settings, our LocalMOML not only
removes the unrealistic requirement of the large batch size
$\mathcal{O}(1/\epsilon^{2})$ in every iteration for the convergence of Per-
FedAvg, but also improves the sample complexity.

##  6 Experiments

We testify our proposed algorithms MOML ($\text{MOML}^{\text{v1}}$) and
LocalMOML on the sinewave regression task in the centralized setting and the
image classification task in the simulated federated learning setting.

###  6.1 Sinewave Regression in the Centralized Setting

First, we compare our proposed MOML and LocalMOML with baselines NASA, MAML
(BSGD) and BSpiderBoost on the sinusoid regressiom problem (Finn et al., 2017)
in the centralized setting. We generate 25 tasks for training in total, each
of which is to fit the function $f(x)=A\sin(\phi+x)$, where
$A=\\{1,2,3,4,5\\}$, $\phi=\frac{i\pi}{5}$, $i=1,2,3,4,5$. Similar to Finn et
al. (2017), we choose the feedforward neural network with ReLU nonlinearities
and 2 hidden layers of size 40 as the model and the mean-square error as the
loss function. The training and validation data of each task are randomly
sampled in an online manner $x\sim\text{Unif}(-5,5)$. MAML (BSGD) and MOML
sample $B=3$ out of $n$ tasks, NASA uses all $n$ tasks, and BSpiderBoost uses
all $n$ tasks or $\sqrt{n}$ tasks following their theory. We consider two
possible minibatch sizes of data points: $K=1$ and $K=5$ in one iteration. The
meta-learned model is adapted to 5 randomly sampled unseen task thats
$A\sim\text{Unif}(1,5)$, $\phi\sim\text{Unif}(\frac{\pi}{5},\pi)$,
$x\in[-5,5]$ after 10 steps of gradient descent with learning rate 0.01 and 10
test data points (in other words, 10-shots). After the adaptation, we evaluate
the final test error on another 100 data points from the unseen task. The
inner step size $\alpha$ is set to 0.01 for all algorithms. The outer step
size $\eta$ is decayed 10 times at 75% of the total iterations and the initial
step sizes are tuned for the algorithms separately by grid search in
$\\{0.001,0.005,0.01,0.05,0.1\\}$. We also tune $\beta$ for MOML and LocalMOML
and it turns out that $\beta=0.5$ and $\beta=0.7$ work reasonably well for
$K=1$ and $K=5$, respectively. For LocalMOML, we set the size of the initial
number of samples $K_{0}$ of each round to be 2 times of $K$ and $H=5$.

$K=1$
---
Metrics | MAML (BSGD) | NASA | BSpiderBoost | MOML | LocalMOML
Avg. Test Error | 0.651 $\pm$ 0.123 | 0.622 $\pm$ 0.032 | 1.044 $\pm$ 0.040 | 0.310 $\pm$ 0.025 | 0.381 $\pm$ 0.059
Avg. Time Per Iteration (ms) |  2.083 $\pm$ 0.024 | 20.302 $\pm$ 0.139 | 7.633 $\pm$ 0.036 | 2.732 $\pm$ 0.032 | 3.182 $\pm$ 0.026
$K=5$
Metrics | MAML (BSGD) | NASA | BSpiderBoost | MOML | LocalMOML
Avg. Test Error | 0.188 $\pm$ 0.016 | 0.367 $\pm$ 0.003 | 0.760 $\pm$ 0.061 | 0.128 $\pm$ 0.011 | 0.133 $\pm$ 0.020
Avg. Time Per Iteration (ms) | 2.449 $\pm$ 0.237 | 22.438 $\pm$ 0.057 | 8.145 $\pm$ 0.307 | 3.170 $\pm$ 0.150 | 3.649 $\pm$ 0.181

Table 2: Comparison of final test error for the Sinewave regression task in
the centralized setting. For each metric, the best one is highlighted in black
and the second-best one is highlighted in gray.

Figure 1: Left two: training errors curves; Right two: the fitted curves on
the unseen task.

We plot the training errors across the iterations in Figure 1 (left two), and
the fitted sinusoids on the unseen task in Figure 1 (right two) for different
settings, and report the final test errors and wall-clock time per iteration
in Table 2. The results are averaged over 5 trials with different random
seeds.

From the results, we can observe that NASA and BSpiderBoost might not be good
solvers for the test problem: 1) Since the NASA algorithm does not support
task sampling, the forward/backward propagations are needed on all $n$ tasks
(instead of only 3 sampled tasks for MAML/MOML/LocalMOML), which leads to
longer per-iteration running time (as shown in the Table). Moreover, NASA
requires $n/3$ times more data samples per iteration compared to
MAML/MOML/LocalMOML given the same batch size $K$; 2) The performance of
BSpiderBoost in our experiment is poor may be due to the small batch size $K$,
considering its theory requires $K=\mathcal{O}(1/\epsilon^{2})$ to converge,
which is also observed in their original paper (Hu et al., 2020). As a
comparison, MOML and LocalMOML lead to superior performance, especially when
fewer train samples per iteration, which is in accordance with our theory.

###  6.2 Image Classification in a (Simulated) Federated Learning Setting.

Second, we compare our LocalMOML with Per-FedAvg (Fallah et al., 2020b) on the
image classification problem. We consider three data sets, MNIST, CIFAR-10,
CIFAR-100. In order to create heterogeous data distribution, we follow the
setup in Fallah et al. (2020b). In particular, For MNIST and CIFAR-10 (10
classes), we distribute the training data between $N=50$ clients (tasks) as
follows: (i) half of the clients, each have $a$ images of each of the first
five classes; (ii) The rest half clients, each have $a/2$ images from only one
of the first five classes and $2a$ images from only one of the other five
classes. For CIFAR-100, we consider the 20 superclasses and distribute the
data similarly by dividing them into first 10 classes and other ten classes to
run the same procedure. Similarly, we divide the test data over the clients
with the same distribution as the one for the training data. We set $a=68$ for
constructing the distributed training sets of MNIST, CIFAR-10, CIFAR-100, and
set $a=34$ for constructing the test sets of MNIST and CIFAR-10 and $a=15$ for
constructing the test sets of CIFAR-100.

|  | $H=4$ | $H=10$
---|---|---|---
|  # Sampled Tasks ($B^{\prime}$) per-GPU | Per-FedAvg | LocalMOML | Per-FedAvg | LocalMOML
MNIST | 1 | 90.97$\pm$0.75 | 91.02$\pm$0.70 | 93.03$\pm$0.56 | 93.04$\pm$0.50
5 | 91.18$\pm$1.01 | 91.27$\pm$0.85 | 93.31$\pm$0.60 | 93.33$\pm$0.60
CIFAR-10 | 1 | 65.57$\pm$1.06 | 67.03$\pm$1.20 | 66.89$\pm$2.14 | 68.58$\pm$1.20
5 | 65.69$\pm$1.96 | 67.24$\pm$1.38 | 67.43$\pm$1.48 | 68.68$\pm$1.53
CIFAR-100 | 1 | 51.79$\pm$2.17 | 53.64$\pm$0.51 | 52.50$\pm$1.31 | 55.71$\pm$1.61
5 | 52.19$\pm$2.10 | 54.05$\pm$1.15 | 53.48$\pm$1.47 | 55.95$\pm$1.39

Table 3: Comparison of final test accuracy (percentage) for 5-shot learning on
three image classification data sets in a cross-device federated learning
setting. The results as averged over three runs with different random seeds.

We conduct experiments on four GPUs to mimic the cross-device federated
learning setting, where all 50 tasks are distributed to the four GPUs roughly
evenly. At each round, all four GPUs participate the learning but each GPU
only samples a batch of $B^{\prime}$ tasks from its owned tasks, and we
consider two settings and $B^{\prime}=1,B^{\prime}=5$. It means every round a
total of $B=4B^{\prime}$ tasks are sampled for updating. We train a neural
network with 2 hidden layers with each layer having 40 neurons, and use the
relu activation function. We use $\alpha=0.001$ and the step size for the
considered algorithms is tuned in a range similar as before. For Per-FedAvg
and LocalMOML, we consider two settings of $H=4$ and $H=10$. The mini-batch
size at every iteration (including the initial one at each round) is set to
$5$, that is, $K=5,K_{0}=5$. We tune the $\beta$ in a range $[0.1,0.9]$, and
run a total of $5000$ iterations. We evaluate the test accuracy after $5$-shot
learning with 10 steps of fine-tuning on the test set. The final results are
reported in Table 3. We can see that the proposed LocalMOML outperforms Per-
FedAvg in almost all settings with substantial improvements on the most
difficult CIFAR-100 data.

##  7 Conclusions and Discussion

In this paper, we have considered stochastic optimization for model-agnostic
meta-learning and proposed two algorithms in the centralized setting and the
federated learning setting. Our algorithms improve the existing algorithms for
model-agnostic meta-learning in several perspectives. In particular, our
convergence theory guarantees that proposed algorithms converge to a
stationary solution by sampling the constant number of tasks and a constant
number of samples per iteration. Our federated learning algorithms not only
improve the computational complexity but also the communication complexity
over an existing federated learning algorithm for solving the same
formulation. One limitation of the proposed MOML algorithm is that they need
to maintain an individualized model for each task, which makes it not
applicable to the problem (1) with a large/infinite number of tasks or
embedded systems with small memory for learning large models. However, our
second algorithm LocalMOML does not suffer from the same issue.

Basic
---
$n$ | total number of tasks/clients |
$\mathcal{T}_{i}$ | the $i$-th task |
$\mathbf{w}$ | the trainable parameter |
$\ell_{i}(\cdot,\mathbf{z})$ | loss function with the data sample $\mathbf{z}$ of $i$-th task |
$\mathcal{L}_{i}(\cdot)$ | risk function for the $i$-th task |
$\mathbf{D}_{i}$ | distribution of data for the $i$-th task |
$F_{i}(\cdot)$ | the meta-objective of task $i$ | Lemma 10
$F(\cdot)$ | the total meta-objective | (1)
$\epsilon$ | target accuracy for an approximate stationary point | Table 1
$L_{i}$, $\rho_{i}$ | Lipschitz constants of gradient $\nabla\mathcal{L}_{i}(\cdot)$ and Hessian $\nabla^{2}\mathcal{L}_{i}(\cdot)$ | 1
$\sigma_{G}^{2}$, $\sigma_{H}^{2}$ |  variance upper bounds of stochastic gradient $\nabla\ell(\cdot,\mathbf{z})$ and stochastic Hessian $\nabla^{2}\ell(\cdot,\mathbf{z})$ | 2
$\gamma_{G}^{2}$, $\gamma_{H}^{2}$ |  heterogeneity constants of gradients $\nabla\mathcal{L}_{i}(\cdot)$ and Hessian $\nabla^{2}\mathcal{L}_{i}(\cdot)$ among the tasks/clients |  5 6
$G$ | upper bound of the norm of gradient $\nabla\mathcal{L}_{i}(\cdot)$ | 4
Centralized Learning
$B,K$ | batch size of tasks and batch size of data samples per batch |
$\mathcal{B}$ | size-$B$ batch of tasks |
$\mathcal{T}_{i}$ | the $i$-th task |
$\mathcal{S}_{1}^{i},\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}$ | independent batches of data points of task $i$ |
$\widehat{\nabla}_{\mathcal{S}}\mathcal{L}_{i}(\cdot)$ | stochastic gradient of $\mathcal{L}_{i}(\cdot)$ estimated on batch $\mathcal{S}$ |
$\widehat{\Delta}_{\mathcal{B}}$ |  stochastic estimator of meta-gradient estimated on a batch of tasks $\mathcal{B}$ | (2)
$\mathbf{u}^{i}$ | personalized model for task $i$ |
$\mathbf{v}_{i}(\mathbf{w})$ |  updated model by one step of gradient descent $\mathbf{v}_{i}(\mathbf{w})\coloneqq\mathbf{w}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w})$ |
$\widehat{\bf v}^{i}$ |  updated model by one step of stochastic gradient descent $\widehat{\bf v}^{i}\coloneqq\mathbf{w}-\alpha\widehat{\nabla}\mathcal{L}_{i}(\mathbf{w})$ |
$\eta_{t}$ | the step size in iteration $t$ |
$\beta_{t}$ | the momentum factor in iteration $t$ |
$p_{i}$ | the probability of sampling task $i$ in $\mathcal{B}_{t}^{\prime}$ |
$C_{p}$ | the constant $\max_{i}1/p_{i}-1$ |
Federated Learning
$\mathcal{C}_{i}$ | the $i$-th client |
$R,H$ |  total number of communication rounds and total number of iterations in each round |
$\Delta\mathbf{w}_{r}$ |  the average of stochastic meta-gradients over the sampled clients $i\in\mathcal{B}_{r}$ and local steps of round $r$ |
$\tilde{\eta}$ |  the effective stepsize per round, i.e., $\tilde{\eta}\coloneqq\eta H$ |
Table 4: Notations we use throughout the paper.

##  Appendix A Preliminary Lemmas

######  Lemma 10 (Lemma A.3 in Fallah et al. (2020a))

If $\alpha\in[0,\frac{\sqrt{2}-1}{L})$ and Assumptions 1, 2, and 5 are
satisfied. Then, for any $\mathbf{w}\in\mathbb{R}^{d}$, we have

$\displaystyle\left\|\nabla\mathcal{L}(\mathbf{w})\right\|^{2}\leq
C_{1}^{2}\left\|\nabla F(\mathbf{w})\right\|+C_{2}^{2}\gamma_{G}^{2},$ (9)
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla
F_{i}(\mathbf{w})\right\|^{2}\leq 2(1+\alpha L)^{2}C_{1}^{2}\left\|\nabla
F(\mathbf{w})\right\|^{2}+(1+\alpha L)^{2}(2C_{2}^{2}+1)\gamma_{G}^{2},$ (10)

where $L\coloneqq\max_{i}L_{i}$,
$\nabla\mathcal{L}(\mathbf{w})\coloneqq\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}_{i}(\mathbf{w})$,
$F_{i}(\mathbf{w})\coloneqq\mathcal{L}_{i}(\mathbf{w}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}))$,
$F(\mathbf{w})\coloneqq\frac{1}{n}\sum_{i=1}^{n}F_{i}(\mathbf{w})$,
$C_{1}\coloneqq\frac{1}{1-2\alpha L-\alpha^{2}L^{2}}$, and
$C_{2}\coloneqq\frac{2\alpha L+\alpha^{2}L^{2}}{1-2\alpha L-\alpha^{2}L^{2}}$.

###### Lemma 11

Assume that $X=\frac{1}{n}\sum_{i=1}^{n}X_{i}$. If we sample a size-$B$
minibatch $\mathcal{B}$ from $\\{1,\dotsc,n\\}$ uniformly at random, we have
$\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}}X_{i}\right]=X$ and

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}}X_{i}\right\|^{2}\right]\leq\frac{n-B}{B(n-1)}\frac{1}{n}\sum_{i=1}^{n}\left\|X_{i}\right\|^{2}+\frac{n(B-1)}{B(n-1)}\left\|X\right\|^{2}.$
(11)

Proof  First, we define random variables $\xi_{i}$ and $\xi_{ii^{\prime}}$ as

$\displaystyle\xi_{i}=\begin{cases}1&i\in\mathcal{B}\\\
0&i\notin\mathcal{B}\end{cases},\quad\xi_{ii^{\prime}}=\begin{cases}1&i\in\mathcal{B}~{}\text{and}~{}i^{\prime}\in\mathcal{B}\\\
0&\text{otherwise}\end{cases}$

Note that $\xi_{ii^{\prime}}=\xi_{i}\xi_{i^{\prime}}$,
$\mathbb{E}\left[\xi_{i}\right]=\Pr(i\in\mathcal{B})=\frac{B}{n}$ and
$\mathbb{E}\left[\xi_{ii^{\prime}}\right]=\Pr(i\in\mathcal{B},i^{\prime}\in\mathcal{B})=\frac{B(B-1)}{n(n-1)}$.
Thus, it is clear that

$\displaystyle\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}}X_{i}\right]=\mathbb{E}\left[\frac{1}{B}\sum_{i=1}^{n}\xi_{i}X_{i}\right]=\frac{1}{B}\sum_{i=1}^{n}\mathbb{E}\left[\xi_{i}\right]X_{i}=\frac{1}{B}\sum_{i=1}^{n}\Pr(i\in\mathcal{B})X_{i}=\frac{1}{B}\sum_{i=1}^{n}\frac{B}{n}X_{i}=X.$

Moreover,

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}}X_{i}\right\|^{2}\right]$
$\displaystyle=\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i=1}^{n}\xi_{i}X_{i}\right\|^{2}\right]=\frac{1}{B^{2}}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\xi_{i}X_{i}\right\|^{2}\right]+\frac{1}{B^{2}}\sum_{i\neq
i^{\prime}}\mathbb{E}[\xi_{i}\xi_{i^{\prime}}]\langle
X_{i},X_{i^{\prime}}\rangle$
$\displaystyle=\frac{1}{B^{2}}\sum_{i=1}^{n}\frac{B}{n}\left\|X_{i}\right\|^{2}+\frac{1}{B^{2}}\sum_{i\neq
i^{\prime}}\frac{B(B-1)}{n(n-1)}\langle X_{i},X_{i}^{\prime}\rangle$
$\displaystyle=\frac{n-B}{B(n-1)}\frac{1}{n}\sum_{i=1}^{n}\left\|X_{i}\right\|^{2}+\frac{n(B-1)}{B(n-1)}\left\|X\right\|^{2}.$

######  Lemma 12 (Corollary A.1 of Fallah et al. 2020a)

If $\alpha\in[0,1/L]$, for any
$\mathbf{w},\mathbf{w}^{\prime}\in\mathbb{R}^{d}$,

$\displaystyle F(\mathbf{w}^{\prime})-F(\mathbf{w})-\langle\nabla
F(\mathbf{w}),\mathbf{w}^{\prime}-\mathbf{w}\rangle\leq\frac{L(\mathbf{w})}{2}\left\|\mathbf{w}^{\prime}-\mathbf{w}\right\|^{2},$
(12)

where $L(\mathbf{w})\coloneqq
4L+\frac{2\rho\alpha}{n}\sum_{i=1}^{n}\left\|\nabla\mathcal{L}_{i}(\mathbf{w})\right\|$.
If $\mathcal{L}_{i}(\cdot)$ is $G$-Lipschitz-continuous, that is,
$\left\|\nabla\mathcal{L}_{i}(\mathbf{w})\right\|\leq G$ for any
$\mathbf{w}\in\mathbb{R}^{d}$, then $F(\cdot)$ is $L_{F}$-smooth with
$L_{F}\coloneqq 4L+2\rho\alpha G$.

######  Lemma 13 (Lemma 4.4 of Fallah et al. 2020b)

If $\alpha\in(0,1/L]$ and Assumptions 5, 4, and 6 are satisfied, we have

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla F_{i}(\mathbf{w})-\nabla
F(\mathbf{w})\right\|^{2}\leq\gamma_{F}^{2}\coloneqq
3G^{2}\alpha^{2}\gamma_{H}^{2}+192\gamma_{G}^{2}.$ (13)

###### Lemma 14

For $\widehat{L}(\mathbf{w})\coloneqq
4L+\frac{2\rho\alpha}{B}\sum_{i\in\mathcal{B}}\left\|\widehat{\nabla}_{\mathcal{S}}\mathcal{L}_{i}(\mathbf{w})\right\|$
and $\frac{(n-1)B}{n-B}\geq\lceil(8\rho\alpha\gamma_{G}/L)^{2}\rceil$,
$S\geq\lceil(8\rho\alpha\sigma_{G}/L)^{2}\rceil$, we have:

$\displaystyle\frac{4}{5L(\mathbf{w})}\leq\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})}\right]\leq\frac{11\eta_{0}}{8L(\mathbf{w})}\leq\frac{\eta_{0}}{4L},\quad\quad\frac{L(\mathbf{w})}{2}\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})^{2}}\right]\leq\frac{2}{L(\mathbf{w})}\leq\frac{1}{2L},$
(14)

where $L(\mathbf{w})\coloneqq
4L+\frac{2\rho\alpha}{n}\sum_{i=1}^{n}\left\|\nabla\mathcal{L}_{i}(\mathbf{w})\right\|$.

Proof  The results
$\frac{4}{5L(\mathbf{w})}\leq\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})}\right]$
and
$\frac{L(\mathbf{w})}{2}\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})^{2}}\right]\leq\frac{2}{L(\mathbf{w})}\leq\frac{1}{2L}$
can be found in Lemma 5.9 of Fallah et al. (2020a). Now we prove the rest
part. Utilize Theorem A.2 of Fallah et al. (2020a) with
$X=\frac{2\rho\alpha}{B^{\prime}}\sum_{i\in\mathcal{B}^{\prime}}\left\|\widehat{\nabla}_{\mathcal{S}^{\prime}}\mathcal{L}_{i}(\mathbf{w})\right\|$,
$c=4L$, $k=1$:

$\displaystyle
L(\mathbf{w})\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})}\right]\leq\frac{\sigma_{X}^{2}/4L+\mu_{X}^{2}\frac{\mu_{X}}{\sigma_{X}^{2}+\mu_{X}(\mu_{X}+c)}}{\sigma_{X}^{2}+\mu_{X}^{2}}L(\mathbf{w}),$

where $\mu_{X}$ and $\sigma_{X}^{2}$ are the mean and variance of $X$.
Consider that $\sigma_{X}^{2}+\mu_{X}(\mu_{X}+4L)\geq\mu_{X}(\mu_{X}+4L)$ and
$L(\mathbf{w})\leq\mu_{X}+5L$, which is shown in (60) of Fallah et al.
(2020a).

$\displaystyle
L(\mathbf{w})\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})}\right]\leq\frac{\sigma_{X}^{2}/4L+\mu_{X}^{2}/(\mu_{X}+4L)}{\sigma_{X}^{2}+\mu_{X}^{2}}(\mu_{X}+5L)=\frac{\frac{5}{4}\sigma_{X}^{2}+\frac{\mu_{X}^{2}(\mu_{X}+5L)}{\mu_{X}+4L}+\frac{\sigma_{X}^{2}\mu_{X}}{4L}}{\sigma_{X}^{2}+\mu_{X}^{2}}.$

Note that
$\frac{\mu_{X}+5L}{\mu_{X}+4L}=1+\frac{L}{\mu_{X}+4L}\leq\frac{5}{4}$ and
$\sigma_{X}^{2}\leq\frac{4\rho^{2}\alpha^{2}(n-B)}{B(n-1)}\left(\gamma_{G}^{2}+\frac{\sigma_{G}^{2}}{S}\right)$

$\displaystyle
L(\mathbf{w})\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})}\right]\leq\frac{5}{4}+\frac{\sigma_{X}^{2}\mu_{X}}{4L(\sigma_{X}^{2}+\mu_{X}^{2})}\leq\frac{5}{4}+\frac{\sigma_{X}^{4}/L^{2}+\mu_{X}^{2}}{8(\sigma_{X}^{2}+\mu_{X}^{2})},$

where the last inequality above uses Young’s inequality. We only require
$\sigma_{X}^{2}\leq L^{2}$ to make
$L(\mathbf{w})\mathbb{E}\left[\frac{1}{\widehat{L}(\mathbf{w})}\right]\leq\frac{11\eta_{0}}{8}$,
which is satisfied if
$\frac{(n-1)B}{n-B}\geq\lceil(8\rho\alpha\gamma_{G}/L)^{2}\rceil$,
$S\geq\lceil(8\rho\alpha\sigma_{G}/L)^{2}\rceil$.

##  Appendix B Convergence Analaysis of $\text{MOML}^{\text{v1}}$

###### Lemma 15

For the stochastic meta-gradient estimator
$\widehat{\Delta}_{\mathcal{B}_{t}}$ defined in (3), it holds that

$\displaystyle\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}$
$\displaystyle\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\right]\leq
C_{\Delta},$

where $C_{\Delta}\coloneqq(\alpha^{2}\sigma_{H}^{2}/K+(1+\alpha
L)^{2})(\sigma_{G}^{2}/K+G^{2})$.

Proof  Based on Assumption 1 and 4, we have

$\displaystyle\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t})\right\|^{2}\right]=\alpha^{2}\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t})-\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t})\right\|^{2}\right]+(1+\alpha
L)^{2}\leq\frac{\alpha^{2}\sigma_{H}^{2}}{K}+(1+\alpha L)^{2},$
$\displaystyle\mathbb{E}_{\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\right]=\mathbb{E}_{\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})-\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\right]+\left\|\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\leq\frac{\sigma_{G}^{2}}{K}+G^{2}.$

Consider that $\mathcal{S}_{2}^{i}$, $\mathcal{S}_{3}^{i}$ are mutually
independent.

$\displaystyle\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\right]$
$\displaystyle\leq\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t})\right\|^{2}\right]\mathbb{E}_{\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\right]\leq\left(\frac{\alpha^{2}\sigma_{H}^{2}}{S_{2}}+(1+\alpha
L)^{2}\right)\left(\frac{\sigma_{G}^{2}}{S_{3}}+G^{2}\right).$

Proof [Proof of Lemma 1] Based on the $L_{F}$-smoothness of $F$, we have

$\displaystyle F(\mathbf{w}_{t+1})$ $\displaystyle\leq
F(\mathbf{w}_{t})+\langle\nabla
F(\mathbf{w}_{t}),\mathbf{w}_{t+1}-\mathbf{w}_{t}\rangle+\frac{L_{F}}{2}\left\|\mathbf{w}_{t+1}-\mathbf{w}_{t}\right\|^{2}$
$\displaystyle=F(\mathbf{w}_{t})-\eta\langle
F(\mathbf{w}_{t}),\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\rangle+\frac{\eta^{2}L_{F}}{2}\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}$
$\displaystyle=F(\mathbf{w}_{t})-\eta\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta\langle\nabla F(\mathbf{w}_{t}),\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\rangle+\frac{\eta^{2}L_{F}}{2}\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}.$

Take expectation on both sides conditioned on $\mathcal{F}_{t}$, where
$\mathcal{F}_{t}$ denotes all the randomness before the $t$-th iteration.

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{t+1})\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq F(\mathbf{w}_{t})-\eta\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}$ $\displaystyle+\eta\langle\nabla
F(\mathbf{w}_{t}),\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]\rangle+\frac{\eta^{2}L_{F}}{2}\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}\mid\mathcal{F}_{t}\right].$
(15)

Consider that
$\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))\mid\mathcal{F}_{t}\right]=\nabla
F(\mathbf{w}_{t})$.

$\displaystyle\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]$
$\displaystyle=\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))-\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\mid\mathcal{F}_{t}\right]$
$\displaystyle=\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))-\mathbb{E}\left[(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\mid\mathcal{F}_{t},\mathcal{B}_{t}\right]\right)\mid\mathcal{F}_{t}\right]$
$\displaystyle=\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t})\right)\left(\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))-\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right)\mid\mathcal{F}_{t}\right].$

Then, Young’s inequality implies that

$\displaystyle\langle\nabla F(\mathbf{w}_{t}),\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]\rangle$
$\displaystyle=\mathbb{E}\left[\langle\nabla
F(\mathbf{w}_{t}),\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t})\right)\left(\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))-\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right)\rangle\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq\frac{\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}}{2}+\frac{(1+\alpha
L)^{2}L^{2}}{2}\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left\|\mathbf{v}_{i}(\mathbf{w}_{t})-\mathbf{u}_{t+1}^{i}\right\|^{2}\mid\mathcal{F}_{t}\right].$
(16)

Considering Lemma 15 and (16) and setting $\alpha\leq 1/L$, the R.H.S. of (15)
can be upper bounded as

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{t+1})\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq F(\mathbf{w}_{t})-\frac{\eta}{2}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\frac{\eta^{2}L_{F}C_{\Delta}}{2}+\frac{4\eta
L^{2}}{2}\mathbb{E}\left[\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left\|\mathbf{v}_{i}(\mathbf{w}_{t})-\mathbf{u}_{t+1}^{i}\right\|^{2}\mid\mathcal{F}_{t}\right].$

Use the tower property of conditional expectation, re-arrange the terms, and
unwrap the recursion above from iteration $0$ to $T-1$

$\displaystyle\mathbb{E}\left[\sum_{t=0}^{T-1}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]$
$\displaystyle\leq\frac{2F(\mathbf{w}_{0})}{\eta}+\eta
L_{F}C_{\Delta}T+\frac{8L^{2}}{B}\mathbb{E}\left[\sum_{t=0}^{T-1}\sum_{i\in\mathcal{B}_{t}}\left\|\mathbf{v}_{i}(\mathbf{w}_{t})-\mathbf{u}_{t+1}^{i}\right\|^{2}\right].$

Proof [Proof of Lemma 2] Based on (v1), the following equation holds

$\displaystyle\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{u}_{t_{k_{i}}^{i}+1}^{i}\right\|^{2}$
$\displaystyle=\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k_{i}}^{i}})-(1-\beta)\mathbf{u}_{t_{k-1}^{i}+1}^{i}-\beta\widehat{\bf
v}_{t_{k}^{i}}^{i}\right\|^{2}$
$\displaystyle=\mathbb{E}\left[\left\|(1-\beta)(\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}})-\mathbf{u}_{t_{k-1}^{i}+1}^{i})+(1-\beta)(\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}}))+\beta(\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\widehat{\bf
v}_{t_{k}^{i}}^{i})\right\|^{2}\right]$

We take expectation on both sides condition on $\mathcal{F}_{t_{k-1}^{i}}$

$\displaystyle\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{u}_{t_{k}^{i}+1}^{i}\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]$
$\displaystyle=\beta^{2}\alpha^{2}\mathbb{E}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t_{k}^{i}})-\nabla\mathcal{L}_{i}(\mathbf{w}_{t_{k}^{i}})\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]$
$\displaystyle\quad\quad+(1-\beta)^{2}\mathbb{E}\left[\left\|(\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}})-\mathbf{u}_{t_{k-1}^{i}+1}^{i})+(\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}}))\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]$
$\displaystyle\leq\frac{\beta^{2}\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}+(1-\beta)\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}})-\mathbf{u}_{t_{k-1}^{i}+1}^{i}\right\|^{2}$
$\displaystyle\quad\quad+(1-\beta)^{2}(1+1/\beta)(1+\alpha
L)^{2}\mathbb{E}\left[\left\|\mathbf{w}_{t_{k}^{i}}-\mathbf{w}_{t_{k-1}^{i}}\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right],$
$\displaystyle\leq\frac{\beta^{2}\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}+(1-\beta)\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}})-\mathbf{u}_{t_{k-1}^{i}+1}^{i}\right\|^{2}+8\eta^{2}\mathbb{E}\left[\left\|\sum_{\tau=t_{k-1}^{i}}^{t_{k}^{i}-1}\widehat{\Delta}_{\mathcal{B}_{\tau}}\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]\mathbb{I}[\beta\in(0,1)],$

where $t_{k-1}^{i}$ is the latest iteration before $t_{k}^{i}$ that task
$\mathcal{T}_{i}$ is also sampled, in other words,
$t_{k-1}^{i}=\max\\{\tau\mid\tau\in\mathcal{T}_{i}\land\tau<t_{k}^{i}\\}$.
Lemma 15 implies that

$\displaystyle\mathbb{E}\left[\left\|\sum_{\tau=t_{k-1}^{i}}^{t_{k}^{i}-1}\widehat{\Delta}_{\mathcal{B}_{\tau}}\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]$
$\displaystyle\leq\mathbb{E}\left[(t_{k}^{i}-t_{k-1}^{i})\sum_{\tau=t_{k-1}^{i}}^{t_{k}^{i}-1}\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{\tau}}\right\|^{2}\mid\mathcal{B}_{\tau},\mathcal{F}_{t_{k-1}^{i}}\right]\mid\mathcal{F}_{t_{k-1}^{i}}\right]$
$\displaystyle\leq\mathbb{E}\left[(t_{k}^{i}-t_{k-1}^{i})^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]C_{\Delta}.$

It is worth noting that $t_{k}^{i}-t_{k-1}^{i}$ follows the geometric
distribution. Thus, the second moment satisfies
$\mathbb{E}\left[(t_{k}^{i}-t_{k-1}^{i})^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]\leq\frac{2n^{2}}{B^{2}}$.
Then,

$\displaystyle\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k}^{i}})-\mathbf{u}_{t_{k}^{i}+1}^{i}\right\|^{2}\mid\mathcal{F}_{t_{k-1}^{i}}\right]$
$\displaystyle\quad\quad\quad\leq(1-\beta)\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{k-1}^{i}})-\mathbf{u}_{t_{k-1}^{i}+1}^{i}\right\|^{2}+\frac{16\eta^{2}n^{2}C_{\Delta}}{\beta
B^{2}}\mathbb{I}[\beta\in(0,1)]+\frac{\beta^{2}\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}.$

Re-arranging the terms, summing over $k=1,\dotsc,T_{i}$, and using the tower
property of conditional expectation leads to

$\displaystyle\mathbb{E}\left[\sum_{k=0}^{T_{i}-1}\left\|\mathbf{v}_{i}(\mathbf{w}_{k}^{i})-\mathbf{u}_{t_{k}^{i}+1}^{i}\right\|^{2}\right]$
$\displaystyle\leq\left(\frac{\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{t_{0}^{i}})-\mathbf{u}_{t_{0}^{i}+1}^{i}\right\|^{2}\right]}{\beta}+\frac{16\eta^{2}n^{2}C_{\Delta}}{\beta^{2}B^{2}}\mathbb{E}\left[T_{i}\right]\right)\mathbb{I}[\beta\in(0,1)]+\frac{\beta\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}\mathbb{E}\left[T_{i}\right].$

Initializing the personalized model $\mathbf{u}^{i}$ as
$\mathbf{u}_{t_{0}^{i}+1}^{i}=\mathbf{w}_{t_{0}^{i}}-\alpha\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t_{0}^{i}})$
and summing over $i=1,\dotsc,n$ results in

$\displaystyle\mathbb{E}\left[\sum_{i=1}^{n}\sum_{k=0}^{T_{i}-1}\left\|\mathbf{v}_{i}(\mathbf{w}_{k}^{i})-\mathbf{u}_{t_{k}^{i}+1}^{i}\right\|^{2}\right]$
$\displaystyle\leq\frac{n\sigma_{G}^{2}}{\beta|\mathcal{S}_{1}^{i}|}\mathbb{I}[\beta\in(0,1)]+\left(\frac{16\eta^{2}n^{2}C_{\Delta}}{\beta^{2}B^{2}}\mathbb{I}[\beta\in(0,1)]+\frac{\beta\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}\right)\mathbb{E}\left[\sum_{i=1}^{n}T_{i}\right].$

Note that $\sum_{i=1}^{n}T_{i}=T$ based on the definition.

######  Theorem 16 (Detailed Version of Theorem 4)

Under Assumptions 1, 2 3, 4, _$\text{MOML}^{\text{v1}}$_ with stepsizes
$\eta_{t}=\frac{B^{2/5}}{n^{2/5}T^{3/5}}$,
$\beta_{t}=\frac{n^{2/5}}{B^{2/5}T^{2/5}}<1$ and constant batch sizes
$|\mathcal{S}_{1}^{i}|=|\mathcal{S}_{2}^{i}|=|\mathcal{S}_{3}^{i}|=K=1$,
$|\mathcal{B}_{t}|=B=1$ can find a stationary point $\mathbf{w}_{\tau}$ in
$T=\mathcal{O}(n\epsilon^{-5})$ iterations.

Proof  Based on Lemma 1 and Lemma 2, we have

$\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]$
$\displaystyle\leq\frac{2F(\mathbf{w}_{0})}{\eta T}+\eta
L_{F}C_{\Delta}+\frac{8L^{2}}{B}\left(\frac{n\sigma_{G}^{2}}{\beta
KT}+\frac{16\eta^{2}n^{2}C_{\Delta}}{\beta^{2}B^{2}}+\frac{\beta\alpha^{2}\sigma_{G}^{2}}{K}\right).$

Choosing $B=1$, $K=1$, $\eta=\frac{B^{2/5}}{n^{2/5}T^{3/5}}$ and
$\beta=\frac{n^{2/5}}{B^{2/5}T^{2/5}}<1$ leads to

$\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]$
$\displaystyle\leq\frac{2n^{2/5}F(\mathbf{w}_{0})}{T^{2/5}}+\frac{L_{F}C_{\Delta}}{n^{2/5}T^{3/5}}+\frac{8L^{2}\sigma_{G}^{2}n^{3/5}}{T^{3/5}}+\frac{144L^{2}C_{\Delta}n^{2/5}}{T^{2/5}}+\frac{8L^{2}\alpha^{2}\sigma_{G}^{2}n^{2/5}}{T^{2/5}}.$

For $\mathbf{w}_{\tau}$ and $\tau$ is sampled from $0,\dotsc,T-1$ uniformly at
random, $\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{\tau})\right\|^{2}\right]\leq\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]$. Making the R.H.S. of the upper bound of
$\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]$ be equal to or smaller than
$\epsilon^{2}$ finishes the proof.

##  Appendix C Convergence Analysis of $\text{MOML}^{\text{v2}}$

###### Lemma 17

For the stochastic estimator
$\widehat{\Delta}_{\mathcal{B}_{t}}=\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\nabla_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})$
of the meta-gradient, we have

$\displaystyle\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}\mid\mathcal{F}_{t}\right]\leq
C_{3}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right]+C_{4}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+C_{5},$ (17)

where $\mathcal{F}_{t}$ denotes all randomness occurred before the $t$-th
iteration, and the constants are defined as $C_{3}\coloneqq
2L^{2}\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)$,
$C_{4}\coloneqq 4\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)(1+\alpha
L)^{2}C_{1}$, and $C_{5}\coloneqq\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)\left(\frac{\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+2(1+\alpha
L)^{2}(2C_{2}^{2}+1)\gamma_{G}^{2}\right)$.

Proof  The definition of the stochastic estimator
$\widehat{\Delta}_{\mathcal{B}_{t}}$ implies that

$\displaystyle\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}\mid\mathcal{F}_{t}\right]=\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\nabla_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle=\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left((I-\alpha\nabla_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})-(I-\alpha\nabla_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right)\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\quad\quad\quad+\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}\left((I-\alpha\nabla_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right)\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\quad\quad\quad+\mathbb{E}\left[\left\|\frac{1}{B}\sum_{i\in\mathcal{B}_{t}}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)\left(\frac{\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\right\|^{2}\mid\mathcal{F}_{t}\right]\right)$
$\displaystyle\leq\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)\left(\frac{\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+\frac{2}{n}\sum_{i=1}^{n}\left\|\nabla
F_{i}(\mathbf{w}_{t}))\right\|^{2}+\frac{2L^{2}}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right]\right),$

where the last inequality uses the fact
$\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))=\nabla
F_{i}(\mathbf{w}_{t})$. Lemma 10 shows that

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\left\|\nabla
F_{i}(\mathbf{w}_{t})\right\|^{2}\leq 2(1+\alpha L)^{2}C_{1}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+(1+\alpha L)^{2}(2C_{2}^{2}+1)\gamma_{G}^{2}.$

Then,

$\displaystyle\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle=\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)\left(\frac{\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+2(1+\alpha
L)^{2}(2C_{2}^{2}+1)\gamma_{G}^{2}\right)$ $\displaystyle+4\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)(1+\alpha
L)^{2}C_{1}\left\|\nabla F(\mathbf{w}_{t})\right\|^{2}$
$\displaystyle+2L^{2}\left((1+\alpha
L)^{2}+\frac{\alpha^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\right)\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right].$

Apart from the notations in Table 4, we define that
$\mathbf{I}_{i}\coloneqq\left(0_{d\times d},\dotsc,I_{d\times
d},\dotsc,0_{d\times d}\right)^{\top}\in\mathbb{R}^{nd\times d}$ (where the
$i$-th block in $\mathbf{I}_{i}$ is an identity matrix while the others are
zeros),
$\bar{\mathbf{w}}_{t}\coloneqq\left(\mathbf{w}_{t}^{\top},\dotsc,\mathbf{w}_{t}^{\top}\right)^{\top}\in\mathbb{R}^{nd}$,
$\mathbf{u}_{t}=\left([\mathbf{u}_{t}^{1}]^{\top},\dotsc,[\mathbf{u}_{t}^{n}]^{\top}\right)^{\top}\in\mathbb{R}^{nd}$,
$\widehat{\mathbf{g}}_{t}\coloneqq\sum_{i\in\mathcal{B}_{t}^{\prime}}\frac{1}{\alpha
p_{i}}\mathbf{I}_{i}\left(\mathbf{w}_{t}-\widehat{\bf v}_{t}^{i})\right)$,
$\bar{\mathbf{g}}_{t}\coloneqq\sum_{i=1}^{n}\mathbf{I}_{i}(\mathbf{w}_{t}-\widehat{\bf
v}_{t}^{i}))/\alpha$,
$\mathbf{g}_{t}\coloneqq\sum_{i=1}^{n}\mathbf{I}_{i}(\mathbf{w}_{t}-\mathbf{v}_{i}(\mathbf{w}_{t}))/\alpha$,
$\widetilde{\mathbf{w}}_{t}\coloneqq\bar{\mathbf{w}}_{t}-\alpha\mathbf{g}_{t}$.
Then, we can re-write the update rule of the personalized models
$\mathbf{u}_{t}^{i}$ for all tasks $i\in[n]$ in a more succinct expression
$\mathbf{u}_{t+1}=(1-\beta_{t})\mathbf{u}_{t}+\beta_{t}\left(\bar{\mathbf{w}}_{t}-\alpha\widehat{\mathbf{g}}_{t}\right)$.

###### Lemma 18

For
$\widehat{\mathbf{g}}_{t}\coloneqq\sum_{i\in\mathcal{B}_{t}^{\prime}}\frac{1}{\alpha
p_{i}}\mathbf{I}_{i}\left(\mathbf{w}_{t}-\widehat{\bf v}_{t}^{i})\right)$ and
$\mathbf{g}_{t}\coloneqq\sum_{i=1}^{n}\mathbf{I}_{i}(\mathbf{w}_{t}-\mathbf{v}_{i}(\mathbf{w}_{t}))/\alpha$,
we have

$\displaystyle\mathbb{E}\left[\left\|\mathbf{g}_{t}-\widehat{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq 2nC_{p}C_{1}^{2}\left\|\nabla
F(\mathbf{w}_{t-1})\right\|^{2}$
$\displaystyle+nC_{p}\left((2C_{2}^{2}+1)\gamma_{G}^{2}+\frac{2\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}\right)+2nLC_{p}\left\|\mathbf{w}_{t}-\mathbf{w}_{t-1}\right\|^{2},$
(18)

where $C_{p}\coloneqq\max_{i}\left(\frac{1}{p_{i}}-1\right)$.

Proof  Consider that
$\widehat{\mathbf{g}}_{t}\coloneqq\sum_{i\in\mathcal{B}_{t}^{\prime}}\frac{1}{\alpha
p_{i}}\mathbf{I}_{i}\left(\mathbf{w}_{t}-\widehat{\bf v}_{t}^{i})\right)$,
$\bar{\mathbf{g}}_{t}\coloneqq\sum_{i=1}^{n}\mathbf{I}_{i}(\mathbf{w}_{t}-\widehat{\bf
v}_{t}^{i}))/\alpha$,
$\mathbf{g}_{t}\coloneqq\sum_{i=1}^{n}\mathbf{I}_{i}(\mathbf{w}_{t}-\mathbf{v}_{i}(\mathbf{w}_{t})))/\alpha$.

$\displaystyle\mathbb{E}\left[\left\|\mathbf{g}_{t}-\widehat{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]=\mathbb{E}\left[\left\|\widehat{\mathbf{g}}_{t}-\bar{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]+\mathbb{E}\left[\left\|\mathbf{g}_{t}-\bar{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right].$
(19)

The first term on the right hand side of (19) can be upper bounded as

$\displaystyle\mathbb{E}\left[\left\|\widehat{\mathbf{g}}_{t}-\bar{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq
C_{p}\mathbb{E}\left[\left\|\bar{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]=C_{p}\left(\left\|\mathbf{g}_{t}\right\|^{2}+\mathbb{E}\left[\left\|\mathbf{g}_{t}-\bar{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]\right),$
(20)

where we define $C_{p}\coloneqq\max_{i}\left(\frac{1}{p_{i}}-1\right)$. Note
that
$\left\|\mathbf{g}_{t}\right\|^{2}=\sum_{i=1}^{n}\left\|\nabla\mathcal{L}_{i}(\mathbf{w}_{t})\right\|^{2}$.
Then,

$\displaystyle\left\|\mathbf{g}_{t}\right\|^{2}$
$\displaystyle=n\left\|\nabla\mathcal{L}(\mathbf{w}_{t})\right\|^{2}+n\gamma_{G}^{2}\leq
2n\left\|\nabla\mathcal{L}(\mathbf{w}_{t-1})\right\|^{2}+2nL\left\|\mathbf{w}_{t}-\mathbf{w}_{t-1}\right\|^{2}+n\gamma_{G}^{2}$
$\displaystyle\leq 2nC_{1}^{2}\left\|\nabla
F(\mathbf{w}_{t-1})\right\|^{2}+n(2C_{2}^{2}+1)\gamma_{G}^{2}+2nL\left\|\mathbf{w}_{t}-\mathbf{w}_{t-1}\right\|^{2}.$
(21)

The last inequality above utilizes Lemma 10. Besides,

$\displaystyle\mathbb{E}\left[\left\|\mathbf{g}_{t}-\bar{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle=\sum_{i=1}^{n}\mathbb{E}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{t})-\nabla\mathcal{L}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right]\leq\frac{n\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}.$
(22)

According to (19), (20), (21), and (22), we have

$\displaystyle\mathbb{E}\left[\left\|\mathbf{g}_{t}-\widehat{\mathbf{g}}_{t}\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq 2nC_{p}C_{1}^{2}\left\|\nabla
F(\mathbf{w}_{t-1})\right\|^{2}$
$\displaystyle+nC_{p}\left((2C_{2}^{2}+1)\gamma_{G}^{2}+\frac{2\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}\right)+2nLC_{p}\left\|\mathbf{w}_{t}-\mathbf{w}_{t-1}\right\|^{2}.$

Proof [Proof of Lemma 6] Recall that the notations
$\mathbf{v}_{i}(\mathbf{w}_{t})\coloneqq\mathbf{w}_{t}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}_{t})$,
$\bar{\mathbf{w}}_{t}\coloneqq\left(\mathbf{w}_{t}^{\top},\dotsc,\mathbf{w}_{t}^{\top}\right)^{\top}\in\mathbb{R}^{nd}$,
$\mathbf{u}_{t}=\left([\mathbf{u}_{t}^{1}]^{\top},\dotsc,[\mathbf{u}_{t}^{n}]^{\top}\right)^{\top}$,
$\widehat{\mathbf{g}}_{t}\coloneqq\sum_{i\in\mathcal{B}_{t}^{\prime}}\frac{1}{\alpha
p_{i}}\mathbf{I}_{i}\left(\mathbf{w}_{t}-\widehat{\bf v}_{t}^{i})\right)$,
$\mathbf{g}_{t}\coloneqq\sum_{i=1}^{n}\mathbf{I}_{i}(\mathbf{w}_{t}-\mathbf{v}_{i}(\mathbf{w}_{t}))/\alpha$,
$\widetilde{\mathbf{w}}_{t}\coloneqq\bar{\mathbf{w}}_{t}-\alpha\mathbf{g}_{t}$.
The update rule of the personalized models is
$\mathbf{u}_{t+1}=(1-\beta_{t})\mathbf{u}_{t}+\beta_{t}\left(\bar{\mathbf{w}}_{t}-\alpha\widehat{\mathbf{g}}_{t}\right)$.
We define that
$\Upsilon_{t}\coloneqq\frac{1}{n}\sum_{i=1}^{n}\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}$.

$\displaystyle\mathbb{E}\left[\Upsilon_{t+1}\mid\mathcal{F}_{t+1}\right]=\frac{1}{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+2}-\widetilde{\mathbf{w}}_{t+1}\right\|^{2}\mid\mathcal{F}_{t+1}\right]$
$\displaystyle=\frac{1}{n}\mathbb{E}\left[\left\|\widetilde{\mathbf{w}}_{t+1}-(1-\beta_{t+1})\mathbf{u}_{t+1}+\beta_{t+1}\left(\bar{\mathbf{w}}_{t+1}-\alpha\widehat{\mathbf{g}}_{t+1}\right)\right\|^{2}\mid\mathcal{F}_{t+1}\right]$
$\displaystyle=\frac{1}{n}\mathbb{E}\left[\left\|(1-\beta_{t+1})(\widetilde{\mathbf{w}}_{t}-\mathbf{u}_{t+1})+(1-\beta_{t+1})\left(\widetilde{\mathbf{w}}_{t+1}-\widetilde{\mathbf{w}}_{t}\right)+\alpha\beta_{t+1}(\mathbf{g}_{t+1}-\widehat{\mathbf{g}}_{t+1})\right\|^{2}\mid\mathcal{F}_{t+1}\right]$
$\displaystyle\leq(1-\beta_{t+1})\frac{1}{n}\left\|\widetilde{\mathbf{w}}_{t}-\mathbf{u}_{t+1}\right\|^{2}+\frac{8(1+\alpha
L)^{2}}{\beta_{t+1}n}\left\|\bar{\mathbf{w}}_{t+1}-\bar{\mathbf{w}}_{t}\right\|^{2}+\frac{\alpha^{2}\beta_{t+1}^{2}}{n}\mathbb{E}\left[\left\|\mathbf{g}_{t+1}-\widehat{\mathbf{g}}_{t+1}\right\|^{2}\mid\mathcal{F}_{t+1}\right]$
$\displaystyle=(1-\beta_{t+1})\frac{1}{n}\sum_{i=1}^{n}\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}+\frac{8(1+\alpha
L)^{2}}{\beta_{t+1}}\left\|\mathbf{w}_{t+1}-\mathbf{w}_{t}\right\|^{2}+\frac{\alpha^{2}\beta_{t+1}^{2}}{n}\mathbb{E}\left[\left\|\mathbf{g}_{t+1}-\widehat{\mathbf{g}}_{t+1}\right\|^{2}\mid\mathcal{F}_{t+1}\right].$

Based on Lemma 18, we have

$\displaystyle\mathbb{E}\left[\Upsilon_{t+1}\mid\mathcal{F}_{t+1}\right]\leq(1-\beta_{t+1})\Upsilon_{t}+2\left(\frac{4(1+\alpha
L)^{2}}{\beta_{t+1}}+\beta_{t+1}^{2}\alpha^{2}LC_{p}\right)\eta_{t}^{2}\left\|\Delta_{\mathcal{B}_{t}}\right\|^{2}$
$\displaystyle\quad\quad\quad\quad\quad\quad+2\beta_{t+1}^{2}\alpha^{2}C_{p}C_{1}^{2}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\beta_{t+1}^{2}\alpha^{2}C_{p}(2C_{2}^{2}+1)\gamma_{G}^{2}+\frac{2\beta_{t+1}^{2}\alpha^{2}C_{p}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}.$

We choose $\beta_{t+1}=6L^{2}\eta_{0}^{-1/3}\eta_{t}$. Since we need to ensure
$\beta_{t}\leq 1$ for any $t$, we only need to maintain
$\eta_{0}\leq\left(\frac{2}{3L}\right)^{\frac{3}{2}}$. Then,
$\beta_{t+1}^{2}=36L^{4}\eta_{0}^{-2/3}\eta_{t}^{2}\leq
3L^{2}\eta_{0}^{\frac{4}{3}}$.

$\displaystyle\mathbb{E}\left[\Upsilon_{t+1}\mid\mathcal{F}_{t+1}\right]\leq(1-6L^{2}\eta_{0}^{-1/3}\eta_{t})\Upsilon_{t}+C_{6}\eta_{0}^{1/3}\eta_{t}\left\|\Delta_{\mathcal{B}_{t}}\right\|^{2}+\eta_{0}^{1/3}\eta_{t}C_{7}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta_{0}^{4/3}C_{8},$

where $C_{6}\coloneqq\frac{4(1+\alpha L)^{2}+\alpha^{2}LC_{p}}{3L^{2}}$,
$C_{7}\coloneqq 18L^{3}\alpha^{2}C_{p}C_{1}^{2}$, $C_{8}\coloneqq
3L^{2}\left(\alpha^{2}C_{p}(2C_{2}^{2}+1)\gamma_{G}^{2}+\frac{2\alpha^{2}C_{p}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}\right)$.
Then, in view of the tower property of conditional expectation, we have

$\displaystyle\mathbb{E}\left[\Upsilon_{t+1}\mid\mathcal{F}_{t}\right]\leq(1-6L^{2}\eta_{0}^{-1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right])\mathbb{E}\left[\Upsilon_{t}\mid\mathcal{F}_{t}\right]$
$\displaystyle\quad\quad\quad+\eta_{0}^{1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]C_{7}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta_{0}^{4/3}C_{8}$
$\displaystyle\quad\quad\quad+C_{6}\eta_{0}^{1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\left(C_{3}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\Upsilon_{t}\mid\mathcal{F}_{t}\right]+C_{4}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+C_{5}\right)$
$\displaystyle=(1-6L^{2}\eta_{0}^{-1/3}(1-C_{3}C_{6}\eta_{0}^{2/3}/6L^{2})\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right])\mathbb{E}\left[\Upsilon_{t}\mid\mathcal{F}_{t}\right]$
$\displaystyle\quad\quad\quad+\eta_{0}^{1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\left(C_{7}+C_{4}C_{6}\right)\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta_{0}^{4/3}C_{8}+\eta_{0}^{1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]C_{5}C_{6}$
$\displaystyle\leq\left(1-3L^{2}\eta_{0}^{-1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\right)\mathbb{E}\left[\Upsilon_{t}\mid\mathcal{F}_{t}\right]$
$\displaystyle\quad\quad\quad+\eta_{0}^{1/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\left(C_{7}+C_{4}C_{6}\right)\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta_{0}^{4/3}\left(C_{8}+\frac{C_{5}C_{6}}{4L}\right),$
(23)

where the last step holds when
$\eta_{0}\leq\left(\frac{3L^{2}}{C_{3}C_{6}}\right)^{3/2}$.

###### Lemma 19

If we set
$\eta_{0}\leq\min\left\\{\frac{2L^{2}}{5C_{3}},\frac{1}{8C_{9}^{3/2}},\frac{1}{20C_{4}}\right\\}$
and define the potential function $\Phi_{t}$ as
$\Phi_{t}\coloneqq\eta_{0}^{1/3}\frac{1}{n}\sum_{i=1}^{n}\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}+F(\mathbf{w}_{t})$,
we have

$\displaystyle\mathbb{E}\left[\Phi_{t+1}\right]$
$\displaystyle\leq\mathbb{E}\left[\Phi_{t}\right]-\frac{\eta_{0}}{80}\min\left\\{\frac{\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]}{L+\rho\alpha\sigma},\frac{\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|\right]}{\rho\alpha}\right\\}+\eta_{0}^{5/3}\left(\frac{\eta_{0}^{1/3}C_{5}}{2}+C_{10}\right).$

Proof  Based on Lemma 12, we have

$\displaystyle F(\mathbf{w}_{t+1})$ $\displaystyle\leq
F(\mathbf{w}_{t})+\langle\nabla
F(\mathbf{w}_{t}),\mathbf{w}_{t+1}-\mathbf{w}_{t}\rangle+\frac{L(\mathbf{w}_{t})}{2}\left\|\mathbf{w}_{t+1}-\mathbf{w}_{t}\right\|^{2}$
$\displaystyle=F(\mathbf{w}_{t})-\eta_{t}\langle
F(\mathbf{w}_{t}),\widehat{\Delta}_{\mathcal{B}_{t}}\rangle+\frac{\eta_{t}^{2}L(\mathbf{w}_{t})}{2}\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}$
$\displaystyle=F(\mathbf{w}_{t})-\eta_{t}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\eta_{t}\langle\nabla F(\mathbf{w}_{t}),\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\rangle+\frac{\eta_{t}^{2}L(\mathbf{w}_{t})}{2}\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}.$

Consider the step size $\eta_{t}$ and $\widehat{\Delta}_{\mathcal{B}_{t}}$ are
independent. Take expectation on both sides conditioned on $\mathcal{F}_{t}$,
where $\mathcal{F}_{t}$ denotes all randomness occurred before the $t$-th
iteration.

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{t+1})\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq
F(\mathbf{w}_{t})-\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}$
$\displaystyle\quad\quad\quad+\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]\langle\nabla
F(\mathbf{w}_{t}),\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]\rangle+\frac{\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]L(\mathbf{w}_{t})}{2}\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}\mid\mathcal{F}_{t}\right]$

Consider the fact
$\mathbb{E}\left[\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]=\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i})\mid\mathcal{F}_{t}\right]$.

$\displaystyle\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{t}))(\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{t}))-\nabla\mathcal{L}_{i}(\mathbf{u}_{t+1}^{i}))\mid\mathcal{F}_{t}\right]$

Since $\left\|\cdot\right\|$ is a convex function, we have the following
equation based on the Jensen’s and Cauchy-Schwarz inequalities

$\displaystyle\langle\nabla F(\mathbf{w}_{t}),\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]\rangle\leq\left\|\nabla
F(\mathbf{w}_{t})\right\|\left\|\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]\right\|$
$\displaystyle\leq\frac{\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}}{2}+\frac{\left\|\mathbb{E}\left[\nabla
F(\mathbf{w}_{t})-\widehat{\Delta}_{\mathcal{B}_{t}}\mid\mathcal{F}_{t}\right]\right\|^{2}}{2}$
$\displaystyle\leq\frac{\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}}{2}+\frac{(1+\alpha
L)^{2}L^{2}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right]}{2}$
$\displaystyle\leq\frac{\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}}{2}+\frac{4L^{2}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right]}{2},$

where the last inequality holds when $\alpha\leq 1/L$. Thus,

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{t+1})\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq
F(\mathbf{w}_{t})-\frac{\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]}{2}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\frac{4L^{2}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]}{2}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\quad\quad\quad+\frac{L(\mathbf{w}_{t})\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]}{2}\mathbb{E}\left[\left\|\widehat{\Delta}_{\mathcal{B}_{t}}\right\|^{2}\mid\mathcal{F}_{t}\right]$
$\displaystyle\leq
F(\mathbf{w}_{t})-\frac{\left(\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]-C_{4}L(\mathbf{w}_{t})\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]\right)}{2}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}+\frac{C_{5}L(\mathbf{w}_{t})\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]}{2}$
$\displaystyle\quad\quad\quad+\frac{\left(4L^{2}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]+L(\mathbf{w}_{t})\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]C_{3}\right)}{2}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\mid\mathcal{F}_{t}\right].$
(24)

Based on Lemma 14, we can derive that

$\displaystyle-3L^{2}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]+\frac{\left(4L^{2}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]+L(\mathbf{w}_{t})\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]C_{3}\right)}{2}=-\frac{4L^{2}\eta_{0}}{5L(\mathbf{w}_{t})}+\frac{2C_{3}\eta_{0}^{2}}{L(\mathbf{w}_{t})}\leq
0,$

where we need $\eta_{0}\leq\frac{2L^{2}}{5C_{3}}$. Besides, if
$\eta_{0}\leq\frac{1}{8C_{9}^{3/2}}$ and $\eta_{0}\leq\frac{1}{20C_{4}}$, we
have

$\displaystyle
C_{9}\eta_{0}^{2/3}\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]-\frac{\left(\mathbb{E}\left[\eta_{t}\mid\mathcal{F}_{t}\right]-C_{4}L(\mathbf{w}_{t})\mathbb{E}\left[\eta_{t}^{2}\mid\mathcal{F}_{t}\right]\right)}{2}\leq-\frac{\eta_{0}}{5L(\mathbf{w}_{t})}+\frac{2\eta_{0}^{2}C_{4}}{L(\mathbf{w}_{t})}\leq-\frac{\eta_{0}}{10L(\mathbf{w}_{t})}.$

Multiplying (23) by $\eta_{0}^{1/3}$ and summing it to (24) leads to

$\displaystyle\eta_{0}^{1/3}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+2}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t+1})\right\|^{2}\right]+\mathbb{E}\left[F(\mathbf{w}_{t+1})\right]$
$\displaystyle\leq\eta_{0}^{1/3}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}\right]+\mathbb{E}\left[F(\mathbf{w}_{t})\right]-\mathbb{E}\left[\frac{\eta_{0}}{10L(\mathbf{w}_{t})}\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]+\eta_{0}^{5/3}\left(\eta_{0}^{1/3}C_{5}/2+C_{10}\right).$

Define that
$\Phi_{t}\coloneqq\eta_{0}^{1/3}\frac{1}{n}\sum_{i=1}^{n}\left\|\mathbf{u}_{t+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{t})\right\|^{2}+F(\mathbf{w}_{t})$.
Besides, utilize (103) $\sim$ (106) in Fallah et al. (2020a):

$\displaystyle\mathbb{E}\left[\Phi_{t+1}\right]$
$\displaystyle\leq\mathbb{E}\left[\Phi_{t}\right]-\frac{\eta_{0}}{80}\min\left\\{\frac{\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]}{L+\rho\alpha\sigma},\frac{\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|\right]}{\rho\alpha}\right\\}+\eta_{0}^{5/3}\left(\eta_{0}^{1/3}C_{5}/2+C_{10}\right).$

######  Theorem 20 (Detailed Version of Thoerem 7)

Under Assumptions 1, 2, and 5, it is guaranteed that
_$\text{MOML}^{\text{v2}}$_ can find an $\epsilon$-stationary point in
$\frac{160(L+\rho\alpha(\sigma+\epsilon))\Phi_{0}}{C_{11}\epsilon^{5}}$
iterations, where $C_{11}=\Omega(1/C_{p})$, $C_{p}=\max_{i}1/p_{i}-1$.

Proof  Based on Lemma 19, we can derive that:

$\displaystyle\mathbb{E}\left[\Phi_{t+1}\right]$
$\displaystyle\leq\mathbb{E}\left[\Phi_{t}\right]-\frac{\eta_{0}}{80}\min\left\\{\frac{\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]}{L+\rho\alpha\sigma},\frac{\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|\right]}{\rho\alpha}\right\\}+\eta_{0}^{5/3}\left(\eta_{0}^{1/3}C_{5}/2+C_{10}\right).$

Suppose that $\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|\right]\geq\epsilon$ and
$\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|^{2}\right]\geq\left(\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{t})\right\|\right]\right)^{2}\geq\epsilon^{2}$, $\forall t\in
1,\dotsc,T$. Otherwise, we can find an $\epsilon$-stationary point in the
first $T$ iterations. Thus, choosing $\eta_{0}=C_{11}\epsilon^{3}$
($C_{11}>0$) and telescoping over the $T$ iterations leads to:

$\displaystyle T\frac{\epsilon^{2}}{L+\rho\alpha(\sigma+\epsilon)}$
$\displaystyle\leq
T\min\left\\{\frac{\epsilon^{2}}{L+\rho\alpha\sigma},\frac{\epsilon}{\rho\alpha}\right\\}\leq\frac{80\Phi_{0}}{C_{11}\epsilon^{3}}+80TC_{11}^{2/3}\epsilon^{2}\left(C_{9}^{1/3}\epsilon
C_{5}/2+C_{10}\right),$

where
$C_{11}\coloneqq\min\left\\{\frac{1}{160C_{5}(L+\rho\alpha(\sigma+\epsilon))},\frac{1}{\left(320C_{9}(L+\rho\alpha(\sigma+\epsilon))\right)^{3/2}}\right\\}$.
Note that $\eta_{0}=C_{11}\epsilon^{3}$ can satisfy the requirements on
$\eta_{0}$ in Lemma 6 and Lemma 19. Thus, we can find at least one
$\epsilon$-stationary point if
$T\geq\frac{160(L+\rho\alpha(\sigma+\epsilon))\Phi_{0}}{C_{11}\epsilon^{5}}$.

##  Appendix D Convergence Analysis of LocalMOML

To tackle with partial client sampling, we follow the ideas of Li et al.
(2019); Karimireddy et al. (2020): in each round, the global model
$\mathbf{w}_{r}$ is sent to the sampled clients $i\in\mathcal{B}_{r}$ and the
sampled clients run local step. Here we assume that $\mathbf{w}_{r}$ is also
_virtually_ sent to the other clients $i\not\in\mathcal{B}_{r}$ and those
clients also _virtually_ run local step. After $H$ iterations, only
$\Delta\mathbf{w}_{r}^{i}$, $i\in\mathcal{B}_{r}$ are aggregated to compute
$\mathbf{w}_{r+1}$. It is worth noting that the extra communication of
$\mathbf{w}_{r}$ to clients $i\not\in\mathcal{B}_{r}$ and the local steps on
clients $i\not\in\mathcal{B}_{r}$ are only used in the proof and not actually
executed when running Algorithm 3. The lemma below is key to our analysis.

###### Lemma 21

After one round of _LocalMOML_ , it satisfies that:

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{r+1})\right]\leq\mathbb{E}\left[F(\mathbf{w}_{r})\right]-\frac{\tilde{\eta}}{2}\left(1-8\tilde{\eta}-\frac{8\tilde{\eta}(n-B)}{B(n-1)}\right)\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}\right]$
$\displaystyle\quad\quad+\frac{\tilde{\eta}}{2}\left(C_{\rho,G,L}+8\tilde{\eta}L_{F}^{2}+\frac{8\tilde{\eta}(n-B)L_{F}^{2}}{B(n-1)}\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\right\|^{2}\right]$
$\displaystyle\quad\quad+8\tilde{\eta}\left(1+L^{2}\tilde{\eta}\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\right]+\frac{4(n-B)}{B(n-1)}\tilde{\eta}^{2}\gamma_{F}^{2}+\frac{\tilde{\eta}^{2}\hat{\sigma}^{2}}{BH}.$
(25)

where
$\hat{\sigma}^{2}\coloneqq\frac{2\sigma_{G}^{2}}{S_{3}}+\frac{2\alpha^{2}\sigma_{G}^{2}}{S_{3}}\left(\frac{\sigma_{H}^{2}}{S_{2}}+L^{2}\right)+\frac{\alpha^{2}G^{2}\sigma_{H}^{2}}{S_{2}}$,
$\tilde{\eta}=\eta H$, $C_{\rho,G,L}\coloneqq 2G^{2}\rho^{2}/L^{2}+16L^{2}$.

Proof  Based on the smoothness of $F$ shown in Lemma 12, we have

$\displaystyle F(\mathbf{w}_{r+1})$ $\displaystyle\leq
F(\mathbf{w}_{r})-\langle\nabla
F(\mathbf{w}_{r}),\Delta\mathbf{w}_{r}\rangle+\frac{L_{F}}{2}\left\|\Delta\mathbf{w}_{r}\right\|^{2}.$

Let $\tilde{\eta}\coloneqq H\eta$ and $\mathcal{F}_{r}$ denotes all randomness
occurred before the communication round $r$. Note that
$\Delta\mathbf{w}_{r}=\frac{\tilde{\eta}}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}\widehat{\Delta}_{r,h}^{i}$
and
$\mathbb{E}\left[\Delta\mathbf{w}_{r}\right]=\frac{\tilde{\eta}}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\widehat{\Delta}_{r,h}^{i}\right]$
because $\widehat{\Delta}_{r,h}^{i}$ does not depend on the client sampling
$\mathcal{B}_{r}$.

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{r+1})\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq
F(\mathbf{w}_{r})-\tilde{\eta}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\langle\nabla
F(\mathbf{w}_{r}),\mathbb{E}\left[\widehat{\Delta}_{r,h}^{i}\mid\mathcal{F}_{r}\right]\rangle+\mathbb{E}\left[\left\|\Delta\mathbf{w}_{r}\right\|^{2}\mid\mathcal{F}_{r}\right].$
(26)

The second term on the right hand side can be decomposed as

$\displaystyle-\tilde{\eta}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\langle\nabla
F(\mathbf{w}_{r}),\mathbb{E}\left[\widehat{\Delta}_{r,h}^{i}\mid\mathcal{F}_{r}\right]\rangle$
$\displaystyle=-\tilde{\eta}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\langle\nabla
F(\mathbf{w}_{r}),\mathbb{E}\left[\widehat{\Delta}_{r,h}^{i}-\nabla
F(\mathbf{w}_{r})\mid\mathcal{F}_{r}\right]\rangle-\tilde{\eta}\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}.$

Based on the definition of $\widehat{\Delta}_{r,h}^{i}$, we have

$\displaystyle-\tilde{\eta}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\langle\nabla
F(\mathbf{w}_{r}),\mathbb{E}\left[\widehat{\Delta}_{r,h}^{i}-\nabla
F(\mathbf{w}_{r})\mid\mathcal{F}_{r}\right]\rangle$
$\displaystyle=-\tilde{\eta}\mathbb{E}\left[\langle\nabla
F(\mathbf{w}_{r}),\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r}))\right)\rangle\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq\frac{\tilde{\eta}}{2}\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}$
$\displaystyle\quad\quad+\frac{\tilde{\eta}}{2}\mathbb{E}\left[\left\|\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r}))\right)\right\|^{2}\mid\mathcal{F}_{r}\right].$

The second term on the right hand side can be upper bounded as

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r}))\right)\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq
2\mathbb{E}\left[\left\|\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right)\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\quad+2\mathbb{E}\left[\left\|\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r}))\right)\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq
2\alpha^{2}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})-\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r})\right\|^{2}\left\|\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\quad+4\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right\|^{2}\left\|I-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}_{r})\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\quad+4\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))-\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r}))\right\|^{2}\left\|I-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}_{r})\right\|^{2}\mid\mathcal{F}_{r}\right].$

Based on Assumption 1, 4, we have

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left((I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{w}_{r})\right)\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq
2\alpha^{2}G^{2}\rho^{2}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\right\|^{2}\mid\mathcal{F}_{r}\right]+4(1+\alpha
L)^{2}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\quad\quad+4(1+\alpha
L)^{2}L^{2}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle=C_{\rho,G,L}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\right\|^{2}\mid\mathcal{F}_{r}\right]+16\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\mid\mathcal{F}_{r}\right],$

where $C_{\rho,G,L}\coloneqq 2\alpha^{2}G^{2}\rho^{2}+4(1+\alpha
L)^{2}L^{2}\leq 2G^{2}\rho^{2}/L^{2}+16L^{2}$ when $\alpha\leq 1/L$. Besides,
the last term on the right hand side of (26) can be upper bounded as

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}\widehat{\Delta}_{r,h}^{i}\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq\mathbb{E}\left[\frac{1}{B^{2}H^{2}}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\Delta}_{r,h}^{i}-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]\mid\mathcal{F}_{r}\right]$
$\displaystyle+\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right],$
(27)

which is due to
$\widehat{\Delta}_{r,h}^{i}=(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})$
and

$\displaystyle\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right]=0.$

Next, we can decompose
$\widehat{\Delta}_{r,h}^{i}-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r}))\nabla\mathcal{L}_{i}(\mathbf{w}_{r})$
as

$\displaystyle(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})$
$\displaystyle=\left(\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right)-\alpha\left(\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right)$
$\displaystyle\quad\quad-\alpha\left(\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right),$

Then, the term
$\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\Delta}_{r,h}^{i}-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]$
can be upper bounded as

$\displaystyle\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\Delta}_{r,h}^{i}-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]$
$\displaystyle=\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})+\alpha(\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i}))\right\|^{2}\right]$
$\displaystyle\quad\quad+\alpha^{2}\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|\left(\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})-\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right)\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]$
$\displaystyle\leq
2\mathbb{E}_{\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]+\alpha^{2}G^{2}\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})-\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\right]$
$\displaystyle\quad\quad+2\alpha^{2}\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\right]\mathbb{E}_{\mathcal{S}_{3}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]$
$\displaystyle\leq\frac{2\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+\frac{2\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}\mathbb{E}_{\mathcal{S}_{2}^{i}}\left[\left\|\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})-\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}+L^{2}\right]+\frac{\alpha^{2}G^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}$
$\displaystyle\leq\frac{2\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}+\frac{2\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{3}^{i}|}\left(\frac{\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}+L^{2}\right)+\frac{\alpha^{2}G^{2}\sigma_{H}^{2}}{|\mathcal{S}_{2}^{i}|}\coloneqq\hat{\sigma}^{2}.$

Besides, the last term on the right hand side of (27) can be upper bounded as

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]$
$\displaystyle\leq
2\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\quad\quad+2\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\left(\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))-\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right)\right\|^{2}\mid\mathcal{F}_{r}\right].$

Apply Lemma 11 to the first term on the right hand side:

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\leq\frac{(n-B)}{B(n-1)}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle+\mathbb{E}\left[\left\|\frac{1}{n}\sum_{i=1}^{n}\sum_{h=1}^{H}(1-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\mid\mathcal{F}_{r}\right].$

Then, we have

$\displaystyle\mathbb{E}\left[\left\|\frac{1}{BH}\sum_{i\in\mathcal{B}_{r}}\sum_{h=1}^{H}(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\right\|^{2}\right]$
$\displaystyle\leq\frac{2(n-B)}{B(n-1)}\frac{1}{nH}\sum_{i=1}^{N}\sum_{h=1}^{H}\mathbb{E}\left[\left\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right\|^{2}\mid\mathcal{F}_{r}\right]+4\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}$
$\displaystyle\quad\quad+4E\left[\left\|\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\left(\nabla
F_{i}(\mathbf{w}_{r,h}^{i})-\nabla
F_{i}(\mathbf{w}_{r})\right)\right\|^{2}\mid\mathcal{F}_{r}\right]$
$\displaystyle\quad\quad+\frac{2(1+\alpha
L)^{2}L^{2}}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})-\mathbf{u}_{r,h}^{i}\right\|^{2}\mid\mathcal{F}_{r}\right].$

Based on Lemma 13, we have

$\displaystyle\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\left(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right)\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right\|^{2}\right]=\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]$
$\displaystyle\leq\frac{2}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r})\|^{2}\right]+\frac{2L_{F}^{2}}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]$
$\displaystyle\leq 2\gamma_{F}^{2}+2\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right]+\frac{2L_{F}^{2}}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right],$

Put them together and use the tower property of expectation on both sides.

$\displaystyle\mathbb{E}\left[F(\mathbf{w}_{r+1})\right]$
$\displaystyle\leq\mathbb{E}\left[F(\mathbf{w}_{r})\right]-\frac{\tilde{\eta}}{2}\left(1-8\tilde{\eta}-\frac{8\tilde{\eta}(n-B)}{B(n-1)}\right)\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}\right]$
$\displaystyle\quad\quad+\frac{\tilde{\eta}}{2}\left(C_{\rho,G,L}+8\tilde{\eta}L_{F}^{2}+\frac{8\tilde{\eta}(n-B)L_{F}^{2}}{B(n-1)}\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\right\|^{2}\right]$
$\displaystyle\quad\quad+8\tilde{\eta}\left(1+L^{2}\tilde{\eta}\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\left\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\right]+\frac{4(n-B)}{B(n-1)}\tilde{\eta}^{2}\gamma_{F}^{2}+\frac{\tilde{\eta}^{2}\hat{\sigma}^{2}}{BH}.$

###### Lemma 22

If $\eta\leq\frac{1}{2HL_{F}}$, it is satisfied that:

$\displaystyle\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]$
$\displaystyle\leq
16\eta^{2}H(H-1)\left(\hat{\sigma}^{2}+2\gamma_{F}^{2}\right)+32\eta^{2}H^{2}L^{2}\left(\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]\right)$
$\displaystyle\quad\quad\quad+32\eta^{2}H(H-1)\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right].$ (28)

Proof  For the stochastic estimator
$\widehat{\Delta}_{r,h}^{i}\coloneqq(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}^{i}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}^{i}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})$
of the meta-gradient and
$\Delta_{r,h}^{i}\coloneqq(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})$,
we have

$\displaystyle\mathbb{E}\left[\widehat{\Delta}_{r,h}^{i}-\Delta_{r,h}^{i}\right]=\mathbb{E}\left[\mathbb{E}_{\mathcal{S}_{2}^{i},\mathcal{S}_{3}^{i}}\left[\widehat{\Delta}_{r,h}^{i}-\Delta_{r,h}^{i}\right]\right]=0.$

Then, the local drift at iteration $h+1$, round $r$ can be upper bounded as

$\displaystyle\mathbb{E}\left[\|\mathbf{w}_{r,h+1}^{i}-\mathbf{w}_{r}\|^{2}\right]$
$\displaystyle\leq\left(1+\frac{1}{H}\right)\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]+\eta^{2}(1+H)\mathbb{E}\left[\|(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\|^{2}\right]$
$\displaystyle=\left(1+\frac{1}{H}\right)\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]+\eta^{2}(1+H)\mathbb{E}\left[\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\|^{2}\right]$
$\displaystyle\quad\quad+\eta^{2}(1+H)\mathbb{E}\left[\|(I-\alpha\widehat{\nabla}_{\mathcal{S}_{2}}^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\widehat{\nabla}_{\mathcal{S}_{3}}\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\|^{2}\right]$
$\displaystyle\leq\left(1+\frac{1}{H}\right)\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]+2\eta^{2}H\mathbb{E}\left[\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\|^{2}\right]+2\eta^{2}H\hat{\sigma}^{2}.$

Next, we upper bound
$\mathbb{E}\left[\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\|^{2}\right]$.
Note that we define
$\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\coloneqq\mathbf{w}_{r,h}^{i}-\alpha\nabla\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})$
and $\nabla
F_{i}(\mathbf{w}_{r,h}^{i})\coloneqq\left(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i})\right)\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))$.

$\displaystyle\mathbb{E}\left[\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})\|^{2}\right]$
$\displaystyle\leq\mathbb{E}\left[\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\left(\nabla\mathcal{L}_{i}(\mathbf{u}_{r,h}^{i})-\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\right)\|^{2}\right]$
$\displaystyle\quad\quad\quad+\mathbb{E}\left[\|(I-\alpha\nabla^{2}\mathcal{L}_{i}(\mathbf{w}_{r,h}^{i}))\nabla\mathcal{L}_{i}(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i}))\|^{2}\right]$
$\displaystyle\leq(1+\alpha
L)^{2}L^{2}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]+\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]$ $\displaystyle\leq(1+\alpha
L)^{2}L^{2}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]+2L_{F}^{2}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]+2\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r})\|^{2}\right].$

Thus, we have the following equation if $\eta\leq\frac{1}{2HL_{F}}$

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\mathbf{w}_{r,h+1}^{i}-\mathbf{w}_{r}\right\|^{2}\right]$
$\displaystyle\leq\left(1+\frac{2}{H}\right)\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]+2\eta^{2}H\hat{\sigma}^{2}$
$\displaystyle\quad\quad+2\eta^{2}H\left(\frac{(1+\alpha
L)^{2}L^{2}}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]+\frac{2}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r})\|^{2}\right]\right).$

Based on Lemma 13, we have

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r})\|^{2}\right]$ $\displaystyle=\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right]+\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r})-\nabla F(\mathbf{w}_{r})\|^{2}\right]$
$\displaystyle\leq\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right]+\gamma_{F}^{2}.$

Given that $\mathbf{w}_{r,1}^{i}=\mathbf{w}_{r}$ and $\alpha\leq 1/L$, and
$H\geq 2$, we can obtain that

$\displaystyle\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]$
$\displaystyle\leq
16\eta^{2}H(H-1)\left(\hat{\sigma}^{2}+2\gamma_{F}^{2}\right)+32\eta^{2}H^{2}L^{2}\left(\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]\right)$
$\displaystyle\quad\quad\quad+32\eta^{2}H(H-1)\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right].$

###### Lemma 23

If $\alpha\leq 1/L$ and
$\eta\leq\min\left\\{\frac{1}{4HL_{F}},\frac{\beta}{16L\mathbb{I}[\beta\in(0,1)]}\right\\}$,
we have

$\displaystyle\frac{1}{nHR}\sum_{i=1}^{n}\sum_{r=1}^{R}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]$
$\displaystyle\leq\frac{2\sigma_{G}^{2}}{\beta\left(\mathbb{I}[B<n]HS_{0}+\mathbb{I}[B=n]HR\right)}\mathbb{I}[\beta\in(0,1)]+\frac{96\eta^{2}(\hat{\sigma}^{2}+2\gamma_{F}^{2})}{\beta^{2}}\mathbb{I}[\beta\in(0,1)]+\frac{2\beta\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}$
$\displaystyle\quad\quad\quad+\frac{192\eta^{2}}{\beta^{2}R}\sum_{r=1}^{R}\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right]\mathbb{I}[\beta\in(0,1)].$ (29)

Proof  We define that $\widehat{\bf
v}_{r,h}^{i}\coloneqq\mathbf{w}_{r,h}^{i}-\alpha\widehat{\nabla}_{\mathcal{S}_{1}^{i}}\mathcal{L}_{i}(\mathbf{w}_{r,h})$.

$\displaystyle\mathbb{E}\left[\left\|\mathbf{u}_{r,h+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})\right\|^{2}\right]$
$\displaystyle=\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})-(1-\beta)\mathbf{u}_{r,h}^{i}-\beta\widehat{\bf
v}_{r,h+1}^{i}\right\|^{2}\right]$
$\displaystyle=\mathbb{E}\left[\left\|(1-\beta)\left(\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})-\mathbf{u}_{r,h}^{i}+\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\right)+\beta(\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})-\widehat{\bf
v}_{r,h+1}^{i})\right\|^{2}\right]$

Since $\mathbb{E}\left[\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})-\widehat{\bf
v}_{r,h+1}^{i}\right]=0$ and
$\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})-\widehat{\bf
v}_{r,h+1}^{i}\right\|^{2}\right]\leq\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}$,
we have

$\displaystyle\mathbb{E}\left[\left\|\mathbf{u}_{r,h+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})\right\|^{2}\right]$
$\displaystyle\leq(1-\beta)^{2}\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})-\mathbf{u}_{r,h}^{i}+\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\right\|^{2}\right]+\beta^{2}\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}.$
$\displaystyle\leq(1-\beta)^{2}(1+\beta)\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})-\mathbf{u}_{r,h}^{i}\right\|^{2}\right]$
$\displaystyle\quad\quad+4(1-\beta)^{2}(1+1/\beta)(1+\alpha
L)^{2}\mathbb{E}\left[\left\|\mathbf{w}_{r,h+1}^{i}-\mathbf{w}_{r,h}^{i}\right\|^{2}\right]\mathbb{I}[\beta\in(0,1)]+\beta^{2}\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}$
$\displaystyle\leq(1-\beta)\mathbb{E}\left[\left\|\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})-\mathbf{u}_{r,h}^{i}\right\|^{2}\right]+\frac{4\eta^{2}(1+\alpha
L)^{2}}{\beta}\mathbb{E}\left[\left\|\widehat{\Delta}_{r,h}^{i}\right\|^{2}\right]\mathbb{I}[\beta\in(0,1)]+\beta^{2}\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}.$

As done in the proof of Lemma 22, the term
$\mathbb{E}\left[\left\|\widehat{\Delta}_{r,h}^{i}\right\|^{2}\right]$ can be
upper bounded as

$\displaystyle\mathbb{E}\left[\|\widehat{\Delta}_{r,h}^{i}\|^{2}\right]$
$\displaystyle\leq\hat{\sigma}^{2}+(1+\alpha
L)^{2}L^{2}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]+2L_{F}^{2}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]+2\mathbb{E}\left[\|\nabla
F_{i}(\mathbf{w}_{r})\|^{2}\right].$

If $\alpha\leq 1/L$, we have

$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\mathbf{u}_{r,h+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})\|^{2}\right]$
$\displaystyle\leq\left(1-\beta+\frac{64\eta^{2}L^{2}}{\beta}\mathbb{I}[\beta\in(0,1)]\right)\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]+\beta^{2}\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}+\frac{16\eta^{2}}{\beta}\hat{\sigma}^{2}\mathbb{I}[\beta\in(0,1)]$
$\displaystyle\quad\quad+\frac{32L_{F}^{2}\eta^{2}}{\beta}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]\mathbb{I}[\beta\in(0,1)]+\frac{32\eta^{2}}{\beta}\left(\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right]+\gamma_{F}^{2}\right)\mathbb{I}[\beta\in(0,1)].$

Telescope it from $h=1$ to $H$ and divide both sides by $H$.

$\displaystyle\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})\|^{2}\right]$
$\displaystyle\leq\left(1-\beta+\frac{64\eta^{2}L^{2}}{\beta}\mathbb{I}[\beta\in(0,1)]\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]+\beta^{2}\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}+\frac{16\eta^{2}}{\beta}\hat{\sigma}^{2}\mathbb{I}[\beta\in(0,1)]$
$\displaystyle\quad\quad+\frac{32L_{F}^{2}\eta^{2}}{\beta}\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{w}_{r,h}^{i}-\mathbf{w}_{r}\|^{2}\right]\mathbb{I}[\beta\in(0,1)]+\frac{32\eta^{2}}{\beta}\left(\mathbb{E}\left[\|\nabla
F(\mathbf{w}_{r})\|^{2}\right]+\gamma_{F}^{2}\right)\mathbb{I}[\beta\in(0,1)].$

Applying Lemma 22 and setting
$\eta\leq\min\left\\{\frac{1}{4HL_{F}},\frac{\beta}{16L\mathbb{I}[\beta\in(0,1)]}\right\\}$
leads to

$\displaystyle\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h+1}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h+1}^{i})\|^{2}\right]$
$\displaystyle\leq\left(1-\beta+\frac{64\eta^{2}L^{2}}{\beta}(1+16\eta^{2}H^{2}L_{F}^{2})\mathbb{I}[\beta\in(0,1)]\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]$
$\displaystyle\quad\quad\quad+\frac{16\eta^{2}}{\beta}(1+32L_{F}^{2}\eta^{2}H(H-1))\left(\hat{\sigma}^{2}+2\gamma_{F}^{2}\right)\mathbb{I}[\beta\in(0,1)]+\beta^{2}\frac{\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}$
$\displaystyle\quad\quad\quad+\frac{32\eta^{2}}{\beta}\left(1+32L_{F}^{2}\eta^{2}H(H-1)\right)\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}\right]\mathbb{I}[\beta\in(0,1)]$
$\displaystyle\leq\left(1-\frac{\beta}{2}\right)\frac{1}{nH}\sum_{i=1}^{n}\sum_{h=1}^{H}\mathbb{E}\left[\|\mathbf{u}_{r,h}^{i}-\mathbf{v}_{i}(\mathbf{w}_{r,h}^{i})\|^{2}\right]\mathbb{I}[\beta\in(0,1)]$
$\displaystyle\quad\quad\quad+\frac{48\eta^{2}\left(\hat{\sigma}^{2}+2\gamma_{F}^{2}\right)}{\beta}\mathbb{I}[\beta\in(0,1)]+\frac{\beta^{2}\alpha^{2}\sigma_{G}^{2}}{|\mathcal{S}_{1}^{i}|}+\frac{96\eta^{2}}{\beta}\mathbb{E}\left[\left\|\nabla
F(\mathbf{w}_{r})\right\|^{2}\right]\mathbb{I}[\beta\in(0,1)].$