

\section{The Proposed Approach}  \label{section_framework}


In this section, we propose the {FedSGDA} framework for nonconvex-PL minimax optimization in the cross-device federated learning setting.
FedSGDA accepts a gradient estimator as input.
By incorporating {FedSGDA} with two gradient estimators, we derive two specific algorithms.


\subsection{The {FedSGDA} framework}
In each round $t = 1, \ldots, T$, {FedSGDA} proceeds in the following two phases:
\begin{enumerate}[(i)]
    \item \textbf{Gradient collection phase.} The server sends $(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$ to a subset of clients, receives local gradients from clients that successfully respond (denoted by $\mathcal{S}_t'$), and then computes the global gradient estimate $(\mathbf{u}_t, \mathbf{v}_t)$ by using some estimator $\mathcal{E}$.
    \item \textbf{Parameter update phase.} The server sends $(\mathbf{u}_t, \mathbf{v}_t)$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$ to a subset of clients that can be different from $S_t'$.
    Active clients in this subset perform $K$ local steps and sends the last iterate to the server.
    Upon receiving the local iterates from clients that successfully respond (denoted by $\mathcal{S}_t$), the server computes the new parameter $(\bar{\mathbf{x}}_{t+1}, \bar{\mathbf{y}}_{t+1})$ by averaging the local iterates.
\end{enumerate}
For ease of exposition, we assume $|\mathcal{S}_t'| = |\mathcal{S}_t| = S$ for some fixed integer $S$.

\begin{framework}[t]
    \caption{{FedSGDA} on the server.}
    \label{algorithm_fed_sgda_server}
    \begin{algorithmic}[1]
    \SUB{Input:}
    initial point $(\bar{\mathbf{x}}_{-1}, \bar{\mathbf{y}}_{-1}) = (\bar{\mathbf{x}}_0, \bar{\mathbf{y}}_0)$,
    the number of rounds $T$,
    the batch size of clients $S$, 
    and gradient estimator $\mathcal{E}$.

    \FOR{$t = 0, 1, \ldots, T - 1$}
        \STATE{\bfseries Gradient collection phase :}
        \STATE Send $(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$ to a subset of clients;
        
        \STATE Receive $(\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}), \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}))$ and $(\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t), \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t))$ from $i \in \mathcal{S}_t'$;
        \STATE Compute $(\mathbf{u}_t, \mathbf{v}_t)$ by using the estimator $\mathcal{E}$;
    
        \STATE{\bfseries Parameter update phase:}
        \STATE Send $(\mathbf{u}_t$, $\mathbf{v}_t)$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$ to a subset of clients;

        \STATE Receive $(\mathbf{x}_{t, i}^{(K)}, \mathbf{y}_{t, i}^{(K)})$ from $i \in \mathcal{S}_t$;

        \STATE $\bar{\mathbf{x}}_{t + 1} \leftarrow \frac{1}{S_t} \sum_{i \in \mathcal{S}_t} \mathbf{x}_{t, i}^{(K)}$, $\bar{\mathbf{y}}_{t + 1} \leftarrow \frac{1}{S_t} \sum_{i \in \mathcal{S}_t} \mathbf{y}_{t, i}^{(K)}$;
    \ENDFOR
\end{algorithmic}
\end{framework}






In each local iteration (line~\ref{line_local_iteration_start} to~\ref{line_local_iteration_end} of Framework~\ref{algorithm_fed_sgda_client_parameter_update}), we simultaneously take a descent step on the local primal variable $\mathbf{x}_{t, i}^{(k)}$ and an ascent step on the dual variable $\mathbf{y}_{t, i}^{(k)}$.
Instead of using local stochastic gradients as local update directions, we use the combination of local stochastic gradients and a global gradient estimate,
similar to the techniques in~\citep{johnson2013accelerating, karimireddy2020mime}.
Specifically, we define the local primal and dual update directions as 
$$
\begin{aligned}
\begin{cases}
    \mathbf{d}_{\mathbf{x}, t, i}^{(k)}
    := \nabla_{\mathbf{x}} F_i(\mathbf{x}_{t, i}^{(k)}, \mathbf{y}_{t, i}^{(k)}; \mathcal{B}_{t, i}^{(k)})
        - \nabla_{\mathbf{x}} F_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t; \mathcal{B}_{t, i}^{(k)}) + \mathbf{u}_t, \\
    \mathbf{d}_{\mathbf{y}, t, i}^{(k)}
    := \nabla_{\mathbf{y}} F_i(\mathbf{x}_{t, i}^{(k)}, \mathbf{y}_{t, i}^{(k)}; \mathcal{B}^{(k)})
       - \nabla_{\mathbf{y}} F_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t; \mathcal{B}_{t, i}^{(k)}) + \mathbf{v}_t,
\end{cases}
\end{aligned}
$$
where $F_i(\cdot, \cdot; \mathcal{B}_{t, i}^{(k)})$ is the average of local losses over the minibatch of data samples $\mathcal{B}_{t, i}^{(k)}$, i.e.,
$$
    F_i(\mathbf{x}, \mathbf{y}, \mathcal{B}) := \frac{1}{|\mathcal{B}|} \sum_{\zeta \in \mathcal{B}} F_i(\mathbf{x}, \mathbf{y}; \zeta).
$$


The steps of {FedSGDA} on the server, clients in the gradient collection phase, and clients in the parameter update phase are detailed in Framework~\ref{algorithm_fed_sgda_server}, \ref{algorithm_fed_sgda_client_gradient_collection} and~\ref{algorithm_fed_sgda_client_parameter_update}, respectively.



\subsection{Algorithms derived from {FedSGDA}}  \label{section_fedsgda_instances}
In what follows, we give two examples of the gradient estimators $\mathcal{E}$: (i) the vanilla minibatch gradient estimator, and (ii) the state-of-the-art gradient estimator for nonconvex optimization, STORM~\citep{cutkosky2019momentum}.
We refer to {FedSGDA} incorporated with these estimators as {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM}.




\paragraph{FedSGDA-MB.}
The minibatch gradient estimator simply averages the local gradients of the sampled clients. 
It is the most widely used and the simplest estimator in the stochastic optimization literature, but it usually requires a large minibatch to reduce variance.
The minibatch gradient estimator is defined as
$$
    \mathbf{u}_t = \frac{1}{S} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t), 
    \ \textrm{and} \
    \mathbf{v}_t = \frac{1}{S} \sum_{i \in \mathcal{S}_t} \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t).
$$
We note that if this estimator is used in FedSGDA, the gradient collection phase of FedSGDA can be simplified.
Specifically, the communication of $(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ and the computation of $\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ and $\nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ can be omitted.


\paragraph{FedSGDA-STORM.}
To reduce variance in stochastic gradients, Variance Reduced (VR) gradient estimators have been developed.
The STORM estimator is an efficient VR estimator, which has been successfully applied in stochastic optimization problems.
STORM leverages local gradients at the previous iterate $(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ to reduce variance.
It is constructed in a recursive manner
$$
\begin{aligned}
\begin{cases}
    \mathbf{u}_{t} ={}&
            (1 - \alpha_t) (\mathbf{u}_{t-1} - \frac{1}{S} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}))
            + \frac{1}{S} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t), \\
    \mathbf{v}_t ={}&
            (1 - \beta_t) (\mathbf{v}_{t-1} - \frac{1}{S} \sum_{i \in \mathcal{S}_t} \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}))
            + \frac{1}{S} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_{t}, \bar{\mathbf{y}}_{t}),
\end{cases}
\end{aligned}
$$
where $\alpha_t$ and $\beta_t$ are real numbers in $[0, 1]$, and $\mathbf{u}_0$ and $\mathbf{v}_0$ are computed by
$$
    \mathbf{u}_0 = \frac{1}{S} \sum_{i \in \mathcal{S}_0'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_0, \bar{\mathbf{y}}_0), 
    \ \textrm{and} \
    \mathbf{v}_0 = \frac{1}{S} \sum_{i \in \mathcal{S}_0} \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_0, \bar{\mathbf{y}}_0).
$$
Note that when $\alpha_t = \beta_t = 1$, the STORM estimator is the same as the minibatch gradient estimator. 

\begin{remark}
Many estimators other than the minibatch estimator and STORM can be incorporated into FedSGDA, e.g., SVRG~\citep{johnson2013accelerating}, averaged stochastic gradient~\citep{mokhtari2018stochastic}, SARAH/SPIDER~\citep{nguyen2017sarah, fang2018spider}.
\end{remark}



\begin{framework}[t]
    \caption{{FedSGDA} on a client $i$ in the \textbf{gradient collection phase} of the $t$-th round.}
    \label{algorithm_fed_sgda_client_gradient_collection}
    \begin{algorithmic}[1]
        \SUB{Input:}
        $(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$.

        \STATE Compute local gradients at $(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$;
        \STATE Send the above gradients to the server; 
\end{algorithmic}
\end{framework}

\begin{framework}[t]
    \caption{{FedSGDA} on a client $i$ in the \textbf{parameter update phase} of the $t$-th round.}
    \label{algorithm_fed_sgda_client_parameter_update}
    \begin{algorithmic}[1]
        \SUB{Input:}
        the number of local steps $K$,
        step sizes $\eta_{t}$ and $\gamma_{t}$,
        $(\mathbf{u}_t, \mathbf{v}_t)$ and $(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$.


        \STATE Initialize local model $\mathbf{x}_{t, i}^{(0)} \leftarrow \bar{\mathbf{x}}_t$, $\mathbf{y}_{t, i}^{(0)} \leftarrow \bar{\mathbf{y}}_t$; \
    
        \FOR{$k = 0, 1, \ldots, K - 1$}
            \STATE Sample a minibatch $\mathcal{B}_{t, i}^{(k)}$ from local data;  \label{line_local_iteration_start}
    
            \STATE $\mathbf{x}_{t, i}^{(k + 1)} \leftarrow
                \mathbf{x}_{t, i}^{(k)} - \eta_t  \mathbf{d}_{\mathbf{x}, t, i}^{(k)}$;
    
            \STATE $\mathbf{y}_{t, i}^{(k + 1)} \leftarrow
                \mathbf{y}_{t, i}^{(k)} + \gamma_t \mathbf{d}_{\mathbf{y}, t, i}^{(k)}$; \label{line_local_iteration_end}
        \ENDFOR
        \STATE Send $(\mathbf{x}_{t, i}^{(K)}, \mathbf{y}_{t, i}^{(K)})$ to the server; 
\end{algorithmic}
\end{framework}

\section{Conclusion}
We present the first algorithmic framework for general nonconvex-PL minimax problems in the cross-device federated learning setting.
Our framework FedSGDA performs multiple local SGDA steps on a random subset of clients in each round and accepts a global gradient estimator as input.  
We further propose two efficient algorithms by plugging the minibatch gradient estimator and STORM into FedSGDA.
Theoretical analysis and experimental results show the effectiveness of the proposed algorithms.  

\section{Related Work}
\textbf{Minimax Optimization.}
In recent years, minimax problems has received an increasing interest in the machine learning community~\citep{goodfellow2014generative, madry2017towards, nouiehed2019solving,rafique2018non, lin2019gradient,wang2020zeroth}.
Generally,  nonconvex-nonconcave problems are notoriously hard to solve~\citep{jin2019local}.
Hence, many existing works consider special minimax problems in which $f(\mathbf{x},  \mathbf{y})$ is nonconvex w.r.t.\ $\mathbf{x}$ and the inner maximization problems $\max_{\mathbf{y}} f(\mathbf{x}, \mathbf{y})$ are tractable, e.g., when $f(\mathbf{x}, \mathbf{y})$ satisfies the PL condition on $\mathbf{y}$ for any $\mathbf{x}$.

There are two types of algorithms for minimax problems with tractable inner maximization:
(i) GDmax-type algorithms~\citep{sanjabi2018solving, nouiehed2019solving, lin2019gradient, wang2020zeroth, luo2020stochastic, xu2020enhanced},
and (ii) GDA-type algorithms~\citep{lin2019gradient, liu2019min, wang2020zeroth, huang2020accelerated, qiu2020single, yang2020global}.
In each iteration of GDmax-type algorithms, they perform single update on the primal variable $\mathbf{x}$ and require to solve a maximization problem w.r.t. the dual variable $\mathbf{y}$.
In contrast, GDA-type algorithms  only perform both a descent step on $\mathbf{x}$ and an ascent step on $\mathbf{y}$ in each iteration.
Hence, the GDA-type methods are usually more computationally practical for real-world problems with complex structures~\citep{qiu2020single}.   

\textbf{Federated learning.}
Generally, federated learning settings can be categorized into two types:
(i) cross-silo setting, and (ii) cross-device setting.
The former focuses on a medium-scale cluster with relatively reliable network. The latter
usually involves a swarm of unreliable devices with low communication frequency and varying client availability, 
thus is more challenging.

The principal cross-device federated learning algorithm for minimization problems is FedAvg~\citep{mcmahan2017communication}, which performs multiple local SGD steps on a random subset of clients before averaging the local variables on the server in each round.
While FedAvg has succeeded  in a few applications, it may suffer from the ``client drift'' issue (i.e., on heterogeneous data, the local variables move towards the individual client optima instead of a global optimum), which can  be corrected by adopting global statistics in local updates~\citep{karimireddy2019scaffold,karimireddy2020mime}.

Recently, a few federated learning algorithms, such as Stochastic AFL~\citep{mohri2019agnostic},  FedRobust~\citep{reisizadeh2020robust}, CoDA~\citep{guo2020communication}, have been proposed for minimax problems. 
\citet{mohri2019agnostic} consider a distributionally robust optimization problem and develop a stochastic GDA-type algorithm, Stochastic AFL, which performs only one local step in each round and hence suffers from dense communication.
\citet{reisizadeh2020robust} propose to learn a centralized model robust to affine perturbations to the training data.
and develop a stochastic GDA-type algorithm named FedRobust.
CoDA is designed for the AUC maximization problem over a network, which repeatedly solves strongly-convex-strongly-concave minimax subproblems~\citep{guo2020communication}.
However, they require either dense communications or participation of all clients in each round and thus do not apply to the cross-device federated learning setting.

\section{Introduction}  \label{section_introduction}
\textit{Cross-device federated learning}~\citep{kairouz2019advances} is a distributed machine learning setting that has received significant interest recently.
In this setting, numerous clients (mobile or IoT devices) collaborate to train a machine learning model under the coordination of a central server without transferring client's local data to the server.
In contrast with conventional centralized distributed learning, cross-device federated learning has the following typical constraints:
(i) the communication between the server and a client is usually slow and unreliable;
(ii) only a fraction of clients are available for training at any time and each client will likely participate at most once in the training procedure;
(iii) data distributions among clients are generally different.

While existing work in cross-device federated learning focuses on minimization problems,
there is increasing awareness that machine learning systems are embedded in real-world settings subject to game-theoretic constraints.
This is calling for renewed interest on more complicated minimax problems.
A typical example is AUC (area under the ROC curve) maximization, which is an important metric for measuring the performance of a classification model on imbalanced data and also equivalent to a minimax problem ~\citep{ying2016stochastic, liu2020stochastic}. However, directly maximizing AUC over a network of clients will suffer from a significant or even infeasible blow-up  in communication costs~\citep{shen2018towards, guo2020communication}, which precludes AUC maximization from the cross-device federated learning setting.
Similar cases can also be found in generative model learning~\citep{goodfellow2014generative} and robust adversarial learning~\citep{namkoong2016stochastic,madry2017towards}.


In this paper, we consider the problem of training nonconvex-PL minimax models under the constraints imposed by the cross-device federated setting, i.e.,
\begin{equation}  \label{eq_problem}
    \min_{\mathbf{x} \in \mathbb{R}^{p}} \max_{\mathbf{y} \in \mathbb{R}^{q}} 
    \left\{f(\mathbf{x}, \mathbf{y}) := \mathbb{E}_{i \sim \mathcal{D}} [f_i(\mathbf{x}, \mathbf{y})] \right\},
\end{equation}
where $f(\cdot, \mathbf{y})$ is nonconvex for any $\mathbf{y}$,
$f(\mathbf{x}, \cdot)$ satisfies the Polyak-{\L}ojasiewicz (PL) condition (see Assumption~\ref{assumption_pl}) for any $\mathbf{x}$,
$\mathcal{D}$ denotes the distribution of numerous unreliable clients that have limited communication bandwidth.
Each $f_i$ is defined as $f_i(\mathbf{x}, \mathbf{y}) := \frac{1}{n_i} \sum_{j = 1}^{n_i} F_i(\mathbf{x}, \mathbf{y}; \zeta_{i, j})$, where $\{\zeta_{i, 1}, \ldots, \zeta_{i, n_i}\}$ denotes the set of local data samples held by client $i$ and $F_i(\cdot, \cdot; \zeta_{i, j})$ denotes the local loss function associated with the data sample $\zeta_{i, j}$.

Algorithms for general minimization problems in federated learning are not directly applicable to~\eqref{eq_problem}, although a minimax problem can be rephrased as a minimization problem of the form $\min_{\mathbf{x} \in \mathbb{R}} \Phi(\mathbf{x}) := \max_{\mathbf{y} \in \mathbb{R}^q} f(\mathbf{x}, \mathbf{y})$.
The reason is that the global loss function $\Phi$ in general cannot be decomposed into the expectation of local loss functions over clients, while federated learning algorithms for general minimization problems highly rely on the decomposibility of the global loss function.
On the other hand, algorithms designed for minimax problems in the single-machine or conventional distributed setting cannot be easily extended to the cross-device federated learning setting.
The major barriers are summarized as follows.
\begin{enumerate}[(i)]
    \item \textbf{Communication bottleneck.} Directly parallelizing an algorithm designed for the single-machine setting over multiple clients usually requires global communication in each local update, which entails frequent communication.
    This will significantly deteriorate the performance when the network is slow and unreliable.
    \item \textbf{Data heterogeneity.}
    Since the local data distribution on each client is different from the overall data distribution, each client cannot access unbiased data samples from the global distribution.
    This together with the limited communication bandwidth makes it difficult to train a single global model~\citep{kairouz2019advances, li2019convergence}.
    \item \textbf{Low client availability.} Since each client will likely participate in at most one round of training, 
    many strategies that are effective in conventional centralized distributed learning do not apply to cross-device federated learning, e.g., synchronization over all clients and gradient compression with error compensation~\citep{stich2018sparsified}.
\end{enumerate}

\textbf{Contribution.} To overcome the above barriers, we develop a novel algorithmic framework named {Federated Stochastic Gradient Descent-Ascent} ({FedSGDA}) for the problem~\eqref{eq_problem}.
To mitigate the communication bottleneck, {FedSGDA} performs multiple local stochastic gradient descent-ascent steps on clients before communication to the server in each round. 
In addition, only a fraction of clients participate in one round of the optimization procedure to satisfy the low client availability constraint.
To reduce the bias of local updates introduced by data heterogeneity, {FedSGDA} utilizes both global gradient estimates and local stochastic gradients to compute local update directions.
By incorporating {FedSGDA} with two global gradient estimators--—the minibatch estimator and STORM~\citep{cutkosky2019momentum}, we develop two specific algorithms named {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM}, respectively.

We provide theoretical analysis and establish convergence guarantees of the proposed algorithms.
Our theoretical results are listed as follows.
\begin{enumerate}[(i)]
    \item FedSGDA-MB achieves an $\mathcal{O}(1/T + 1/S)$ convergence rate, where $T$ denotes the number of communication rounds and $S$ denotes the number of clients sampled in each round.
    If $S$ is large, FedSGDA-MB converges to a close neighborhood of a stationary point.

    \item {{\FrameworkAbbr}-STORM} achieves a convergence rate of $\tilde{\mathcal{O}}(1/T + 1/(T^{2/3} S^{1/3}))$\footnote{The symbol $\tilde{\mathcal{O}}()$ suppresses a logarithmic factor in $T$.},
    which implies the convergence to a stationary point for any positive $S$ when $T \rightarrow +\infty$.
\end{enumerate}
Experimental results on synthetic and real datasets demonstrate the effectiveness of the proposed algorithms.

\section{Convergence Analysis}  \label{section_analysis}
We first provide several key lemmas for the FedSGDA framework with an arbitrary gradient estimator.
Then we establish convergence guarantees of FedSGDA-MB and FedSGDA-STORM.
We note that in our analysis, each client in $\mathcal{S}_t'$ or $\mathcal{S}_t$ in Framework~\ref{algorithm_fed_sgda_server} is assumed to be randomly drawn from the underlying distribution $\mathcal{D}$, which is a common assumption in the federated learning literature~\citep{li2018federated, li2019convergence, karimireddy2020mime}.
All missing proofs are deferred to the Appendix due to the limit of space.

For ease of notation, we omit the subscript $t, i$ of the iterates in FedSGDA and denote $\mathbf{x}_{t, i}^{(k)}$ (resp., $\mathbf{y}_{t, i}^{(k)}$) as $\mathbf{x}^{(k)}$ (resp., $\mathbf{y}^{(k)}$).
We denote the error of gradient estimations as
$$
\begin{aligned}[b]
\begin{cases}
    \varepsilon_{\mathbf{x}, t} =& \mathbb{E}[\|\mathbf{u}_t - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_{t}, \bar{\mathbf{y}}_{t})\|^2], \\
    \varepsilon_{\mathbf{y}, t} =& \mathbb{E}[\|\mathbf{v}_t - \nabla_{\mathbf{y}} f(\bar{\mathbf{x}}_{t}, \bar{\mathbf{y}}_{t})\|^2],
\end{cases}
\end{aligned}
$$
where the expectation is taken over all randomness. 
For $(\mathbf{x}, \mathbf{y}) \in \mathbb{R}^{p} \times \mathbb{R}^{q}$, we denote $\mathbf{y}^*(\mathbf{x})$ as the projection of $\mathbf{y}$ onto the solution set of $\max_{\mathbf{y}} f(\mathbf{x}, \mathbf{y})$.
We denote $\mathbf{z} := (\mathbf{x}, \mathbf{y}) \in \mathbb{R}^{p \times q}$, $\tilde{L} = 2 (L_{\Phi} + L_f)$.

The next lemma shows that the expected loss value $\mathbb{E}[\Phi(\mathbf{x}^{(k)})]$ decreases as $k$ increases if error terms are controlled in a desirable level. 
\begin{lemma}[One-iteration descent]  \label{lemma_descent}
Suppose that $\eta_t \le \frac{1}{L_{\Phi}}$.
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, $\forall k \in \{0, \ldots, K - 1\}$, we have
$$
\begin{aligned}[b]
    \mathbb{E}[ \Phi(\mathbf{x}^{(k+1)}) - \Phi(\mathbf{x}^{(k)})]
    \le& - \frac{\eta_t}{2} \mathbb{E}[ \| \nabla \Phi(\mathbf{x}^{(k)})\|^2]
        + \frac{3 L_{12}^2}{4 \mu}  \eta_t \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})] \\
        &+ 6 L_f^2 \eta_t \mathbb{E}[ \| \mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2]
        + \frac{3}{2} \eta_t \varepsilon_{\mathbf{x}, t}.
