PaperlayerView exact revision on arXiv

Source-linked AI summary

Personalized Federated Learning with Moreau Envelopes

Canh T. Dinh, Nguyen H. Tran, Tuan Dung Nguyen

arXiv:2006.08848v3cs.LGcs.DCstat.ML

TL;DR

Statistical diversity makes a single federated global model poorly generalized across clients, while isolated client training suffers from insufficient data. pFedMe uses Moreau envelopes in a bi-level personalized FL formulation to decouple personalized optimization from global learning, achieving state-of-the-art convergence speedups and outperforming FedAvg and Per-FedAvg experimentally.

  • Problem

    Statistical diversity among clients limits the global model’s generalization to each client, while individual learning suffers from insufficient data.

  • Method

    pFedMe uses Moreau envelopes as regularized client losses in a bi-level problem that decouples personalized model optimization from global model learning.

  • Results

    pFedMe achieves quadratic speedup for strongly convex objectives, order-2/3 sublinear speedup for smooth nonconvex objectives, and outperforms FedAvg and Per-FedAvg experimentally.

  • Takeaways & Limitations

    Personalized models can adapt to clients’ statistical diversity while retaining federated aggregation, supporting better predictors for client-specific conditions.

  • Takeaways & Limitations

    The theoretical formulation allows non-i.i.d. client distributions and avoids assuming uniformly bounded gradients, limiting the stated scope to those analyzed conditions.

Abstract

from arXiv · show

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.

1 Introduction

Federated learning coordinates privacy-preserving model training across clients, but heterogeneous data distributions limit one global model’s performance on individual tasks. pFedMe addresses this through Moreau-envelope regularization, achieving faster theoretical convergence and stronger empirical results than FedAvg and Per-FedAvg.

  • Federated learning setting: FL repeatedly broadcasts a global model, performs local client updates, and averages models from sampled clients at the server.This preserves client data locally while reducing communication through iterative aggregation.
  • Challenge: Statistical diversity makes the global model poorly generalized to individual clients, while isolated client training suffers from insufficient data.The generalization error of the global model increases significantly as client diversity grows.
  • Proposed approach: pFedMe minimizes clients’ Moreau envelopes in a bi-level formulation to separate personalized-model optimization from global-model learning.Clients optimize local personalized models while contributing to a shared reference model.
  • Theory: Quadratic speedup is obtained for strongly convex objectives, while smooth nonconvex objectives achieve sublinear speedup of order 2/3.These rates improve on existing linear and order-1/2 speedups, respectively.
  • Empirical evaluation: pFedMe outperforms vanilla FedAvg and Per-FedAvg in convergence rate and local accuracy on real and synthetic datasets.The experiments target datasets exhibiting statistical diversity among clients.

2 Related Work

Related work addresses federated learning’s communication, privacy, statistical-heterogeneity, and personalization challenges through multiple algorithmic strategies. Personalization methods include model mixing, contextualization, meta-learning, and multi-task learning.

  • FL and challenges: FedAvg builds a global model from local SGD updates, while later methods reduce communication through one-shot learning, quantization, or multiple local optimization rounds.The related work also includes methods targeting statistical diversity and privacy preservation.
  • Personalized FL: Personalized FL approaches include mixing global and local models, user clustering, data interpolation, model interpolation, and multi-task learning.These strategies differ in how they represent or exploit client-specific information.
  • Personalized FL: Some contextualization approaches require meta-features from all clients, which is incompatible with federated privacy constraints.Model interpolation was used in MOCHA to address systems and statistical heterogeneity.

3 Personalized Federated Learning with Moreau Envelopes (pFedMe)

