\section{Introduction}
Federated Learning denotes the process in which a set of distributed located clients jointly perform a machine learning task under the coordination of a central server over their privately-held data. A widely used method in FL is the FedAvg(Local-SGD)~\cite{mcmahan2017communication} algorithm. As indicated by its name, FedAvg performs (stochastic) gradient descent steps on each client and averages local states periodically. This method can be shown to converge~\cite{stich2018local, haddadpour2019convergence, woodworth2020local} when the distributions of the clients are homogeneous or with bounded heterogeneity. Recently, a large amount of literature has focused on accelerating FedAvg. In particular, many research works use momentum-based methods to accelerate FL, and significant progress has been made in this direction with improved gradient complexity and communication complexity~\cite{das2020faster,karimireddy2019scaffold, khanduri2021stem}. However, another important category of methods: adaptive gradient methods have received much less attention, and there is still no general framework to incorporate adaptive gradient methods into the federated setting.

Adaptive gradient methods such as Adagrad~\cite{duchi2011adaptive}, Adam~\cite{kingma2014adam} and AMSGrad~\cite{reddi2018convergence} are widely used in the non-distributed setting. The gradient descent method uses either a fixed learning rate or a fixed learning rate schedule. In contrast, adaptive gradient methods set the learning rate to be inversely proportional to the magnitude of the gradient; this can incorporate the local curvature structure of the problem. Adaptive gradient methods perform well in practice; meanwhile, they also enjoy useful theoretical implications that make them outperform the vanilla gradient descent method~\cite{duchi2011adaptive,guo2021stochastic}. For example, a recent study~\cite{staib2019escaping} showed that adaptive gradients help escape saddle points. Furthermore, some studies~\cite{loshchilov2018decoupled, chen2018closing} showed that adaptive gradients improve the generalization performance of the model. 

Adaptive gradient methods can be viewed as a type of generalized mirror descent~\cite{huang2021super} methods, where the associated mirror map is defined according to adaptive learning rates. However, the mirror map is dynamic and changes at every training step. As a special case, the gradient descent method can be viewed as a mirror descent method with the mirror map being the $L_2$ distance function. Following the convention in the mirror descent literature, we denote the parameter space as the primal space and the gradient space as the dual space. The primal and dual space differ in adaptive gradient methods, but they coincide in the gradient descent method. We can exploit this primal-dual view to understand existing FL algorithms and design new algorithms. FedAvg actually exploits the usefulness of average dual states. In FedAvg, the gradient average approximates the true gradient evaluated at an average point of the client states, and the approximation error is upper-bounded by the client states difference; therefore, clients can perform multiple local steps without communication. Although in FedAvg, we do not differentiate the primal and dual space as they are the same, but the dual state average and primal state average are not equivalent for adaptive gradient methods.

Current federated adaptive gradient methods in the literature either only perform adaptive gradient steps on the server side, or ignore this primal-dual nuance when supporting local adaptive gradient steps. An early work is~\cite{reddi2020adaptive}, the authors proposed applying adaptive gradients in the server-average step, while performing normal gradient descent updates locally. This method is simple to implement and gets better performance than FedAvg, but the adaptive information is not exploited during local updates; this weakens the usage of adaptive gradients. Recently, some work ~\cite{karimireddy2020mime, chen2020toward} exploited adaptive information during local update steps; however, a common characteristic of these methods is that they average the primal states (parameters) of the problem during the synchronization step. This will cause some problems. Firstly, since adaptive learning rates define the mirror map, updating adaptive learning rates locally makes the dual space not aligned, thus we can not average the primal states directly. Then even if the adaptive learning rates are fixed locally, the primal space might be nonlinear \emph{w.r.t.} the dual space, \emph{e.g.}, when we solve a constrained optimization problem.  In summary, we propose two principles to apply adaptive gradients in FL. First, the local dual spaces should be aligned with each other; Second, we should average dual states.

\begin{table*}
  \centering
  \caption{ Comparisons of representative Federated Learning algorithms for finding
  an $\epsilon$-stationary point of Objective~\eqref{Eq: Problem}
  \emph{i.e.}, $\|\nabla f(x)\|^2 \leq \epsilon$ or its equivalent variants.
  $Gc(f,\epsilon)$ denotes the number of gradient queries \emph{w.r.t.} $f^{(k)}(x)$ for $k\in [K]$; $Cc(f,\epsilon)$ denotes the number of communication rounds; \textbf{State} means what state the algorithm maintains locally (Primal/Dual); \textbf{Local-Adaptive} means whether the algorithm performs adaptive gradient descent locally or not; \textbf{Constrained} means whether the algorithm can solve both constrained and unconstrained problems or not. The first three algorithms are not adaptive gradient methods, and the last four methods support some form of adaptive gradients.
   } \label{tab:1}
  \resizebox{\textwidth}{!}{
 \begin{tabular}{{c|c|c|c|c|c|c}}\toprule
  \textbf{Algorithm} & $Gc(f,\epsilon)$ & $Cc(f,\epsilon)$ & \textbf{State}  & \textbf{Local-Adaptive} & \textbf{Constrained}\\\hline
  FedAvg~\cite{mcmahan2017communication} & $O(\epsilon^{-2})$ & $O(\epsilon^{-1.5})$ & Primal/Dual  &  &  \\ \hline
  FedCM~\cite{xu2021fedcm} & $\tilde{O}(\epsilon^{-1.5})$ & $\tilde{O}(\epsilon^{-1})$ &  Primal/Dual&  &  \\ \hline
  STEM~\cite{khanduri2021stem} & $\tilde{O}(\epsilon^{-1.5})$ & $\tilde{O}(\epsilon^{-1})$ &  Primal/Dual&  &  \\ \midrule
  FedAdam~\cite{reddi2020adaptive} & $O(\epsilon^{-2})$ & $O(\epsilon^{-1})$ & Primal  &  & \\ \hline
  Local-AMSGrad~\cite{chen2020toward} & $O(\epsilon^{-2})$ & $O(\epsilon^{-1.5})$ & Primal & $\surd$  & \\  \hline
  MIME-MVR~\cite{karimireddy2020mime} & $\tilde{O}(\epsilon^{-1.5})$ & $O(\epsilon^{-1.5})$ & Primal & $\surd$ &  \\   \hline
  \textbf{FedDA-MVR(Ours)} & $\tilde{O}(\epsilon^{-1.5})$ & { $\tilde{O}(\epsilon^{-1})$} & Dual & $\surd$ & $\surd$  \\ \bottomrule
 \end{tabular}
 }
