Federated Conditional Stochastic Optimization

TL;DR

This paper introduces the first federated conditional stochastic optimization algorithms, FCSG and FCSG-M, with accelerated variants, achieving optimal complexity for large-scale distributed nonconvex optimization tasks in machine learning.

Contribution

It proposes novel federated algorithms for conditional stochastic optimization, incorporating variance reduction and momentum, with theoretical guarantees matching lower bounds.

Findings

Algorithms achieve optimal sample and communication complexity.

Experimental results validate efficiency across various tasks.

First to address federated conditional stochastic optimization with theoretical guarantees.

Abstract

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.