Primal-Dual Algorithms in Federated Learning

In traditional optimization, primal-dual algorithms are a classic tool for problems with an explicit primal-dual structure. Consider the consensus optimization problem with an explicit global-variable constraint (note the subtle difference from the consensus set formulation in the operator splitting and skipping sections),

(1)\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \DeclareMathOperator*{\maximize}{maximize} \DeclareMathOperator*{\argmin}{arg\,min} \newcommand{\col}{\operatorname{col}} \newcommand{\R}{\mathbb{R}}\\\begin{split}\begin{array}{cl} \minimize & F(\Theta) := \sum\limits_{k=1}^K f_k(\theta_k), \\ \text{subject to} & \theta_k = \theta, ~ k = 1, \ldots, K, \end{array}\end{split}\end{aligned}\end{align} \]

with \(\Theta = \col(\theta_1, \cdots, \theta_K)\) and \(\theta, \theta_1, \ldots, \theta_K \in \R^d\). The augmented Lagrangian (AL) of (1) is

(2)\[\mathcal{L}(\theta, \Theta, \Lambda) = \sum\limits_{k=1}^K \mathcal{L}_k(\theta, \theta_k, \lambda_k), \qquad \mathcal{L}_k(\theta, \theta_k, \lambda_k) = f_k(\theta_k) + \langle \lambda_k, \theta_k - \theta \rangle + \frac{1}{2s} \lVert \theta_k - \theta \rVert^2,\]

where \(\Lambda = \col(\lambda_1, \ldots, \lambda_K)\) collects the dual variables (the Lagrangian multipliers) and \(s > 0\) is the penalty parameter. Applying a gradient ascent-type method to the dual problem \(\maximize_\Lambda \ G(\Lambda) := \inf_{\theta, \Theta} \mathcal{L}(\theta, \Theta, \Lambda)\) yields the generic iteration scheme for this class of problems,

(3)\[\begin{split}\begin{aligned} \theta_k^{(t+1)} & = \argmin_{\theta_k} \mathcal{L}_k\left( \theta^{(t)}, \theta_k, \lambda_k^{(t)} \right), \\ \lambda_k^{(t+1)} & = \lambda_k^{(t)} + \frac{1}{s} \left( \theta_k^{(t+1)} - \theta^{(t)} \right), \\ \theta^{(t+1)} & = \argmin_{\theta} \sum\limits_{k=1}^K \left\{ \frac{1}{2s} \lVert \theta_k^{(t+1)} - \theta \rVert^2 - \langle \lambda_k^{(t+1)}, \theta \rangle \right\} = \frac{1}{K} \sum\limits_{k=1}^K \left( \theta_k^{(t+1)} + s\lambda_k^{(t+1)} \right). \end{aligned}\end{split}\]

FedPD: Federated Primal-Dual Algorithm

Based on this scheme, the FedPD algorithm [Zhang et al.[1]] alternates the client-side local updates, the client-side dual updates, and the server-side global averaging. Two design choices are worth highlighting.

A flexible local solver. For the local update in (3) (the minimization of the local AL \(\mathcal{L}_k\)), FedPD allows an arbitrary client-side solver \(\operatorname{\mathbf{Oracle}}_k\) – plain SGD, or combined with variance-reduction techniques as in SCAFFOLD. This flexibility is a hallmark of federated optimization (the same idea as the server-side FedOpt). In the variance-reduced case, one can further linearize \(f_k\) at the current inner iterate \(\theta_k^{(t, r)}\) (with local gradient \(g_k^{(t, r)}\)), under which the local AL becomes a quadratic whose minimizer has a closed form,

(4)\[\argmin \ \widetilde{\mathcal{L}}(\theta_{k}) = \frac{s}{s+\gamma} \, \theta_k^{(t,r)} + \frac{\gamma}{s+\gamma} \, \theta_{k, 0}^{(t)} - \frac{s\gamma}{s+\gamma}\left( g_k^{(t,r)} + \lambda_k^{(t)} \right),\]

