Personalized Federated Learning Algorithms

In personalized federated learning (PFL), each client seeks a model tailored to its own data distribution while still benefiting from the knowledge of the other clients. A simple and general way to formulate PFL is via weak consensus: instead of enforcing full consensus \(\theta_1 = \cdots = \theta_K\), one adds a regularizer \(\mathcal{R}\) that penalizes the deviation between the local models. This page collects the PFL algorithms (implemented in fl_sim or closely related) together with their modeling, theory, and pseudocodes.

L2SGD: Loopless Local SGD

The simplest algorithm of this family is the loopless stochastic gradient descent (Loopless Local SGD, L2SGD) introduced in [Hanzely and Richtárik[1]]. It considers the following weak-consensus problem

(1)\[\begin{split}\begin{array}{cl} \minimize & F(\Theta) = f(\Theta) + \lambda \varphi(\Theta) \\ \text{where} & f(\Theta) = \frac{1}{K} \sum\limits_{k=1}^K f_k(\theta_k), \\ & \varphi(\Theta) = \frac{1}{2K} \sum\limits_{k=1}^K \left\lVert \theta_k - \bar{\theta} \right\rVert^2 = \frac{1}{2K} \sum\limits_{k=1}^K \left\lVert \theta_k - \frac{1}{K} \sum\limits_{j=1}^K \theta_j \right\rVert^2, \\ & \Theta = \col(\theta_1, \ldots, \theta_K), \end{array}\end{split}\]

where \(\lambda \geqslant 0\) is the penalty coefficient and \(\bar{\theta} = \frac{1}{K} \sum_{k=1}^K \theta_k\) is the global average model. For \(0 < \lambda < \infty\), the solution \(\bar{\theta}, \theta_1, \ldots, \theta_K\) is called a set of mixed models. It is easy to see that this problem is equivalent to the problem considered by FedProx (cf. the constrained problem in the proximal algorithms section).

Instead of evaluating the full gradient \(\nabla F = \nabla f + \lambda \nabla \varphi\) at every iteration, [Hanzely and Richtárik[1]] defines an unbiased stochastic gradient of \(F\)

(2)\[\begin{split}G(\Theta) := \begin{cases} \frac{\nabla f(\Theta)}{1 - p}, & \text{with probability } 1 - p, \\ \frac{\lambda \nabla \varphi(\Theta)}{p}, & \text{with probability } p, \end{cases}\end{split}\]

with \(p \in (0, 1)\). Note that computing \(\nabla f(\Theta) = \col(\nabla f_1(\theta_1), \ldots, \nabla f_K(\theta_K))\) is completely block-separable: in the federated setting, this step is carried out locally on the clients without any communication with the server, which is a significant advantage in communication-constrained scenarios. In this sense, L2SGD already incorporates the loopless (skipping) idea; cf. the ProxSkip algorithm in the skipping algorithms section. Under the additional finite-sum assumption \(f_k(\theta_k) = \sum_{j=1}^m f_{k,j}(\theta_k)\), one can further combine the stochastic gradient with randomized Kaczmarz-type (importance-sampled) updates; see [Kovalev et al.[2]] for the analysis of loopless methods. The resulting method is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{L2SGD}}
\begin{algorithmic}
\REQUIRE penalty coefficient $\lambda,$ learning rate $\eta,$ probability $p \in (0, 1)$

\STATE {Initiation:}
\STATE local (clients) model parameters $\theta_k^{(0)} \in \R^d, ~ \forall k \in [K]$

\FOR{each round $t = 0, 1, \cdots, T-1$}
\STATE \COMMENT{with probability $1-p$: no global communication, only clients update}
\STATE \FOR{each client $k \in [K]$ {in parallel}}
\STATE sample $j \in [m]$ (uniformly at random)
\STATE $g_k^{(t)} \gets \frac{1}{K(1-p)} \left( \nabla f_{k,j}(\theta_k^{(t)}) \right)$
\STATE $\theta_k^{(t+1)} \gets \theta_k^{(t)} - \eta g_k^{(t)}$
\STATE \ENDFOR

\STATE \COMMENT{with probability $p$: global communication}
\STATE client $k$ sends $\theta_{k}^{(t)}$ to the server, $\forall k \in [K]$
\STATE {Server Update:}
\STATE $\hspace{1.3em}$ $\theta^{(t)} \gets \frac{1}{K} \sum\limits_{k=1}^K \theta_{k}^{(t)}$ \COMMENT{compute the global average}
\STATE server broadcasts $\theta^{(t)}$ to clients $k \in [K]$
\STATE {Clients Update:}
\STATE \FOR{each client $k \in [K]$ {in parallel}}
\STATE $g_k^{(t)} \gets \frac{\lambda}{Kp} \left( \theta_k^{(t)} - \theta^{(t)} \right)$
\STATE $\theta_k^{(t+1)} \gets \theta_k^{(t)} - \eta g_k^{(t)}$
\STATE \ENDFOR
\ENDFOR
\end{algorithmic}
\end{algorithm}

pFedMe: Personalized Federated Mean

The pFedMe algorithm [Dinh et al.[3]] considers a weak-consensus problem of the same family as (1), with the global model appearing as a single variable \(\theta \in \R^d\),

