Algorithm unrolling is ubiquitous in machine learning, particularly in hyperparameter optimization and meta-learning, where Jacobians of solution mappings are computed by differentiating through iterative algorithms. Although unrolling is known to yield asymptotically correct Jacobians under suitable conditions, recent work has shown that the derivative iterates may initially diverge from the true Jacobian, a phenomenon known as the curse of unrolling. In this work, we provide a non-asymptotic analysis that explains the origin of this behavior and identifies the algorithmic factors that govern it. We show that truncating early iterations of the derivative computation mitigates the curse while simultaneously reducing memory requirements. Finally, we demonstrate that warm-starting in bilevel optimization naturally induces an implicit form of truncation, providing a practical remedy. Our theoretical findings are supported by numerical experiments on representative examples.
Figures & tables
Figure 1: Iterate x(k)(u) vs derivative Dx(k)(u) error plot for gradient descent applied to f(x,u):=∥Ax−b∥2/2+u∥x∥2/2 . Unlike x(k)(u) , Dx(k)(u) initially drifts away from its limit before eventually coming back to it. We provide a non-asymptotic understanding of this transient behavior, called the curse of unrolling, and study simple ways to mitigate it.
Iterations
Iterates
Propagation
Unrolled Algorithm
(x(k))k=0K
0→K
Forward Mode
(X˙(k))k=0K
0→K
Reverse Mode
(UˉK(k))k=K0
K→0
Table 1: Propagation of algorithm and derivative iterates.
Figure 2: Error evolution of e(k)(u) , e˙(k)(u) , and eˉ(k)(u) generated by gradient descent applied to f(x,u):=∥Ax−b∥2/2+u∥x∥2/2 . The dashed lines denote the bounds given in ( 8 ) and ( 19 ). The vertical lines denote k˙ and kˉ defined in Sections 3.1 and 3.2 respectively.
Iterations
Iterates
Propagation
Idle Algorithm
(x(k))k=0T′
0→T′
Unrolled Algorithm
(x(k))k=T′K′
T′→K′
Forward Mode
(X˙(k))k=T′K′
T′→K′
Reverse Mode
(UˉK(k))k=K′T′
K′→T′
Table 2: Propagation of the iterates under truncation.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Late-start / truncation behavior for N=2 , α=2/(L+m) , and ρ≈0.763526 .
Figure 4: Late-start / truncation behavior for N=5 , α=2/(L+m) , and ρ≈0.921979 .
Figure 5: Late-start / truncation behavior for N=10 , α=2/(L+m) , and ρ≈0.974362 .
Figure 6: Late-start / truncation behavior for N=20 , α=2/(L+m) , and ρ≈0.993876 .
Figure 7: Late-start / truncation behavior for N=30 , α=2/(L+m) , and ρ≈0.998509 .
Figure 8: Late-start / truncation behavior for N=40 , α=2/(L+m) , and ρ≈0.999722 .
Figure 9: Late-start / truncation behavior for N=2 , α=1/(3L) , and ρ≈0.763526 .
Figure 10: Late-start / truncation behavior for N=5 , α=1/(3L) , and ρ≈0.921979 .
Figure 11: Late-start / truncation behavior for N=10 , α=1/(3L) , and ρ≈0.974362 .
Figure 12: Late-start / truncation behavior for N=20 , α=1/(3L) , and ρ≈0.993876 .
Figure 13: Late-start / truncation behavior for N=30 , α=1/(3L) , and ρ≈0.998509 .
Figure 14: Late-start / truncation behavior for N=40 , α=1/(3L) , and ρ≈0.999722 .
Bilevel optimization (BO) is widely applicable to many machine learning problems. Scaling BO, however, requires repeatedly computing hypergradients, which involves solving inverse Hessian-vector products (IHVPs). In practice, these operations are often approximated using crude surrogates such as one-step gradient unrolling or identity/short Neumann expansions, which discard curvature information. We build on implicit function theorem-based algorithms and propose to incorporate Kronecker-factored approximate curvature (KFAC), yielding curvature-aware hypergradients with a better performance efficiency trade-off than Conjugate Gradient (CG) or Neumann methods and consistently outperforming unrolling. We evaluate this approach across diverse tasks, including meta-learning and AI safety problems. On models up to BERT, we show that curvature information is valuable at scale, and KFAC can provide it with only modest memory and runtime overhead. Our implementation is available at https://github.com/liaodisen/NeuralBo.
Deep unfolding (DU) accelerates iterative optimizers by introducing learnable components and training them through unrolled iterations, but extending DU to the large-scale semidefinite programs (SDPs) common in robotics has remained limited. Unrolling a full-update conic solver such as COSMO exposes two obstacles that prior work on learned conic solvers has not: backpropagating through the per-iteration linear-system solve incurs memory quadratic in the problem size once the coefficient matrix is formed explicitly, and backpropagating through the positive semidefinite (PSD) cone projection becomes numerically unstable when eigenvalues coincide. We address the first obstacle with a matrix-free implicit differentiation rule that operates entirely through matrix-vector products, reducing memory from O(n2) to O(n) and enabling backpropagation at scales where direct factorization runs out of memory. We address the second with a backward rule based on the Dalečkii--Krein representation of the Fréchet derivative, which remains well-defined under repeated eigenvalues. Together these make it possible to learn lightweight hyperparameter policies and warm-starts for a full-update conic solver. We evaluate on nonlinear covariance steering problems solved via sequential convex programming (SCP), as well as standalone SDPs and second-order cone programs ranging from max-cut and Lovász ϑ SDPs to robust estimation and control problems. The learned policies outperform state-of-the-art solvers across all problems, and can provide up to a 50× speedup depending on the class. When used as a subroutine in SCP, the learned approach delivers over a 30× speedup compared to COSMO.
Alex Oshin, Rahul Vodeb Ghosh, Evangelos A. Theodorou
Bilevel optimization underpins many machine learning applications, including hyperparameter optimization, meta-learning, neural architecture search, and reinforcement learning. While hypergradient-based methods have advanced significantly, a gap persists between theoretical guarantees and practical single-loop implementations required for efficiency. We bridge this gap by establishing sharper convergence results for single-loop approximate implicit differentiation (AID) and iterative differentiation (ITD) methods, leveraging our proposed analytical framework, decoupled norm analysis (DNA). For AID, we improve the convergence rate from O(κ6/K) to O(κ5/K), where κ is the condition number of the inner-level problem. For ITD, we prove that the asymptotic error is O(κ2), exactly matching the known lower bound and improving upon the previous O(κ3) guarantee. Numerical experiments on synthetic and real tasks corroborate our theoretical findings.
Yubo Zhou, Jun Shu, Luo Luo +4
School of Mathematics and Statistics, Xi’an Jiaotong University · School of Data Science, Fudan University · 3SGIT AI Lab, State Grid Corporation of China +1