pFedMe formulates personalized federated learning as a bi-level optimization problem using Moreau-envelope regularization. It jointly learns a shared reference model and client-specific models through local inner optimization and federated outer updates.

  • Problem formulation: Conventional FL seeks a global model from client losses, with each client loss defined as an expected loss over its data distribution.The formulation aggregates information across N clients.
  • Problem formulation: Client data distributions may be non-i.i.d. because samples arise from different environments, contexts, and applications.This heterogeneity motivates personalized rather than purely global optimization.
  • pFedMe formulation: pFedMe regularizes each personalized model toward the global reference model, with λ controlling the strength of that coupling.Large λ favors aggregation for clients with unreliable data, whereas small λ favors personalization for clients with sufficient useful data.
  • pFedMe formulation: The bi-level formulation optimizes each θ_i with respect to local data at the inner level and aggregates clients to learn w at the outer level.The inner solution is the unique personalized model and corresponds to the proximal operator associated with the Moreau envelope.
  • Comparison with Per-FedAvg: Compared with Per-FedAvg’s one-step initialization update, pFedMe can use any iterative inner optimizer with multi-step updates.Both methods give the global model a meta-model interpretation, but pFedMe jointly pursues global and personalized models.
  • Comparison with Per-FedAvg: pFedMe directly minimizes the client loss, whereas Per-FedAvg optimizes a first-order approximation and generally requires Hessian computation or estimation.pFedMe uses a first-order gradient approach.
  • Assumptions: The convergence analysis assumes strong convexity or smooth nonconvexity, bounded stochastic-gradient variance, and bounded client diversity.The diversity and variance assumptions represent sampling noise and differences among client data distributions.
  • Moreau-envelope properties: Moreau envelopes preserve convexity and provide smooth gradients, enabling convergence analysis under the stated objective conditions.For nonconvex L-smooth losses, the smoothness result requires λ > 2L.

4 pFedMe: Convergence Analysis

The analysis establishes pFedMe’s convergence under strongly convex and smooth nonconvex settings, with rates affected by client sampling, local-update drift, and approximation error. It derives quadratic speedup for strongly convex objectives and sublinear speedup of order 2/3 for nonconvex objectives.

  • Intermediate results: The convergence analysis begins with an intermediate Moreau-envelope result that bounds client-diversity variance for strongly convex and smooth nonconvex losses.The bounds characterize diversity through σ2_F,1 and σ2_F,2; σ2_F,1 is finite at the unique pFedMe solution in the strongly convex case.
  • Strongly convex objectives: Under Assumption 1(a), the strongly convex analysis establishes convergence of pFedMe’s global and personalized models.The theorem requires Assumptions 1(a) and 2, while the step size is bounded using ˆη1, β, and R.
  • Error structure: The strongly convex convergence result shows that client sampling, local-update drift, and approximation error contribute distinct terms to the overall error.When all clients participate, the client-sampling term is zero; the approximation error leaves convergence within a λ2δ2/µF-neighbourhood, and λ trades off errors involving δ2 and σ2.
  • Strongly convex objectives: Quadratic speedup is obtained for strongly convex objectives, exceeding the linear speedup available in most existing analyses.With suitable client-sampling and parameter choices, the analysis also identifies linear speedup for client-sampling error and communication rounds.
  • Smooth nonconvex objectives: Theorem 2 establishes pFedMe convergence for smooth nonconvex objectives under Assumptions 1(b), 2, and 3.The result uses a step-size condition involving ˆη2, β, R, λ, and L, and samples t* uniformly from the communication rounds.
  • Smooth nonconvex objectives: For smooth nonconvex objectives, pFedMe achieves sublinear speedup of order 2/3, while existing analyses cited achieve only a sublinear speedup of order 1/2.Corollary 2 specifies parameter choices yielding O((TRN)^-2/3).

5 Experimental Results and Discussion

Experiments on heterogeneous MNIST and Synthetic data evaluate pFedMe’s hyperparameter sensitivity and compare it with FedAvg and Per-FedAvg. pFedMe’s personalized model achieves the strongest reported accuracy and convergence behavior, while several hyperparameters create computation–stability trade-offs.

  • Experimental setup: MNIST distributes heterogeneous handwritten-digit data across 20 clients, while Synthetic data uses heterogeneous client distributions for evaluation.MNIST clients receive different data sizes and only 2 of 10 labels; experiments use both strongly convex and nonconvex models.
  • Hyperparameter effects: Larger R improves convergence of personalized and global models but increases local computation and trades off against communication.R denotes local computation rounds; the experiments vary R on MNIST under fixed |D|, λ, K, and β.
  • Hyperparameter effects: K around 3 to 5 steps approximately computes personalized models, while larger K adds computation without improving convergence; K = 5 is selected thereafter.The study identifies increased user computation and energy consumption as costs of larger K.
  • Hyperparameter effects: Larger mini-batches initially improve convergence, but very large |D| slow convergence and increase local computation; experiments fix |D| = 20.The mini-batch size controls the approximation parameter δ.
  • Hyperparameter effects: Larger λ accelerates convergence, but excessively large λ can cause divergence, so λ requires dataset-specific tuning; MNIST uses λ = 15.β values above 1 can also accelerate convergence, especially for the global model, but require careful η adjustment for stability.
  • Performance comparison: pFedMe’s personalized model is the best performer across reported MNIST and Synthetic settings, outperforming FedAvg and Per-FedAvg in accuracy and convergence comparisons.With shared hyperparameters, MNIST personalized models exceed the global model, Per-FedAvg, and FedAvg by 1.1%, 1.3%, and 1.5% in the strongly convex setting, and by 0.9%, 0.9%, and 1.3% in the nonconvex setting; Synthetic gains are also reported.