\end{table*}

More specifically, we propose the FL adaptive gradients framework \textbf{FedDA}, which is short for \textbf{Fed}erated \textbf{D}ual-averaging \textbf{A}daptive-gradient. In each global round of \textbf{FedDA}, the clients aggregate gradients (dual states) locally, and the server averages the dual states of the clients in the synchronization step. Local weights (primal states) are used as gradient query points in local updates and are recovered through the inverse mirror map (defined by the adaptive gradients). The global primal state is updated on the basis of the averaged dual states and the inverse mirror map. In addition, we utilize a restarting technique to make sure that all clients share the same dual space during local updates; more precisely, we refresh the adaptive gradients at every global epoch and use a fixed one in the local update. Our \textbf{FedDA} framework is general and can incorporate a large family of adaptive gradient methods to the FL setting. In particular, \textbf{FedDA-MVR}, an instantiation of our framework, achieves the best-known gradient complexity and communication complexity in the FL setting. Finally, we highlight our \textbf{contribution} as follows:
\begin{itemize}
\vspace*{-4pt}
\item[(i)] We propose \textbf{FedDA}, a framework for federated adaptive gradient methods. The framework uses a restarted dual averaging technique and adapts a large family of adaptive gradient methods to the FL setting;
\item[(ii)] \textbf{FedDA-MVR}, an instantiation of our framework, obtains the gradient complexity of $\tilde{O}(\epsilon^{-1.5})$ and communication complexity of $\tilde{O}(\epsilon^{-1})$. This matches the optimal rate of non-adaptive federated algorithms and outperforms existing adaptive federated algorithms. \textbf{FedDA-MVR} uses the momentum-based variance-reduction gradient estimation, and exponential moving average of the gradient square as adaptive learning rates;
\item[(iii)] We empirically verify the efficacy of the framework \textbf{FedDA} by performing a colorrectal cancer prediction task and a classification task over the CIFAR10 and FEMNIST datasets.
\vspace*{-4pt}
\end{itemize}

\noindent\textbf{Notations.} $\nabla f(x)$ ($\nabla f^{(k)} (x)$)
denotes the first-order derivatives of the function $f(x)$ ($f^{(k)} (x)$) \emph{w.r.t.} variable $x$. $\xi$ denotes a random sample and $\nabla f (x;\xi)$($\nabla f^{(k)} (x;\xi)$) is the stochastic estimate $\nabla f(x)$ ($\nabla f^{(k)} (x)$). $O(\cdot)$ is the big O notation, and $\tilde{O}(\cdot)$ hides logarithmic terms. $I_d$ denotes a $d$-dimensional identity matrix. $Diag(x)$ denotes the matrix whose diagonal is the vector $x$. $\|\cdot\|$ denotes the $\ell_2$ norm for vectors and the spectral norm for matrices, respectively. $\langle \cdot,\cdot\rangle$ denotes the Euclidean inner product. [K] denotes the set of $\{1,2,...,K\}$. For a random variable $X$, $\mathbb{E}[X]$ denotes its expectation.

\section{Related Works}
\textbf{Optimization Algorithms in Federated Learning.} The term Federated Learning was first coined in~\cite{mcmahan2017communication}, where the task is learned from a set of distributed located clients under the coordination of a server. In the paper~\cite{mcmahan2017communication}, the authors proposed the FedAvg algorithm, in which each client performs multiple steps of gradient descent with its local data and then sends the updated model to the server for averaging. The idea of FedAvg algorithm resembles the Local-SGD algorithm, which is studied in a more general distributed setting for a longer time~\cite{mangasarian1993backpropagation}. The convergence of the local-SGD method has been heavily analyzed in the literature~\cite{stich2018local, karimireddy2019error, dieuleveut2019communication, khaled2020tighter, yu2019parallel, woodworth2020local, woodworth2021minimax,glasgow2022sharp}. Recently, \cite{glasgow2022sharp} proved a convergence rate of the Local-SGD under convex setting that matches the lower bound. On top of the vanilla Local-SGD, various acceleration methods are considered; we list a few representatives here. \cite{karimireddy2020scaffold} adopted the idea of variance reduction technique for non-distributed finite sum problems: a 'control variate' which contains historical full gradient information is used to correct the bias of local gradients. Then in~\cite{karimireddy2020mime}, the authors proposed a general framework (MIME) to translate a centralized optimizer into the FL setting, including adaptive gradient methods. In MIME, the states of an optimizer are fixed during local update steps and only updated at the server-average step. In~\cite{das2020faster, khanduri2021near}, momentum-based variance reduction is applied to the FL setting to control the noise of the stochastic gradients. In \cite{das2020faster}, the authors maintained a server momentum state and a client momentum state, while in \cite{khanduri2021near}, the authors maintained a momentum state and the momentum was averaged periodically similar to the primal state.

Adaptive gradient methods are also studied in the FL setting. The 'Adaptive Federated Optimization'~\cite{reddi2020adaptive} method proposed to use adaptive gradients on the server side while the local gradients are used to update the states of the adaptive gradient methods. In~\cite{chen2020toward}, the authors first showed the divergence of a naive local AMSGrad method that directly averages the primal states periodically. The authors then proposed Local-AMSGrad, a method in which clients update adaptive learning rates locally and average at the synchronization step. At the server average step, both primal states and local adaptive learning rates are averaged to replace the old states. 
Finally, another line of research~\cite{tang2020apmsqueeze, tang20211,lu2022maximizing, chen2020efficient} considers federated adaptive learning rates through the compression approach, these methods communicate local gradients at every step, but the compression techniques are used to reduce the communication cost.

