Non-Convex Federated Optimization under Cost-Aware Client Selection

2025-12-04

AI Summary

This work addresses the long-overlooked impact of client selection strategies on communication and computation costs in federated optimization. We propose the first unified analytical framework that explicitly models heterogeneous communication overheads to precisely quantify resource consumption across diverse selection policies. To reduce overall complexity, we design RG-SAGA—a novel gradient estimator integrating recursive gradient updates with SAGA-type variance reduction, combined with inexact composite gradient methods and functional similarity modeling—thereby significantly tightening the error bound. Theoretically, RG-SAGA achieves state-of-the-art communication and local computation complexity for non-convex objectives: O(1/ε²) communication rounds and O(1/ε) per-client gradient computations. Extensive experiments demonstrate its superior convergence speed and resource efficiency over leading federated optimization algorithms.

Technology Category

• Search and Optimization: Non-convex Optimization
• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Distributed Machine Learning & Federated Learning

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs

Abstract

Different federated optimization algorithms typically employ distinct client-selection strategies: some methods communicate only with a randomly sampled subset of clients at each round, while others need to periodically communicate with all clients or use a hybrid scheme that combines both strategies. However, existing metrics for comparing optimization methods typically do not distinguish between these strategies, which often incur different communication costs in practice. To address this disparity, we introduce a simple and natural model of federated optimization that quantifies communication and local computation complexities. This new model allows for several commonly used client-selection strategies and explicitly associates each with a distinct cost. Within this setting, we propose a new algorithm that achieves the best-known communication and local complexities among existing federated optimization methods for non-convex optimization. This algorithm is based on the inexact composite gradient method with a carefully constructed gradient estimator and a special procedure for solving the auxiliary subproblem at each iteration. The gradient estimator is based on SAGA, a popular variance-reduced gradient estimator. We first derive a new variance bound for it, showing that SAGA can exploit functional similarity. We then introduce the Recursive-Gradient technique as a general way to potentially improve the error bound of a given conditionally unbiased gradient estimator, including both SAGA and SVRG. By applying this technique to SAGA, we obtain a new estimator, RG-SAGA, which has an improved error bound compared to the original one.

Problem

• Addresses client selection strategies in federated optimization.
• Quantifies communication and computation costs in federated learning.
• Proposes an algorithm for non-convex optimization with improved complexities.

Innovation

• Introduces cost-aware client selection model
• Proposes RG-SAGA gradient estimator technique
• Optimizes communication and local computation complexities