Optimal Client Sampling for Federated Learning

Wenlin Chen; Samuel Horvath; Peter Richtarik

Source-linked AI summary

Optimal Client Sampling for Federated Learning

Wenlin Chen, Samuel Horvath, Peter Richtarik

arXiv:2010.13723v3Oct 26, 2020cs.LGcs.DC

TL;DR

Federated learning communication can be a primary bottleneck, so the paper develops a privacy-preserving strategy that restricts communication to clients with important updates. Its optimal sampling formulation and aggregation-only approximation perform close to full participation and better than uniform sampling, while remaining compatible with other communication-reduction methods.

Problem
Client-master communication is a primary federated-learning bottleneck, motivating methods that reduce communicated bandwidth under limited client participation.
Method
The paper selects clients adaptively using update norms, derives optimal participation probabilities, and approximates them using only secure aggregation.
Results
The method is superior to uniform sampling and close to full participation for DSGD and FedAvg, including strong communication-budget results.
Takeaways & Limitations
Importance-based client sampling can reduce federated-learning communication costs while preserving performance near full participation and supporting secure aggregation.
Takeaways & Limitations
Under high communication latency, the method may not reduce real communication time because it does not account for client computation speed, bandwidth, or latency constraints.

Abstract

from arXiv · show

It is well understood that client-master communication can be a primary bottleneck in Federated Learning. In this work, we address this issue with a novel client subsampling scheme, where we restrict the number of clients allowed to communicate their updates back to the master node. In each communication round, all participating clients compute their updates, but only the ones with "important" updates communicate back to the master. We show that importance can be measured using only the norm of the update and give a formula for optimal client participation. This formula minimizes the distance between the full update, where all clients participate, and our limited update, where the number of participating clients is restricted. In addition, we provide a simple algorithm that approximates the optimal formula for client participation, which only requires secure aggregation and thus does not compromise client privacy. We show both theoretically and empirically that for Distributed SGD (DSGD) and Federated Averaging (FedAvg), the performance of our approach can be close to full participation and superior to the baseline where participating clients are sampled uniformly. Moreover, our approach is orthogonal to and compatible with existing methods for reducing communication overhead, such as local methods and communication compression methods.

1 Introduction

Federated learning communication is constrained by client bandwidth and participation limits, motivating selective client sampling. The paper proposes adaptive importance sampling with privacy-preserving aggregation, convergence guarantees, and compatibility with other communication-reduction methods.

Motivation: Communication Bottleneck in Federated Learning: Client-master communication can bottleneck federated learning because end-user connections are slower, costly, unreliable, and bandwidth-limited.System constraints also limit the number of clients participating in each communication round.
Related approaches: Local-update methods reduce communication frequency, while compression methods reduce the bits per transmitted update through quantization or sparsification.FedAvg is a representative local-update method, whereas compression transmits C(g) instead of the full update vector g.
Client sampling: Adaptive client sampling targets more informative updates, allowing partial participation to reduce communicated bits compared with ignoring clients uniformly.The approach is motivated by the observation that clients can produce updates with different informativeness in a communication round.
Contributions: The paper derives optimal sampling probabilities that minimize master-update variance for any client budget m, extending prior results restricted to m = 1.The strategy is obtained by applying importance-sampling results to the federated learning setting.
Contributions: An aggregation-only approximation identifies important clients while supporting secure aggregation and stateless clients without exposing individual updates.The paper presents this as a principled importance client-sampling strategy compatible with both privacy and stateless-client requirements.
Contributions: For DSGD and FedAvg, the method has convergence guarantees, outperforms uniform sampling, and can approach full-participation performance.The sampling strategy is also compatible with local updates and communication compression methods.

2 Smart Client Sampling for Reducing Communication

The paper formulates client selection as an optimal partial-participation problem and derives a norm-based sampling strategy under a communication budget. An approximate implementation uses iterative probability rescaling and secure aggregation while communicating only update norms before selecting updates.

