SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {x x} ^ {2} f (x, \theta) J _ {x} F (x, \theta) = \nabla^ {3} f (x, \theta) [ x, \cdot , \cdot ] - \nabla^ {3} f (x, \theta) [ F (x, \theta), \cdot , \cdot ]
$$
so that  $J_{x}F(\bar{x} (\theta),\theta) = 0$  and Lipschitz continuity of the derivatives of  $f$  implies Lipschitz continuity of  $J_{x}F(x,\theta)$  using the fact that  $\nabla_{xx}^2 f(x,\theta)\succeq \mu I$  and that matrix inversion is smooth and Lipschitz for such matrices. The quadratic convergence of the three derivative estimates for Newton's method is illustrated on a logistic regression example in Figure 2.
3 Hypergradient descent for bilevel problems
Consider the following bilevel optimization problem
$$
\min _ {\theta} g (x (\theta)) \mathrm {s . t .} x (\theta) \in \arg \min _ {y} f (y, \theta),
$$
where  $g$  and  $f$  are  $C^1$  functions. We will consider bilevel problems such that the inner minimum is uniquely attained and can be described as a fixed point equation  $x = F(x, \theta)$  where  $F$  is as in Assumption 1. The problem may then be rewritten as
$$
\min  _ {\theta} \quad g (x (\theta)) \quad \text {s . t .} \quad x (\theta) = F (x (\theta), \theta), \tag {9}
$$
Method	Time	Memory	Error sources
Piggyback recursion	kn(ωCF+n m)	n(n+m)	suboptimality + burn-in
AD forward-mode	kωCFm	n+m	suboptimality + burn-in
AD reverse-mode	kωCF	kn	suboptimality + burn-in
Implicit differentiation	ωCFn+n3	n	suboptimality
One-step differentiation	ωCF	n	suboptimality + lack of contractivity
Forward algorithm	kCF	n	suboptimality
Table 2: Time and memory complexities of the estimators in Section 2.1 (up to multiplicative constant).  $F$  has time complexity denoted by  $C_F$  and we consider  $k$  iterations in  $\mathbb{R}^n$  with  $m$  parameter.  $\omega$  is the multiplicative overhead of evaluating a gradient (cheap gradient principle).
see illustrations in Remark 1. Gradient descent (or hyper-gradient) on (9) using our one-step estimator in (4) consists in the following recursion
$$
\theta_ {l + 1} = \theta_ {l} - \alpha J ^ {\mathrm {O S}} x _ {k} \left(\theta_ {l}\right) ^ {T} \nabla g \left(x _ {k} \left(\theta_ {l}\right)\right), \tag {10}
$$
where  $\alpha > 0$  is a step size parameter. Note that the quantity  $J^{\mathrm{OS}}x_k(\theta)^T\nabla g(x_k)$  is exactly what is obtained by applying backpropagation to the composition of  $g$  and Algorithm 3, without any further custom variation on backpropagation. Note that one in practice may use amortized algorithms, such as [18]. This section is dedicated to the theoretical guarantees which can be obtained using such a procedure, proofs are postponed to Appendix A.
3.1 Complexity analysis of different hypergradient strategies
We essentially follow the complexity considerations in [27, Section 4.6]. Let  $C_F$  denote the computation time cost of evaluating the fixed-point map  $F$  and  $\omega > 0$  be the multiplicative overhead of gradient evaluation, in typical applications,  $\omega \leq 5$  (cheap gradient principle [6]). The time cost of evaluating the Jacobian of  $F$  is  $n\omega C_F$  ( $n$  gradients). Forward algorithm evaluation (i.e.,  $x_k$ ) has computational time cost  $kC_F$  with a fixed memory complexity  $n$ . Vanilla piggyback recursion (2) requires  $k - 1$  full Jacobians and matrix multiplications of costs  $n^2 m$ . The forward-mode of AD has time complexity  $k\omega C_F m$  (compute  $m$  partial derivatives each of them cost  $\omega$  times the time to evaluate the forward algorithm), and requires to store the iterate vector of size  $n$  and  $m$  partial derivatives. The reverse-mode of AD has time complexity  $k\omega C_F$  (cheap gradient principle on  $F_\theta^k$ ) and requires to store  $k$  vectors of size  $n$ . Implicit differentiation requires one full Jacobian ( $\omega C_F n$ ) and solution of one linear system of size  $n \times n$ , that is roughly  $n^3$ . Finally, one-step differentiation is given by only differentiating a single step of the algorithm at cost  $\omega C_F$ . For each estimate, distance to the solution will result in derivative errors. In addition, automatic differentiation based estimates may suffer from the burn-in effect [41] while one-step differentiation will suffer from a lack of contractivity as in Corollary 1. We summarize the discussion in Table 2. Let us remark that, if  $C_F \geq n^2$ , then reverse AD has a computational advantage if  $k \leq n$ , which makes sense for fast converging algorithms, but in this case, one-step differentiation has a small error and a computational advantage compared to reverse AD.
Remark 3 (Implicit differentiation: but on which equation?) Table 2 is informative yet formal. In practical scenarios, implicit differentiation should be performed using the simplest equation available, not necessarily  $F = 0$ . This can significantly affect the computational time required. For instance, when using Newton's method  $F = -[\nabla^2 f]^{-1}\nabla f$ , implicit differentiation should be applied to the gradient mapping  $\nabla f = 0$ , not  $F$ . In typical application  $C_{\nabla f} = O(n^{2})$ , and the dominant cost of implicit differentiation is  $O(n^{3})$ , which is of the same order as the one-step differentiation as  $C_F = O(n^3)$  (a linear system needs to be inverted). However, if the implicit step was performed on  $F$  instead of  $\nabla f$ , it would incur a prohibitive cost of  $O(n^{4})$ . In conclusion, the implicit differentiation phase is not only custom in terms of the implementation, but also in the very choice of the equation.
3.2 Approximation analysis of one step differentiation
The following corollary, close to [42, Prop. 3.1], provides a bound on the one-step differentiation (Algorithm 3) gradient estimator for (9). The bound depends on  $\rho$ , the contraction factor, and distance to the solution for the inner algorithm. The proof is given in Appendix A.
Corollary 4 Let  $F \colon \mathbb{R}^n \times \mathbb{R}^m \to \mathbb{R}^n$  be as in Corollary 1 and consider the bilevel problem (9), where  $g$  is a  $C^1$ ,  $l_g$  Lipschitz function with  $l_{\overline{\nu}}$  Lipschitz gradient. Then,
$$
\left| \left| \nabla_ {\theta} (g \circ \bar {x}) (\theta) - J ^ {\mathrm {O S}} x _ {k} (\theta) ^ {T} \nabla g (x _ {k}) \right| \right| \leq \frac {\rho L _ {F} l _ {g}}{1 - \rho} + L _ {J} l _ {g} \| x _ {k - 1} - \bar {x} (\theta) \| + L _ {F} l _ {\nabla} \| \bar {x} (\theta) - x _ {k} \|.
$$
3.3 Approximate critical points
The following lemma is known, but we provide a proof for completeness in Appendix A.
Lemma 3 Assume that  $h \colon \mathbb{R}^m \to \mathbb{R}$  is  $C^1$  with  $L$  Lipschitz gradient and lower bounded by  $h^*$ . Assume that for some  $\epsilon > 0$ , for all  $l \in \mathbb{N}$ ,  $\left\| \theta_{l+1} - \theta_l + \frac{1}{L} \nabla h(\theta_l) \right\| \leq \frac{\epsilon}{L}$ . Then, for all  $K \in \mathbb{N}$ ,  $K \geq 1$ , we have
$$
\min _ {l = 0, \ldots , K} \| \nabla h (\theta_ {l}) \| ^ {2} \leq \epsilon^ {2} + \frac {2 L (h (x _ {0}) - h ^ {*})}{K + 1}.
$$
Combining with Corollary 4, it provides a complexity estimate for the recursion in (10).
Corollary 5 Under the setting of Corollary 4, consider iterates in (10) with  $k \geq 2$  and assume the following
-  $\sup_{\theta}\| x_0(\theta) - F_\theta (x_0(\theta))\| \leq M$  , for some  $M > 0$
-  $F$  is  $L_{F}$  Lipschitz and  $J_{(x,\theta)}F$  is  $L_{J}$  Lipschitz jointly in operator norm.
-  $g$  is  $C^1$ ,  $l_g$  Lipschitz with  $l_{\nabla}$  Lipschitz gradient.
$\frac{1}{\alpha} \geq \left(\frac{L_J}{1 - \rho}\left(\frac{L_F}{1 - \rho} + 1\right)l_g + l_{\nabla}\frac{L_F}{1 - \rho}\right)\frac{L_F}{1 - \rho}$
$g\circ \bar{x}$  is lower bounded by  $g^{*}$
Then setting  $\epsilon = \frac{\rho}{1 - \rho} (L_F l_g + (L_J l_g + L_F l_\nabla)M\rho^{k - 2})$ , for all  $K\in \mathbb{N}$
$$
\min  _ {l = 0, \dots , K} \| \nabla_ {\theta} (g \circ \bar {x}) (\theta_ {l}) \| ^ {2} \leq \epsilon^ {2} + \frac {2 L ((g \circ \bar {x}) (\theta_ {0}) - g ^ {*})}{K + 1}.
$$
The level of approximate criticality is the sum of a term proportional to  $\rho$  and a term inversely proportional to  $K$ . For large values of  $k$  (many steps on the inner problem),  $K$  (many steps on the outer problem), approximate criticality is essentially proportional to  $\rho L_F l_g / (1 - \rho)$  which is small if  $\rho$  is close to 0 (e.g. superlinear algorithms).
We illustrate our findings on three different problems. First, we consider Newton's method applied to regularized logistic regression, as well as interior point solver for quadratic problems. These are two fast converging algorithms for which the results of Section 2.3 can be applied and the one-step procedure provides accurate estimations of the derivative with a computational overhead negligible with respect to solution evaluation, as for implicit differentiation. We then consider the gradient descent algorithm applied to a ridge regression problem to illustrate the behavior of the one step procedure for linearly convergent algorithms.
Logistic regression using Newton's algorithm. Let  $A \in \mathbb{R}^{N \times n}$  be a design matrix, the first column being made of 1s to model an intercept. Rows are denoted by  $(a_1, \ldots, a_N)$ . Let  $x \in \mathbb{R}^n$  and  $y \in \{-1, 1\}^N$ . We consider the regularized logistic regression problem
$$
\min  _ {x \in \mathbb {R} ^ {n}} \sum_ {i = 1} ^ {N} \theta_ {i} \ell \left(\left\langle a _ {i}, x \right\rangle y _ {i}\right) + \lambda \| x _ {- 1} \| ^ {2}, \tag {11}
$$
Figure 3: Timing experiment to evaluate one gradient. Left: Differentiable QP in (12), one step and implicit estimators agree up to an error of order  $10^{-16}$ . Right: Newton's method on logistic regression (11), one step and implicit estimators agree up to an error of order  $10^{-12}$ . Label "None" represent solving time and "Autodiff", "Implicit" and "One step" represent solving time and gradient evaluation for each estimator in Section 2.1. The mean is depicted using shading to indicate standard deviation estimated over 10 runs.
where  $\ell$  is the logistic loss,  $\ell \colon t\mapsto \log (1 + \exp (-t))$ ,  $\lambda >0$  is a regularization parameter, and  $x_{-1}$  denotes the vector made of entries of  $x$  except the first coordinate (we do not penalize intercept). This problem can be solved using Newton's method which we implement in jax using backtracking line search (Wolfe condition). Gradient and Hessian are evaluated using jax automatic differentiation, and the matrix inversion operations are performed with an explicit call to a linear system solver.
We denote by  $x(\theta)$  the solution to problem (11) and try to evaluate the gradient of  $\theta \mapsto \|x(\theta)\|^2/2$  using the three algorithms presented in Section 2.1. We simulate data with Gaussian class conditional distributions for different values of  $N$  and  $n$ . The results are presented in Figure 3 where we represent the time required by algorithms as a function of the number of parameters required to specify problem (11), in our case size of  $A$  and size of  $y$ , which is  $(n + 1)N$ .
Figure 3 illustrates that both one-step and implicit differentiation enjoy a marginal computational overhead, contrary to algorithmic differentiation. In this experiment, the one-step estimator actually has a slight advantage in terms of computation time compared to implicit differentiation.
Interior point solver for quadratic programming: The content of this section is motivated by elements described in [4], which is associated with a pytorch library implementing a standard interior point solver. Consider the following quadratic program (QP):
$$
\min  _ {x \in \mathbb {R} ^ {n}} \quad \frac {1}{2} x ^ {T} Q x + q ^ {T} x \quad \text {s . t .} \quad A x = \theta , \quad G x \leq h, \tag {12}
$$
where  $Q \in \mathbb{R}^{n \times n}$  is positive definite,  $A \in \mathbb{R}^{m \times n}$  and  $G \in \mathbb{R}^{p \times n}$  are matrices,  $q \in \mathbb{R}^n$ ,  $\theta \in \mathbb{R}^m$  and  $h \in \mathbb{R}^p$  are vectors. We consider  $x(\theta)$  the solution of problem (12) as a function of  $\theta$ , the right-hand side of the equality constraint. We implemented in jax a standard primal-dual Interior Point solver for problem (12). Following [4], we use the implementation described in [43], and we solve linear systems with explicit calls to a dedicated solver. For generic inputs, this algorithm converges very fast, which we observed empirically. Differentiable programming capacities of jax can readily be used to implement the automatic differentiation and one-step derivative estimators without requiring custom interfaces as in [4]. Indeed, implicit differentiation for problem (12) was proposed in [4] with an efficient pytorch implementation. We implemented these elements in jax in order to evaluate  $J_{\theta}x(\theta)$  using implicit differentiation. More details on this experiment are given in Appendix B.
We consider evaluating the gradient of the function  $\theta \mapsto \| x(\theta)\|^2 / 2$  using the three algorithms proposed in Section 2.1. We generate random instances of QP in (12) of various sizes. The number of parameters needed to describe each instance is  $n(n + 1) + (n + 1)m + (n + 1)p$ . The results are presented in Figure 3 where we represent the time required by algorithms as a function of the number of parameters required to specify problem (12). In all our experiments, the implicit and one-step estimates agree up to order  $10^{-6}$ . From Figure 3, we
Figure 4: Differentiation of gradient descent for solving weighted Ridge regression on cpusmall. Top line: condition number of 1000. Bottom line: condition number of 100. Left column: small learning rate  $\frac{1}{L}$ . Right column: big learning rate  $\frac{2}{\mu + L}$ . Dotted (resp. filled) lines represent the error of the iterates (resp. of the Jacobians).
see that both one-step and implicit differentiation enjoy a marginal additional computational overhead, contrary to algorithmic differentiation.
Weighted ridge using gradient descent. We consider a weighted ridge problem with  $A \in \mathbb{R}^{N \times n}$ ,  $y \in \mathbb{R}^N$ ,  $\lambda > 0$  and a vector of weights  $\theta \in \mathbb{R}^N$ :
$$
\bar {x} (\theta) = \arg \min _ {x \in \mathbb {R} ^ {n}} f _ {\theta} (x) = \frac {1}{2} \sum_ {i = 1} ^ {N} \theta_ {i} (y _ {i} - \langle a _ {i}, x \rangle) ^ {2} + \frac {\lambda}{2} \| x \| ^ {2}.
$$
We solve this problem using gradient descent with adequate step-size  $F(x, \theta) = x - \alpha \nabla f_{\theta}(x)$  with  $x_0(\theta) = 0$ , and we consider the  $K$ -step truncated Jacobian propagation  $\tilde{F} = F_{\theta}^{K}$  with  $K = 1 / \kappa$  where  $\kappa$  is the effective condition number of the Hessian. Figure 4 benchmarks the use of automatic differentiation, one-step differentiation, and implicit differentiation on the data set cpusmall provided by LibSVM [17] for two types of step-sizes. We monitor both quantities  $\| x_k(\theta) - \bar{x}(\theta)\|^2$  for the iterates, and  $\| J_\theta x_k(\theta) - J_\theta \bar{x}(\theta)\|^2$  for the Jacobian matrices. As expected, implicit differentiation is faster and more precise, it is our gold standard which requires custom implicit system to be implemented (or the use of an external library). For large steps, autodiff suffers from the burn-in phenomenon described in [41], which does not impact the one-step estimator. Therefore, for a fixed time budget, the one-step strategies allow to obtain higher iterate accuracy and similar, or better, Jacobian accuracy. In the small step regime, one step differentiation provides a trade-off, for a fixed time budget, one obtains better estimates for the iterate and worse estimates for the Jacobian matrices. Our results suggest that  $K$ -step truncated backpropagation allows to save computation time, at the cost of possibly degraded derivatives compared to full backpropagation, in line with [42].
We studied the one-step differentiation, also known as Jacobian-free backpropagation, of a generic iterative algorithm, and provided convergence guarantees depending on the initial rate of the algorithm. In particular, we show that one-step differentiation of a quadratically convergent algorithm, such as Newton's method, leads to a quadratic estimation of the Jacobian. A future direction of research would be to understand how to extend our findings to the nonsmooth world as in [14] for linearly convergent algorithms.
The authors acknowledge the support of the AI Interdisciplinary Institute ANITI funding, through the French "Investments for the Future - PIA3" program under the grant agreement ANR-19-PI3A0004, Air Force Office of Scientific Research, Air Force Material Command, USAF, under grant numbers FA8655-22-1-7012, ANR MaSDOL 19-CE23-0017-0, ANR Chess, grant ANR-17-EURE-0010, ANR Regulia, ANR GraVa ANR-18-CE40-0005. Jérôme Bolte, Centre Lagrange, and TSE-P.
[1] Martin Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. Software available from tensorflow.org.
[2] Pierre Ablin, Gabriel Peyré, and Thomas Moreau. Super-efficiency of automatic differentiation for functions defined as a minimum. In ICML, volume 119, pages 32-41, 13-18 Jul 2020.
[3] Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019.
[4] Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as a layer in neural networks. In International Conference on Machine Learning, pages 136-145. PMLR, 2017.
[5] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32, 2019.
[6] Walter Baur and Volker Strassen. The complexity of partial derivatives. Theoretical computer science, 22(3):317-330, 1983.
[7] Amir Beck. First-order methods in optimization. SIAM, 2017.
[8] Thomas Beck. Automatic differentiation of iterative processes. Journal of Computational and Applied Mathematics, 50(1-3):109-118, 1994.
[9] Yoshua Bengio. Gradient-Based Optimization of Hyperparameters. Neural Computation, 12(8):1889-1900, 2000.
[10] Q. Bertrand, Q. Klopfenstein, M. Blondel, S. Vaiter, A. Gramfort, and J. Salmon. Implicit differentiation of lasso-type models for hyperparameter optimization. In ICML, 2020.
[11] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint arXiv:2105.15183, 2021.
[12] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. Advances in Neural Information Processing Systems, 35:5230-5242, 2022.
[13] Jérôme Bolte, Tam Le, Edouard Pauwels, and Tony Silveti-Falls. Nonsmooth implicit differentiation for machine-learning and optimization. In NeurIPS, volume 34, pages 13537-13549, 2021.
[14] Jerome Bolte, Edouard Pauwels, and Samuel Vaiter. Automatic differentiation of nonsmooth iterative algorithms. In Advances in Neural Information Processing Systems, 2022.
[15] Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
[16] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[17] Chih-Chung Chang and Chih-Jen Lin. Libsvm: A library for support vector machines. ACM transactions on intelligent systems and technology (TIST), 2(3):1-27, 2011.
[18] Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. In NeurIPS, 2022.
[19] Matthias J Ehrhardt and Lindon Roberts. Analyzing inexact hypergradients for bilevel learning. arXiv preprint arXiv:2301.04764, 2023.
[20] Laurent El Ghaoui, Fangda Gu, Bertrand Travacca, Armin Askari, and Alicia Tsai. Implicit deep learning. SIAM Journal on Mathematics of Data Science, 3(3):930-958, 2021.
[21] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In ICML, volume 70, pages 1126–1135, 2017.
[22] Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin. Jfb: Jacobian-free backpropagation for implicit models. Proceedings of the AAAI Conference on Artificial Intelligence, 2022.
[23] Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is attention better than matrix decomposition? In ICLR, 2021.
[24] Zhengyang Geng, Xin-Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. Advances in Neural Information Processing Systems, 34:24247-24260, 2021.
[25] Jean Charles Gilbert. Automatic differentiation and iterative processes. Optimization methods and software, 1(1):13-21, 1992.
[26] Andreas Griewank and Christele Faure. Piggyback differentiation and optimization. In Large-scale PDE-constrained optimization, pages 148-164. Springer, 2003.
[27] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
[28] Fangda Gu, Heng Chang, Wenwu Zhu, Somayeh Sojoudi, and Laurent El Ghaoui. Implicit graph neural networks. Advances in Neural Information Processing Systems, 33:11984-11995, 2020.
[29] Howard Heaton, Daniel McKenzie, Qiuwei Li, Samy Wu Fung, Stanley Osher, and Wotao Yin. Learn to predict equilibria via fixed point networks. arXiv preprint arXiv:2106.00906, 2021.
[30] Seppo Linnainmaa. The representation of the cumulative rounding error of an algorithm as a Taylor expansion of the local rounding errors. Master's thesis, Univ. Helsinki, 1970.
[31] Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In ICLR, 2019.
[32] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In AISTATS, pages 1540-1552, 2020.