A framework for bilevel optimization that enables stochastic and global variance reduction algorithms

1 Introduction

• Computing the value-function gradient requires approximating both the inner solution and a linear-system solution, making first-order methods challenging to implement.Bilevel optimization applies to settings including hyperparameter selection, data augmentation, implicit deep learning, and neural architecture search.Source: §1 · pp. 1–2
• For large empirical-mean objectives, stochastic methods use only a few samples, but the inverse Hessian in the gradient makes unbiased estimates difficult.The introduction notes that such methods are the only scalable algorithms when evaluating F or G is extremely expensive.Source: §1 · p. 2
• The framework jointly evolves the inner variable, linear-system solution, and outer variable, expressing their directions as sums that yield unbiased stochastic estimators.It gives rise to SOBA, an SGD extension, and SABA, a SAGA extension.Source: §1 · p. 2
• SABA matches gradient descent’s convergence rate on the value function and achieves linear convergence when the value function satisfies a Polyak-Łojasciewicz inequality.The authors describe these as the first such guarantees for a stochastic bilevel algorithm.Source: §1 · p. 2
• The paper reports benchmarks of stochastic bilevel methods on hyperparameter selection and data hyper-cleaning to illustrate the proposed approach’s usefulness.The introduction also lists those tasks among its numerical experiments.Source: §1 · p. 2
• Prior approaches include implicit differentiation, which can require storing inner iterates, and stochastic methods using nested or single-loop updates.Approximate implicit differentiation can use approximate inner and linear-system solutions when their errors decrease sufficiently quickly.Source: §1 · p. 3

2 Proposed framework

• When G(·, x) is strongly convex, any zero of the three directions gives z = z*(x), v = v*(x), and ∇h(x) = 0.The framework jointly evolves z, v, and x along directions written as sums of derivatives of F_j and G_i, making unbiased estimates straightforward.Source: §2 · pp. 3–4

2.1 First example: the SOBA algorithm

• SOBA estimates each direction using one uniformly sampled G_i and F_j, with updates computed at the same point and in parallel.The method uses different step sizes for the inner and outer updates to account for their potentially different conditioning.Source: §2.1 · p. 4

2.2 Global variance reduction with the SABA algorithm

• SABA adapts SAGA by maintaining memories for the joint directions; its variance-reduced estimates have variance tending to zero and O(1/T) convergence.Its convergence rate matches gradient descent on h, while storing computed per-sample quantities requires additional memory.Source: §2.2 · pp. 4–5

3.1 Background and assumptions

• The guarantees use a finite-sample setting and regularity assumptions, including strong convexity of G in z and conditions ensuring smoothness of v*.The paper notes that its regularity assumptions are stronger than those in some prior stochastic bilevel work.Source: §3 · p. 5, §3.1 · pp. 5–6

3.2 Fundamental descent lemmas

• The descent analysis tracks inner-solution and linear-system errors; when both vanish, the outer direction is ∇h and the inequality resembles SGD's.The descent lemma is stated for the joint-update framework with arbitrary unbiased directions.Source: §3.2 · pp. 6–7

3.3 Analysis of SOBA

• SOBA reaches O(epsilon^-2) oracle complexity, matching non-convex smooth SGD; decreasing step sizes are needed for actual convergence.The analysis also gives a fixed-step result with steps depending on the iteration horizon.Source: §3.3 · p. 7

3.4 SABA: a stochastic method with optimal rates

• SABA improves the epsilon scaling from epsilon^-2 for SOBA to epsilon^-1, and converges linearly when h satisfies the PL condition.The linear-convergence guarantee is stated for global variance reduction across the variables.Source: §3.4 · p. 8

4 Experiments

• The comparison includes stocBiO, AmiGO, FSLA, MRBO, TTSA, BSA, and SUSTAIN, alongside the proposed SOBA and SABA.The hyperparameter task selects regularization parameters for ℓ2 logistic regression; the data-cleaning task learns weights for training samples, with corrupted labels intended to receive zero weight.Source: §4 · p. 8, §4.1 · p. 8, §4.2 · p. 9

4.1 Hyperparameters selection

• On IJCNN1 hyperparameter selection, SABA reaches the lowest suboptimality, and is the only single-loop method to go below 10^-3.SOBA finishes slightly better than TTSA and FSLA, while variance reduction gives a lower plateau and fixed step sizes enable faster convergence.Source: §4.1 · pp. 8–9

4.2 Data hyper-cleaning

