Federated Conditional Stochastic Optimization
This addresses the need for communication-efficient distributed optimization in applications like invariant learning and meta-learning, but it is incremental as it builds on existing federated learning and conditional stochastic optimization methods.
The paper tackles the problem of nonconvex conditional stochastic optimization in federated learning by proposing new algorithms (FCSG, FCSG-M, Acc-FCSG-M) that achieve improved sample and communication complexity, matching lower bounds in single-machine settings.
Conditional stochastic optimization has found applications in a wide range of machine learning tasks, such as invariant learning, AUPRC maximization, and meta-learning. As the demand for training models with large-scale distributed data grows in these applications, there is an increasing need for communication-efficient distributed optimization algorithms, such as federated learning algorithms. This paper considers the nonconvex conditional stochastic optimization in federated learning and proposes the first federated conditional stochastic optimization algorithm (FCSG) with a conditional stochastic gradient estimator and a momentum-based algorithm (FCSG-M). To match the lower bound complexity in the single-machine setting, we design an accelerated algorithm (Acc-FCSG-M) via the variance reduction to achieve the best sample and communication complexity. Compared with the existing optimization analysis for MAML in FL, federated conditional stochastic optimization considers the sample of tasks. Extensive experimental results on various tasks validate the efficiency of these algorithms.