\textbf{Adaptive Gradients in the Non-distributed Learning.} Adaptive gradient methods are widely used in the non-distributed machine learning setting. The first adaptive gradient method \emph{i.e.} Adagrad was proposed in~\cite{duchi2011adaptive}, where the method was shown to outperform SGD in the sparse gradient setting. Since Adagrad does not perform well under dense gradient setting and non-convex setting, some of its variants are proposed, such as SC-Adagra~\cite{mukkamala2017variants} and SAdagrad~\cite{chen2018sadagrad}. Furthermore, Adam~\cite{kingma2014adam} and YOGI~\cite{zaheer2018adaptive} proposed to use the exponential moving average
instead of the arithmetic average used in Adagrad. Adam/YOGI is widely used and very successful in deep learning applications; however, Adam diverges in some settings and the gradient information quickly disappears, so AMSGrad~\cite{reddi2018convergence} is proposed, and it applies an extra `long term memory' variable to preserve the past gradient information to handle the convergence issue of Adam. The convergence of Adam-type methods is also studied in the literature~\cite{chen2019convergence, zhou2018convergence, liu2019variance, guo2021stochastic, huang2021super}. Adaptive gradient methods with good generalization performance are also proposed, such as AdamW~\cite{loshchilov2018decoupled}, Padam~\cite{chen2018closing}, Adabound~\cite{luo2019adaptive}, Adabelief~\cite{zhuang2020adabelief} and AaGrad-Norm~\cite{ward2019adagrad}.


\section{Preliminaries}
In this section, we introduce some preliminaries before introducing our framework. First, we consider the following formulation of Federated Learning:
\begin{align}
\label{Eq: Problem}
     \min_{x \in \mathcal{X} \subset \mathbb{R}^d} ~&\bigg\{  f(x) \coloneqq \frac{1}{K} \sum_{k = 1}^K \big\{ f^{(k)}(x) \coloneqq \mathbb{E}_{\xi^{(k)} \sim \mathcal{D}^{(k)}}[ f^{(k)}(x; \xi^{(k)})] \big\}\bigg\}.
\end{align}
which considers $K$ clients. For the $k_{th}$ client, we optimize the loss objective $f^{(k)} (x):\mathcal{X} \to \mathbb{R} $ which is smooth and possibly non-convex, and $x$ denotes the variable of interest. $\mathcal{X}\subset \mathbb{R}^{d}$ is a compact and convex set. $\xi^{(k)} \sim \mathcal{D}^{(k)}$ is a random example that follows an unknown data distribution $\mathcal{D}^{(k)}$. The formulation in~\eqref{Eq: Problem} includes both the homogeneous case \emph{i.e.} $f^{(k)} (x) = f^{(j)} (x)$ for any $k,j \in [K]$, and the heterogeneous case \emph{i.e.} $f^{(k)} (x) \neq f^{(j)} (x)$ for some $k,j \in [K]$.

Next, we introduce some basics of adaptive gradient methods from a mirror-descent perspective. Generally, mirror descent is associated with a mirror map $\Phi(x)$. Given the objective $f(x)$ and the primal state $x_t \in \mathcal{X}$ at $t_{th}$ step, we first map the primal state to the mirror space as $y_t = \nabla \Phi(x_t)$, then we perform the gradient descent step in the mirror space: $y_{t+1} = y_t - \eta\nabla f(x)$, where $\eta$ is the learning rate, finally, we map $y_{t+1}$ back to the primal space as $x_{t+1} = \underset{x \in \mathcal{X}}{\arg\min} D_{\Phi}(x, y_{t+1})$, where $D_{\Phi}(x,y)$ denotes the Bregman Divergence associated to $\Phi$, \emph{i.e.} $D_{\Phi}(x,y) = f(x) - f(y) - \langle \nabla f(y), x - y\rangle$, In summary, the mirror descent step can be written as a Bregman proximal gradient step as follows: \[x_{t+1} = \underset{x \in \mathcal{X}}{\arg\min}\ \eta \langle \nabla f(x_t), x \rangle + D_{\Phi}(x, x_{t})\]
For the adaptive gradient methods, we uses the following mirror map: $\Phi(x) = \frac{1}{2}x^THx$, where $H$ is the adaptive matrix and is positive definite. Many adaptive gradient methods can be written in the following proximal gradient descent form:
\begin{align}\label{eq:adapt_update}
x_{t+1} = \underset{x \in \mathcal{X}}{\arg\min}\ \eta \langle \nu_t, x \rangle +  \frac{1}{2}(x - x_t)^TH_t(x - x_t),
\end{align}
we replace the gradient $\nabla f(x)$ with the generalized gradient estimation $\nu_t$, besides, we replace $H$ with $H_t$ based on the fact that the adaptive matrix is updated at every step. Next, we show some examples of adaptive gradients methods that can be phrased as the above formulation. For the Adagrad~\cite{duchi2011adaptive} method, we set 
\begin{align}\label{eq:ada1}
\nu_t = \nabla f(x_t, \xi_t),\; H_t = Diag(\sqrt{\mu_t}),\; \mu_t =\frac{1}{t}\sum_{i=1}^{t}\nu_i^2
\end{align}
For Adam~\cite{kingma2014adam}, we have:
\begin{align}\label{eq:ada2}
&\hat{\nu}_t = (1 - \beta_1)\nabla f(x_t, \xi_t) + \beta_1 \hat{\nu}_{t-1},\nonumber\\
&\hat{\mu}_t =(1 - \beta_2) \nabla f(x_t, \xi_t)^2 + \beta_2 \hat{\mu}_{t-1}\nonumber \\
&\nu_t = \hat{\nu}_t/(1 - \gamma_1^t),\; \mu_t = \hat{\mu}_t/(1 - \gamma_2^t),\nonumber\\
&H_t = Diag(\sqrt{\mu_t} + \epsilon)
\end{align}
where $\beta_1, \beta_2, \gamma_1, \gamma_2$ are some constants. For other adaptive gradient methods, please refer to~\cite{huang2021super}.



\begin{algorithm}[t]
\caption{\textbf{FedDA}-Server}
\label{alg:fedada-server}
\begin{algorithmic}[1]
\STATE {\bfseries Input:} Number of global epochs $E$, tuning parameters $\{\beta_{\tau}\}_{i=1}^{E}$;
\STATE {\bfseries Initialize:} Choose $x_0 \in \mathcal{X}$ and compute $\nu_0 = \frac{1}{K} \sum_{j=1}^{K} \nabla f^{(j)}(x_0, \mathcal{\mathcal{B}}^{(k)}_0)$ where $\{\mathcal{\mathcal{B}}^{(k)}_0\}_{k=1}^K$ are a mini-batch of random points selected from each of $K$ clients;
\FOR{$\tau=0$ \textbf{to} $E - 1$}
\STATE Server selects a set $\mathcal{S}_{\tau}$  of $r$ clients chosen uniformly at random w/o replacement;
\FOR{the client $k \in \mathcal{S}_\tau$ in parallel}
\STATE  ($z^{(k)}_{\tau+1, I}$, $\nu^{(k)}_{\tau + 1, I}$) = \textbf{FedDA}-client($x_\tau$, $\nu_{\tau}$, $H_\tau$)
\ENDFOR
\STATE Compute $z_{\tau + 1} = \frac{1}{r} \sum_{k \in \mathcal{S}_{\tau}}  z^{(k)}_{\tau+1, I}$;
\STATE Compute $x_{\tau+1} = \underset{x \in \mathcal{X}}{\arg\min} \{ -\langle x, z_{\tau + 1}\rangle + \frac{1}{2\lambda} (x - x_{\tau})^{T}H_{\tau}(x - x_{\tau})$ \};
\STATE Compute $\nu_{\tau + 1} = \frac{1}{r} \sum_{k \in \mathcal{S}_{\tau}}  \nu^{(k)}_{\tau+1, I}$;
\STATE Compute $H_{\tau+1} = \mathcal{V}(H_{\tau}, z_{\tau+1})$;
\ENDFOR
\end{algorithmic}
\end{algorithm}


\begin{algorithm}[t]
\caption{\textbf{FedDA}-Client ($x_\tau, \nu_\tau$, $H_\tau$)}
\label{alg:fedada-client}
\begin{algorithmic}[1]
\STATE {\bfseries Input:} Number of local steps $I$, tuning parameters $\{\eta_{\tau + 1,i}\}_{i=0}^{I-1}$, $\{\alpha_{\tau + 1,i}\}_{i=1}^{I}$;
\STATE {\bfseries Initialize:} $x^{(k)}_{\tau + 1,0} = x_\tau$; $\nu^{(k)}_{\tau + 1,0} = \nu_\tau$; $z^{(k)}_{\tau+1,0} = 0$;
\FOR{$i=0$ \textbf{to} $I-1$}
\STATE Compute $z^{(k)}_{\tau + 1,i+1} = z^{(k)}_{\tau + 1, i} - \eta_{\tau + 1, i} \nu_{\tau + 1, i}^{(k)}$;
\STATE Compute $x^{(k)}_{\tau + 1, i+1} = \underset{x \in \mathcal{X}}{\arg\min} \{ -\langle x, z^{(k)}_{\tau + 1,i+1}\rangle + \frac{1}{2\lambda} (x - x^{(k)}_{\tau + 1, 0})^{T}H_\tau(x - x^{(k)}_{\tau + 1, 0})$ \};
\STATE Compute $\nu^{(k)}_{\tau + 1, i+1} = \mathcal{U}(\nu^{(k)}_{\tau + 1, i}, x^{(k)}_{\tau + 1, i+1}, x^{(k)}_{\tau + 1, i}; \alpha_{\tau+1, i+1}, \mathcal{\mathcal{B}}_{\tau + 1, i+1}^{(k)})$, where $\mathcal{\mathcal{B}}_{\tau + 1, i+1}^{(k)}$ is a minibatch of random samples from the client $k$;
\ENDFOR
\STATE {\bfseries Output:} Send $z^{(k)}_{\tau + 1, I}$, $\nu^{(k)}_{\tau + 1, I}$ to the server.
\end{algorithmic}
\end{algorithm}

\section{Local Adaptive Gradients via Dual Averaging}
In this section, we introduce \textbf{FedDA}, a framework of federated adaptive gradient methods.  The procedure of \textbf{FedDA} is summarized in Algorithm~\ref{alg:fedada-server}. 

In Algorithm~\ref{alg:fedada-server}, we perform $E$ global steps and at each global step, we select a subset of clients for training. All selected clients at each step will run Algorithm~\ref{alg:fedada-client}. In Algorithm~\ref{alg:fedada-client}, clients receive the current model weight $x_\tau$, gradient estimation $\nu_\tau$ and adaptive gradient matrix $H_\tau$. The clients then perform $I$ local training steps: line 3- line 7 in Algorithm~\ref{alg:fedada-client}. For each step, we first accumulate the dual state in the variable $z^{(k)}_{\tau, i}$ (line 4), then we calculate the local primal state $x^{(k)}_{\tau, i} $ (line 5), which is a proximal gradient step similar to~\eqref{eq:adapt_update}. The function of this step is to map the aggregated dual state $z^{(k)}_{\tau, i}$ back to the primal space, and we use the primal state to query the gradient to update the estimation of the gradient $\nu^{(k)}_{\tau, i}$ (line 6). Note, we use a fixed adaptive matrix $H_{\tau}$ during local steps, this makes the clients share the same dual space. In line 6 of Algorithm~\ref{alg:fedada-client}, we update the gradient estimation $\nu^{(k)}_{\tau, i}$. The update rule $\mathcal{U}(\cdot)$ is general, \emph{e.g.},the momentum-based variance reduction update~\eqref{eq:nu-case1} and the momentum update~\eqref{eq:nu-case2} as follows ($\alpha_{\tau,i}$ is some constant):
\begin{align}
\label{eq:nu-case1}
    &\nu_{\tau+1, i+1}^{(k)} = \nabla f^{(k)} (x^{(k)}_{\tau+1, i+1}, \mathcal{\mathcal{B}}_{\tau+1, i+1}^{(k)}) + (1 - \alpha_{\tau+1, i+1})(\nu_{\tau+1, i}^{(k)} - \nabla f^{(k)} (x^{(k)}_{\tau+1, i}, \mathcal{\mathcal{B}}_{\tau+1, i+1}^{(k)}))
\end{align}
and
\begin{align}
\label{eq:nu-case2}
   &\nu_{\tau+1, i+1}^{(k)} = \alpha_{\tau+1, i+1}\nabla f^{(k)} (x^{(k)}_{\tau+1, i+1}, \mathcal{\mathcal{B}}_{\tau+1, i+1}^{(k)}) + (1 - \alpha_{\tau+1, i+1})\nu_{\tau+1, i}^{(k)}
\end{align}

After the client runs Algorithm~\ref{alg:fedada-client}, it returns the aggregated local dual states $z^{(k)}_{\tau + 1, I}$ and the local gradient estimation $\nu^{(k)}_{\tau + 1, I}$ to the server. The server first averages the local dual states (line 8 of Algorithm~\ref{alg:fedada-server}) to get $z_{\tau+1}$. We can average local dual states as all clients have a common dual space. The server then calculates the new primal states $x_{\tau+1}$ as in line 9 of Algorithm~\ref{alg:fedada-server}. Next, the gradient estimation $\nu_\tau$ is also updated by averaging local states (line 10 of Algorithm~\ref{alg:fedada-server}). Finally, we update the adaptive matrix $H_\tau$ (line 11 of Algorithm~\ref{alg:fedada-server}). The update rule $\mathcal{V}$ is general, \emph{e.g.}, 
\begin{align}
\label{eq:mu-case1}
    \mu_{\tau+1} &= \beta_{\tau+1}z_{\tau+1}^2/\eta_{\tau+1,I-1}^2 + (1 - \beta_{\tau+1})\mu_{\tau+1},\nonumber\\
    H_{\tau+1} &= Diag(\sqrt{\mu_{\tau+1}} + \epsilon)
\end{align}
and
\begin{align}
\label{eq:mu-case2}
    \mu_{\tau+1} &= \beta_{\tau+1}||z_{\tau+1}||/\eta_{\tau+1, I-1} + (1 - \beta_{\tau+1})\mu_{\tau+1},\nonumber\\
    H_{\tau+1} &= (\mu_{\tau+1} + \epsilon)I_d
\end{align}
where we set $\mu_0 = 0$, $\epsilon$ is some constant. In summary, Algorithm~\ref{alg:fedada-server} aggregates and averages dual states at each global round. The adaptive matrix $H_\tau$ is fixed during local updates and is refreshed on the server side at each global round. Since the algorithm uses a new mirror map (adaptive gradient matrix) at each global round, we call our framework to be restarted dual averaging.


\begin{remark}
In contrast to our dual-averaging strategy, some existing adaptive FL algorithms~\cite{praneeth2020mime} average the local primal states. In the unconstrained case, the primal and dual spaces are linear with each other, but in the constrained case, the linearity does not exist, and the averaging in the primal space and dual space is not equivalent. As we show in the subsequent theoretical analysis, dual averaging leads to the convergence in the constrained case.
\end{remark}

\begin{remark}
Note that we use the averaged dual states $z_{\tau+1}$ when we update the adaptive matrix (line 11 of Algorithm~\ref{alg:fedada-server}). An alternative choice is to use the most recent gradient~\cite{praneeth2020mime}. In comparison, the dual state aggregates information of whole round and offers smoother estimation of the problem's local curvature. Another possible choice, as used in the Local-AMSGrad method~\cite{chen2020toward}, is to update the state of the adaptive matrix $\mu_\tau$ (see~\eqref{eq:mu-case1}) locally and then average in the server synchronization step. The limitation of this approach is that $\mu_\tau$ is not linear \emph{w.r.t} gradient, and thus averaging $\mu_\tau$ does not offer a linear speed-up \emph{w.r.t.} the number of clients; in contrast, the dual state satisfies linearity.
\end{remark}

\begin{remark}
By choosing different update rules $\mathcal{U}$ and $\mathcal{V}$, we can create many variants of FedDA. An representative is \textbf{FedDA-MVR}, in which we update $\nu_{\tau,i}^{(k)}$ with momentum-based variance reduction (~\eqref{eq:nu-case1}) and the adaptive matrix $H_{\tau}$ with an exponential average of the square of the gradient (~\eqref{eq:mu-case1}). In the subsequent discussion, we focus on this variant and perform both theoretical and empirical analysis.
\end{remark}

\section{Theoretical Analysis}
In this section, we provide the theoretical analysis of our \textbf{FedDA} framework; more specifically, we focus on the analysis of \textbf{FedDA-MVR}. FedDA-MVR uses ~\eqref{eq:mu-case1} to update the adaptive matrix $H_\tau$ and~\eqref{eq:nu-case1} to update the gradient estimation $\nu^{(k)}_{\tau,i}$. We first state the assumptions we need in our analysis:
\subsection{Some Mild Assumptions}
\begin{assumption}[Bounded Client Heterogeneity]
\label{ass: inter_variance} The difference of gradients between different workers are bounded: \[\| \nabla f^{(k)} (x) - \nabla f^{(\ell)} (x) \|^2 \leq \zeta^2,\; \forall k,\ell \in [K].\]
\end{assumption}
We measure the heterogeneity of the clients in terms of gradient dissimilarity. The above assumption or its similar form is also exploited in the analysis of other Federated Learning Algorithms, such as in~\cite{khanduri2021stem, das2020faster}.
\begin{assumption} \label{ass:value}
The function $f(x)$ is bounded from below in $\mathcal{X}$, \emph{i.e.,} $f^* = \inf_{x\in \mathcal{X}} f(x)$.
\end{assumption}
\begin{assumption}[Unbiased and Bounded-variance Stochastic Gradient]
\label{ass: unbiased_bouned} The stochastic gradients are unbiased with bounded variance, \emph{i.e.}
$$\mathbb{E}[\nabla f^{(k)} (x ; \xi^{(k)})] = \nabla f^{(k)} (x)$$ and there exists a constant $\sigma$ such that 
\begin{align*}
&\mathbb{E}\| \nabla f^{(k)} (x ; \xi^{(k)}) - \nabla f^{(k)} (x) \|^2 \leq \sigma^2, \nonumber\\
&\qquad \forall~\xi^{(k)} \sim \mathcal{D}^{(k)},\; \forall~k\in[K]
\end{align*}
\end{assumption}
Assumption~\ref{ass:value} guarantees the feasibility of the Federated Learning problem~\eqref{Eq: Problem}, and Assumption~\ref{ass: unbiased_bouned} is widely used in stochastic optimization analysis.
\begin{assumption} \label{ass:postive_definite}
The adaptive matrix $H_\tau$ is symmetric positive definite, \emph{i.e.} there exists a constant $\rho > 0$ such that \[ H_\tau \succeq \rho I_d \succ 0,\; \forall t\geq 1,\] 
\end{assumption}
In our analysis, we assume the adaptive matrix is positive definite, and this requirement can be easily satisfied by many adaptive gradient methods. Firstly, most adaptive gradient methods always have non-negative adaptive learning rates, such as~\eqref{eq:ada1} and~\eqref{eq:ada2}. To make it positive, we can add a bias term $\epsilon$ such as in the Adam update rule~\eqref{eq:ada2}.

\begin{assumption}[{Sample Gradient Lipschitz Smoothness}]
\label{Ass: Lip_Smoothness}
The stochastic functions $f^{(k)}(x, \xi^{(k)})$ with $\xi^{(k)} \sim \mathcal{D}^{(k)}$ for all $k \in [K]$, satisfy the mean squared smoothness property, i.e, we have 
\begin{align*}
&\mathbb{E} \| \nabla f^{(k)} (x ; \xi^{(k)}) -  \nabla f^{(k)} (y ; \xi^{(k)}) \|^2 \leq L^2   \| x - y \|^2 \nonumber\\
&\qquad \text{for all}~x,y \in \mathbb{R}^d
\end{align*}
\end{assumption}
The smoothness assumption above is a slightly stronger requirement than the standard smooth condition, but this assumption is widely used in the analysis of variance reduction methods, such as SPIDER~\cite{fang2018spider} and STORM~\cite{cutkosky2019momentum}.

\begin{assumption}
\label{Ass: Full_part}
All clients participate in the training at each step, \emph{i.e.} choose $r = K$ in Algorithm~\ref{alg:fedada-server}.
\end{assumption}
We make the full participation assumption to simplify the exposition of the theoretical results. All the results presented can be easily generalized to the partial participation case.

\subsection{Convergence Property of \textbf{Fed-MVR}}
In this subsection, we provide the convergence property of our \textbf{FedDA-MVR} variant. For convenience of discussion, we redefine the subscript $t = \tau I + i$, \emph{i.e.} we denote the $t$ step as the $i$ local step in the $\tau$ global round. Similarly, we denote the total number of running steps as $T = EI$. We analyze our algorithm through the following measure:
\begin{align}\label{eq:measure}
    \mathcal{G}_t = \frac{\rho^2}{\lambda^2\eta_t^2}||\tilde{x}_t - \tilde{x}_{t+1}||^2 + ||\bar{\nu}_t - \nabla f(\tilde{x}_t)||^2
\end{align}
where $\bar{\nu}_t$ denotes the average gradient estimation at the $t$ step and $\tilde{x}_t$ denotes the virtual global primal state at the $t$ step (see Section~\ref{sec:pre-local} in the appendix for formal definitions). In Remark~\ref{rem:1} of the appendix, we discuss the intuition of the measure $\mathcal{G}_t$. In particular, in the unconstrained case \emph{i.e.} when $\mathcal{X} = R^d$, the measure upper-bounds the square norm of the gradient. Therefore, the convergence of our measure $\mathcal{G}_t$ means the convergence to a first-order stationary point. Now, we are ready to provide the main result of our convergence theorem.
\begin{theorem} 
\label{Thm: PR_Convergence_Main} In Algorithm~\ref{alg:fedada-server}, we choose the parameters as $\displaystyle \kappa = \frac{\rho K^{2/3}}{\lambda L}$, $\displaystyle c = \frac{96\lambda^2L^2}{K\rho^2} + \frac{ \rho }{72 \kappa^3\lambda LI^2}$, $w_t=\max \bigg\{48^3 I^6 K^{2} - t - I,  14^3K^{0.5} \bigg\}$, $\lambda > 0$, and choose $\eta_{t} = \frac{\kappa}{(\omega_{t} + t + I)^{1/3}}$, then we have:
\begin{align*}
    \frac{1}{T} \sum_{t = 0}^{T-1} \mathbb{E}[\mathcal{G}_t] &\leq   \bigg[\frac{96LI^2}{T} + \frac{2L}{K^{2/3} T^{2/3}}\bigg] (f(x_0) -   f^\ast )  + \bigg[\frac{72 I^4}{bT}   + \frac{3I^2}{2b K^{2/3}T^{2/3}} \bigg] \sigma^2 \\
 &  \quad  + 192^2\times\bigg( \frac{48I^2}{T} + \frac{1}{K^{2/3} T^{2/3}} \bigg)\bigg(\frac{\sigma^2}{4b_1} + \frac{2\zeta^2}{21}\bigg) \log(T+1).
\end{align*}
\end{theorem}
Note, by choosing a proper value of local updates $I$ and using a minibatch of samples for the first iteration to decrease the noise, our result matches the best known convergence rate for stochastic federated gradient methods~\cite{khanduri2021stem}, \emph{i.e.} our algorithms has gradient complexity of $\tilde{O}(\epsilon^{-1.5})$ and communication complexity of $\tilde{O}(\epsilon^{-1})$, moreover we achieve linear speed up \emph{w.r.t} the number of clients $K$. More formally, we have the following corollary:
\begin{corollary}
Suppose in Algorithm~\ref{alg:fedada-server}, we set $I = O((T/K^2)^{1/6})$, and use sample minibatch of size O($I^2$) in the initialization, then we have:
\begin{align*}
   \frac{1}{T} \sum_{t = 1}^T \bigg( \mathbb{E} [\mathcal{G}_t]\bigg) = \tilde{O}(\frac{1}{K^{2/3}T^{2/3}}) 
\end{align*}
and to reach an $\epsilon$-stationary point, we need to make $\tilde{O}(\epsilon^{-1.5}/K)$ number of steps and need $\tilde{O}(\epsilon^{-1})$ number of communication rounds.
\end{corollary}


\begin{figure}[t]
	\begin{center}
		\includegraphics[width=0.3\linewidth]{figs/bio-train-acc.png}
		\includegraphics[width=0.3\linewidth]{figs/bio-test-acc.png}
		\includegraphics[width=0.3\linewidth]{figs/bio-dens.png}
		\caption{Results for the PATHMNIST~\cite{medmnistv2} dataset. Plots show the Train Accuracy, Test Accuracy, Density vs Number of Rounds ($E$ in Algorithm~\ref{alg:fedada-server}) respectively. The post-fix of $L_1$ means we consider the $L_1$ constraints. $I$ is chosen as 5.}
		\label{fig:bio}
	\end{center}
\end{figure}

\section{Numerical Experiments}
In this section, we perform numerical experiments to verify the efficacy of the proposed adaptive federated learning framework \emph{i.e.} \textbf{FedDA}. More specifically, we consider the variant of \textbf{FedDA-MVR} here, and defer experiments for other variants to Section~\ref{sec:more_exp} of the appendix. We performed two sets of experiments. In the first experiment, we consider a biomedical prediction task: predicting the survival of colorectal cancer. In this task, we impose a $L_1$ sparsity constraint. $L_1$ constraint improves the explainability of the model, which is essential for biomedical applications. Then in the second experiment, we consider a federated multiclass image classification task. More specifically, we consider two datasets: CIFAR10~\cite{krizhevsky2009learning} and FEMNIST~\cite{caldas2018leaf}.  All experiments are run on a machine with an Intel Xeon Gold 6248 CPU and 4 Nvidia Tesla V100 GPUs. The code is written in Pytorch. We simulate the Federated Learning environment through the Pytorch.distributed package.

\subsection{Colorrectal Cancer Survival Prediction with Sparse Constraints}
In this subsection, we consider a colorrectal cancer prediction task on the PATHMNST dataset~\cite{medmnistv2,kather2019predicting}, which contains 9 different classes. It has 89996 training images, and we equally randomly split the training set into 10 clients. We used the original test set for the metric. In this task, we impose the $L_1$ sparsity constraint to improve the explainability of the model.

In this task, we compare with the following baselines: FedAvg~\cite{mcmahan2017communication} and FedDualAvg~\cite{yuan2021federated}. FedDualAvg is a recently proposed federated algorithm that deals with composite optimization problems. In FedDualAvg, clients maintain dual states locally, but adaptive gradients are not applied. For our FedDA-MVR, we train with and without the $L_1$ constraint. In this task, we use a 4-layer convolutional neural network with 32 filters at each layer. We have 10 clients and run 20000 steps ($T$), average states with interval 5 ($I$) and use mini-batch size of 16. Besides, we calculate density with threshold 0.01. For other hyper-parameters, we perform grid search and choose the best setting for each method. More specifically, for the SGD method, we use learning rate 0.01; for the FedDualAvg algorithm, we use local learning rate 0.1, global learning rate 0.1, $L_1$ constraint 0.01; for our FedDA-MVR, we use learning rate 0.01, $w$ as 100000, $c$ as 5000000, $\beta$ as 0.999 and $\tau$ as 0.01, for the $L_1$ regularized version FedDA-MVR-$L_1$, we also add $L_1$ constraint 0.01. For other variants of FedDA: for FedDA-2-1, we use learning rate 0.001, $\alpha$ as 0.9, $\beta$ as 0.999, $\tau$ as 0.01; for FedDA-1-2, we use learning rate 1, $w$ as 10000, $c$ as 200, $\beta$ as 0.999, $\tau$ as 0.001, $L_1$ constraint 0.01; for FedDA-2-2, we use learning rate 0.01, $\alpha$ 0.9, $\beta$ as 0.999, $\tau$ as 0.01, $L_1$ constraint 0.01. 

The results are summarized in Figure~\ref{fig:bio}, the plots are averaged over 5 independent runs and then smoothed. In Figure~\ref{fig:bio}, FedDualAvg and FedDA-MVR-$L_1$ consider the $L_1$ constraint, while FedAvg and FedDA-MVR do not. We show results of Train/Test Accuracy and also the number of non-zero (below a threshold) elements in the parameter (\emph{i.e.} the rightmost plot in Figure~\ref{fig:bio}). As shown in the plots, FedDA-MVR-$L_1$ outperforms unconstrained FedDA-MVR in all metrics. This shows the importance of considering constrained problems in Federated Learning. Furthermore, FedDA-MVR-$L_1$ also outperforms FedAvg and FedDualAvg in all metrics. This shows that our algorithm can effectively exploit adaptive gradient information in the constrained case. For more details of this experiment, such as the hyper-parameter choices, please refer to Section~\ref{sec:more_exp} of the appendix.



\subsection{Image Classification Task with CIFAR10 and FEMNIST}
In this subsection, we consider an unconstrained image classification task for both homogeneous and heterogeneous cases. More specifically, we consider two datasets: CIFAR10~\cite{krizhevsky2009learning} and FEMNIST~\cite{caldas2018leaf}. CIFAR10 is a widely used image classification benchmark dataset which contains 50000 training images, and we construct both homogeneous and heterogeneous cases based on it. For the homogeneous case, we uniformly randomly distribute them into 10 clients. For the heterogeneous case, we create heterogeneity in the training set as follows: Suppose we have 10 clients, for $i_{th}$ client, we distribute $\rho$-percent samples of $i_{th}$ class, and $(1 - \rho)/9$-percent samples of other classes, where $0 < \rho \leq 1$. Note for $\rho$ close to 1, the $i_{th}$ client will be dominated by images of $i_{th}$ class, thus the data distribution among clients will be very different. In our experiments, we choose $\rho = 0.8$. This means the $i_{th}$ client has 4000 images of $i_{th}$ class and 111 images of other classes. This creates a high level of heterogeneity. Note that we use the original test set of CIFAR10. FEMNIST is a Federated dataset of hand-written digits; it contains hand-written digits of 3550 users (we randomly sample 500 users in our experiments). Data distribution of FEMNIST is heterogeneous for different writing styles of people. Note we randomly sample 50 users at each global round.  We run 20000 steps ($T$), average states with interval 5 ($I$) and use mini-batch size of 16. We use a 4-layer convolutional neural network with 64 filters at each layer.
For other hyper-parameters, we perform grid search and choose the best setting for each method.

\begin{figure}[t]
	\begin{center}
		\includegraphics[width=0.24\columnwidth]{figs/cifar-train-loss.png}
		\includegraphics[width=0.24\columnwidth]{figs/cifar-train-acc.png}
		\includegraphics[width=0.24\columnwidth]{figs/cifar-test-loss.png}
		\includegraphics[width=0.24\columnwidth]{figs/cifar-test-acc.png}
		\caption{Results for CIFAR10 dataset. From left to right, we show Train Loss, Train Accuracy, Test Loss, Test Accuracy \emph{w.r.t} the number of rounds (E in Algorithm~\ref{alg:fedada-server}), respectively. $I$ is chosen as 5.}
		\label{fig:cifar}
	\end{center}
\end{figure}


\begin{figure}[t]
	\begin{center}
		\includegraphics[width=0.24\columnwidth]{figs/femnist-train-loss.png}
		\includegraphics[width=0.24\columnwidth]{figs/femnist-train-acc.png}
		\includegraphics[width=0.24\columnwidth]{figs/femnist-test-loss.png}
		\includegraphics[width=0.24\columnwidth]{figs/femnist-test-acc.png}
		\caption{Results for FEMNIST dataset. From left to right, we show Train Loss, Train Accuracy, Test Loss, Test Accuracy \emph{w.r.t} the number of rounds (E in Algorithm~\ref{alg:fedada-server}), respectively. $I$ is chosen as 5.}
		\label{fig:femnist}
	\end{center}
\end{figure}

In this task, we compare our method with the following baselines: the non-adaptive methods: FedAvg~\cite{mcmahan2017communication}, FedCM~\cite{xu2021fedcm}, STEM~\cite{khanduri2021stem} and adaptive methods: FedAdam~\cite{reddi2020adaptive}, Local-Adapt~\cite{wang2021local}, Local-AMSGrad~\cite{chen2020toward}, MIME-MVR~\cite{praneeth2020mime}. For all methods, we tune their hyper-parameters to find the best setting. The results are summarized in Figure~\ref{fig:cifar} and Figure~\ref{fig:cifar-hete} (CIFAR10) and Figure~\ref{fig:femnist} (FEMNIST), the plots are averaged over 5 runs and then smoothed. Note in Figure~\ref{fig:cifar-hete}, \textbf{FedDA-$i$-$j$} represents different variants of our framework, please refer to Appendix~\ref{sec:more_exp} for their definition. In particular, \textbf{FedDA-$1$-$1$} represents \textbf{FedDA-MVR}.  As shown in the figures, our FedDA-MVR outperforms all baselines. In addition, the FedAvg algorithm has competitive training performance; however, it tends to overfit the training data severely and suffers most from the heterogeneity. Then we observe that adaptive methods in general get better train and test performance. Finally, the superior performance of our method compared with the three adaptive baselines shows that our method exploits adaptive information better; for example, MIME-MVR also exploits the momentum-based variance reduction technique, but it fixes all optimizer states during local updates, in contrast, we only fix the adaptive matrix but update the momentum $\nu_t^{(k)}, k \in [K]$ at every step. For more details, including the hyper-parameter selection, please refer to Section~\ref{sec:more_exp} of the appendix.

\section{Conclusion}
In this paper, we proposed the FedDA framework to incorporate adaptive gradients into the Federated Learning environment. More specifically, we adopted the Mirror Descent view of adaptive gradients, furthermore, we proposed to maintain and average the dual states in the training, meanwhile we fixed the adaptive matrix during local training such that the dual space is shared by all clients. We also analyze the convergence property of our Framework: for the variant FedDA-MVR, we proved that it reaches an $\epsilon$-optimal stationary point with $\tilde{O}(\epsilon^{-1.5})$ gradient queries and $\tilde{O}(\epsilon^{-1})$ communication rounds, these results match the best known gradient complexity and communication complexity of stochastic federated algorithms under the non-convex case. Finally, we validate our algorithm for both constrained and unconstrained tasks. The numerical results show the superior performance of our algorithm compared to various baseline methods.