• In both experiments, SABA achieves the best median performance over 10 runs; on data hyper-cleaning it reaches a smaller error faster than the other methods.The data hyper-cleaning comparison includes variance-reduced methods FSLA, MRBO, and SUSTAIN.Source: Figure 2 · p. 9, §4.2 · p. 9
• For data hyper-cleaning, SABA is fastest to its smallest test error, which is below the errors of the other methods; SOBA finishes below stocBiO and BSA.The experiment uses MNIST with corruption probability p = 0.5; for higher corruption probabilities, SABA is still reported as fastest to its final accuracy.Source: §4.2 · p. 9, Figure 2 · p. 9

• The framework writes its update directions as simple sums of sample derivatives, enabling straightforward development of stochastic bilevel algorithms.Source: §5 · p. 9
• SOBA extends SGD and SABA extends SAGA within the framework, with both achieving rates similar to their single-level counterparts.Source: §5 · p. 9
• The authors identify extrapolation, variance reduction, and momentum as potential directions for further stochastic bilevel methods.Source: §5 · p. 9

• The authors report that limitations were discussed and theoretical assumptions and proofs were provided.Source: Checklist · p. 13
• They report providing code, data, instructions, training details, error bars, and compute-resource information for the experiments.Source: Checklist · p. 13
• The checklist marks asset citation and licensing as addressed, while human-subject and societal-impact items are marked not applicable.Source: Checklist · p. 13

A Extensive comparison between stochastic methods for bilevel optimization

• The comparison groups prior solvers by one-loop or two-loop design, describes HIA and SHIA as Neumann-based Hessian-inverse approximations, and defines complexity by oracle calls to an ϵ-stationary point.Source: A Extensive comparison between… · p. 14
• The appendix describes STORM momentum updates as requiring two evaluations per iteration and memory for the previous estimates and running gradient estimate.Source: A Extensive comparison between… · p. 14

B.1 Generalities

• Experiments use contiguous batches of 64, typically 10 inner steps for methods with inner loops and 10 Hessian-approximation steps; step-size exponents follow the original methods’ theoretical choices.Source: §B.1 · p. 16

B.2 Hyperparameter selection on a toy problem

• The toy hyperparameter-selection experiment uses Ridge regression and generates 1,000 samples in 10 dimensions, with 750 pairs used for training and the remaining pairs for validation.Source: §B.2 · p. 16

B.3 Hyperparameters selection on IJCNN1

• The IJCNN1 study selects one regularization hyperparameter per feature for multiregularized logistic regression and uses grid search to choose step sizes.Source: §B.3 · p. 17
• In both IJCNN1 plots, SABA achieves the best performance; the plots compare median results over 10 runs by running time and sampled gradient/Hessian-vector products.Source: Figure B.1 · p. 17

B.4 Data hyper-cleaning

• The regularization parameter Cr is set to 0.2 after manual search for the best final test accuracy.Source: §B.4 · p. 17
• The experiment selects α and β using grid search over logarithmically spaced candidate values.α uses 11 values from 10^-3 to 100, and β is set to αr with r selected from 11 values from 10^-5 to 1.Source: §B.4 · p. 17
• SABA is the fastest algorithm to reach its final accuracy at corruption probabilities 0.5, 0.7, and 0.9.The main-text figure corresponds to p = 0.5.Source: §B.4 · p. 17

B.5 Additional experiment: Hyperparameter selection on the covtype dataset

• The experiment selects one regularization hyperparameter per class for ℓ2-regularized multinomial logistic regression on covtype, using batch size 64 and 525 CPU hours.The dataset has 581,012 samples, 54 features, and 7 classes; the split uses 371,847 training, 92,962 validation, and 116,203 test samples.Source: §B.5 · pp. 17–18
• SABA and SOBA achieve the best performances on covtype; SABA’s initial runtime gap is due to memory initialization and can be reduced with a larger batch size.Figure B.3 compares median performance over 10 runs by running time and sampled gradient/Hessian-vector products.Source: §B.5 · p. 18, Figure B.3 · p. 19

C.1 Proof of Proposition 2.1

• A zero of (Dz, Dv, Dx) has z = z*(x), v = v*(x), and ∇h(x) = 0.The proof uses strong convexity of G in z to identify z as the inner minimizer, then the linear-system equation to identify v.Source: §C.1 · p. 19

C.2 Proof of Lemma 3.4

• The proof of Lemma 3.4 bounds the residual involving ∇²₁₁G and ∇₁F by decomposing it into terms controlled by smoothness and strong convexity.The displayed derivation uses LG₁-smoothness of G, LF₁-smoothness of F, and µG-strong convexity of G.Source: §C.2 · p. 19

