Federated Conditional Stochastic Optimization

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 (\emph{i. e. }, 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 Meta-Learning in FL, federated conditional stochastic optimization considers the sample of tasks. Extensive experimental results on various tasks validate the efficiency of these algorithms.

Authors

• Xidong Wu Electrical and Computer Engineering, University of Pittsburgh, PA, USA
• Jianhui Sun Computer Science, University of Virginia, VA, USA
• Zhengmian Hu University of Pittsburgh
• Junyi Li Harbin Institute of Technology
• Aidong Zhang University of Virginia
• Heng Huang University of Maryland at College Park

Keywords

No keywords are indexed for this paper.