[Paper Review] Federated Learning under Importance Sampling

Elsa Rizk, Stefan Vlaski+1 more|arXiv (Cornell University)|Dec 14, 2020

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes a non-uniform sampling strategy in federated learning using importance sampling for both agents and local data, improving convergence by reducing variance. The method adaptively selects agents and data points based on local data variability and model heterogeneity, achieving O(μ) convergence error under asynchronous, non-IID settings with partial participation.

ABSTRACT

Federated learning encapsulates distributed learning strategies that are managed by a central unit. Since it relies on using a selected number of agents at each iteration, and since each agent, in turn, taps into its local data, it is only natural to study optimal sampling policies for selecting agents and their data in federated learning implementations. Usually, only uniform sampling schemes are used. However, in this work, we examine the effect of importance sampling and devise schemes for sampling agents and data non-uniformly guided by a performance measure. We find that in schemes involving sampling without replacement, the performance of the resulting architecture is controlled by two factors related to data variability at each agent, and model variability across agents. We illustrate the theoretical findings with experiments on simulated and real data and show the improvement in performance that results from the proposed strategies.

Motivation & Objective

• To address the performance limitations of uniform sampling in federated learning by introducing non-uniform sampling strategies for agents and local data.
• To analyze the impact of importance sampling on convergence in asynchronous, non-IID, and partially participating federated learning systems.
• To derive theoretical bounds on the error variance of the stochastic gradient under non-uniform sampling, linking performance to data and model heterogeneity.
• To design adaptive sampling policies that minimize the expected error in the presence of data and model variability across agents.
• To validate the theoretical findings with experiments on both simulated and real-world datasets, demonstrating improved convergence speed and accuracy.

Proposed method

• Proposes a two-layer importance sampling scheme: non-uniform selection of agents via distribution πₖ and non-uniform sampling of local data at each agent via πₙ⁽ᵏ⁾.
• Derives theoretical convergence bounds for asynchronous SGD in federated learning under non-uniform sampling, showing error scales as O(μ) when sampling is optimized.
• Introduces a performance metric based on data variability per agent and model variability across agents to guide sampling distribution design.
• Uses a recursive error propagation model to bound the expected squared norm of the gradient noise term qᵢ, linking it to local data variance and model drift.
• Applies Jensen’s inequality and Lemma 4 to simplify the error bound expression, enabling analysis of the impact of sampling probabilities on convergence.
• Derives adaptive sampling probabilities that minimize the upper bound on the expected gradient noise, leading to faster convergence.

Figure 1: MSD plots of the regression problem: blue curve is the standard mini-batch implementation, green curve is the importance sampling implementation with the true probabilities, red curve is the importance sampling implementation with approximate probabilities ( 49 )–( 50 ), purple curve is th

Experimental results

Research questions

• RQ1 How does non-uniform sampling of agents and local data affect the convergence rate of federated learning under non-IID data and asynchronous operation?
• RQ2 What are the key factors controlling the performance of importance sampling in federated learning, particularly in terms of data and model heterogeneity?
• RQ3 Can adaptive sampling strategies based on local data and model variability reduce the expected gradient noise and improve convergence?
• RQ4 What theoretical bounds can be established for the expected squared norm of the stochastic gradient under non-uniform sampling in asynchronous federated learning?
• RQ5 How do the proposed sampling schemes compare to uniform sampling in terms of convergence speed and final model accuracy on real and simulated data?

Key findings

• The convergence error of the proposed importance sampling scheme scales as O(μ), which is an improvement over uniform sampling in non-IID and asynchronous settings.
• Performance is governed by two main factors: data variability within each agent and model variability across agents, both of which influence the optimal sampling distribution.
• Adaptive sampling based on local data variance and model drift significantly reduces the expected gradient noise, leading to faster convergence.
• Theoretical analysis shows that the expected squared norm of the gradient noise term qᵢ is bounded by O(μ), with constants depending on data distribution and sampling probabilities.
• Experiments on simulated and real data confirm that the proposed importance sampling strategy outperforms uniform sampling in terms of convergence speed and final model accuracy.

Figure 2: Testing error plots of the classification problem.