PaperlayerView exact revision on arXiv

Source-linked AI summary

Personalized Federated Learning: A Meta-Learning Approach

Alireza Fallah, Aryan Mokhtari, Asuman Ozdaglar

arXiv:2002.07948v4cs.LGmath.OCstat.ML

TL;DR

Federated Learning’s shared model does not adapt well to heterogeneous users, motivating a personalized objective. The paper connects this objective to MAML and develops Per-FedAvg, characterizing its non-convex convergence and dependence on distribution similarity. Numerical experiments compare Per-FedAvg approximations with FedAvg.

  • Problem

    Federated Learning produces a common model that does not adapt to users with heterogeneous data distributions.

  • Method

    The paper formulates personalization as learning a shared initialization that users adapt locally, then develops Per-FedAvg to optimize this objective.

  • Results

    The paper characterizes Per-FedAvg’s communication complexity for reaching an ϵ-approximate first-order stationary point in nonconvex settings and evaluates first-order approximations numerically.

  • Takeaways & Limitations

    A decentralized MAML implementation can solve the proposed personalized Federated Learning problem while retaining the Federated Learning architecture.

  • Takeaways & Limitations

    The Wasserstein-distance bound has exponentially slow convergence with dimension, and the analysis assumes bounded-support user distributions.

Abstract

from arXiv · show

In Federated Learning, we aim to train models across multiple computing units (users), while users can only communicate with a common central server, without exchanging their data samples. This mechanism exploits the computational power of all users and allows users to obtain a richer model as their models are trained over a larger set of data points. However, this scheme only develops a common output for all the users, and, therefore, it does not adapt the model to each user. This is an important missing feature, especially given the heterogeneity of the underlying data distribution for various users. In this paper, we study a personalized variant of the federated learning in which our goal is to find an initial shared model that current or new users can easily adapt to their local dataset by performing one or a few steps of gradient descent with respect to their own data. This approach keeps all the benefits of the federated learning architecture, and, by structure, leads to a more personalized model for each user. We show this problem can be studied within the Model-Agnostic Meta-Learning (MAML) framework. Inspired by this connection, we study a personalized variant of the well-known Federated Averaging algorithm and evaluate its performance in terms of gradient norm for non-convex loss functions. Further, we characterize how this performance is affected by the closeness of underlying distributions of user data, measured in terms of distribution distances such as Total Variation and 1-Wasserstein metric.

1 Introduction

Federated Learning combines users’ computational resources without sharing local data, but its single shared model does not adapt to heterogeneous user distributions. The paper proposes a personalized formulation grounded in MAML and studies its theoretical foundations.

  • Federated Learning trains over users’ combined data through a central server without exchanging local samples.
  • Averaging local losses improves access to distributed data but produces one common model that may not fit heterogeneous users.
  • The proposed formulation seeks a shared initialization that each user can adapt with a few gradient-based updates.
  • Per-FedAvg is a personalized FedAvg method, with convergence properties studied for non-convex losses.
  • The analysis focuses on Federated Learning MAML convergence, where users perform multiple local updates before communicating with the server.
  • Prior work connects personalization in Federated Learning with multi-task learning, mixture models, and empirical MAML-type approaches.

2 Personalized Federated Learning via Model-Agnostic Meta-Learning

The paper formulates personalized Federated Learning as a MAML-style problem: learn an initialization that adapts quickly to each user’s heterogeneous data. This formulation supports one or a few local gradient steps and motivates a non-convex convergence analysis.

  • MAML seeks an initialization that performs well after limited updating on a new task rather than a model optimal in expectation across tasks.
  • Users adapt the shared initialization to their own datasets by taking one or a few gradient-descent steps.
  • Unlike minimizing the average loss, the personalized objective targets an initial model that becomes effective for each user after local adaptation.
  • The paper studies a provably convergent method for the personalized objective when user loss functions are nonconvex.

3 Personalized FedAvg

