Back

Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach

Alireza FallahAryan MokhtariA. Ozdaglar

article2020NeurIPS1,345 citations

Proposes Per-FedAvg, a personalized federated learning algorithm built on model-agnostic meta-learning, establishing non-convex convergence guarantees and quantifying how statistical data heterogeneity directly impacts performance.

Listen

Federated learning enables multiple distributed clients—such as mobile devices or medical institutions—to collaboratively train a shared machine learning model without centralizing private data. However, standard federated learning methods produce a single global model designed to minimize average network-wide error. In real-world environments where users generate highly heterogeneous data, this uniform model often performs poorly for individual users because it fails to capture unique local characteristics.

The article develops and evaluates Personalized Federated Averaging, a framework designed to train a shared initial meta-model that individual users can rapidly specialize to their own local datasets through one or a few quick local gradient updates. The primary objective is to mathematically formulate this personalization objective using model-agnostic meta-learning, establish provable convergence guarantees for general non-convex loss functions, and evaluate the method against existing federated learning standards.

To establish these properties, the authors conduct a theoretical optimization analysis of the algorithm under standard assumptions, including bounded gradient variance, smooth loss functions, and varying degrees of data heterogeneity characterized by statistical metrics such as Total Variation and 1-Wasserstein distances. They also conduct empirical validation on multi-class image classification benchmarks across a network of simulated heterogeneous users, comparing the framework against standard federated averaging and evaluating two practical computational approximations: a first-order variant that omits second-order derivatives and a Hessian-free variant that approximates curvature through gradient differences.

The analysis yields three key findings. First, the proposed personalized federated learning algorithm achieves provable convergence to an approximate stationary point under non-convex objectives, requiring a number of communication rounds proportional to the inverse three-halves power of the target error tolerance. Second, task diversity directly slows convergence as a function of the statistical distance between local user data distributions and the global population average. Third, in experimental benchmarks under severe data heterogeneity, the Hessian-free implementation substantially outperforms both standard federated averaging and first-order meta-learning approximations, achieving average test accuracy improvements of roughly 10 to 13 percentage points on complex classification tasks.

These findings demonstrate that organizations do not need to choose between data privacy and model customization. Training an adaptable base model rather than a static global solution allows edge devices to achieve high individual performance with minimal local computing overhead. Furthermore, because the first-order approximation degrades significantly when adaptation step sizes increase, deploying the Hessian-free approximation provides a reliable balance of computational efficiency and personalization quality across heterogeneous environments.

For practical deployment, organizations should adopt the Hessian-free formulation when building personalized distributed systems and calibrate local batch sizes to mitigate estimation bias. Future work should evaluate the method in scaled production pilots involving real-world edge hardware, asymmetric device capabilities, and non-image domains such as conversational text or clinical datasets.

Confidence in the mathematical convergence guarantees is high, as they build upon established non-convex optimization principles. However, decision-makers should note that the empirical evaluations were conducted on simulated partitions of standard vision datasets rather than live edge deployments, and the performance guarantees assume bounded gradient variances and smooth underlying objective functions.

Massachusetts Institute of TechnologyUniversity of Texas at Austincomputational powerdata pointsfederated learningPersonalized Federated LearningModel-Agnostic Meta-Learninglocal stochastic gradient descent
arXiv: 2002.07948
  • Paper: Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks, Chelsea Finn et al. (2017). This foundational paper introduces the Model-Agnostic Meta-Learning (MAML) framework upon which the source's personalized federated learning formulation directly builds.
  • Paper: Communication-Efficient Learning of Deep Networks from Decentralized Data, H. B. McMahan et al. (2016). This paper establishes the canonical Federated Averaging (FedAvg) algorithm that the source adapts into a personalized meta-learning variant.
  • Paper: Federated Learning with Non-IID Data, Yue Zhao et al. (2018). This work characterizes the severe performance degradation of FedAvg caused by non-IID client data distributions and weight divergence, establishing the core problem that personalized federated learning aims to solve.
  • Paper: Federated Optimization in Heterogeneous Networks, Tian Li et al. (2018). This paper analyzes optimization challenges and non-convex convergence under statistical heterogeneity in federated settings, providing key background for theoretical federated analysis.
  • Paper: Federated Multi-Task Learning, Virginia Smith et al. (2017). This work introduces multi-task formulation for federated learning to handle statistical heterogeneity, framing the alternative paradigm to global model consensus that motivates client personalization.
  • Paper: Federated Learning with Personalization Layers, Manoj Ghuhan Arivazhagan et al. (2019). This paper provides an essential early approach to personalized federated learning by splitting base and personalization layers, highlighting the benefits of local model adaptation.
  • Paper: On the Convergence of FedAvg on Non-IID Data, Xiang Li et al. (2019). This study derives rigorous convergence rates of FedAvg on non-IID data, providing theoretical foundations relevant to the source's convergence analysis.
  • Paper: Federated Learning: Challenges, Methods, and Future Directions, Tian Li et al. (2019). This comprehensive survey outlines the foundational statistical and systems challenges in federated learning, contextualizing the need for personalization methods.
  • Paper: Personalized Federated Learning with Moreau Envelopes, Canh T. Dinh et al. (2020). This work proposes pFedMe using Moreau envelope regularization as an alternative personalized federated learning framework and benchmarks directly against the source's Per-FedAvg algorithm.
  • Paper: Towards Personalized Federated Learning, Alysa Ziying Tan et al. (2021). This survey systematically categorizes personalized federated learning paradigms, evaluating the trade-offs and computational costs of MAML-based personalization approaches established in the source.
  • Paper: Ditto: Fair and Robust Federated Learning Through Personalization, Tian Li et al. (2020). This paper extends personalized federated optimization by demonstrating that local personalization frameworks can simultaneously guarantee robustness against adversarial attacks and ensure fairness across clients.
  • Paper: Federated Learning on Non-IID Data: A Survey, Hangyu Zhu et al. (2021). This survey provides a broader review of techniques handling non-IID data in federated learning, synthesizing meta-learning and personalization strategies alongside other algorithmic interventions.
  • Paper: Model-Contrastive Federated Learning, Qinbin Li et al. (2021). This work develops model-contrastive federated learning to correct local drift under heterogeneous data, offering a complementary strategy to meta-learning-based personalization.
  • Paper: Federated Learning on Non-IID Data Silos: An Experimental Study, Qinbin Li et al. (2021). This empirical benchmark provides a systematic evaluation of federated optimization algorithms across structured non-IID data silos, contextualizing the conditions under which personalized adaptation is necessary.
