Source-linked AI summary
Truncated Back-propagation for Bilevel Optimization
Amirreza Shaban, Ching-An Cheng, Nathan Hatch, Byron Boots
TL;DR
High-dimensional bilevel optimization makes exact gradient evaluation computationally difficult, while truncated back-propagation lacked clear theoretical guarantees. The paper establishes convergence conditions for these approximate gradients and finds competitive task performance with roughly half the computation time and substantially lower memory use, while identifying settings where convergence can fail.
Problem
Exact gradients are computationally challenging for high-dimensional bilevel problems, and the theoretical properties of truncated back-propagation remain unclear.
Method
The paper analyzes truncated back-propagation through the iterative lower-level optimization procedure and derives sufficient conditions for approximate or exact stationary-point convergence.
Results
Across hyperparameter optimization and meta learning tasks, truncated back-propagation usually has competitive application-specific performance, roughly twice the speed, and significantly lower memory use than full back-propagation.
Takeaways & Limitations
Truncated back-propagation is a reasonable choice for hyperparameter optimization or meta learning applications with memory constraints.
Takeaways & Limitations
Exact convergence appears uncommon in practice, and a full-rank assumption used in one proof may be relaxable only with further work.
Abstract
from arXiv · showhide
Bilevel optimization has been recently revisited for designing and analyzing algorithms in hyperparameter tuning and meta learning tasks. However, due to its nested structure, evaluating exact gradients for high-dimensional problems is computationally challenging. One heuristic to circumvent this difficulty is to use the approximate gradient given by performing truncated back-propagation through the iterative optimization procedure that solves the lower-level problem. Although promising empirical performance has been reported, its theoretical properties are still unclear. In this paper, we analyze the properties of this family of approximate gradients and establish sufficient conditions for convergence. We validate this on several hyperparameter tuning and meta learning tasks. We find that optimization with the approximate gradient computed using few-step back-propagation often performs comparably to optimization with the exact gradient, while requiring far less memory and half the computation time.
1 INTRODUCTION
Bilevel optimization supports hyperparameter optimization and meta learning, but exact gradient computation becomes difficult in high dimensions. The paper studies truncated back-propagation, establishing convergence conditions and competitive empirical performance with lower computational costs.
- Motivation: Bilevel optimization formulates hyperparameter optimization and meta learning as nested optimization problems involving an approximate lower-level solution.The lower-level solution is produced by a prespecified iterative algorithm that may itself be parametrized by hyperparameters such as step size.
- Motivation: The upper-level objective depends on the lower-level solution, making gradient computation challenging when both hyperparameters and lower-level variables are high-dimensional.This difficulty can preclude black-box approaches such as grid search and Bayesian optimization.
- Related approaches: Full differentiation through the lower-level optimization can learn algorithmic hyperparameters, but its superlinear time or space complexity limits applicability in high-dimensional settings.This limitation motivates alternatives that avoid tracking long-term dependencies.
- Truncated back-propagation: Truncated back-propagation reduces approximate-gradient time and space costs by ignoring long-term dependencies, although its theoretical properties had remained unclear.The approach was proposed in both few-step reverse-mode and few-step forward-mode forms.
- Contributions: When the lower-level problem is locally strongly convex near its approximate solution, O(log 1/ϵ)-step truncated back-propagation guarantees on-average convergence to an ϵ-approximate stationary point.Additional problem structures can guarantee asymptotic convergence to an exact stationary point.
- Contributions: Across hyperparameter optimization and meta learning tasks, truncated back-propagation usually achieves competitive performance with full back-propagation while using half the computation time and significantly less memory.The paper evaluates the strategy empirically on these application families.
2 BILEVEL OPTIMIZATION
The paper computes bilevel hypergradients by differentiating through the lower-level iterative dynamics using reverse- or forward-mode automatic differentiation. These modes trade memory against computation, motivating truncated reverse-mode back-propagation for large problems.
- Hypergradient computation: The hypergradient combines the direct hyperparameter derivative with the lower-level solution's total derivative contribution.The computational bottleneck is the matrix-vector product involving the Jacobian of the approximate lower-level solution.
- Iterative dynamics: The lower-level optimizer is represented as a dynamical system whose transition function updates the state for T iterations from an initial condition depending on λ.For gradient descent, the transition uses a potentially hyperparameter-dependent step size.
- Iterative dynamics: Unrolling the dynamics enables computation of the required derivative by applying the chain rule through the transition Jacobians.The matrices A_t and B_t encode derivatives with respect to the lower-level state and hyperparameters, respectively.
- Reverse and forward modes: Reverse-mode differentiation computes the hypergradient by back-propagating state sensitivities, whereas forward-mode differentiation propagates a matrix of hyperparameter sensitivities.The two procedures compute the same derivative through opposite traversal directions.
- Complexity trade-offs: Reverse-mode differentiation requires storing all intermediate lower-level states, while forward-mode avoids that storage but propagates an M×N matrix and is N-times slower than reverse mode.Checkpointing can reduce reverse-mode memory but doubles its computation time.
3 TRUNCATED BACK-PROPAGATION
K-step truncated back-propagation provides a computationally cheaper hypergradient whose bias can decay exponentially under local strong convexity, with stronger assumptions yielding exact-stationary-point convergence.
- Method: K-step truncated back-propagation uses the intermediate variable hT−K to construct an approximate gradient while storing only the last K iterates.The method can also be combined with checkpointing, although that combination is not investigated here.
- General properties: Under local strong convexity, the bias of the truncated gradient becomes exponentially small in K.This property means that a small truncation depth can provide a good gradient approximation in finite precision.
- Convergence: K-step truncated back-propagation converges on-average to an ϵ-approximate stationary point with ϵ = O((1 − γα)^−K).Choosing K = O(log 1/ϵ) is sufficient to update λ under the stated assumptions.
- Convergence: With second-order continuous differentiability, full column rank of Bt, and additional non-interference assumptions, every K ≥ 1 can yield convergence to an exact stationary point.The result requires T large enough and γ small enough; the non-interference assumption permits convergence even with one-step back-propagation.
- Convergence: The non-interference assumption is necessary for this guarantee because a problem satisfying the other assumptions can still fail to converge when it is violated.The paper also notes that this assumption holds when ∇λf = 0, including many listed practical applications.
- Relationship with implicit differentiation: In the limit, hT−K approximates the implicit-differentiation matrix inverse with an order-K Taylor series, while truncated back-propagation remains usable with approximate lower-level solutions.Compared with K-step conjugate gradients, truncated back-propagation has a slower local bias rate but does not require convergence of wt and can optimize lower-level-process hyperparameters.
4 EXPERIMENTS
Experiments on a toy problem, data hypercleaning, and task interaction examine truncated back-propagation’s convergence behavior, bias, and computational trade-offs. Few-step methods generally provide useful optimization behavior, with performance close to full back-propagation while reducing computation and memory.
- 4.1 Toy problem: In the toy problem, the approximation error decreases according to the theoretical bound as reverse depth K increases.Figure 1 visualizes Proposition 3.1 using the approximation error and its theoretical bound.
- 4.1 Toy problem: For f, hT −1 is a descent direction and the ratio h⊤T −Kdλf/∥dλf∥2 stays well away from zero along the optimization trajectory.The original objective satisfies the conditions guaranteeing exact convergence, although the finite lower-level horizon offsets the optimal λ from the origin.
- 4.1 Toy problem: For ef, truncated back-propagation can converge to a suboptimal point, while increasing K rapidly decreases the bias.This biased-convergence behavior occurs after adding a direct λ-dependent term so that ∇λef ≠ 0.
- 4.2.1 Data hypercleaning: In hypercleaning, 1-RMD performs somewhat worse, whereas K = 5 already maximizes F1 score and achieves results close to full back-propagation at about twice the speed.The experiment tunes 5000 hyperparameters for a classifier with 7850 weights and evaluates validation loss, generalization, and corrupted-sample identification.
- 4.2.2 Task interaction: In task interaction, increasing K generally improves test accuracy, but the gaps between truncated methods and full RMD are small.For CIFAR100, increasing K rapidly reduces the stationary-point bias; for CIFAR10, K > 1 eventually yields nearly identical truncated and full gradient directions.
- 4.3 Meta-learning: One-shot classification: All truncated methods are faster than full RMD, sometimes by five times, while regularization produces exponential decay of gradient error as K increases.Without regularization, the observed exponential decay is absent; cosine similarities remain positive, indicating descent directions.
5 CONCLUSION
The conclusion presents K-RMD as a first-order heuristic for bilevel problems with iteratively approximated lower-level optimization. Its theoretical and empirical results support truncated back-propagation as a practical alternative to full RMD, while noting current experimental scale limits.
- K-RMD is analyzed as a first-order heuristic for bilevel optimization with an iteratively approximated lower-level problem.
- The analysis establishes sufficient conditions for convergence to an approximate or exact stationary point of the upper-level objective.
- Gradient approximation error decays exponentially with reverse depth when the lower-level iterate is near a strict local minimum.
- K-RMD achieves application-specific performance close to full RMD and is roughly twice as fast.
- Experiments use modest parameter, hyperparameter, and horizon sizes because both K-RMD and full RMD must be computed for comparison.
A Proof of Proposition 3.1
The proposition bounds truncated back-propagation under smoothness and local strong convexity assumptions. It contrasts contraction near a strongly convex region with a looser worst-case bound for nonconvex lower-level objectives.
- Under β-smoothness, twice differentiability, and local α-strong convexity around the final iterates, the proposition analyzes the gradient approximation.
- The lower-level update is modeled as Ξt+1(wt, λ) = wt −γ∇wg(wt, λ), with a step-size restriction on γ.
- When the lower-level objective is globally α-strongly convex, the resulting bounds simplify relative to the locally convex case.
- For a β-smooth nonconvex objective, the worst case gives ∥At∥ = 1 + γβ ≤ 2, producing the stated bound.
B Proof of Lemma 3.2
The proof establishes sufficient descent for the truncated gradient under global strong convexity and regularity assumptions. It decomposes the relevant inner product and controls its error terms using contraction and linear convergence of gradient descent.
- Under global strong convexity, smoothness, continuous second derivatives, and full column rank of Bt, hT−K is a sufficient descent direction for large T and small γ.
- The proof first expands the K = 1 case and notes that analogous steps apply for K > 1.
- The error analysis assumes Lipschitz continuity of At and Bt and uses linear convergence of gradient descent on strongly convex, smooth functions.
- The inner-product bound combines a contraction term, (1 − γα)T−t∥BT∇ŵ∗f∥2, with error terms Δ1, Δ2, and Δ3.
- The final result follows by combining bounds on Δ1, Δ2, and Δ3, with full column rank ensuring the required nonsingularity.
C Proof of Theorem 3.3
The theorem treats truncated back-propagation as a biased stochastic first-order oracle for updating hyperparameters. With bounded gradient-estimation error and decaying step sizes, learning converges to an approximate stationary point.
- If ∥hT−K − dλf∥ ≤ ϵ, using hT−K as a stochastic first-order oracle with ητ = O(1/√τ) yields convergence guarantees.
- Under Proposition 3.1’s assumptions, learning with hT−K converges to an ϵ-approximate stationary point.
- The proof uses standard nonconvex optimization analysis for biased gradient estimates.
- The argument assumes F is L-smooth and bounds both the exact and truncated gradients by G almost surely.
- The gradient-estimation error is defined as eτ = dλf(τ) − hT−K,(τ) and enters the recursive update analysis.
D Proof of Theorem 3.4
The proof establishes exact-stationary-point convergence under structural assumptions by decomposing the truncated gradient and verifying sufficient descent conditions.
- Theorem 3.4 assumes second-order continuous differentiability, full column rank of B_t near w_T, a non-interference condition, and a deterministic problem.
- For all K ≥1, sufficiently large T, and sufficiently small γ, the limit point is an exact stationary point.
- The proof decomposes the full gradient into four components to relate the truncated back-propagation direction to the exact hypergradient.
- The argument verifies the non-interference condition together with an exponentially decaying error bound for the truncated gradient.
- The proof applies a lemma to the resulting update sequence after establishing the required gradient and error conditions.
E Proof of Theorem 3.5
The proof constructs a deterministic scalar counterexample showing that truncated back-propagation can fail to converge when the non-interference assumption is violated.
- Theorem 3.5 states that a problem satisfying all but assumption 3 of Theorem 3.4 can prevent convergence to a stationary point.
- The counterexample uses a scalar deterministic bilevel problem with a perturbation function φ and a finite-step lower-level gradient-descent solution.
- The construction satisfies smoothness, strong convexity, differentiability, full-rank Jacobian, and deterministic-upper-objective assumptions while violating non-interference.
- A suitable φ creates stationary points where the approximate 1-RMD gradient remains nonzero.
- Because the truncated update cannot vanish near any stationary point, no iterate sequence can converge to one.
- The contradiction follows because convergence would require arbitrarily small successive steps, whereas the update norm remains bounded below.
F Proof of Proposition 3.6
The proposition analyzes the limiting matrices induced by lower-level gradient descent and uses a contraction lemma to establish their limiting relationship.
- Proposition 3.6 considers lower-level iterates converging to a stationary point and defines the limiting matrices A∞ and B∞.
- The proof evaluates the relevant derivatives at the limiting lower-level solution and passes to the limit.
- The equality in the proposition is established using Lemma F.1 for matrices whose norm is less than 1.
- For sufficiently small γ, the Hessian bounds imply ∥I −γ∇w,wg∥< 1, enabling the contraction argument.
- The appendix reports vanilla gradient descent for the lower objective and denotes the result after T steps by ˆw∗.
G.2 Task interaction
The task-interaction experiment optimizes a structured covariance-like hyperparameterization in a one-shot learning model, with practical limits on data and training duration.
- The experiment uses T = 100 lower-level gradient-descent iterations with learning rate 0.1 to obtain ˆw∗.
- The matrix C and scalar ρ are reparameterized with softplus transformations so the optimized hyperparameters are λ = {B, ν}.
- The one-shot architecture freezes the first two convolutional layers as a hyper-representation and fine-tunes the last three layers per task.
- Testing uses 10 independent trials with balanced training and validation sets, and early stopping after 500 hyper-iterations without testing-error improvement.
- The model contains approximately 110k hyperparameters and 75k parameters.
- All algorithms are stopped after 5000 hyper-iterations, while data augmentation, larger meta-batches, and longer training could further improve results.