\end{aligned}
$$
\end{lemma}

We provide upper bounds of the error terms $\mathbb{E}[\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2]$ and $\mathbb{E}[\Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]$ in the next two lemmas.
\begin{lemma}  \label{lemma_iterates_recursion}
Consider the FedSGDA framework.
Suppose that $\eta_t \le \frac{1}{8 \sqrt{3} K L_f}$ and $\gamma_t \le \frac{1}{8 \sqrt{3} K L_f}$.
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, for any $k \in \{0, \ldots, K - 1\}$, we have
$$
\begin{aligned}[b]
    \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_{t}\|^2 ]
    \le{}& (1 + \frac{1}{K}) \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 6 K \eta_t^2 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ] \\
        &+ 6 K \gamma_t^2 \mathbb{E}[ \|\nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ]
        + 12 K \eta_t^2 \varepsilon_{\mathbf{x}, t}
        + 12 K \gamma_t^2 \varepsilon_{\mathbf{y}, t}.
\end{aligned}
$$
\end{lemma}


\begin{lemma}  \label{lemma_max_residue}
Consider the FedSGDA framework.
Suppose that $\eta_t \le \min\{\frac{1}{2 L_1}, \frac{1}{42 L_{21}}\}$ and $\gamma_t \le \min\{\frac{1}{2 L_2}, \frac{1}{L_{21}}, \frac{1}{\mu}\}$.
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, for any $k \in \{0, \ldots, K - 1\}$, it holds that
$$
\begin{aligned}[b]
    \mathbb{E}[ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{z}^{(k+1)})]
    \le{}& (1 - \frac{\mu \gamma_t}{2})(1 + \frac{3 L_{12}^2}{4 \mu} \eta_t) \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})] \\
        &+ \Big(10 \eta_t (1 - \frac{\mu \gamma_t}{2}) + 6 \gamma_t + 21 L_{21}^2 \eta_t^2 \gamma_t \Big) L_f^2 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ] \\
        &+ \big( (1 - \frac{\mu \gamma_t}{2}) \frac{5 \eta_t}{2} + \frac{21}{4} L_{21}^2 \eta_t^2 \gamma_t \big) \varepsilon_{\mathbf{x}, t}
        + \frac{3 \gamma_t}{2} \varepsilon_{\mathbf{y}, t} \\
        &- \frac{\eta_t}{2} (1 - \frac{\mu \gamma_t}{2}) \mathbb{E}[ \| \nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        - \frac{\eta_t}{8} \mathbb{E}[ \| \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ] \\
        &- \frac{\gamma_t}{8} \mathbb{E}[ \| \nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ]
        - \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2].
\end{aligned}
$$
\end{lemma}


Lemma~\ref{lemma_max_residue} provides a tight upper bound of this function value gap $\mathbb{E}[\Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]$, which involves negative terms related to $\mathbb{E}[\|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2]$, $\mathbb{E}[ \| \nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ]$, $\mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]$, and $\mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2]$.






\subsection{Convergence of FedSGDA-MB}  \label{section_sgd}
To analyze the convergence of FedSGDA-MB, we construct the following potential function
\begin{equation}  \label{eq_potential_simple_definition}
\begin{aligned}[b]
    \mathcal{L}^{(k)} :={}& \mathbb{E}[ 2 \Phi(\mathbf{x}^{(k)}) - \Phi(\mathbf{x}^*) - f(\mathbf{z}^{(k)})]
        + \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k} \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2 ].
\end{aligned}
\end{equation}
We note that $\mathcal{L}^{(k)} \ge 0$ since $\Phi(\mathbf{x}^{(k)}) \ge \Phi(\mathbf{x}^*)$ and $\Phi(\mathbf{x}^{(k)}) \ge f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})$.
If $\mathcal{L}^{(k)} = 0$, then $\mathbf{x}^{(k)}$ is a minimizer of $\Phi$ and $\mathbf{y}^{(k)} \in \mathrm{argmax}_{\mathbf{y} \in \mathbb{R}^{q}} f(\mathbf{x}^{(k)}, \mathbf{y})$.

The following lemma shows that the potential function $\mathcal{L}^{(k)}$ decreases as $k$ increases if the error terms $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$ are properly controlled.
Note that Lemma~\ref{lemma_potential_recursion_decay_simple} relies heavily on Lemma~\ref{lemma_descent}, \ref{lemma_iterates_recursion}, and \ref{lemma_max_residue} to bound the terms in $\mathcal{L}^{(k+1)} - \mathcal{L}^{(k)}$.
\begin{lemma}  \label{lemma_potential_recursion_decay_simple}
Consider FedSGDA-MB.
Suppose that
$$
\begin{aligned}[b]    
\eta_t
    \le& \min\left\{
    \frac{1}{4 L_1},
    \frac{1}{L_{\Phi}},
    \frac{1}{42 L_{21}},
    \gamma_t,
    \frac{\mu^2}{3 L_{12}^2} \gamma_t
\right\}, \\
\gamma_t
    \le& \min\left\{
    \frac{1}{2 L_2},
    \frac{1}{\mu},
    \frac{1}{L_{21}},
    \frac{1}{192 \tilde{L} K}
\right\}.
\end{aligned}
$$
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, we have
$$
\begin{aligned}[b]
    \mathcal{L}^{(k+1)}
    \le{}& \mathcal{L}^{(k)}
        - \frac{3 \eta_t}{4} \mathbb{E}[ \|\nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        + \frac{35 \eta_t}{8} \varepsilon_{\mathbf{x}, t}
        + \frac{7 \gamma_t}{4} \varepsilon_{\mathbf{y}, t}
        - \frac{1}{8 \gamma_t} \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \mathbf{z}^{(k)}\|^2 ].
\end{aligned}
$$
\end{lemma}

Based on the result in Lemma~\ref{lemma_potential_recursion_decay_simple}, we provide an upper bound of the average squared gradient norm in the $t$-th round in Lemma~\ref{lemma_potential_recursion_sum}.

\begin{lemma}  \label{lemma_potential_recursion_sum}
Under the conditions in Lemma~\ref{lemma_potential_recursion_decay_simple}, we have
\begin{equation}  \label{eq_potential_recursion_sum}
\begin{aligned}[b]
    \frac{3 \eta_t}{4 S K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le{}& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_t) - f(\bar{\mathbf{z}}_t)]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})] \\
        &+ \frac{35 \eta_t}{8} \varepsilon_{\mathbf{x}, t}
        + \frac{7 \gamma_t}{4} \varepsilon_{\mathbf{y}, t}
        - \frac{1}{8 K^2 \gamma_t} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ].
\end{aligned}
\end{equation}
\end{lemma}

It remains to bound the gradient estimation errors $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$.
By Assumption~\ref{assumption_bgd}, the gradient approximation errors $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$ for the minibatch gradient estimator can be simply bounded by $\sigma_1^2 / S$ and $\sigma_2^2 / S$, respectively.
Combining these bounds and Lemma~\ref{lemma_potential_recursion_sum}, we establish the convergence rate of FedSGDA-MB in the following theorem.
\begin{theorem}  \label{theorem_fed_sgda_mb}
Consider FedSGDA-MB.
Denote $C_0 := 2 \Phi(\mathbf{x}_0) - \Phi(\mathbf{x}^*) - f(\mathbf{x}_0, \mathbf{y}_0)$, where $\mathbf{x}^* \in \mathop{\rm argmin}_{\mathbf{x} \in \mathbb{R}^{p}} \Phi(\mathbf{x})$.
Suppose that $C_0 < \infty$,
$
\eta_t
    = \eta
    = \min\left\{
    \frac{1}{4 L_1},
    \frac{1}{L_{\Phi}},
    \frac{1}{42 L_{21}},
    \gamma,
    \frac{\mu^2}{3 L_{12}^2} \gamma
    \right\}
$, and
$
\gamma_t
    = \gamma
    = \min\left\{
    \frac{1}{2 L_2},
    \frac{1}{\mu},
    \frac{1}{L_{21}},
    \frac{1}{192 \tilde{L} K}
\right\}
$.
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, we have
\begin{equation}  \label{eq_rate_fedsgda_mb}
\begin{aligned}[b]
    \frac{1}{S K T} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le{}& \mathcal{O} \Big(
        \frac{C_0}{K T} \big( L_{\Phi} + L_{21} + \frac{L_{12}^2}{\mu^2} (L_2 + \mu + L_{21} + \tilde{L} K) \big) \\
        &+ \frac{\sigma_1^2}{S} + \frac{\sigma_2^2}{S} \big( \frac{L_{\Phi}}{\tilde{L} K} + \frac{L_{21}}{\tilde{L} K} + \frac{L_{12}^2}{\mu^2} \big)
    \Big).
\end{aligned}
\end{equation}
\end{theorem}


\begin{remark}  \label{remark_mb}
    Theorem~\ref{theorem_fed_sgda_mb} implies that FedSGDA-MB converges to a close neighborhood of a stationary point as $T \rightarrow +\infty$ if we choose a large $S$.
    Furthermore, it can be observed that the r.h.s. of~\eqref{eq_rate_fedsgda_mb} decreases as $K$ increases.
    However, when $K$ reaches certain threshold, $\frac{C_0 L_{12}^2 \tilde{L}}{\mu^2} T$ or $\frac{\sigma_2^2 L_{12}^2}{\mu^2 S}$ dominates the convergence rate, especially when $\mu$ is tiny or $\sigma_2^2$ is large.
\end{remark}














\subsection{Convergence of FedSGDA-STORM}  \label{section_storm}
To analyze the convergence of FedSGDA-STORM, we construct the following potential function.
$$
\begin{aligned}
    \tilde{\mathcal{L}}^{(k)}
    :={}& \mathbb{E}[ 2 \Phi(\mathbf{x}^{(k)}) - \Phi(\mathbf{x}^*) - f(\mathbf{z}^{(k)})]
        + \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{(K - k)} \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2].
       
       
\end{aligned}
$$
Note that the potential function $\tilde{\mathcal{L}}^{(k)}$ is more complicated than the one used by FedSGDA-MB.
It involves two extra error terms $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$ which can be bounded as follows. 
\begin{lemma} \label{lemma_gradient_error_storm}
The gradient estimation errors $\varepsilon_{\mathbf{x}, t+1}$ and $\varepsilon_{\mathbf{y}, t+1}$ of the STORM estimator can be bounded by
$$
\begin{aligned}[b]
    \varepsilon_{\mathbf{x}, t+1}
    \le& (1 - \alpha_{t+1})^2 \varepsilon_{\mathbf{x}, t}
        + \frac{2 \alpha_{t+1}^2 \sigma_1^2}{S}
        + \frac{8 (1 - \alpha_{t+1})^2 L_f^2}{S} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_{t}\|^2 ], \\
    \varepsilon_{\mathbf{y}, t+1}
    \le& (1 - \beta_{t+1})^2 \varepsilon_{\mathbf{y}, t}
        + \frac{2 \beta_{t+1}^2 \sigma_2^2}{S}
        + \frac{8 (1 - \beta_{t+1})^2 L_f^2}{S} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_{t}\|^2 ].
\end{aligned}
$$
\end{lemma}


According to the result in Lemma~\ref{lemma_gradient_error_storm}, we provide the following lemma to show that the potential function $\tilde{\mathcal{L}}^{(k)}$ decreases as $k$ increases.

\begin{lemma}  \label{lemma_potential_recursion_decay_storm}
Consider FedSGDA-STORM.
Define $\tilde{L} = 2 (L_{\Phi} + L_f)$.
Suppose that $\alpha_t = \beta_t = 2304 \tilde{L}^2 K^2 \gamma_t^2$,
$
\eta_t
= \min\left\{
\frac{1}{4 L_{1}},
\frac{1}{L_{\Phi}},
\frac{1}{42 L_{21}},
\gamma_t,
\frac{\mu^2}{3 L_{12}^2} \gamma_t
\right\}
$, and 
$
\gamma_t
    = \min\left\{
    \frac{1}{2 L_2},
    \frac{1}{\mu},
    \frac{1}{L_{21}},
    \frac{1}{192 \tilde{L} K},
    \frac{S^{1/3}}{\tilde{L} K (t+1)^{1/3}}
\right\}
$.
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, we have
$$
\begin{aligned}[b]
    \tilde{\mathcal{L}}^{(k+1)}
        + \frac{48 \eta_t}{\alpha_t} \varepsilon_{\mathbf{x}, t+1}
        + \frac{24 \gamma_t}{\alpha_t} \varepsilon_{\mathbf{y}, t+1}
    \le{}& \tilde{\mathcal{L}}^{(k)}
        + \frac{48 \eta_t}{\alpha_t} \varepsilon_{\mathbf{x}, t}
        + \frac{24 \gamma_t}{\alpha_t} \varepsilon_{\mathbf{y}, t}
        - \frac{3 \eta_t}{4} \mathbb{E}[\|\nabla \Phi(\mathbf{x}^{(k)})\|^2] \\
        &+ 96 \eta_{t+1} \alpha_{t+1} \frac{\sigma_1^2}{S}
        + 48 \gamma_{t+1} \alpha_{t+1} \frac{\sigma_2^2}{S} \\
        &+ \frac{1}{8 S K^2 \gamma_t} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ]
        - \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2].
\end{aligned}
$$
\end{lemma}



Now, we establish the convergence of FedSGDA-STORM based on the one-iteration descent result in Lemma~\ref{lemma_potential_recursion_decay_storm}.
\begin{theorem}  \label{theorem_fed_sgda_storm}
Consider FedSGDA-STORM.
Define $\tilde{L} = 2 (L_{\Phi} + L_f)$.
Suppose that $\alpha_t = \beta_t = 2304 \tilde{L}^2 K^2 \gamma_t^2$
$
\eta_t
    = \min\left\{
    \frac{1}{4 L_{1}},
    \frac{1}{L_{\Phi}},
    \frac{1}{42 L_{21}},
    \gamma_t,
    \frac{\mu^2}{3 L_{12}^2} \gamma_t
\right\}
$, and
$\gamma_t
    = \min\left\{
    \frac{1}{2 L_2},
    \frac{1}{\mu},
    \frac{1}{L_{21}},
    \frac{1}{192 \tilde{L} K},
    \frac{S^{1/3}}{\tilde{L} K (t+1)^{1/3}}
\right\}
$.
Let $C_0 := \Phi(\mathbf{x}_0) - \Phi(\mathbf{x}^*) + \Phi(\mathbf{x}_0) - f(\mathbf{x}_0, \mathbf{y}_0)$ and assume that $C_0 < \infty$, where $\mathbf{x}^* \in \mathop{\rm argmin}_{\mathbf{x} \in \mathbb{R}^{p}} \Phi(\mathbf{x})$.
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl} and the conditions in Lemma~\ref{lemma_potential_recursion_decay_storm}, we have
$$
\begin{aligned}[b]
    \frac{1}{S K T} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le& \mathcal{O} \Big( \Big(
        \frac{1}{T} \big( L_{\Phi} + \frac{L_{12}^2 (L_2 + L_{21} + \mu + \tilde{L} K)}{\mu^2} \big)
        + \frac{L_{12}^2 \tilde{L} K }{T^{2/3} S^{1/3} \mu^2} \Big) \\
    & \times \Big( \frac{C_0}{K}
        + (\sigma_1^2 + \sigma_2^2) \big( \frac{L_2 + \mu + L_{21}}{\tilde{L}^2 K^2} + \frac{\log T}{\tilde{L} K} \big)
    \Big) \Big).
\end{aligned}
$$
\end{theorem}

\begin{remark}  \label{remark_storm}
FedSGDA-STORM converges to a stationary point as $T \rightarrow +\infty$ for arbitrary $S$.
    Similar to the result of FedSGDA-MB, the error bound in Theorem~\ref{theorem_fed_sgda_storm} decreases as $K$ increases
    and  
    $\frac{C_0}{K} (\frac{L_{12}^2 \tilde{L} K}{\mu^2 T} + \frac{L_{12}^2 \tilde{L} K}{\mu^2 T^{2/3} S^{1/3}})$
    or $(\sigma_1^2 + \sigma_2^2)\frac{\log T L_{12}^2 \tilde{L}}{\tilde{L} T}$
    dominates the convergence rate when $K$ reaches certain threshold.