Per-FedAvg adapts FedAvg to optimize personalized meta-functions rather than the conventional global objective. It samples users, performs local stochastic updates, and averages their resulting models at the server.

  • Per-FedAvg is designed to solve the personalized objective by representing its global function as an average of user-specific meta-functions.
  • Each round samples rn users uniformly, sends them the current model, and has them perform τ local stochastic-gradient steps.
  • Per-FedAvg is specified as an iterative algorithm with an initial iterate, random active-user selection, model transmission, local computation, and server aggregation.
  • Local meta-gradient estimation uses stochastic gradients and Hessian estimates computed from data batches.
  • The implementation includes nested stochastic-gradient computations and can be organized into two equivalent stages.
  • The server updates the global model by averaging the models returned by participating users.

4 Theoretical Results

The paper analyzes Per-FedAvg in nonconvex settings, deriving communication complexity under smoothness, variance, and user-similarity assumptions. It also connects function similarity to Total Variation and 1-Wasserstein distances between user distributions.

  • Convergence Analysis: Per-FedAvg convergence is analyzed for finding an ϵ-approximate first-order stationary point in nonconvex settings.The target condition is E[∥∇F(wϵ)∥2] ≤ ϵ.
  • Assumptions: The analysis assumes smooth local losses, bounded gradients, Lipschitz Hessians, bounded gradient and Hessian-estimation variance, and controlled diversity across users.The user-similarity assumption separately captures differences among local gradients and Hessians.
  • Task Similarity and Distribution Distances: Assumption 5 links user-function similarity to the closeness of data distributions under Total Variation and 1-Wasserstein distances.The paper studies these connections when users share the same loss function over a common input-label space.
  • Task Similarity and Distribution Distances: For homogeneous data distributions with empirical samples of size m, the Wasserstein-based similarity parameters satisfy γG, γH = O(1/√m).This result uses W1(pi, pu) = O(1/√m).
  • Task Similarity and Distribution Distances: The Wasserstein bound has exponentially slow convergence with dimension, limiting its usefulness in high-dimensional settings.This caveat concerns the dependence of Wasserstein distance on the number of samples.
  • Convergence Analysis: The convergence bound depends on stochastic-estimation error, and increasing batch sizes can make σF arbitrarily small.In the exact-gradient and exact-Hessian case, σF = 0.
  • Convergence Analysis: Per-FedAvg requires K = O(ϵ^-3/2) communication rounds under the stated local-update and stepsize choices.The corollary sets τ = O(ϵ^-1/2) and β = ϵ, while the result reaches an O(ϵ + α2σ2G/D)-first-order stationary point.

5 Numerical Experiments

The experiments compare FedAvg with two efficient Per-FedAvg approximations under heterogeneous MNIST and CIFAR-10 distributions. Per-FedAvg (HF) consistently provides the strongest personalization and accuracy, while Per-FedAvg (FO) is more sensitive to larger adaptation steps.

  • Experimental setup: The study evaluates FedAvg, Per-FedAvg (FO), and Per-FedAvg (HF) on heterogeneous MNIST and CIFAR-10 user distributions.The experiments use a neural network across 50 users and report test accuracy with 95% confidence intervals.
  • Algorithms: Per-FedAvg (FO) ignores the Hessian term, whereas Per-FedAvg (HF) approximates Hessian-vector products using differences of gradients.Both methods avoid the computational cost of accessing second-order information directly.
  • Results: For α = 0.001 and τ = 10, both Per-FedAvg approximations perform almost similarly and better than FedAvg.Decreasing τ decreases the performance of all three algorithms because the total number of iterations is reduced.
  • Results: Increasing α from 0.001 to 0.01 improves Per-FedAvg (HF) for τ = 4, while Per-FedAvg (FO) drops significantly at larger α.The authors suggest that HF may adapt better to user data at test time under this setting.
  • Results: Per-FedAvg (HF) outperforms FedAvg in all tested cases and yields a more personalized solution, especially under the modified heterogeneous distributions.Per-FedAvg (FO) continues to suffer from the performance issue observed at larger α.

