Smoothed gradient method improves federated bilevel optimization performance
Smoothed Gradient Method for Nonconvex Federated Stochastic Bilevel Optimization
Machine Learning
Summary
Optimizing complex machine learning models often involves solving two related problems at once, which can be slow and tricky when done across many devices. The authors developed a new method that speeds up this process by smoothing out gradients and allowing different learning speeds for each problem level. Their method works without strict conditions on the lower-level problem and improves the speed and efficiency of finding good model settings. They also confirmed through experiments that their approach works well in practice.
What this means in practice
- •For distributed machine learning engineers: Train complex hierarchical models more efficiently across many devices with improved convergence and reduced communication overhead.
- •For autonomous system developers: Implement bilevel optimization for real-time adaptation in systems without requiring strict assumptions on problem convexity or learning rates.
Authors
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.