C.3 Smoothness constant of h

• The appendix invokes Ghadimi and Wang [19, Lemma 2.2] to establish the Lh-smoothness of h.The supplied passages give components of the smoothness constant but not a complete legible expression.Source: §C.3 · p. 20

C.4 Lemmas on the regularity of z∗and v∗

• Under the stated regularity assumptions, z* is Lipschitz and smooth, while v* is differentiable and its differential is Lipschitz continuous.The proof also states that v* is Lvx-smooth.Source: §C.4 · pp. 20–23

C.5 Proof of Lemma 3.9

• The proof of Lemma 3.9 derives bounds for δz and δv using unbiasedness, strong convexity, and regularity of z* and v*.For δz, it uses a linear approximation of z*; for δv, it uses the zero of Dv at (z*(xt), v*(xt), xt).Source: §C.5 · pp. 24–25
• The C.5 proof applies similar manipulations to v* using Lemma C.3 and a bound analogous to Lemma C.1 for z, for any η > 0.Source: §C.5 · p. 26

C.6 Proof of Lemma 3.10

• The C.6 proof uses smoothness of h and Lemma 3.4 to derive a bound involving the expected squared norm of Dx and delta terms.Source: §C.6 · p. 26

C.7 Proof of Theorem 1

• Theorem 1 considers SOBA under Assumptions 3.1–3.7 with fixed step sizes ρ_t = ρ√T and γ_t = ξρ_t.The supplied theorem statement does not include a legible bound for its conclusion.Source: §C.7 · p. 27
• The proof of Theorem 1 adapts descent lemmas for SOBA, bounds the change in its Lyapunov quantity, and sums the resulting inequalities.Source: §C.7 · pp. 27–29

C.8 Proof of Theorem 2

• Theorem 2's proof uses step sizes that decrease with iteration and sums and rearranges the resulting inequality.The supplied passages do not state the theorem's conclusion.Source: §C.8 · p. 29

C.9 Proof of Theorem 3

• Theorem 3 concerns SABA in the smooth case, under Assumptions 3.1–3.3 and 3.7–3.8, with ρ = ρ′N^-2/3 and γ = ξρ.The supplied theorem statement does not include a legible bound for its conclusion.Source: §C.9 · p. 29

C.9.1 Control of distance from memory to iterates

• SABA maintains parallel memories for its variables, and the proof controls the distance between iterates and those memories.Source: §C.9.1 · p. 30

C.9.2 Bounds on the variances

• The variance analysis bounds errors for SABA's variable directions using memory-to-iterate differences and smoothness assumptions.Source: §C.9.2 · pp. 31–32

C.9.3 Putting it all together

• The proof combines the bounds using a Lyapunov function and scales its step sizes and coefficients with N = n + m.It then sums and telescopes inequalities to complete the argument.Source: §C.9.3 · pp. 35, 36, 38

C.10 Proof of Theorem 4

• Theorem 4 assumes that h satisfies the PL inequality and that Assumptions 3.1–3.3 and 3.7–3.8 hold.Source: §C.10 · p. 38
• The proof sets ρ = ρ′N^-2 and γ = ξρ′N^-1, with ρ′ and ξ depending only on F and G.Source: §C.10 · p. 38
• The convergence bound’s constant C0 depends on the initialization of z, v, x, and the memory.Source: §3.4 · p. 8
• The proof uses the Lyapunov function from Equation (143) and then unrolls the resulting inequalities.Source: §C.10 · pp. 39, 42

D Convergence rates with weaker regularity assumptions

• The analysis requires stronger assumptions than those used in the stochastic bilevel optimization literature.Source: D Convergence rates with weaker… · p. 42
• The assumptions require G(·, x) to be μG-strongly convex and its first and second derivatives to be Lipschitz continuous.Source: D Convergence rates with weaker… · p. 42
• The analysis does not ensure that v* is smooth, so its descent lemmas use a different bound based on Young’s inequality.Source: D Convergence rates with weaker… · p. 42
• The weaker-regularity analysis requires ρ to go to zero for convergence, preventing rates that match single-level algorithms.Source: D Convergence rates with weaker… · p. 43
• The reported convergence rate is O(T^-2/5).Source: D Convergence rates with weaker… · p. 43
• SABA has sample complexity O((n + m)ε^-1), matching the sample complexity of SOBA with full-batch estimated directions.Source: D Convergence rates with weaker… · p. 43