[Paper Review] Optimal Client Sampling for Federated Learning

Wenlin Chen, Samuel Horváth+1 more|arXiv (Cornell University)|Oct 26, 2020

Privacy-Preserving Technologies in Data47 references95 citations

TL;DR

Proposes an adaptive, privacy-friendly optimal client sampling scheme that reduces communication in Federated Learning, with convergence guarantees for DSGD and FedAvg and performance nearing full participation.

ABSTRACT

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.

Motivation & Objective

• Address the communication bottleneck in cross-device Federated Learning
• Develop an adaptive client sampling strategy that picks informative updates
• Ensure compatibility with secure aggregation and stateless clients
• Provide convergence guarantees for DSGD and FedAvg under convex and non-convex settings

Proposed method

• Model the partial participation as a random set S with inclusion probabilities p_i and independent sampling
• Derive a closed-form optimal sampling probability p_i^k that minimizes the gradient estimator variance under a budget m (Equation 7)
• Provide an approximate, privacy-preserving algorithm (Algorithm 2) that avoids revealing individual update norms while maintaining proximal optimality
• Show compatibility with secure aggregation and stateless clients via an aggregation-only approach
• Present convergence analyses for DSGD and FedAvg under convex and non-convex assumptions (Theorems 13–18)

Experimental results

Research questions

• RQ1 How can we minimize the variance of the master update under a fixed participation budget m?
• RQ2 Can the proposed sampling be implemented with secure aggregation and stateless clients without leaking private information?
• RQ3 What are the convergence guarantees for DSGD and FedAvg when using optimal client sampling, in both convex and non-convex settings?
• RQ4 How does the proposed method compare to full participation and uniform sampling in practice?

Key findings

• The optimal sampling scheme minimizes the gradient estimator variance under a participation budget via independent sampling.
• A closed-form solution for p_i^k is provided (Equation 7), with an efficient approximate Algorithm 2 suitable for privacy constraints.
• The method achieves performance close to full participation and outperforms uniform sampling in theory and empirically.
• The sampling is orthogonal to and compatible with local updates and gradient compression methods, and with secure aggregation and stateless clients.
• The approach allows larger learning rates than the uniform sampling baseline, improving communication efficiency and convergence.