(3)\[ \begin{align}\begin{aligned}\DeclareMathOperator*{\minimize}{minimize} \DeclareMathOperator*{\argmin}{arg\,min} \newcommand{\R}{\mathbb{R}} \newcommand{\prox}{\mathbf{prox}}\\\minimize_{\theta, \theta_1, \ldots, \theta_K} \quad \sum\limits_{k=1}^K \left\{ f_k(\theta_k) + \frac{\lambda}{2} \left\lVert \theta_k - \theta \right\rVert^2 \right\},\end{aligned}\end{align} \]

where \(\lambda \geqslant 0\) is the penalty coefficient. Exploiting the center–client structure of federated learning, the problem is equivalently reformulated as a bi-level optimization problem,

(4)\[\begin{split}\begin{array}{cl} \minimize & F(\theta) := \frac{1}{K} \sum\limits_{k=1}^K F_k(\theta), \\ \text{where} & F_k(\theta) = \min\limits_{\theta_k \in \R^d} \left\{ f_k(\theta_k) + \frac{\lambda}{2} \left\lVert \theta_k - \theta \right\rVert^2 \right\}, \end{array}\end{split}\]

whose inner objective \(F_k(\theta)\) is exactly the Moreau envelope of \(f_k\), \(\mathcal{M}_{f_k, \lambda}(\theta) := \inf_{\theta_k} \{ f_k(\theta_k) + \frac{\lambda}{2} \lVert \theta_k - \theta \rVert^2 \}\), and the inner minimizer is the proximal operator \(\prox_{f_k, \lambda}(\theta) := \argmin_{\theta_k} \{ f_k(\theta_k) + \frac{\lambda}{2} \lVert \theta_k - \theta \rVert^2 \}\) evaluated at the global model. This proximal point is precisely the personalized model of client \(k\). By the Moreau envelope–proximal operator relation, the gradient of the outer objective is

(5)\[\nabla F(\theta) = \frac{1}{K} \sum\limits_{k=1}^K \nabla \mathcal{M}_{f_k, \lambda} (\theta) = \frac{\lambda}{K} \sum\limits_{k=1}^K \left( \theta - \prox_{f_k, \lambda} (\theta) \right),\]

so that one (stochastic) gradient step on \(F\) decomposes into the average of the client-side partial updates \(\theta^{(t)} - \eta \lambda \left( \theta^{(t)} - \prox_{f_k, \lambda} (\theta^{(t)}) \right)\). Operationally, each client holds a copy \(\omega_k^{(t, r)}\) of the global model and, between two communications, alternates for \(R\) rounds

(6)\[\theta_k^{(t, r)} \gets \prox_{f_k, \lambda} \left( \omega_k^{(t, r)}; \ b_r \right), \qquad \omega_k^{(t, r+1)} \gets \omega_k^{(t, r)} - \eta \lambda \left( \omega_k^{(t, r)} - \prox_{f_k, \lambda} \left( \omega_k^{(t, r)} \right) \right),\]

where \(b_r\) is the minibatch of round \(r\); the server then aggregates

(7)\[\theta^{(t+1)} \gets (1-\beta)\,\theta^{(t)} + \frac{\beta}{\# \mathcal{S}^{(t)}} \sum\limits_{k \in \mathcal{S}^{(t)}} \omega_{k}^{(t, R)}.\]

Note that the two parts of (6) are executed for \(R\) local rounds before any communication happens – yet another instance of the loopless (skipping) idea; cf. L2SGD and the skipping algorithms section. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{pFedMe}}
\begin{algorithmic}
\REQUIRE penalty coefficient $\lambda,$ learning rate $\eta,$ global update smoothing factor $\beta,$ {\bfseries proximal solvers} $\texttt{prox\_update}_k$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Init server:} global model parameters $\theta^{(0)} \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Init clients:} personalized model parameters $\theta_k^{(0)} \in \mathbb{R}^d, ~ \forall k \in [K]$

