[论文解读] 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 Data被引用 4

一句话总结

该论文提出 SimFBO，一种通信高效的联邦双层优化框架，通过联合更新客户端和服务器变量，消除了子循环，并在每次迭代中实现单轮通信。该方法在部分客户端参与的情况下实现线性收敛加速，同时提升了样本复杂度和通信复杂度。其鲁棒变体（ShroFBO）在系统级异构性下仍能保证收敛。

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.

研究动机与目标

• 为解决现有联邦双层优化（FBO）方法依赖多个子循环所导致的高通信成本和复杂实现问题。
• 在嵌套优化结构下，实现 FBO 中无需放回的客户端采样，这在以往极具挑战。
• 提升对客户端计算能力与存储差异等系统级异构性的鲁棒性。
• 在部分客户端参与下实现可证明的线性收敛加速，并降低样本复杂度与通信复杂度。

提出的方法

• SimFBO 每轮迭代仅使用一轮通信，同时更新上层变量 $x$、下层变量 $y$ 和超梯度近似变量 $v$。
• 提出一种广义的服务器端聚合机制，用于 $x$、$y$ 和 $v$，支持灵活的聚合策略，并实现无需放回的客户端采样。
• 通过单循环中的递归更新近似 Hessian 逆向量乘积，避免了为下层优化和超梯度计算分别设置独立子循环。
• ShroFBO 通过引入归一化、基于校正的服务器更新机制，增强 SimFBO 的鲁棒性，确保在客户端计算异构性下仍能收敛至原始目标函数。
• 理论分析表明，在部分客户端参与下，该方法可实现线性收敛加速，并改善样本复杂度与通信复杂度。

实验结果

研究问题

• RQ1 是否可以在不使用子循环的前提下，实现每轮迭代仅一轮通信的联邦双层优化？
• RQ2 在嵌套优化结构下，如何在 FBO 中支持无需放回的客户端采样？
• RQ3 是否可以对服务器端聚合进行泛化，以提升通信效率并支持灵活的聚合策略？
• RQ4 所提方法在系统级异构性下是否仍能保持对原始目标函数的收敛性？
• RQ5 与先前工作相比，该方法的理论样本复杂度与通信复杂度如何？

主要发现

• SimFBO 和 ShroFBO 的样本复杂度为 $\mathcal{O}(\epsilon^{-2}P^{-1})$，与 FedMBO 中的最佳已知结果一致。
• 两种方法均在部分客户端参与及无需放回采样下实现了线性收敛加速。
• 由于消除了子循环，与基于 AID 和 ITD 的 FBO 方法相比，通信复杂度显著降低。
• ShroFBO 在系统级异构性下可证明收敛至原始目标函数的驻点，而朴素聚合方法则无法做到。
• 实验结果表明，SimFBO 和 ShroFBO 在收敛速度与通信效率方面均优于现有 FBO 算法。