Federated Multi-Sequence Stochastic Approximation with Local Hypergradient Estimation
This addresses a bottleneck in federated learning for complex optimization problems like bilevel optimization, offering significant communication efficiency improvements, though it is incremental in extending existing MSA frameworks to federated settings.
The paper tackles the problem of designing efficient federated algorithms for multi-sequence stochastic approximation (MSA), which includes bilevel and multi-level optimization, by introducing FedMSA, the first federated algorithm for MSA that achieves near-optimal communication complexity and enables provable local hypergradient estimation, with experiments showing order-of-magnitude savings in communication rounds compared to prior methods.
Stochastic approximation with multiple coupled sequences (MSA) has found broad applications in machine learning as it encompasses a rich class of problems including bilevel optimization (BLO), multi-level compositional optimization (MCO), and reinforcement learning (specifically, actor-critic methods). However, designing provably-efficient federated algorithms for MSA has been an elusive question even for the special case of double sequence approximation (DSA). Towards this goal, we develop FedMSA which is the first federated algorithm for MSA, and establish its near-optimal communication complexity. As core novelties, (i) FedMSA enables the provable estimation of hypergradients in BLO and MCO via local client updates, which has been a notable bottleneck in prior theory, and (ii) our convergence guarantees are sensitive to the heterogeneity-level of the problem. We also incorporate momentum and variance reduction techniques to achieve further acceleration leading to near-optimal rates. Finally, we provide experiments that support our theory and demonstrate the empirical benefits of FedMSA. As an example, FedMSA enables order-of-magnitude savings in communication rounds compared to prior federated BLO schemes.