SCAFFOLD: Stochastic Controlled Averaging for Federated Learning (1910.06378v4)

Abstract: Federated Averaging (FedAvg) has emerged as the algorithm of choice for federated learning due to its simplicity and low communication cost. However, in spite of recent research efforts, its performance is not fully understood. We obtain tight convergence rates for FedAvg and prove that it suffers from client-drift' when the data is heterogeneous (non-iid), resulting in unstable and slow convergence. As a solution, we propose a new algorithm (SCAFFOLD) which uses control variates (variance reduction) to correct for theclient-drift' in its local updates. We prove that SCAFFOLD requires significantly fewer communication rounds and is not affected by data heterogeneity or client sampling. Further, we show that (for quadratics) SCAFFOLD can take advantage of similarity in the client's data yielding even faster convergence. The latter is the first result to quantify the usefulness of local-steps in distributed optimization.

• The paper's main contribution is the introduction of SCAFFOLD, which uses control variates to effectively reduce client-drift in federated learning.
• It provides a rigorous theoretical analysis demonstrating faster convergence and improved communication efficiency compared to FedAvg.
• Empirical results validate SCAFFOLD's robustness over non-iid data, establishing tighter convergence rates in diverse decentralized settings.

Stochastic Controlled Averaging for Federated Learning

The paper presents a comprehensive analysis and introduction of a new algorithm designed to improve the efficiency and robustness of federated learning in environments with heterogeneous data — a common scenario in real-world applications. Federated Averaging (FedAvg) has been the traditional choice due to its simplicity and low communication cost. However, its performance can degrade significantly when faced with heterogeneous or non-iid data distributions due to a phenomenon termed "client-drift." This paper investigates this issue and proposes an enhanced method to mitigate these limitations.

Key Contributions

1. Analysis of FedAvg: The paper offers an in-depth analysis of FedAvg's performance, which experiences a convergence slowdown titled "client-drift" on heterogeneous data. It establishes theoretical bounds and demonstrates that client-drift persists even when full batch gradients are used and all clients participate in training. This analysis justifies the need for an improved algorithm that can effectively manage data heterogeneity.
2. Stochastic Controlled Averaging (SCAFFOLD): The new algorithm, SCAFFOLD, uses "control variates" for variance reduction, effectively addressing the client-drift issue. SCAFFOLD corrects the local updates' directionality, thereby aligning them better with the global model's optimal path. This approach reduces the number of communication rounds needed and shows robustness to data heterogeneity and client sampling.
3. Theoretical Guarantees and Empirical Validation: The authors provide rigorous theoretical analysis demonstrating that SCAFFOLD converges faster than FedAvg and large-batch SGD under certain conditions. For scenarios where client data shows high similarity, SCAFFOLD achieves even faster convergence. The paper also establishes tighter convergence rates for FedAvg than previously known. Additionally, empirical results confirm the theoretical insights, showcasing SCAFFOLD's superior performance using both simulated and real datasets.
4. Implications for Similarity and Heterogeneity: A significant insight arising from the research is the distinction in algorithmic advantages between gradient and Hessian similarity. SCAFFOLD leverages Hessian similarity, allowing it to outperform competitors even when the clients' optimal points are diverse. This understanding extends the applicability of variational reduction techniques to federated learning.
5. Improved Communication Complexity: The paper also demonstrates that with SCAFFOLD, communication rounds are optimized better than with existing methods such as FedAvg and DANE. This offers practical improvements for large-scale federated systems, minimizing the overhead of client-server interactions and thus reducing overall resource expenditure.

Implications and Future Directions

The introduction of SCAFFOLD as a resilient method to bypass the inefficiencies in FedAvg brings significant theoretical and practical implications. It offers a pathway to deploy federated learning in diverse and challenging environments with decentralized data while maintaining communication efficiency. This development is particularly essential in applications involving privacy-sensitive data where direct data aggregation is infeasible.

Theoretical advancements in understanding the interplay between different notions of function similarity — namely, gradient versus Hessian similarity — pave the way for further optimized distributed learning algorithms. Moreover, the detailed analysis and improvements in handling client sampling variability signify a robust step forward in adapting federated learning for real-world constraints.

Future work could focus on extending these methodologies to more complex models and data distributions, integrating concepts from recent advances in deep learning optimization. Additionally, adapting SCAFFOLD for asynchronous or partially active client systems could further enhance its practical utility. The exploration of alternative control variate mechanisms could also yield new insights into distributed optimization, potentially unlocking new performance bounds and widening the scope for federated learning applications.