\FOR{each round $t = 0, \cdots, T-1$}
    \STATE Server sends $\theta^{(t)}$ to all clients
    \FOR{each client $k = 1, \cdots, K$ in parallel}
        \STATE $\omega_{k}^{(t, 0)} = \theta^{(t)}$ \COMMENT{a copy of the global model $\theta^{(t)}$}
        \FOR{$r = 0,\cdots, R-1$}
            \COMMENT{solve inner problem $\min\limits_{\theta_k \in \mathbb{R}^d} \left\{ f_k(\theta_k) + \frac{\lambda}{2} \left\lVert \theta_k - \theta \right\rVert^2 \right\}$}
            \STATE sample a mini-batch $b_r$
            \COMMENT{use proximal solver $\texttt{prox\_update}_k$ to solve $\underset{\theta_k}{\text{argmin}} \left\{ \ell_k(\theta_k; b_r) + \frac{\lambda}{2} \left\lVert \theta_k - \omega_{k}^{(t, r)} \right\rVert^2 \right\}, \text{ where } f_k(\theta_k) = \mathbb{E}_{b}[\ell_k(\theta_k; b)]$}
            \STATE $\theta_k^{(t, r)} \gets \texttt{prox\_update}_k (\omega_{k}^{(t, r)}; b_r)$
            \COMMENT{update the local copy of the global model}
            \STATE $\omega_{k}^{(t, r+1)} \gets \omega_{k}^{(t, r)} - \eta \lambda \left( \omega_{k}^{(t, r)} - \theta_k^{(t, r)} \right)$
        \ENDFOR
        \STATE $\theta_k^{(t+1)} \gets \theta_k^{(t, R)}$ \COMMENT{update the personalized model}
    \ENDFOR
    \STATE Server uniformly samples a subset of clients $\mathcal{S}^{(t)}$
    \STATE each client in $\mathcal{S}^{(t)}$ sends the local model $\omega_{k}^{(t, R)}$ to the server
    \STATE Sever update $\theta^{(t+1)} \gets (1-\beta)\theta^{(t)} + \frac{\beta}{\# \mathcal{S}^{(t)}} \sum\limits_{k \in \mathcal{S}^{(t)}} \omega_{k}^{(t, R)}$
\ENDFOR
\STATE final global model: $\theta^* \gets \theta^{(T)}$
\STATE final personalized models: $\theta_k^* \gets \theta_k^{(T)}$
\end{algorithmic}
\end{algorithm}

Ditto: Fair and Robust Personalized Models

The Ditto algorithm [Li et al.[4]] develops the proximal-subproblem idea of FedProx into a bi-level formulation seen from the client’s perspective (APFL, proposed slightly earlier, adopts a similar viewpoint; see below),

(8)\[\begin{split}\begin{array}{cl} \minimize\limits_{\omega_k} & h_k(\omega_k, \theta^*) := f_k(\omega_k) + \frac{\mu}{2} \lVert \omega_k - \theta^* \rVert^2, \\ \text{subject to} & \theta^* \in \argmin\limits_{\theta} G\left( f_1(\theta), \cdots, f_K(\theta) \right), \end{array}\end{split}\]

where \(G\) is the aggregation rule applied by the server to the client models, e.g. plain averaging as in FedAvg. Reading the formulation: the inner problem is a plain federated optimization problem on the global model \(\theta\) alone, while the outer problem attaches to each client a proximal task for personalization – and the personalized model \(\omega_k\) does not participate in the global update. The global and the personalized models are thereby decoupled, and their updates can be scheduled flexibly.

Consequently, each client solves two problems per round: the global update \(\theta_k^{(t)} \gets \operatorname{\mathbf{UpdateGlobal}}(\theta^{(t)}, f_k)\) (the inner problem), and the proximal task whose solution tracks the moving center \(\theta^{(t)}\),

(9)\[\omega_k^{(t+1)} \gets \omega_k^{(t)} - \eta \left( \nabla f_k(\omega_k^{(t)}) + \mu \left( \omega_k^{(t)} - \theta^{(t)} \right) \right).\]

Unlike FedSplit or pFedMe, where the proximal operator is solved in an inner loop to (approximate) optimality, Ditto executes only a single (stochastic) gradient step per round for the outer problem, letting the proximal center \(\theta^{(t)}\) evolve alongside the global model. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{Ditto}}
\begin{algorithmic}
\REQUIRE penalty coefficient $\mu,$ learning rate $\eta,$ methods $\textbf{UpdateGlobal}, \textbf{Aggregate}$

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

\FOR{each round $t = 0, 1, \cdots, T-1$}
    \STATE $\mathcal{S}^{(t)} \gets$ (random set of clients) $\subseteq [K]$
    \STATE broadcast $\theta^{(t)}$ to clients $k \in \mathcal{S}^{(t)}$
    \FOR{each client $k \in \mathcal{S}^{(t)}$ {\bfseries in parallel}}
        \COMMENT{solve the local sub-problem of $G(\dot)$ inexactly starting from $\theta^{(t)}$ to obtain $\theta_k^{(t)}$}
        \STATE $\theta_k^{(t)} \gets \textbf{UpdateGlobal}(\theta^{(t)}, f_k)$
        \COMMENT{update the personalized model via solving the proximal problem}
        \STATE $\omega_k^{(t+1)} \gets \omega_k^{(t)} - \eta \left( \nabla f_k(\omega_k^{(t)}) + \mu (\omega_k^{(t)} - \theta^{(t)}) \right)$
        \STATE send $\Delta_k^{(t)} \gets \theta_k^{(t)} - \theta^{(t)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $\theta^{(t+1)} \gets \textbf{Aggregate} (\theta^{(t)}, \{ \Delta_k^{(t)} \}_{k \in \mathcal{S}^{(t)}})$
\ENDFOR
\end{algorithmic}
\end{algorithm}

APFL: Adaptive Personalized Federated Learning

Proposed slightly before Ditto [Deng et al.[5]], APFL considers the bi-level problem

(10)\[\begin{split}\begin{array}{cl} \minimize\limits_{\omega_k \in \mathrm{dom} f_k} & h_k(\omega_k, \theta^*) := f_k\left( \alpha_k \omega_k + (1 - \alpha_k) \theta^* \right), \\ \text{subject to} & \theta^* \in \argmin\limits_{\theta} \frac{1}{K} \sum\limits_{k=1}^K f_k(\theta), \end{array}\end{split}\]

with the mixture weight \(\alpha_k \in (0, 1)\). The main difference from (8) is the personalization mechanism: instead of a quadratic proximal term, the local objective is evaluated at the convex combination \(\alpha_k \omega_k + (1 - \alpha_k) \theta^*\) of an auxiliary local model and the global model, and \(\alpha_k\) plays the same “personalization degree” role as the penalty \(\mu\) in a proximal term. The resulting personalized model is

(11)\[\hat{\omega}_k = \alpha_k \omega_k^* + (1 - \alpha_k) \theta^*, \qquad \omega_k^* = \argmin\limits_{\omega_k \in \mathrm{dom} f_k} f_k\left( \alpha_k \omega_k + (1 - \alpha_k) \theta^* \right).\]

A caveat worth noting: when \(\mathrm{dom} f_k\) is a proper subset of \(\R^d\), the convex combination in (11) may fail to attain a minimizer of \(f_k\) – the range of \(\alpha_k \omega_k^* + (1 - \alpha_k) \theta^*\) can be a strict subset of the search domain that excludes the minimizer. Enlarging the search domain of \(\omega_k\) restores attainability, but then \(\hat{\omega}_k\) could equally be obtained by minimizing \(f_k\) alone, i.e. by pure local training – which contradicts the very goal of personalized federated learning. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{APFL}}
\begin{algorithmic}
\REQUIRE mixture weights $\alpha_1, \ldots, \alpha_K,$ synchronization gap $\tau$

