Smoothed Gradient Method for Nonconvex Federated Stochastic Bilevel Optimization
Xinwen Zhang, Peiran Yu, Zhaosong Lu, Hongchang Gao
Abstract
In recent years, federated stochastic bilevel optimization has attracted increasing attention due to its wide range of applications in machine learning. To reduce the computational overhead associated with second-order Hessian and Jacobian matrices, several first-order methods have been proposed. However, existing methods typically impose restrictive assumptions on the lower-level function, suffer from a strong dependence on the condition number in their convergence rates, and require different learning-rate scales for variables across the upper- and lower-level problems, limiting their practical applicability and complicating hyperparameter tuning. To address these challenges, we propose a stochastic doubly smoothed gradient method for nonconvex federated stochastic bilevel optimization problems, which decouples the learning rates of upper- and lower-level variables and does not require a strongly-convex lower-level loss function. We establish rigorous theoretical guarantees for the proposed algorithm, demonstrating an improved convergence rate of $O(κ^{15/2}/ε^5)$ and a communication complexity of $O(κ^{4}/ε^3)$, where $κ$ denotes the condition number and $ε$ represents the solution accuracy. Notably, these bounds exhibit significantly better dependence on the condition number $κ$ than those of existing methods. Extensive experiments validate the effectiveness of our algorithm.