Title: Similarity-aware Non-Convex Federated Optimization
Primary area: optimization
Keywords: Client Sampling, SAGA, Second-order Similarity, Composite Gradient Method, Variance Reduction

Abstract: Federated optimization is a distributed machine learning approach where a model is trained collaboratively across a network of devices without sharing their local datasets. Reducing communication costs is essential for scalable federated optimization. Recent advances show that SARAH-type variance reduction methods  achieve the best-known communication complexity depending only on second-order similarity constants—often substantially better than the dependence on individual smoothness derived from incremental gradient methods like SAGA. However, SARAH-type methods require periodic participation from all clients, which may be impractical when some clients are intermittently inactive. In contrast, SAGA-based methods require full participation only during initialization and can sample clients afterward, offering greater flexibility. In this work, we show for the first time that the SAGA estimator can also exploit functional similarity, and its approximation error can be further reduced when equipped with the recursive gradient estimator. We analyze the composite gradient method using this enhanced estimator and establish the state-of-the-art communication complexity for optimizing non-convex functions. Under additional smoothness assumptions, a specific instantiation of our framework outperforms all existing methods in both communication and local computation complexity, making it well-suited for real-world federated learning with partial client participation.

TL;DR: We propose a composite gradient method with SAGA and recursive gradient estimator for non-convex federated optimization that can exploit second-order similarity and achieve high communication and local computation efficiency.