where \(\gamma\) is the linearization (proximal) coefficient and \(\theta_{k, 0}^{(t)}\) is the value of the local model at the beginning of the round.

Skipping the global step. The outer (server) step in (3) is executed only with probability \(1 - p\). When it is skipped, the clients use their own \(\widetilde{\theta}^{(t)} = \theta_k^{(t-1)} + s \lambda_k^{(t-1)}\) in place of the global model,

(5)\[\begin{split}\theta^{(t+1)} = \begin{cases} \frac{1}{K} \sum\limits_{k=1}^K \left( \theta_k^{(t+1)} + s\lambda_k^{(t+1)} \right), & \text{with probability } 1 - p, \\ \theta^{(t)}, & \text{with probability } p, \end{cases}\end{split}\]

so the expensive server round – exactly the communication bottleneck – is skipped with probability \(p\). Moreover, the skipping probability can be made heterogeneity-aware: [Zhang et al.[1]] (TABLE III) relates the achievable \(p\) to the degree of data heterogeneity across clients \(G\) (quantified via the bounded gradient dissimilarity), the target accuracy \(\varepsilon\) of the (epsilon-)stationary point, and the Lipschitz constant \(L\) of \(f, f_k\), with the step size \(s = \frac{\sqrt{5}-1}{8L}\):

(6)\[\begin{split}p \approx \begin{cases} \sqrt{\frac{1}{36s} \cdot \frac{\varepsilon}{G^2}}, & \text{if } 0 \leqslant p < \frac{1-2Ls}{1+Ls}, \\ 1 - 2 \Big/ \log\left(\frac{1}{42s} \cdot \frac{\varepsilon}{G^2}\right), & \text{if } \frac{1-2Ls}{1+Ls} < p < 1. \end{cases}\end{split}\]

The more homogeneous the client data (smaller \(G\)), the more aggressively the communication rounds can be skipped. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{FedPD}}
\begin{algorithmic}
\REQUIRE step size $s = \frac{1}{\mu} > 0,$ skip probability $p \in [0, 1)$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Init server:} global parameters $\theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Init clients:} local parameters$\theta_{0,k}^{(0)} \in \mathbb{R}^d,$ dual variables $\lambda_k^{(0)} \in \mathbb{R}^d,$ $\forall k \in [K]$

\FOR{each round $t = 0, 1, \cdots$}
    \FOR{each client $k = 1, \cdots, K$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t+1)} \gets \textbf{Oracle}_k(\mathcal{L}_k(\theta_{k, 0}^{(t)}, \theta_k, \lambda_k^{(t)}), \theta_k^{(t)})$ \COMMENT{$\textbf{Oracle}_k$ can be SGD, etc.}
        \STATE $\lambda_k^{(t+1)} \gets \lambda_k^{(t)} + \frac{1}{s} (\theta_k^{(t+1)} - \theta_{k, 0}^{(t)})$ \COMMENT{dual update step}
        \STATE $\theta_{k, 0}^{(t+\frac{1}{2})} \gets \theta_k^{(t+1)} + s \lambda_k^{(t+1)}$
    \ENDFOR
    \STATE with probability $1 - p$ do global communication
    \STATE $\hspace{1.3em}$ client $k$ send $\theta_{k, 0}^{(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, 0}^{(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,0} \gets \theta^{(t+1)}, ~ \forall k \in [K]$
    \STATE with probability $p$ skip global communication:
    \STATE $\hspace{1.3em}$ {\bfseries Client Update:} $\theta^{(t+1)}_{k,0} \gets \theta_{k, 0}^{(t+\frac{1}{2})}, ~ \forall k \in [K]$
    \STATE $\hspace{1.3em}$ \COMMENT{On server, $\theta^{(t+1)} \gets \theta^{(t)}$}
\ENDFOR
\end{algorithmic}
\end{algorithm}

References

[1] (1,2)

Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: A Federated Learning Framework With Adaptivity to Non-IID Data. IEEE Transactions on Signal Processing, 69:6055–6070, 2021. doi:10.1109/tsp.2021.3115952.