Personalized Federated Learning: A Unified Framework and Universal Optimization Techniques

We outline the main contributions of this work.

Single personalized FL objective. We propose a single objective (1) capable of recovering many existing convex personalized FL approaches by carefully constructing the local loss $f_{m}(w,\beta_{m})$. Consequently, training different personalized FL models is equivalent to solving a particular instance of (1).

Recovering best-known complexity and novel guarantees. We develop algorithms for solving (1) and prove sharp convergence rates for strongly convex objectives. Specializing our rates from the general setting to the individual personalized FL objectives, we recover the best-known optimization guarantees from the literature or advance beyond the state-of-the-art with a single exception: objective (11) with $\lambda>L^{\prime}$. Therefore, our results often render optimization tailored to solve a specific personalized FL unnecessary in many cases.

Universal (convex) optimization methods and theory for personalized FL. To develop an optimization theory for solving (1), we impose particular assumptions on the objective: $\mu-$strong convexity of $F$ and convexity and $(L^{w},ML^{\beta})$-smoothness of $f_{m}$ for all $m\in[M]$ (see Assumptions 1, 2). These assumptions are naturally satisfied for the vast majority of personalized FL objectives in the literature, with the exception of personalized FL approaches that are inherently nonconvex, such as MAML (Finn et al., 2017). Under these assumptions, we propose three algorithms for solving the general personalized FL objective (1): i) Local Stochastic Gradient Descent for Personalized FL (LSGD-PFL), ii) Accelerated Block Coordinate Descent for Personalized FL (ACD-PFL), and iii) Accelerated Stochastic Variance Reduced Coordinate Descent for Personalized FL (ASVRCD-PFL). The convergence rates of these methods are summarized in Table 1. We emphasize that these optimizers can be used to solve many (convex) personalized FL objectives from the literature by casting a given objective as a special case of (1), oftentimes matching or outperforming algorithms originally designed for the particular scenario.

Minimax optimal rates. We provide lower complexity bounds for solving (1). Using the construction of Hendrikx et al. (2021), we show that to solve (1), one requires at least ${\cal O}\left(\sqrt{{L^{w}}/{\mu}}\log\epsilon^{-1}\right)$ communication rounds. Note that communication is often the bottleneck when training distributed and personalized FL models. Furthermore, one needs at least ${\cal O}\left(\sqrt{{L^{w}}/{\mu}}\log\epsilon^{-1}\right)$ evaluations of $\nabla_{w}F$ and at least ${\cal O}\left(\sqrt{{L^{\beta}}/{\mu}}\log\epsilon^{-1}\right)$ evaluations of $\nabla_{\beta}F$. Given the $n$-finite sum structure of $f_{m}$ with $({\cal L}^{w},M{\cal L}^{\beta})$-smooth components, we show that one requires at least ${\cal O}\left(n+\sqrt{{n{\cal L}^{w}}/{\mu}}\log{\epsilon^{-1}}\right)$ stochastic gradient evaluations with respect to $w$-parameters and at least ${\cal O}\left(n+\sqrt{{n{\cal L}^{\beta}}/{\mu}}\log{\epsilon^{-1}}\right)$ stochastic gradient evaluations with respect to $\beta$-parameters. We show that ACD-PFL is always optimal in terms of communication and local computation when the full gradients are available, while ASVRCD-PFL can be optimal either in terms of the number of evaluations of the $w$-stochastic gradient or the $\beta$-stochastic gradient. However, note that ASVRCD-PFL cannot achieve optimal rate for both evaluations of the $w$-stochastic gradient and the $\beta$-stochastic gradient simultaneously, which we leave for future research.

Personalization and communication complexity. Given that a specific FL objective contains a parameter that determines the amount of personalization, we observe that the value of $\sqrt{L^{w}/\mu}$ is always a non-increasing function of this parameter. Since the communication complexity of (1) is equal to $\sqrt{L^{w}/\mu}$ up to constant and log factors, we conclude that personalization has a positive effect on the communication complexity of training an FL model.

