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

TL;DR

The paper addresses stochastic bilevel optimization where an upper-level objective $f$ depends on the lower-level solution $y^*(x)$ of $g(x,y)$. It introduces PnPBO, a plug-and-play, single-loop framework that can freely combine biased or unbiased stochastic estimators across the three update directions ($x$, $y$, $z$), augmented with a moving-average mechanism for unbiased UL estimators and a clipping operation on the implicit variable. A unified convergence and complexity analysis shows that integrating modern estimators within PnPBO achieves optimal or near-optimal sample complexity, matching single-level optimization in the finite-sum setting; instantiations SFFBA and MSEBA demonstrate these rates using ZeroSARAH/ PAGE-based estimators. Empirical results on data hyper-cleaning and hyperparameter selection tasks validate the theoretical findings and illustrate practical gains from MA and clipping, signaling strong potential for scalable, plug-and-play BLO in ML applications.

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.

Figures (5)

• Figure 1: PnPBO and Unified Analysis Framework Schematic
• Figure 2: Roadmap of the Analysis Framework. Above, "BE" represents biased stochastic estimator and "UE" represents unbiased stochastic estimator.In the figure, we have kept only the key terms from the formulas.
• Figure 3: Comparison of SPABA, SFFBA and MSEBA with other variance-reduction-based stochastic BLO methods. Left: Test error for data hyper-cleaning on MNIST with $\tilde{p} = 0.9$ corruption rate, Right: Suboptimality gap for hyperparameter optimization for $l^2$ penalized logistic regression on IJCNN1 data set.
• Figure 4: Top: Compare SPABA, SFFBA and MSEBA with other benchmark algorithms in the problem of data hyper-cleaning with different corruption probability $\tilde{p}$. Bottom: Compare SPABA, SFFBA and MSEBA with other benchmark algorithms in the problem of hyperparameter selection experiment on the covtype and IJCNN1 data set.
• Figure 5: Left: Compare the performance of SOBA and MA-SOBA, and SABA and MA-SABA in the problem of hyperparameter selection experiment on the IJCNN1 data set. Right: Compare the performance of SPABA, SFFBA, and MSEBA with their versions without the clipping technique in the problem of hyperparameter selection experiment on the covtype data set.

Theorems & Definitions (15)

• Remark 1
• Remark 2
• Lemma 3
• Remark 8
• Remark 11
• Theorem 12
• Theorem 13
• Theorem 14
• Remark 15
• Theorem 16