A Provably Convergent Plug-and-Play Framework for Stochastic Bilevel Optimization

TL;DR

This paper introduces PnPBO, a flexible and provably convergent plug-and-play framework for stochastic bilevel optimization that unifies various estimators and achieves optimal complexity, validated through theoretical analysis and experiments.

Contribution

It presents a novel unified framework for stochastic bilevel optimization that incorporates multiple estimators and proves optimal convergence complexity.

Findings

Achieves optimal sample complexity comparable to single-level optimization.

Unifies multiple stochastic estimators within a single framework.

Empirically validates the effectiveness of PnPBO on benchmark problems.

Abstract

Bilevel optimization has recently attracted significant attention in machine learning due to its wide range of applications and advanced hierarchical optimization capabilities. In this paper, we propose a plug-and-play framework, named PnPBO, for developing and analyzing stochastic bilevel optimization methods. This framework integrates both modern unbiased and biased stochastic estimators into the single-loop bilevel optimization framework introduced in [9], with several improvements. In the implementation of PnPBO, all stochastic estimators for different variables can be independently incorporated, and an additional moving average technique is applied when using an unbiased estimator for the upper-level variable. In the theoretical analysis, we provide a unified convergence and complexity analysis for PnPBO, demonstrating that the adaptation of various stochastic estimators (including PAGE, ZeroSARAH, and mixed strategies) within the PnPBO framework achieves optimal sample complexity, comparable to that of single-level optimization. This resolves the open question of whether the optimal complexity bounds for solving bilevel optimization are identical to those for single-level optimization. Finally, we empirically validate our framework, demonstrating its effectiveness on several benchmark problems and confirming our theoretical findings.