New personalized FL objectives. The universal personalized FL objective (1) enables us to obtain a range of novel personalized FL formulations as special cases. While we study various (parametric) extensions of known models, we believe that the objective (1) can lead to easier development of brand new objectives as well. However, we stress that proposing novel personalized FL models is not the main focus of our work; the paper’s primary focus is on providing universal optimization guarantees for (convex) personalized FL.

Despite the aforementioned benefits of our proposed unified framework, we acknowledge that this is neither the only nor the universally best approach for personalized federated learning. However, providing a general framework that can include many existing methods as special cases can help us gain a clear understanding and motivate us to propose new personalized methods.

Complexity Notations. For two sequences $\{a_{n}\}$ and $\{b_{n}\}$, $a_{n}={\cal O}(b_{n})$ if there exists $C>0$ such that $|a_{n}/b_{n}|\leq C$ for all $n$ large enough; $a_{n}=\Theta(b_{n})$ if $a_{n}={\cal O}(b_{n})$ and $b_{n}={\cal O}(a_{n})$ simultaneously. Similarly, $a_{n}=\tilde{\cal O}(b_{n})$ if $a_{n}={\cal O}(b_{n}\log^{k}b_{n})$ for some $k\geq 0$; $a_{n}=\tilde{\Theta}(b_{n})$ if $a_{n}=\tilde{\cal O}(b_{n})$ and $b_{n}=\tilde{\cal O}(a_{n})$ simultaneously.

Local Objective. We assume three different ways to access the gradient of the local objective $f_{m}$. The first, and the simplest case, corresponds to having access to the full gradient of $f_{m}$ with respect to either $w$ or $\beta_{m}$ for all $m\in[M]$ simultaneously. The second case corresponds to a situation where $\nabla f_{m}(w,\beta_{m})$ is the expectation itself, i.e.,

while having access to stochastic gradients with respect to either $w$ or $\beta_{m}$ simultaneously for all $m\in M$, where $\hat{f}_{m}$ represents the loss function on a single data point. The third case corresponds to a finite sum $f_{m}$:

having access to $\nabla_{w}f_{m,i}(w,\beta_{m})$ or to $\nabla_{\beta}f_{m,i}(w,\beta_{m})$ for all $m\in[M]$ and $i\in[n]$ selected uniformly at random.

Assumptions. We argue that the objective (1) is capable of recovering virtually any (convex) personalized FL objective. Since the structure of the individual personalized FL objectives varies greatly, it is important to impose reasonable assumptions on the problem (1) in order to obtain meaningful rates in the special cases.

When $f_{m}$ is a finite sum (3), we require the smoothness of the finite sum components.

In Section 2 we justify Assumptions 1 and 2 and characterize the constants $\mu$, $L^{w}$, $L^{\beta}$, ${\cal L}^{w}$, ${\cal L}^{\beta}$ for special cases of (1). Table 3 summarizes these parameters.

Price of generality. Since Assumption 1 is the only structural assumption we impose on (1), one cannot hope to recover the minimax optimal rates, that is, the rates that match the lower complexity bounds, for all individual personalized FL objectives as a special case of our general guarantees. Note that any given instance of (1) has a structure that is not covered by Assumption 1, but can be exploited by an optimization algorithm to improve either communication or local computation. Therefore, our convergence guarantees are optimal in light of Assumption 1 only. Despite this, our general rates specialize surprisingly well as we show in Section 2: our complexities are state-of-the-art in all scenarios with a single exception: the communication/computation complexity of (11).

Individual treatment of $w$ and $\beta$. Throughout this work, we allow different smoothness of the objective with respect to global parameters $w$ and local parameters $\beta$. At the same time, our algorithm is allowed to exploit the separate access to gradients with respect to $w$ and $\beta$, given that these gradients can be efficiently computed separately. Without such a distinction, one might not hope for the communication complexity better than $\Theta\left(\max\{L^{w},L^{\beta}\}/\mu\log\epsilon^{-1}\right)$, which is suboptimal in the special cases. Similarly, the computational guarantees would be suboptimal as well. See Section 2 for more details.