6 Conclusion

The paper formulates personalized federated learning as learning an initialization that users can quickly adapt locally. It connects this formulation to MAML and presents Per-FedAvg as a decentralized solution with theoretical and numerical support.

  • Conclusion: The proposed personalized FL formulation learns an initialization that users can quickly adapt to their local data after training.This targets heterogeneous users while retaining the federated architecture.
  • Conclusion: The formulation is connected to MAML, and its decentralized implementation is called Per-FedAvg.The paper characterizes Per-FedAvg for achieving first-order optimality in nonconvex settings.
  • Conclusion: Numerical experiments compare two first-order Per-FedAvg approximations with FedAvg and show that Per-FedAvg produces a more personalized solution.The conclusion summarizes both the algorithmic comparison and the personalization outcome.

B Proofs of results in Subsection 4.1

These proof passages establish distribution-dependent bounds by applying Lipschitz assumptions, Kantorovich-Rubinstein duality, and norm inequalities. The argument separately handles gradient and Hessian-related quantities.

  • Proof strategy: The proof uses gradient Lipschitzness and bounded gradients to derive conditions on the local losses.The argument plugs the resulting bound into an aggregate gradient-difference expression.
  • Wasserstein bound: Kantorovich-Rubinstein duality expresses W1(pi, p) through expectations of continuous 1-Lipschitz functions.The proof then uses this representation to control a gradient-based function of the data point.
  • Hessian bound: The same duality argument is applied to Hessian quadratic forms to obtain the corresponding γH result.The proof considers both positive and negative Hessian-based functions.
  • Conclusion of proof: Smoothness, the stepsize condition α ≤ 1/L, and triangle inequalities combine the intermediate bounds into the desired result.The final step puts equations (26) and (27) together.

D Proof of Lemma 4.3

The proof of Lemma 4.3 bounds stochastic-gradient terms by decomposing their errors and controlling moments under smoothness, independence, and bounded-Hessian assumptions. Variance is then related to a centered second moment.

  • Gradient bounds: The proof begins from the stochastic-gradient expression and bounds its norm using matrix-norm inequalities.The bound uses ∥I − α∇2fi(w)∥ ≤ 1 + αL ≤ 2.
  • Error decomposition: The stochastic error is decomposed into terms e1 and e2, whose moments are analyzed using independence and zero-mean conditioning.Conditioning on D makes the first term zero mean while the second term is deterministic.
  • Moment bounds: Smoothness, Assumption 4, and the condition αL ≤ 1 control the first and second moments of the decomposed errors.Cauchy-Schwarz and matrix-norm inequalities are used repeatedly in these bounds.
  • Proof conclusion: Comparing the relevant expressions and substituting the intermediate bounds yields the stated results of Lemma 4.3.The proof combines the decomposed-term estimates with the Hessian norm bound.
  • Variance bound: The variance of ∇Fi(w) is bounded through the second moment of its centered difference from ∇f(w − α∇f(w)).Subtracting a constant leaves variance unchanged, after which the proof applies the previously established bounds.

F An Intermediate Result

This section establishes bounds on local-update deviations and their averages under the stated smoothness and variance assumptions. The proposition yields a corollary under a sufficiently small local stepsize.

  • Local-update averaging: The average local update is defined as wk,t = 1/n Σ_i wi_k,t, linking individual updates to the aggregate iterate.This average is used throughout the deviation analysis.
  • Proposition F.1: Under Assumptions 2–4, Proposition F.1 bounds local-update quantities for any α ∈ [0, 1/L] and t ≥ 0.The constants LF, σF, and γF are defined through supporting lemmas.
  • Corollary F.2: For β ≤ 1/(10τLF), Corollary F.2 converts Proposition F.1 into simplified bounds on the update deviations.The proof uses the stepsize restriction and t ≤ τ − 1.
  • Proof of Proposition F.1: The analysis recursively controls deviations across local steps using smoothness, Cauchy–Schwarz, and auxiliary lemmas.The resulting recursion completes the bounds in Proposition F.1.
  • Proof of Corollary F.2: The proof combines bounds on gradient differences and local-update variance to obtain the stated results for both deviation terms.These steps use Lemmas 4.2–4.4 and standard norm inequalities.