Cover for Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach

Abstract

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.

Topics

computational powerdata pointsfederated learningmeta-learningprivacy preservinguser information

Keywords

Personalized Federated LearningModel-Agnostic Meta-Learninglocal stochastic gradient descentFedAvg algorithmmulti-task learningnon-convex loss functionstheoretical guarantees

Table of Contents

  • 1 Introduction
  • 2 Personalized Federated Learning via Model-Agnostic Meta-Learning
  • 3 Personalized FedAvg
  • 4 Theoretical Results
  • 4.1 On the Connections of Task Similarity and Distribution Distances
  • 4.2 Convergence Analysis of Per-FedAvg Algorithm
  • 5 Numerical Experiments
  • 6 Conclusion
  • Broader Impact
  • Acknowledgments and Disclosure of Funding
  • References

Knowls

  1. Knowl 1 — Personalized Federated Learning Objective Formulation

    model/method

    In personalized federated learning with nn heterogeneous users, where each user i∈{1,…,n}i \in \{1, \dots, n\} has a local expected loss function fi:Rd→Rf_i: \mathbb{R}^d \to \mathbb{R} over its local data distribution pip_i, the goal is to learn an initial shared meta-model w∈Rdw \in \mathbb{R}^d such that each user achieves low loss after performing one step of local gradient descent using its own dataset. The personalized federated learning problem is formulated within the Model-Agnostic Meta-Learning (MAML) framework as:

    min⁡w∈RdF(w):=1n∑i=1nFi(w),where Fi(w):=fi(w−α∇fi(w))\min_{w \in \mathbb{R}^d} F(w) := \frac{1}{n} \sum_{i=1}^n F_i(w), \quad \text{where } F_i(w) := f_i(w - \alpha \nabla f_i(w))

    Here, α≥0\alpha \ge 0 is the local adaptation stepsize (meta-stepsize). The gradient of each user's meta-loss function FiF_i with respect to the initial model parameter ww is:

    ∇Fi(w)=(I−α∇2fi(w))∇fi(w−α∇fi(w))\nabla F_i(w) = \left( I - \alpha \nabla^2 f_i(w) \right) \nabla f_i(w - \alpha \nabla f_i(w))

    where I∈Rd×dI \in \mathbb{R}^{d \times d} is the identity matrix, and ∇2fi(w)\nabla^2 f_i(w) is the Hessian matrix of fif_i evaluated at ww.

  2. Knowl 2 — Personalized Federated Averaging Algorithm

    algorithm

    The Personalized Federated Averaging (Per-FedAvg) algorithm optimizes the personalized federated learning objective min⁡w∈Rd1n∑i=1nfi(w−α∇fi(w))\min_{w \in \mathbb{R}^d} \frac{1}{n} \sum_{i=1}^n f_i(w - \alpha \nabla f_i(w)). In each round kk, the server randomly selects a subset of active users AkA_k of size ∣Ak∣=rn|A_k| = rn, where r∈(0,1]r \in (0, 1] is the user participation ratio, and broadcasts the current global model wkw_k. Each selected user i∈Aki \in A_k performs τ≥1\tau \ge 1 local stochastic meta-gradient steps using independent mini-batches Di,t,Di,t′,Di,t′′D_{i,t}, D'_{i,t}, D''_{i,t} of sizes D,D′,D′′D, D', D'' sampled from pip_i. The stochastic meta-gradient is computed via a two-stage update: first computing an adapted intermediate parameter w~k+1,ti=wk+1,t−1i−α∇~fi(wk+1,t−1i,Di,t)\tilde{w}_{k+1, t}^i = w_{k+1, t-1}^i - \alpha \tilde{\nabla} f_i(w_{k+1, t-1}^i, D_{i,t}), and then updating with the stochastic gradient ∇~fi(w~k+1,ti,Di,t′)\tilde{\nabla} f_i(\tilde{w}_{k+1, t}^i, D'_{i,t}) premultiplied by the Hessian estimate (I−α∇~2fi(wk+1,t−1i,Di,t′′))(I - \alpha \tilde{\nabla}^2 f_i(w_{k+1, t-1}^i, D''_{i,t})). Active users send their local models back to the server, which computes their average to form wk+1w_{k+1}.

    Input: Initial iterate w0∈Rdw_0 \in \mathbb{R}^d, fraction of active users r∈(0,1]r \in (0, 1], local learning rate β>0\beta > 0, adaptation stepsize α>0\alpha > 0, local steps τ≥1\tau \ge 1, total rounds K≥1K \ge 1, batch sizes D,D′,D′′D, D', D''.
    Output: Final global model wK∈Rdw_K \in \mathbb{R}^d.
    for k=0k = 0 to K−1K - 1 do
        Server chooses a subset of users Ak⊆{1,…,n}A_k \subseteq \{1, \dots, n\} uniformly at random with size ∣Ak∣=rn|A_k| = rn
        Server sends wkw_k to all users i∈Aki \in A_k
        for each user i∈Aki \in A_k in parallel do
            Set local initial iterate wk+1,0i=wkw_{k+1, 0}^i = w_k
            for t=1t = 1 to τ\tau do
                Sample independent mini-batches Di,tD_{i,t}, Di,t′D'_{i,t}, and Di,t′′D''_{i,t} of sizes D,D′,D′′D, D', D'' from pip_i
                Compute stochastic gradient ∇~fi(wk+1,t−1i,Di,t)=1D∑(x,y)∈Di,t∇li(wk+1,t−1i;x,y)\tilde{\nabla} f_i(w_{k+1, t-1}^i, D_{i,t}) = \frac{1}{D} \sum_{(x,y) \in D_{i,t}} \nabla l_i(w_{k+1, t-1}^i; x, y)
                Compute intermediate adapted iterate: w~k+1,ti=wk+1,t−1i−α∇~fi(wk+1,t−1i,Di,t)\tilde{w}_{k+1, t}^i = w_{k+1, t-1}^i - \alpha \tilde{\nabla} f_i(w_{k+1, t-1}^i, D_{i,t})
                Compute stochastic Hessian estimator ∇~2fi(wk+1,t−1i,Di,t′′)=1D′′∑(x,y)∈Di,t′′∇2li(wk+1,t−1i;x,y)\tilde{\nabla}^2 f_i(w_{k+1, t-1}^i, D''_{i,t}) = \frac{1}{D''} \sum_{(x,y) \in D''_{i,t}} \nabla^2 l_i(w_{k+1, t-1}^i; x, y)
                Compute stochastic meta-gradient: ∇~Fi(wk+1,t−1i)=(I−α∇~2fi(wk+1,t−1i,Di,t′′))∇~fi(w~k+1,ti,Di,t′)\tilde{\nabla} F_i(w_{k+1, t-1}^i) = (I - \alpha \tilde{\nabla}^2 f_i(w_{k+1, t-1}^i, D''_{i,t})) \tilde{\nabla} f_i(\tilde{w}_{k+1, t}^i, D'_{i,t})
                Update local model: wk+1,ti=wk+1,t−1i−β∇~Fi(wk+1,t−1i)w_{k+1, t}^i = w_{k+1, t-1}^i - \beta \tilde{\nabla} F_i(w_{k+1, t-1}^i)
            end for
            User ii sends wk+1,τiw_{k+1, \tau}^i to the server
        end for
        Server updates global model: wk+1=1rn∑i∈Akwk+1,τiw_{k+1} = \frac{1}{rn} \sum_{i \in A_k} w_{k+1, \tau}^i
    end for
    return wKw_K

  3. Knowl 3 — First-Order and Hessian-Free Approximations of Per-FedAvg

    model/method

    To reduce the computational burden of computing full second-order derivatives in Per-FedAvg, two efficient approximations can be utilized:

    1. First-Order Per-FedAvg (Per-FedAvg (FO)): Ignores the Hessian term by setting (I−α∇2fi)≈I(I - \alpha \nabla^2 f_i) \approx I. The local stochastic meta-gradient is approximated as:

    ∇~Fi(wk+1,t−1i)≈∇~fi(wk+1,t−1i−α∇~fi(wk+1,t−1i,Di,t), Di,t′)\tilde{\nabla} F_i(w_{k+1, t-1}^i) \approx \tilde{\nabla} f_i\left( w_{k+1, t-1}^i - \alpha \tilde{\nabla} f_i(w_{k+1, t-1}^i, D_{i,t}), \, D'_{i,t} \right)

    1. Hessian-Free Per-FedAvg (Per-FedAvg (HF)): Approximates the Hessian-vector product ∇2fi(u)v\nabla^2 f_i(u) v using a finite difference of gradients with a small perturbation scalar δ>0\delta > 0:

    ∇2fi(u)v≈∇fi(u+δv)−∇fi(u−δv)2δ\nabla^2 f_i(u) v \approx \frac{\nabla f_i(u + \delta v) - \nabla f_i(u - \delta v)}{2\delta}

    By taking v=∇~fi(w~k+1,ti,Di,t′)v = \tilde{\nabla} f_i(\tilde{w}_{k+1, t}^i, D'_{i,t}), Per-FedAvg (HF) avoids explicit Hessian matrix computations while retaining second-order curvature information.

  4. Knowl 4 — Gradient and Hessian Diversity Bounds via Distribution Distances

    theoretical result

    Let fi(w):=Ez∼pi[l(z;w)]f_i(w) := \mathbb{E}_{z \sim p_i}[l(z; w)] be the local loss of user i∈{1,…,n}i \in \{1, \dots, n\} over sample z=(x,y)∈Zz = (x, y) \in \mathcal{Z}, and let f(w):=1n∑i=1nfi(w)f(w) := \frac{1}{n} \sum_{i=1}^n f_i(w) with average distribution p:=1n∑i=1npip := \frac{1}{n} \sum_{i=1}^n p_i. The task diversity parameters γG2\gamma_G^2 and γH2\gamma_H^2, defined such that for all w∈Rdw \in \mathbb{R}^d:

    1n∑i=1n∥∇fi(w)−∇f(w)∥2≤γG2,1n∑i=1n∥∇2fi(w)−∇2f(w)∥2≤γH2\frac{1}{n} \sum_{i=1}^n \|\nabla f_i(w) - \nabla f(w)\|^2 \le \gamma_G^2, \quad \frac{1}{n} \sum_{i=1}^n \|\nabla^2 f_i(w) - \nabla^2 f(w)\|^2 \le \gamma_H^2

    are bounded in terms of probability distribution distances as follows:

    1. Total Variation (TV) Distance: If ∥∇wl(z;w)∥≤B\|\nabla_w l(z; w)\| \le B and ∥∇w2l(z;w)∥≤L\|\nabla_w^2 l(z; w)\| \le L for all z∈Z,w∈Rdz \in \mathcal{Z}, w \in \mathbb{R}^d, then:

    γG2=4B21n∑i=1n∥pi−p∥TV2,γH2=4L21n∑i=1n∥pi−p∥TV2\gamma_G^2 = 4 B^2 \frac{1}{n} \sum_{i=1}^n \|p_i - p\|_{TV}^2, \quad \gamma_H^2 = 4 L^2 \frac{1}{n} \sum_{i=1}^n \|p_i - p\|_{TV}^2

    where ∥q1−q2∥TV:=12∑z∈Z∣q1(z)−q2(z)∣\|q_1 - q_2\|_{TV} := \frac{1}{2} \sum_{z \in \mathcal{Z}} |q_1(z) - q_2(z)|.

    1. 1-Wasserstein Distance: If the sample gradient and Hessian are Lipschitz continuous with respect to zz on metric space (Z,d)(\mathcal{Z}, d) with constants LZL_Z and ρZ\rho_Z (i.e., ∥∇wl(z1;w)−∇wl(z2;w)∥≤LZd(z1,z2)\|\nabla_w l(z_1; w) - \nabla_w l(z_2; w)\| \le L_Z d(z_1, z_2) and ∥∇w2l(z1;w)−∇w2l(z2;w)∥≤ρZd(z1,z2)\|\nabla_w^2 l(z_1; w) - \nabla_w^2 l(z_2; w)\| \le \rho_Z d(z_1, z_2)), and each pip_i has bounded support, then:

    γG2=LZ21n∑i=1nW1(pi,p)2,γH2=ρZ21n∑i=1nW1(pi,p)2\gamma_G^2 = L_Z^2 \frac{1}{n} \sum_{i=1}^n W_1(p_i, p)^2, \quad \gamma_H^2 = \rho_Z^2 \frac{1}{n} \sum_{i=1}^n W_1(p_i, p)^2

    where W1(q1,q2):=inf⁡q∈Q(q1,q2)∫Z×Zd(z1,z2) dq(z1,z2)W_1(q_1, q_2) := \inf_{q \in \mathcal{Q}(q_1, q_2)} \int_{\mathcal{Z} \times \mathcal{Z}} d(z_1, z_2) \, dq(z_1, z_2).

  5. Knowl 5 — Meta-Loss Smoothness and Stochastic Estimator Properties

    theoretical result

    Let each fif_i be twice continuously differentiable, LL-smooth (∥∇fi(w)−∇fi(u)∥≤L∥w−u∥\|\nabla f_i(w) - \nabla f_i(u)\| \le L \|w - u\|), with bounded gradient (∥∇fi(w)∥≤B\|\nabla f_i(w)\| \le B), and ρ\rho-Lipschitz continuous Hessian (∥∇2fi(w)−∇2fi(u)∥≤ρ∥w−u∥\|\nabla^2 f_i(w) - \nabla^2 f_i(u)\| \le \rho \|w - u\|). Let the local stochastic gradient and Hessian estimators have bounded variance E[∥∇li(x,y;w)−∇fi(w)∥2]≤σG2\mathbb{E}[\|\nabla l_i(x,y; w) - \nabla f_i(w)\|^2] \le \sigma_G^2 and E[∥∇2li(x,y;w)−∇2fi(w)∥2]≤σH2\mathbb{E}[\|\nabla^2 l_i(x,y; w) - \nabla^2 f_i(w)\|^2] \le \sigma_H^2. For the meta-loss Fi(w):=fi(w−α∇fi(w))F_i(w) := f_i(w - \alpha \nabla f_i(w)) and F(w):=1n∑i=1nFi(w)F(w) := \frac{1}{n} \sum_{i=1}^n F_i(w) with adaptation stepsize α∈[0,1/L]\alpha \in [0, 1/L]:

    1. Smoothness: Both FiF_i and FF are smooth with Lipschitz constant:

    LF:=4L+αρBL_F := 4L + \alpha \rho B

    1. Estimator Bias and Variance: The stochastic meta-gradient estimator ∇~Fi(w)\tilde{\nabla} F_i(w) computed using independent batches D,D′,D′′D, D', D'' satisfies:

    ∥E[∇~Fi(w)−∇Fi(w)]∥≤2αLσGD\|\mathbb{E}[\tilde{\nabla} F_i(w) - \nabla F_i(w)]\| \le \frac{2\alpha L \sigma_G}{\sqrt{D}}

    E[∥∇~Fi(w)−∇Fi(w)∥2]≤σF2:=12(B2+σG2)(1D′+(αL)2D)(1+σH2α24D′′)−12B2\mathbb{E}[\|\tilde{\nabla} F_i(w) - \nabla F_i(w)\|^2] \le \sigma_F^2 := 12(B^2 + \sigma_G^2) \left( \frac{1}{D'} + \frac{(\alpha L)^2}{D} \right) \left( 1 + \frac{\sigma_H^2 \alpha^2}{4 D''} \right) - 12 B^2

    1. Meta-Gradient Diversity: If 1n∑i=1n∥∇fi(w)−∇f(w)∥2≤γG2\frac{1}{n} \sum_{i=1}^n \|\nabla f_i(w) - \nabla f(w)\|^2 \le \gamma_G^2 and 1n∑i=1n∥∇2fi(w)−∇2f(w)∥2≤γH2\frac{1}{n} \sum_{i=1}^n \|\nabla^2 f_i(w) - \nabla^2 f(w)\|^2 \le \gamma_H^2, then:

    1n∑i=1n∥∇Fi(w)−∇F(w)∥2≤γF2:=3B2α2γH2+192γG2\frac{1}{n} \sum_{i=1}^n \|\nabla F_i(w) - \nabla F(w)\|^2 \le \gamma_F^2 := 3 B^2 \alpha^2 \gamma_H^2 + 192 \gamma_G^2

  6. Knowl 6 — Non-Convex Convergence of Personalized FedAvg

    theoretical result

    Suppose each fif_i is bounded below with min⁡wfi(w)>−∞\min_w f_i(w) > -\infty (let F∗=min⁡wF(w)F^* = \min_w F(w)), LL-smooth, with gradient norm bounded by BB, ρ\rho-Lipschitz continuous Hessian, and bounded stochastic gradient and Hessian variances σG2,σH2\sigma_G^2, \sigma_H^2. Let the user gradient and Hessian diversities be bounded by γG2\gamma_G^2 and γH2\gamma_H^2. If Per-FedAvg is run for KK communication rounds with τ\tau local steps per round, local learning rate β≤110τLF\beta \le \frac{1}{10 \tau L_F}, meta-stepsize α∈(0,1/L]\alpha \in (0, 1/L], and active user fraction r∈(0,1]r \in (0, 1], then:

    1τK∑k=0K−1∑t=0τ−1E[∥∇F(wˉk+1,t)∥2]≤4(F(w0)−F∗)βτK+O(1)(βLF(1+βLFτ(τ−1))σF2+βLFγF2(1−rr(n−1)+βLFτ(τ−1))+α2L2σG2D)\frac{1}{\tau K} \sum_{k=0}^{K-1} \sum_{t=0}^{\tau-1} \mathbb{E}[\|\nabla F(\bar{w}_{k+1, t})\|^2] \le \frac{4(F(w_0) - F^*)}{\beta \tau K} + \mathcal{O}(1) \left( \beta L_F \left(1 + \beta L_F \tau(\tau - 1)\right) \sigma_F^2 + \beta L_F \gamma_F^2 \left( \frac{1-r}{r(n-1)} + \beta L_F \tau(\tau-1) \right) + \frac{\alpha^2 L^2 \sigma_G^2}{D} \right)

    where wˉk+1,t:=1rn∑i∈Akwk+1,ti\bar{w}_{k+1, t} := \frac{1}{rn} \sum_{i \in A_k} w_{k+1, t}^i is the average of active users' local iterates at step tt of round kk, LF=4L+αρBL_F = 4L + \alpha \rho B, and σF2\sigma_F^2 and γF2\gamma_F^2 are the meta-gradient estimator variance and diversity bounds.

  7. Knowl 7 — Communication Complexity of Per-FedAvg for First-Order Stationary Points

    theoretical result

    For non-convex objective function F(w)F(w), an ϵ\epsilon-approximate First-Order Stationary Point (FOSP) is a random vector wϵ∈Rdw_\epsilon \in \mathbb{R}^d satisfying E[∥∇F(wϵ)∥2]≤ϵ\mathbb{E}[\|\nabla F(w_\epsilon)\|^2] \le \epsilon.

    Setting the number of local updates per round as τ=O(ϵ−1/2)\tau = \mathcal{O}(\epsilon^{-1/2}), the number of communication rounds with the central server as K=O(ϵ−3/2)K = \mathcal{O}(\epsilon^{-3/2}), and the local stepsize as β=ϵ\beta = \epsilon, Per-FedAvg finds an O(ϵ+α2σG2D)\mathcal{O}\left(\epsilon + \frac{\alpha^2 \sigma_G^2}{D}\right)-first-order stationary point of FF.

    To eliminate the bias term α2σG2D\frac{\alpha^2 \sigma_G^2}{D} and achieve an ϵ\epsilon-stationary point (E[∥∇F(wϵ)∥2]≤ϵ\mathbb{E}[\|\nabla F(w_\epsilon)\|^2] \le \epsilon), the algorithm requires setting either:

    1. The local batch size D=O(ϵ−1)D = \mathcal{O}(\epsilon^{-1}), or
    2. The adaptation stepsize α=O(ϵ1/2)\alpha = \mathcal{O}(\epsilon^{1/2}).

    Under either condition, the total number of communication rounds between the server and users is K=O(ϵ−3/2)K = \mathcal{O}(\epsilon^{-3/2}).

  8. Knowl 8 — Experimental Setup for Personalized Federated Learning Evaluation

    experimental setup

    The empirical performance of Personalized FedAvg is evaluated on multi-class image classification using MNIST and CIFAR-10 across n=50n = 50 users under non-i.i.d. data distributions:

    1. Data Partitioning:
    • For MNIST, a=196a = 196; for CIFAR-10, a=68a = 68.
    • 25 users (half of the network) each hold aa images from each of the first 5 classes.
    • 25 users (the other half) each hold a/2a/2 images from only one of the first 5 classes and 2a2a images from only one of the remaining 5 classes.
    • For a higher-heterogeneity setting ("diff. hetero."), the a/2a/2 images from the first 5 classes are completely removed for the second half of users.
    • Test data is partitioned across users following identical proportions.
    1. Model Architecture and Hyperparameters:
    • Fully connected neural network with two hidden layers of sizes 80 and 60 with Exponential Linear Unit (ELU) activations.
    • Communication rounds K=1000K = 1000, user sampling ratio r=0.2r = 0.2 (10 active users per round), batch sizes D=D′=40D = D' = 40, learning rate β=0.001\beta = 0.001.
    • Evaluation protocol: Test accuracy is averaged across all users after each user takes one step of local stochastic gradient descent on its test data starting from the trained global model.

  9. Knowl 9 — Test Accuracy Comparison of FedAvg and Personalized FedAvg Variants

    data/table

    The table below compares the average test accuracy (with 95% confidence intervals) across all users for FedAvg (updated by 1 step of SGD at test time), First-Order Personalized FedAvg (Per-FedAvg (FO)), and Hessian-Free Personalized FedAvg (Per-FedAvg (HF)) on MNIST and CIFAR-10 across different numbers of local steps τ\tau, adaptation stepsizes α\alpha, and data heterogeneity conditions.

    Dataset Parameters FedAvg + update Per-FedAvg (FO) Per-FedAvg (HF)
    MNIST τ=10,α=0.01\tau = 10, \alpha = 0.01 75.96% ±\pm 0.02% 78.00% ±\pm 0.02% 79.85% ±\pm 0.02%
    τ=4,α=0.01\tau = 4, \alpha = 0.01 60.18% ±\pm 0.02% 64.55% ±\pm 0.02% 70.94% ±\pm 0.03%
    CIFAR-10 τ=10,α=0.001\tau = 10, \alpha = 0.001 40.49% ±\pm 0.07% 46.98% ±\pm 0.10% 50.44% ±\pm 0.15%
    τ=4,α=0.001\tau = 4, \alpha = 0.001 38.38% ±\pm 0.07% 34.04% ±\pm 0.08% 43.73% ±\pm 0.11%
    τ=4,α=0.01\tau = 4, \alpha = 0.01 35.97% ±\pm 0.17% 25.32% ±\pm 0.18% 46.32% ±\pm 0.12%
    τ=4,α=0.01\tau = 4, \alpha = 0.01, diff. hetero. 58.59% ±\pm 0.11% 37.71% ±\pm 0.23% 71.25% ±\pm 0.05%

    Key takeaways:

    • Per-FedAvg (HF) outperforms FedAvg and Per-FedAvg (FO) across all configurations, reaching up to 79.85% on MNIST and 71.25% on CIFAR-10 (under high heterogeneity).
    • Increasing the adaptation stepsize from α=0.001\alpha = 0.001 to α=0.01\alpha = 0.01 at τ=4\tau = 4 on CIFAR-10 improves Per-FedAvg (HF) accuracy from 43.73% to 46.32%, but degrades Per-FedAvg (FO) accuracy from 34.04% down to 25.32% because the first-order approximation errors increase with larger α\alpha.
    • Under extreme data heterogeneity ("diff. hetero."), Per-FedAvg (HF) achieves 71.25% accuracy compared to 58.59% for FedAvg + update and 37.71% for Per-FedAvg (FO).

Coverage note — None was omitted; all key theoretical bounds, algorithmic formulations, distribution distance connections, experimental setups, and empirical findings are covered.

References

  1. 1.J. Konečnỳ, H. B. McMahan, F. X. Yu, P. Richtárik, A. T. Suresh, and D. Bacon, “Federated learning: Strategies for improving communication efficiency,” arXiv preprint arXiv:1610.05492, 2016.
  2. 2.C. Finn, P. Abbeel, and S. Levine, “Model-agnostic meta-learning for fast adaptation of deep networks,” in Proceedings of the 34th International Conference on Machine Learning, (Sydney, Australia), 06–11 Aug 2017.
  3. 3.B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-Efficient Learning of Deep Networks from Decentralized Data,” in Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, vol. 54 of Proceedings of Machine Learning Research, (Fort Lauderdale, FL, USA), pp. 1273–1282, PMLR, 20–22 Apr 2017.
  4. 4.P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings, et al., “Advances and open problems in federated learning,” arXiv preprint arXiv:1912.04977, 2019.
  5. 5.T. Li, A. K. Sahu, A. Talwalkar, and V. Smith, “Federated learning: Challenges, methods, and future directions,” IEEE Signal Process. Mag., vol. 37, no. 3, pp. 50–60, 2020.
  6. 6.J. C. Duchi, M. I. Jordan, and M. J. Wainwright, “Privacy aware learning,” Journal of the ACM (JACM), vol. 61, no. 6, p. 38, 2014.
  7. 7.H. B. McMahan, D. Ramage, K. Talwar, and L. Zhang, “Learning differentially private recurrent language models,” arXiv preprint arXiv:1710.06963, 2017.
  8. 8.N. Agarwal, A. T. Suresh, F. X. X. Yu, S. Kumar, and B. McMahan, “cpsgd: Communication-efficient and differentially-private distributed sgd,” in Advances in Neural Information Processing Systems, pp. 7564–7575, 2018.
  9. 9.W. Zhu, P. Kairouz, B. McMahan, H. Sun, and W. Li, “Federated heavy hitters discovery with differential privacy,” in International Conference on Artificial Intelligence and Statistics, pp. 3837–3847, 2020.
  10. 10.A. Reisizadeh, A. Mokhtari, H. Hassani, A. Jadbabaie, and R. Pedarsani, “Fedpaq: A communication-efficient federated learning method with periodic averaging and quantization,” in International Conference on Artificial Intelligence and Statistics, pp. 2021–2031, 2020.
  11. 11.X. Dai, X. Yan, K. Zhou, K. K. Ng, J. Cheng, and Y. Fan, “Hyper-sphere quantization: Communication-efficient sgd for federated learning,” arXiv preprint arXiv:1911.04655, 2019.
  12. 12.D. Basu, D. Data, C. Karakus, and S. Diggavi, “Qsparse-local-sgd: Distributed sgd with quantization, sparsification and local computations,” in Advances in Neural Information Processing Systems, pp. 14668–14679, 2019.
  13. 13.Z. Li, D. Kovalev, X. Qian, and P. Richtárik, “Acceleration for compressed gradient descent in distributed and federated optimization,” arXiv preprint arXiv:2002.11364, 2020.
  14. 14.S. U. Stich, “Local sgd converges fast and communicates little,” arXiv preprint arXiv:1805.09767, 2018.
  15. 15.J. Wang and G. Joshi, “Cooperative sgd: A unified framework for the design and analysis of communication-efficient sgd algorithms,” arXiv preprint arXiv:1808.07576, 2018.
  16. 16.F. Zhou and G. Cong, “On the convergence properties of a k-step averaging stochastic gradient descent algorithm for nonconvex optimization,” in Proceedings of the 27th International Joint Conference on Artificial Intelligence, pp. 3219–3227, 2018.
  17. 17.T. Lin, S. U. Stich, K. K. Patel, and M. Jaggi, “Don’t use large mini-batches, use local SGD,” in 8th International Conference on Learning Representations, ICLR, 2020.
  18. 18.Y. Zhao, M. Li, L. Lai, N. Suda, D. Civin, and V. Chandra, “Federated learning with non-iid data,” arXiv preprint arXiv:1806.00582, 2018.
  19. 19.A. K. Sahu, T. Li, M. Sanjabi, M. Zaheer, A. Talwalkar, and V. Smith, “On the convergence of federated optimization in heterogeneous networks,” arXiv preprint arXiv:1812.06127, 2018.
  20. 20.S. P. Karimireddy, S. Kale, M. Mohri, S. J. Reddi, S. U. Stich, and A. T. Suresh, “Scaffold: Stochastic controlled averaging for on-device federated learning,” arXiv preprint arXiv:1910.06378, 2019.
  21. 21.F. Haddadpour and M. Mahdavi, “On the convergence of local descent methods in federated learning,” arXiv preprint arXiv:1910.14425, 2019.
  22. 22.X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the convergence of fedavg on non-iid data,” arXiv preprint arXiv:1907.02189, 2019.
  23. 23.A. K. R. Bayoumi, K. Mishchenko, and P. Richtarik, “Tighter theory for local sgd on identical and heterogeneous data,” in International Conference on Artificial Intelligence and Statistics, pp. 4519–4529, 2020.
  24. 24.A. Antoniou, H. Edwards, and A. Storkey, “How to train your MAML,” in International Conference on Learning Representations, 2019.
  25. 25.Z. Li, F. Zhou, F. Chen, and H. Li, “Meta-SGD: Learning to learn quickly for few-shot learning,” arXiv preprint arXiv:1707.09835, 2017.
  26. 26.E. Grant, C. Finn, S. Levine, T. Darrell, and T. Griffiths, “Recasting gradient-based meta-learning as hierarchical bayes,” in International Conference on Learning Representations, 2018.
  27. 27.A. Nichol, J. Achiam, and J. Schulman, “On first-order meta-learning algorithms,” arXiv preprint arXiv:1803.02999, 2018.
  28. 28.L. Zintgraf, K. Shiarli, V. Kurin, K. Hofmann, and S. Whiteson, “Fast context adaptation via meta-learning,” in Proceedings of the 36th International Conference on Machine Learning, pp. 7693–7702, 2019.
  29. 29.H. S. Behl, A. G. Baydin, and P. H. S. Torr, “Alpha MAML: adaptive model-agnostic meta-learning,” 2019.
  30. 30.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, pp. 1534–1544, Curran Associates, Inc., 2019.
  31. 31.A. Fallah, A. Mokhtari, and A. Ozdaglar, “On the convergence theory of gradient-based model-agnostic meta-learning algorithms,” in International Conference on Artificial Intelligence and Statistics, pp. 1082–1092, 2020.
  32. 32.F. Chen, M. Luo, Z. Dong, Z. Li, and X. He, “Federated meta-learning with fast convergence and efficient communication,” arXiv preprint arXiv:1802.07876, 2018.
  33. 33.Y. Jiang, J. Konečnỳ, K. Rush, and S. Kannan, “Improving federated learning personalization via model agnostic meta learning,” arXiv preprint arXiv:1909.12488, 2019.
  34. 34.T. Li, M. Sanjabi, and V. Smith, “Fair resource allocation in federated learning,” arXiv preprint arXiv:1905.10497, 2019.
  35. 35.S. Lin, G. Yang, and J. Zhang, “A collaborative learning framework via federated meta-learning,” arXiv preprint arXiv:2001.03229, 2020.
  36. 36.M. Khodak, M.-F. F. Balcan, and A. S. Talwalkar, “Adaptive gradient-based meta-learning methods,” in Advances in Neural Information Processing Systems, pp. 5915–5926, 2019.
  37. 37.J. Li, M. Khodak, S. Caldas, and A. Talwalkar, “Differentially private meta-learning,” arXiv preprint arXiv:1909.05830, 2019.
  38. 38.V. Smith, C.-K. Chiang, M. Sanjabi, and A. S. Talwalkar, “Federated multi-task learning,” in Advances in Neural Information Processing Systems, pp. 4424–4434, 2017.
  39. 39.F. Hanzely and P. Richtárik, “Federated learning of a mixture of global and local models,” arXiv preprint arXiv:2002.05516, 2020.
  40. 40.Y. Deng, M. M. Kamani, and M. Mahdavi, “Adaptive personalized federated learning,” arXiv preprint arXiv:2003.13461, 2020.
  41. 41.E. del Barrio, E. Giné, and C. Matrán, “Central limit theorems for the wasserstein distance between the empirical and the true distributions,” Annals of Probability, pp. 1009–1071, 1999.
  42. 42.Y. Arjevani, Y. Carmon, J. C. Duchi, D. J. Foster, N. Srebro, and B. Woodworth, “Lower bounds for non-convex stochastic optimization,” arXiv preprint arXiv:1912.02365, 2019.
  43. 43.Y. LeCun, “The mnist database of handwritten digits,” http://yann. lecun. com/exdb/mnist/, 1998.
  44. 44.A. Krizhevsky, G. Hinton, et al., “Learning multiple layers of features from tiny images,” 2009.
  45. 45.J. Langelaar, “Mnist neural network training and testing,” MATLAB Central File Exchange, 2019.
  46. 46.C. Villani, Optimal transport: old and new, vol. 338. Springer Science & Business Media, 2008.

Citation

MLA
Fallah, A., et al. “Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach”. Advances in Neural Information Processing Systems, vol. 33, 2020, pp. 3557–68, https://proceedings.neurips.cc/paper_files/paper/2020/file/24389bfe4fe2eba8bf9aa9203a44cdad-Paper.pdf.
APA
Fallah, A., Mokhtari, A., & Ozdaglar, A. (2020). Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach. Advances in Neural Information Processing Systems, 33, 3557–3568. https://proceedings.neurips.cc/paper_files/paper/2020/file/24389bfe4fe2eba8bf9aa9203a44cdad-Paper.pdf
Chicago
Fallah, A., A. Mokhtari, and A. Ozdaglar. 2020. “Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach”. Advances in Neural Information Processing Systems 33: 3557–68. https://proceedings.neurips.cc/paper_files/paper/2020/file/24389bfe4fe2eba8bf9aa9203a44cdad-Paper.pdf.
Harvard
Fallah, A., Mokhtari, A. and Ozdaglar, A. (2020) “Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach”, Advances in Neural Information Processing Systems. Curran Associates, Inc., pp. 3557–3568. Available at: https://proceedings.neurips.cc/paper_files/paper/2020/file/24389bfe4fe2eba8bf9aa9203a44cdad-Paper.pdf.
Vancouver
1. Fallah A, Mokhtari A, Ozdaglar A (2020) Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach. In: Advances in Neural Information Processing Systems. Curran Associates, Inc., pp 3557–3568

BibTeX

@inproceedings{fallah2020personalized,
  title = {Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach},
  author = {Fallah, Alireza and Mokhtari, Aryan and Ozdaglar, Asuman},
  year = {2020},
  booktitle = {Advances in Neural Information Processing Systems},
  publisher = {Curran Associates, Inc.},
  volume = {33},
  pages = {3557-3568},
  url = {https://proceedings.neurips.cc/paper_files/paper/2020/file/24389bfe4fe2eba8bf9aa9203a44cdad-Paper.pdf}
}
Metadata:DOI registry

Access the Paper

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

Open PDF
License: Authors