Data heterogeneity. While the convergence rate of LSGD-PFL relies on data heterogeneity (See Theorem 1), we allow for an arbitrary dissimilarity among the individual clients for analyzing ACD-PFL and ASVRCD-PFL (see Theorem 7 and Theorem 8). Our experimental results also support that ASCD-PFL (ACD-PFL with stochastic gradient to reduce computation) and ASVRCD-PFL are more robust to data heterogeneity compared to the widely used Local SGD.

The rest of the paper is organized as follows. In Section 2, we show how (1) can be used to recover various personalized federated learning objectives in the literature. In Section 3, we propose a local-SGD based algorithm, LSGD-PFL, for solving (1). We further establish computational upper bounds for LSGD-PFL in strongly convex, weakly convex, and nonconvex cases. In Section 4, we discuss the minimax optimal algorithms for solving (1). We first show the minimax lower bounds in terms of the number of communication rounds, number of evaluations of the gradient of global parameters, and number of evaluations of the gradient of local parameters, respectively. We subsequently propose two coordinate-descent based algorithms, ACD-PFL and ASVRCD-PFL, which can match the lower bounds. In Section 5 and Section 6, we use experiments on synthetic and real data to illustrate the performance of the proposed algorithms and empirically validate the theorems. Finally, we conclude the paper with Section7. Technical proofs are deferred to the Appendix.

The traditional, non-personalized FL objective (McMahan et al., 2017) is given as

where $f^{\prime}_{m}$ corresponds to the loss on the $m$-th client’s data. Assume that $f^{\prime}_{m}$ is $L^{\prime}$-smooth and $\mu^{\prime}-$ strongly convex for all $m\in[M]$. The minimax optimal communication to solve (4) up to $\epsilon$-neighborhood of the optimum is $\tilde{\Theta}\left(\sqrt{{L^{\prime}}/{\mu^{\prime}}}\log{\epsilon^{-1}}\right)$ (Scaman et al., 2018). When $f_{m}^{\prime}=\frac{1}{n}\sum_{j=1}^{n}f_{m,j}^{\prime}(w)$ is an $n$-finite sum with convex and ${\cal L}^{\prime}-$smooth components, the minimax optimal local stochastic gradient complexity is $\tilde{\Theta}\left(\left(n+\sqrt{{n{\cal L}^{\prime}}/{\mu}}\right)\log{\epsilon^{-1}}\right)$ (Hendrikx et al., 2021). The FL objective (4) is a special case of (1) with $d_{1}=\dots=d_{M}=0$ and our theory recovers the aforementioned rates.

At the other end of the spectrum lies the fully personalized FL where the $m$-th client trains their own model without any influence from other clients:

The above objective is a special case of (1) with $d_{0}=0$. As the objective is separable in $\beta_{1},\dots,\beta_{M}$, we do not require any communication to train it. At the same time, we need $\tilde{\Theta}\left(\left(n+\sqrt{{n{\cal L}^{\prime}}/{\mu}}\right)\log{\epsilon^{-1}}\right)$ local stochastic oracle calls to solve it (Lan & Zhou, 2018) – which is what our algorithms achieve.

The objective is given as

where $(w^{\prime})^{*}$ is a solution of the traditional FL in (4) and $\lambda\geq 0$. Assuming that $(w^{\prime})^{*}$ is known (which Li et al. (2020) does), the problem (6) is a particular instance of (5); thus our approach achieves the optimal complexity.

A more challenging objective (in terms of optimization) is the following relaxed version of (6):