\end{remark}

\section{Experiments}

In this section, we conduct experiments to demonstrate the effectiveness of the proposed algorithms FedSGDA-MB and FedSGDA-STORM.
We consider two problems: (i) a simple nonconvex-PL minimax problem on synthetic data, and (ii) an AUC maximization problem on real-world dataset.
For the first problem, the computations are performed on a desktop with Intel i7-6800K (3.40GHz) CPU and 64GB memory, where each client is assigned to a single process.
For the second problem, we use a cluster of $5$ workstations each with $4$ Geforce RTX 2080Ti GPUs and 378GB memory.
The algorithms were implemented in Pytorch.
Throughout the experiments, we use the same random seed for all algorithms.
We report the performance of the algorithms with respect to the number of communication rounds, since communication is often the major bottleneck.






\subsection{Simulation}  \label{section_experiment_toy}
In the simulation, we run our algorithms on a cluster with $500$ clients and one server.
We let $\mathcal{D}$ in~\eqref{eq_problem} be the uniform distribution over the $500$ clients.  
Each client holds a local dataset $\mathcal{Z}_i = \{\zeta_{i, j}\}_{j=1}^{n_i}$, where each data point $\zeta_{i, j}$ is of the form $\zeta_{i, j} = \{(\mathbf{a}_{i, j}, \mathbf{b}_{i, j}, \mathbf{c}_{i, j}) \in \mathbb{R}^p \times \mathbb{R}^p \times \mathbb{R}^p\}_{j=1}^{n_i}$.
The component function $F_i(\mathbf{x}, \mathbf{y}; \zeta_{i, j})$ is chosen as
$$
\begin{aligned}[b]
    F_i(\mathbf{x}, \mathbf{y}; \zeta_{i, j})
    ={}& \sum_{\ell = 1}^p \Big( 1 - \exp(\frac{-([\mathbf{x}]_{\ell} - [\mathbf{a}_{i, j}]_{\ell})^2}{2 \nu}) \Big) \\
        &+ (\mathbf{x} - \mathbf{a}_{i, j})^{\intercal} \mathbf{b}_{i, j} \mathbf{b}_{i, j}^{\intercal} (\mathbf{y} - \mathbf{c}_{i, j}) - \frac{\mu}{2} \| \mathbf{y} - \mathbf{c}_{i, j}\|^2,
\end{aligned}
$$
where $\mathbf{x}, \mathbf{y} \in \mathbb{R}^p$,
$[\mathbf{x}]_{\ell}$ denotes the $\ell$-th coordinate of $\mathbf{x}$,
and $\mu$ and $\nu$ are positive real numbers.
Our goal is to solve the global problem~\eqref{eq_problem} with $f_i(\mathbf{x}, \mathbf{y}) = \frac{1}{n_i} \sum_{j=1}^{n_i} F_i(\mathbf{x}, \mathbf{y}; \zeta_{i, j})$.
It is easy to verify that $F_i(\cdot, \cdot; \zeta_{i, j})$ satisfies Assumption~\ref{assumption_smooth}. 
In addition, since $F_i(\mathbf{x}, \cdot; \zeta_{i, j})$ is $\mu$-strongly concave for any $\mathbf{x} \in \mathbb{R}^p$, it satisfies Assumption~\ref{assumption_pl} with constant $\mu$ (see, e.g.,~\cite{karimi2016linear}).
It follows that $f_i$ and $f = \mathbb{E}_{i \sim \mathcal{D}} [f_i]$ also satisfies Assumption~\ref{assumption_pl} with constant $\mu$.

We generate a synthetic dataset as follows.
The dimensionality $p$ is set to $100$.
For each local dataset $\mathcal{Z}_i$, we set $n_i = 100$ and sample the local data points $\{(\mathbf{a}_{i, j}, \mathbf{b}_{i, j}, \mathbf{c}_{i, j})\}_{j=1}^{100}$ from $\mathcal{N}(\mathbf{w}_i, 0.1 \cdot \mathbf{I})$,
where $\mathbf{w}_i \in \mathbb{R}^{300}$ is sampled from the multidimensional Gaussian distribution $\mathcal{N}(\mathbf{0}, 0.5 \cdot \mathbf{I})$.
Both the constants $\nu$ and $\mu$ are set to $1$.
One can verify that Assumption~\ref{assumption_bgd} is satisfied for $f$ with this synthetic dataset.

In this experiment, we report the squared norm of the gradient $\nabla \Phi(\bar{\mathbf{x}}_t)$ defined in~\eqref{eq_outer_gradient}, which can be computed in a closed form.
We did a grid-search on hyperparameters for the algorithms and select the best hyperparameters with the lowest gradient norm when terminated.
The range of the tuning values is listed in Appendix~\ref{section_specification_of_hyperparameter}.
The initial variable $(\bar{\mathbf{x}}_0, \bar{\mathbf{y}}_0)$ is set to $\mathbf{0}$. 

To show the influence of $S$ (corresponding to the client availability) on the performance, we fix $K$ and $T$, and vary $S$ in $\{5, 10, 20\}$.
The experimental results with different $S$ are presented in Figure~\ref{figure_toy}.
It can be observed that as $S$ increases, both FedSGDA-MB and FedSGDA-STORM converge faster.
Moreover,  FedSGDA-STORM outperforms FedSGDA-MB for all $S$.
We can also see that the squared gradient norm of FedSGDA-MB first decreases and then oscillates around a point.
In contrast, the squared gradient norm of FedSGDA-STORM decreases as $t$ increases.
This corroborates our theoretical results in the previous section.














\subsection{AUC Maximization}  \label{section_experiment_auc}
In the second experiment, we consider $\ell_2$-relaxed AUC maximization problem,
which is equivalent to a minimax problem~\citep{ying2016stochastic, liu2020stochastic}.
AUC is a widely used performance metric for binary classification models on imbalanced data.
For completeness, we refer the reader to Appendix~\ref{section_auc_minimax_formulation} for the $\ell_2$-relaxed AUC maximization problem and its minimax formulation.
Similar to the simulation, we run our algorithms on a cluster with $500$ clients and one server and the client distribution $\mathcal{D}$ is the uniform distribution over the $500$ clients.  


We conduct this experiment on the CIFAR-10\footnote{\url{https://www.cs.toronto.edu/~kriz/cifar.html}} dataset, which consists of $50,000$ training images and $10,000$ test images of size $32 \! \times \! 32 \! \times \! 3$ in $10$ classes.
To make the data imbalanced, we split the original dataset into two classes by treating the original ``airplane'' class as the positive class and the remaining as negative.
In addition, we rescale each pixel value to the range $[-1, 1]$.
The training data is equally partitioned into $500$ clients, where the $i$-th client contains the $(100 i - 99)$-th to $(100 i)$-th data points in the original training dataset.
Following~\cite{mcmahan2017communication}, we use a convolutional neural network\footnote{\url{https://www.tensorflow.org/tutorials/images/cnn}} as the classification model.
The parameters of this network are randomly initialized by default in Pytorch. 

We report the AUC value on the test dataset versus the communication rounds.
We did a grid-search on hyperparameters for the algorithms and select the best hyperparameters with the highest test AUC when terminated.
The range of the tuning values is listed in Appendix~\ref{section_specification_of_hyperparameter}.

In this experiment, we fix $S$ and $T$ and run all algorithms with different $K$ from $\{1, 10, 20\}$ to show the effect of local updates.
Besides, we also vary the training minibatch size in each local iteration to see the influence of the minibatch size for FedSGDA with different estimators.
In Figure~\ref{figure_auc}, we show the experimental results for different $K$ and minibatch sizes.
We can see from the left and middle columns that both FedSGDA-MB and FedSGDA-STORM with multiple local iterations achieve significantly higher AUC value than with only $1$ local iteration given the same $S$ and $T$.
Both of them achieve the best performance when $K = 20$.
In addition, the right column shows that FedSGDA-STORM outperforms FedSGDA-MB when $K$ is set to $20$.
Furthermore, it can be observed that the performance of FedSGDA-MB varies with different minibatchs, while FedSGDA-STORM is insensitive to the minibatch size.
This also indicates the effectiveness of VR.

\section{Notation and Preliminaries}  \label{section_preliminaries}
\textbf{Notation.} We use bold lowercase symbols (e.g., $\mathbf{x}$) to denote vectors.
For a function $f(\cdot, \cdot): \mathbb{R}^{p} \times \mathbb{R}^{q} \rightarrow \mathbb{R}$, $\nabla_{\mathbf{x}} f$ and $\nabla_{\mathbf{y}} f$ denote its partial derivatives with respect to the first and second variables, respectively.
The Euclidean norm for a vector $\mathbf{x}$ is denoted by $\|\mathbf{x}\|$.
For problem~\eqref{eq_problem}, we denote $\Phi(\mathbf{x}) := \max_{\mathbf{y} \in \mathbb{R}^{q}} f(\mathbf{x}, \mathbf{y})$.

We make the following common assumptions.

\begin{assumption}[Bounded gradient dissimilarity]  \label{assumption_bgd}
There exists $\sigma_1 > 0$ and $\sigma_2 > 0$ such that $\forall \mathbf{x} \in \mathbb{R}^{p}, \ \forall \mathbf{y} \in \mathbb{R}^{q}$,
$$
\begin{aligned}[b]
\begin{cases}
    \mathbb{E}_{i \sim \mathcal{D}} [\|\nabla_{\mathbf{x}} f_i(\mathbf{x}, \mathbf{y}) - \nabla_{\mathbf{x}} f(\mathbf{x}, \mathbf{y}) \|^2] \le& \sigma_1^2, \\
    \mathbb{E}_{i \sim \mathcal{D}} [\|\nabla_{\mathbf{y}} f_i(\mathbf{x}, \mathbf{y}) - \nabla_{\mathbf{y}} f(\mathbf{x}, \mathbf{y}) \|^2] \le& \sigma_2^2.
\end{cases}
\end{aligned}
$$
\end{assumption}

\begin{assumption}[Lipschitz continuous gradient]  \label{assumption_smooth}
There exists $L_1, L_2, L_{12}, L_{21} > 0$ such that for any $\zeta \in \{\zeta_{i, 1}, \ldots, \zeta_{i, n_i}\}$, $\mathbf{x}, \mathbf{x}_1, \mathbf{x}_2 \in \mathbb{R}^{p}$ and $\mathbf{y}, \mathbf{y}_1, \mathbf{y}_2 \in \mathbb{R}^{q}$,
$$
\begin{aligned}[b]
\begin{cases}
    \| \nabla_{\mathbf{x}} F_i(\mathbf{x}_1, \mathbf{y}; \zeta) - \nabla_{\mathbf{x}} F_i(\mathbf{x}_2, \mathbf{y}; \zeta) \|
    \le& L_1 \| \mathbf{x}_1 - \mathbf{x}_2 \|, \\
    \| \nabla_{\mathbf{x}} F_i(\mathbf{x}, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{x}} F_i(\mathbf{x}, \mathbf{y}_2; \zeta) \|
    \le& L_{12} \| \mathbf{y}_1 - \mathbf{y}_2 \|, \\
    \| \nabla_{\mathbf{y}} F_i(\mathbf{x}_1, \mathbf{y}; \zeta) - \nabla_{\mathbf{y}} F_i(\mathbf{x}_2, \mathbf{y}; \zeta) \|
    \le& L_{21} \| \mathbf{x}_1 - \mathbf{x}_2 \|, \\
    \| \nabla_{\mathbf{y}} F_i(\mathbf{x}, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{y}} F_i(\mathbf{x}, \mathbf{y}_2; \zeta) \|
    \le& L_2 \| \mathbf{y}_1 - \mathbf{y}_2 \|.
\end{cases}
\end{aligned}
$$
\end{assumption}

\begin{assumption}[Polyak-{\L}ojasiewicz (PL) condition]  \label{assumption_pl}
There exists a constant $\mu > 0$ such that $\forall \mathbf{x} \in \mathbb{R}^{p}, \mathbf{y} \in \mathbb{R}^{q}$,
\begin{equation}
    \frac{1}{2 \mu} \| \nabla_{\mathbf{y}} f(\mathbf{x}, \mathbf{y}) \|^2
    \ge \max_{\mathbf{y}' \in \mathbb{R}^{q}} f(\mathbf{x}, \mathbf{y}') - f(\mathbf{x}, \mathbf{y}).
\end{equation}
\end{assumption}

Under the above assumptions, both $f$ and $\Phi$ have Lipschitz continuous gradients, as shown in the two lemmas below.
\begin{lemma}[Smoothness of $f$]  \label{lemma_f_smooth}
    Under Assumption~\ref{assumption_smooth}, $\forall \zeta \in \{\zeta_{i, 1}, \ldots, \zeta_{i, n_i}\}$, $F_i(\mathbf{x}, \mathbf{y}; \zeta)$ is $L_f$-smooth with respect to $(\mathbf{x}, \mathbf{y})$,
    i.e., $\forall \mathbf{x}_1, \mathbf{x}_2 \in \mathbb{R}^{p}$ and $\forall \mathbf{y}_1, \mathbf{y}_2 \in \mathbb{R}^{q}$,
    $$
        \left\|
            \begin{bmatrix}
                \nabla_{\mathbf{x}} F_i(\mathbf{x}_1, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{x}} F_i(\mathbf{x}_2, \mathbf{y}_2; \zeta) \\
                \nabla_{\mathbf{y}} F_i(\mathbf{x}_1, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{y}} F_i(\mathbf{x}_2, \mathbf{y}_2; \zeta)
            \end{bmatrix}
        \right\|
        \le L_f \left\|
            \begin{bmatrix}
                \mathbf{x}_1 - \mathbf{x}_2 \\
                \mathbf{y}_1 - \mathbf{y}_2
            \end{bmatrix} \right\|,
    $$
    where the constant $L_f$ is defined as $L_f := \max\{\sqrt{2 (L_1^2 + L_{21}^2)}, \sqrt{2 (L_2^2 + L_{12}^2)}\}$.
    Furthermore, both $f_i$ and $f$ are also $L_f$-smooth.
\end{lemma}



The next lemma shows that $\Phi$ is differentiable and $L_{\Phi}$-smooth.
\begin{lemma}  \label{lemma_phi_smooth}
\citep[Lemma 22]{nouiehed2019solving}
Under Assumption~\ref{assumption_smooth} and~\ref{assumption_pl}, $\Phi(\mathbf{x})$ is differentiable and its gradient is given by
\begin{equation}  \label{eq_outer_gradient}
    \nabla \Phi(\mathbf{x}) = \nabla_{\mathbf{x}} f(\mathbf{x}, \mathbf{y}^*),
\end{equation}
where $\mathbf{y}^* \in \mathrm{argmax}_{\mathbf{y} \in \mathbb{R}^{q}} f(\mathbf{x}, \mathbf{y})$.
Moreover, $\Phi$ has $L_{\Phi}$-Lipschitz continuous gradients with $L_{\Phi} := L_1 + \frac{L_{12} L_{21}}{2 \mu}$.
\end{lemma}

Since $\Phi$ is differentiable, a natural metric for measuring the performance of an algorithm is the squared gradient norm of $\Phi$, i.e., $\|\nabla \Phi(\mathbf{x})\|^2$.
This performance metric is commonly used in analyzing algorithms for nonconvex-PL minimax problems~\citep{lin2019gradient, luo2020stochastic, wang2020zeroth}.

\section{Deferred proofs}  \label{appendix}
In our proofs, we denote the pair of primal and dual variables $(\mathbf{x}_{t, i}^{(k)}, \mathbf{y}_{t, i}^{(k)})$ as $\mathbf{z}_{t, i}^{(k)}$ and omit the subscript $t, i$ when clear from context.   


\subsection{Proof of Lemma~\ref{lemma_f_smooth}}  \label{proof_lemma_f_smooth}
\begin{proof}
Let $\mathbf{z} := (\mathbf{x}, \mathbf{y})$. We have
\begin{equation}
\begin{aligned}[b]
    &\| \nabla F_i(\mathbf{z}_1; \zeta) - \nabla F_i(\mathbf{z}_2; \zeta) \|^2 \\
    ={}& \| \nabla_{\mathbf{x}} F_i(\mathbf{x}_1, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{x}} F_i(\mathbf{x}_2, \mathbf{y}_2; \zeta) \|^2
        + \| \nabla_{\mathbf{y}} F_i(\mathbf{x}_1, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{y}} F_i(\mathbf{x}_2, \mathbf{y}_2; \zeta) \|^2 \\
    \le{}& 2 \| \nabla_{\mathbf{x}} F_i(\mathbf{x}_1, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{x}} F_i(\mathbf{x}_1, \mathbf{y}_2; \zeta) \|^2
        + 2 \| \nabla_{\mathbf{x}} F_i(\mathbf{x}_1, \mathbf{y}_2; \zeta) - \nabla_{\mathbf{x}} F_i(\mathbf{x}_2, \mathbf{y}_2; \zeta) \|^2 \\
        &+ 2 \| \nabla_{\mathbf{y}} F_i(\mathbf{x}_1, \mathbf{y}_1; \zeta) - \nabla_{\mathbf{y}} F_i(\mathbf{x}_1, \mathbf{y}_2; \zeta) \|^2
        + 2 \| \nabla_{\mathbf{y}} F_i(\mathbf{x}_1, \mathbf{y}_2; \zeta) - \nabla_{\mathbf{y}} F_i(\mathbf{x}_2, \mathbf{y}_2; \zeta) \|^2 \\
    \le{}& 2 (L_1^2 + L_{21}^2) \|\mathbf{x}_1 - \mathbf{x}_2\|^2 + 2 (L_2^2 + L_{12}^2) \|\mathbf{y}_1 - \mathbf{y}_2\|^2 \\
    \le{}& \max\{2 (L_1^2 + L_{21}^2), 2 (L_2^2 + L_{12}^2)\} \|\mathbf{z}_1 - \mathbf{z}_2\|^2
    = L_f^2 \|\mathbf{z}_1 - \mathbf{z}_2\|^2.
\end{aligned}
\end{equation}
Thus, $F_i(\mathbf{x}, \mathbf{y}, \zeta)$ is $L_f$-smooth w.r.t.\ $(\mathbf{x}, \mathbf{y})$.
Recall that $f_i = \frac{1}{n_i} \sum_{j=1}^{n_i} F_i(\mathbf{x}, \mathbf{y}; \zeta_{i,j})$ and $f(\mathbf{x}, \mathbf{y}) = \mathbb{E}_{i \sim \mathcal{D}}[f_i(\mathbf{x}, \mathbf{y})]$.
By Jensen's inequality and the convexity of $\|\cdot\|^2$, both $f_i$ and $f$ are $L_f$-smooth.
\end{proof}















\subsection{Proof of Lemma~\ref{lemma_descent}}  \label{proof_lemma_descent}
Before we prove Lemma~\ref{lemma_descent}, we need the following lemma which relates the approximation error of $\mathbf{d}_{\mathbf{x}}^{(k)}$ (resp., $\mathbf{d}_{\mathbf{y}}^{(k)}$) to $\varepsilon_{\mathbf{x}, t}$ (resp., $\varepsilon_{\mathbf{y}, t}$).
\begin{lemma}  \label{lemma_gradient_similarity}
Consider {FedSGDA}. Under Assumptions~\ref{assumption_smooth}-\ref{assumption_pl}, we have
$$
\begin{aligned}[b]
    \mathbb{E}[ \|\mathbf{d}_{\mathbf{x}}^{(k)} - \nabla \Phi(\mathbf{x}^{(k)})\|^2]
    \le{}& \frac{3 L_{12}^2}{2 \mu} \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)}) ]
        + 12 L_f^2 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 3 \varepsilon_{\mathbf{x}, t}, \\
    \mathbb{E}[ \|\mathbf{d}_{\mathbf{x}}^{(k)} - \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2]
    \le{}& 8 L_f^2 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 2 \varepsilon_{\mathbf{x}, t}, \\
    \mathbb{E}[ \|\mathbf{d}_{\mathbf{y}}^{(k)} - \nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2]
    \le{}& 8 L_f^2 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 2 \varepsilon_{\mathbf{y}, t}.
