Operator Splitting Algorithms in Federated Learning

Why operator splitting? The correctness of the limit

Most existing federated learning algorithms deal with the convergence rate and stability in non-I.I.D. scenarios, while the correctness of the point they converge to has long been overlooked. A simple least-squares example shows that the limits of FedAvg and FedProx generally differ from the true solution of the federated problem [Pathak and Wainwright[1]].

Consider fitting a linear model \(y = mx\) (with \(m\) the unknown) by least squares, with two clients: client 1 holds the data \(\{(0, 2), (1, 2)\}\), hence

\[f_1(m_1) = \frac{1}{2} \left( 4 + (m_1 - 2)^2 \right);\]

client 2 holds the data \(\{(2, 0), (2, 1)\}\), hence

\[f_2(m_2) = \frac{1}{2} \left( 4 m_2^2 + (2 m_2 - 1)^2 \right).\]

Fitting on each client separately yields \(m_1 = 2\) and \(m_2 = \frac{1}{4}\); FedAvg returns the average \(m = \frac{9}{8}\), while FedProx (minimizing the objective with the proximal term) returns

(1)\[m = \frac{18 + 4\mu}{16 + 9\mu},\]

where \(\mu\) is the proximal coefficient. However, fitting all the data jointly gives \(m = \frac{4}{9}\). The discrepancy arises because the objective is quadratic in \(m\): the average of the individually fitted models generally differs from the jointly fitted model.

More generally, [Pathak and Wainwright[1]] analyzed the limit points of FedAvg (in its simplified full-gradient form, termed FedGD) and FedProx: assuming each \(f_k\) is convex and \(L\)-smooth, and that the sequences generated by the algorithm converge, then all clients share a common limit \(\theta^*\) characterized by the fixed-point conditions

(2)\[\sum\limits_{i=1}^R \sum\limits_{k=1}^K \nabla f_k\big( \mathcal{G}_k^{i-1}(\theta^*) \big) = 0, \qquad \mathcal{G}_k(\theta_k) := \theta_k - \eta \nabla f_k(\theta_k),\]

for FedGD (\(R\) being the number of local steps), and

(3)\[\sum\limits_{k=1}^K \nabla \mathcal{M}_{f_k, \mu}(\theta^*) = 0, \qquad \mathcal{M}_{f_k, \mu}(\theta) := \inf_{\theta_k} \left\{ f_k(\theta_k) + \frac{\mu}{2} \lVert \theta - \theta_k \rVert^2 \right\},\]

for FedProx, where \(\mathcal{M}_{f_k, \mu}\) is the Moreau envelope of \(f_k\). Applied to a least-squares problem \(f_k(\theta_k) = \frac{1}{2} \lVert A_k \theta_k + b_k \rVert^2\) (with \(A_k\) of full column rank), these fixed points admit closed forms that generally differ from the true solution \(\theta^* = \big( \sum_{k=1}^K A_k^{\mathrm{T}} A_k \big)^{-1} \sum_{k=1}^K A_k^{\mathrm{T}} b_k\) of the joint problem.

This observation motivates reformulating the federated problem in a form whose splitting-based solution converges to the correct point, which is exactly the route taken by FedSplit and FedDR below.

From monotone inclusions to splitting methods

Start from the constrained form of the federated problem, \(\minimize \ F(\Theta) := \frac{1}{K} \sum_{k=1}^K f_k(\theta_k)\) subject to \(\Theta = \col(\theta_1, \ldots, \theta_K) \in \mathcal{E}\) with the consensus set \(\mathcal{E} = \{ \Theta \mid \theta_1 = \cdots = \theta_K \}\), and penalize the constraint with its indicator function \(\iota_{\mathcal{E}}\) (\(\iota_{\mathcal{E}}(\Theta) = 0\) on \(\mathcal{E}\), \(+\infty\) otherwise). The penalized problem \(\minimize \ F(\Theta) + \iota_{\mathcal{E}}(\Theta)\) has the first-order optimality condition