6 Conclusion

The paper proposes pFedMe to address statistical diversity in federated learning by separating personalized-model optimization from global-model learning. It reports state-of-the-art convergence speedup and empirical advantages over FedAvg and Per-FedAvg, while noting that personalization requires dataset- and setting-specific parameter adaptation.

  • pFedMe uses Moreau envelopes to decompose personalized-model optimization from global-model learning.The approach updates the global model similarly to FedAvg while optimizing personalized models for each client’s local data distribution.
  • pFedMe achieves state-of-the-art convergence speedup rates.The paper reports theoretical speedup results for the method.
  • pFedMe outperforms vanilla FedAvg and Per-FedAvg in convex and non-convex settings on real and synthetic datasets.The reported comparisons cover both convergence and empirical performance across the evaluated settings.
  • The usefulness of personalization depends on adapting parameters to each dataset and federated setting.

A Proof of the Results

This section introduces existing results needed for the paper’s proofs and then presents proofs for two lemmas and two theorems.

  • The section first provides existing results used in the subsequent proofs.
  • The section then presents proofs of Lemma 1, Lemma 2, Theorem 1, and Theorem 2.

A.1 Review of useful existing results

This section states useful inequalities for smooth, strongly convex functions and a Jensen-inequality result for client-indexed vectors.

  • For an LF-smooth and µF-strongly convex function Fi, Proposition 2 gives two useful inequalities.
  • For vectors xi across clients, Proposition 3 applies Jensen’s inequality.

A.2 Proof of Lemma 1

The proof of Lemma 1 applies smoothness and Jensen-type inequalities, takes expectations, and uses independence and unbiasedness properties; case (b) follows similarly.

  • The proof applies Proposition 2 to derive an inequality involving the relevant iterates.
  • The proof uses Proposition 3 for its second inequality.
  • The proof takes expectations of both sides after establishing the inequality.
  • One inequality is justified by independent random variables and an unbiased estimate.
  • The proof of case (b) follows by a similar argument while considering hi(θi; wt.

A.3 Proof of Lemma 2

The proof treats two cases separately and derives bounds through expectation inequalities and cited propositions. It concludes by rearranging the resulting inequality.

  • The proof begins by separating cases (a) and (b).
  • Expectation-based inequalities and Propositions 2 and 3 provide the intermediate bounds in both cases.
  • The proof concludes by averaging over clients and rearranging the final inequality.

A.4 Proof of Theorem 1

The proof of Theorem 1 introduces analysis notation, rewrites local and global updates, and combines auxiliary lemmas under stated smoothness, convexity, and step-size conditions. The argument separately handles multiple cases before concluding the theorem's parts.

  • A.4 Proof of Theorem 1: The analysis defines additional notation and supporting lemmas before combining them to prove Theorem 1.
  • A.4 Proof of Theorem 1: The local and global updates are rewritten so the global step has an interpretable step size and approximate stochastic gradient.
  • A.4 Proof of Theorem 1: Client sampling, bounded diversity, and bounded client drift are controlled through separate lemmas and expectation bounds.
  • A.4 Proof of Theorem 1: The proof applies smoothness, convexity, Cauchy–Schwarz, and AM-GM inequalities while imposing step-size and parameter conditions.
  • A.4 Proof of Theorem 1: The analysis considers separate cases, combines them, and uses telescoping and uniform sampling to establish the theorem's parts.
  • A.4 Proof of Theorem 1: The broader FL setting remains constrained by explainability, malicious communication behavior, and uneven client data or computational resources.
Loading 2006.08848v3…