\STATE {Initiation:}
    \STATE $\hspace{1.3em}$ {Init server:} global model parameters $\theta^{(0)} \in \mathbb{R}^d,$ random set of clients $S^{(t)} \subseteq [K]$
    \STATE $\hspace{1.3em}$ {Init clients:} local model parameters $\omega_k^{(0)} \in \mathbb{R}^d, ~ \theta_k^{(0)} \gets \theta^{(0)}, ~ \forall k \in [K]$
    \STATE $\hspace{1.3em}$ {Constants:} condition number $\kappa \gets \frac{L}{\mu},$ $a \gets \max\{128\kappa, \tau\}$

\FOR{each round $t = 0, 1, \cdots, T-1$}
    \FOR{each client $k \in \mathcal{S}^{(t)}$ {in parallel}}
        \STATE $\bar{\omega}_k^{(t)} \gets \alpha_k \omega_k^{(t)} + (1-\alpha_k) \theta_k^{(t)}$ \COMMENT{mixture model}
        \STATE $\eta^{(t)} \gets \frac{16}{\mu(t+a)}$ \COMMENT{decay learning rate}
        \STATE $\theta_k^{(t+1)} \gets \theta_k^{(t)} - \eta^{(t)} \nabla f_k(\theta_k^{(t)})$ \COMMENT{inner problem (global model) update}
        \STATE $\omega_k^{(t+1)} \gets \omega_k^{(t)} - \eta^{(t)} \nabla f_k(\bar{\omega}_k^{(t)})$ \COMMENT{outer problem (personalized model) update}
        \COMMENT{Optional: adaptive mixture weights update}
        \COMMENT{$\alpha_k \gets \alpha_k - \eta^{(t)} \nabla_{\alpha_k} f_k(\bar{\omega}_k^{(t)}) = \alpha_k - \eta^{(t)} \left\langle \omega_k^{(t)} - \theta_k^{(t)}, \nabla f_k(\bar{\omega}_k^{(t)}) \right\rangle$}
    \ENDFOR

    \IF{$t$ not divides synchronization gap $\tau$}
        \STATE {Server updates:} $\mathcal{S}^{(t+1)} \gets \mathcal{S}^{(t)}$
    \ELSE
        \STATE each client $k \in \mathcal{S}^{(t)}$ sends $\theta_k^{(t+1)}$ to server
        \STATE {Server updates:}
            \STATE $\hspace{1.3em}$ $\theta^{(t+1)} \gets \frac{1}{\# \mathcal{S}^{(t)}} \sum_{k \in \mathcal{S}^{(t)}} \theta_k^{(t+1)}$
            \STATE $\hspace{1.3em}$ $\mathcal{S}^{(t+1)} \gets$ (random set of clients) $\subseteq [K]$
            \STATE $\hspace{1.3em}$ broadcast $\theta^{(t+1)}$ to clients $k \in S^{(t+1)}:$ $\theta_k^{(t+1)} \gets \theta^{(t+1)}$
    \ENDIF
\ENDFOR

\STATE final personalized model: $\hat{\omega}_k \gets \frac{1}{S_T}\sum\limits_{t=1}^{T} p_t \left( \alpha_k \omega_k^{(t)} + (1-\alpha_k)\frac{1}{\#\mathcal{S}^{(t-1)}}\sum\limits_{k\in\mathcal{S}^{(t-1)}}\theta_k^{(t)} \right)$
\STATE final global model: $\hat{\theta} \gets \frac{1}{S_T}\sum\limits_{t=1}^{T} \frac{p_t}{\#\mathcal{S}^{(t-1)}} \sum\limits_{k\in\mathcal{S}^{(t-1)}}\theta_k^{(t)}$
\STATE where $p_t = (t+a)^2, S_T = \sum\limits_{t=1}^{T}p_t$
\end{algorithmic}
\end{algorithm}

FedDyn: Dynamic Regularized Federated Learning

The algorithms above (and FedProx alike) regularize the local objective with a static proximal term: in round \((t+1)\), client \(k\) minimizes \(f_k(\theta_k) + \frac{\mu}{2} \lVert \theta_k - \theta^{(t)} \rVert^2\). The FedDyn algorithm [Acar et al.[6]] refines this into a dynamically regularized objective,

(12)\[F_k(\theta_k) := f_k(\theta_k) - \left\langle \nabla f_k\left( \theta_k^{(t)} \right), \theta_k \right\rangle + \frac{\mu}{2} \lVert \theta_k - \theta^{(t)} \rVert^2,\]

which is equivalent to

(13)\[F_k(\theta_k) = \left( f_k(\theta_k) + \frac{\mu}{2} \lVert \theta_k - \theta^{(t)} \rVert^2 \right) - \left\langle \nabla f_k\left( \theta_k^{(t)} \right), \theta_k - \theta_k^{(t)} \right\rangle.\]

In words: each (stochastic) step on \(F_k\) follows the descent direction of the prox-augmented objective, corrected by the linearization of \(f_k\) at the previous local iterate \(\theta_k^{(t)}\) (with \(\mathfrak{g}_k^{(t)} := \nabla f_k(\theta_k^{(t)})\)). Another remarkable property is the first-order optimality condition of (12),

(14)\[\nabla f_k\left( \theta_k^{(t+1)} \right) - \nabla f_k\left( \theta_k^{(t)} \right) + \mu \left( \theta_k^{(t+1)} - \theta^{(t)} \right) = 0,\]

which yields the recursion

(15)\[\mathfrak{g}_k^{(t+1)} \gets \mathfrak{g}_k^{(t)} - \mu \left( \theta_k^{(t+1)} - \theta^{(t)} \right),\]

i.e. the local gradient estimate is maintained without additional function or gradient evaluations – a substantial computational and implementational convenience. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{FedDyn}}
\begin{algorithmic}
\REQUIRE penalty coeffecient $\mu$

\STATE {\bfseries Initiation:}
\STATE $\hspace{1.3em}$ {\bfseries Init server:} global model parameters $\theta^{(0)} \in \mathbb{R}^d,$ $h = 0 \in \mathbb{R}^d$
\STATE $\hspace{1.3em}$ {\bfseries Init clients:} local gradient $\mathfrak{g}_k^{(0)} \gets 0 \in \mathbb{R}^d, ~ \forall k \in [K]$

\FOR{each round $t = 0, 1, \cdots, T-1$}
    \STATE $\mathcal{S}^{(t)} \gets$ (random set of clients) $\subseteq [K]$
    \STATE broadcast $\theta^{(t)}$ to clients $k \in \mathcal{S}^{(t)}$
    \FOR{each client $k \in \mathcal{S}^{(t)}$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t+1)} \gets \underset{\theta_k}{\text{argmin}} \left\{ f_k(\theta_k) - \langle \mathfrak{g}_k^{(t)}, \theta_k \rangle + \frac{\mu}{2} \lVert \theta_k - \theta^{(t)} \rVert^2 \right\}$
        \STATE $\mathfrak{g}_k^{(t+1)} \gets \mathfrak{g}_k^{(t)} - \mu (\theta_k^{(t+1)} - \theta^{(t)})$ \COMMENT{update local gradient}
        \STATE send $\theta_k^{(t+1)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $h^{(t+1)} \gets h^{(t)} - \frac{\mu}{K} \left(\sum\limits_{k \in \mathcal{S}^{(t)}} \theta_k^{(t+1)} - \theta^{(t)} \right)$
    \STATE $\hspace{1.3em}$ $\theta^{(t+1)} \gets \left( \frac{1}{\# \mathcal{S}^{(t)}}\sum\limits_{k \in \mathcal{S}^{(t)}} \theta_k^{(t+1)} \right) - \frac{1}{\mu} h^{(t+1)}$
\ENDFOR
\end{algorithmic}
\end{algorithm}

pFedMac: Personalized FL via Maximizing Correlation

Related to the “direction correction” of FedDyn, the pFedMac algorithm [Li et al.[7]] considers the bi-level problem

(16)\[\begin{split}\begin{array}{cl} \minimize & F(\theta) = \frac{1}{K} \sum\limits_{k=1}^K F_k(\theta), \\ \text{where} & F_k(\theta) := \min\limits_{\theta_k} \left\{ f_k(\theta_k) - \mu \langle \theta_k, \theta \rangle + \frac{\mu}{2} \lVert \theta \rVert^2_2 \right\}, \end{array}\end{split}\]

in which the inner problem corrects the personalized model \(\theta_k\) using the global model \(\theta\) through their inner product. In probability and statistics, \(\langle \theta_k, \theta \rangle\) quantifies the correlation between the personalized and the global models; hence the technique is known as maximizing correlation. Two elementary rewritings connect it to the formulations above: moving the global term out of the inner minimization gives \(F_k(\theta) = F_k'(\theta) + \frac{\mu}{2} \lVert \theta \rVert^2_2\) with \(F_k'(\theta) := \min_{\theta_k} \{ f_k(\theta_k) - \mu \langle \theta_k, \theta \rangle \}\), and completing the square shows

(17)\[F_k(\theta) = \min\limits_{\theta_k} \left\{ f_k(\theta_k) + \frac{\mu}{2} \lVert \theta_k - \theta \rVert^2_2 - \frac{\mu}{2} \lVert \theta_k \rVert^2_2 \right\},\]

i.e. the prox-augmented objective corrected by a quadratic term in \(\theta_k\) itself – exactly the flavor of (13), so the same one-step direction-correction updates apply. pFedMac further studies sparse personalized models via an additional \(\ell_1\) term,

(18)\[F_k(\theta) := \min\limits_{\theta_k} \left\{ f_k(\theta_k) - \mu \langle \theta_k, \theta \rangle + \frac{\mu}{2} \lVert \theta \rVert^2_2 + \gamma \lVert \theta_k \rVert_1 \right\},\]

and extends the center–client hierarchy to a three-level structure with sub-centers (edges or hubs), each aggregating a cluster \(\{ \theta_{k,1}, \ldots, \theta_{k, J_k} \}\) of clients before communicating with the center – a structure expected to become increasingly common as federated networks scale, and closely related to clustered FL. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{pFedMac}}
\begin{algorithmic}
\REQUIRE learning rate $\eta$, penalty coefficient $\lambda$, $\beta$

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

\FOR{each round $t = 0, 1, \cdots, T-1$}
    \STATE $\mathcal{S}^{(t)} \gets$ (random set of clients) $\subseteq [K]$
    \STATE broadcast $\theta^{(t)}$ to clients $k \in \mathcal{S}^{(t)}$
    \FOR{each client $k \in \mathcal{S}^{(t)}$ {\bfseries in parallel}}
        \STATE $\theta_k^{(t)} \gets$ {\bfseries ClientUpdate}$(k, \theta^{(t)})$
        \STATE send $\theta_k^{(t)}$ to server
    \ENDFOR
    \STATE {\bfseries Server Update:}
    \STATE $\hspace{1.3em}$ $\theta^{(t+1)} \gets (1 - \beta) \theta^{(t)} + \frac{\beta}{\lvert \mathcal{S}^{(t)} \rvert} \sum\limits_{k\in \mathcal{S}^{(t)}} \theta_k^{(t)}$
\ENDFOR

\PROCEDURE{ClientUpdate}{$(k, \theta)$} \COMMENT{on client $k$}
\STATE $\omega_k^{(t,0)} = \theta_k^{(t,0)} = \theta^{(t)}$
\FOR{local step $r = 0, 1, \cdots, R-1$}
  \STATE $\mathcal{D}_{k, r} \gets$ (sample a mini-batch data)
  \STATE $\omega_k^{(t,r)} \gets \argmin_{\omega_k} \left\{ \ell_k(\omega_k; \mathcal{D}_{k, r}) - \lambda \langle \omega_k, \theta_k^{(t,r)} \rangle \right\}$
  \STATE $\theta_k^{(t,r+1)} \gets \theta_k^{(t,r)} - \eta\lambda \left( \theta_k^{(t,r)} - \omega_k^{(t,r)} \right)$
\ENDFOR
\RETURN{$\theta_k^{(t,R)}$}
\ENDPROCEDURE
\end{algorithmic}
\end{algorithm}

FPFC: Fusion Penalized Federated Clustering

Beyond regularizing the deviation between each pair of local models (as in L2SGD) or between the local model and the global one, one can also encourage clusters of clients to share identical models, without knowing the clusters a priori. The fusion penalized federated clustering (FPFC) algorithm (see arXiv:2211.04218) achieves this by imposing a nonconvex pairwise fusion penalty on the model differences. Introducing the auxiliary variables \(\theta_{ij} \equiv \omega_i - \omega_j\) for all pairs \(i, j \in [m]\), it considers

(19)\[\begin{split}\begin{array}{cl} \minimize\limits_{\omega, \theta} & \sum\limits_{i=1}^m f_i(\omega_i) + \frac{1}{2m} \sum\limits_{i=1}^m \sum\limits_{j=1}^m \tilde{g}(\lVert \theta_{ij} \rVert) \\ \text{subject to} & \omega_i - \omega_j = \theta_{ij}, \qquad i, j \in [m], \end{array}\end{split}\]

where \(\tilde{g}(\cdot)\) is a smoothed SCAD-type penalty (parameterized by \(\xi, \lambda, a\)) whose proximal operator shrinks small pairwise differences towards zero (clients in the same cluster) while keeping large ones (clients in different clusters) almost untouched. The problem is solved via an ADMM-style (Douglas–Rachford) splitting with the augmented Lagrangian penalty \(\rho\) and dual variables \(v_{ij}\): each round, a random subset of devices performs \(T_i\) local gradient epochs on \(f_i(\omega) + \frac{\rho}{2} \lVert \omega - \zeta_i \rVert^2\), and the server applies the piecewise proximal update to \(\theta_{ij}\) and ascends the duals for the pairs of active devices. After training, devices \(i\) and \(j\) are assigned to the same cluster whenever \(\lVert \theta_{ij} \rVert\) is (nearly) zero. The pseudocode is summarized below; note that neither the number of clusters nor the cluster membership needs to be specified in advance.

\begin{algorithm}
\caption{PseudoCode for \texttt{FPFC}}
\begin{algorithmic}
\REQUIRE number of devices $m,$ number of rounds $K,$ local epochs $T_i,$ stepsize $\alpha,$ penalty parameters $\xi, \lambda, a,$ and $\rho$

\STATE {Initiation:}
\STATE $\omega_1^0 = \cdots = \omega_m^0$ on the clients
\STATE $\zeta_i^0 = \omega_i^0$ for $i \in [m]$, $\theta_{ij}^0 = 0$ and $v_{ij}^0 = 0$ for $i, j \in [m]$ on the server

\FOR{each round $k = 0, 1, \cdots, K-1$}
\STATE the server randomly chooses a subset of devices $\mathcal{A}_k$
\STATE the server downloads $\zeta_i^k$ to each device $i \in \mathcal{A}_k$
\STATE \FOR{each device $i \in \mathcal{A}_k$ {in parallel}}
\STATE $\omega_i^{k,0} = \omega_i^k$
\STATE \FOR{$t = 0, 1, \cdots, T_i - 1$}
\STATE $\hspace{1.3em}$ $\omega_i^{k,t+1} = \omega_i^{k,t} - \alpha \left[ \nabla f_i(\omega_i^{k,t}) + \rho \left( \omega_i^{k,t} - \zeta_i^k \right) \right]$
\STATE \ENDFOR
\STATE $\omega_i^{k+1} = \omega_i^{k,T_i}$
\STATE \ENDFOR
\STATE each device $i \in \mathcal{A}_k$ uploads $\omega_i^{k+1}$ to the server
\STATE {Server Update:}
\STATE \FOR{$i, j \in \mathcal{A}_k$ ($i < j$)}
\STATE $\hspace{1.3em}$ $\delta_{ij}^{k+1} = \omega_i^{k+1} - \omega_j^{k+1} + v_{ij}^k / \rho$
\STATE $\hspace{1.3em}$ $\theta_{ij}^{k+1} \gets \underset{\theta}{\text{prox}}$ of the (scaled) fusion penalty at $\delta_{ij}^{k+1}$ (Eq. (6) of the paper)
\STATE $\hspace{1.3em}$ $v_{ij}^{k+1} = v_{ij}^k + \rho \left( \omega_i^{k+1} - \omega_j^{k+1} - \theta_{ij}^{k+1} \right)$
\STATE $\hspace{1.3em}$ $\theta_{ji}^{k+1} = -\theta_{ij}^{k+1}, ~ v_{ji}^{k+1} = -v_{ij}^{k+1}$
\STATE \ENDFOR
\STATE for $i \notin \mathcal{A}_k$ or $j \notin \mathcal{A}_k$: $\theta_{ij}^{k+1} = \theta_{ij}^k, ~ v_{ij}^{k+1} = v_{ij}^k$
\STATE for $i \in [m]$: $\hspace{0.5em}$ $\zeta_i^{k+1} = \frac{1}{m} \sum\limits_{j=1}^m \left( \omega_j^{k+1} + \theta_{ij}^{k+1} - v_{ij}^{k+1} / \rho \right)$
\ENDFOR
\end{algorithmic}
\end{algorithm}

FedCR: Personalized Federated Learning Based on Across-Client Common Representation

The methods above personalize the whole local model. FedCR (see ICML 2023) instead decouples each local model \(w_i = [w_i^f, w_i^p]\) into a feature extractor (body) \(w_i^f\), shared across clients and aggregated at the server, and a predictor (head) \(w_i^p\), kept personal. The extractors are aligned through a common representation: for each class \(c\), the local class-conditional mutual information is constrained to be close to the global one, which (via the Markov chain \(y \to x \to z\)) amounts to a class-wise feature alignment

(20)\[I(z; x \mid y_i) - I_i(z; x_i \mid y_i) = \mathbb{E}_{p(x_i, y_i)} \mathbb{E}_{p(x \mid x_i, y_i)} \left[ \mathrm{KL}\left[ p(z \mid x) \,\middle\|\, p(z \mid x_i) \right] \right],\]

weighted by a Lagrange multiplier \(\beta \geqslant 0\) (the CMI regularizer). The representation mapping is modeled as a diagonal Gaussian \(p(z \mid x_i) = \mathcal{N}\left(z \,\middle|\, \mu(x_i), \Sigma(x_i)\right)\), under which the intractable global posterior is estimated by a product of experts (PoE) over the class-wise local posteriors of the active clients \(\mathcal{P}_t\), with prior \(p(z) = \mathcal{N}(0, 1)\):

(21)\[p(z^c \mid x) \propto p(z) \prod_{i \in \mathcal{P}_t} p^c(z \mid x_i), \qquad p^c(z \mid x_i) = \prod_{n = 1}^{N_i^{c}} p\left(z \,\middle|\, x_i^{(n)}\right)^{1 / N_i^{c}},\]

where \(p^c(z \mid x_i)\) is the (geometric-mean) local PoE over the \(N_i^{c}\) samples of class \(c\) at client \(i\), computed on the fly from the latest forward passes of local training. As a product of Gaussians is itself Gaussian, both PoEs reduce to class-wise means and covariances: each client uploads \((\mu_i^c, \Sigma_i^c)\) for its classes only, and the server combines them analytically via the precisions,

(22)\[\Sigma^c = \left( I + \sum_{i \in \mathcal{P}_t} \left( \Sigma_i^c \right)^{-1} \right)^{-1}, \qquad \mu^c = \Sigma^c \sum_{i \in \mathcal{P}_t} \mu_i^c \left( \Sigma_i^c \right)^{-1},\]

besides averaging the extractors, \(w^f = \frac{1}{\lvert \mathcal{P}_t \rvert} \sum_{i \in \mathcal{P}_t} w_i^f\). Classes missed by all active clients keep the statistics of the previous round. The local objective of client \(i\) is its empirical loss plus the alignment term above,

(23)\[L_i = f_i(w_i) + \beta \, \mathbb{E}_{p(x_i, y_i)} \left[ \mathrm{KL}\left[ p(z^{c = y_i} \mid x) \,\middle\|\, p(z \mid x_i) \right] \right],\]

optimized with the re-parameterization trick for the stochastic features. Since the PoE posterior can be sharper than any single expert, the common representation carries not only the common but also the complementary information across clients; the alignment term doubles as noise injection into the local features, which improves generalization and uncertainty calibration. Note that for \(\beta = 0\) and \(\Sigma(x_i) = 0\) (deterministic features), the method degenerates to FedPer. The pseudocode is summarized below.

\begin{algorithm}
\caption{PseudoCode for \texttt{FedCR}}
\begin{algorithmic}
\REQUIRE number of clients $m,$ number of rounds $T,$ local steps $K,$ learning rate $\eta,$ CMI coefficient $\beta$; class-wise initialization $\mu_0^c = 0,$ $\Sigma_0^c = I$ for all classes $c$

\FOR{each round $t = 1, 2, \cdots, T$}
\STATE the server samples a subset of active clients $\mathcal{P}_t$ and sends $w^f$, $\mu_t^c$ and $\Sigma_t^c$ ($\forall c \in \mathcal{C}$) to each $i \in \mathcal{P}_t$
\STATE \FOR{each client $i \in \mathcal{P}_t$ {in parallel}}
\STATE initialize the local model as $w_{i,0} = [w^f, w_i^p]$
\FOR{$k = 1, 2, \cdots, K$}
\STATE $\hspace{1.3em}$ update the whole network by SGD on the local objective:
\STATE $\hspace{1.3em}$ $w_{i,k} = w_{i,k-1} - \eta \nabla L_i$, where
\STATE $\hspace{2.3em}$ $L_i \approx \frac{1}{N_i} \sum\limits_{n=1}^{N_i} \left[ -\log \hat{p}\left( y_i^{(n)} \,\middle|\, z^{(n)} \right) + \beta \, \mathrm{KL}\left[ p(z^{c = y_i^{(n)}} \mid x) \,\middle\|\, p(z \mid x_i^{(n)}) \right] \right]$
\STATE \ENDFOR
\STATE aggregate the class-wise local features via the local PoE: $\mu_i^c, \Sigma_i^c$ for $c \in \mathcal{C}_i$
\STATE keep the personal head $w_i^p = w_{i,K}^p$
\STATE send $w_{i,K}^f$ and $\mu_i^c, \Sigma_i^c$ ($\forall c \in \mathcal{C}_i$) to the server
\STATE \ENDFOR
\STATE {Server Update:}
\STATE average the feature extractors: $w^f = \frac{1}{\lvert \mathcal{P}_t \rvert} \sum\limits_{i \in \mathcal{P}_t} w_{i,K}^f$
\STATE for classes $c$ contained in no active client: keep $\mu_t^c = \mu_{t-1}^c$, $\Sigma_t^c = \Sigma_{t-1}^c$
\STATE aggregate the global class-wise features via the PoE with prior $p(z) = \mathcal{N}(0, 1)$, i.e. for $c$ contained in some active client:
\STATE $\hspace{1.3em}$ $\Sigma_t^c = \left( I + \sum\limits_{i \in \mathcal{P}_t} \left( \Sigma_i^c \right)^{-1} \right)^{-1}$, \quad $\mu_t^c = \Sigma_t^c \sum\limits_{i \in \mathcal{P}_t} \mu_i^c \left( \Sigma_i^c \right)^{-1}$
\STATE send $w^f$, $\mu_t^c$ and $\Sigma_t^c$ ($\forall c \in \mathcal{C}$) to the clients
\ENDFOR
\end{algorithmic}
\end{algorithm}

References

[1] (1,2)

Filip Hanzely and Peter Richtárik. Federated Learning of a Mixture of Global and Local Models. arXiv preprint arXiv:2002.05516, 2020.

[2]

Dmitry Kovalev, Samuel Horváth, and Peter Richtárik. Don’t Jump Through Hoops and Remove Those Loops: SVRG and Katyusha are Better Without the Outer Loop. In Aryeh Kontorovich and Gergely Neu, editors, Proceedings of the 31st International Conference on Algorithmic Learning Theory, volume 117 of Proceedings of Machine Learning Research, 451–467. PMLR, 2 2020.

[3]

Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. Personalized Federated Learning with Moreau Envelopes. In Proceedings of the 34th International Conference on Neural Information Processing Systems. Red Hook, NY, USA, 2020. Curran Associates Inc.

[4]

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and Robust Federated Learning Through Personalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, 6357–6368. PMLR, 7 2021.

[5]

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive Personalized Federated Learning. arXiv preprint arXiv:2003.13461v3, 3 2020. doi:10.48550/ARXIV.2003.13461.

[6]

Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated Learning Based on Dynamic Regularization. In International Conference on Learning Representations. 2021.

[7]

Yinchuan Li, Xiaofeng Liu, Xu Zhang, Yunfeng Shao, Qing Wang, and Yanhui Geng. Personalized Federated Learning via Maximizing Correlation with Sparse and Hierarchical Extensions. arXiv preprint arXiv:2107.05330, 2021.