where $\Lambda\geq 0$ is the relaxation parameter, recovering the original objective for $\Lambda\rightarrow\infty$. Note that, since $\Lambda\rightarrow\infty$, finding a minimax optimal method for the optimization of (6) is straightforward. First, one has to compute a minimizer $(w^{\prime})^{*}$ of the classical FL objective (4), which can be done with a minimax optimal complexity. Next, one needs to compute the local solution $\beta_{m}^{*}=\operatorname*{arg\,min}_{\beta_{m}\in\mathbb{R}^{d}}f^{\prime}_{m}(\beta_{m})+\lambda\|w^{*}-\beta_{m}\|^{2}$, which only depends on the local data and thus can also be optimized with a minimax optimal algorithm.

A more interesting scenario is obtained when we do not set $\Lambda\rightarrow\infty$ in (7), but rather consider a finite $\Lambda>0$ that is sufficiently large. To obtain the right smoothness/strong convexity parameter (according to Assumption 1), we scale the global parameter $w$ by a factor of $M^{-\frac{1}{2}}$ and arrive at the following objective:

where

The next lemma determines parameters $\mu,L^{w},L^{\beta},{\cal L}^{w},{\cal L}^{\beta}$ in Assumption 1. See the proof in Appendix B.1.

Evaluating gradients. Note that evaluating $\nabla_{w}f_{m}(x,\beta_{m})$ under the objective (8) can be perfectly decoupled from evaluating $\nabla_{\beta}f_{m}(x,\beta_{m})$. Therefore, we can make full use of our theory and take advantage of different complexities w.r.t. $\nabla_{w}$ and $\nabla_{\beta}$. Resulting communication and computation complexities for solving (8) are presented in Table 3.

In its simplest form, the multi-task personalized objective (Smith et al., 2017; Wang et al., 2018) is given as

where $\bar{\beta}\coloneqq\frac{1}{M}\sum_{m=1}^{M}\beta_{m}$ and $\lambda\geq 0$ (Hanzely & Richtárik, 2020). On the other hand, the goal of implicit MAML (Rajeswaran et al., 2019; Dinh et al., 2020) is to minimize

By reparametrizing (1), we can recover an objective that is simultaneously equivalent to both (9) and (10). In particular, by setting

the objective (1) becomes

It is a simple exercise to notice the equivalence of (11) to both (9) and (10).²²2To the best of our knowledge, we are the first to notice the equivalence of (9) and (10). Indeed, we can always minimize (11) in $w$, arriving at $w^{*}=M^{\frac{1}{2}}\bar{\beta}$, and thus recovering the solution of (9). Similarly, by minimizing (11) in $\beta$ we arrive at (10).

Next, we establish the parameters in Assumptions 1 and 2.

The proof is given in Appendix B.2. Hanzely et al. (2020a) showed that the minimax optimal communication complexity to solve (9) (and therefore to solve (10) and (11)) is $\Theta\left(\sqrt{{\min(L^{\prime},\lambda)}/{\mu^{\prime}}}\log\epsilon^{-1}\right)$. Furthermore, they showed that the minimax optimal number of gradients w.r.t. $f^{\prime}$ is $\tilde{\Theta}\left(\left(\sqrt{{L^{\prime}}/{\mu^{\prime}}}\right)\log\epsilon^{-1}\right)$ and proposed a method that has the complexity $\Theta\left(\left(n+\sqrt{{n({\cal L}^{\prime}+\lambda)}/{\mu^{\prime}}}\right)\log\epsilon^{-1}\right)$ w.r.t. the number of $f^{\prime}_{m,j}$-gradients. We match the aforementioned communication guarantees when $\lambda={\cal O}(L^{\prime})$ and computation guarantees when $L^{\prime}={\cal O}(\lambda)$. Furthermore, when $\lambda={\cal O}(L^{\prime})$, our complexity guarantees are strictly better compared to the guarantees for solving the implicit MAML objective (10) directly (Rajeswaran et al., 2019; Dinh et al., 2020).

The objective is given as

