Federated Stochastic Bilevel Optimization with Fully First-Order Gradients

Federated stochastic bilevel optimization (FedSBO) is critical for machine learning tasks like meta-learning and hyperparameter optimization. However, existing algorithms for this problem typically require computing second-order Hessian and Jacobian matrices, which are computationally prohibitive in high-dimensional federated settings. The authors aim to develop a fully first-order algorithm that avoids these second-order costs while maintaining convergence guarantees.

The authors introduce the Federated Stochastic Variance-Reduced Bilevel Gradient Descent with Fully First-Order Oracles (FedSVRBGD-FO). Their approach reformulates the bilevel problem into an unconstrained minimax optimization problem using a penalty method. To address the challenges of federated learning, they implement:

• A novel single-timescale constant learning rate mechanism to coordinate variable updates without requiring iteration-dependent tuning.
• A variance-reduced gradient estimator that relies exclusively on first-order oracles.
• A new potential function for convergence analysis that allows for a linear speedup with respect to the number of devices.

The study demonstrates that FedSVRBGD-FO achieves an O(1/Nε^5) convergence rate, providing a linear speedup relative to the number of devices (N). Experimental results on hyperparameter optimization and hyper-representation learning tasks show that this first-order method significantly reduces running time compared to state-of-the-art second-order methods like FedNEST and FedMBO, while achieving comparable test accuracy.

This work provides a practical, efficient solution for federated bilevel optimization. By removing the dependency on second-order information, the algorithm makes complex bilevel optimization tasks feasible in high-dimensional, distributed environments where computational resources on edge devices are limited.