(4)\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \DeclareMathOperator*{\argmin}{arg\,min} \DeclareMathOperator*{\expectation}{\mathbb{E}} \newcommand{\col}{\operatorname{col}} \newcommand{\R}{\mathbb{R}} \newcommand{\prox}{\mathbf{prox}}\\0 \in \nabla F(\Theta) + \partial \iota_{\mathcal{E}}(\Theta) = \nabla F(\Theta) + \mathcal{N}_{\mathcal{E}}(\Theta),\end{aligned}\end{align} \]

a monotone inclusion problem: \(\nabla F\) and the normal cone \(\mathcal{N}_{\mathcal{E}}\) (here \(\mathcal{E}^{\perp}\) on \(\mathcal{E}\), empty elsewhere) are both maximal monotone operators [Ryu and Yin[2]]. The workhorses for splitting such an inclusion \(0 \in \mathcal{A}(x) + \mathcal{B}(x)\) are the resolvents \(\mathcal{R}_{s\mathcal{T}} := (\mathcal{I} + s\mathcal{T})^{-1}\) and the reflected resolvents (Cayley operators) \(\mathcal{C}_{s\mathcal{T}} := 2\mathcal{R}_{s\mathcal{T}} - \mathcal{I}\). For a gradient or a subdifferential the resolvent is a proximal operator, \(\mathcal{R}_{s\partial f} = \prox_{s f}\), and for the indicator function it is the projection, \(\mathcal{R}_{s\partial\iota_{\mathcal{E}}} = \Pi_{\mathcal{E}}\). The key equivalence is that solving the inclusion is the same as finding a fixed point of the nonexpansive composition \(\mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}\):

(5)\[0 \in \mathcal{A}(x) + \mathcal{B}(x) \quad \Longleftrightarrow \quad \mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}(z) = z, \quad x = \mathcal{R}_{s\mathcal{B}}(z).\]

Since \(\mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}\) is merely nonexpansive, its averaged version \((1 - \tau)\mathcal{I} + \tau\, \mathcal{C}_{s\mathcal{A}} \mathcal{C}_{s\mathcal{B}}\) (\(\tau \in (0, 1]\)) is used for the fixed-point iteration, giving the family of splitting methods

(6)\[\begin{split}\begin{aligned} x^{(t+1/2)} & = \mathcal{R}_{s\mathcal{B}}\left( z^{(t)} \right), \\ z^{(t+1/2)} & = 2 x^{(t+1/2)} - z^{(t)}, \\ x^{(t+1)} & = \mathcal{R}_{s\mathcal{A}}\left( z^{(t+1/2)} \right), \\ z^{(t+1)} & = z^{(t)} + 2\tau \left( x^{(t+1)} - x^{(t+1/2)} \right). \end{aligned}\end{split}\]

The case \(\tau = 1\) is the Peaceman–Rachford splitting method (PRSM); the case \(\tau = \frac{1}{2}\) is the Douglas–Rachford splitting method (DRSM) [Douglas and Rachford[3], Lions and Mercier[4]]. The two differ only in the last (relaxation) step – precisely in the choice of \(\tau\) – and this difference is decisive: unlike PRSM, DRSM converges for general maximal monotone operators without any additional assumptions beyond monotonicity [Ryu and Yin[2]].

FedSplit: Federated Peaceman–Rachford Splitting

Apply the machinery above with \(\mathcal{A} = \nabla F\), \(\mathcal{B} = \partial\iota_{\mathcal{E}} = \mathcal{N}_{\mathcal{E}}\) (and \(s = \frac{1}{\mu}\)). The two resolvents are \(\mathcal{R}_{s\nabla F} = \prox_{F, \mu}\) and \(\mathcal{R}_{s\mathcal{N}_{\mathcal{E}}} = \Pi_{\mathcal{E}}\), and – crucially for federated learning – the proximal operator of the block-separable objective decomposes into client-local proximal steps,

(7)\[\prox_{F, \mu}(\Theta) = \argmin_{\omega} \left\{ \sum_{k=1}^K f_k(\omega_k) + \frac{\mu}{2} \sum_{k=1}^K \lVert \omega_k - \theta_k \rVert^2 \right\} = \left( \prox_{f_k, \mu}(\theta_k) \right)_{k=1}^K,\]