where $(w^{\prime})^{*}=\operatorname*{arg\,min}_{w\in\mathbb{R}^{d}}F^{\prime}(w)$ is a solution to (4) and $0<\alpha_{1},\dots\alpha_{M}<1$. Assuming that $(w^{\prime})^{*}$ is known, as was done in Deng et al. (2020), the problem (12) is an instance of (5); thus our approach achieves the optimal complexity.

A more interesting case (in terms of optimization) is when considering a relaxed variant of (12), given as

where $\Lambda\geq 0$ is the relaxation parameter that allows recovering the original objective when $\Lambda\rightarrow\infty$. Such a choice, alongside with the usual rescaling of the parameter $w$ results in the following objective:

where

The proof is given in Appendix B.3.

The most typical example of the weight sharing setting is when parameters $w,\beta$ correspond to different layers of the same neural network. For example, $\beta_{1},\dots,\beta_{M}$ could be the weights of first few layers of a neural network, while $w$ are the weights of the remaining layers (Liang et al., 2020). Or, alternatively, each of $\beta_{1},\dots,\beta_{M}$ can correspond to the weights of last few layers, while the remaining weights are included in the global parameter $w$ (Arivazhagan et al., 2019). Overall, we can write the objective as follows:

where $d_{w}+d_{\beta}=d$. Using an equivalent reparameterization of the $w-$space, we aim to minimize

which is an instance of (1) with $f_{m}(w,\beta_{m})=f^{\prime}_{m}([M^{-\frac{1}{2}}w,\beta_{m}])$.

The proof is straightforward and, therefore, omitted. A distinctive characteristic of the explicit weight sharing paradigm is that evaluating a gradient w.r.t. $w$-parameters automatically grants either free or highly cost-effective access to the gradient w.r.t. $\beta$-parameters (and vice versa).

Agarwal et al. (2020) proposed federated residual learning:

where $(x^{w}_{m},x^{\beta}_{m})$ is a local feature vector (there may be an overlap between $x^{w}_{m}$ and $x^{\beta}_{m}$), $A^{w}(w,x^{w}_{m})$ represents the model prediction using global parameters/features, $A^{\beta}(\beta,x^{\beta}_{m})$ denotes the model prediction using local parameters/features, and $l(\cdot,\cdot)$ is a loss function. Clearly, we can recover (17) with

Unlike the other objectives, here we cannot relate constants $\mu^{\prime},L^{\prime},{\cal L}^{\prime}$ to $F_{R}$, since we do not write $f_{m}$ as a function of $f^{\prime}m$. However, it seems natural to assume $L^{w}_{R}$ (or $L^{\beta}_{R}$)-smoothness of $l(A^{w}(w,x^{w}{m}),a^{\beta}_{m})$ (or $l(a^{w}_{m},A^{\beta}(\beta_{m},x^{\beta}_{m}))$) as a function of $w$ (or $\beta$) for any $a^{\beta}_{m},x^{\beta}_{m},a^{w}_{m},x^{w}_{m}$. Let us define ${\cal L}^{w}_{R}$, ${\cal L}^{\beta}_{R}$ analogously, given that $l$ has an $n$-finite sum structure. Assuming, furthermore, that $F$ is $\mu$-strongly convex and $f_{m}$ is convex (for each $m\in M$), we can apply our theory.

Meta-learning has recently been employed for personalization (Chen et al., 2018; Khodak et al., 2019; Jiang et al., 2019; Fallah et al., 2020; Lin et al., 2020). Notably, the model-agnostic meta-learning (MAML) (Finn et al., 2017) based personalized FL objective is given as

Although we can recover (19) as a special case of (1) by setting $f_{m}(w,\beta_{m})=f^{\prime}_{m}(w-\alpha\nabla f^{\prime}_{m}(w))$, our (convex) convergence theory does not apply due to the inherent non-convex structure of (19). Specifically, objective $F_{MAML}$ is non-convex even if function $f^{\prime}_{m}$ is convex. In this scenario, only our non-convex rates of Local SGD apply.