Optimal sampling: The sampling problem asks which procedure minimizes the master-update variance when the expected number of communicating clients is bounded by m.The framework permits arbitrary random subsets, with client inclusion probabilities defining expected participation.
Algorithm: The optimal client-sampling algorithm has clients compute updates, send their norms, receive probabilities, and transmit updates probabilistically with inverse-probability scaling.The procedure uses the norm of each update to compute probabilities and preserves an unbiased estimator through scaling by 1/p_i.
Optimal sampling: Independent sampling minimizes the relevant variance bound among samplings with fixed client inclusion probabilities.For independent sampling, the bound is attained with v_i = 1 − p_i.
Optimal sampling: The closed-form solution selects client probabilities from update norms under the constraint that expected participation does not exceed m.The resulting strategy can be interpreted as uniform sampling with an update-dependent effective sample size between m and n.
Optimality scope: The method is optimal when only client norms are communicated as extra information, whereas full general-case optimality would require revealing all full updates.This establishes the privacy-aware scope of the optimality claim.
Algorithm: The approximate algorithm iteratively rescales probabilities after truncation and requires secure aggregation, making it compatible with privacy-preserving FL infrastructure.Its communication includes extra synchronous rounds; these may be less useful under metrics other than total communication cost.

3 Convergence Guarantees

The convergence analysis covers DSGD and FedAvg under convex and non-convex assumptions and positions optimal client sampling between full and uniform partial participation. The position depends on the distribution of client updates, with edge cases recovering the corresponding benchmark guarantees.

Analysis setup: The analysis compares optimal client sampling with full participation and independent uniform sampling of sample size m.The convergence bounds are matched to existing literature for direct comparison.
Assumptions: For non-convex objectives, bounded sampling variance requires an assumption on similarity among local gradients.The paper notes that counterexamples can diverge without bounded sampling variance, while convex objectives do not require this assumption.
Interpretation: The improvement factor α_k lies between 0 and 1: α_k = 0 matches full participation, while α_k = 1 matches uniform sampling with m updates.Its value depends on the client updates, and the worst case cannot theoretically outperform uniform sampling for the given update structure.
DSGD: DSGD convergence guarantees under optimal sampling lie between the best-known full-participation and uniform-sampling rates.When γ_k = 1, the strongly-convex analysis recovers the full-participation rate; when γ_k = m/n, it recovers the uniform-sampling rate.
FedAvg: FedAvg convergence guarantees likewise lie between full and uniform partial participation, with γ_k = 1 and γ_k = m/n recovering the respective edge-case complexity guarantees.The stated FedAvg results cover strongly-convex and non-convex settings under the paper's smoothness and oracle assumptions.

4 Related Work

Prior federated-learning client-sampling methods use heuristic or partial information and can conflict with privacy requirements. The proposed method instead provides principled variance-minimizing sampling compatible with core privacy requirements and flexible participation budgets.

Existing efficient client-sampling methods rely on heuristics, historical losses, or partial information as proxies for optimal sampling.
Some prior selection strategies violate federated-learning privacy requirements by revealing client identities or lacking secure aggregation and stateless clients.
The proposed strategy is principled and optimal because it minimizes the variance of the master update while remaining compatible with core privacy requirements.
Importance sampling research establishes optimal sampling proportional to per-sample gradient norms, but obtaining that distribution can cost as much computation as the full gradient.
The federated-learning method adapts importance-sampling results to communication reduction, avoiding their computational limitation and allowing any participation budget m < n.

5 Experiments

The experiments evaluate optimal client sampling on modified FEMNIST and Shakespeare datasets against full participation and uniform sampling. Across datasets, optimal sampling approaches full-participation performance with fewer communicated bits and outperforms uniform sampling.