while the projection \(\Pi_{\mathcal{E}}\) is exactly the averaging aggregation at the server. The splitting is thus fully decoupled across clients and fits the center–client communication pattern. Applying the PRSM (\(\tau = 1\)) yields the FedSplit algorithm [Pathak and Wainwright[1]] – the first application of splitting methods to federated learning – whose \(\texttt{prox\_update}_k\) is an inexact client-side solver of \(\prox_{f_k, \mu}\) (e.g. stochastic gradient descent).

FedSplit fixes the correctness problem above: the fixed point of its iteration yields the genuine solution of the federated problem, not a shifted one as for FedAvg/FedProx. Its convergence, however, inherits the caveat of PRSM: for general nonexpansive Cayley operators PRSM carries no convergence guarantee, so additional assumptions (strong convexity and \(L\)-smoothness of \(f_k\)) are needed, under which the rate is

(8)\[\lVert \theta^{(t+1)} - \theta^* \rVert \leqslant \left( 1 - \frac{2}{\sqrt{\kappa}+1} \right)^t \frac{\lVert \theta^{(1)} - \theta^* \rVert}{\sqrt{K}} + (\sqrt{\kappa}+1)\, b,\]

with \(\ell_* = \min_k \ell_k\) the strong convexity modulus, \(L^* = \max_k L_k\), \(\kappa = L^* / \ell_*\), and \(b\) the uniform error bound of the inexact proximal solves. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{FedSplit}}
\begin{algorithmic}
\REQUIRE {\bfseries proximal solvers} $\texttt{prox\_update}_k: \mathbb{R}^d \to \mathbb{R}^d$

\STATE {\bfseries Initiation:} parameters $\theta^{(0)} \in \mathbb{R}^d$

