A Near-Optimal Algorithm for Stochastic Bilevel Optimization via Double-Momentum

TL;DR

This paper introduces SUSTAIN, a novel single-timescale stochastic algorithm with double momentum for bilevel optimization, achieving near-optimal iteration complexity for both convex and non-convex upper-level objectives.

Contribution

The paper presents a new single-timescale stochastic bilevel optimization algorithm with double momentum, improving convergence analysis and complexity bounds over prior two-timescale methods.

Findings

Achieves $ ilde{O}(rac{1}{ ext{epsilon}^{3/2}})$ iteration complexity for non-convex cases.

Matches the best-known stochastic gradient complexity for single-level problems.

Effectively controls stochastic gradient errors in bilevel settings.

Abstract

This paper proposes a new algorithm -- the \underline{S}ingle-timescale Do\underline{u}ble-momentum \underline{St}ochastic \underline{A}pprox\underline{i}matio\underline{n} (SUSTAIN) -- for tackling stochastic unconstrained bilevel optimization problems. We focus on bilevel problems where the lower level subproblem is strongly-convex and the upper level objective function is smooth. Unlike prior works which rely on \emph{two-timescale} or \emph{double loop} techniques, we design a stochastic momentum-assisted gradient estimator for both the upper and lower level updates. The latter allows us to control the error in the stochastic gradient updates due to inaccurate solution to both subproblems. If the upper objective function is smooth but possibly non-convex, we show that {\aname}~requires $\mathcal{O}(\epsilon^{-3/2})$ iterations (each using ${\cal O}(1)$ samples) to find an $\epsilon$-stationary solution. The $\epsilon$-stationary solution is defined as the point whose squared norm of the gradient of the outer function is less than or equal to $\epsilon$. The total number of stochastic gradient samples required for the upper and lower level objective functions matches the best-known complexity for single-level stochastic gradient algorithms. We also analyze the case when the upper level objective function is strongly-convex.