Bilevel Learning via Strategically Inexact Stochastic Gradient Descent
Authors: Mohammad Sadegh Salehi, Subhadip Mukherjee, Lindon Roberts, Matthias J. Ehrhardt
Organizations: Independent Researcher, Bath, UK · Department of Electronics & Electrical Communication Engineering, Indian Institute of Technology (IIT), Kharagpur, India · School of Mathematics and Statistics, University of Melbourne, Parkville VIC 3010, Australia · Department of Mathematical Sciences, University of Bath, Bath, BA2 7AY, UK
Bilevel optimization is a central tool in computational imaging for high-dimensional hyperparameter tuning; prominent applications include learning data-adaptive regularizers and optimizing forward operators in variational regularization. However, due to the nested nature of the problem, exact gradients with respect to the parameters (so-called hypergradients) are not available: the precision of the hypergradient is linearly related to the computational cost, thus requiring algorithms which can solve the problem without unnecessary precision. In learned-regularizer pipelines, the inner reconstruction solver is typically run for a fixed number of iterations, leaving the hypergradient accuracy uncontrolled and existing convergence theory inapplicable. In this work, we analyze the schedules that practitioners actually run: we prove convergence of inexact stochastic bilevel optimization and establish rates under both decaying accuracy and step size schedules, showing that, with optimal configurations, convergence occurs at an O(k−1/4) rate in expectation. This is the optimal rate of nonconvex stochastic gradient descent and matches single-loop stochastic bilevel methods; with warm-started inner solves, total costs also agree up to constants. Experiments on image denoising and inpainting with convex ridge regularizers and input-convex neural networks confirm our analysis: decreasing step sizes improve stability, accuracy scheduling is more critical than step size strategy, and adaptive preconditioning (e.g., Adam) further boosts performance, particularly for input-convex neural networks. In a budget-matched comparison with four state-of-the-art single-loop methods, accuracy-scheduled stochastic gradient descent and its Adam variant attain the highest test peak signal-to-noise ratio.
Figures & tables
θk+1=θk−αkzk
(4)
Algorithm 1 Inexact Stochastic Gradient Descent (ISGD) with varying step size and accuracy
Method
EK
Upper iters.
Stochastic model
Additional assumptions
SGD (single-level) [ 38 ]
K−1/2 (optimal [ 37 ] )
ϵ−4
unbiased, bounded variance
L -smoothness
ISGD (this work, Corollary 1 )
K−(1−q)→K−1/2
ϵ−4−δ′
certified biased , ES/ABC
Hessian-Lipschitz hi ( Lemma 1 )
ISGD, fixed horizon ( Corollary 3 )
K−1/2
ϵ−4
certified biased , ES/ABC
as above
TTSA [ 27 ]
K−2/5
ϵ−5
unbiased, bounded var.
Lipschitz ∇,∇2 ; projection
ALSET [ 28 ]
K−1/2
ϵ−4
unbiased, bounded var.
Lipschitz ∇,∇2
SOBA [ 31 ]
K−1/2
ϵ−4
unbiased, bounded var.
∇3 Lipschitz; bounded ∇1F
TABLE I: Convergence rates and assumptions for stochastic bilevel methods (nonconvex upper level, μ -strongly-convex lower level), unified to the squared-gradient metric. “Upper iters.” is the iteration count for E∥∇f∥≤ϵ (norm convention); papers quoting O(ϵ−2) use the squared convention E∥∇f∥2≤ϵ , which coincides with ϵ−4 here ( Remark 2 ). SGD row shown for reference.
Fig. 1: The accuracy schedule, more than the step size, determines finite-budget performance (denoising with CRR). Columns: initial accuracy ϵ0 . (a)–(c) Training loss with a fixed step size for different accuracy decay exponents p and initial step sizes α0 : a loose ϵ0 needs an aggressive decay, a tight one hardly any. (d)–(f) Test PSNR of the best fixed and decreasing step size configurations: once p is chosen well, a fixed step size gives the highest PSNR within the budget, and a decreasing one the smallest fluctuations. Training losses with decreasing step sizes are shown in Fig. 5 of the supplementary material.
Fig. 2: Inpainting with the CRR regularizer ( 70% of the pixels missing). Top: ground truth and observed test image, each with the enlarged region. Rows: the enlarged reconstruction at training checkpoints of computational cost 2,500 , 5,000 , 10,000 , and 20,000 (left to right) for ϵ0=10−1 , α0=10−2 , p=1 , and step size decay exponents q∈{0,0.25,0.5} ; inset numbers are test PSNR (dB). The fixed step size is ahead at every checkpoint.
Fig. 3: Denoising with the ICNN regularizer. Top: ground truth and noisy test image, with the red region enlarged 4× . Rows: the enlarged region at five training checkpoints of increasing computational cost for the best configurations of ISGD (rows 1–2) and of IAdam, its Adam variant (rows 3–4); inset numbers are test peak signal-to-noise ratio (PSNR, dB). All runs use α0=10−2 ; rows 1, 2, and 4 use (ϵ0,p)=(1,1) and row 3 uses (10−2,0.25) . At the first checkpoint, IAdam is 2.0 – 3.5 dB ahead of ISGD.
Fig. 4: Budget-matched comparison ( 2×104 total cost) of scheduled ISGD/IAdam with single-loop stochastic bilevel methods (TTSA, SOBA, SABA, MA-SOBA) on denoising with a CRR regularizer. (a) Running-average training loss and (b) test PSNR, both against the shared computational cost meter (lower-level iterations + Hessian-system iterations); axes cropped to the informative range, legend in (a).
Method (Configuration)
PSNR (dB)
SSIM
IAdam ( α0=10−2 , ϵ0=10−1 , p=1 )
29.34
0.828
ISGD ( α0=10−3 , ϵ0=10−1 , p=1 )
29.00
0.814
ISGD decaying α (per “ISGD” with q=0.5 )
28.73
0.800
MA-SOBA [ 32 ] ( α=10−3,β=10−2 )
28.06
0.758
SOBA [ 31 ] ( α=10−3,β=10−2 )
27.99
0.775
SABA [ 31 ] ( α=3⋅10−4,β=3⋅10−3 )
26.79
0.695
TABLE II: Best test PSNR/SSIM within a total computational budget of 2×104 (denoising with CRR on a fixed dataset). All methods use warm-started lower-level iterates. Baselines run in their published form under their own step-size schedules; constants of all methods are tuned by long-horizon grid search (over (α,β) for the baselines and over (α0,ϵ0,p) for the scheduled methods).
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Step size αk
Accuracy ϵk
Gradient rate
Note
O(k−q),21<q<1
O(k−p),2p<1−q
O(k−p)
limited by accuracy
O(k−p),2p=1−q
O(k1−qlogk)
boundary: slower with logk
O(k−p),2p>1−q
O(k−(1−q)/2)
limited by step size
O((logk)−p)
O((logk)−p)
any p>0 ; slowest
Appendix
TABLE III: Convergence rates for different step size and accuracy decay regimes.
Fig. 5: Denoising with CRR: training loss with a decreasing upper-level step size ( q>0 ) for the three initial accuracies of Fig. 1 in the main text. Decreasing the step size reduces the variance, but the ranking of the accuracy schedules is the same as with a fixed step size.
Fig. 6: Best-performing fixed and decreasing step size configurations across all accuracy settings ϵ0∈{100,10−2,10−4} . (a) Training loss plotted against computational cost. (b) Test PSNR plotted against computational cost.
Fig. 7: Denoising with the CRR regularizer at training checkpoints of computational cost 5,000 , 10,000 , 50,000 , and 100,000 (left to right) for two configurations. Top: ground truth and noisy test image with the red region, which is shown enlarged in all rows; inset numbers are test PSNR (dB).
Fig. 8: Running average training loss and average test PSNR for inpainting, plotted against total computational cost. Results are shown for upper-level step size α0=10−2 , initial accuracy ϵ0=10−1 , accuracy schedule exponent p∈{0.5,1} , and step size decay exponent q∈{0,0.25,0.5} . All settings with p=1 appear to perform faster than their p=0.5 counterparts.
Fig. 9: Comparison between SGD and Adam with inexact stochastic hypergradients across different initial accuracies, using the most suitable values of accuracy schedule exponent p and step size schedule exponent q for each, found by trial and error. Top row: (a) shows configurations with fixed step size; (b) corresponds to decreasing step size. Bottom row: best-performing configurations of ISGD and IAdam in terms of training loss and test PSNR in (c) and (d), respectively.
Fig. 10: Empirical validation of the theoretical refinements, with the rate-optimal accuracy schedule ϵk=O(k−1/4) . (a) A horizon-tuned constant step size (Corollary 3) outperforms the anytime decaying schedule at equal cost. (b) With persistent warm starts the lower-level iterations per upper-level step decay towards the O(1) amortized regime of Remark 3, whereas cold starts grow as predicted by the logk term of Corollary 2.
School of Mathematics and Statistics, Xi’an Jiaotong University · School of Data Science, Fudan University · 3SGIT AI Lab, State Grid Corporation of China +1
Mathematical Institute, University of Oxford, United Kingdom · Department of Business Analytics and Decision Sciences, University of Economics and Business, Vienna, Austria · Bernouilli Institute, University of Groningen, The Netherlands