\FOR{each round $t = 0, 1, \cdots$}
    \STATE broadcast $\theta^{(t)}$ to clients $k \in [K]$
    \FOR{each client $k = 1, \cdots, K$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t)} \gets \theta^{(t)}$
        \STATE $\theta_k^{(t+1/2)} \gets$ $\texttt{prox\_update}_k(2\theta^{(t)} - \theta_k^{(t)})$ \COMMENT{local prox step}
        \STATE $\theta_k^{(t+1)} \gets$ $\theta_k^{(t)} + 2(\theta_k^{(t+1/2)} - \theta^{(t)})$ \COMMENT{local centering step}
        \STATE send $\theta_k^{(t+1)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $\theta^{(t+1)} \gets \frac{1}{K} \sum\limits_{k=1}^K \theta_k^{(t+1)}$ \COMMENT{global averaging}
    \IF{meet convergent criteria}
        \STATE $\theta^* \gets \theta^{(t+1)}$
        \STATE {\bfseries break}
    \ENDIF
\ENDFOR
\end{algorithmic}
\end{algorithm}

FedDR: Federated Douglas–Rachford Splitting

The FedDR algorithm [Tran-Dinh et al.[5]] replaces PRSM by the convergent DRSM (\(\tau = \frac{1}{2}\)) and, in doing so, enlarges the admissible problem class to nonconvex smooth objectives, with an explicit regularizer built into the model,

(9)\[\minimize \quad F(\theta) = f(\theta) + g(\theta) = \frac{1}{K} \sum_{k=1}^K f_k(\theta) + g(\theta),\]

where \(f_k\) are \(L\)-smooth (possibly nonconvex) and \(g\) is a closed, proper, convex (CCP) regularizer. Introducing one more block \(\theta_{K+1}\) for \(g\) and penalizing the consensus constraint \(\theta_1 = \cdots = \theta_{K+1}\) with \(\iota_{\mathcal{E}}\) as before, the first-order condition becomes the three-operator inclusion

(10)\[0 \in \nabla f(\Theta) + \partial \left( g + \iota_{\mathcal{E}} \right)(\Theta),\]

with \(\mathcal{B} = \nabla f\) and \(\mathcal{A} = \partial (g + \iota_{\mathcal{E}})\). The resolvent of \(\mathcal{B}\) is \(\mathcal{R}_{s\mathcal{B}} = \prox_{K s f}\) – blockwise, the client-local proximal steps – and the resolvent of \(\mathcal{A}\) handles the regularizer and the consensus constraint jointly at the server. The DRSM iteration ((6) with \(\tau = \frac{1}{2}\), written with relaxation coefficient \(\alpha\)) reads

(11)\[\begin{split}\begin{aligned} \omega_k^{(t+1)} & = \omega_k^{(t)} + \alpha \left( \overline{\theta}^{(t)} - \theta_k^{(t)} \right), & k \in [K], \\ \theta_k^{(t+1)} & = \prox_{f_k, \mu} \left( \omega_k^{(t+1)} \right), & k \in [K], \\ \widehat{\theta}_k^{(t+1)} & = 2\theta_k^{(t+1)} - \omega_k^{(t+1)}, & k \in [K], \\ \widetilde{\theta}^{(t+1)} & = \frac{1}{K}\sum_{k=1}^K \widehat{\theta}_k^{(t+1)}, & \\ \omega^{(t+1)} & = \omega^{(t)} + \alpha \left( \overline{\theta}^{(t)} - \omega^{(t)} \right), & \\ \overline{\theta}^{(t+1)} & = \prox_{g, \frac{K+1}{Ks}} \left( \frac{K}{K+1} \widetilde{\theta}^{(t+1)} + \frac{1}{K+1} \omega^{(t+1)} \right). & \end{aligned}\end{split}\]

The first three lines are the clients’ local work; the last line is the server, which averages the reflected blocks and applies the proximal step of the regularizer \(g\).

As for convergence [Tran-Dinh et al.[5]], assuming \(L\)-smooth \(f_k\) and letting \(e_k^{(t)}\) denote the error of the inexact client-side proximal solve at round \(t\): (i) the gap between the (regularizer-corrected) global model and each local model is bounded by the consecutive local differences and the proximal errors \(\lVert e_k^{(t)} \rVert\); (ii) the squared norm of the gradient mapping \(\mathcal{G}_{\beta}(\theta) := \frac{1}{\beta} \left( \theta - \prox_{\beta g}\left( \theta - \beta \nabla f(\theta) \right) \right)\) of the global objective at \(\overline{\theta}^{(t)}\) is bounded by the same quantities (and the drift of the auxiliary variable \(\omega^{(t)}\)); and (iii) the averaged gradient-mapping norm enjoys an \(\mathcal{O}(1/T)\) bound,

(12)\[\frac{1}{T} \sum_{t=0}^{T-1} \expectation \left[ \lVert \mathcal{G}_{\frac{Ks}{K+1}} \left( \overline{\theta}^{(t)} \right) \rVert^2 \right] \leqslant \frac{C_1 \left( F(\theta^{(0)}) - F^* \right)}{T} + \frac{1}{KT} \sum_{t=0}^{T-1} \sum_{k=1}^K \left( C_2 \varepsilon_{k,t}^2 + C_3 \varepsilon_{k,t+1}^2 \right),\]

with \(F^* = \inf_\theta F(\theta)\); in particular, after

(13)\[T = \left\lfloor \frac{C_1 \left[ F(\theta^{(0)}) - F^* \right] + (C_2 + C_3) M}{\varepsilon^2} \right\rfloor, \qquad M \geqslant \frac{1}{K} \sum_{t=0}^{T-1} \sum_{k=1}^K \varepsilon_{k,t}^2,\]

iterations the output \(\overline{\theta}^{(T)}\) is an \(\varepsilon\)-stationary point of (9). The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{FedDR}}
\begin{algorithmic}
\REQUIRE step size $s = \frac{1}{\mu} > 0$, $\alpha \in (0, 2)$, error bounds $\varepsilon_{k,0} \geqslant 0$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Init server:} global model parameters $\theta^{(0)} \in \text{dom}(f)$, $\overline{\theta}^{(0)} = \widetilde{\theta}^{(0)} = \omega^{(0)} = \theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Init clients:} $\omega_k^{(0)} = \theta^{(0)}$, $\theta_k^{(0)} \approx \mathbf{prox}_{f_k, \mu}(\omega_k^{(0)})$, $\widehat{\theta}_k^{(0)} = 2\theta_k^{(0)} - \omega_k^{(0)}, ~ \forall k \in [K]$