\end{aligned}
$$
\end{lemma}



\begin{proof}
By Lemma~\ref{lemma_f_smooth}, both $F_i(\cdot, \cdot; \mathcal{B}^{(k)})$ and $f(\cdot, \cdot)$ are $L_f$-smooth.
Thus, we can bound
\begin{equation}
\begin{aligned}[b]
    &\mathbb{E}[\| \nabla_{\mathbf{x}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{x}} F_i(\bar{\mathbf{z}}_{t}; \mathcal{B}^{(k)}) + \mathbf{u}_t - \nabla \Phi(\mathbf{x}^{(k)}) \|^2] \\
    ={}& \mathbb{E}[\| \nabla_{\mathbf{x}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{x}} F_i(\bar{\mathbf{z}}_{t}; \mathcal{B}^{(k)})
        - \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) + \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t})
        - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}) + \mathbf{u}_t
        + \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) - \nabla \Phi(\mathbf{x}^{(k)}) \|^2] \\
    \le{}& 3 \mathbb{E} [\|\nabla_{\mathbf{x}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{x}} F_i(\bar{\mathbf{z}}_{t}; \mathcal{B}^{(k)})
        - (\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}))\|^2 ]
        + 3 \varepsilon_{\mathbf{x}, t}
        + 3 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) - \nabla \Phi(\mathbf{x}^{(k)})\|^2] \\
    \le{}& 6 \mathbb{E} [\|\nabla_{\mathbf{x}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{x}} F_i(\bar{\mathbf{z}}_{t}; \mathcal{B}^{(k)}) \|^2
        + \| \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t})\|^2 ]
        + 3 \varepsilon_{\mathbf{x}, t}
        + 3 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) - \nabla \Phi(\mathbf{x}^{(k)})\|^2] \\
    \le{}& 12 L_f^2 \mathbb{E} [\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 3 \varepsilon_{\mathbf{x}, t}
        + 3 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) - \nabla \Phi(\mathbf{x}^{(k)})\|^2] \\
    \le{}& 12 L_f^2 \mathbb{E} [\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 3 \varepsilon_{\mathbf{x}, t}
        + 3 L_{12}^2 \mathbb{E}[ \|\mathbf{y}^{(k)} - \mathbf{y}^*(\mathbf{x}^{(k)})\|^2],
\end{aligned}
\end{equation}
where the third inequality follows from Lemma~\ref{lemma_f_smooth} and the last inequality follows from Assumption~\ref{assumption_smooth}.
By Theorem 2 in~\cite{karimi2016linear}, the PL condition implies that
\begin{equation}
    \Phi(\mathbf{x}^{(k)}) - f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \ge 2 \mu \| \mathbf{y}^{(k)} - \mathbf{y}^*(\mathbf{x}^{(k)})\|^2.
\end{equation}
Combining the above inequalities yields the first inequality of Lemma~\ref{lemma_gradient_similarity}.
Similarly, we can bound
\begin{equation}
\begin{aligned}[b]
    \mathbb{E}[ \|\mathbf{d}_{\mathbf{x}}^{(k)} - \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2]
    ={}& \mathbb{E}[ \| \nabla_{\mathbf{x}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{x}} F_i(\bar{\mathbf{z}}_{t})
        - \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) + \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t})
        + \mathbf{u}_t - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}) \|^2 ] \\
    \le{}& 4 L_f^2 \mathbb{E} [\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ] + 2 \varepsilon_{\mathbf{x}, t}, \\
    \mathbb{E}[ \|\mathbf{d}_{\mathbf{y}}^{(k)} - \nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2]
    ={}& \mathbb{E}[ \| \nabla_{\mathbf{y}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{y}} F_i(\bar{\mathbf{z}}_{t})
        - \nabla_{\mathbf{y}} f(\mathbf{z}^{(k)}) + \nabla_{\mathbf{y}} f(\bar{\mathbf{z}}_{t})
        + \mathbf{v}_t - \nabla_{\mathbf{y}} f(\bar{\mathbf{z}}_{t}) \|^2 ] \\
    \le{}& 4 L_f^2 \mathbb{E} [\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ] + 2 \varepsilon_{\mathbf{y}, t}.
\end{aligned}
\end{equation}
\end{proof}







We can now prove Lemma~\ref{lemma_descent}.
\begin{proof}[Proof of Lemma~\ref{lemma_descent}]
By Lemma~\ref{lemma_phi_smooth}, we have
\begin{equation}  \label{eq_smooth_expand}
\begin{aligned}[b]
    \mathbb{E}[\Phi(\mathbf{x}^{(k + 1)}) - \Phi(\mathbf{x}^{(k)})]
    \le{}& \mathbb{E}[\langle \nabla \Phi(\mathbf{x}^{(k)}), \mathbf{x}^{(k + 1)} - \mathbf{x}^{(k)} \rangle
        + \frac{L_{\Phi}}{2} \|\mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}\|^2] \\
    ={}& - \eta_t \mathbb{E}[ \langle \nabla \Phi(\mathbf{x}^{(k)}), \mathbf{d}_{\mathbf{x}}^{(k)} \rangle]
        + \frac{L_{\Phi} \eta_t^2}{2} \mathbb{E}[ \| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2] \\
    ={}& \frac{\eta_t}{2} \mathbb{E} [ \| \mathbf{d}_{\mathbf{x}}^{(k)} - \nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        - \frac{\eta_t}{2} \mathbb{E} [ \|\nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        - (\frac{\eta_t}{2} - \frac{L_{\Phi} \eta_t^2}{2}) \mathbb{E}[\|\mathbf{d}_{\mathbf{x}}^{(k)}\|^2] \\
    \le{}& \frac{\eta_t}{2} \mathbb{E} [ \| \mathbf{d}_{\mathbf{x}}^{(k)} - \nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        - \frac{\eta_t}{2} \mathbb{E} [ \|\nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
\end{aligned}
\end{equation}
where the first equality follows from the update rule $\mathbf{x}^{(k+1)} = \mathbf{x}^{(k)} - \eta_t \mathbf{d}_{\mathbf{x}}^{(k)}$,
the second equality follows from the fact that $-2 \langle \mathbf{a}, \mathbf{b} \rangle = \|\mathbf{a} - \mathbf{b}\|^2 - \|\mathbf{a}\|^2 - \|\mathbf{b}\|^2$,
and the last inequality follows from the condition $\eta_t \le 1 / L_{\Phi}$.
Combining~\eqref{eq_smooth_expand} and Lemma~\ref{lemma_gradient_similarity} completes the proof.
\end{proof}









\subsection{Proof of Lemma~\ref{lemma_iterates_recursion}}  \label{proof_lemma_iterates_recursion}
\begin{proof}
First, we bound $\mathbb{E}[ \|\mathbf{x}^{(k+1)} - \bar{\mathbf{x}}_{t}\|^2 ]$ as follows
\begin{equation}
\begin{aligned}[b]
    \mathbb{E}[ \|\mathbf{x}^{(k+1)} - \bar{\mathbf{x}}_{t}\|^2 ]
    =& \mathbb{E}[ \|\mathbf{x}^{(k)} - \eta_t \mathbf{d}_{\mathbf{x}}^{(k)} - \bar{\mathbf{x}}_{t}\|^2 ] \\
    \le& (1 + \frac{1}{2 K}) \mathbb{E}[ \|\mathbf{x}^{(k)} - \bar{\mathbf{x}}_{t}\|^2 ]
        + (1 + 2 K) \eta_t^2 \mathbb{E}[ \|\mathbf{d}_{\mathbf{x}}^{(k)}\|^2 ] \\
    \le& (1 + \frac{1}{2 K}) \mathbb{E}[ \|\mathbf{x}^{(k)} - \bar{\mathbf{x}}_{t}\|^2 ]
        + 6 K \eta_t^2 \mathbb{E}[ \|\mathbf{d}_{\mathbf{x}}^{(k)} - \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ]
        + 6 K \eta_t^2 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ] \\
    \le& (1 + \frac{1}{2 K}) \mathbb{E}[ \|\mathbf{x}^{(k)} - \bar{\mathbf{x}}_{t}\|^2 ]
        + 6 K \eta_t^2 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ]
        + 48 K L_f^2 \eta_t^2 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 12 K \eta_t^2 \varepsilon_{\mathbf{x}, t},
\end{aligned}
\end{equation}
where the first inequality follows from Young's inequality,
the second inequality follows from Cauchy-Schwarz and the fact that $K \ge 1$,
and the last inequality follows from Lemma~\ref{lemma_gradient_similarity}.
Similarly, $\mathbb{E}[ \|\mathbf{y}^{(k+1)} - \bar{\mathbf{y}}_{t}\|^2 ]$ can be bounded by
\begin{equation}
\begin{aligned}
    \mathbb{E}[ \|\mathbf{y}^{(k+1)} - \bar{\mathbf{y}}_{t}\|^2 ]
    \le& (1 + \frac{1}{2 K}) \mathbb{E}[ \|\mathbf{y}^{(k)} - \bar{\mathbf{y}}_{t}\|^2 ]
        + (1 + 2 K) \gamma_t^2 \mathbb{E}[ \|\mathbf{d}_{\mathbf{y}}^{(k)}\|^2 ] \\
    \le& (1 + \frac{1}{2 K}) \mathbb{E}[ \|\mathbf{y}^{(k)} - \bar{\mathbf{y}}_{t}\|^2 ]
        + 6 K \gamma_t^2 \mathbb{E}[ \|\nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ] \\
        &+ 48 K L_f^2 \gamma_t^2 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 12 K \gamma_t^2 \varepsilon_{\mathbf{y}, t}
\end{aligned}
\end{equation}
where the last inequality follows from Lemma~\ref{lemma_gradient_similarity}.
Combining the above inequalities yields that
\begin{equation}
\begin{aligned}[b]
    \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_{t}\|^2 ]
    \le& (1 + \frac{1}{2 K}
        + 48 K L_f^2 \eta_t^2 + 48 K L_f^2 \gamma_t^2) \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 6 K \eta_t^2 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ] \\
        &+ 6 K \gamma_t^2 \mathbb{E}[ \|\nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ]
        + 12 K \eta_t^2 \varepsilon_{\mathbf{x}, t}
        + 12 K \gamma_t^2 \varepsilon_{\mathbf{y}, t} \\
    \le& (1 + \frac{1}{K}) \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 ]
        + 6 K \eta_t^2 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ] \\
        &+ 6 K \gamma_t^2 \mathbb{E}[ \|\nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ]
        + 12 K \eta_t^2 \varepsilon_{\mathbf{x}, t}
        + 12 K \gamma_t^2 \varepsilon_{\mathbf{y}, t}
\end{aligned}
\end{equation}
where the last inequality follows from the condition $\eta_t \le \frac{1}{8 \sqrt{3} K L_f}$ and $\gamma_t \le \frac{1}{8 \sqrt{3} K L_f}$.
\end{proof}








\subsection{Proof of Lemma~\ref{lemma_max_residue}}  \label{proof_lemma_max_residue}
\begin{proof}
\begin{equation}  \label{eq_max_residule_expand}
\begin{aligned}[b]
    &\mathbb{E}[ \Phi(\mathbf{x}^{k + 1}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k+1)}) ] \\
    \le& \mathbb{E} \Big[ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)})
        - \gamma_t \langle \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}), \mathbf{d}_{\mathbf{y}}^{(k)} \rangle
        + \frac{L_2 \gamma_t^2}{2} \| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2 \Big] \\
    \le& \mathbb{E} \Big[ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)})
        - \frac{\gamma_t}{2} \| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) \|^2 \\
        &+ \frac{\gamma_t}{2} \| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) - \mathbf{d}_{\mathbf{y}}^{(k)} \|^2
        - (\frac{\gamma_t}{2} - \frac{L_2 \gamma_t^2}{2}) \| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2 \Big] \\
    \le& (1 - \frac{\mu \gamma_t}{2}) \mathbb{E} [ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) ]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) \|^2] \\
        &+ \frac{\gamma_t}{2} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) - \mathbf{d}_{\mathbf{y}}^{(k)} \|^2]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2 ],
\end{aligned}
\end{equation}
where the first inequality follows from Assumption~\ref{assumption_smooth} and the update rule of $\mathbf{y}^{(k)}$,
the second inequality follows from the fact that $- 2 \langle \mathbf{a}, \mathbf{b} \rangle = \|\mathbf{a} - \mathbf{b}\|^2 - \|\mathbf{a}\|^2 - \|\mathbf{b}\|^2$,
and the last inequality follows from Assumption~\ref{assumption_pl} and the condition $\gamma_t \le \frac{1}{2 L_2}$.

We proceed to bound the first term on the RHS of~\eqref{eq_max_residule_expand}.
Notice that
\begin{equation}  \label{eq_term_1_max_residule_expand}
\begin{aligned}[b]
    \mathbb{E} [ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) ]
    =& \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})
        + \Phi(\mathbf{x}^{(k+1)}) - \Phi(\mathbf{x}^{(k)}) \\
        &+ f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) ].
\end{aligned}
\end{equation}
We can bound $\mathbb{E}[ f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) ]$ by
\begin{equation}  \label{eq_f_diff_bound}
\begin{aligned}[b]
    &\mathbb{E}[ f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) ] \\
    \overset{(a)}{\le}& \mathbb{E} \big[- \langle \nabla_{\mathbf{x}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}), \mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}) \rangle
        + \frac{L_1}{2} \|\mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}\|^2 \big] \\
    \overset{(b)}{=}& \mathbb{E} \big[ \frac{\eta_t}{2} \| \nabla_{\mathbf{x}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) - \mathbf{d}_{\mathbf{x}}^{(k)}\|^2
        - \frac{\eta_t}{2} \| \nabla_{\mathbf{x}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})\|^2
        - (\frac{\eta_t}{2} - \frac{L_1 \eta_t^2}{2}) \| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2 \big], \\
    \overset{(c)}{\le}&
    4 L_f^2 \eta_t \mathbb{E}[\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2]
        + \eta_t \varepsilon_{\mathbf{x}, t}
        - \frac{\eta_t}{2} \mathbb{E}[\| \nabla_{\mathbf{x}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})\|^2]
        - \frac{\eta_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2].
\end{aligned}
\end{equation}
where (a) follows from Assumption~\ref{assumption_smooth},
(b) follows from the fact that $- 2 \langle \mathbf{a}, \mathbf{b} \rangle = \|\mathbf{a} - \mathbf{b}\|^2 - \|\mathbf{a}\|^2 - \|\mathbf{b}\|^2$,
and (c) follows from Lemma~\ref{lemma_gradient_similarity} and the condition $\eta_t \le \frac{1}{2 L_1}$.
Combining~\eqref{eq_term_1_max_residule_expand}, \eqref{eq_f_diff_bound}, and Lemma~\ref{lemma_descent}, we have
\begin{equation}  \label{eq_term_1_max_residule_bound}
\begin{aligned}[b]
    \mathbb{E}[ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)})]
    \le& \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]
        + 4 L_f^2 \eta_t \mathbb{E}[\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2]
        + \eta_t \varepsilon_{\mathbf{x}, t} \\
        &- \frac{\eta_t}{2} \mathbb{E}[ \| \nabla \Phi(\mathbf{x}^{(k)})\|^2]
        + \frac{3 L_{12}^2}{4 \mu}  \eta_t \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]
        + 6 L_f^2 \eta_t \mathbb{E}[ \| \mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2] \\
        &+ \frac{3}{2} \eta_t \varepsilon_{\mathbf{x}, t}
        - \frac{\eta_t}{2} \mathbb{E}[\| \nabla_{\mathbf{x}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})\|^2]
        - \frac{\eta_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2] \\
    =& (1 + \frac{3 L_{12}^2}{4 \mu} \eta_t) \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]
        - \frac{\eta_t}{2} \mathbb{E}[ \| \nabla \Phi(\mathbf{x}^{(k)})\|^2]
        + 10 L_f^2 \eta_t \mathbb{E}[ \| \mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2] \\
        &+ \frac{5}{2} \eta_t \varepsilon_{\mathbf{x}, t}
        - \frac{\eta_t}{2} \mathbb{E}[\| \nabla_{\mathbf{x}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})\|^2]
        - \frac{\eta_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2].
\end{aligned}
\end{equation}

It remains to bound the second term and the third term on the RHS of~\eqref{eq_max_residule_expand}.
Notice that
\begin{equation}  \label{eq_term_2_max_residule_expand}
\begin{aligned}[b]
    - \frac{\gamma_t}{4} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) \|^2]
    \le& - \frac{\gamma_t}{8} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2]
        + \frac{\gamma_t}{4} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) - \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2] \\
    \le& - \frac{\gamma_t}{8} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2]
        + \frac{L_{21}^2 \gamma_t}{4} \mathbb{E}[\|\mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}\|^2],
\end{aligned}
\end{equation}
where the first inequality follows from the fact that $- \|\mathbf{a}\|^2 \le - \frac{1}{2} \|\mathbf{b}\|^2 + \|\mathbf{b} - \mathbf{a}\|^2$
and the last inequality follows from Assumption~\ref{assumption_smooth}.
Thus, we have
\begin{equation}  \label{eq_term_2n3_max_residule_expand}
\begin{aligned}[b]
    &- \frac{\gamma_t}{4} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) \|^2]
    + \frac{\gamma_t}{2} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) - \mathbf{d}_{\mathbf{y}}^{(k)} \|^2] \\
    \le{}& - \frac{\gamma_t}{8} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2]
        + \frac{L_{21}^2 \gamma_t}{4} \mathbb{E}[\|\mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}\|^2]
        + \frac{\gamma_t}{2} \mathbb{E} \Big[ \| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k)}) - \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \\
            &+ \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) - \nabla_{\mathbf{y}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)
            - \big( \nabla_{\mathbf{y}} F_i(\mathbf{z}^{(k)}; \mathcal{B}^{(k)}) - \nabla_{\mathbf{y}} F_i(\bar{\mathbf{z}}_{t}; \mathcal{B}^{(k)}) \big)
            - \mathbf{v}_t + \nabla_{\mathbf{y}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) \|^2 \Big] \\
    \le{}& - \frac{\gamma_t}{8} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2]
        + \frac{L_{21}^2 \gamma_t}{4} \mathbb{E}[\|\mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}\|^2]
        + \frac{3 \gamma_t}{2} \mathbb{E} \Big[ L_{21}^2 \| \mathbf{x}^{(k+1)} - \mathbf{x}^{(k)}\|^2
            + \varepsilon_{\mathbf{y}, t}
            + 4 L_f^2 \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2 \Big] \\
    ={}& - \frac{\gamma_t}{8} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2]
        + \frac{7 L_{21}^2 \eta_t^2 \gamma_t}{4} \mathbb{E}[\|\mathbf{d}_{\mathbf{x}}^{(k)} - \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)}) + \nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2]
        + \frac{3 \gamma_t}{2} \varepsilon_{\mathbf{y}, t}
        + 6 L_f^2 \gamma_t \mathbb{E}[\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2] \\
    \le{}& - \frac{\gamma_t}{8} \mathbb{E}[\| \nabla_{\mathbf{y}} f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)}) \|^2]
        + \frac{3 \gamma_t}{2} \varepsilon_{\mathbf{y}, t}
        + 6 L_f^2 \gamma_t \mathbb{E}[\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2]
        + \frac{21 L_{21}^2 \eta_t^2 \gamma_t}{4} \mathbb{E}[4 L_f^2 \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_{t}\|^2 + \varepsilon_{\mathbf{x}, t} + \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2],
