SimFBO: Towards Simple, Flexible and Communication-efficient Federated Bilevel Learning

Published: 21 Sept 2023, Last Modified: 27 Dec 2023NeurIPS 2023 spotlightEveryoneRevisionsBibTeX
Keywords: Federated bilevel optimization, federated hypergradient, communication efficiency, system-level heterogeneity, linear speedup
Abstract: Federated bilevel optimization (FBO) has shown great potential recently in machine learning and edge computing due to the emerging nested optimization structure in meta-learning, fine-tuning, hyperparameter tuning, etc. However, existing FBO algorithms often involve complicated computations and require multiple sub-loops per iteration, each of which contains a number of communication rounds. In this paper, we propose a simple and flexible FBO framework named SimFBO, which is easy to implement without sub-loops, and includes a generalized server-side aggregation and update for improving communication efficiency. We further propose System-level heterogeneity robust FBO (ShroFBO) as a variant of SimFBO with stronger resilience to heterogeneous local computation. We show that SimFBO and ShroFBO provably achieve a linear convergence speedup with partial client participation and client sampling without replacement, as well as improved sample and communication complexities. Experiments demonstrate the effectiveness of the proposed methods over existing FBO algorithms.
Supplementary Material: zip
Submission Number: 8751