Faster federated learning with first order bilevel optimization
Federated stochastic bilevel optimization with fully first-order gradients
Machine Learning
Summary
Training complex machine learning models often requires updates that involve heavy math, like second-order derivatives, which slow things down. The authors introduce a new method that avoids these slow calculations by using only simpler first-order information, making training quicker in federated (distributed) setups. They also design a stable learning rate approach that improves how different parts of the model update together. Their experiments show this method works well while saving time.
What this means in practice
- •For machine learning engineers: Speed up federated training of complex models by reducing the need for costly second-order derivatives.
- •For distributed systems developers: Implement efficient coordination of variable updates in federated systems using a stable, single-timescale learning rate.
Authors
Yihan Zhang, Rohit Dhaipule, Chiu C Tan, Haibin Ling, Hongchang Gao
Abstract
Federated stochastic bilevel optimization has been actively studied in recent years due to its widespread applications in machine learning. However, most existing federated stochastic bilevel optimization algorithms require the computation of second-order Hessian and Jacobian matrices, which leads to longer running times in practice. To address these challenges, we propose a novel federated stochastic variance-reduced bilevel gradient descent algorithm that relies solely on first-order oracles. Specifically, our approach does not require the computation of second-order Hessian and Jacobian matrices, significantly reducing running time. Furthermore, we introduce a novel learning rate mechanism, i.e., a constant single-timescale learning rate, to coordinate the update of different variables. We also present a new strategy to establish the convergence rate of our algorithm. Finally, the extensive experimental results confirm the efficacy of our proposed algorithm.