Personalized Federated Learning with Moreau Envelopes

Canh T. DinhNguyen H. TranTuan Dung Nguyen

article2020NeurIPS1,509 citations

Introduces pFedMe, a personalized federated learning algorithm that uses Moreau envelopes to decouple client-specific models from global aggregation, delivering proven convergence speedups and superior accuracy over standard meta-learning baselines.

Listen

Federated learning is an increasingly important technique for training artificial intelligence models across distributed user devices without centralizing private personal data. In real-world deployments such as mobile applications, healthcare systems, and financial platforms, a fundamental bottleneck is statistical diversity: each participant generates distinct, non-uniform data. As a result, standard global models perform poorly when deployed locally on individual client devices, while training strictly isolated local models fails due to insufficient client data.

The article evaluates a personalized federated learning algorithm termed pFedMe, which uses Moreau envelope mathematical regularization to simultaneously build a high-quality central reference model while optimizing tailored personalized models for every participating client. The primary objective is to demonstrate that this bi-level optimization approach outperforms standard and meta-learning-based federated frameworks in convergence speed and localized task accuracy.

To evaluate this framework, the authors conducted mathematical convergence analyses across strongly convex and smooth non-convex objective spaces, alongside simulated experiments on non-uniform real (MNIST digit classification distributed across 20 nodes) and synthetic benchmarks (60-dimensional classification across 100 nodes). The method was directly compared against standard Federated Averaging (FedAvg) and a leading meta-learning personalization algorithm (Per-FedAvg).

The article establishes several key findings. First, the personalized models generated by the proposed framework achieved the highest overall accuracy across both benchmarks: reaching 95.62% in linear classification and 99.46% in neural networks on the real dataset, and up to 86.36% on the synthetic benchmark. Second, under strongly convex objectives, the personalized model outperformed standard FedAvg and Per-FedAvg by 1.5% and 1.3% on real data, and by 5.2% and 3.8% on synthetic data, respectively. Third, theoretical convergence proofs confirmed that the proposed framework achieves state-of-the-art convergence rates—specifically a quadratic speedup for strongly convex models and a two-thirds order sublinear speedup for non-convex models—outperforming conventional linear and square-root rates. Finally, the framework decouples personalized and global optimization, requiring only standard first-order gradient calculations and roughly three to five local approximation steps, thereby avoiding the heavy computational cost of second-order matrix evaluations.

These findings indicate that organizations can deliver highly accurate, tailored user-end machine learning models while retaining strong data privacy and reducing communication overhead between clients and servers. This lowers operational risks and computing costs on edge devices compared to prior meta-learning architectures.

For technical leaders seeking to deploy personalized federated architectures, the authors recommend tuning the regularization parameter carefully to match local data diversity and selecting moderate local iteration steps (around 3 to 5 internal steps) to minimize client energy consumption without sacrificing accuracy. Further pilot testing in live edge computing environments with heterogeneous hardware power and potential network disruptions is recommended before large-scale production adoption.

Key limitations include reliance on bounded variance assumptions, standard dataset simulations rather than live commercial edge deployments, and the sensitivity of the regularization hyperparameter, which must be tuned per dataset to prevent divergence. Nevertheless, the theoretical proofs and empirical validations provide high confidence in the framework's superior balance of personalization, convergence speed, and computational efficiency.

University of Melbourne University of Sydney application servers federated learning meta-learning Personalized Federated Learning Moreau Envelopes pFedMe

arXiv: 2006.08848

Paper: Communication-Efficient Learning of Deep Networks from Decentralized Data, H. B. McMahan et al. (2016). Introduces the foundational FederatedAveraging (FedAvg) algorithm and the decentralized client-server architecture that pFedMe extends to achieve personalization.
Paper: Federated Optimization in Heterogeneous Networks, Tian Li et al. (2018). Introduces proximal regularization (FedProx) to handle statistical and systems heterogeneity in federated optimization, directly preceding pFedMe's Moreau envelope approach.
Paper: Federated Multi-Task Learning, Virginia Smith et al. (2017). Establishes multi-task learning formulations for heterogeneous clients in federated settings, laying the conceptual groundwork for personalized federated learning.
Paper: Federated Learning with Non-IID Data, Yue Zhao et al. (2018). Analyzes the detrimental impact of non-IID client data distributions on global model convergence, motivating the need for personalized regularized objectives.
Paper: Advances and Open Problems in Federated Learning, Peter Kairouz et al. (2019). Provides a comprehensive survey framing non-IID challenges and the necessity of personalization and multi-task techniques in federated learning.
Paper: LEAF: A Benchmark for Federated Settings, Sebastian Caldas et al. (2018). Presents standard benchmark datasets and evaluation suites (LEAF) for measuring optimization performance under heterogeneous federated distributions.

Paper: Model-Contrastive Federated Learning, Qinbin Li et al. (2021). Proposes model-contrastive learning to correct local client drift under severe heterogeneity, building upon the personalized and regularized local training paradigms.
Paper: Federated Learning on Non-IID Data Silos: An Experimental Study, Qinbin Li et al. (2021). Provides a systematic empirical benchmark across diverse non-IID data skew scenarios to evaluate advanced heterogeneous and personalized federated algorithms.

Cover for Personalized Federated Learning with Moreau Envelopes

Abstract