\end{aligned}
\end{equation}
where the second inequality follows from Cauchy-Schwarz and Assumption~\ref{assumption_smooth};
the last inequality follows from the update rule of $\mathbf{x}^{(k)}$, Assumption~\ref{assumption_smooth}, and the same argument as Lemma~\ref{lemma_gradient_similarity}.
Plugging~\eqref{eq_term_1_max_residule_bound} and~\eqref{eq_term_2n3_max_residule_expand} into~\eqref{eq_max_residule_expand}
and applying the conditions $\eta_t \le \frac{1}{42 L_{21}}, \gamma_t \le \min\{\frac{1}{\mu}, \frac{1}{L_{21}}\}$ leads to the desired result.
\end{proof}










\subsection{Proof of Lemma~\ref{lemma_potential_recursion_decay_simple}}  \label{proof_lemma_potential_recursion_decay_simple}
\begin{proof}
We note that $1 \le (1 + \frac{2}{K})^{K - k} \le 8$ for any $k \in \{0, \ldots, K\}$.
Define $\tilde{L} = 2 (L_{\Phi} + L_f)$.
Adding $\mathbb{E}[ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k+1)})]
    + \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k - 1} \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_t\|^2 ]$
to both sides of the inequality in Lemma~\ref{lemma_descent} and applying Lemma~\ref{lemma_iterates_recursion} and~\ref{lemma_max_residue}, we get
\begin{equation}  \label{eq_sum_descent_lemma_potential_simple}
\begin{aligned}[b]
    &\mathbb{E}[ 2 \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k+1)}) ]
        + \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k - 1} \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_t\|^2 ] \\
    \le& \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) ]
        - a_1 \mathbb{E}[ \|\nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        + a_2 \mathbb{E}[ \Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]
        + a_3 \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2 ] \\
        &- a_4 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2 ]
        - a_5 \mathbb{E}[ \|\nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2 ]
        + a_6 \varepsilon_{\mathbf{x}, t}
        + a_7 \varepsilon_{\mathbf{y}, t}
        - \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2],
\end{aligned}
\end{equation}
where $a_1, \ldots, a_6$ are defined as
\begin{equation}
\begin{aligned}[b]
    a_1 =& \frac{\eta_t}{2} ( 2 - \frac{\mu \gamma_t}{2}), \\
    a_2 =& (1 - \frac{\mu \gamma_t}{2}) (1 + \frac{3 L_{12}^2}{4 \mu} \eta_t)
        + \frac{3 L_{12}^2}{4 \mu} \eta_t, \\
    a_3 =& 6 L_f^2 \eta_t + 10 (1 - \frac{\mu \gamma_t}{2}) L_f^2 \eta_t
        + 6 L_f^2 \gamma_t + 21 L_f^2 L_{21}^2 \eta_t^2 \gamma_t + \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k - 1} (1 + \frac{1}{K}), \\
    a_4 =& \frac{\eta_t}{8} - 3 \tilde{L} (1 + \frac{2}{K})^{K - k - 1} K \eta_t^2, \\
    a_5 =& \frac{\gamma_t}{8} - 3 \tilde{L} (1 + \frac{2}{K})^{K - k - 1} K \gamma_t^2, \\
    a_6 =& \frac{3}{2} \eta_t + \frac{5}{2} (1  - \frac{\mu \gamma_t}{2}) \eta_t + \frac{21}{4} L_{21}^2 \eta_t^2 \gamma_t + 6 \tilde{L} (1 + \frac{2}{K})^{K - k - 1} K \eta_t^2, \\
    a_7 =& \frac{3}{2} \gamma_t + 6 \tilde{L} (1 + \frac{2}{K})^{K - k - 1} K \gamma_t^2.
\end{aligned}
\end{equation}
Recall the conditions
$
\eta_t
    \le \min\{
    \frac{1}{4 L_1},
    \frac{1}{L_{\Phi}},
    \frac{1}{42 L_{21}},
    \gamma_t,
    \frac{\mu^2}{3 L_{12}^2} \gamma_t
\},
\
\gamma_t
    \le \min\{
    \frac{1}{2 L_2},
    \frac{1}{\mu},
    \frac{1}{L_{21}},
    \frac{1}{192 \tilde{L} K}
\}
$
and note that $1 \le (1 + \frac{2}{K})^{K - k} \le 8$.
We have the following bounds
$$
\begin{aligned}[b]
    a_1 \ge{}& \eta_t - \frac{\mu \gamma_t \eta_t}{4}
    \ge \eta_t - \frac{\eta_t}{4}
    = \frac{3 \eta_t}{4}, \\
    a_2 \le{}& 1 + \frac{3 L_{12}^2}{2 \mu} \eta_t - \frac{\mu \gamma_t}{2}
    \le 1 + \frac{3 L_{12}^2}{2 \mu} \frac{\mu^2 \gamma_t}{3 L_{12}^2} - \frac{\mu \gamma_t}{2}, \\
    a_3 \le{}& \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k}
        - \frac{\tilde{L}}{2 K}
        + L_f^2(6 \eta_t + 10 \eta_t + 6 \gamma_t + 21 L_{21}^2 \cdot \frac{\eta_t}{42 L_{21}^2}) \\
    \le{}& \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k}
        - \frac{\tilde{L}}{2 K}
        + \frac{\tilde{L}}{4} \cdot \frac{6 + 10 + 6 + 0.5}{192 K}
    \le \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k}, \\
    a_4 \ge{}& \frac{\eta_t}{8} - 24 \tilde{L} K \eta_t^2
    \ge \frac{\eta_t}{8} - 24 \tilde{L} K \cdot \frac{\eta_t}{192 \tilde{L} K}
    = 0, \\
    a_5 \ge{}& \frac{\gamma_t}{8} - 24 \tilde{L} K \gamma_t^2
    \ge \frac{\gamma_t}{8} - 24 \tilde{L} K \cdot \frac{\gamma_t}{192 \tilde{L} K}
    = 0, \\
    a_6 \le{}& \frac{3 \eta_t}{2} + \frac{5 \eta_t}{2}
        + \frac{21}{4} L_{21}^2 \frac{\eta_t}{42 L_{21}^2}
        + 48 \tilde{L} K \frac{\eta_t}{192 \tilde{L} K} 
    \le \eta_t \big( \frac{3}{2} + \frac{5}{2} + \frac{1}{8} + \frac{1}{4} \big)
    = \frac{35 \eta_t}{8}, \\
    a_7 \le{}& \frac{3 \gamma_t}{2} + 48 \tilde{L} K \frac{\gamma_t}{192 \tilde{L} K}
    = \gamma_t \big( \frac{3}{2} + \frac{1}{4}\big) = \frac{7 \gamma_t}{4}.
\end{aligned}
$$
Since $\mathbf{x}^{(k + 1)} = \mathbf{x}^{(k)} - \eta_t \mathbf{d}_{\mathbf{x}}^{(k)}$
and $\mathbf{y}^{(k + 1)} = \mathbf{y}^{(k)} - \gamma_t \mathbf{d}_{\mathbf{y}}^{(k)}$,
we have
$$
\begin{aligned}[b]
    \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        + \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2]
    = \mathbb{E} \left[ \frac{1}{8 \eta_t} \left\|\mathbf{x}^{(k + 1)} - \mathbf{x}^{(k)} \right\|^2
        + \frac{1}{4 \gamma_t} \left\|\mathbf{y}^{(k + 1)} - \mathbf{y}^{(k)} \right\|^2 \right]
    \ge \frac{1}{8 \gamma_t} \mathbb{E} \left[ \left\|\mathbf{z}^{(k + 1)} - \mathbf{z}^{(k)} \right\|^2 \right],
\end{aligned}
$$
where the inequality follows from the condition $\eta_t \le \gamma_t$.

Applying the above bounds into~\eqref{eq_sum_descent_lemma_potential_simple} leads to
\begin{equation}  \label{eq_potential_recursion_simple}
\begin{aligned}[b]
    & \frac{3 \eta_t}{4} \mathbb{E}[ \|\nabla \Phi(\mathbf{x}^{(k)})\|^2 ]
        + \mathbb{E}[ 2 \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{x}^{(k+1)}, \mathbf{y}^{(k+1)})]
        + \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k - 1} \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_t\|^2 ] \\
    \le& \mathbb{E}[ 2 \Phi(\mathbf{x}^{(k)}) - f(\mathbf{x}^{(k)}, \mathbf{y}^{(k)})] + \frac{\tilde{L}}{2} (1 + \frac{2}{K})^{K - k} \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2 ]
        + \frac{35 \eta_t}{8} \varepsilon_{\mathbf{x}, t}
        + \frac{7 \gamma_t}{4} \varepsilon_{\mathbf{y}, t}
        - \frac{1}{8 \gamma_t} \mathbb{E}[ \|\mathbf{z}^{(k + 1)} - \mathbf{z}^{(k)}\|^2 ],
\end{aligned}
\end{equation}
which is the desired result.
\end{proof}



\subsection{Proof of Lemma~\ref{lemma_potential_recursion_sum}}
\begin{proof}
By Cauchy-Schwarz inequality, we have
\begin{equation}
\begin{aligned}[b]
    \frac{1}{S K} \sum_{i \in \mathcal{S}_t} \sum_{k = 0}^{K - 1} \left\|\mathbf{z}_{t, i}^{(k + 1)} - \mathbf{z}_{t, i}^{(k)} \right\|^2
    \ge& \frac{1}{S K^2} \sum_{i \in \mathcal{S}_t} \left\| \sum_{k=0}^{K - 1} \mathbf{z}_{t, i}^{(k + 1)} - \mathbf{z}_{t, i}^{(k)} \right\|^2
    = \frac{1}{S K^2} \sum_{i \in \mathcal{S}_t} \left\| \mathbf{z}_{t, i}^{(K)} - \mathbf{z}_{t, i}^{(0)} \right\|^2 \\
    \ge& \frac{1}{K^2} \left\| \frac{1}{S} \sum_{i \in \mathcal{S}_t} \mathbf{z}_{t, i}^{(K)} - \bar{\mathbf{z}}_t \right\|^2
    = \frac{1}{K^2} \left\| \bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t \right\|^2,
\end{aligned}
\end{equation}
where the second inequality follows from Jensen's inequality and the convexity of $\|\cdot\|^2$.

Recall that $\bar{\mathbf{x}}_{t+1} = \frac{1}{S} \sum_{i \in \mathcal{S}_t} \mathbf{x}_{t, i}^{(K)}$ and $\bar{\mathbf{y}}_{t+1} = \frac{1}{S} \sum_{i \in \mathcal{S}_t} \mathbf{y}_{t, i}^{(K)}$.
Since $\Phi$ is $L_{\Phi}$-smooth and $f$ is $L_f$-smooth, the function $2 \Phi(\mathbf{x}) - f(\mathbf{x}, \mathbf{y})$ is $(2 L_{\Phi} + L_f)$-smooth w.r.t.\ $(\mathbf{x}, \mathbf{y})$.
Thus, the function $(\mathbf{x}, \mathbf{y}) \rightarrow 2 \Phi(\mathbf{x}) - f(\mathbf{x}, \mathbf{y})$ is $(2 L_{\Phi} + L_f)$-weakly convex~\citep{davis2019stochastic},
which implies that for any $\lambda \ge 2 L_{\Phi} + L_f$ and given some $(\mathbf{x}', \mathbf{y}') \in \mathbb{R}^{p} \times \mathbb{R}^{q}$, the function $(\mathbf{x}, \mathbf{y}) \rightarrow 2 \Phi(\mathbf{x}) - f(\mathbf{x}, \mathbf{y}) + \frac{\rho}{2} (\| \mathbf{x} - \mathbf{x}' \|^2 + \|\mathbf{y} - \mathbf{y}'\|^2)$ is convex.
By taking $\rho = \tilde{L}$ and $(\mathbf{x}', \mathbf{y}') = (\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$, we obtain 
\begin{equation}  \label{eq_weakly_convex}
\begin{aligned}[b]
&\frac{1}{S} \sum_{i \in \mathcal{S}_t} \Big( 2 \Phi(\mathbf{x}_{t, i}^{(K)}) - f(\mathbf{x}_{t, i}^{(K)}, \mathbf{y}_{t, i}^{(K)})
    + \frac{\tilde{L}}{2} ( \|\mathbf{x}_{t, i}^{(K)} - \bar{\mathbf{x}}_t\|^2 + \|\mathbf{y}_{t, i}^{(K)} - \bar{\mathbf{y}}_t\|^2 ) \Big) \\
\ge& 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{x}}_{t+1}, \bar{\mathbf{y}}_{t+1})
    + \frac{\tilde{L}}{2} ( \|\bar{\mathbf{x}}_{t+1} - \bar{\mathbf{x}}_t\|^2 + \|\bar{\mathbf{y}}_{t+1} - \bar{\mathbf{y}}_t\|^2 )
\ge 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{x}}_{t+1}, \bar{\mathbf{y}}_{t+1}).
\end{aligned}
\end{equation}
Combining the above two inequalities and~\eqref{eq_potential_recursion_simple}, we get
\begin{equation}
\begin{aligned}[b]
    \frac{3 \eta_t}{4 S K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_t) - f(\bar{\mathbf{z}}_t)]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})] \\
        &+ \frac{35 \eta_t}{8} \varepsilon_{\mathbf{x}, t}
        + \frac{7 \gamma_t}{4} \varepsilon_{\mathbf{y}, t}
        - \frac{1}{8 K^2 \gamma_t} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ].
\end{aligned}
\end{equation}
\end{proof}












