math.OCOct 5, 2026

A Single-Loop, Constant-Batch First-Order Penalty Method for Stochastic Bilevel Optimization

Authors: Xingyu Chen, Ming Yang, Quanqi Hu, Tianbao Yang

Organizations: Department of Computer Science and Engineering Texas A&M University, College Station, USA

Abstract

Recent advances in penalty-based methods for stochastic bilevel optimization (SBO) have eliminated the need for second-order derivative oracles. However, for stochastic nonconvex-strongly convex bilevel problems, existing first-order methods typically rely on nested loops and/or large batch sizes for attaining O(ε−6)O(ε^{-6}) or O(ε−4)O(ε^{-4}) sample complexity under standard bounded-variance assumption or mean-square smoothness assumption. Achieving these rates with a single-loop penalty method and a constant batch size remains challenging due to a large penalty value needed for an accurate approximation. To address this challenge, we develop a stochastic SIngle-loop COnstant-Batch first-order penalty method (SICO) that combines two complementary ingredients. First, it performs one stochastic-gradient update per-iteration for both the original lower-level and penalized problems, with a projection that controls the separation between their iterates. Second, it applies an exponential moving average to stabilize the upper-level gradient estimator. We show that this combination achieves O(ε−6)O(ε^{-6}) sample complexity using only O(1)O(1) stochastic-gradient samples per iteration under unbiased, bounded-variance stochastic gradients. Under the additional mean-square smoothness assumption on the lower-level stochastic gradients, the same algorithm improves the complexity to O(ε−4)O(ε^{-4}) also with O(1)O(1) batch size. To the best of our knowledge, this is the first work to match the best-known convergence rate for fully first-order SBO methods using a single loop and a constant batch size. This result addresses an open problem posed in the literature.

Figures & tables

Explore similar work

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

    May 2, 2025Tianshu Chu, Dachuan Xu, Wei Yao +2Bilevel OptimizationStochastic Optimization

  2. Optimal Stochastic Bilevel Optimization with First-Order Oracles

    Oct 1, 2026Linxuan Pan, Junchi YangBilevel OptimizationFirst Order Oracle Complexity