Experimental setup: The evaluation uses LEAF FEMNIST and Shakespeare datasets, comparing optimal sampling with full participation and uniform partial participation.The experiments use TensorFlow Federated and report validation accuracy and local training loss against communication rounds and client-to-master bits.
Experimental setup: Three unbalanced FEMNIST training sets are created by removing or reducing examples on selected clients to test settings favoring optimal sampling.The unchanged EMNIST validation set contains 40,832 images, and experiments sample n = 32 clients per round.
Experimental setup: The approximation used in the experiments produces identical results to the optimal algorithm, with negligible performance loss attributed to approximation.Consequently, the reported FedAvg results use Algorithm 2.
Results: Optimal sampling nearly reaches full participation with fewer than 10% of available clients communicating when m = 2 or 3.Increasing m improves optimal-sampling performance and quickly makes it almost identical to full participation in communication rounds.
Results: m = 6 for n = 32 and m = 12 for n = 128 match full participation in validation accuracy by communication rounds, motivating the conjecture m = O(√n).This is stated as a conjecture based on the Shakespeare experiments.
Results: On FEMNIST Dataset 1, optimal sampling with m = 3 reaches around 85% accuracy after 26 × 10^8 bits, while full and uniform sampling remain below 40%.Full participation requires more than 29 × 10^8 bits—8× more—to reach the same 85% accuracy; the other datasets show similar qualitative results.
Results: Optimal sampling supports larger learning rates than uniform sampling, and it outperforms uniform sampling even on balanced Federated CIFAR100.The larger-step-size observation is reported for Sections 5.2 and 5.3; the CIFAR100 result appears in Appendix G.

6 Conclusion and Future Work

The paper presents principled optimal client sampling for federated learning and reports convergence guarantees and empirical performance close to full participation. It also identifies high communication latency as a setting where the method may not reduce real communication time.

Conclusion: The proposed strategy computes optimal client sampling from aggregated update norms and is compatible with stateless clients and secure aggregation.The paper reports convergence guarantees for DSGD and FedAvg under relaxed assumptions.
Conclusion: Empirical results show that the method is superior to uniform sampling and close to full participation.These results are reported on federated datasets from the LEAF database and corroborate the theoretical analysis.
Future Work: The approach can be combined with communication compression methods to further reduce communicated-update sizes.This combination is proposed as a direction for future work.
Future Work: In high-latency settings, the method may not reduce real communication time.Future extensions could account for client computational speed, network bandwidth, and communication latency.

A Proof of Lemma 1

The proof represents client inclusion with indicator variables, computes the mean and variance of the sampled aggregate, and bounds the variance using a matrix inequality. It then shows that independent sampling attains the bound with equality.

Proof: Indicator variables encode whether individual clients or client pairs belong to the sampled set.Their expectations are identified with the corresponding marginal and joint inclusion probabilities.
Proof: The proof rewrites the variance of the sampled aggregate in a form suitable for establishing an upper bound.The vector of all ones and the matrix of weighted client vectors are introduced in this derivation.
Proof: The matrix assumption bounds the sampling covariance by a diagonal matrix involving client probabilities and variance terms.The proof compares diagonal elements to obtain v_i ≥ 1 − p_i.
Proof: For independent sampling, the covariance matrix is diagonal and v_i = 1 − p_i.Independent sampling includes client i by an independent coin flip with probability p_i.

B The Improvement Factor for Optimal Client Sampling

This section derives a closed-form solution for the client-sampling optimization problem and connects it to the improvement factor. Subsequent bounds support the paper’s convergence analysis under smoothness, convexity, and step-size assumptions.

