[Paper Review] SimFBO: Towards Simple, Flexible and Communication-efficient Federated Bilevel Learning

Y. F. Yang, Peiyao Xiao+1 more|arXiv (Cornell University)|May 30, 2023

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes SimFBO, a communication-efficient federated bilevel optimization framework that eliminates sub-loops and enables single-round communication per iteration by jointly updating client and server variables. It achieves linear convergence speedup with partial client participation and improved sample and communication complexity, with a robust variant (ShroFBO) ensuring convergence under system-level heterogeneity.

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.

Motivation & Objective

• To address the high communication cost and complex implementation of existing federated bilevel optimization (FBO) methods that rely on multiple sub-loops.
• To enable practical client sampling without replacement in FBO, which is challenging due to nested optimization structures.
• To improve resilience to system-level heterogeneity such as varying client computation power and storage.
• To achieve provably linear convergence speedup with partial client participation and reduced sample and communication complexity.

Proposed method

• SimFBO employs a single communication round per iteration, simultaneously updating the upper-level variable $x$, lower-level variable $y$, and hypergradient approximation variable $v$.
• It introduces a generalized server-side aggregation for $x$, $y$, and $v$ that supports flexible aggregation strategies and enables client sampling without replacement.
• The method approximates the Hessian-inverse-vector product via a recursive update in a single loop, avoiding the need for separate sub-loops for lower-level optimization and hypergradient computation.
• ShroFBO enhances SimFBO with a normalized, correction-based server update to ensure convergence to the original objective despite client computation heterogeneity.
• Theoretical analysis establishes linear convergence speedup under partial client participation and improved sample and communication complexities.

Experimental results

Research questions

• RQ1 Can federated bilevel optimization be achieved with a single communication round per iteration without sub-loops?
• RQ2 How can client sampling without replacement be supported in FBO under nested optimization structures?
• RQ3 Can server-side aggregation be generalized to improve communication efficiency and support flexible aggregation strategies?
• RQ4 Does the proposed method maintain convergence to the original objective under system-level heterogeneity?
• RQ5 What is the theoretical sample and communication complexity of the proposed method compared to prior work?

Key findings

• SimFBO and ShroFBO achieve a sample complexity of $\mathcal{O}(\epsilon^{-2}P^{-1})$, matching the best-known result in FedMBO.
• Both methods achieve linear convergence speedup with partial client participation and sampling without replacement.
• The communication complexity is significantly reduced compared to AID- and ITD-based FBO methods due to elimination of sub-loops.
• ShroFBO provably converges to a stationary point of the original objective even under system-level heterogeneity, unlike naive aggregation.
• Experiments demonstrate superior performance of SimFBO and ShroFBO over existing FBO algorithms in convergence speed and communication efficiency.