\subsection{Proof of Theorem~\ref{theorem_fed_sgda_mb}}  \label{proof_theorem_fed_sgda_mb}
\begin{proof}
By Assumption~\ref{assumption_bgd}, $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$ can be bounded by
\begin{equation}  \label{eq_gradient_error_simple}
\begin{aligned}[b]
    \varepsilon_{\mathbf{x}, t}
    ={}& \mathbb{E}[ \|\frac{1}{S} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)\|^2 ]
    = \frac{1}{S^2} \sum_{i \in \mathcal{S}_t'} \| \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) \|^2
    \le \frac{\sigma_1^2}{S}, \\
    \varepsilon_{\mathbf{y}, t}
    ={}& \mathbb{E}[ \|\frac{1}{S} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{y}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)\|^2 ]
    = \frac{1}{S^2} \sum_{i \in \mathcal{S}_t'} \| \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{y}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) \|^2
    \le \frac{\sigma_2^2}{S}.
\end{aligned}
\end{equation}
Summing~\eqref{eq_potential_recursion_sum} from $t = 0$ to $T-1$ and applying~\eqref{eq_gradient_error_simple}, we get
\begin{equation}  \label{eq_sum_all_rounds_simple}
\begin{aligned}[b]
    \frac{1}{S K T} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le& \frac{4 C_0}{3 K \eta T}
        + \frac{35 \sigma_1^2}{6 S}
        + \frac{7 \sigma_2^2 \gamma}{S \eta}.
\end{aligned}
\end{equation}
Applying the choice of $\eta_t$ and $\gamma_t$ into~\eqref{eq_sum_all_rounds_simple}, we conclude that
\begin{equation}
\begin{aligned}[b]
    &\frac{1}{S K T} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \mathcal{O} \Big(
        \frac{C_0}{K T} \big( L_{\Phi} + L_{21} + \frac{L_{12}^2}{\mu^2} (L_2 + \mu + L_{21} + \tilde{L} K) \big)
        + \frac{\sigma_1^2}{S} + \frac{\sigma_2^2}{S} \big( \frac{L_{\Phi}}{\tilde{L} K} + \frac{L_{21}}{\tilde{L} K} + \frac{L_{12}^2}{\mu^2} \big)
    \Big),
\end{aligned}
\end{equation}
which is the desired result.
\end{proof}























\subsection{Proof of Lemma~\ref{lemma_gradient_error_storm}}  \label{proof_lemma_gradient_error_storm}
\begin{proof}
By the definition of $\varepsilon_{\mathbf{x}, t+1}$, we have
\begin{equation}
\begin{aligned}[b]
    \varepsilon_{\mathbf{x}, t+1} =& \mathbb{E}[\|\mathbf{u}_{t+1} - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t+1})\|^2] \\
    =& \mathbb{E}[\|(1 - \alpha_{t+1}) \big(\mathbf{u}_t - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}) \big)
        + \alpha_{t+1} \big( \frac{1}{S} \sum_{i \in \mathcal{S}_{t+1}'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t+1}) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t+1}) \big) \\
        &+ (1 - \alpha_{t+1}) \big( \frac{1}{S} \sum_{i \in \mathcal{S}_{t+1}'} (\nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t+1}) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t})) - \nabla_{\bar{\mathbf{x}}} f(\bar{\mathbf{z}}_{t+1}) + \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}) \big)
        \|^2] \\
    \overset{(a)}{=}& (1 - \alpha_{t+1})^2 \varepsilon_{\mathbf{x}, t}
        + \mathbb{E} \big[ \| \alpha_{t+1} \big( \frac{1}{S} \sum_{i \in \mathcal{S}_{t+1}'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t+1}) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t+1}) \big) \\
        &+ (1 - \alpha_{t+1}) \big( \frac{1}{S} \sum_{i \in \mathcal{S}_{t+1}'} (\nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t+1}) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t})) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t+1}) + \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}) \big)
        \|^2 \big] \\
     \le& (1 - \alpha_{t+1})^2 \varepsilon_{\mathbf{x}, t}
        + 2 \alpha_{t+1}^2 \mathbb{E} \big[ \| \big( \frac{1}{S} \sum_{i \in \mathcal{S}_{t+1}'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t+1}) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t+1}) \big) \|^2 \big]\\
        &+ 2 (1 - \alpha_{t+1})^2 \mathbb{E} \big[ \|
            \big( \frac{1}{S} \sum_{i \in \mathcal{S}_{t+1}'} (\nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t+1}) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{z}}_{t})) - \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t+1}) + \nabla_{\mathbf{x}} f(\bar{\mathbf{z}}_{t}) \big)
        \|^2 \big] \\
    \overset{(b)}{\le}& (1 - \alpha_{t+1})^2 \varepsilon_{\mathbf{x}, t}
        + \frac{8 (1 - \alpha_{t+1})^2 L_f^2}{S} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_{t}\|^2 ]
        + \frac{2 \alpha_{t+1}^2 \sigma_1^2}{S},
\end{aligned}
\end{equation}
where (a) holds because $\mathbb{E}_{i \in \mathcal{S}_{t+1}'} f_i(\mathbf{z}) = f(\mathbf{z})$ and $\varepsilon_{\mathbf{x}, t}$ is independent of $\mathcal{S}_{t+1}'$;
(b) follows from Assumption~\ref{assumption_bgd} and Lemma~\ref{lemma_f_smooth}.
This completes the first part of Lemma~\ref{lemma_gradient_error_storm}.
The second part of the lemma follows from the same argument.
\end{proof}






\subsection{Proof of Lemma~\ref{lemma_potential_recursion_decay_storm}}  \label{proof_lemma_potential_recursion_decay_storm}


\begin{proof}
First, we recall that our goal is to prove the following inequality
\begin{equation}  \label{eq_potential_STORM_descent}
\begin{aligned}[b]
    &\mathbb{E}[ 2 \Phi(\mathbf{x}^{(k+1)}) - \Phi(\mathbf{x}^*) - f(\mathbf{z}^{(k+1)})]
        + \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{(K - k - 1)} \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_t\|^2]
        + \frac{48 \eta_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{x}, t+1}
        + \frac{24 \gamma_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{y}, t+1} \\
    \le{}&
        \mathbb{E}[ 2 \Phi(\mathbf{x}^{(k)}) - \Phi(\mathbf{x}^*) - f(\mathbf{z}^{(k)})]
        + \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{(K - k)} \mathbb{E}[ \|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2]
        + \frac{48 \eta_t}{\alpha_t} \varepsilon_{\mathbf{x}, t}
        + \frac{24 \gamma_t}{\alpha_t} \varepsilon_{\mathbf{y}, t} \\
        &- \frac{3 \eta_t}{4} \mathbb{E}[\|\nabla \Phi(\mathbf{x}^{(k)})\|^2]
        + 96 \eta_{t+1} \alpha_{t+1} \frac{\sigma_1^2}{S}
        + 48 \gamma_{t+1} \alpha_{t+1} \frac{\sigma_2^2}{S}
        + \frac{1}{8 S K^2 \gamma_t} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ]
        - \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2].
\end{aligned}
\end{equation}


We note that $1 \le (1 + \frac{2}{K})^{K - k} \le 8$ for any $k \in \{0, \ldots, K\}$.
Adding $\mathbb{E}[ \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{z}^{(k+1)})]
        + \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{(K - k - 1)} \mathbb{E}[ \|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_t\|^2]
        + \frac{48 \eta_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{x}, t+1}
        + \frac{24 \gamma_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{y}, t+1}$
to both sides of the inequality in Lemma~\ref{lemma_descent} and applying Lemma~\ref{lemma_iterates_recursion}, \ref{lemma_max_residue}, and~\ref{lemma_gradient_error_storm}, we have
\begin{equation}  \label{eq_sum_descent_lemma_potential_storm}
\begin{aligned}[b]
    &\mathbb{E}[ 2 \Phi(\mathbf{x}^{(k+1)}) - f(\mathbf{z}^{(k+1)})]
        + \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{(K - k - 1)} \mathbb{E}[\|\mathbf{z}^{(k+1)} - \bar{\mathbf{z}}_t\|^2]
        + \frac{48 \eta_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{x}, t+1}
        + \frac{24 \gamma_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{y}, t+1} \\
    \le& \mathbb{E}[ \Phi(\mathbf{x}^{(k)})]
        - a_1 \mathbb{E}[\|\nabla \Phi(\mathbf{x}^{(k)})\|^2]
        + a_2 \mathbb{E}[\Phi(\mathbf{x}^{(k)}) - f(\mathbf{z}^{(k)})]
        + a_3 \mathbb{E}[\|\mathbf{z}^{(k)} - \bar{\mathbf{z}}_t\|^2]
        + a_4 \varepsilon_{\mathbf{x}, t}
        + a_5 \varepsilon_{\mathbf{y}, t} \\
        &- a_6 \mathbb{E}[ \|\nabla_{\mathbf{x}} f(\mathbf{z}^{(k)})\|^2]
        - a_7 \mathbb{E}[ \|\nabla_{\mathbf{y}} f(\mathbf{z}^{(k)})\|^2]
        + 96 \eta_{t+1} \alpha_{t+1} \frac{\sigma_1^2}{S}
        + 48 \gamma_{t+1} \alpha_{t+1} \frac{\sigma_2^2}{S} \\
        &+ a_8 \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ]
        - \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        - \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2],
\end{aligned}
\end{equation}
where $a_1, \ldots, a_8$ are defined as
\begin{equation}
\begin{aligned}[b]
    a_1 =& \frac{\eta_t}{2} (2 - \frac{\mu \gamma_t}{2}),  \\
    a_2 =& \frac{3 L_{12}^2 \eta_t}{4 \mu}
        + (1 - \frac{\mu \gamma_t}{2}) (1 + \frac{3 L_{12}^2 \eta_t}{4 \mu}), \\
    a_3 =& L_f^2 \Big( 6 \eta_t
        + 10 (1 - \frac{\mu \gamma_t}{2}) \eta_t
        + 6 \gamma_t
        + 21 L_{21}^2 \eta_t^2 \gamma_t \Big)
        + \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{K - k - 1} (1 + \frac{1}{K}), \\
    a_4 =& \frac{3 \eta_t}{2}
        + \frac{48 (1 - \alpha_{t+1})^2 \eta_{t+1}}{\alpha_{t+1}}
        + \frac{5}{2} (1 - \frac{\mu \gamma_t}{2}) \eta_t
        + \frac{21}{4} L_{21}^2 \eta_t^2 \gamma_t
        + \frac{72 \tilde{L}^2 K^2 \eta_t^2 \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{K - k - 1} \\
    a_5 =& \frac{3 \gamma_t}{2} + \frac{24 (1 - \alpha_{t+1})^2 \gamma_{t+1}}{\alpha_{t+1}}
        + \frac{72 \tilde{L}^2 K^2 \gamma_t^3}{\alpha_t} (1 + \frac{2}{K})^{K - k - 1} \\
    a_6 =& \frac{\eta_t}{8} - \frac{36 \tilde{L}^2 K^2 \eta_t^2 \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{K - k - 1}, \\
    a_7 =& \frac{\gamma_t}{8} - \frac{36 \tilde{L}^2 K^2 \gamma_t^3}{\alpha_t} (1 + \frac{2}{K})^{K - k - 1} \\
    a_8 =& 192 L_f^2 \Big( \frac{2 \eta_{t+1}}{S \alpha_{t+1}} (1 - \alpha_{t+1})^2 + \frac{\gamma_{t+1}}{S \alpha_{t+1}} (1 - \alpha_{t+1})^2 \Big).
\end{aligned}
\end{equation}
Recall that $1 \le (1 + \frac{2}{K})^{K - k} \le 8$, $\alpha_t \le \frac{2304 \tilde{L}^2 K^2}{192^2 \tilde{L}^2 K^2} \le 1$, and $L_f \le \tilde{L} / 2$.
We can bound
\begin{equation}  \label{eq_a1_a2_a3_bound_storm}
\begin{aligned}[b]
    a_1 \ge& \frac{\eta_t}{2} (2 - 1 /2) = \frac{3 \eta_t}{4}, \\
    a_2 \le& 1 - \frac{\mu  \gamma_t}{2} + \frac{3 L_{12}^2}{2 \mu} \eta_t \le 1, \\
    a_3 \le& \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{K - k}
        - \frac{6 \tilde{L}^2 \gamma_t}{\alpha_t}
        + (16 \eta_t + 6 \gamma_t + \frac{21 \gamma_t}{1764}) L_f^2 \\
        \le& \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{K - k}
        - 6 \tilde{L}^2 \gamma_t
        + 5.51 \tilde{L}^2 \gamma_t
        \le \frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} (1 + \frac{2}{K})^{K - k}.
\end{aligned}
\end{equation}
Since $(t + 2)^{1/3} \le (t + 1)^{1/3} + 1 \le 2 (t + 1)^{1/3}$ for any $t \ge 0$, we have
\begin{equation}
    \gamma_{t+1} =
    \begin{cases}
        \gamma_{t},
            \ \textrm{if} \ \frac{S^{1/3}}{\tilde{L} K (t + 2)^{1/3}} > \min\{    \frac{1}{L_2}, \frac{1}{\mu}, \frac{1}{L_{21}}, \frac{1}{192 \tilde{L} K}\}, \\
        \frac{S^{1/3}}{\tilde{L} K (t + 2)^{1/3}}
            \ge \frac{S^{1/3}}{2 \tilde{L} K (t + 1)^{1/3}}
            \ge \frac{\gamma_t}{2},
        \ \textrm{otherwise.}
    \end{cases}
\end{equation}
Hence, $\gamma_t \ge \gamma_{t+1} \ge \gamma_t / 2$ always holds.
Similarly, $\eta_t \ge \eta_{t+1} \ge \eta_t / 2$.
We also have
\begin{equation}
    \frac{1}{\gamma_{t+1}^2} =
    \begin{cases}
    \frac{1}{\gamma_t^2},
        \ \textrm{if} \ \frac{S^{1/3}}{\tilde{L} K (t + 2)^{1/3}} > \min\{    \frac{1}{L_2}, \frac{1}{\mu}, \frac{1}{L_{21}}, \frac{1}{192 \tilde{L} K}\}, \\
    \frac{\tilde{L}^2 K^2 (t + 2)^{2/3}}{S^{2/3}}
        \le \frac{\tilde{L}^2 K^2 ( (t + 1)^{2/3} + 1 )}{S^{2/3}}
        \le \frac{1}{\gamma_t^2} + \frac{\tilde{L}^2 K^2}{S^{2/3}},
    \ \textrm{otherwise.}
    \end{cases}
\end{equation}
In either case, $\frac{1}{\gamma_{t+1}^2} \le \frac{1}{\gamma_t^2} + \frac{\tilde{L}^2 K^2}{S^{2/3}}$.
Thus, we can bound $a_4$ as follows
\begin{equation}  \label{eq_a4_bound_storm}
\begin{aligned}[b]
    a_4
    \le& \frac{48 \eta_{t+1}}{2304 \tilde{L}^2 K^2 \gamma_{t+1}^2}
        - 48 \eta_{t+1}
        + 4 \eta_t + \frac{21 L_{21}^2}{4} \cdot \frac{\eta_t}{42 L_{21}^2}
        + \frac{576 \tilde{L}^2 K \eta_t^2 \gamma_t}{2304 \tilde{L}^2 K^2 \gamma_t^2} \\
    \le{}& \frac{\eta_{t+1}}{48 \tilde{L}^2 K^2 \gamma_t^2}
        + \frac{\eta_{t+1}}{48 S^{2/3}}
        - 48 \eta_{t+1}
        + 4.375 \eta_t \\
    \le& \frac{\eta_{t+1}}{48 \tilde{L}^2 K^2 \gamma_t^2}
        + \frac{\eta_{t+1}}{48}
        - 48 \eta_{t+1}
        + 4.375 \eta_t
    \le \frac{\eta_{t+1}}{48 \tilde{L}^2 K^2 \gamma_t^2}
       - (48 - \frac{1}{48}) \frac{\eta_{t}}{2}
        + 4.375 \eta_t \\
    \le& \frac{\eta_{t+1}}{48 \tilde{L}^2 K^2 \gamma_t^2}
    \le \frac{\eta_{t}}{48 \tilde{L}^2 K^2 \gamma_t^2}
    = \frac{48 \eta_t}{\alpha_t}.
\end{aligned}
\end{equation}
Similarly, $a_5$ can be bounded by
\begin{equation}  \label{eq_a5_bound_storm}
\begin{aligned}[b]
    a_5
    \le& \frac{24 \gamma_{t+1}}{2304 \tilde{L}^2 K^2 \gamma_{t+1}^2}
        - 24 \gamma_{t+1}
        + \frac{3 \gamma_t}{2}
        + \frac{576 \tilde{L}^2 K \gamma_t^3}{2304 \tilde{L}^2 K^2 \gamma_t^2}
    \le \frac{24 \gamma_{t}}{\alpha_{t}}
        + \frac{\gamma_{t+1}}{96 S^{2/3}}
        - 24 \gamma_{t+1}
        + \frac{7 \gamma_t}{4} \\
    \le& \frac{24 \gamma_{t}}{\alpha_{t}}
        - (24 - \frac{1}{96}) \frac{\gamma_{t}}{2}
        + \frac{7 \gamma_t}{4}
    \le \frac{24 \gamma_{t}}{\alpha_{t}}.
\end{aligned}
\end{equation}
Next, $a_6$ and $a_7$ can be bounded by
\begin{equation}  \label{eq_a6_a7_bound_storm}
\begin{aligned}[b]
    a_6
    =& \eta_t \Big(\frac{1}{8} - \frac{\eta_t}{8 \gamma_t} \Big)
    \ge \eta_t \Big(\frac{1}{8} - \frac{\gamma_t}{8 \gamma_t} \Big)
    = 0, \\
    a_7
    =& \gamma_t \Big(\frac{1}{8} - \frac{1}{8} \Big) = 0.
\end{aligned}
\end{equation}
Finally, we bound $a_8$ as follows
\begin{equation}  \label{eq_a8_bound_storm}
    a_8
    \le \frac{384 L_f^2 \eta_{t+1}}{S \alpha_{t+1}} + \frac{192 L_f^2 \gamma_{t+1}}{S \alpha_{t+1}}
    \le \frac{576 L_f^2 \gamma_{t+1}}{S \alpha_{t+1}}
    = \frac{576 L_f^2}{2304 S \tilde{L}^2 K^2 \gamma_{t+1}}
    \le \frac{1}{8 S K^2 \gamma_{t}}.
\end{equation}
Plugging~\eqref{eq_a1_a2_a3_bound_storm}, \eqref{eq_a4_bound_storm}, \eqref{eq_a5_bound_storm},
\eqref{eq_a6_a7_bound_storm}, and~\eqref{eq_a8_bound_storm} into~\eqref{eq_sum_descent_lemma_potential_storm}, we obtain the desired result.
\end{proof}










\subsection{Proof of Theorem~\ref{theorem_fed_sgda_storm}}  \label{proof_theorem_fed_sgda_storm}

\begin{proof}
Since $\gamma_t \le 1 / (192 \tilde{L} K)$, $\frac{6 \tilde{L}^2 K \gamma_t}{\alpha_t} = \frac{6}{2304 K \gamma_t} \ge \frac{\tilde{L}}{2}$.
Thus, we can bound
\begin{equation}
\begin{aligned}[b]
    &\frac{1}{S} \sum_{i \in \mathcal{S}_t} \Big( \mathbb{E}[2 \Phi(\mathbf{x}_{t, i}^{(K)}) - f(\mathbf{z}_{t, i}^{(K)})]
        + \frac{6 \tilde{L}^2 K \gamma}{\alpha_t} \mathbb{E}[ \|\mathbf{z}_{t, i}^{(K)} - \bar{\mathbf{z}}_t\|^2 ] \Big) \\
    \ge& \mathbb{E}[2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})]
        + \frac{6 \tilde{L}^2 K \gamma}{\alpha_t} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2
    \ge \mathbb{E}[2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})],
\end{aligned}
\end{equation}
where the first inequality follows from~\eqref{eq_weakly_convex} and the fact that $\bar{\mathbf{z}}_{t+1} = \frac{1}{S} \sum_{i \in \mathcal{S}_t} \mathbf{z}_{t, i}^{(K)}$.
Note that $\mathbb{E}[ \|\mathbf{z}_{t, i}^{(0)} - \bar{\mathbf{z}}_t\|^2] = 0$ since $\mathbf{z}_{t, i}^{(0)} = \bar{\mathbf{z}}_t$.
Averaging~\eqref{eq_potential_STORM_descent} for $i \in \mathcal{S}_t$ and $k \in \{0, \ldots, K-1\}$, we get
\begin{equation}
\begin{aligned}[b]
    &\frac{3 \eta_t}{4 S K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \Big( \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_t) - f(\bar{\mathbf{z}}_t)]
        + \frac{48 \eta_t}{\alpha_t} \varepsilon_{\mathbf{x}, t}
        + \frac{24 \gamma_t}{\alpha_t} \varepsilon_{\mathbf{y}, t} \Big)
        + 96 \eta_{t+1} \alpha_{t+1} \frac{\sigma_1^2}{S}
        + 48 \gamma_{t+1} \alpha_{t+1} \frac{\sigma_2^2}{S} \\
        &- \Big( \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})]
        + \frac{48 \eta_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{x}, t+1}
        + \frac{24 \gamma_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{y}, t+1} \Big) \\
        &+ \frac{1}{8 S K^2 \gamma_t} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ]
        - \frac{1}{S K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} (\frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        + \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2]).
\end{aligned}
\end{equation}
Since $\mathbf{x}_{t, i}^{(k + 1)} = \mathbf{x}_{t, i}^{(k)} - \eta_t \mathbf{d}_{\mathbf{x}}^{(k)}$
and $\mathbf{y}_{t, i}^{(k + 1)} = \mathbf{y}_{t, i}^{(k)} - \gamma_t \mathbf{d}_{\mathbf{y}}^{(k)}$,
we have
\begin{equation}
\begin{aligned}[b]
    \frac{\eta_t}{8} \mathbb{E}[\| \mathbf{d}_{\mathbf{x}}^{(k)}\|^2]
        + \frac{\gamma_t}{4} \mathbb{E}[\| \mathbf{d}_{\mathbf{y}}^{(k)} \|^2]
    = \mathbb{E} \left[ \frac{1}{8 \eta_t} \left\|\mathbf{x}_{t, i}^{(k + 1)} - \mathbf{x}_{t, i}^{(k)} \right\|^2
        + \frac{1}{4 \gamma_t} \left\|\mathbf{y}_{t, i}^{(k + 1)} - \mathbf{y}_{t, i}^{(k)} \right\|^2 \right]
    \ge \frac{1}{8 \gamma_t} \mathbb{E} \left[ \left\|\mathbf{z}_{t, i}^{(k + 1)} - \mathbf{z}_{t, i}^{(k)} \right\|^2 \right],
\end{aligned}
\end{equation}
where the inequality follows from the condition $\eta_t \le \gamma_t$.
By Cauchy-Schwarz inequality, we have
\begin{equation}
\begin{aligned}[b]
    \frac{1}{S K} \sum_{i \in \mathcal{S}_t} \sum_{k = 0}^{K - 1} \left\|\mathbf{z}_{t, i}^{(k + 1)} - \mathbf{z}_{t, i}^{(k)} \right\|^2
    \ge& \frac{1}{S K^2} \sum_{i \in \mathcal{S}_t} \left\| \sum_{k=0}^{K - 1} \mathbf{z}_{t, i}^{(k + 1)} - \mathbf{z}_{t, i}^{(k)} \right\|^2
    = \frac{1}{S K^2} \sum_{i \in \mathcal{S}_t} \left\| \mathbf{z}_{t, i}^{(K)} - \mathbf{z}_{t, i}^{(0)} \right\|^2 \\
    \ge& \frac{1}{K^2} \left\| \frac{1}{S} \sum_{i \in \mathcal{S}_t} \mathbf{z}_{t, i}^{(K)} - \bar{\mathbf{z}}_t \right\|^2
    = \frac{1}{K^2} \left\| \bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t \right\|^2,
\end{aligned}
\end{equation}
where the second inequality follows from Jensen's inequality and the convexity of $\|\cdot\|^2$.
Combining the above results, we obtain
\begin{equation}
\begin{aligned}[b]
    &\frac{3 \eta_t}{4 S K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \Big( \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_t) - f(\bar{\mathbf{z}}_t)]
        + \frac{48 \eta_t}{\alpha_t} \varepsilon_{\mathbf{x}, t}
        + \frac{24 \gamma_t}{\alpha_t} \varepsilon_{\mathbf{y}, t} \Big)
        + 96 \eta_{t+1} \alpha_{t+1} \frac{\sigma_1^2}{S}
        + 48 \gamma_{t+1} \alpha_{t+1} \frac{\sigma_2^2}{S} \\
        &- \Big( \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})]
        + \frac{48 \eta_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{x}, t+1}
        + \frac{24 \gamma_{t+1}}{\alpha_{t+1}} \varepsilon_{\mathbf{y}, t+1} \Big).
