One-step differentiation of iterative algorithms
Authors: Jerome Bolte
Edouard Pauwels
Samuel Vaiter
Venue: NeurIPS 2023 spotlight (2023)
Keywords: automatic differentiation
implicit differentiation
super-linear algorithms
bilevel optimization.
TL;DR: One-step differentiation is as easy as autodiff, and as performant as implicit differentiation.

Abstract: In appropriate frameworks, automatic differentiation is transparent to the user, at the cost of being a significant computational burden when the number of operations is large. For iterative algorithms, implicit differentiation alleviates this issue but requires custom implementation of Jacobian evaluation. In this paper, we study one-step differentiation, also known as Jacobian-free backpropagation, a method as easy as automatic differentiation and as performant as implicit differentiation for fast algorithms (e.g. superlinear optimization methods). We provide a complete theoretical approximation analysis with specific examples (Newton's method, gradient descent) along with its consequences in bilevel optimization. Several numerical examples illustrate the well-foundness of the one-step estimator.

One-step differentiation of iterative algorithms
Jérôme Bolte Toulouse School of Economics, Université Toulouse Capitole, Toulouse, France.
Edouard Pauwels Toulouse School of Economics (IUF), Toulouse, France.
Samuel Vaiter CNRS & Université Côte d'Azur, Laboraloire J. A. Dieudonné. Nice, France.
Abstract
In appropriate frameworks, automatic differentiation is transparent to the user at the cost of being a significant computational burden when the number of operations large. For iterative algorithms,, implicit, differentiation alle- viates this issue but requires custom implementation of Jacobian evaluation. this paper, we study one-step differentiation, also known as Jacobian- free backpropagation, a method easy as automatic differentiation and as efficient as implicit differentiation for fast algorithms (e.g., superlinear optimization methods). We provide complete theoretical approximation analysis with specific examples (Newton's 's method, gradient descent) along with its consequences in bilevel optimization. Several numerical examples illustrate the fell-founeedness of the one-step estimator.
1 Introduction
Differentiating the solution of machine learning problem is an important task, e.g., in hyperparameters optimization [9], in neural architecture search [31] and when using convex layers [3]. There are two main ways to achieve this goal: automatic differentiation (AD) and implicit differentiation (ID). Automatic differentiation implements the idca of evaluating derivatives through the compositional rules of differential calculus in txer-transparent way. is a mature concept [27] implemented in several machine learning frameworks [36, 16, 1). However, the lime and memory complexity incurred may become prohibilive as soon the computational graph becomes bigger, a typical example being unrolling iterative optimization algorithms such as gradient descent (5). The alternative, implicit differentiation, is not always accessible: it does not solely rely on the compositional rules of differential calculus (Jacobian multiplication) and usually requires solving linear system. The user needs to implement custom rules an automatic differentiation framework (as done, for example, in [4]) or use dedicated libraries such as [11, 3, 10] implementing these rules for given models. Provided that the implementation carefully done, this is most the time the gold standard for the task of differentiating problem solutions.
Contributions. We study a one-step Jacobian approximator based a simple principle: differentiate only the last itcration of an itcrative process and drop previous derivatives. The idca of dropping derivatives or single-step differentiation was explored, for example, in [24, 23, 39, 22, 40, 29] and our main contribution is a general account and approximation analysis for Jacobians of iteralive algorithms. One-step estimation constitutes a rough approximation at first sight and our motivation study is its ease-of-use within automatic differentiation frameworks: custom Jacobian-vector products vector-Jacobian products needed as long as a stop_gradient primitive available (see Table 1).
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

Algorithm 1: Auto-
Algorithm 2: Implicit
Algorithm 3: One-step
matic
Input:
Input: 30 X, k
Input:
Io E X k 0.
Eval:
++ (0) E R k2
Eval:
stop-gradis gradieadtttt...
Eval:
stop_gradient for I,..., do
for = I...., k
do
with bith_gradient
2; F1
for i= 1...... k do
2j = P ((;;-1,0)
with_ aethhgrraient
= F(x;-1(0), Return: (0)
Return: Ik
2& (0) F(
Differentiation:
Return: (x) Differentiation:
Differentiation:
Custom implicit VJP/ JVP.
nalive autodiff on Eval.
native autodiff on Eval
Table 1: Qualitative comparison of differen-
tiation strategies. Native autodiff refers to widespread primitives in differentiable pro- gramming (e.g. grad in JAX). Custom JVP/VJP refers to specialized libraries such as jaxopt [12] or qpth [4] implicit differenti- ation in specific contexts.
A Algorithm
Implementation
Efficient
Automatic Implicit One step
Native autodiff
no
Custom
yes
Native autodiff
yes
We conduct an approximation analysis of one-step Jacobian (Corollary 1). true Jacobian are produced by the distance to the solution (i.e., the quality of completion of the optimization phase) and the lack of contractivity of the iteration mapping. This imprecision has to be balanced with the ease of implementation of the one-step Jacobian. This suggests that one-step differentiation is efficient for small contraction factors, which corresponds to fast algorithms. We indeed show that one-step Jacobian is asymptotically correct for super-linearly convergent algorithms (Corollary 2) and provide similar approximation rate as implicit differentiation for quadratically convergent algorithms (Corollary 3). We exemplify these results with hypergradients in bilevel optimization, conduct a detailed complexity analysis and highlight Corollary 4 the estimation of approximate critical points. Finally, mumerical illustrations are provided to show the practicality of the method on logistic regression using Newton's algorithm, interior point solver for quadratic programming and weighted ridge regression using gradient descent.
The distance to the
Related works. Automatic differentiation [27] was first proposed in the forward mode in [44] and its reverse mode in [30]. The study of the behaviour of differentiating iterative procedure by automatic differentiation was first analyzed in [25] and [8] in the optimization community. It was studiod the machine lcarning community for smooth methods [35, 32, 37), and nonsmooth methods [14]. Implicit differentiation is recent highlight in machine learning. was shown be good way to estimate the Jacobian of problem solutions, for deep equilibrium network [5], optimal transport [33], and also for nousmooth problems [13], such as sparse models [10]. Inexact implicit diflerentiation was explored in [19]. Truncated estimation of the 'Neumann series" for the implicit differentiation is routinely used [32, 34] and truncated backpropagation was investigated in [42] for the gradient descent iterations. In the very specific case of min-min problems, [2] studied the speed of convergence of automatic, implicit, and analytical differentiation.
The closest work to ours is [22] under the name Jacobian-free backpropagation but differs significantly in the following ways. Their focus is on single implicit layer networks, and guarantees are qualitative (descent direction of [22, Theorem 3.1]). In contrast, we provide quantitative results on abstract fixed points and applicable to any architecture. The idea of "dropping derivatives" was proposed in [21] for meta-learning, one-step differentiation was also investigated to train Transformer architectures [23), and to solve bilevel problems with quasi- Newtons methods [39].

---

2 One-step differentiation
2.1 Automatic, implicit and one-step differentiation
Throughout the text F: I x R. - R denoles a recursive algorithmic map from I to I" with parameters. For any 0 E R", we write F6 : F(x, 0) and let F& denote & recursive composition of F9, for E N. The map Fg defines a recursive algorithm as follows
zo(0) R" and Ix+1(0) = (0),
(1)
We denote by Jo the Jacobian matrix with respect to the variable x. The following assumption is sufficient to ensure a degenerate asymptotic behavior.
Assumption 1 (Contraction) Let F: R" R R" be C² 0S p< and X' c R" be nonempty convex closed, such that for any O E R", (o( c X' and Folloo S p.
Remark The main algorithms considered in this paper fall in the scope of smooth optimization. The algorithmic map Fo is associated to smooth parametric optimization problem given by f: I2" x T - IR such that fo: x f(x, 0) is strongly convex, uniformly in 0. Two examples of algorithmic maps are given by gradient descent, Fo(z) x a fo(o,,), or Newton's Fa(x) = x V (x))) ff((z) for positive stop > O. For small stop sizes, gradient descent provides a contraction and Newton's method provides local contraction, both fall in the scope of Assumption 1. Other examples include inertial algorit such as the Heavy Ball method [38], which has to be considered in phase space and for which single iteration is not contracting, but a large number of iteration is (see e.g. [38, 35]).
The following lemma gathers known properties regarding the fixed point of Fo, denoted by x(0) and for which we will be interested in estimating derivatives.
Lemma 1 Under Assumption 1, for each 0 in R there is a unique fized point of Fo in X denoted by æ(f), which is C. function of 9. Furthermore, for all k N, we have (0)x =(0)l S (=1x=(0)) - i(0)|| S p² (Fn-Fa(zo)))
This is well known, we briefly sketch the proof. The mapping F6 is a p contraction on X use the convexity of X and the intermediate value theorem. Banach fixed point theorem ensures existence and uniqueness, differentiability is due to the implicit function thcorem. The convergence rate is classical. We are interested in the mumerical evaluation of the Jacobian Jnz(O), thus well-defined under Assumption 1.
The automatic differentiation estimator JAPzz (0) Jæææ (0) propagates the derivatives (either in a forward or reverse mode) through iterations based on the piggyback recursion [26], for i 1....,
10z;+1(0) = (0,) Jo
(2)
Under assumption 1 we have JAD (0) + (o) and the convergence is asymptotically linear [25, 8, 35, 41, 14). (0) is available differentiable programming framework implementing common primitives such as backpropagation.
xx(() The implicit differentiation estimator ''Dxx (0) is given by application of the implicit function theorem using æy as a surrogate for the fixed point x,
j (0) = (1 - Jg (xx -1)00)))--(aaaaa (eF((x(())
(3)
By continuity of the derivatives of F, we also have )1Dxx1() - (oi(() as k co, (scc e.g. [27, Lemma 15.1) [12, Theorem 1J). Implementing J (0) requires either manual implementation or dedicated techniques or libraries [4, 3, 45, 28, 20, 12] as the matrix inversion operation is not directly expressed using common differentiable programming primitives. A related estimator is the Inexact Automatic Differentiation (IAD) estimator which implements (2) but with Jacobians evaluated at the last iterates Ik, which can be seen as an approximation of JID [35, 19].
The one-step estimator jos (0) is the Jacobian of the fixed point map for the last iteration
x0)) J& (eF(
(4)
3

---

def F(,, theta):
Inplseit dsfferentiation
# here a gradsent step
1_3d a Ouston sop implementation (e.g. jazopt)
alpha grad_r(x, thata)
# Ose-step defferentiation
iterative_ craatvvw_.ooccauuraattttttt x0):
iterative peraativeeprocdureee httpeaaddrrruutt_aeeeeeeeee with_ aith_atoo_graactthee x0):
x0 for in r10g)
x0
for in renge
P(x, theta)
F(stop_gradi gradient(x)
raturn x
stop graaten
F(x, theta)
Autonatic differentiation
return
J_ad jacted (iterative. pewpouuuueeppeeeee (theta, x0)
J_o8 eerrvattteratvveprooelloo (theta, x0)
Figure 1: Implementation of Algorithms 1, and 3 in jax. F is gradient step of some function f. The custom implementation implicit differentiation is not explicitly stated. The function sttp_graaieeeeeee gradient is present in jax. lax and jacfwd computes the full Jacobian using forward-mode AD.
Contrary to automatic differentiation or implicit differentiation estimales, we do not have j0Szz(0) + JBz(O) in general as & - 00, but will see that the error is essentially proportional to p, and thus negligible for fast algorithms for which the estimate is accurate.
From a practical viewpoint, the three estimators JAD, JID and jos are implemented in differentiable programming framework, such as jax, thanks to a primitive stop urp__grrieent. as primi-
illustrated by Algorithms 1, 2 and 3. The computational effect of the stop_gradie live to replace the actual Jacobian (x; 0) by zero for chosen iterations i (1,...,k) Using it for all iterations except the last one, allows one to implement joS in (4) using Algo- rithm 3. This illustrates the main interest of the one-step estimator: it can be implemented using any differentiable programming framework which provides a stop_gradi ent primitive and does not require custom implementation of implicit differentiation. Figure 1 illustrates an implementation in jax for gradient descent.
2.2 Approximation analysis of one step differentiation for linearly convergent algorithms
The following lemma is elementary. It describes the main mathematical mechanism at stake behind our analysis of one-step differentiation.
Lemma 2 Let A E Rax= with |||| S < and B, B E R=xm then
A) B- A(I - A) B+B- B-B
Moreover, have the following estimate,
- B) BB- Bllop
P p
Blllop ||B Bllop-
Proof: First for any E R" we have |1(1 A)oll 2
2 (1- - p|l|u||, which shows
that I - A is invertible (the kernel of I - A is trivial). We also deduce that ||(1 A) lop <
1/(1 - p). Second, we have (T - A)-1 - = A(T - A)-1, since ((T - A) - T) - A) A, and therefore
A(I ((--A) - B A) B) B (I A) 1 B.
The norm bound follows using the submuliiplicatiiity of operator norm, the triangular inequality and the fact that AA S p and ([[-A) |I llo 1/(1 pp.
Corollary 1 Let and X be as in Assumption 1 such that 0 + F(r. 0) is Lp Lipschitz and + JoF 0) is L, Lipschitz (in operator norm) for all E R Then, for all A E R
pLf
1s00 Joo(O) S 1-p L -00(]
Proof The result follows from Lemma 2 with A = J)F()((),0), B JeF()((), and B JoF FoF((x-1, 0) using the fact that || LF and | BBlop - ( 24-1l 0

---

Remark 2 (Comparison with implicit differentiation) In [12] similar bound is de-
scribed for jLD, roughly under the assumption that x + F(r, 0) also has Ly Lipschitz Jacobian, one has
|a x0) -o)0))) S LFLF
LJ | x()]
(1 p)² 12x =(0)]l
(6)
For small and large k the main difference between the two bounds (5) and (6) lies in their first term which is of the same order whenever p and L || ---il are of the same order.
Corollary 1 provides a bound on to p. This means that for fast linearly convergent algorithms, meaning p < 1, one-step differentiation provides good approximation of the actual derivative. Besides, given F which satisfies Assumption 1, with a given p < 1, not specially small, one can set Fe = F for some K E N. In this case, F satisfics assumption with ý = pp and the one-stop estimator in (4) applied to F becomes K-steps estimator on F itself, we only differentiate through the K last steps of the algorithm which amounts to truncated backpropagation [42].
which is asymptotically proport
Example 1 (Gradient. descent) Let f: IR x Ih - IK be such that f(: is p-strongly
convex (p > 0) with I. Lipschitz gradient for all E Rmm then the gradient mapping F: (x,0) + aV..( (z, 0) satisfies Assumption 1 with p max - ay,, aL - 1}, smaller than 1 as long as 0 < a < 2/L. The optimal a = 2/(L + μJ) leads to a contraction factor = 1 2p/(L + p)) Assuming that V is also L Lipschitz, Corollary 1 holds with LF = LJ = 2L/(p + L) 2. For step size 1/L, Corollary holds with LE = L. 1 and p == p/L. In both cases, the contraction factor p is close to 0 when L/μ 1, for well conditioned problems. As outlined above, we may consider Fg FK in which case Corollary applies with a smaller value of p, recovering the result of [42, Proposition 3.1]
2.3 Superlinear and quadratic algorithms
The one-step Jacobian estimator in (4) as implemented in Algorithm 3 obviously not an exact estimator in the sense that one does not necessarily have (0) Joā(0) as k + 20. However, it is casy to see that this estimator is exact in the case of exact single-step algorithms,, meaning satisfics F(x, 0) æ(0) for all 2, Indeed, in this case, one has Fa, 0) 0 and JaF 0) = (0) for all x, O. Such situation occurs, for example, when applying Newlon's 's method to an unconstrained quadratic problem. This is a very degenerate situation as it does not really make sense to talk about an "iterative algorithm" in this case. It turns out that this property of being "asymptotically correct" " remains valid for very fast algorithms, that is, algorithms that require few iterations to converge, the archetypal example being Newton's method for which we obtain quantitative estimates.
2.3.1 Super-linear algorithms
The following is a typical property of fast converging algorithms.
Assumption 2 (Vanishing Jacobian) Assume that F: I2 x Rm R. C² and that the recursion 2x+1 = Fo(x)) converges globally, locally uniformly in 0 to the unique fixed point x0)) of Fe such that (x))((()).)))) 0.
Note that under Assumption 2, is always possible to find a small neighborhood of such that || llop remains small, that is, Assumption 1 holds locally and Lemma 1 applics. Furthermore, it is possible to show that the derivative estimate is asymptotically correct as follows.
Corollary 2 (Jacobian convergence) Let F: Ra x RT + Rn be as in Assumption 2. Then jos z (0) - (t) as k + oc, and jos x(0) = JgJ(
Proof Since J, (7(), 0) = 0, implicit differentiation of the fixed point equation reduces to Joz( JgF (z (0), 0), and the result follows by continuity of the derivatives.
5

---

Example 2 (Superlinearly convergent algorithm) Assume that F is C' and for each P > 0, there is R > 0 such that [F)(z)) x(0) Σ e|| z(0] for all x, 0 such that | 2(0)| < R. Then Corollary 2 applies as for any
F( + tc, 0) 3(0) 0.
Ja JaF()(0),0) 0)v
lim t05
2.3.2 Quadratically convergent algorithms
10a
Under additional quantitative assumptions, it is possible to obtain more precise convergence estimates similar to those obtained for implicit. differentiation, see Remark 2.
10-
One-step diff
10
Forward diff
10
Dist optimum
Implicit
Corollary 3 Let F be as in Assumption 2 such that x + J(a,0)F(2, (joint jacobian in (z,0)) is LJ Lipschitz (in oper- ator norm). Then, the recursion is asymptotically quadratically convergent and for each 2 1,
Iteratio
Figure 2: Newton's 's method (7)
2z) Joz L3 (=)) z(0)l.
quadratic convergence.
Proof: Following the same argument as in the proof of Corollary 2, we have
(|]0)z(00) dz = (l))F(zx(0,,0) - (=)(( ol S Lo(l=x=1(0)) - =(0)]l
(8)
As for the quadratic convergence, we may assume that F(f) and drop the variable to simplify notations. We have F(0) = 0 and for all
S( J.F((t) s ||||| / [(x)(tt s [>>||> [ tdt L
F(x) = F(o)=f(
Thus ///2[|0x+1|| S |L,/2| III and asymptotic quadratic convergence follows.
Example (Newton's algorithm) Assume that f: R. x Rm R is C with Lipschitz derivatives, and for each 0, x f(x, 0) is j-strongly convex. Then Newton's 's algorithu with backtracking line search satisfies the hypotheses of Corollary 3, see [15, Sec. 9.5.3). Indeed, it takes unit steps after finite number of iterations, denoting by z=)) the unique solution to V fz, 0) = 0, for x locally around x0))
F(x, 0) = x - fz, 0) (z)(x,0).
We have, V2 (2, 0) V2 f(r 0)x V V.f (z, 0), differentiating using tensor notations,
V2 (2,f(,,0),FF(,..
that (.F( 0 and Lipschitz continuity of the derivatives of implics Lipschitz continuity Jz, 8) using the fact that v2 fz, ž pl 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.
Hypergradient descent for bilevel problems
Consider the following bilevel optimization problem
min g(x(0)) s.t. x(0) arg fin 0),
where g and I are C functions. We will consider bilevel problems such that the inner fixed point equation F)x,0)
minimum is uniquely attained and be described
where F' is as in Assumption 1. The problem may then be rewritten as
min g(x(f)) s.t. x00) = ((x()),)).
(9)

---

Method
Time
Memory
Error sources
Piggyback recursion
kn(wCp + nm) KwOFm kwOF
n(n + m) n +
suboptimality + burn-in
AD forward-mode
suboptimality burn-in
AD reverse-mode
kn
suboptimality + burn-in
Implicit differentiation
wCr n²
suboptimalitt
One-step differentiation
wCF
n
suboptimality + lack of contractivity
Forward algorithm
kCr
suboptimalitty
Table 2: Time and memory complexities of the estimators in Section 2.1 (up to multiplicative constant). P' has time complexity denoted by Of and we consider iterations in R" with m parameter. w 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
8141 = ft - a.J In Vg))
(10)
where a >> 0 is a step size parameter. Note that the quantity 20 (0)) (0)" is exactly what is obtained by applying backpropagation to the composition of 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 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 CF denote the computation time cost of evaluating the fixed-point map F and > 0 be the multiplicative overhead of gradient evaluation, in typical applications, w S 5 (cheap gradient principle [6]). The time cost of evaluating the Jacobian of F is nwCr (n gradients). Forward algorithm evaluation (i.e., xx) has computational time cost kCF with a fixed memory complexity n. Vanilla piggyback recursion (2) requires k - 1 full Jacobians and matrix multiplications of costs n m. The forward-mode of AD has time complexity kwCrm (compute m partial derivatives each of them cost times the time to evaluate the forward algorithm), and requires to store the iterate vector of size n and m partial derivatives. The reverso-mode of AD has time complexity kwCr (cheap gradient principle on Ft)) and requires to store vectors of size n. Implicit differentiation requires one full Jacobian (wCpn) and solution of one linear system of size n X n, that is roughly n². Finally, one-step differentiation is given by only differentiating a single step of the algorithm at cost wCp. 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 Cr 2 n², then reverse AD has a computational advantage if A: S 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 informa- tive 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' -VV implicit differentiation should be applicd to the gradient mapping Vf = 0, not F. In typical application Cvy = O(n²), and the dominant cost of implicit diffcrentiation is O(n²), which of the same order the one-step differentiation to be inverted). However, if the implicit step was performed on F instead of V., it would incur a prohibitive cost of O(nt). In conclusion, the implicit differentiation phase is not only custom in terms of the implementation, but also in the very choice of the equation.
CF O(n²) (a linear system needs

---

3.2 Approximation analysis of one step differentiation
The following corollary, close to [42, Prop. 3.1], provides bound the one-step differenti- ation (Algorithm 3) gradient estimator for (9). The bound depends on P, the contraction factor, and distance the solution for the inner algorithm. The proof is given in Appendix A.
Corollary 4 Let F: R x Rm - R as in Corollary and consider the bilevel problem (9), where g is C. ig Lipschitz function with lo Lipschitz gradient. Then,
V Vo (0) Jo8 (V)V)( pLrtg +L = + Lply [T](]] x.
3.3 Approximate critical points
The following lemma is known, but we provide proof for completeness in Appendix A.
Lemma Assume that h: Tum - is C with I, Läpschitz gradient and lower bounded by htt Assume that for some > 0, for all l E N, |0041 01 + Vh(Ot)) & E. Then, for all KEN, K2K 1, we have
21(h(r0)-
min D
[h((0)]
K+1
Combining with Corollary 4, it provides complexity estimate for the recursion in (10).
Corollary 5 Under the selting of Covollary 4, consider ilerates in (10) with k: 2 and assume the following
supe ([=o(0))
S M, for some M > 0.
is LE Lipschitz and J(z,0)F is Ls Lipschitz jointly in operator norm.
g is C² ly Lipschitz with ly Lipschitz gradient.
Li - (1=, + 11) + +iv I-p
LF
g goo is lower bounded by g"
Then setting T + (Lobg + LFlv M for all K E,N,
min v ()))1 < 2L((g T) (00) - -g¹)
K +1
The level of approximate criticality the sum of a term proportional to p and term inversely proportional to K. For large values of & (many steps on the inner problem), K (many steps the outer problem), approximate criticality is essentially proportional to pLplg/(1 p) which is small if p is close to 0 (e.g. superlinear algorithms).
Numerical experiments
We illustrate our findings on three different problems. First, we consider Newton's '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 computational overhead negligible with respect to solution evaluation, 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 R Nxn be a design matrix, the first column being made of 1s to model an intercept. Rows are denoted by (d1....,ax). Let x e R and y E [ We consider the regularized logistic regression problem
N
min RERn
Billia, x) u.) + A|
(11)
8

---

Intertor point for QP
Newton for logistic regression
12.5
Gradient
10.0
6-
Autoditf
7.5
Time
Implicit
None
2.5
2-
One step
0.0
500
1000
1500
500
1000
1500
2000
Total number of parameters (1e3)
Figure 3: Timing experiment to evaluate step and implicit estimators agree up to an error order 10-16, 16, Right: Newton's method on logistic regression (11), one step and implicit estimators agrec up to error of order 10-12 Label "None" represent solving time and " "Autodiff", "Implicit" and "One stop" represent solving time and gradient evaluation for cach estimator in Section 2.1. The mcan is depicted using shading to indicate standard deviation estimated over 10 runs.
gradient. Left: Differentiable QP in (12), one
where l is the logistic loss, t+ log(1 + exp(-t)), À >> regularization parameter, and T-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 evaluated using jax automatic differentiation, and the matrix inversion operations are performed with an explicit call to linear system solver.
We denote by z(0) the solution problem (11) and try to evaluate the gradient of 0 +> |z<(0)]|//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 specify problem (11), our case size of and size of y, which is (n + I)N.
Figure: illustrates that both one-step and implicit differentiation enjoy marginal compu- tational 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 moti- vated by elements described in [4], which is associated with pytorch library implementing a standard interior point solver. Consider the following quadratic program (QP):
min BER
Qr 2" xx"Qu + q" s.t. =x=0, 0. Gx S h,
(12)
where Q E Ruxa is positive definite, A E Ru and G E RRx are matrices, E R", e R and IR are voctors. We consider z(0) the solution of problem (12) as a function of 8, the right-hand side of the equality constraint. We implemented in jax 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 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 Jaz(() using implicit differentiation. More details this experiment are given in Appendix B.
We consider evaluating the gradient of the function 0 + [[2 112 using the throe 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 n(n + 1) (n + 1)m + (n + 1)p. The results are presented in Figure 3 where we represent the lime required by algorithms as 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,

---

autcd (bcobronn
SUIDOM (ttr)
stop (acocioni
step (her))
Impleit dir (acoblaa))
imbictt liter)
10"1
10
10"
200
300
100
Reltiiee sub0pemalty
10
10a4
1077
10ah
10
20
30
-00
50
60
70
10
20
30
-00
50
60
70
Time 150
Time 15)
Figure 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 +. Right column: big learning rate 2 lines represent the error the iterates (resp. of the Jacobians).
w+I Dotted (resp. filled)
that both one-step and implicit differentiation enjoy a marginal additional computational overhead, contrary to algorithmic differentiation.
Weighted ridge using gradient descent. We consider weighted ridge problem with E Nxn E R N and vector of weights f R.;;
N
(0) = arg min fo(x)
Σ 0.10: - x))² 2 11=112. 2
#ER"
We solve this problem using gradient descent with adequate step-size F(z, ((z,0)=r-aV/e with =0( O, and we consider the K-step truncated Jacobian propagation F= = F² with K = 1/K where к is the effective condition number of the Hessian. Figure 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 | (8) [2
[+) for the Jacobian matrices. As expected, implicit differentiation is faster and more precise,
2 for the iterates, and
our gold standard which requires cuslom implicil 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 fixed time budget, the one step strategies allows to obtain higher iterate accuracy and similar, or better, Jacobian accuracy. In the small step regime, 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 dograded dcrivatives compared to full backpropagation, in line with [42].
step differentiation provides trade-off, for fixed time budget, one
Conclusion
We studied the one-step differentiation, also known Jacobian-free backpropagation, of generic iterative algorithm, and provided convergence guarantees depending on the initial rate of the algorithm. In particular, show that one-step differentiation of 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.
10

---

Acknowledgements
The authors acknowledge the support of the Al Interdisciplinary Institute ANITI funding, through the French "Investments for the Future PIA3" program under the grant agreement ANR-19-P13A0004, Air Force Office of Scientific Rescarch, Air Foree Material Command, USAF, under grant numbers PA8655-22-1-7012 ANR MaSDOL 19-CE23-0017-0, ANR Chess, grant
ANR Regulia, ANR GraVa ANR-18 CE40-0005. Jérôme Bolte, Centre Lagrange, and TSE-P.
References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeflrey Dean, Matthieu Devin, Sanjay Ghemawal, 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, Iyya 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 heterogencous systems, 2015. Software available from tensorflow.org
[2] Pierre Ablin, Gabriel Peyré, and Thomas Moreau. Super-efficiency of automatic diffcrentiation 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 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):10--118, 1994.
9 Yoshua Bengio. Gradient-Bacd Based Optimization of Hyperparameters. Neural Computation, (189)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 Wlinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficcintt and modular implicit differentiation. arXio preprint arXiv:2005.15 15183, 2021.
[12] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Tlinares-T.6pez, Fabian Pedregosa, and Jean-Philippe Vert. Efficient. and modular implicit differentiation. Advances in Neural Information Processing Systems, 35:5230-522 5242, 2022.
[13] Jérôme Bolte, Tam Le, Edouard Pauwels, and Tony Silveti-Falls. Nonsmooth implicit differentiation for machine-learning and optimization. In NewrIPS, volume 34, pages 13537-13549, 2021.
11

---

[14] Jerome Bolte, Edouard Pauwels, and Samuel Vaiter. Automatic differentiation of nonsmooth iterative algorithms. In Advonces in Neural Information Processing Systems, 2022.
[15] Stephen Boyd and Lieven Vandenberghe. Conver 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] Chin-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 NearIPS, 2022.
[19] Matthias Ehrhardt and Lindon Roberts. Aualyzing inexact hypergradients for bilevel learning. ar Xiv preprint arXiv:2301. 04 764, 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-955) 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 AAAT Conference on Artificial Intelligence, 2022.
[23] Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is atlention 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-soale PDE-construined oplimization, pages 148-164. Springer, 2003.
[27] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and tech- niques 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 ar Xiv: 200.00006.00006, 2021.
[30] Seppo Linnainmas. The representation of the cumulative rounding error of an algorithm 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 Lorraïne, Paul Vicol, and David Duvenaud. Optimizing millions of hyperpa- rameters by implicit differentiation. In AISTATS, pages 1540-1552 1552, 2020.
12

---

[33] Giulia Luise, Alessandro Rudi, Massimiliano Pontil, and Carlo Ciliberto. Differential properties of sinkhorn approximation for learning with wasserstein distance. In NeurIPS, volume 31, 2018.
[34] Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani Raiko. Scalable gradient- based tuning of contimuous regularization hyperparameters. In ICML, pages 2952-2960) 2016.
[35] Shcheryar Mchmood and Peter Ochs. Automatic differentiation of some first-order meth- ods in parametric optimization. In International Conference on Artificial Intelligence and Slatislics, pages 1584 1594. PMLR, 2020.
[36] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances Neural Information Processing Systems 32, pages
2019.
[37] Edouard Pauwels and Samuel Vaiter. The derivatives of sinkhorn-knopp converge. SIAM
J Optim, 33(3):1494-1517, 2023.
[38] Boris Polyak.
Introduction to optimization. In Optimization Software, Publications Division. Citesecr, 1987.
[39] Zaccharie Ramzi, Florian Mannel, Shaojie Bai, Jean-Lue Starck, Philippe Ciuciu, and Thomas Moreau. SHINE: SHaring the INverse estimate from the forward pass for bi-level optimization and implicit models. In International Conference on Learning Representations, 2022.
[40] Subham Sekhar Sahoo, Anselm Paulus, Marin Vlastelica, Vit Musil, Volodymyr Kuleshov, and Georg Martius. Backpropagation through combinatorial algorithms: Identity with projection works. In The Eleventh International Conference Learning Representations, 2023.
[41] Damien Scieur, Gauthier Gidel, Quentin Bertrand, and Fabian Pedregosa. The curse of unrolling: Rate of differentiating through optimization. Advances in Neural Information Processing Systems, 35:17133- 17145, 2022.
[42] Amirreza Shaban, Ching An Cheng, Nathan Hatch, and Byron Boots. Truncated back propagatoonnn for bilevel optimization. In AISTATS, pages 1723- 1732. PMLR, 2019.
[43] Lieven Vandenberghe. The cvxopt lincar and quadratic cone program solvers. Online: htll:///cwopp org/ dddoe//eeaalion/ooeeeeeeoo
pdf. 2010.
[44] R. E. Wengert. A simple automatic derivative evaluation program. Communications of the ACM, 7(8):463 464, August 1964.
[45] Matt Johnson Zico Kolter, David Duvenaud. Deep implicit layers - neural odes, deep equilibirum models, and beyond, 2020. NeurIPS Tutorial.
13

---

A Proof of Section
Proof of Corollary 4: We have (o(goz)(0) Jez(0) and
J - Jos xx) =(JoI(0)= (()0))))))))0)0)))))))0000000100 100
- - (0) vg(z(0)).
The result follows from the triangular incquality combined with Corollary 1 and the following
II(Gz((0) Jos x00 v [})ā((0) xz (0))lop (0)00)))))((((( Vg()()))
- [111J00(00) - Jos ax xxx
and
0os xx (Vg(())) Vg(E(0)) & 1 x (V)(e))))) Vg(T(())) S [ xz
Proof of Lemma 3: We have, using the "descent lemma" (scc [7]), for all 1 E O...... K,
L
h(01+1) - 1(O)) < (Vh(0t), 01+1 - 81) + || - 0:112
2
Vh(O,)
2 2
0111 - 01
+ +1101+1 - 0:11
L
[1+1 01 + Vh()t)
Vh((i) 12)
L
(e ( min 1=0...,K K
<
2L
Summing for = O,... , K, we have
ht h- S h(0k+1) h(Bb) S
K +
(2- (
min
[VV(()))1))
2L
I=0...,K
and deduce, by using concavity of the square root, that
[h(00)) 2 2L(h(ro) h²)
min
K+1
Lemma Let A: - and B: Rd + MXP, be MA and MB bounded and La and La Lipschilz respectively, in operator norm, then A B is MALB LAMB Lipschilz in operalor norm.
Proof
We have for all x. y E Rdd
[[A(x)B(x) - A()) = [[A((x) - - B(x))llop
S [[(A(x) - Bz))) + (B(())))()) B(x))
(|
< A((x) A(v)llop [B())) + |1A(y) B B(x)
A ML + MMLR)]
Lemma The function M (I M) -1 is p)- Lipschitz in operator norm on the set of square matrices M such that |M Mlop p <
Proof : For any two square matrices A, B with operator norm bounded by p p<1, T- A and I- B are invertible and the operator norm of the inverse is at most (1-p)-1 (see proof of
14

---

Lemma 2). We have
A) (1-- (
(1 -
- 'lop
|| (A -B)
I |1(I A)oollop II(A - B)llop : - B) llop
(A Blop
(1 p)2
Proof Corollary 5: Let us evaluate a Lipschitz constant of Vo(g o). We have
o = =,z(()"VV(((00))
Vg((()))
=
This can be seen as a composition of a matrix function with E. Combining Lemma and Lemma 5 with the fact that J,F and JoF are both L, Lipschitz and LF bounded, the function (x, 0) + (1 - (z, 0))-0))---- -o f 0) is Lobe Lu - (+z Lipschitz.
(T-)) (1=
Invoking Lemma 4 and using the fact that the operator norm of a vector is the euclidean norm, we have that the function (x, 0) + (I J,F(x, ()) lv Le. Lipschitz.
is Le 1-p (t,, + 1) lg +
I-p
From the implicit function theorem, [| dpop & Lp/(1 - p) so that to is LF/(1 p) i) is (1ep [-p, + +1) ly + 1-,,
Lipschitz and i)
Le Lipschitz so the choice of a is I-p
such that 1/a a Lipschitz constant for the gradient.
Using Corollary 4, we are in the setting of Lemina 3 and the result follows.
B Experimental details Section 4
B.1 QP experiment
Implicit differentiation: We explicitly formed the Jacobians of the iterative process (function f) and solved the underlying linear systems using solver in jax. The inversion was not done based on VJP (using fixed point iterations or conjugate gradient) because, for the relatively small scale of our problems, pure lincar algebra was more efficient. We did not form the full implicit Jacobian in equation (3) using full matrix inversion; we only solved a linear system involving left incoming vector in equation (3). this sense, we did implement VJP for implicit different iation although we explicitly formed Jacobians because this was the most efficient.
Goneration of random QPs: To generate the QP instances, the quantitics n. m. p are varied on a logarithmic scale. We set A where M of size n x n with entries uniform in [-,, Constraint matrices A and G also have entries uniform in [-1, and we chose m and p to be smaller than n so that feasibility generic and occurs with probability one.
15


Reviews:
[{"review_id": "zQ9CqbaCE0", "reviewer": "Reviewer_WGTU", "summary": "This paper studies the gradient of iterative algorithms, specifically examining one-step differentiation of these algorithms. This approach replaces the complex Jacobian inverse in implicit differentiation with a simple and fast identity approximation. The work refines the theoretical approximation analysis for one-step differentiation and characterizes its efficiency under fast convergent algorithms. Numerical examples are provided to illustrate the well-foundness of the one-step estimator.\n", "questions": "Please see above.", "limitations": "No limitations are included.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- Organization. The paper is well-structured overall. Each section has clear motivation, and there is good coherence in the transition of analysis. The presentation of the algorithms and analysis is easy to follow.\n- Detailed approximation analysis. The work includes the approximation analysis for different algorithms under the one-step setting.", "weaknesses": "- Missing references. The work [1] also discusses the theoretical aspects of the differentiation approximation of iterative processes and should be included.\n- Notations. The paper could benefit from more explicit notation to enhance readability.\n- Novelty. While the novelty of the algorithm proposal is not a major concern, it becomes a considerable issue regarding the analysis part. The concept of one-step differentiation has been proposed for a while. Although this paper does not claim novelty in the algorithm proposal and instead focuses on the theoretical aspects, the approximation analysis is not sufficiently novel from a theoretical standpoint. I would expect more insights related to optimization and generalization when applying the one-step differentiation to learning problems.\n\n[1] Zhengyang Geng, Xin-Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. Advances in Neural Information Processing Systems, 2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly thank the referee for his constructive comments. Regarding the weaknesses pointed out by the reviewer, we propose the following modification which we hope will address the concerns of the reviewers so that the reviewer can update his evaluation. \n\n- One-step differentiation has indeed connections with phantom gradients proposed in [1], we will definitely add and discuss the missing reference.\n\n- We will sharpen the notations throughout the text and propose a deep rewriting of Section 3.3 to avoid notational confusion.\n\n- Following the reviewer's comment, as well as the comments of other reviewers, we will rework Section 3.3 and add a corollary providing explicit optimization guarantees for the bilevel problem. This will combine Corollary 4 and Lemma 3. This constitutes a specific instantiation of our results for learning problems which are formulated as bilevel problems with applications in hyper-parameter tuning or meta-learning."}, {"review_id": "o5JhZCeU1Y", "reviewer": "Reviewer_G3kk", "summary": "This paper develops a convergence-rate analysis for the approximation of the Jacobian under the technique of one-step differentiation. This technique is similar to that of iterative differentiation (aka unrolling, differentiating through optimization), except only the last iterate of the algorithm is used to compute the Jacobian.\n\nThe authors develop a quite generic framework that contains many common algorithms as special cases. The paper also contains several examples of how to apply this framework to specific examples such as gradient descent and Newton's method.", "questions": "L20: I didn't understand the point of the phrase \"it does not solely relies on the compositional rules of differential calculus\" (and it has at least a grammatical error relies -> rely)", "limitations": "In general the paper is quite honest about its limitations. However, I suggest the authors expand a bit more on which class of algorithms verify and which don't Assumption 1 (see Weaknesses)", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is clearly written, notation is largely standard. Main results are clearly explained. Overall, this paper was a pleasure to read.\n\n2. As far as I can tell, the results in this paper are sound.\n\n3. Experiments are to the point and clearly illustrate the theoretical results of the paper.\n\n4. The topic considered is an important one. The authors take a technique that has been shown to work empirically and develop a sound theory around it.", "weaknesses": "1. The current framework doesn't apply to algorithms that either i) depend on past iterates, such as gradient descent with moment, or that ii) are not always contractive such as accelerated gradient descent.\n\nIt could be possible to easily extend the results to ii) by considering an enlarged space where x now contains the current iterate and past information (such as momentum), although I haven't checked whether its possible to get back the bound on the Jacobian after this transformation.\n\n2. One of the most interesting results IMO is the convergence of a bi-level scheme, but the analysis done in 3.3 seems more of an afterthought. To start, the statement is highly confusing: its presented in terms of f, but the f here is not (if I understood correctly) the f of the inner optimization. Instead, it should be understood as any f that is L Lipschitz. In fact, (and this is nowhere written), to obtain a bound on the suboptimality of the bilevel problem, we should take f = g. Even with this in mind, casting the result of Lemma 3 into a bound for the bilevel problem requires to make a number of substitutions. Why not state clearly the rate for the bilevel problem in term of the quantities and constants that are defined for this problem? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the referee for his feedback. We provide below detailed answers.\n\n1. We agree with the referee that some algorithms may appear not directly in the form of a contraction or even not directly in the iterative form given in (1). A typical example is the heavy ball method for strongly convex problems: it should be considered in the phase space (because it is a second-order algorithm), and iterations are not contractive. However, the composition of a certain number of iterations is contractive in the phase space. We will add a remark and cite this example with a bibliographic reference to illustrate it.\n\n2. We agree that Section 3.3 needs rewriting, and we will work on it in a revised version. More precisely, we will add a corollary for the bilevel problem by combining the results of Corollary 4 and Lemma 3. Furthermore, we will modify the notations of Lemma 3 as it is not harmonized with the rest of the text (x and f do not have the same meaning as above). We prefer to keep Lemma 3 as a separate result since it applies beyond bi-level problems."}, {"review_id": "3e6DsLiBkm", "reviewer": "Reviewer_USU8", "summary": "The paper presents a method called one-step differentiation, also known as Jacobian-free backpropagation, as an alternative to automatic and implicit differentiation. The authors analyze the theoretical approximation of one-step Jacobian and provide specific examples such as Newton's method and gradient descent, on which the authors showed explicit convergence of the one-step case to ground truth solution.. They demonstrate the efficiency of one-step differentiation in bilevel optimization and provide numerical illustrations using logistic regression, interior point solver for quadratic programming, and weighted ridge regression with gradient descent. On these examples, the authors demonstrate the efficiency of the one step approach while demonstrating clear speed advantage over the autodiff approach.", "questions": "line 134 can the author please discuss more when the practioners are advised to use the K last step? in other words, K-step differentiation of iterative algorithm.", "limitations": "The paper lacks a discussion of potential limitations or scenarios where the one-step differentiation method may not be suitable or may exhibit suboptimal performance.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Automatic differentiation of iterative system is long known for its slowness, rendering many differentiable systems useless in practice. Of course, one can re-route to implicit differentiation, also known as sensitivity analysis, adjoint method, etc. However, implementing them is not fun, especially for complex systems. The authors proposed a competitive alternative that is both easy to implement, fast to compute (figure 3), and, without any loss for accuracy (line 284). I can see it laying the foundation for many downstream applications.\n\nLike author pointed out, using just the jacobian of the last step seems like a too naive approximation. However, the author proved several crucial theorems in the paper showing that it's in fact quite accurate in many cases (of course not all the cases).\n\nThe result is general in the sense that it applies to all iterative methods that have a fixed point.\n\nThe paper is well written, with a good balance of text, code, and proofs, making the paper easy to follow. That said, I admit that I did not check every step of the proof carefully.", "weaknesses": "As a theoretical paper, I really don't see much weakness. That said, I would be interested in seeing validation and experiments on more complex engineering system, potentially involving highly non-convex neural networks.\n\nMinor:\nFigure 1 caption, \"is explicited\". explicit is not a verb.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the referee for his positive feedback. Experiments and validation on highly nonconvex neural network is something that we are currently investigating. From our intuition so far (and as predicted by the theory we developed), one-step differentiation is not a panacea that will perfectly work under all circumstances, but it may have some important benefits for specific architectures under specific settings. We do not have solid enough findings in this direction to include in the present work, but we will for sure invest efforts in this direction.\n\nOf course, we will correct the typo and include a discussion regarding situations for which using K-last steps differentiation is advisable. This will also provide an answer to the reviewer's last: the typical scenario for which one-step differentiation may fail is for slow algorithms, which corresponds to a contraction factor close to 1 (or even equal to 1, in this case, even implicit differentiation has no guarantee). This is actually the reason why we propose to differentiate K steps instead of a single step in this situation. This will be more properly discussed in the numerical and conclusion sections."}, {"review_id": "KpVPZuBmDF", "reviewer": "Reviewer_agCr", "summary": "In this paper, the authors consider the problem of differentiating the fixed-point of an algorithm with recursive update, with respect to some parameter of the algorithm. They propose a one-step automatic differentiation technique where, once having approximated the fixed-point through the recursive algorithm, they back-propagate only through the last step. They provide a theoretical and numerical comparison of their one-step technique with the existing Automatic and Implicit Differentiation techniques. They show that the technique converges for superlinear algorithms, and incurs an error for linear algorithms which depends on the rate of convergence. For linear algorithms, they propose to use K-step technique where K depends on the condition number. They compare the time complexity of the three strategies on the Newton's method applied to weighted logistic regression and the interior point method applied to constrained quadratic programming problem. They also show the K-step technique on gradient descent applied to weighted ridge regression.", "questions": "-> In implementations of Algorithms 1, 2 and 3 for Bilevel Optimization, did the authors first construct the jacobians and then evaluate J^T \\grad g? Because that is quite inefficient. For bilevel Optimization, implementing VJP for each algorithm is the suitable way. The authors should mention that.\n\n-> In Table 2, I think the authors should clarify that the time complexity of Piggyback recursion (2) is not the same as that of the forward mode AD recursion (1). I believe that is kn\\omega C_F. Similarly, the memory complexity is of forward mode AD is n.\n\n-> Figure 4 lacks explaination. I understand that the weighted ridge regression example in Section 4 is aimed to depict the behaviour of the three techniques on gradient descent applied to an ill-conditioned problem. But I don't really understand what is shown in Figure 4.", "limitations": "Not applicable.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "-> For super linearly convergent algorithms and for fast linear algorithms, the authors demonstrate\n\t* Easy implementation of Automatic Differentiation, and\n\t* Time and Memory complexity of Implicit differentiation.\n\n-> The derivative error of their one-step technique is shown to be depended on the error in the estimation of the solution and the convergence rate (which is zero for superlinear algorithms).\n\n-> Except for the third example in Section 4 and in particular Figure~4, the paper is well-written and structured and very easy to follow. The proofs are quite easy to understand.", "weaknesses": "-> For large scale applications (n >> 1), memory and time complexities of one update step of super-linear algorithms do not scale well with n:\n\n- Time: O(n^2) for quasi-newton methods, O(n^3) for Newton methods.\n- Memory: O(n^2) for both.\n\n-> In such cases, we resort to the first order methods where the convergence rate is non-zero. For ill-conditioned problems, the authors suggest K-step back-propagation with K = 1/\\kappa but then that is just truncated back-propagation and we still have a memory overhead.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the referee for his thoughtful comments.\n\nIndeed, super linearly convergent algorithms / second-order like methods are fast, but the price to pay is that each step is very expensive (compared to first-order methods, for example). There is no free lunch, and in that case, our analysis only makes sense practically for relatively small-scale scenarios or very well-conditioned problems, but this is a problem-dependent feature, and one does not really have control of this. \nRegarding the memory overhead of K step differentiation, there is indeed an overhead, but it could be relatively small if K remains not too large (e.g., well-conditioned problems).\n\nQuestions: \n- We did form Jacobians of the iterative process (F) and solved linear systems using a dedicated solver (in Jax). The inversion was not done based on VJP (using fixed point iterations or conjugate gradient) because, for the relatively small scale of our problems, linear algebra was more efficient.\n- We did not form the full implicit Jacobian in equation (3) using full matrix inversion; we only solved a linear system involving a left incoming vector in equation (3). In this sense, we did implement VJP for implicit differentiation.\n- We will make this precise in a dedicated appendix in the revised version.\n\nPoint 2: The reviewer is right; we will add a line with an evaluation of forward AD, which should be slightly more favorable than the piggyback recursion as Jacobian matrices are not explicitly formed and multiplied.\n\nPoint 3: We agree and will revise the discussion in the weighted ridge regression section and better explain the content of Figure 4 with more detailed comments."}, {"review_id": "ben4MqQrJA", "reviewer": "Reviewer_qFET", "summary": "In bilevel optimization, or optimization problems with equilibrium constraints, computing a derivative of the upper-level problem is a well-known stumbling block as this requires \"differentiating through\" the solution to the lower level problem. \n\nThis paper provides a theoretical study for one approach to overcoming this issue, __one-step differentiation__, also known as Jacobian-free Backpropagation. Although one-step differentiation is deceptively simple, the core contribution of this paper is that it is safe to use in many situations.\n\nTheoretical results are bolstered by some numerical results. In addition there are several nice examples, remarks, and corollaries that could be of use to practitioners (for example, how to improve the quality of the one-step derivative by using a k-step approach instead).", "questions": "In [3] implicit networks with a head and tail ($S_{Theta}$ and $Q_{Theta}$ in their notation) are considered. This complicates the analysis of the hypergradient, see the proof of Theorem 3.1, particularly when analyzing the inner product between the true hypergradient and the approximation $p_{\\Theta}$.\n\n Your analysis of hypergradients (i.e. Corollary 4) works out nicely, partly because you consider the distance between true and approx. hypergradient, not the inner product. Can you apply this to get quantitative bounds for the setting considered in [3]?\n\n\n[3] _JFB: Jacobian-Free Backpropagation for Implicit Networks_ by Wu Fung et al (2022)", "limitations": "Yes.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": " - In my opinion, this paper's biggest strength is that it has a simple message (\"One-step differentiation works in many situations\") and conveys this message in a clear and accessible to non-experts way.\n  - There are several little gems of wisdom sprinkled throughout this paper, e.g. applying one-step differentiation to $F^K$ instead of $F$ (top of pg. 4), using different operators for the forward and backward pass for implicit differentation (Remark 2).\n - I like the bound on hypergradient approximation (Corollary 4).\n - The clear comparisons between the complexities of various gradient estimation approaches contained in Table 2 is great.", "weaknesses": "\n1. In addition to Automatic Differentiation (AD) and Implicit Differentiation (ID), there is a third benchmark approach the authors should consider, namely _Inexact Automatic Differentiation_ (IAD) (see for example [1,2]). In particular, [2] makes a surprising connection between IAD and ID, showing they are essentially the same. I'd like to see some discussion of the IAD approach in this paper, and perhaps even an inclusion of this approach into the numerical experiments.\n2. See \"Questions\" section below for some questions on the relationship between your work and that of [3].\n\n*Minor Issues*\n\n3. the phrase \"piggyback recursion\" is used for the first time in line 207. I suggest mentioning this terminology immediately after equation (2) and providing a reference.\n4. In addition to [34] (Vlastelica et al), I'd suggest citing [4], which predates [34] and also considers one-step differentiation for polytope-constrained lower-level problems.\n5. The notation in Lemma 3 is confusing. I'd strongly suggest using $\\theta_i$ instead of $x_k$ and $g$ instead of $f$, as the Lemma is intended to be applied to the upper level problem.\n6. I have several comments regarding Fig. 3:\n      - The time displayed is to calculate a single gradient, right? If yes, this should be made clearer in the caption. \n      - The x-axis needs to be changed. I'd suggest $5\\times 10^{5}$ instead of 50000 and so on.\n      - What does the shading represent?\n      - See comment 1 about Inexact Automatic Differentiation.\n7. Could you say more about how you randomly generate QP instances (pg 8--9) in particular, how do you ensure feasibility?\n\n*Typos etc.*\n\n8. \"well-foundness\" in line 10 should be \"well-foundedness\".\n9. \"a important\" in line 12 should be \"an important\".\n10. \"relies\" on line 20 should be \"rely\".\n11. \"distance with\" on lone 33 should be \"distance to\".\n12. \"quantiative\" on line 58 should be \"quantitative\"\n13. \"parameteric\" on line 72 should be \"parametric\"\n14. \"Superlinar\" on line 144 should be \"Superlinear\"\n15. \"guaranties\" on line 199 should be \"guarantees\"\n \n[1] _Automatic Differentiation of Some First-Order Methods in Parametric Optimization_ by Mehmood and Ochs (2019)\n\n[2] _Analyzing Inexact Hypergradients for Bilevel Learning_ by Ehrhardt and Roberts (2023)\n\n[3] _JFB: Jacobian-Free Backpropagation for Implicit Networks_ by Wu Fung et al (2022)\n\n[4] _Learn to Predict Equilibria via Fixed Point Networks_ by Heaton et al (2021) \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly thank the referee for his detailed comments and evaluation of our manuscript.\n\n1. We will mention Inexact AD and discuss the references proposed by the referee in the related work section. In our situation, it is not really clear how to apply inexact AD. Indeed, inexact AD is often applied to first-order methods, where the evaluation of the Jacobian of a single iteration is cheap, and many iterations are performed. In the context of Figure 3, we consider super-linearly convergent algorithms, typically variations of Newton's method. In this setting, evaluation of a single iteration is costly (typically similar as implicit differentiation) with very few iterations performed. One step differentiation is \"exact\" in the sense that it finds the same derivative as implicit differentiation up to numerical errors (in our experiments) while inexact AD will necessarily have some time overhead as it requires more operations compared to one-step differentiation (form the full Jacobian and perform fixed point iterations). We, therefore, believe that the settings where one-step differentiation has good performances are not favorable settings for inexact AD (and vice versa). We see the two approaches as complementary, which we will discuss in a revised version of the paper. For this reason, we will not include timing experiments for inexact AD as the comparison would not be really fair (for the same reason, we did not include forward AD, or unrolling, in our experiments).\n\n3. Indeed, we will mention the term piggyback right after (2).\n4. We will include the cited references in the related work section.\n5. The referee is right, Lemma 3 will be modified, and we will add a corollary to specify it for bilevel problems.\n6. We will take the remarks of the referee into account regarding Figure 3. In particular, the shaded area is the standard deviation over 10 repetitions of the experiment.\n7. To generate the QP instances, the quantities $n, m, p$ are varied on a logarithmic scale. We set $Q = M^TM$ where $M$ is of size $n \\times n$ with entries uniform in [-1,1]. Constraint matrices $A$ and $G$ also have entries uniform in [-1,1], and we chose $m$ and $p$ to be smaller than $n$ so that feasibility is generic and occurs with probability one. This will be discussed in a dedicated appendix.\nWe will correct all the typos pointed out by the referee.\n\nQuestion:\nThanks for this nice question. The ideas developed in [3] are indeed similar, but the theoretical results are quite different. The main assumption in [3] is (3.3), which is pretty different from ours:\n- It does not put any constraint on the magnitude of the derivative with respect to parameters (theta).\n- It requires good conditioning: Jacobian matrices (with respect to parameters theta) should be close to identity, and the level of closedness should be balanced with respect to the contraction factor ($\\gamma$ in [3]).\n\nThe results are, therefore, quite different in nature (descent direction versus quantitative estimation). We will not be able to obtain quantitative estimation under the exact same setting as [3] (since assumptions are different), but we will be able to obtain quantitative estimates under suitable assumptions. These will be exactly the same as in [3] and will be essentially the same as what we developed in Corollary 4. We will add a remark regarding this after Corollary 4 and in comparison with the existing work section."}]

Author rebuttal:
We thank the referees for the feedback on our work. We are glad that the referees found our paper *well-written* (WGTU, G3kk, USU8, qFET) with a *clear message* (qFET). Both the theoretical analysis and our experiments seems to have been appreciated by the reviewers.

We propose to include remarks and discussions following their comments. We will also modify section 3.3: use different notations for Lemma 3 and add a corollary dedicated to bilevel problems (following remarks of WGTU, G3kk and qFET). We will also include more details on the numerical experiments (implicit differentiation implementation details and random generation details, following the remarks of agCr and qFET). We provide below a more detailed response to each comment of the reviewers.

Meta-review: Theory and experiments complement one another well in this paper, which tackles automatically differentiating iterative algorithms (consider differentiating through the iterative optimization procedure, but using only the final iterate). All reviewers here recommend acceptance with positive reviews (some strongly so, and with high confidence!), many highlighting that it is both well presented and technically solid (with a detailed convergence analysis). During discussion, one reviewer underscored that the paper also synthesizes related work well.

The discussion with reviewers was productive. I will highlight especially the thread with Reviewer G3kk, which led to some clarifying remarks on heavy ball methods and contractivity. The exchanges with several reviewers also led to technical revisions and improvements in Section 3.3 that I agree with, including stating guarantees for the bilevel problem, as well as several new references to discuss, clarifications on the "K-step" variant of the estimator, and more. I recommend indeed following through on all of these edits, as the authors have indicated they will. That will strengthen the paper further.
Decision: Accept (spotlight)