Skipping Algorithms in Federated Learning

Skipping Algorithms in Federated Learning#

Even the earliest federated learning algorithm, FedAvg [McMahan et al.[1]], performs multiple gradient-descent steps in the local inner loop of its clients – which, in a sense, is already a skipping update: the communication round is skipped with probability \(1 - p = 0\). The ProxSkip algorithm [Mishchenko et al.[2], Zhang and Loizou[3]] (named Scaffnew in the original paper) studies the skipping technique systematically and makes the communication round skipped with a nonzero probability. It is instructive to read ProxSkip as SCAFFOLD combined with the skipping technique.

ProxSkip/Scaffnew#

From the optimization-theoretic viewpoint, model the federated learning problem as the consensus optimization problem over the consensus set \(\mathcal{E} = \{ \Theta \mid \theta_1 = \cdots = \theta_K \}\),

(1)#\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \newcommand{\col}{\operatorname{col}} \newcommand{\R}{\mathbb{R}} \newcommand{\prox}{\mathbf{prox}}\\\begin{split}\begin{array}{cl} \minimize & F(\Theta) = \frac{1}{K} \sum\limits_{k=1}^K f_k(\theta_k), \\ \text{subject to} & \Theta \in \mathcal{E}, \\ \text{where} & \Theta = \col(\theta_1, \cdots, \theta_K) \in \R^{Kd}. \end{array}\end{split}\end{aligned}\end{align} \]

As in the operator splitting section, one can convert this constrained problem into an unconstrained one with the indicator function \(\iota_{\mathcal{E}}\) of the consensus set and split the first-order optimality condition \(0 \in \nabla F(\Theta) + \partial \iota_{\mathcal{E}}(\Theta)\). More generally, replacing \(\iota_{\mathcal{E}}\) by any CCP function \(g\), the fixed-point solver is the proximal gradient descent (ProxGD)

(2)#\[\Theta^{(t+1)} = \prox_{s g}\left( \Theta^{(t)} - s \nabla F(\Theta^{(t)}) \right),\]

for the problem \(\minimize \ F(\Theta) + g(\Theta)\). The practical difficulty of ProxGD is that the proximal operator \(\prox_{s g}\) of the CCP function \(g\) is typically expensive to evaluate [Friedlander and Goh[4]]. In federated learning, this proximal step is the communication: the model aggregation from the clients to the server followed by the broadcast back. Skipping this performance-bottleneck step with some probability is therefore a natural remedy, and this is exactly what ProxSkip does.

A remark on the control variates: unlike SCAFFOLD, ProxSkip does not use (explicit) gradient variance reduction. The primary role of the control variates \(c_k\) is to guarantee the correct convergence of the overall iteration: one can show \(c_k \to \nabla f_k(\theta^*)\) at a solution \(\theta^*\) of the original problem, so the fixed point of the iteration is a solution of the problem – this is the key property. ProxSkip maintains no global control variate communicated between the server and the clients; instead, whenever a communication happens, the local control variate \(c_k\) is implicitly updated by \(\frac{p}{\gamma} \left( \theta^{(t+1)} - \theta_k^{(t + \frac{1}{2})} \right)\), which acts as a global control variate. This is reminiscent of the Option II update rule of the local control variates in SCAFFOLD: it saves the extra parameter exchange between the server and the clients, at the cost of the flexibility of the control variate updates. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{ProxSkip}}
\begin{algorithmic}
\REQUIRE learning rate $\gamma,$ probability $p \in [0, 1)$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Server init:} global model parameters $\theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Clients init:} local model parameters $\theta_k^{(0)} \in \mathbb{R}^d,$ control variates $c_k^{(0)} \in \mathbb{R}^d, ~ \forall k \in [K]$

\FOR{each round $t = 0, 1, \cdots, T-1$}
    \FOR{each client $k \in [K]$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t, 0)} \gets \theta_k^{(t)}$
        \FOR{local step $r = 0, 1, \cdots, R-1$}
            \COMMENT{$R$ iterates of SGD with variance reduction}
            \STATE compute mini-batch gradient $g_k^{(t, r)}$ of $\nabla f_k(\theta_k^{(t, r)})$
            \STATE $\theta_k^{(t, r+1)} \gets \theta_k^{(t, r)} - \gamma \left( g_k^{(t, r)} - c_k^{(t)} \right)$
        \ENDFOR
        \STATE $\theta_k^{(t+\frac{1}{2})} \gets \theta_k^{(t, R)}$
    \ENDFOR
    \STATE with probability $1 − p$ do global communication
    \STATE $\hspace{1.3em}$ client $k$ send $\theta_{k}^{(t+\frac{1}{2})}$ to server $\forall k \in [K]$
    \STATE $\hspace{1.3em}$ {\bfseries Server Update:} $\theta^{(t+1)} \gets \frac{1}{K} \sum\limits_{k=1}^K \theta_{k}^{(t+\frac{1}{2})}$ \COMMENT{compute global average}
    \STATE $\hspace{1.3em}$ server broadcast $\theta^{(t+1)}$ to clients $k \in [K]$
    \STATE $\hspace{1.3em}$ on client $k:$ $\theta^{(t+1)}_{k} \gets \theta^{(t+1)}, ~ c_k^{(t+1)} \gets c_k^{(t)} + \frac{p}{\gamma}(\theta^{(t+1)}_{k} - \theta^{(t+\frac{1}{2})}_{k}), ~ \forall k \in [K]$
    \STATE with probability $p$ skip global communication:
    \STATE $\hspace{1.3em}$ {\bfseries Client Update:} $\theta^{(t+1)}_{k} \gets \theta_{k}^{(t+\frac{1}{2})}, ~ c_k^{(t+1)} \gets c_k^{(t)}, ~ \forall k \in [K]$
    \STATE $\hspace{1.3em}$ \COMMENT{on server, $\theta^{(t+1)} \gets \theta^{(t)}$}
\ENDFOR
\end{algorithmic}
\end{algorithm}

References

Contents