\end{aligned}
\end{equation}

Summing the above inequality from $t = 0$ to $T-1$, we get
\begin{equation}  \label{eq_sum_all_rounds_storm}
\begin{aligned}[b]
    &\frac{3}{4 S K} \sum_{t=0}^{T-1} \eta_t \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \frac{C_0}{K}
        + \frac{48 \eta_0}{\alpha_0} \varepsilon_{\mathbf{x}, 0}
        + \frac{24 \gamma_0}{\alpha_0} \varepsilon_{\mathbf{y}, 0}
        + \frac{96 \sigma_1^2}{S} \sum_{t=0}^{T-1} \eta_{t+1} \alpha_{t+1}
        + \frac{48 \sigma_2^2}{S} \sum_{t=0}^{T-1} \gamma_{t+1} \alpha_{t+1} \\
    =& \frac{C_0}{K}
        + \frac{48 \sigma_1^2 \eta_0}{S \alpha_0}
        + \frac{24 \sigma_2^2 \gamma_0}{S \alpha_0}
        + \frac{96 \sigma_1^2}{S} \sum_{t=0}^{T-1} \eta_{t+1} \alpha_{t+1}
        + \frac{48 \sigma_2^2}{S} \sum_{t=0}^{T-1} \gamma_{t+1} \alpha_{t+1}.
\end{aligned}
\end{equation}
Since the sequence $\{\eta_t\}$ is non-increasing,
\begin{equation}  \label{eq_eta_lower_bound}
    \frac{3 \eta_{T-1}}{4 S K} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le \frac{3}{4 S K} \sum_{t=0}^{T-1} \eta_t \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2].
\end{equation}
Plugging~\eqref{eq_eta_lower_bound} into~\eqref{eq_sum_all_rounds_storm} implies that
\begin{equation}
\begin{aligned}[b]
    &\frac{1}{S K T} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \frac{4 C_0}{3 K T \eta_{T-1}}
        + \frac{64 \sigma_1^2 \eta_0}{S \alpha_0 T \eta_{T-1}}
        + \frac{32 \sigma_2^2 \gamma_0}{S \alpha_0 T \eta_{T-1}}
        + \frac{128 \sigma_1^2}{S T \eta_{T-1}} \sum_{t=0}^{T-1} \eta_{t+1} \alpha_{t+1}
        + \frac{64 \sigma_2^2}{S T \eta_{T-1}} \sum_{t=0}^{T-1} \gamma_{t+1} \alpha_{t+1}.
\end{aligned}
\end{equation}
Plugging the choice of $\eta_t, \gamma_t$, and $\alpha_t$ into the above inequality, we obtain
\begin{equation}
\begin{aligned}[b]
    &\frac{1}{S K T} \sum_{t=0}^{T-1} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}^{(k)})\|^2] \\
    \le& \mathcal{O} \Big(
        \frac{1}{T} \big( L_{\Phi} + \frac{L_{12}^2 (L_2 + L_{21} + \mu + \tilde{L} K)}{\mu^2} \big)
        + \frac{L_{12}^2 \tilde{L} K }{T^{2/3} S^{1/3} \mu^2} \Big)
    \Big( \frac{C_0}{K}
        + (\sigma_1^2 + \sigma_2^2) \big( \frac{L_2 + \mu + L_{21}}{\tilde{L}^2 K^2} + \frac{\log T}{\tilde{L} K} \big)
    \Big),
\end{aligned}
\end{equation}
which is the desired result.
\end{proof}


\section{{FedSGDA} incorporated with the SPIDER Estimator}
In addition to the estimators introduced in Section~\ref{section_fedsgda_instances}, {FedSGDA} can be incorporated with other estimators.
In this section, we propose {{\FrameworkAbbr}-SPIDER} by incorporating {FedSGDA} with the SPIDER estimator~\citep{nguyen2017sarah, fang2018spider}.
Unlike the minibatch gradient estimator and STORM which use a constant batch size of clients, SPIDER will use varying batch sizes. 
We denote the batch size of clients in the $t$-th round as $S_t := |\mathcal{S}_t'| = |\mathcal{S}_t|$.
The SPIDER estimator is detailed in Estimator~\ref{estimator_spider}.
 
\begin{estimator}[h]
\caption{SPIDER}
\label{estimator_spider}
\begin{algorithmic}
    \IF{$\mathrm{mod}(t, m) = 0$}
        \STATE $\mathbf{u}_t = \frac{1}{S_t} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$;

        \STATE $\mathbf{v}_t = \frac{1}{S_t} \sum_{i \in \mathcal{S}_t'} \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)$;
    \ELSE
        \STATE $\mathbf{u}_{t} \leftarrow
            \mathbf{u}_{t-1} + \frac{1}{S_t} \sum_{i \in \mathcal{S}_t'} \big(\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}) \big)$;

        \STATE $\mathbf{v}_{t} \leftarrow
            \mathbf{v}_{t-1} + \frac{1}{S_t} \sum_{i \in \mathcal{S}_t'} \big(\nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{y}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}) \big)$;
    \ENDIF
\end{algorithmic}
\end{estimator}

The gradient estimation errors $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$ of SPIDER can be bounded as follows.

\begin{lemma}  \label{lemma_gradient_error_spider}
Let $S_t = S$ if $\mod(T, m) \neq 0$, and otherwise $S_t = \hat{S}$.
The gradient estimation errors $\varepsilon_{\mathbf{x}, t}$ and $\varepsilon_{\mathbf{y}, t}$ of the SPIDER estimator can be bounded from above by
$$
\begin{aligned}[b]
    \varepsilon_{\mathbf{x}, t}
    \le&
    \begin{cases}
        \frac{\sigma_1^2}{\hat{S}}, \ \mathrm{if} \ \mathrm{mod}(t, m) = 0, \\
        \varepsilon_{\mathbf{x}, t-1} + \frac{4 L_f^2}{S} \mathbb{E}[ \|\bar{\mathbf{z}}_t - \bar{\mathbf{z}}_{t-1}\|^2 ], \ \mathrm{otherwise.}
    \end{cases} \\
    \varepsilon_{\mathbf{y}, t}
    \le&
    \begin{cases}
        \frac{\sigma_2^2}{\hat{S}}, \ \mathrm{if} \ \mathrm{mod}(t, m) = 0, \\
        \varepsilon_{\mathbf{y}, t-1} + \frac{4 L_f^2}{S} \mathbb{E}[ \|\bar{\mathbf{z}}_t - \bar{\mathbf{z}}_{t-1}\|^2 ], \ \mathrm{otherwise.}
    \end{cases}
\end{aligned}
$$
\end{lemma}

\begin{proof}
When $\mathrm{mod}(t, m) = 0$, we can bound
\begin{equation}
\begin{aligned}[b]
    \varepsilon_{\mathbf{x}, t}
    = \mathbb{E}[ \|\frac{1}{\hat{S}} \sum_{i \in \mathcal{S}_t'} \nabla f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t)\|^2 ]
    \le \frac{\sigma_1^2}{\hat{S}}.
\end{aligned}
\end{equation}
When $\mathrm{mod}(t, m) \neq 0$, we can bound
\begin{equation}
\begin{aligned}[b]
    \varepsilon_{\mathbf{x}, t}
    =& \mathbb{E}[ \|
        \mathbf{u}_{t-1}
        + \frac{1}{S} \sum_{i \in \mathcal{S}_t'} \big(\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}) \big)
        - \nabla f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) \|^2 ] \\
    =& \mathbb{E}[ \|
        \mathbf{u}_{t-1} - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})
        + \frac{1}{S} \sum_{i \in \mathcal{S}_t'} \big(\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}) \big)
        + \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})
        - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) \|^2 ] \\
    =& \mathbb{E}[ \|
        \mathbf{u}_{t-1} - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1}) \|^2 ]
        + \frac{1}{S} \mathbb{E}[ \|\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})
        + \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})
        - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) \|^2 ] \\
    \le{}& \varepsilon_{\mathbf{x}, t-1}
        + \frac{4 L_f^2}{S} \mathbb{E}[\|\bar{\mathbf{x}}_t - \bar{\mathbf{x}}_{t-1}\|^2 + \|\bar{\mathbf{y}}_t - \bar{\mathbf{y}}_{t-1}\|^2],
\end{aligned}
\end{equation}
where the second equality holds because $\mathbb{E}_{\mathcal{S}_t'} [\nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f_i(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})] = \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_t, \bar{\mathbf{y}}_t) - \nabla_{\mathbf{x}} f(\bar{\mathbf{x}}_{t-1}, \bar{\mathbf{y}}_{t-1})$
and the inequality follows from Lemma~\ref{lemma_f_smooth}.
The upper bound of $\varepsilon_{\mathbf{y}, t}$ follows from the same argument.
\end{proof}



By combining Lemma~\ref{lemma_potential_recursion_sum} and Lemma~\ref{lemma_gradient_error_spider} and constructing another potential function, we establish the convergence rate of {{\FrameworkAbbr}-SPIDER} in the following theorem.
\begin{theorem}  \label{theorem_fed_sgda_spider}
Consider {{\FrameworkAbbr}-SPIDER}.
Denote $C_0 := \Phi(\mathbf{x}_0) - \Phi(\mathbf{x}^*) + \Phi(\mathbf{x}_0) - f(\mathbf{x}_0, \mathbf{y}_0)$, where $\mathbf{x}^* \in \mathop{\rm argmin}_{\mathbf{x} \in \mathbb{R}^{p}} \Phi(\mathbf{x})$.
Let $S_t = \hat{S}$ if $\mathrm{mod}(t, m) = 0$, otherwise $S_t = S = m$.
Suppose that $C_0 < \infty$, $\mathrm{mod}(T, m) = 0$, and $\eta_t, \gamma_t$ satisfy
$$
\begin{aligned}
\eta_t
    ={}& \eta
    = \min\left\{
    \frac{1}{4 L_1},
    \frac{1}{L_{\Phi}},
    \frac{1}{42 L_{21}},
    \gamma_t,
    \frac{\mu^2}{3 L_{12}^2} \gamma_t
\right\}, \\
\gamma_t
    ={}& \gamma
    = \min\left\{
    \frac{1}{2 L_2},
    \frac{1}{\mu},
    \frac{1}{L_{21}},
    \frac{1}{192 \tilde{L} K}
\right\}.
\end{aligned}
$$
Under Assumptions~\ref{assumption_bgd}-\ref{assumption_pl}, we have
$$
\begin{aligned}[b]
    &\frac{1}{K T} \sum_{t = 0}^{T - 1} \frac{1}{S_t} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t}
    \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \mathcal{O} \Big(
        \frac{C_0}{K T} \big( L_{\Phi} + L_{21} + \frac{L_{12}^2}{\mu^2} (L_2 + \mu + L_{21} + \tilde{L} K) \big) \\
        &+ \frac{\sigma_1^2}{\hat{S}} + \frac{\sigma_2^2}{\hat{S}} \big( \frac{L_{\Phi}}{\tilde{L} K}  + \frac{L_{21}}{\tilde{L} K} + \frac{L_{12}^2}{\mu^2} \big)
    \Big).
\end{aligned}
$$
\end{theorem}

\begin{proof}
First, we define  
\begin{equation}
    \hat{\mathcal{L}}_t
    := \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{t}) - f(\bar{\mathbf{z}}_{t})]
        +\frac{35 \eta (m - \mathrm{mod}(t, m))}{8} \varepsilon_{\mathbf{x}, t}
        + \frac{7 \gamma (m - \mathrm{mod}(t, m))}{4} \varepsilon_{\mathbf{y}, t}.
\end{equation}
When $\mathrm{mod}(t + 1, m) \neq 0$, by adding $\frac{35 \eta (m - \mathrm{mod}(t + 1, m))}{8} \varepsilon_{\mathbf{x}, t+1}
+ \frac{7 \gamma (m - \mathrm{mod}(t + 1, m))}{4} \varepsilon_{\mathbf{y}, t+1}$
to both sides of~\eqref{eq_potential_recursion_sum} and applying Lemma~\ref{lemma_gradient_error_spider}, we get
\begin{equation}
\begin{aligned}[b]
    \hat{\mathcal{L}}_{t+1}
    \le& \hat{\mathcal{L}}_t
        + (\frac{35 \eta}{8} + \frac{7 \gamma}{4}) \mathrm{mod}(t + 1, m) \frac{4 L_f^2}{S} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2] \\
        &- \frac{3 \eta}{4 S_t K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
        - \frac{1}{8 K^2 \gamma} \mathbb{E}[ \|\bar{\mathbf{z}}_{t+1} - \bar{\mathbf{z}}_t\|^2 ]
\end{aligned}
\end{equation}
Since $S = m$, $\eta \le \gamma$, and $\gamma \le \frac{1}{192 \tilde{L} K}$, we have
\begin{equation}
    (\frac{35 \eta}{8} + \frac{7 \gamma}{4}) \mathrm{mod}(t + 1, m) \frac{4 L_f^2}{S} - \frac{1}{8 K^2 \gamma}
    \le \frac{49 \gamma L_f^2}{2} - \frac{1}{8 K^2 \gamma}
    = \frac{196 \gamma^2 L_f^2 K^2 -1}{8 K^2 \gamma}
    \le 0.
\end{equation}
Combining the above two inequalities, we have, for any $t$ such that $\mathrm{mod}(t+1, m) \neq 0$,
\begin{equation}  \label{eq_potential_recursion_spider}
\begin{aligned}[b]
    &\frac{3 \eta}{4 S_t K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    + \frac{35 \eta (m - \mathrm{mod}(t + 1, m))}{8} \varepsilon_{\mathbf{x}, t+1}
    + \frac{7 \gamma (m - \mathrm{mod}(t + 1, m))}{4} \varepsilon_{\mathbf{y}, t+1} \\
    \le& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_t) - f(\bar{\mathbf{z}}_t)]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{t+1}) - f(\bar{\mathbf{z}}_{t+1})]
        + \frac{35 \eta (m - \mathrm{mod}(t, m))}{8} \varepsilon_{\mathbf{x}, t}
        + \frac{7 \gamma (m - \mathrm{mod}(t, m))}{4} \varepsilon_{\mathbf{y}, t}.
\end{aligned}
\end{equation}
Summing~\eqref{eq_potential_recursion_spider} from $t = \ell m$ to $t = (\ell + 1) m -2$ for some integer $\ell \ge 0$, we obtain
\begin{equation}
\begin{aligned}[b]
    &\frac{3 \eta}{4 K} \sum_{t = \ell m}^{(\ell + 1) m - 2} \frac{1}{S_t} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    + \frac{35 \eta}{8} \varepsilon_{\mathbf{x}, (\ell + 1) m - 1}
    + \frac{7 \gamma}{4} \varepsilon_{\mathbf{y}, (\ell + 1) m - 1} \\
    \le& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{\ell m}) - f(\bar{\mathbf{z}}_{\ell m})]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{(\ell + 1) m - 1}) - f(\bar{\mathbf{z}}_{(\ell + 1) m - 1})]
        + \frac{35 \eta m}{8} \varepsilon_{\mathbf{x}, \ell m}
        + \frac{7 \gamma m}{4} \varepsilon_{\mathbf{y}, \ell m} \\
    \le& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{\ell m}) - f(\bar{\mathbf{z}}_{\ell m})]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{(\ell + 1) m - 1}) - f(\bar{\mathbf{z}}_{(\ell + 1) m - 1})]
        + \frac{35 \eta m \sigma_1^2}{8 \hat{S}}
        + \frac{7 \gamma m \sigma_2^2}{4 \hat{S}}.
\end{aligned}
\end{equation}
When $t = (\ell + 1) m - 1$, by~\eqref{eq_potential_recursion_sum}, we have
\begin{equation}
\begin{aligned}[b]
    \frac{3 \eta}{4 S_t K} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2]
    \le& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{(\ell + 1) m - 1}) - f(\bar{\mathbf{z}}_{(\ell + 1) m - 1})]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{(\ell + 1) m}) - f(\bar{\mathbf{z}}_{(\ell + 1) m})] \\
        &+ \frac{35 \eta}{8} \varepsilon_{\mathbf{x}, (\ell + 1) m - 1}
        + \frac{7 \gamma}{4} \varepsilon_{\mathbf{y}, (\ell + 1) m - 1}.
\end{aligned}
\end{equation}
Summing the above two inequalities yields that
\begin{equation}  \label{eq_potential_recursion_block}
\begin{aligned}[b]
    &\frac{3 \eta}{4 K} \sum_{t = \ell m}^{(\ell + 1) m - 1} \frac{1}{S_t} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{\ell m}) - f(\bar{\mathbf{z}}_{\ell m})]
        - \frac{1}{K} \mathbb{E}[ 2 \Phi(\bar{\mathbf{x}}_{(\ell + 1) m}) - f(\bar{\mathbf{z}}_{(\ell + 1) m})]
        + \frac{35 \eta m \sigma_1^2}{8 \hat{S}}
        + \frac{7 \gamma m \sigma_2^2}{4 \hat{S}}.
\end{aligned}
\end{equation}
Summing~\eqref{eq_potential_recursion_block} for $\ell = 0$ to $T / m - 1$ and rearranging terms, we arrive at
\begin{equation}
\begin{aligned}[b]
    &\frac{1}{K T} \sum_{t = 0}^{T - 1} \frac{1}{S_t} \sum_{k=0}^{K-1} \sum_{i \in \mathcal{S}_t} \mathbb{E}[\|\nabla \Phi(\mathbf{x}_{t, i}^{(k)})\|^2] \\
    \le& \frac{4}{3 \eta K T} C_0
        + \frac{4}{3 \eta T} (\frac{35 \eta m \sigma_1^2}{8 \hat{S}}
        + \frac{7 \gamma m \sigma_2^2}{4 \hat{S}}) \frac{T}{m}
    \le \frac{4}{3 \eta K T} C_0
        + \frac{35 \sigma_1^2}{6 \hat{S}}
        + \frac{7 \gamma \sigma_2^2}{3 \eta \hat{S}} \\
    \le& \mathcal{O} \Big(
        \frac{C_0}{K T} \big( L_{\Phi} + L_{21} + \frac{L_{12}^2}{\mu^2} (L_2 + \mu + L_{21} + \tilde{L} K) \big)
        + \frac{\sigma_1^2}{\hat{S}} + \frac{\sigma_2^2}{\hat{S}} \big( \frac{L_{\Phi}}{\tilde{L} K}  + \frac{L_{21}}{\tilde{L} K} + \frac{L_{12}^2}{\mu^2} \big)
    \Big),
\end{aligned}
\end{equation}
where the last inequality follows from the choice of $\eta$ and $\gamma$.
This completes the proof.
\end{proof}



\begin{remark}
    Similar to {{\FrameworkAbbr}-MB}, {{\FrameworkAbbr}-SPIDER} converges to a close neighborhood of a stationary point as $T \rightarrow +\infty$ if $\hat{S}$ is large.
    The average number of clients sampled in each round is $\bar{S} = \frac{\hat{S}}{S} + S - 1$.
    The minimum value of $\bar{S}$ is $2 S - 1$ when $\hat{S} = S^2$.
\end{remark}

