Optimal Client Sampling for Federated Learning (2010.13723v3)

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.

• The paper introduces an optimal client subsampling framework that leverages update norms to reduce communication costs in federated learning.
• It proposes an efficient algorithm for stateless, privacy-preserving aggregation that outperforms uniform random sampling in convergence speed.
• Empirical evaluations on LEAF benchmark datasets show that the method achieves near-full participation performance with substantially lower communication overhead.

Optimal Client Sampling for Federated Learning: A Detailed Overview

The paper addresses the significant communication bottleneck present in federated learning (FL) setups, particularly in cross-device configurations where numerous devices, such as mobile phones or IoT devices, possess constrained bandwidth. This constraint limits the number of clients that can effectively participate in each communication round without excess communication overhead. The authors propose an innovative client subsampling strategy aimed at optimizing client participation by selectively communicating updates that are deemed "important." This importance is objectively measured based on the norm of each client's update.

Key Contributions and Methodology

1. Subsampling Strategy: The paper introduces a mathematical framework to compute optimal client participation probabilities. This framework aims to minimize the variance between complete client participation (where all client updates are communicated) and restricted participation. The strategy primarily leverages the norms of updates as a criterion for importance.
2. Algorithmic Innovation: A practical algorithm that approximates the optimal sampling probabilities is proposed. This algorithm allows for secure aggregation and supports stateless clients, thereby not compromising privacy requirements integral to FL systems. Importantly, the algorithm does not necessitate the sharing of raw updates, adhering to privacy protocols.
3. Convergence and Performance: The paper provides rigorous theoretical convergence guarantees for two pivotal approaches in distributed optimization within FL: Distributed Stochastic Gradient Descent (DSGD) and Federated Averaging (FedAvg). The results suggest that the proposed method can approach the performance of full client participation and is demonstrably superior to uniform random sampling of clients.
4. Compatibility with Existing Techniques: The optimal sampling strategy is orthogonal to—and thus can be combined with—other techniques aimed at reducing communication overhead in FL, such as local SGD methods and gradient compression techniques.

Empirical Evaluations

In practice, the optimal client sampling method was evaluated on multiple datasets derived from the LEAF benchmark suite, including the Federated EMNIST and Shakespeare datasets. The experiments showcased that the proposed method offers substantial improvements in terms of communication efficiency compared to both full and uniform partial client participation strategies. Particularly, the method showed faster convergence in communication-efficient terms, achieving high performance with significantly reduced communication costs.

Implications and Future Prospects

The implications of this research are twofold:

• Practical Improvements in FL: The proposed method potentially enables FL systems to operate more efficiently, especially where bandwidth and communication are significant constraints. This aligns with deployment scenarios where devices have limited communication capability, thus extending the applicability of FL to a broader range of tasks.
• Theoretical Advancements: The research presents a refined understanding of the variance embeddings in FL, offering insights into the development of even more sophisticated subsampling strategies that factor in other aspects, such as heterogeneity in client data.

Future research could explore further optimizations, perhaps integrating the subsampling methodology with advanced communication compression techniques or expanding the approach to accommodate dynamically changing client availability. Another direction would be to empirically test the boundaries of the $m=\mathcal{O}(\sqrt{n})$ conjecture, where $m$ is the expected number of sampled clients necessary to maintain performance comparable to full participation.