Source-linked AI summary
A Near-Optimal Algorithm for Stochastic Bilevel Optimization via Double-Momentum
Prashant Khanduri, Siliang Zeng, Mingyi Hong, Hoi-To Wai, Zhaoran Wang, Zhuoran Yang
TL;DR
Stochastic bilevel optimization must optimize an upper objective that depends on an inexact lower-level solution. SUSTAIN addresses this with single-timescale double momentum, achieving single-level-matching complexity for non-convex and strongly-convex cases.
Problem
Stochastic bilevel optimization couples an upper objective with the optimizer of a strongly convex lower-level problem.
Method
SUSTAIN uses momentum-assisted stochastic gradient estimators for both levels with single-loop, same-order updates and O(1) samples per update.
Results
SUSTAIN matches best-known single-level stochastic sample-complexity guarantees, requiring O(ε^-3/2) samples for non-convex and O(ε^-1) for strongly-convex bilevel problems.
Takeaways & Limitations
Improving both lower-level and outer-gradient estimates yields a sample- and computation-efficient stochastic bilevel algorithm.
Takeaways & Limitations
The guarantees apply to unconstrained bilevel problems with strongly convex lower-level subproblems, while Adam-based meta-learning comparisons lack theoretical guarantees.
Abstract
from arXiv · showhide
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}(ε^{-3/2})$ iterations (each using ${\cal O}(1)$ samples) to find an $ε$-stationary solution. The $ε$-stationary solution is defined as the point whose squared norm of the gradient of the outer function is less than or equal to $ε$. 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.
1 Introduction
The paper studies unconstrained bilevel optimization with strongly convex lower-level problems, where upper-level objectives depend on lower-level optimizers. SUSTAIN uses single-timescale double momentum to match single-level stochastic complexity while improving sample and computational efficiency.
- Problem setting: Bilevel optimization models hierarchical problems in which upper-level variables depend on a lower-level optimizer.The paper considers unconstrained problems with stochastic upper- and lower-level objectives.
- Problem setting: The lower-level objective is strongly convex, implying a smooth but possibly non-convex outer objective.This assumption defines the main problem class analyzed in the paper.
- Motivation: O(ε^-2) was the fastest previously available bilevel rate, compared with O(ε^-3/2) for optimal single-level stochastic gradient methods.The paper identifies this iteration/sample-complexity gap as a central motivation.
- SUSTAIN approach: SUSTAIN uses momentum-assisted stochastic gradient estimators for both upper- and lower-level updates, with O(1) samples per update and same-order step sizes.Its single-loop, single-timescale design avoids relying on two-timescale or double-loop updates.
- Results: O(ε^-3/2) samples suffice for non-convex bilevel problems, while O(ε^-1) samples suffice for strongly-convex bilevel problems.The stated bounds match optimal stochastic-gradient sample complexity for corresponding single-level problems.
- Efficiency: SUSTAIN avoids explicit sampled Hessian inversion and has lower per-iteration dimension dependence than methods requiring O(d_lo^3) computation.The paper attributes this advantage to Lipschitz properties of biased stochastic outer-gradient estimates.
2 Preliminaries
The preliminaries specify smoothness, strong-convexity, stochastic-gradient, and Lipschitz assumptions for the bilevel problem. They also explain how practical biased outer-gradient estimators approximate the true gradient and motivate single-loop updates.
- Assumptions: The analysis assumes smooth upper- and lower-level functions, with the lower-level objective strongly convex in y.The lower-level gradient and relevant Hessians are assumed Lipschitz continuous.
- Assumptions: Stochastic versions of the upper- and lower-level functions satisfy the same structural assumptions over their sampling supports.These assumptions define the stochastic problem class used in the analysis.
- Outer-gradient construction: The true outer gradient is obtained through the implicit function theorem and depends on the lower-level minimizer y*(x).Because y*(x) is generally unavailable in closed form, the algorithm uses a gradient surrogate evaluated at an approximate lower-level solution.
- Outer-gradient construction: The practical stochastic outer-gradient estimator is biased, with bias decaying exponentially as the estimator parameter K increases.Each computation uses at most K Hessian-vector products.
- Computational cost: O(log(T)d_lo^2) complexity results from using K = O(log(T)) Hessian-vector products per iteration.This is contrasted with O(d_lo^3) complexity for methods that explicitly estimate Hessian inverses.
- Lipschitz properties: Lipschitz continuity of the approximate outer gradient, true outer gradient, and lower-level solution supports single-loop updates with improved sample efficiency.The approximate outer gradient approaches the true gradient when y is close to y*(x).
3 The proposed SUSTAIN algorithm
SUSTAIN addresses bilevel stochastic optimization by combining recursive momentum estimators for the lower- and upper-level gradients in a single-timescale, single-loop method. Its convergence analysis yields near-single-level sample complexity under smoothness and strong-convexity assumptions.
- Algorithm design: The method uses O(1) samples per update and computes stochastic gradients at consecutive iterates with shared sample sets for recursive estimation.The algorithm initializes x and y, performs separate y- and x-updates, and returns a randomly selected iterate.
- Algorithm design: SUSTAIN uses recursive momentum estimators for both lower-level ∇yg(x, y) and outer-gradient estimates, updating x and y on the same timescale.The method tracks two vector estimators rather than each stochastic component in the implicit outer-gradient formula.
- Relation to prior methods: SUSTAIN avoids the double-loop updates and two-timescale step sizes used by earlier bilevel methods, while retaining a low-complexity sample-based Hessian inversion.Earlier methods improve inner tracking through repeated inner updates, large batches, or step-size separation.
- Estimator design: The estimator tracks the outer gradient directly instead of separately tracking all stochastic vectors and matrices, which avoids costly Hessian inversions.This design leverages cheap stochastic Hessian estimates and stores only two vector estimators.
- Convergence analysis: For smooth possibly non-convex outer objectives, SUSTAIN reaches an ε-stationary point using Õ(ε^-3/2) stochastic-gradient samples from both levels.This matches the best complexity bounds cited for single-level stochastic optimization, subject to the paper’s additional smoothness conditions.
- Strongly-convex case: For strongly-convex outer objectives, SUSTAIN reaches an ε-optimal point using Õ(ε^-1) stochastic-gradient samples from both upper and lower levels.The paper reports this as matching strongly-convex single-level SGD complexity and improving over the cited TTSA and BSA bounds.
4 Numerical experiments
The experiments evaluate SUSTAIN on meta-learning and data hyper-cleaning tasks, comparing it with bilevel optimization baselines under stochastic and Adam-based outer updates. Across the reported settings, SUSTAIN generally achieves strong performance with fewer iterations, samples, or gradient evaluations, while Adam-based comparisons are heuristic.
- Meta learning: SUSTAIN requires the smallest number of iterations, samples, and gradient computations to reach a given training or testing accuracy in the miniImageNet meta-learning experiment.This comparison uses vanilla stochastic-gradient outer updates; each algorithm accesses the same number of samples per iteration, while SUSTAIN uses twice the gradient computations.
- Experimental scope: The Adam-based meta-learning comparisons use heuristic algorithms without theoretical performance guarantees.This caveat applies to the Adam-based bilevel algorithms evaluated in the meta-learning experiments.
- Optimization variants: A single-gradient-computation update direction for SUSTAIN is reported to preserve performance, while Adam can be incorporated to improve performance over vanilla stochastic-gradient methods.The experiments use Option II as the update direction and report that the modification does not compromise performance.
- Meta learning: Adam-based SUSTAIN outperforms ITD-BiO, ANIL, and MAML on miniImageNet, reaching improved performance in fewer iterations and therefore fewer samples and gradient computations.The Adam-based algorithms substantially outperform their vanilla versions in both accuracy and speed.
- Data hyper-cleaning: On Fashion-MNIST data hyper-cleaning with corruption rates p = 0.3 and p = 0.4, SUSTAIN outperforms HOAG and stocBiO in total gradient evaluations and samples.The comparison is based on total computational access because HOAG is deterministic and stocBiO uses large batches, whereas SUSTAIN uses single-sample stochastic updates.
5 Conclusion and future work
The paper concludes that SUSTAIN provides single-timescale, sample- and computation-efficient optimization for unconstrained bilevel problems with strongly convex lower levels. It identifies lower-bound analysis and constrained lower-level problems as future directions.
- SUSTAIN targets unconstrained bilevel optimization with strongly convex lower-level subproblems.
- The algorithm uses neither two-timescale updates, large-batch gradients, nor double-loop procedures.
- SUSTAIN matches the best-known single-level sample-complexity guarantees for non-convex and strongly convex objectives.
- The paper leaves rigorous lower bounds and efficient methods for constrained lower-level problems for future work.
B Proofs of preliminary lemmas
The preliminary lemmas establish the stochastic upper-level gradient estimator’s construction, bias behavior, and Lipschitz continuity. The analysis uses randomized Hessian-inverse approximation and geometric-series bounds.
- The stochastic upper-level gradient estimator combines unbiased gradient and Hessian approximations using a randomly selected truncation depth.
- The estimator’s bias decreases linearly as the number K of Hessian-inverse approximation samples increases.
- The bias can be bounded by choosing the approximation parameters appropriately.
- The stochastic gradient estimate is Lipschitz continuous under the stated assumptions.
- The proofs repeatedly exploit uniform randomization over truncation depth and sums of geometric progressions.
C Proof of Theorem 3.2: smooth (possibly non-convex) outer objective
The theorem proof analyzes descent of the outer objective together with the lower-level iterates under smoothness and strong convexity. Its bounds explicitly track stochastic-gradient and estimation errors.
- The proof begins by establishing outer-objective descent for a smooth, possibly non-convex outer problem.
- The consecutive outer iterates satisfy an expected descent inequality involving the stochastic gradient estimator and its bias.
- The inner iterates satisfy a corresponding expected bound under strong convexity of the lower-level objective.
- The analysis uses the Lipschitz continuity of the solution mapping and Young’s inequality to control error terms.
C.3 Descent in the gradient estimation error of the outer function
This proof section controls the descent of the outer gradient-estimation error across SUSTAIN iterations. It combines recursive-estimator properties, conditional unbiasedness, variance bounds, and Lipschitz continuity.
- The analysis defines the iteration-history sigma algebra to condition expectations on the generated iterates.
- The recursive gradient estimator is used to relate consecutive outer estimation errors.
- Conditional zero-mean stochastic terms eliminate an inner-product contribution in the error analysis.
- The remaining error is bounded using a mean-variance inequality, Lipschitz continuity, and the algorithm’s update steps.
C.4 Descent in the gradient estimation error of the inner function
This section derives descent bounds for the inner gradient-estimation error and combines them through a potential function. Appropriate parameter choices yield the non-convex iterate result used later in the analysis.
- The analysis begins by studying descent in the inner-function gradient-estimation error under the algorithm’s stochasticity.
- A potential function combines the relevant error terms so that its expected change can be bounded across iterations.
- Lemma C.5 specifies parameter choices for Algorithm 1 and states the resulting iterate bound when the outer problem is non-convex.
- The proof obtains the bound by combining descent inequalities, controlling coefficients through step-size and momentum choices, and summing over iterations.
D Proof of Theorem 3.3: strongly-convex outer objective
The strongly-convex outer-case proof establishes descent relations for the outer and inner iterates while accounting for stochastic gradient-estimation errors. Strong convexity and Lipschitz properties supply the inequalities needed to combine these relations.
- The proof analyzes consecutive iterates generated by Algorithm 1 when the outer function is smooth and strongly convex.
- Lemma D.2 gives a stochastic bound for the inner iterates, with the expectation taken over the algorithm’s randomness.
- The inner-iterate bound uses Young’s inequality, the update steps, and the Lipschitzness of the lower-level solution mapping.
- Strong convexity of the lower-level objective and properties of its gradient are used to control the gradient-estimation terms.
D.3 Descent in the gradient estimation error
This section constructs a potential function for the strongly-convex analysis and selects step sizes, momentum parameters, and auxiliary coefficients to show descent. The resulting bounds apply uniformly over the iterations.
- The analysis states a bound for the gradient-estimation quantity with expectation over the algorithm’s stochasticity.
- A new potential function is introduced, and its decrease is established with an appropriate choice of parameters.
- Lemma D.4 specifies parameter choices for the potential-function descent, including η_t = (µ_f + 1)α_t and β_t = ĉ_βα_t.
- The proof combines bounds for the outer and inner terms and imposes upper bounds on α_t to control the resulting coefficients.
D.5 Proof of Theorem 3.3
For the strongly-convex outer objective, fixed step sizes allow recursive application of the descent inequality. The analysis concludes that reaching an ε-optimal solution requires a near-inverse-linear iteration and sample complexity.
- With fixed step sizes, the descent inequality is applied recursively and the resulting geometric progression is summed.
- The iteration requirement is obtained by selecting a constant step size and separately controlling the terms in the descent bound.
- T = ˜O(ε^-1) iterations suffice to reach an ε-optimal solution satisfying E[ℓ(x_t) − ℓ*] ≤ ε.
- Each iteration uses a batch of K = O(log(T)) samples, while the total sample complexity remains bounded as T = ˜O(ε^-1).