\begin{remark}
We note that when $\hat{S} = S^2$, the convergence rate of {{\FrameworkAbbr}-SPIDER} is better than {{\FrameworkAbbr}-MB}.
Specifically, the convergence rate of {{\FrameworkAbbr}-SPIDER} is $\mathcal{O} \Big(
    \frac{C_0}{K T} \big( L_{\Phi} + L_{21} + \frac{L_{12}^2}{\mu^2} (L_2 + \mu + L_{21} + \tilde{L} K) \big)
    + \frac{\sigma_1^2}{\bar{S}^2} + \frac{\sigma_2^2}{\bar{S}^2} \big( \frac{L_{\Phi}}{\tilde{L} K}  + \frac{L_{21}}{\tilde{L} K} + \frac{L_{12}^2}{\mu^2} \big)
\Big)$,
while {{\FrameworkAbbr}-MB} with the same batch size $\bar{S}$ of sampled clients achieves a rate of $\mathcal{O} \Big(
    \frac{C_0}{K T} \big( L_{\Phi} + L_{21} + \frac{L_{12}^2}{\mu^2} (L_2 + \mu + L_{21} + \tilde{L} K) \big)
    + \frac{\sigma_1^2}{\bar{S}} + \frac{\sigma_2^2}{\bar{S}} \big( \frac{L_{\Phi}}{\tilde{L} K} + \frac{L_{21}}{\tilde{L} K} + \frac{L_{12}^2}{\mu^2} \big)
\Big)$.
However, {{\FrameworkAbbr}-SPIDER} is less practical compared to {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} because it periodically draws a fairly large subset of sampled clients (of size $S^2$). 
Hence, we exclude {{\FrameworkAbbr}-SPIDER} from our experiments.
\end{remark}









\begin{figure*}[h]
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 0.8cm 0cm},clip,width=\linewidth]{./imgs/toy/MB_K5}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 0.8cm 0cm},clip,width=\linewidth]{./imgs/toy/MB_K10}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 0.8cm 0cm},clip,width=\linewidth]{./imgs/toy/MB_K20}
    \end{subfigure}
    \\
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 0.8cm 0cm},clip,width=\linewidth]{./imgs/toy/STORM_K5}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 0.8cm 0cm},clip,width=\linewidth]{./imgs/toy/STORM_K10}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 0.8cm 0cm},clip,width=\linewidth]{./imgs/toy/STORM_K20}
    \end{subfigure}
    \caption{
        Results on the simulation experiment where $S \in \{5, 10, 20\}$ and $K \in \{5, 10, 20\}$.
        The training batch size of data samples in each local iteration is set to $5$.
    }
    \label{figure_toy_varing_K}
\end{figure*}

\begin{figure*}[!htbp]
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_10/MB5}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_10/STORM5}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_10/best5}
    \end{subfigure}
    \\
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_10/MB10}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_10/STORM10}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_10/best10}
    \end{subfigure}
    \caption{
        Results on the AUC maximization problem when $S = 10$.
        In the top (resp., bottom) row, the training batch size in each local iteration is set to $5$ (resp., $10$).
        The left (middle) column shows the test AUC of {{\FrameworkAbbr}-MB} ({{\FrameworkAbbr}-STORM}) with different values of $K$.
        The right column compares {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} when $K = 20$.
    }
    \label{figure_auc_S_10}
    \vspace{1em}
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_20/MB5}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_20/STORM5}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_20/best5}
    \end{subfigure}
    \\
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_20/MB10}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_20/STORM10}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_20/best10}
    \end{subfigure}
    \caption{
        Results on the AUC maximization problem when $S = 20$.
        In the top (resp., bottom) row, the training batch size in each local iteration is set to $5$ (resp., $10$).
        The left (middle) column shows the test AUC of {{\FrameworkAbbr}-MB} ({{\FrameworkAbbr}-STORM}) with different values of $K$.
        The right column compares {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} when $K = 20$.
    }
    \label{figure_auc_S_20}
\end{figure*}

\section{Additional Results}
\subsection{Additional Simulation Results}


Figure~\ref{figure_toy_varing_K} shows the simulation results of {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} with different $S$ and $K$.
We observe that as $S$ increases, both {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} converge faster.
Besides, both algorithms achieve the best performance when $S = 20$ and $K = 20$.



\subsection{Additional Results for AUC maximization}
\subsubsection{Results of $S = 10$ and $S = 20$}
Figure~\ref{figure_auc_S_10} and~\ref{figure_auc_S_20} present results of the proposed algorithms with $S = 10$ and $S = 20$ in the same setting as in Section~\ref{section_experiment_auc}, respectively.
The results are similar to those in Figure~\ref{figure_auc}.
Besides, we observe from the last column that when $K = 20$, both {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} achieve higher test AUC values as $S$ increases.


\subsubsection{Results on Sorted Data}
We further conduct experiments on the CIFAR-10 dataset where $10\%$ of the training data are randomly selected and allocated to clients, and the remaining $90\%$ are first sorted by label and then allocated to clients.
We refer the reader to~\citep{hsu2019measuring} for details of this manner of data allocation. 

Figures~\ref{figure_auc_sort_by_label_a} and~\ref{figure_auc_sort_by_label_b} present the results of {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} on a network of $500$ nodes.
We fix $S$ and $T$ and vary $K$ in the range $\{1, 10, 20\}$.
It can be observed that both algorithms with multiple local steps attain higher AUC values than only one local step.
Additionally, the result of $K=10$ is close to $K = 20$.
This phenomenon coincides with the discussion in Remark~\ref{remark_mb} and Remark~\ref{remark_storm}.  

In Figure~\ref{figure_auc_sort_by_label_c}, we compare the performance of the proposed algorithms and CoDA~\citep{guo2020communication}.
Since CoDA requires all clients to participate in each round, we run the algorithms on a network of only $50$ clients.
For fair comparison, we run {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} for $T = 400$ rounds and CoDA for $800$ rounds because {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} involve two communication steps while CoDA only involves one in each round.
The horizontal axis in Figure~\ref{figure_auc_sort_by_label_c} stands for the number of communication steps.
It is shown that both {{\FrameworkAbbr}-MB} and {{\FrameworkAbbr}-STORM} outperform CoDA.
\begin{figure*}[!htbp]
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_5/MB_S_5_B_5_sort_by_label}
    \caption{}
    \label{figure_auc_sort_by_label_a}
    \end{subfigure}
    \hfill
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_5/STORM_S_5_B_5_sort_by_label}
    \caption{}
    \label{figure_auc_sort_by_label_b}
    \end{subfigure}
    \begin{subfigure}{.32\textwidth}
        \centering
        \includegraphics[trim={0cm 0cm 1cm 0cm},clip,width=\linewidth]{./imgs/auc/node_5/FedSGDA_vs_CoDA}
    \caption{}
    \label{figure_auc_sort_by_label_c}
    \end{subfigure}
    \caption{
        Results on the AUC maximization problem where $90\%$ of data points are sorted by label.
        The training batch size in each local iteration is set to $10$.
    }
    \label{figure_auc_sort_by_label}
\end{figure*}



\section{Additional Details of Experiments}  \label{section_additional_experimental_details}

\subsection{Specification of Hyper-parameters}  \label{section_specification_of_hyperparameter}
The hyper-parameters of {{\FrameworkAbbr}-MB} are the constant stepsizes $\eta$ and $\gamma$, while those of {{\FrameworkAbbr}-STORM} are $\eta_t$, $\gamma_t$, $\alpha_t$, and $\beta_t$.
In our implementation, we set $\eta_t = \frac{c_{\eta}}{(t+1)^{\rho}}$, $\gamma_t = \frac{c_{\gamma}}{(t+1)^{\rho}}$, and $\alpha_t = \beta_t = \min\{1, \frac{c_{\alpha}}{(t+1)^{2 \rho}}\}$, where $c_{\eta}$, $c_{\gamma}$, $c_{\alpha}$, and $\rho$ are tunable variables.

\subsubsection{Hyper-parameters for the Simulation}
For the simulation results, the range of the tuning values is listed as follows.
\begin{itemize}
    \item For {{\FrameworkAbbr}-MB}, both $\eta$ and $\gamma$ are selected from $\{\num{1e-4}, \num{1e-3}, \num{1e-2}, \num{1e-1}\}$.
    \item For {{\FrameworkAbbr}-STORM}, both $c_{\eta}$ and $c_{\gamma}$ are tuned from the set $\{\num{1e-4}, \num{1e-3}, \num{1e-2}, \num{1e-1}\}$, $c_{\alpha}$ is selected from $\{1, 10\}$, and $\rho$ is set to $1/3$ as suggested by the theory.
\end{itemize}

The best hyper-parameters for each algorithm are displayed in Table~\ref{table_toy_parameters}.

\begin{table}[!htb]
\centering
\caption{Hyper-parameters for the simulation experiment in Figures~\ref{figure_toy} and~\ref{figure_toy_varing_K}.
$S$ and $K$ stand for the number of sampled clients in each round and the number of local iterations, respectively.}
\label{table_toy_parameters}
\begin{tabular}{cccccccccc}
\toprule
\multicolumn{3}{c}{Setting} & \phantom{a} & \multicolumn{2}{c}{{{\FrameworkAbbr}-MB}} &  \phantom{a} & \multicolumn{3}{c}{{{\FrameworkAbbr}-STORM}}          \\ \cmidrule{1-3} \cmidrule{5-6} \cmidrule{8-10}
S & K & minibatch size && $\eta$ & $\gamma$ && $c_\eta$ & $c_\gamma$ & $c_\alpha$ \\ \midrule
5 & 5 & 5 && $\num{1e-3}$ & $\num{1e-3}$ && $\num{1e-2}$ & $\num{1e-1}$ & 1 \\
5 & 10 & 5 && $\num{1e-4}$ & $\num{1e-3}$ && $\num{1e-3}$ & $\num{1e-2}$ & 10 \\
5 & 20 & 5 && $\num{1e-4}$ & $\num{1e-3}$ && $\num{1e-3}$ & $\num{1e-1}$ & 1 \\ \midrule
10 & 5 & 5 && $\num{1e-3}$ & $\num{1e-3}$ && $\num{1e-2}$ & $\num{1e-1}$ & 1 \\
10 & 10 & 5 && $\num{1e-4}$ & $\num{1e-3}$ && $\num{1e-3}$ & $\num{1e-2}$ & 1 \\
10 & 20 & 5 && $\num{1e-4}$ & $\num{1e-3}$ && $\num{1e-3}$ & $\num{1e-2}$ & 1 \\ \midrule
20 & 5 & 5 && $\num{1e-3}$ & $\num{1e-3}$ && $\num{1e-2}$ & $\num{1e-1}$ & 1 \\
20 & 10 & 5 && $\num{1e-3}$ & $\num{1e-2}$ && $\num{1e-3}$ & $\num{1e-2}$ & 1 \\
20 & 20 & 5 && $\num{1e-4}$ & $\num{1e-3}$ && $\num{1e-3}$ & $\num{1e-3}$ & 1 \\ \bottomrule

\end{tabular}
\end{table}

\begin{table}[!htb]
\centering
\caption{Hyper-parameters for the AUC maximization experiment in Figures~\ref{figure_auc},~\ref{figure_auc_S_10}, and~\ref{figure_auc_S_20}.}
\label{table_auc_parameters}
\begin{tabular}{ccccccccccc}
\toprule
\multicolumn{3}{c}{Setting} & \phantom{a} & \multicolumn{2}{c}{{{\FrameworkAbbr}-MB}} & \phantom{a} & \multicolumn{4}{c}{{{\FrameworkAbbr}-STORM}} \\
\cmidrule{1-3} \cmidrule{5-6} \cmidrule{8-11}
S & K & minibatch size && $\eta$ & $\gamma$       && $c_\eta$ & $c_\gamma$ & $c_\alpha$ & $\rho$ \\
\midrule
5 & 1 & 5 && 1 & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
5 & 1 & 10 && 1 & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
5 & 10 & 5 && $\num{1e-1}$ & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 10 & $1/5$ \\
5 & 10 & 10 && 1 & $\num{1e-1}$ && 1 & 1 & 10 & $1/5$ \\
5 & 20 & 5 && $\num{1e-1}$ & $\num{1e-1}$ && 1 & 1 & 10 & $1/5$ \\
5 & 20 & 10 && 1 & 1 && 1 & $\num{1e-1}$ & 10 & $1/5$ \\
\midrule
10 & 1 & 5 && 1 & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
10 & 1 & 10 && 1 & $\num{1e-3}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
10 & 10 & 5 && 1 & 1 && 1 & 1 & 5 & $1/5$ \\
10 & 10 & 10 && 1 & $\num{1e-1}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
10 & 20 & 5 && 1 & $\num{1e-1}$ && 1 & 1 & 10 & $1/5$ \\
10 & 20 & 10 && 1 & 1 && 1 & 1 & 10 & $1/5$ \\
\midrule
20 & 1 & 5 && 1 & $\num{1e-3}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
20 & 1 & 10 && 1 & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
20 & 10 & 5 && 1 & 1 && 1 & 1 & 5 & $1/5$ \\
20 & 10 & 10 && 1 & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 10 & $1/5$ \\
20 & 20 & 5 && 1 & $\num{1e-1}$ && 1 & $\num{1e-1}$ & 10 & $1/5$ \\
20 & 20 & 10 && 1 & 1 && 1 & 1 & 10 & $1/5$ \\
\bottomrule
\end{tabular}
\end{table}




\subsubsection{Hyper-parameters for AUC Maximization}
For the AUC maximization task, the hyper-parameters are chosen as follows.
\begin{itemize}
    \item For {{\FrameworkAbbr}-MB}, both $\eta$ and $\gamma$ are chosen from the set $\{\num{1e-3}, \num{1e-2}, \num{1e-1}, \num{1}\}$.
    \item For {{\FrameworkAbbr}-STORM}, both $c_{\eta}$ and $c_{\gamma}$ are tuned from $\{\num{1e-3}, \num{1e-2}, \num{1e-1}, \num{1}\}$, $c_{\alpha} \in \{1, 5, 10\}$, and $\rho \in \{1/5, 1/3\}$.
    \item For CoDA, the step size is chosen from $\{\num{1e-3}, \num{1e-2}, \num{1e-1}, \num{1}\}$ and the number of local iterations $K \in \{1, 5, 10, 20\}$.
\end{itemize}

Table~\ref{table_auc_parameters} shows the best hyper-parameters for each algorithm in Figures~\ref{figure_auc},~\ref{figure_auc_S_10}, and~\ref{figure_auc_S_20}.
The hyper-parameters for each algorithm in Figures~\ref{figure_auc_sort_by_label_a} and~\ref{figure_auc_sort_by_label_b} are listed in Table~\ref{table_auc_sort_by_labels_parameters}.
In Figure~\ref{figure_auc_sort_by_label_c}, the chosen hyper-parameters are listed as follows.
\begin{itemize}
    \item For {{\FrameworkAbbr}-MB}, $\eta = \gamma = \num{1e-1}$;
    \item For {{\FrameworkAbbr}-STORM}, $c_{\eta} = c_{\gamma} = 1$, $c_{\alpha} = 10$, and $\rho = 1/5$;
    \item For CoDA, the step size is $\num{1e-2}$ and the number of local steps is $20$.
\end{itemize}


\begin{table}[!htb]
\centering
\caption{Hyper-parameters for the AUC maximization experiment in Figures~\ref{figure_auc_sort_by_label_a} and~\ref{figure_auc_sort_by_label_b}.}
\label{table_auc_sort_by_labels_parameters}
\begin{tabular}{ccccccccccc}
\toprule
\multicolumn{3}{c}{Setting} & \phantom{a} & \multicolumn{2}{c}{{{\FrameworkAbbr}-MB}} & \phantom{a} & \multicolumn{4}{c}{{{\FrameworkAbbr}-STORM}} \\
\cmidrule{1-3} \cmidrule{5-6} \cmidrule{8-11}
S & K & minibatch size && $\eta$ & $\gamma$ && $c_\eta$ & $c_\gamma$ & $c_\alpha$ & $\rho$ \\ \midrule
5 & 1 & 10 && 1 & $\num{1e-2}$ && 1 & $\num{1e-1}$ & 5 & $1/5$ \\
5 & 10 & 10 && $\num{1e-1}$ & $\num{1e-3}$ && $\num{1e-1}$ & $\num{1e-2}$ & 10 & $1/5$ \\
5 & 20 & 10 && $\num{1e-1}$ & $\num{1}$ && $\num{1e-1}$ & $\num{1e-1}$ & 10 & $1/5$ \\
\bottomrule
\end{tabular}
\end{table}



\subsection{The $\ell_2$-relaxed AUC Maximization Problem and its Equivalent Minimax Formulation}  \label{section_auc_minimax_formulation}

Suppose that we are given a dataset $\{\mathbf{w}_i, \ell_i\}_{i=1}^m$ where $\mathbf{w}_i$ denotes a feature vector and $\ell_i \in \{-1, +1\}$ denotes the corresponding label.
For a scoring function $h_{\boldsymbol{\theta}}$ of a classification model parameterized by $\boldsymbol{\theta} \in \mathbb{R}^p$, the AUC is defined as
\begin{equation}  \label{eq_auc_original}
    \max_{\boldsymbol{\theta}} \frac{1}{m^+ m^-} \sum_{\ell_i = +1, \ell_j = -1} \mathbb{I}_{\left\{ h_{\boldsymbol{\theta}}(\mathbf{w}_i) \ge h_{\boldsymbol{\theta}}(\mathbf{w}_j) \right\}},
\end{equation}
where $m^+$ ($m^-$) denotes the number of positive (negative) samples and $\mathbb{I}$ denotes the indicator function.
Since directly maximizing the AUC is NP hard in general, the indicator function in~\eqref{eq_auc_original} is usually replaced by a surrogate function such as the squared loss in practice:
\begin{equation}  \label{eq_auc_relaxed}
    \max_{\boldsymbol{\theta}} \frac{1}{m^+ m^-} \sum_{\ell_i = +1, \ell_j = -1} \left(1 - h_{\boldsymbol{\theta}}(\mathbf{w}_i) + h_{\boldsymbol{\theta}}(\mathbf{w}_j) \right)^2.
\end{equation}
It has been shown that~\eqref{eq_auc_relaxed} has the following equivalent minimax formulation~\citep{ying2016stochastic, liu2020stochastic}
\begin{equation}
    \min_{(\boldsymbol{\theta}, a, b) \in \mathbb{R}^{p+2}} \max_{\lambda \in \mathbb{R}} f(\boldsymbol{\theta}, a, b, \lambda),
\end{equation}
where $f$ is defined as
\begin{equation}
\begin{aligned}[b]
    f(\boldsymbol{\theta}, a, b, \lambda)
    ={}& \frac{1}{m} \sum_{i=1}^m \Big\{ (1 - \tau) (h_{\boldsymbol{\theta}}(\mathbf{w}_i) - a)^2 \mathbb{I}_{\{\ell_i \! = \! 1\}}
        - \tau (1 - \tau) \lambda^2
        + \tau (h_{\boldsymbol{\theta}}(\mathbf{w}_i) - b)^2 \mathbb{I}_{\{\ell_i \! = \! -1\}} \\
        &+ 2 (1 + \lambda) \tau h_{\boldsymbol{\theta}}(\mathbf{w}_i) \mathbb{I}_{\{\ell_i \! = \! -1\}}
        - 2 (1 + \lambda) (1 - \tau) h_{\boldsymbol{\theta}}(\mathbf{w}_i) \mathbb{I}_{\{\ell_i \! = \! 1\}} \Big\}.
\end{aligned}
\end{equation}
Note that $f(\boldsymbol{\theta}, a, b, \cdot)$ is strongly concave for any $(\boldsymbol{\theta}, a, b) \in \mathbb{R}^{p + 2}$ and thus satisfies Assumption~\ref{assumption_pl}.