Federated learning (FL) is a decentralized and privacy-preserving machine learning technique in which a group of clients collaborate with a server to learn a global model without sharing clients' data. One challenge associated with FL is statistical diversity among clients, which restricts the global model from delivering good performance on each client's task. To address this, we propose an algorithm for personalized FL (pFedMe) using Moreau envelopes as clients' regularized loss functions, which help decouple personalized model optimization from the global model learning in a bi-level problem stylized for personalized FL. Theoretically, we show that pFedMe's convergence rate is state-of-the-art: achieving quadratic speedup for strongly convex and sublinear speedup of order 2/3 for smooth nonconvex objectives. Experimentally, we verify that pFedMe excels at empirical performance compared with the vanilla FedAvg and Per-FedAvg, a meta-learning based personalized FL algorithm.

Topics

application servers federated learning meta-learning nonconvex optimization privacy preserving

Keywords

Personalized Federated Learning Moreau Envelopes pFedMe Model-Agnostic Meta-Learning Communication-Efficient SGD Multi-Task Learning Client Drift Vanilla FedAvg Local SGD

1 Introduction
2 Related Work
3 Personalized Federated Learning with Moreau Envelopes (pFedMe)
3.1 pFedMe: Problem Formulation
3.2 pFedMe: Algorithm
4 pFedMe: Convergence Analysis
5 Experimental Results and Discussion
5.1 Experimental Settings
5.2 Effect of hyperparameters
5.3 Performance Comparison
6 Conclusion
A Proof of the Results
A.1 Review of useful existing results
A.2 Proof of Lemma
A.3 Proof of Lemma
A.4 Proof of Theorem
A.4.1 Additional notations
A.4.2 Supporting lemmas
A.4.3 Completing the proof of Theorem
A.5 Theorem
References

Knowls

Knowl 1 — Bi-Level Problem Formulation of Personalized Federated Learning with Moreau Envelopes
model/method
In federated learning with $N$ clients having heterogeneous, non-i.i.d. data distributions, standard federated learning seeks a single global model by minimizing $f(w) = \frac{1}{N} \sum_{i=1}^N f_i(w)$ , where $f_i(w) = \mathbb{E}_{\xi_i}[\tilde{f}_i(w; \xi_i)]$ is the expected loss over client $i$ 's local data distribution $\xi_i$ . When data distributions diverge, this shared model often yields high generalization error on individual clients.

Personalized Federated Learning with Moreau Envelopes (pFedMe) decouples the learning of a shared global reference parameter $w \in \mathbb{R}^d$ from client-specific personalized parameters $\theta_i \in \mathbb{R}^d$ through an $\ell_2$ -regularized bi-level optimization objective:

$\min_{w \in \mathbb{R}^d} \left\{ F(w) := \frac{1}{N} \sum_{i=1}^N F_i(w) \right\}, \quad \text{where } F_i(w) := \min_{\theta_i \in \mathbb{R}^d} \left\{ f_i(\theta_i) + \frac{\lambda}{2} \|\theta_i - w\|^2 \right\}$

The regularizer $\lambda \in (0, \infty)$ controls the strength of the constraint keeping personalized models close to the global reference point $w$ . The inner problem $F_i(w)$ is the Moreau envelope of client loss $f_i$ . The unique optimal personalized model for client $i$ given global model $w$ is defined by the proximal operator:

$\hat{ heta}_i(w) := \text{prox}_{f_i/\lambda}(w) = \arg\min_{\theta_i \in \mathbb{R}^d} \left\{ f_i(\theta_i) + \frac{\lambda}{2} \|\theta_i - w\|^2 \right\}$

If $f_i$ is convex (or nonconvex and $L$ -smooth with $\lambda > 2L$ ), the Moreau envelope $F_i$ is $L_F$ -smooth with $L_F = \lambda$ , and its gradient is given in closed form by:

$\nabla F_i(w) = \lambda (w - \hat{ heta}_i(w))$

Furthermore, if $f_i$ is $\mu$ -strongly convex, $F_i$ is $\mu_F$ -strongly convex with parameter $\mu_F = \frac{\lambda \mu}{\lambda + \mu}$ .
Knowl 2 — pFedMe Personalized Federated Optimization Algorithm
algorithm
The pFedMe algorithm alternates between solving an inner client-level proximal subproblem to compute an approximate personalized model and taking gradient descent steps with respect to the client's Moreau envelope $F_i$ . At each communication round $t$ , the central server broadcasts the current global model $w_t$ . Each client initializes its local outer variable $w_{i,0}^t = w_t$ and executes $R$ local iterations. At local step $r$ , the client draws a mini-batch $\mathcal{D}_i$ of size $|\mathcal{D}|$ and runs $K$ steps of gradient descent to find a $\delta$ -approximate minimizer $\tilde{ heta}_i(w_{i,r}^t)$ of the regularized sample loss $\tilde{h}_i(\theta_i; w_{i,r}^t, \mathcal{D}_i) := \tilde{f}_i(\theta_i; \mathcal{D}_i) + \frac{\lambda}{2}\|\theta_i - w_{i,r}^t\|^2$ such that $\|\nabla \tilde{h}_i(\tilde{ heta}_i; w_{i,r}^t, \mathcal{D}_i)\|^2 \le \nu$ . The local model is then updated via the approximate Moreau gradient: $w_{i,r+1}^t = w_{i,r}^t - \eta \lambda (w_{i,r}^t - \tilde{\theta}_i(w_{i,r}^t))$ . Finally, a uniformly chosen subset $\mathcal{S}^t$ of $S$ clients transmits their local endpoints $w_{i,R}^t$ to the server, which performs a parameterized global update with factor $\beta \ge 1$ .
Input: Communication rounds $T$ , local rounds $R$ , sample size $S$ , regularization $\lambda$ , client learning rate $\eta$ , server update parameter $\beta$ , initial global model $w_0$
Output: Final global model $w_T$ and personalized models $\tilde{\theta}_i(w_{i,R}^{T-1})$
for $t = 0$ to $T - 1$ do
    Server transmits $w_t$ to all $N$ clients
    for each client $i = 1$ to $N$ in parallel do
         $w_{i,0}^t = w_t$
        for $r = 0$ to $R - 1$ do
            Sample fresh mini-batch $\mathcal{D}_i$ with size $|\mathcal{D}|$
            Compute $\delta$ -approximate solution $\tilde{\theta}_i(w_{i,r}^t)$ minimizing $\tilde{h}_i(\theta_i; w_{i,r}^t, \mathcal{D}_i) = \tilde{f}_i(\theta_i; \mathcal{D}_i) + \frac{\lambda}{2}\|\theta_i - w_{i,r}^t\|^2$ to tolerance $\|\nabla \tilde{h}_i(\tilde{\theta}_i; w_{i,r}^t, \mathcal{D}_i)\|^2 \le \nu$
             $w_{i,r+1}^t = w_{i,r}^t - \eta \lambda (w_{i,r}^t - \tilde{\theta}_i(w_{i,r}^t))$
        end for
    end for
    Server uniformly samples client subset $\mathcal{S}^t$ of size $S$
    Sampled clients $i \in \mathcal{S}^t$ send $w_{i,R}^t$ to the server
    Server updates global model: $w_{t+1} = (1 - \beta)w_t + \beta \frac{1}{S} \sum_{i \in \mathcal{S}^t} w_{i,R}^t$