Optimization: The sampling probabilities are chosen by minimizing the variance term in the client-sampling objective.The optimization imposes probability bounds 0 ≤ p_i^k ≤ 1 and a participation budget.
Optimization: Lemma 20 states that the optimization problem has a closed-form solution.The solution is characterized using ordered values and an index satisfying the participation constraints.
Optimization: KKT conditions establish the closed-form solution for the relaxed optimization problem and identify it with the original problem’s solution.The support of the feasible set is convex and the constraints are linear.
Improvement Factor: Substituting the optimal probabilities into the variance expression yields the improvement factor.The proof explicitly combines the resulting expression with the preceding bound.
Convergence Analysis: The convergence derivations use smoothness, strong convexity, bounded-gradient assumptions, and step-size restrictions to obtain expected bounds.The stated proofs take conditional and full expectations and conclude with the desired results.
Convergence Analysis: The paper’s auxiliary inequality combines smoothness and strong convexity through a relaxed triangle inequality, with L ≥ μ.The resulting lemma is used as an ingredient in the convergence proofs.

D.1 Proof of Theorem 17

The theorem proofs analyze the master update by separating sampling effects from local-update drift and bounding both through conditional expectations. The derivations then impose smoothness, moment, and step-size conditions to obtain the stated expected guarantees.

Master Update: The master update is written explicitly and analyzed by taking expectations conditioned on the preceding iterate and sampling outcome.The proofs use sums over clients and local-update rounds.
Local-Update Drift: The analysis separates the drift caused by local client updates from other terms in the convergence bound.The drift is denoted E and is bounded using local-update properties.
Bounding Terms: Smoothness, independence, zero mean, bounded second moments, and Jensen’s inequality are combined to bound the update terms.The parameter γ is defined from the sampling-related quantities in the derivation.
Assumptions: The proof requires step-size restrictions, including η ≤ 1/[8L(2 + M/R)].These restrictions are used to derive intermediate inequalities and rearrange the final bound.
Conclusion: Combining the bounds and taking full expectations yields the theorem’s desired result.The proof concludes after simplifying the iterate and incorporating the bounded terms.

E A Sketch of Results on Partial Participation

The analysis extends to rounds where only a subset of clients is available, modeling availability probabilistically and sampling only from available clients. The resulting variance decomposition supports an analogous convergence analysis.

Client availability is modeled by a known distribution Q, with Q_k denoting the available subset in round k and q_i the probability that client i is available.
Applying the tower property of expectation yields a variance decomposition for the client-availability setting.
The communicating subset S_k must be contained within the available client set Q_k, so sampling is restricted to available clients.
An analogous upper bound on Equation (39) would complete the convergence proof for these settings.

F.1 Federated EMNIST Dataset

The FEMNIST experiments evaluate partial participation across multiple datasets and communication budgets, tracking validation accuracy and training loss over communication rounds. The reported figures compare current best validation accuracy against communication rounds and client-to-master bits.

Experimental setup: Each FEMNIST experiment runs 151 communication rounds, with local training loss reported every round and validation accuracy every 5 rounds.
Experimental setup: In each round, 32 sampled clients perform one local SGD epoch on their training images using batch size 20.
Experimental setup: For partial participation, the expected number of clients allowed to communicate their updates is set to m ∈ {3, 6}.
Results: Figures 8, 9, and 10 plot current best validation accuracy against communication rounds and client-to-master communication bits for FEMNIST Datasets 1, 2, and 3.
Additional comparisons: Figures 11 and 12 provide the same validation-accuracy view for Shakespeare with n = 32 and n = 128 clients.

G Additional Experiment on Federated CIFAR100 Dataset

The Federated CIFAR100 experiment uses balanced client data and compares partial-participation strategies under a fixed local-training setup. Optimal client sampling achieves better performance than uniform sampling on this dataset.

Experimental setup: Federated CIFAR100 is balanced, with every client holding the same number of training images.
Experimental setup: Each round samples 32 clients uniformly, and every client performs several local SGD steps for one epoch with batch size 20.
Results: Figure 13 reports CIFAR100 validation accuracy and local training loss as functions of communication rounds and client-to-master communication bits.
Results: Better performance than uniform sampling is achieved by the optimal client sampling scheme on the balanced CIFAR100 dataset.
Hyperparameters: The optimal local step size is η_l = 1 × 10^-3 for full participation and optimal sampling, versus η_l = 3 × 10^-4 for uniform sampling.