\FOR{$t = 0, 1, \cdots, T-1$}
    \STATE $\mathcal{S}^{(t)} \gets$ (random set of clients) $\subseteq [K]$
    \STATE each client $k \in \mathcal{S}^{(t)}$ receives $\overline{\theta}^{(t)}$ from server \COMMENT{communication}
    \FOR{each client $k \in \mathcal{S}^{(t)}$ {\bfseries in parallel}}
        \STATE choose $\varepsilon_{k,t+1} \geqslant 0$
        \STATE update $\omega_k^{(t+1)} \gets \omega_k^{(t)} + \alpha(\overline{\theta}^{(t)} - \theta_k^{(t)})$
        \STATE $\theta_k^{(t+1)} \approx \mathbf{prox}_{f_k, \mu}(y_k^{(t+1)})$ \COMMENT{inexact local prox step with error bound $\varepsilon_{k,0}$}
        \STATE $\widehat{\theta}_k^{(t+1)} \gets 2\theta_k^{(t+1)} - \omega_k^{(t+1)}$
        \STATE send $\Delta \widehat{\theta}_k^{(t)} = \widehat{\theta}_k^{(t+1)} - \widehat{\theta}_k^{(t)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $\omega^{(t+1)} \gets \omega^{(t)} + \alpha (\overline{\theta}^{(t)} - \omega^{(t)})$
    \STATE $\hspace{1.3em}$ $\widetilde{\theta}^{(t+1)} \gets \widetilde{\theta}^{(t)} + \frac{1}{K}\sum_{k \in \mathcal{S}^{(t)}} \Delta \widehat{\theta}_k^{(t)}$
    \STATE $\hspace{1.3em}$ $\overline{\theta}^{(t+1)} \gets \mathbf{prox}_{g, \frac{K+1}{Ks}} \left( \frac{K}{K+1} \widetilde{\theta}^{(t+1)} + \frac{1}{K+1} \omega^{(t+1)} \right)$
\ENDFOR
\end{algorithmic}
\end{algorithm}

Summary

Splitting methods belong to the family of decomposition algorithms: they break a large problem into smaller subproblems handled separately (here, by the clients and the server), so that each subproblem is easier to solve and easier to parallelize. In theory, they bring many federated algorithms under one umbrella and – unlike FedAvg/FedProx – their limits are guaranteed to be correct (FedSplit) or to approach a stationary point of the true objective (FedDR). The price is bookkeeping: besides the model parameters, every node (server and clients) maintains auxiliary variables (the reflected blocks \(z\), the blended copies \(\omega\)), which adds memory and computation overhead.

References

[1] (1,2,3)

Reese Pathak and Martin J Wainwright. FedSplit: An Algorithmic Framework for Fast Federated Optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, 7057–7066. Curran Associates, Inc., 2020.

[2] (1,2)

Ernest K Ryu and Wotao Yin. Large-Scale Convex Optimization: Algorithms & Analyses via Monotone Operators. Cambridge University Press, 11 2022. doi:10.1017/9781009160865.

[3]

Jim Douglas and H. H. Rachford. On the Numerical Solution of Heat Conduction Problems in Two and Three Space Variables. Transactions of the American Mathematical Society, 82(2):421–439, 1956. doi:10.1090/s0002-9947-1956-0084194-4.

[4]

P. L. Lions and B. Mercier. Splitting Algorithms for the Sum of Two Nonlinear Operators. SIAM Journal on Numerical Analysis, 16(6):964–979, 12 1979. doi:10.1137/0716071.

[5] (1,2)

Quoc Tran-Dinh, Nhan Pham, Dzung T. Phan, and Lam M. Nguyen. FedDR – Randomized Douglas-Rachford Splitting Algorithms for Nonconvex Federated Composite Optimization. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems. 2021.