end for
Knowl 3 — Standard Smoothness, Variance, and Diversity Assumptions for pFedMe Analysis
assumption
The convergence analysis of pFedMe operates under the following conditions on local client loss functions $f_i: \mathbb{R}^d \to \mathbb{R}$ , stochastic loss estimators $\tilde{f}_i(\cdot; \xi_i)$ , and client heterogeneity:
1. Convexity and Smoothness: For all $i \in \{1, \dots, N\}$ and all $w, w' \in \mathbb{R}^d$ , $f_i$ satisfies either:
  - (a) $\mu$ -strong convexity: $f_i(w) \ge f_i(w') + \langle \nabla f_i(w'), w - w' \rangle + \frac{\mu}{2} \|w - w'\|^2$ , where $\mu > 0$ ; or
  - (b) Nonconvex $L$ -smoothness: $\|\nabla f_i(w) - \nabla f_i(w')\| \le L \|w - w'\|$ , with $L < \infty$ .
2. Bounded Stochastic Gradient Variance: The variance of the unbiased local sample gradient $\nabla \tilde{f}_i(w; \xi_i)$ relative to full local gradient $\nabla f_i(w)$ is uniformly bounded across all clients and parameter values: $\mathbb{E}_{\xi_i}\left[\|\nabla \tilde{f}_i(w; \xi_i) - \nabla f_i(w)\|^2\right] \le \gamma_f^2, \quad \forall i, \forall w$
3. Bounded Client Gradient Diversity: The variance between local gradients and the average global gradient across clients is bounded by a constant $\sigma_f \ge 0$ : $\|\nabla f_i(w) - \nabla f(w)\| \le \sigma_f, \quad \forall i \in \{1, \dots, N\}, \forall w \in \mathbb{R}^d$
The analysis explicitly avoids assuming uniformly bounded gradients ( $\|\nabla f_i(w)\| \le G$ ), which is known to fail in unconstrained strongly convex minimization.
Knowl 4 — Approximation Error Bound for the Inner Personalized Model Subproblem
theoretical result
Let $\hat{\theta}_i(w_{i,r}^t) = \arg\min_{\theta_i \in \mathbb{R}^d} \left\{ f_i(\theta_i) + \frac{\lambda}{2}\|\theta_i - w_{i,r}^t\|^2 \right\}$ be the exact proximal point of client $i$ at round $t$ and local step $r$ . Let $\tilde{\theta}_i(w_{i,r}^t)$ be a numerical approximation obtained on mini-batch $\mathcal{D}_i$ of size $|\mathcal{D}|$ satisfying $\|\nabla \tilde{h}_i(\tilde{\theta}_i; w_{i,r}^t, \mathcal{D}_i)\|^2 \le \nu$ , where $\tilde{h}_i(\theta_i; w_{i,r}^t, \mathcal{D}_i) = \tilde{f}_i(\theta_i; \mathcal{D}_i) + \frac{\lambda}{2}\|\theta_i - w_{i,r}^t\|^2$ .

Under bounded stochastic gradient variance $\gamma_f^2$ , the expected mean squared error between the approximate and true personalized models satisfies:

$\mathbb{E}\left[\|\tilde{\theta}_i(w_{i,r}^t) - \hat{\theta}_i(w_{i,r}^t)\|^2\right] \le \delta^2 := \begin{cases} \frac{2}{(\lambda + \mu)^2} \left(\frac{\gamma_f^2}{|\mathcal{D}|} + \nu\right), & \text{if } f_i \text{ is } \mu\text{-strongly convex}, \\ \frac{2}{(\lambda - L)^2} \left(\frac{\gamma_f^2}{|\mathcal{D}|} + \nu\right), & \text{if } f_i \text{ is } L\text{-smooth and } \lambda > L. \end{cases}$

This bound guarantees that the inner personalization error $\delta$ can be made arbitrarily small by increasing the mini-batch size $|\mathcal{D}|$ and lowering the optimization tolerance $\nu$ .
Knowl 5 — Gradient Diversity Bounds for Client Moreau Envelopes
theoretical result
The gradient diversity of the client Moreau envelope functions $F_i(w)$ with respect to the global objective $F(w) = \frac{1}{N}\sum_{i=1}^N F_i(w)$ is bounded as follows:
1. Strongly Convex Case: If each $f_i$ is $\mu$ -strongly convex, then for any $w \in \mathbb{R}^d$ : $\frac{1}{N} \sum_{i=1}^N \|\nabla F_i(w) - \nabla F(w)\|^2 \le 4 L_F (F(w) - F(w^*)) + 2 \sigma_{F,1}^2$ where $w^* = \arg\min_w F(w)$ , $L_F = \lambda$ , and $\sigma_{F,1}^2 := \frac{1}{N} \sum_{i=1}^N \|\nabla F_i(w^*)\|^2$ .
2. Nonconvex Smooth Case: If each $f_i$ is $L$ -smooth with bounded local gradient diversity $\sigma_f$ , and the regularizer satisfies $\lambda > 2\sqrt{2}L$ , then for any $w \in \mathbb{R}^d$ : $\frac{1}{N} \sum_{i=1}^N \|\nabla F_i(w) - \nabla F(w)\|^2 \le \frac{8L^2}{\lambda^2 - 8L^2} \|\nabla F(w)\|^2 + \sigma_{F,2}^2$ where $\sigma_{F,2}^2 := \frac{2\lambda^2}{\lambda^2 - 8L^2} \sigma_f^2$ .
Both $\sigma_{F,1}^2$ and $\sigma_{F,2}^2$ evaluate to zero when client data distributions are identically distributed (i.i.d.).
Knowl 6 — Convergence Rate and Quadratic Speedup for Strongly Convex Objectives in pFedMe
theoretical result
Under $\mu$ -strong convexity and bounded stochastic gradient variance $\gamma_f^2$ , let $\mu_F = \frac{\lambda \mu}{\lambda + \mu}$ , $L_F = \lambda$ , $\kappa_F = \frac{L_F}{\mu_F}$ , $\Delta_0 = \|w_0 - w^*\|^2$ , $\hat{\eta}_1 = \frac{1}{6 L_F (3 + 128 \kappa_F / \beta)}$ with $\beta \ge 1$ , and local learning rate $\eta \le \frac{\hat{\eta}_1}{\beta R}$ . For $T \ge \frac{2}{\hat{\eta}_1 \mu_F}$ , the weighted global model iterate $\bar{w}_T := \frac{\sum_{t=0}^{T-1} \alpha_t w_t}{\sum_{t=0}^{T-1} \alpha_t}$ with weights $\alpha_t = (1 - \eta \beta R \mu_F / 2)^{-(t+1)}$ satisfies:

$\mathbb{E}\left[F(\bar{w}_T) - F(w^*)\right] \le \mathcal{O}\left(\Delta_0 \mu_F e^{-\hat{\eta}_1 \mu_F T / 2}\right) + \tilde{\mathcal{O}}\left(\frac{(N/S - 1)\sigma_{F,1}^2}{\mu_F T N}\right) + \tilde{\mathcal{O}}\left(\frac{(R\sigma_{F,1}^2 + \delta^2 \lambda^2)\kappa_F}{R(T \beta \mu_F)^2}\right) + \mathcal{O}\left(\frac{\lambda^2 \delta^2}{\mu_F}\right)$

The average personalized model error converges to a ball centered at $w^*$ :

$\frac{1}{N}\sum_{i=1}^N \mathbb{E}\left[\|\tilde{\theta}_i^T(w_T) - w^*\|^2\right] \le \frac{1}{\mu_F} \mathcal{O}\left(\mathbb{E}\left[F(\bar{w}_T) - F(w^*)\right]\right) + \mathcal{O}\left(\frac{\sigma_{F,1}^2}{\lambda^2} + \delta^2\right)$

When all clients participate without sampling ( $S = N$ ), setting $\beta = \Theta(N\sqrt{R})$ yields a quadratic speedup of $\mathcal{O}\left(\frac{1}{(TRN)^2}\right)$ with respect to total local computation rounds $TRN$ (when $\sqrt{N} < T$ ), outperforming the standard linear speedup $\mathcal{O}\left(\frac{1}{TRN}\right)$ of conventional federated optimization.
Knowl 7 — Convergence Rate and Sublinear Speedup for Nonconvex Smooth Objectives in pFedMe
theoretical result
Under $L$ -smoothness, bounded stochastic gradient variance $\gamma_f^2$ , and bounded gradient diversity $\sigma_f$ , assume $\lambda \ge \sqrt{8L^2 + 1}$ , $\beta \ge 1$ , $\hat{\eta}_2 = \frac{1}{75 L_F \lambda^2}$ , and $\eta \le \frac{\hat{\eta}_2}{\beta R}$ . For an iteration index $t^* \in \{0, \dots, T-1\}$ chosen uniformly at random, and initial gap $\Delta_F = F(w_0) - F^*$ , the expected squared gradient norm of the global objective satisfies:

$\mathbb{E}\left[\|\nabla F(w_{t^*})\|^2\right] \le \mathcal{O}\left(\frac{\Delta_F}{\hat{\eta}_2 T} + \frac{\left(\Delta_F L_F \sigma_{F,2}^2(N/S - 1)\right)^{1/2}}{\sqrt{TN}} + \frac{(\Delta_F)^{2/3}\left(R\sigma_{F,2}^2 + \lambda^2 \delta^2\right)^{1/3}}{\beta^{4/3} R^{1/3} T^{2/3}} + \lambda^2 \delta^2\right)$

The average distance of personalized models to the global model satisfies:

$\frac{1}{N}\sum_{i=1}^N \mathbb{E}\left[\|\tilde{\theta}_i^{t^*}(w_{t^*}) - w_{t^*}\|^2\right] \le \mathcal{O}\left(\mathbb{E}\left[\|\nabla F(w_{t^*})\|^2\right]\right) + \mathcal{O}\left(\frac{\sigma_{F,2}^2}{\lambda^2} + \delta^2\right)$

When client sampling error is zero ( $S = N$ ), choosing $\beta = \Theta(N^{1/2} R^{1/4})$ and setting $\Theta(T^{1/3}) = \Theta((NR)^{2/3})$ achieves a sublinear speedup of order $\mathcal{O}\left(\frac{1}{(TRN)^{2/3}}\right)$ , compared to the standard $\mathcal{O}\left(\frac{1}{\sqrt{TRN}}\right)$ rate.
Knowl 8 — Experimental Configuration for Heterogeneous Personalized Federated Learning
experimental setup
Empirical validation was performed on both real and synthetic classification datasets:
1. MNIST Dataset: 70,000 handwritten digit images across 10 classes split randomly into 75% training and 25% testing sets. To construct a non-i.i.d. setting, the dataset is distributed across $N = 20$ clients such that each client possesses instances from only 2 distinct digit labels, with unbalanced local dataset sizes ranging between 1,165 and 3,834 samples. The active client subset size per communication round is fixed to $S = 5$ .
2. Synthetic Dataset: A 10-class classification task with 60-dimensional real-valued input vectors distributed among $N = 100$ clients following a power-law distribution. Heterogeneity parameters are set to $\bar{\alpha} = 0.5$ (model parameter variance) and $\bar{\beta} = 0.5$ (feature distribution variance). Local sample sizes vary across clients from 250 to 25,810 samples. The active client sample size per round is $S = 10$ .
3. Model Architectures:
  - Strongly convex objective: $\ell_2$ -regularized multinomial logistic regression (MLR) with softmax activation and cross-entropy loss.
  - Nonconvex objective: Two-layer deep neural network (DNN) with ReLU hidden layer (dimension 100 for MNIST, 20 for Synthetic) and softmax cross-entropy output.
4. Hyperparameters: Mini-batch size $|\mathcal{D}| = 20$ , local outer iterations $R = 20$ , inner gradient descent iterations $K = 5$ , and total communication rounds $T = 800$ (MNIST) or $T = 600$ (Synthetic).
Knowl 9 — Impact of Hyperparameters on pFedMe Optimization Dynamics
empirical result
Empirical evaluation of pFedMe hyperparameters on MNIST demonstrates the following dynamics:
- Inner computation steps ( $K$ ): Values $K \in \{1, 5, 7\}$ show that a small number of inner gradient steps ( $K = 3$ to $5$ ) is sufficient for client personalized models $\tilde{\theta}_i$ to approximate the proximal operator accurately. Increasing $K$ beyond 5 provides negligible convergence or accuracy improvements while increasing client energy consumption.
- Local computation rounds ( $R$ ): Higher values ( $R \in \{10, 20, 30\}$ ) accelerate convergence per communication round for both global and personalized models, demonstrating a direct trade-off between local computation and communication frequency ( $R = 20$ was chosen as balanced).
- Regularization parameter ( $\lambda$ ): Larger $\lambda$ encourages faster initial convergence by anchoring local updates near the global model. However, excessively large $\lambda$ leads to optimization instability and divergence, requiring dataset-dependent tuning ( $\lambda = 15$ for MNIST MLR, $\lambda = 30$ for MNIST DNN).
- Server update step ( $\beta$ ): Setting $\beta > 1$ ( $\beta \in \{1.0, 2.0, 4.0\}$ ) accelerates global and personalized model convergence compared to vanilla FedAvg model averaging ( $\beta = 1$ ). To maintain stability at larger $\beta$ , the client learning rate $\eta$ must be scaled down in inverse proportion.

Knowl 10 — Performance Comparison of pFedMe against FedAvg and Per-FedAvg

data/table

Under fine-tuned hyperparameters with mini-batch size $|\mathcal{D}| = 20$ , local steps $R = 20$ , inner steps $K = 5$ , and server update rate $\beta = 2$ ( $T = 800$ rounds for MNIST, $T = 600$ for Synthetic), the test accuracy (mean $\pm$ standard deviation) was evaluated across models. For Per-FedAvg, $(\hat{\alpha}, \hat{\beta})$ denotes its local and global learning rates. pFedMe evaluates both its Global Model (pFedMe-GM) on $w$ and its Personalized Model (pFedMe-PM) on $\tilde{\theta}_i$ .

Algorithm	Model	MNIST			Synthetic
		$\lambda$	$\eta$ or $(\hat{\alpha}, \hat{\beta})$	Accuracy (%)	$\lambda$	$\eta$ or $(\hat{\alpha}, \hat{\beta})$	Accuracy (%)
FedAvg	MLR	–	0.02	$93.96 \pm 0.02$	–	0.02	$77.62 \pm 0.11$
Per-FedAvg	MLR	–	$(0.03, 0.003)$	$94.37 \pm 0.04$	–	$(0.02, 0.002)$	$81.49 \pm 0.09$
pFedMe-GM	MLR	15	0.01	$94.18 \pm 0.06$	20	0.01	$78.65 \pm 0.25$
pFedMe-PM	MLR	15	0.01	$\mathbf{95.62 \pm 0.04}$	20	0.01	$\mathbf{83.20 \pm 0.06}$
FedAvg	DNN	–	0.02	$98.79 \pm 0.03$	–	0.03	$83.64 \pm 0.22$
Per-FedAvg	DNN	–	$(0.02, 0.001)$	$98.90 \pm 0.02$	–	$(0.01, 0.001)$	$85.01 \pm 0.10$
pFedMe-GM	DNN	30	0.01	$99.16 \pm 0.03$	30	0.01	$84.17 \pm 0.35$
pFedMe-PM	DNN	30	0.01	$\mathbf{99.46 \pm 0.01}$	30	0.01	$\mathbf{86.36 \pm 0.15}$

The personalized model pFedMe-PM outperforms all baselines in test accuracy across both convex (MLR) and nonconvex (DNN) benchmarks on MNIST and Synthetic datasets, exceeding FedAvg by up to 5.58% and Per-FedAvg by up to 1.71% on the heterogeneous Synthetic dataset.

Coverage note — All core contributions—the bi-level Moreau envelope formulation, Algorithm 1, inner problem approximation analysis, gradient diversity bounds, strongly convex and nonconvex convergence theorems and speedup corollaries, and full experimental results—have been faithfully captured. Intermediate proof-only lemmas were omitted.

References

1.H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas, “Communication-Efficient Learning of Deep Networks from Decentralized Data,” arXiv:1602.05629 [cs], Feb. 2017. [Online]. Available: http://arxiv.org/abs/1602.05629
2.M. Mohri, G. Sivek, and A. T. Suresh, “Agnostic Federated Learning,” arXiv:1902.00146 [cs, stat], Jan. 2019. [Online]. Available: http://arxiv.org/abs/1902.00146
3.S. P. Karimireddy et al., “SCAFFOLD: Stochastic Controlled Averaging for Federated Learning,” arXiv:1910.06378 [cs, math, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/1910.06378
4.K. Pillutla, S. M. Kakade, and Z. Harchaoui, “Robust Aggregation for Federated Learning,” arXiv:1912.13445 [cs, stat], Dec. 2019. [Online]. Available: http://arxiv.org/abs/1912.13445
5.D. Li and J. Wang, “FedMD: Heterogenous Federated Learning via Model Distillation,” arXiv:1910.03581 [cs, stat], Oct. 2019. [Online]. Available: http://arxiv.org/abs/1910.03581
6.Y. Deng, M. M. Kamani, and M. Mahdavi, “Adaptive Personalized Federated Learning,” arXiv:2003.13461 [cs, stat], Mar. 2020. [Online]. Available: http://arxiv.org/abs/2003.13461
7.J.-J. Moreau, “Propriétés des applications ‘prox’,” Compte Rendus Acad. Sci., no. 256, pp. 1069–1071, 1963.
8.A. Fallah, A. Mokhtari, and A. Ozdaglar, “Personalized Federated Learning: A Meta-Learning Approach,” arXiv:2002.07948 [cs, math, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/2002.07948
9.N. Guha, A. Talwalkar, and V. Smith, “One-Shot Federated Learning,” arXiv:1902.11175 [cs, stat], Mar. 2019. [Online]. Available: http://arxiv.org/abs/1902.11175
10.A. Reisizadeh, A. Mokhtari, H. Hassani, A. Jadbabaie, and R. Pedarsani, “FedPAQ: A Communication-Efficient Federated Learning Method with Periodic Averaging and Quantization,” arXiv:1909.13014 [cs, math, stat], Mar. 2020. [Online]. Available: http://arxiv.org/abs/1909.13014
11.X. Dai et al., “Hyper-Sphere Quantization: Communication-Efficient SGD for Federated Learning,” arXiv:1911.04655 [cs, stat], Nov. 2019. [Online]. Available: http://arxiv.org/abs/1911.04655
12.J. Wang and G. Joshi, “Cooperative SGD: A unified Framework for the Design and Analysis of Communication-Efficient SGD Algorithms,” arXiv:1808.07576 [cs, stat], Jan. 2019. [Online]. Available: http://arxiv.org/abs/1808.07576
13.T. Lin, S. U. Stich, K. K. Patel, and M. Jaggi, “Don’t Use Large Mini-Batches, Use Local SGD,” arXiv:1808.07217 [cs, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/1808.07217
14.S. U. Stich, “Local SGD Converges Fast and Communicates Little,” arXiv:1805.09767 [cs, math], May 2019. [Online]. Available: http://arxiv.org/abs/1805.09767
15.T. Li et al., “Federated Optimization in Heterogeneous Networks,” arXiv:1812.06127 [cs, stat], Sep. 2019. [Online]. Available: http://arxiv.org/abs/1812.06127
16.Y. Zhao et al., “Federated Learning with Non-IID Data,” arXiv:1806.00582 [cs, stat], Jun. 2018. [Online]. Available: http://arxiv.org/abs/1806.00582
17.F. Haddadpour and M. Mahdavi, “On the Convergence of Local Descent Methods in Federated Learning,” arXiv:1910.14425 [cs, stat], Dec. 2019. [Online]. Available: http://arxiv.org/abs/1910.14425
18.X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the Convergence of FedAvg on Non-IID Data,” arXiv:1907.02189 [cs, math, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/1907.02189
19.A. Khaled, K. Mishchenko, and P. Richtárik, “Tighter Theory for Local SGD on Identical and Heterogeneous Data,” arXiv:1909.04746 [cs, math, stat], Mar. 2020. [Online]. Available: http://arxiv.org/abs/1909.04746
20.V. Smith, C.-K. Chiang, M. Sanjabi, and A. Talwalkar, “Federated Multi-Task Learning,” arXiv:1705.10467 [cs, stat], Feb. 2018. [Online]. Available: http://arxiv.org/abs/1705.10467
21.J. C. Duchi, M. I. Jordan, and M. J. Wainwright, “Privacy Aware Learning,” J. ACM, vol. 61, no. 6, pp. 1–57, Dec. 2014. [Online]. Available: https://dl.acm.org/doi/10.1145/2666468
22.H. B. McMahan, D. Ramage, K. Talwar, and L. Zhang, “Learning Differentially Private Recurrent Language Models,” arXiv:1710.06963 [cs], Feb. 2018. [Online]. Available: http://arxiv.org/abs/1710.06963
23.W. Zhu, P. Kairouz, B. McMahan, H. Sun, and W. Li, “Federated Heavy Hitters Discovery with Differential Privacy,” arXiv:1902.08534 [cs], Feb. 2020. [Online]. Available: http://arxiv.org/abs/1902.08534
24.N. Agarwal, A. T. Suresh, F. X. X. Yu, S. Kumar, and B. McMahan, “cpSGD: Communication-efficient and differentially-private distributed SGD,” p. 12.
25.Z. Li, V. Sharma, and S. P. Mohanty, “Preserving Data Privacy via Federated Learning: Challenges and Solutions,” IEEE Consumer Electronics Magazine, vol. 9, no. 3, pp. 8–16, May 2020. [Online]. Available: https://ieeexplore.ieee.org/document/9055478/
26.F. Hanzely and P. Richtárik, “Federated Learning of a Mixture of Global and Local Models,” arXiv:2002.05516 [cs, math, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/2002.05516
27.Y. Mansour, M. Mohri, J. Ro, and A. T. Suresh, “Three Approaches for Personalization with Applications to Federated Learning,” arXiv:2002.10619 [cs, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/2002.10619
28.M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary, “Federated Learning with Personalization Layers,” arXiv:1912.00818 [cs, stat], Dec. 2019. [Online]. Available: http://arxiv.org/abs/1912.00818
29.A. Hard et al., “Federated Learning for Mobile Keyboard Prediction,” arXiv:1811.03604 [cs], Feb. 2019. [Online]. Available: http://arxiv.org/abs/1811.03604
30.K. Wang et al., “Federated Evaluation of On-device Personalization,” arXiv:1910.10252 [cs, stat], Oct. 2019. [Online]. Available: http://arxiv.org/abs/1910.10252
31.P. Vanhaesebrouck, A. Bellet, and M. Tommasi, “Decentralized Collaborative Learning of Personalized Models over Networks,” arXiv:1610.05202 [cs, stat], Feb. 2017. [Online]. Available: http://arxiv.org/abs/1610.05202
32.C. Finn, P. Abbeel, and S. Levine, “Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks,” arXiv:1703.03400 [cs], Jul. 2017. [Online]. Available: http://arxiv.org/abs/1703.03400
33.A. Nichol, J. Achiam, and J. Schulman, “On First-Order Meta-Learning Algorithms,” arXiv:1803.02999 [cs], Oct. 2018. [Online]. Available: http://arxiv.org/abs/1803.02999
34.A. Fallah, A. Mokhtari, and A. Ozdaglar, “On the Convergence Theory of Gradient-Based Model-Agnostic Meta-Learning Algorithms,” arXiv:1908.10400 [cs, math, stat], Mar. 2020. [Online]. Available: http://arxiv.org/abs/1908.10400
35.M. Khodak, M.-F. Balcan, and A. Talwalkar, “Adaptive Gradient-Based Meta-Learning Methods,” arXiv:1906.02717 [cs, stat], Dec. 2019. [Online]. Available: http://arxiv.org/abs/1906.02717
36.Y. Jiang, J. Konečný, K. Rush, and S. Kannan, “Improving Federated Learning Personalization via Model Agnostic Meta Learning,” arXiv:1909.12488 [cs, stat], Sep. 2019. [Online]. Available: http://arxiv.org/abs/1909.12488
37.F. Chen, Z. Dong, Z. Li, and X. He, “Federated Meta-Learning for Recommendation,” Feb. 2018.
38.T. Li, A. K. Sahu, A. Talwalkar, and V. Smith, “Federated Learning: Challenges, Methods, and Future Directions,” arXiv:1908.07873 [cs, stat], Aug. 2019. [Online]. Available: http://arxiv.org/abs/1908.07873
39.P. Kairouz et al., “Advances and Open Problems in Federated Learning.” arXiv: 1912.04977, Dec. 2019.
40.H. Lin, J. Mairal, and Z. Harchaoui, “Catalyst Acceleration for First-order Convex Optimization: from Theory to Practice,” arXiv:1712.05654 [math, stat], Jun. 2018. [Online]. Available: http://arxiv.org/abs/1712.05654
41.P. Zhou, X. Yuan, H. Xu, S. Yan, and J. Feng, “Efficient Meta Learning via Minibatch Proximal Update,” in Advances in Neural Information Processing Systems 32, H. Wallach et al., Eds. Curran Associates, Inc., 2019, pp. 1534–1544. [Online]. Available: http://papers.nips.cc/paper/8432-efficient-meta-learning-via-minibatch-proximal-update.pdf
42.X. Li, W. Yang, S. Wang, and Z. Zhang, “Communication-Efficient Local Decentralized SGD Methods,” arXiv:1910.09126 [cs, math, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/1910.09126
43.H. Yu, R. Jin, and S. Yang, “On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization,” arXiv:1905.03817 [cs, math], May 2019. [Online]. Available: http://arxiv.org/abs/1905.03817
44.L. Nguyen et al., “New Convergence Aspects of Stochastic Gradient Algorithms,” Journal of Machine Learning Research, vol. 20, Nov. 2019.
45.A. Khaled, K. Mishchenko, and P. Richtárik, “First Analysis of Local GD on Heterogeneous Data,” arXiv:1909.04715 [cs, math, stat], Mar. 2020. [Online]. Available: http://arxiv.org/abs/1909.04715
46.C. Lemaréchal and C. Sagastizábal, “Practical Aspects of the Moreau–Yosida Regularization: Theoretical Preliminaries,” SIAM J. Optim., vol. 7, no. 2, pp. 367–385, May 1997. [Online]. Available: http://epubs.siam.org/doi/10.1137/S1052623494267127
47.C. Planiden and X. Wang, “Strongly Convex Functions, Moreau Envelopes, and the Generic Nature of Convex Functions with Strong Minimizers,” SIAM J. Optim., vol. 26, no. 2, pp. 1341–1364, Jan. 2016. [Online]. Available: http://epubs.siam.org/doi/10.1137/15M1035550
48.T. Hoheisel, M. Laborde, A. Oberman, and ,Department of Mathematics and Statistics, McGill University, Montreal, Canada, “A regularization interpretation of the proximal point method for weakly convex functions,” Journal of Dynamics & Games, vol. 7, no. 1, pp. 79–96, 2020. [Online]. Available: http://aimsciences.org//article/doi/10.3934/jdg.2020005
49.S. Reddi et al., “Adaptive Federated Optimization,” arXiv:2003.00295 [cs, math, stat], Feb. 2020. [Online]. Available: http://arxiv.org/abs/2003.00295
50.S. Bubeck, “Convex Optimization: Algorithms and Complexity,” arXiv:1405.4980 [cs, math, stat], Nov. 2015. [Online]. Available: http://arxiv.org/abs/1405.4980
51.Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner, “Gradient-Based Learning Applied to Document Recognition,” Proceedings of the IEEE, vol. 86, no. 11, pp. 2278–2324, Nov. 1998.
52.A. Paszke et al., “PyTorch: An Imperative Style, High-Performance Deep Learning Library,” in Advances in Neural Information Processing Systems 32, H. Wallach et al., Eds. Curran Associates, Inc., 2019, pp. 8026–8037. [Online]. Available: http://papers.nips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
53.Y. Nesterov, Lectures on convex optimization. New York, NY: Springer Berlin Heidelberg, 2018. [Online]. Available: https://www.springer.com/gp/book/9783319915777
54.Y. Arjevani, O. Shamir, and N. Srebro, “A Tight Convergence Analysis for Stochastic Gradient Descent with Delayed Updates,” arXiv:1806.10188 [cs, math, stat], Jun. 2018. [Online]. Available: http://arxiv.org/abs/1806.10188
55.S. U. Stich, “Unified Optimal Analysis of the (Stochastic) Gradient Method,” Jul. 2019. [Online]. Available: https://arxiv.org/abs/1907.04232v2

Citation

MLA

Dinh, C. T., et al. “Personalized Federated Learning with Moreau Envelopes”. arXiv, 2020, http://arxiv.org/abs/2006.08848v3.

APA

Dinh, C. T., Tran, N. H., & Nguyen, T. D. (2020). Personalized Federated Learning with Moreau Envelopes. arXiv. http://arxiv.org/abs/2006.08848v3

Chicago

Dinh, C. T., N. H. Tran, and T. D. Nguyen. 2020. “Personalized Federated Learning with Moreau Envelopes”. arXiv. http://arxiv.org/abs/2006.08848v3.

Harvard

Dinh, C.T., Tran, N.H. and Nguyen, T.D. (2020) “Personalized Federated Learning with Moreau Envelopes”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2006.08848v3.

Vancouver

1. Dinh CT, Tran NH, Nguyen TD (2020) Personalized Federated Learning with Moreau Envelopes. arXiv

BibTeX

@article{dinh2020personalized,
  title = {Personalized Federated Learning with Moreau Envelopes},
  author = {Dinh, Canh T. and Tran, Nguyen H. and Nguyen, Tuan Dung},
  year = {2020},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2006.08848v3},
  eprint = {2006.08848}
}

Metadata:arXiv

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: Authors

Personalized Federated Learning with Moreau Envelopes

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Bi-Level Problem Formulation of Personalized Federated Learning with Moreau Envelopes

Knowl 2 — pFedMe Personalized Federated Optimization Algorithm

Knowl 3 — Standard Smoothness, Variance, and Diversity Assumptions for pFedMe Analysis

Knowl 4 — Approximation Error Bound for the Inner Personalized Model Subproblem

Knowl 5 — Gradient Diversity Bounds for Client Moreau Envelopes

Knowl 6 — Convergence Rate and Quadratic Speedup for Strongly Convex Objectives in pFedMe

Knowl 7 — Convergence Rate and Sublinear Speedup for Nonconvex Smooth Objectives in pFedMe

Knowl 8 — Experimental Configuration for Heterogeneous Personalized Federated Learning

Knowl 9 — Impact of Hyperparameters on pFedMe Optimization Dynamics

Knowl 10 — Performance Comparison of pFedMe against FedAvg and Per-FedAvg

References

Citation

BibTeX

Access the Paper