G Proof of Theorem 4.5

The proof of Theorem 4.5 decomposes the update error into sampling, local-update, and gradient-related terms. It bounds these terms conditionally, sums the resulting descent inequalities, and obtains the theorem’s convergence result.

  • Analytical setup: For analysis, all agents are assumed to perform local updates, although Algorithm 1 samples only a fraction for computing wk+1.This assumption is analytical and does not change the algorithm’s aggregation rule.
  • Termwise bounds: The proof conditions on Ft and bounds the moments of X, Y, and Z using Cauchy–Schwarz, smoothness, random sampling, and the tower rule.These bounds isolate the stochastic and client-drift contributions to the update.
  • Sampling argument: Uniform random selection of Ak makes the sampled gradient unbiased for the average gradient at the corresponding iterate.This property is used when relating sampled updates to the global objective.
  • Final summation: The descent inequality is summed over local steps t = 0, …, τ − 1 and communication rounds k = 0, …, K − 1 to derive the final convergence statement.The identity w̄k+1,τ = wk+1 connects local-step analysis to the global iterate.
  • Stepsize variation: The analysis can use a diminishing stepsize βk = O(1/√K) while recovering the same complexity bounds.The descent result holds with β replaced by βk at iteration k.

H On First-Order Approximations of Per-FedAvg

This section analyzes first-order approximations of Per-FedAvg that avoid or approximate Hessian–vector products. It characterizes their added error and shows how the convergence analysis changes.

  • Motivation: Per-FedAvg and MAML require Hessian–vector products, motivating first-order approximations for computationally costly applications.The section extends the analysis to dropping the second-order term or approximating it with gradient differences.
  • Convergence analysis: The first-order approximation analysis preserves Theorem 4.5 after replacing σF with σ̃F and modifying the associated variance term.The remaining proof steps are unchanged under this substitution.
  • Approximation methods: FO-MAML drops the second-order term to reduce computation cost, while HF-MAML approximates the Hessian–vector product through gradient differences.The latter approach is called Per-FedAvg (HF).
  • Approximation limitation: For the dropped second-order term, one component of the error can shrink with larger batches, whereas another remains non-vanishing but becomes negligible when α is small.This residual error propagates into the final theorem.
  • HF-MAML analysis: The gradient-difference approximation introduces an additional degree of freedom, δ, that controls an error term not reduced by increasing batch sizes.The approximation error is bounded using the Hessian’s Lipschitz-continuity parameter.

I More on Numerical Experiments

The numerical experiments compare Per-FedAvg variants with federated baselines using accuracy over time and final accuracy. Per-FedAvg (HF) achieves the strongest reported final accuracy and higher accuracy within the same computation time.

  • Runtime comparison: Figure 1a compares average test accuracy with respect to time across the studied algorithms.Per-FedAvg (HF) reaches higher accuracy than regular Fed-Avg with local updates within the same computation time.
  • Final accuracy: 43.73 ± 0.11: Per-FedAvg (HF) final accuracy, compared with 36.74 ± 0.1 for Fed-Avg+ARUBA after refinement.The experiment uses τ = 4 and K = 1000.
  • Final accuracy: 34.04 ± 0.08: Per-FedAvg (FO) final accuracy in the reported comparison.The same comparison reports Fed-Avg+ARUBA at 36.74 ± 0.1 and Per-FedAvg (HF) at 43.73 ± 0.11.
  • Convergence paths: Figure 1b shows one realization of each training path to provide intuition about the methods’ convergence speed.The figure complements the time-based accuracy comparison.
Loading 2002.07948v4…