Source-linked AI summary
Provably Faster Algorithms for Bilevel Optimization
Junjie Yang, Kaiyi Ji, Yingbin Liang
TL;DR
Bilevel optimization requires efficient hypergradient computation, but prior momentum-based methods lacked a provable complexity improvement over SGD-type approaches. This paper proposes MRBO and VRBO, using recursive momentum and nested-loop variance reduction, and reports Ẽ(ε^-1.5) complexity for both with favorable experimental performance.
Problem
Hypergradient-based bilevel optimization is computationally challenging, and prior momentum-based methods had not proved lower complexity than SGD-type approaches.
Method
The paper proposes MRBO, which uses recursive momentum estimators, and VRBO, which uses recursive variance reduction in a double-loop design.
Results
Both algorithms achieve Ẽ(ε^-1.5) computational complexity and outperform existing algorithms in complexity order.
Takeaways & Limitations
The experiments report superior performance and suggest that the double-loop VRBO design may be more suitable for bilevel optimization than the single-loop structure.
Takeaways & Limitations
The CIFAR10 experiments were still ongoing.
Abstract
from arXiv · showhide
Bilevel optimization has been widely applied in many important machine learning applications such as hyperparameter optimization and meta-learning. Recently, several momentum-based algorithms have been proposed to solve bilevel optimization problems faster. However, those momentum-based algorithms do not achieve provably better computational complexity than $\mathcal{\widetilde O}(ε^{-2})$ of the SGD-based algorithm. In this paper, we propose two new algorithms for bilevel optimization, where the first algorithm adopts momentum-based recursive iterations, and the second algorithm adopts recursive gradient estimations in nested loops to decrease the variance. We show that both algorithms achieve the complexity of $\mathcal{\widetilde O}(ε^{-1.5})$, which outperforms all existing algorithms by the order of magnitude. Our experiments validate our theoretical results and demonstrate the superior empirical performance of our algorithms in hyperparameter applications.
1 Introduction
The paper targets faster stochastic bilevel optimization in the nonconvex-strongly-convex setting, where hypergradients require costly Jacobian and Hessian information. It introduces MRBO and VRBO to improve complexity through recursive momentum and variance reduction.
- Bilevel optimization supports applications including hyperparameter optimization, meta-learning, and reinforcement learning.
- Hypergradient computation is challenging because the outer objective depends on the inner optimizer and involves Jacobian and Hessian matrices.
- The paper studies nonconvex-strongly-convex problems, including settings where hyperparameters are outer variables and regularization makes the inner problem strongly convex.
- Existing SGD-type stocBiO achieves Ẽ(ε^-2) complexity, while prior momentum-based bilevel algorithms had not proved an order-wise improvement.
- MRBO uses momentum-based recursive estimators for both inner gradients and hypergradients, achieving Ẽ(ε^-1.5) complexity with Jacobian- and Hessian-vector products.
- VRBO applies recursive variance reduction in a double-loop design, and experiments report higher accuracy and faster convergence than compared single-loop methods.
2 Two New Algorithms
The paper develops MRBO as a single-loop recursive momentum method and VRBO as a double-loop recursive variance-reduction method. Both construct estimators using vector products rather than explicit Hessians, while VRBO periodically uses large batches and recursively updates with smaller batches.
- MRBO: MRBO updates x and y simultaneously in a single loop using STORM-style recursive momentum estimators.The method applies recursive estimators to both the inner gradient and the hypergradient.
- MRBO: MRBO stabilizes gradient estimates by increasing reliance on previous estimates as the momentum coefficients decrease.
- MRBO: MRBO recursively computes hypergradient estimates with Hessian-vector products rather than explicit Hessians to reduce memory and computational cost.
- VRBO: VRBO divides iterations into epochs, periodically computes large-batch gradient and hypergradient estimates, and recursively updates them with small batches inside each epoch.
- VRBO: VRBO is the first bilevel method described here to use recursive variance reduction and achieves the same nearly-optimal complexity as MRBO.
3 Main Results
The paper analyzes MRBO and VRBO under nonconvex-strongly-convex bilevel assumptions, addressing hypergradient estimation, inner-loop tracking, and recursive variance reduction. Both methods obtain near-optimal complexity while improving over existing stochastic bilevel algorithms.
- MRBO convergence analysis: MRBO’s analysis must jointly control hypergradient momentum variance, inner-loop tracking error, and Hessian-vector approximation error.Its hypergradient combines first-order gradients with Hessian-vector products, making the recursive estimator analysis more involved than in standard minimization.
- MRBO convergence analysis: MRBO’s momentum recursive estimator reduces hypergradient variance across iterations, while minibatch size S and Neumann-series length Q control estimator variance.The per-iteration bound includes a variance-reduction term and stochastic-update terms; increasing S reduces variance, while Q affects the Neumann approximation.
- MRBO convergence analysis: Theorem 1 establishes simultaneous convergence of xk, yk, and the momentum-estimator error, including vanishing inner tracking error and estimator variance.These properties facilitate convergence of the outer variable xk.
- MRBO convergence analysis: MRBO reaches an ϵ-stationary point with gradient and Hessian-vector complexity of O(ϵ^-1.5) and Ẽ O(ϵ^-1.5), respectively.The result uses K = O(ϵ^-1.5) and Q = O(log(1/ϵ)).
- VRBO convergence analysis: VRBO recursively reduces variance for both inner gradients and hypergradients through large-batch estimates periodically computed in outer loops.Its analysis bounds cumulative estimator variance and accounts for minibatch error and exponentially decreasing Hessian-vector approximation error.
- VRBO convergence analysis: VRBO attains gradient and Hessian-vector complexity of Ẽ O(ϵ^-1.5), outperforming existing stochastic bilevel algorithms by a factor of Ẽ O(ϵ^-0.5).The paper reports that VRBO empirically performs better than MRBO and several single-loop momentum-based algorithms.
4 Experiments
The experiments compare VRBO and MRBO with deterministic, stochastic, and momentum-based bilevel optimizers using training loss versus running time. VRBO and MRBO show faster convergence than the compared methods, with VRBO also achieving lower and more stable training loss.
- 4 Experiments: The experiments include comparisons with AID-FP, reverse, BSA, MSTSA, SUSTAIN, STABLE, and stocBiO.The listed benchmarks include double-loop deterministic, double-loop stochastic, and single-loop stochastic algorithms.
- 4 Experiments: VRBO and SGD-type stocBiO converge much faster than the other compared algorithms in training loss versus running time.The two methods have comparable performance, while VRBO achieves lower training loss and more stable convergence.
- 4 Experiments: VRBO converges significantly faster than all single-loop momentum-based methods.The comparison provides evidence favoring double-loop algorithms over single-loop algorithms in these experiments.
- 4 Experiments: MRBO achieves the fastest convergence rate among all single-loop momentum-based algorithms, consistent with the theoretical results.In the comparison with MSTSA, STABLE, and SUSTAIN, both MRBO and VRBO are significantly faster than the other methods.
- 4 Experiments: Double-loop VRBO performs best overall, while double-loop stocBiO, reverse, and AID-FP outperform single-loop momentum-accelerated MRBO.Double-loop BSA also converges faster than single-loop momentum-accelerated MSTSA.
5 Conclusion
The paper proposes MRBO and VRBO for nonconvex-strongly-convex bilevel stochastic optimization and reports order-wise complexity improvements over existing algorithms. MRBO introduces momentum-based recursive iterations, while VRBO uses recursive variance reduction; experiments support their superior performance and favor double-loop designs.
- 5 Conclusion: MRBO and VRBO target nonconvex-strongly-convex bilevel stochastic optimization and outperform existing algorithms in computational complexity order-wise.MRBO is identified as the first momentum algorithm with an order-wise improvement over SGD-type algorithms, while VRBO is the first to use recursive variance reduction for bilevel acceleration.
- 5 Conclusion: Experiments demonstrate superior performance for MRBO and VRBO and suggest that double-loop designs may suit bilevel optimization better than single-loop structures.The paper also anticipates extending the analysis to other loss geometries and motivating further double-loop versus single-loop comparisons.
A Hessian Vector Implementation
Algorithm 3 computes MRBO’s hypergradient estimator recursively using Hessian-vector and Jacobian-vector products rather than explicit Hessian matrices. This implementation reduces memory and computational cost while supporting automatic differentiation.
- A Hessian Vector Implementation: Algorithm 3 computes the MRBO hypergradient estimator using Hessian vectors rather than Hessian matrices.The stated purpose is to reduce memory and computational cost.
- A Hessian Vector Implementation: The recursion initializes gradient-related quantities, repeatedly updates G_q and r_q, accumulates M_Q, and returns the hypergradient estimator.The displayed procedure uses Q Hessian-estimation steps and returns ∇xF(x, y; BF ) − ∂(∇yG(x, y; BG)M_Q)/∂x.
- A Hessian Vector Implementation: Instead of explicitly forming a Hessian-vector product, the method computes r_q+1 through the derivative ∂(G_q+1)/∂y.A similar automatic-differentiation implementation computes the Jacobian-vector product in the final return expression.
- A Hessian Vector Implementation: Both the Hessian-vector and Jacobian-vector computations can use torch.grad() for implementation.The recursive procedure therefore uses Hessian-vector products without explicitly computing Hessian matrices.
B Specifications of Experiments
The experiments compare MRBO and VRBO with several bilevel optimization benchmarks on MNIST hyper-cleaning. They specify the dataset sizes, hypergradient implementation, batch sizes, loop settings, and CPU-based training environment.
- B Specifications of Experiments: MRBO and VRBO are compared with stocBiO, reverse, AID-FP, BSA, MSTSA, STABLE, and SUSTAIN on MNIST hyper-cleaning.The experiment uses the data hyper-cleaning formulation and follows the stated benchmark setup.
- B Specifications of Experiments: The experiment uses 20,000 training images, 5,000 validation images, and 10,000 test images.The setup fixes the training and validation sizes and uses 10,000 images for testing.
- B Specifications of Experiments: The hypergradient estimator uses the Hessian-vector implementation with Q = 3 and η = 0.5.MRBO, VRBO, and stocBiO use batch size 1000 for both training and validation procedures.
- B Specifications of Experiments: VRBO uses inner batch size 500 and period q = 3, while double-loop baselines use 200 inner-loop steps except VRBO’s 20 steps.The double-loop settings are tuned for best performance in the reported experiment.
- B Specifications of Experiments: Training runs on an iMac with a 3.8GHz quad-core Intel Core i5 CPU and 32 GB of 2400 MHz DDR4 memory.The code also supports GPU-cluster training.
B.1 Additional Experiments of Hyper-cleaning
Additional hyper-cleaning experiments compare the proposed algorithms with momentum-based and batch-sample baselines using running time, inner-loop settings, and sample efficiency. MRBO and VRBO generally achieve strong training performance, while VRBO is most sensitive to the number of inner-loop steps.
- Momentum-based comparisons: MRBO and VRBO achieve the lowest training loss and fastest convergence among momentum-based algorithms at noise rate p = 0.15.The comparison includes MSTAS, STABLE, and SUSTAIN.
- Inner-loop sensitivity: VRBO becomes more stable and lowers training loss as inner-loop steps increase from 1 to 20.The experiments compare m ∈ {1, 10, 20, 50, 100}.
- Inner-loop sensitivity: Further increasing VRBO’s inner-loop steps beyond 20 worsens stability and increases final training loss.The paper attributes this to unnecessary iterations and less accurate variance-reduced estimators later in the inner loop.
- Sample efficiency: MRBO and VRBO are much more sample efficient than stocBiO, reverse, and AID-FP among batch-sample algorithms.This comparison evaluates training performance against the number of required samples.
- Sample efficiency: Single-sample algorithms use fewer samples than MRBO and VRBO, but batch-sample algorithms converge much faster with respect to running time.The paper explains this difference through minibatch computation efficiency on platforms such as PyTorch.
B.2 Experiments of Logistic Regression
The logistic-regression experiments evaluate test performance against running time at batch sizes 100 and 1000. Double-loop methods generally achieve higher accuracy than single-loop methods, although stocBiO is strongest in one large-batch setting.
- Batchsize 100: At batchsize 100, MRBO and VRBO reach higher accuracy than stocBiO despite its faster initial convergence.The paper attributes this to more accurate hypergradient estimation.
- Batchsize 100: At batchsize 100, MRBO(D) achieves the highest accuracy, while single-loop SUSTAIN and MSTSA do not converge well.This comparison demonstrates an advantage of double-loop updates over single-loop updates.
- Batchsize 1000: At batchsize 1000, double-loop stocBiO and MRBO(D) significantly outperform the other single-loop algorithms.The comparison is based on test accuracy or test loss versus running time.
C.3 Proof of Corollary 1
The proof establishes the complexity guarantee for MRBO by setting the iteration count and Hessian-estimation depth appropriately. With these choices, MRBO reaches an ϵ-stationary point using near-optimal gradient and Hessian-vector complexity up to logarithmic factors.
- Corollary 1: Choosing K = O(ϵ^-1.5) and Q = O(log(1/ϵ)) makes MRBO find an ϵ-stationary point.The result follows under the conditions of Theorem 1.
- Corollary 1: MRBO achieves gradient complexity O(ϵ^-1.5) and Jacobian-Hessian-vector complexity Õ(ϵ^-1.5).The Hessian-vector complexity includes the logarithmic dependence associated with Q.
- Complexity accounting: The proof obtains target accuracy from an error bound of O(1/K^2/3) = O(ϵ) under the selected iteration schedule.The gradient and Jacobian-vector costs are KS = O(ϵ^-1.5), while Hessian-vector cost is KSQ = Õ(ϵ^-1.5).
D.3 Proof of Corollary 2
The proof of Corollary 2 selects batch sizes, epoch length, Hessian-estimation depth, and outer iterations for VRBO. These settings yield an ϵ-stationary point with Õ(ϵ^-1.5) gradient and Hessian-vector complexity.
- Corollary 2: With S1 = O(ϵ^-1), S2 = O(ϵ^-0.5), Q = O(log(1/ϵ^0.5)), and K = O(ϵ^-1), VRBO finds an ϵ-stationary point.These choices are made under the conditions of Theorem 2.
- Complexity accounting: VRBO achieves Õ(ϵ^-1.5) gradient complexity and Õ(ϵ^-1.5) Hessian-vector complexity.The proof derives both rates from the periodic large-batch and inner-loop costs.