Title: One-step differentiation of iterative algorithms
Venue: NeurIPS 2023
Accepted: yes
Keywords: automatic differentiation

implicit differentiation

super-linear algorithms

bilevel optimization.
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

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 Jacobianfree backpropagation, a method as easy as automatic differentiation and as efficient 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-foundedness of the one-step estimator.

1 Introduction

Differentiating the solution of a 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 idea of evaluating derivatives through the compositional rules of differential calculus in a user-transparent way. It is a mature concept [27] implemented in several machine learning frameworks [36, 16, 1]. However, the time and memory complexity incurred may become prohibitive as soon as 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 a linear system. The user needs to implement custom rules in 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 is carefully done, this is most of the time the gold standard for the task of differentiating problem solutions.

Contributions. We study a one-step Jacobian approximator based on a simple principle: differentiate only the last iteration of an iterative process and drop previous derivatives. The idea 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 iterative algorithms. One-step estimation constitutes a rough approximation at first sight and our motivation to study is its ease-of-use within automatic differentiation frameworks: no custom Jacobian-vector products or vector-Jacobian products needed as long as a stop_gradient primitive is available (see Table 1).

Nice, France.

We conduct an approximation analysis of one-step Jacobian (Corollary 1). The distance to the 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 in Corollary 4 the estimation of approximate critical points. Finally, numerical 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.

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 studied in the machine learning community for smooth methods [35, 32, 37], and nonsmooth methods [14]. Implicit differentiation is a recent highlight in machine learning. It was shown to be a good way to estimate the Jacobian of problem solutions, for deep equilibrium network [5], optimal transport [33], and also for nonsmooth problems [13], such as sparse models [10]. Inexact implicit differentiation 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 : R n × R m → R n denotes a recursive algorithmic map from R n to R n with m parameters. For any θ ∈ R m , we write F θ : x ↦→ F ( x, θ ) and let F k θ denote k recursive composition of F θ , for k ∈ N . The map F θ defines a recursive algorithm as follows

We denote by J x F θ the Jacobian matrix with respect to the variable x . The following assumption is sufficient to ensure a non degenerate asymptotic behavior.

Assumption 1 (Contraction) Let F : R n × R m → R n be C 1 , 0 ≤ ρ < 1, and X ⊂ R n be nonempty convex closed, such that for any θ ∈ R m , F θ ( X ) ⊂ X and ‖ J x F θ ‖ op ≤ ρ .

Remark 1 The main algorithms considered in this paper fall in the scope of smooth optimization. The algorithmic map F θ is associated to a smooth parametric optimization problem given by f : R n × R m → R such that f θ : x ↦→ f ( x, θ ) is strongly convex, uniformly in θ . Two examples of algorithmic maps are given by gradient descent, F θ ( x ) = x -α ∇ f θ ( x ), or Newton's F θ ( x ) = x -α ∇ 2 f θ ( x ) -1 ∇ f θ ( x ) for positive step α > 0. For small step sizes, gradient descent provides a contraction and Newton's method provides a local contraction, both fall in the scope of Assumption 1. Other examples include inertial algorithms such as the Heavy Ball method [38], which has to be considered in phase space and for which a 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 F θ , denoted by ¯ x ( θ ) and for which we will be interested in estimating derivatives.

Lemma 1 Under Assumption 1, for each θ in R m there is a unique fixed point of F θ in X denoted by ¯ x ( θ ) , which is a C 1 function of θ . Furthermore, for all k ∈ N , we have ‖ x k ( θ ) -¯ x ( θ ) ‖ ≤ ρ k ‖ x 0 ( θ ) -¯ x ( θ ) ‖ ≤ ρ k ‖ x 0 -F θ ( x 0 ) ‖ 1 -ρ .

This is well known, we briefly sketch the proof. The mapping F θ is a ρ 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 theorem. The convergence rate is classical. We are interested in the numerical evaluation of the Jacobian J θ ¯ x ( θ ), thus well-defined under Assumption 1.

The automatic differentiation estimator J AD x k ( θ ) = J θ x k ( θ ) propagates the derivatives (either in a forward or reverse mode) through iterations based on the piggyback recursion [26], for i = 1 , . . . , k -1,

Under assumption 1 we have J AD x k ( θ ) → J θ ¯ x ( θ ) and the convergence is asymptotically linear [25, 8, 35, 41, 14]. J AD x k ( θ ) is available in differentiable programming framework implementing common primitives such as backpropagation.

The implicit differentiation estimator J ID x k ( θ ) is given by application of the implicit function theorem using x k as a surrogate for the fixed point ¯ x ,

By continuity of the derivatives of F , we also have J ID x k ( θ ) → J θ ¯ x ( θ ) as k → ∞ , (see e.g. [27, Lemma 15.1] or [12, Theorem 1]). Implementing J ID x k ( θ ) 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 x k , which can be seen as an approximation of J ID [35, 19].

The one-step estimator J OS x k ( θ ) is the Jacobian of the fixed point map for the last iteration

Contrary to automatic differentiation or implicit differentiation estimates, we do not have J OS x k ( θ ) → J θ ¯ x ( θ ) in general as k → ∞ , but we will see that the error is essentially proportional to ρ , and thus negligible for fast algorithms for which the estimate is accurate.

From a practical viewpoint, the three estimators J AD , J ID and J OS are implemented in a differentiable programming framework, such as jax , thanks to a primitive stop_gradient , as illustrated by Algorithms 1, 2 and 3. The computational effect of the stop_gradient primitive is to replace the actual Jacobian J x F ( x i ( θ ) , θ ) by zero for chosen iterations i ∈ { 1 , . . . , k } . Using it for all iterations except the last one, allows one to implement J OS in (4) using Algorithm 3. This illustrates the main interest of the one-step estimator: it can be implemented using any differentiable programming framework which provides a stop_gradient 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 ∈ R n × n with ‖ A ‖ op ≤ ρ < 1 and B, ˜ B ∈ R n × m , then

Moreover, we have the following estimate,

Proof : First for any v ∈ R n , we have ‖ ( I -A ) v ‖ ≥ ‖ Iv ‖ - ‖ Av ‖ ≥ (1 -ρ ) ‖ v ‖ , which shows that I -A is invertible (the kernel of I -A is trivial). We also deduce that ‖ ( I -A ) -1 ‖ op ≤ 1 / (1 -ρ ). Second, we have ( I -A ) -1 -I = A ( I -A ) -1 , since (( I -A ) -1 -I )( I -A ) = A , and therefore

The norm bound follows using the submultiplicativity of operator norm, the triangular inequality and the fact that ‖ A ‖ op ≤ ρ and ‖ ( I -A ) -1 ‖ op ≤ 1 / (1 -ρ ). □

Corollary 1 Let F and X be as in Assumption 1 such that θ ↦→ F ( x, θ ) is L F Lipschitz and x ↦→ J θ F ( x, θ ) is L J Lipschitz (in operator norm) for all x ∈ R n . Then, for all θ ∈ R m ,

Proof : The result follows from Lemma 2 with A = J x F (¯ x ( θ ) , θ ), B = J θ F (¯ x ( θ ) , θ ) and ˜ B = J θ F ( x k -1 , θ ) using the fact that ‖ B ‖ op ≤ L F and ‖ B -˜ B ‖ op ≤ L J ‖ ¯ x ( θ ) -x k -1 ‖ . □

Remark 2 (Comparison with implicit differentiation) In [12] a similar bound is described for J ID , roughly under the assumption that x → F ( x, θ ) also has L J Lipschitz Jacobian, one has

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 ρ and L J ‖ ¯ x -x k -1 ‖ are of the same order.

Corollary 1 provides a bound on ‖ J OS x k ( θ ) -J θ ¯ x ( θ ) ‖ op which is asymptotically proportional to ρ . This means that for fast linearly convergent algorithms, meaning ρ /lessmuch 1, one-step differentiation provides a good approximation of the actual derivative. Besides, given F which satisfies Assumption 1, with a given ρ < 1, not specially small, one can set ˜ F θ = F K θ for some K ∈ N . In this case, ˜ F satisfies assumption 1 with ˜ ρ = ρ K and the one-step estimator in (4) applied to ˜ F becomes a K -steps estimator on F itself, we only differentiate through the K last steps of the algorithm which amounts to truncated backpropagation [42].

Example 1 (Gradient descent) Let f : R n × R m → R be such that f ( · , θ ) is µ -strongly convex ( µ > 0) with L Lipschitz gradient for all θ ∈ R m , then the gradient mapping F : ( x, θ ) ↦→ x -α ∇ x f ( x, θ ) satisfies Assumption 1 with ρ = max { 1 -αµ,αL -1 } , smaller than 1 as long as 0 < α < 2 /L . The optimal α = 2 / ( L + µ ) leads to a contraction factor ρ = 1 -2 µ/ ( L + µ ). Assuming that ∇ 2 xθ f is also L Lipschitz, Corollary 1 holds with L F = L J = 2 L/ ( µ + L ) ≤ 2. For step size 1 /L , Corollary 1 holds with L F = L J ≤ 1 and ρ = 1 -µ/L . In both cases, the contraction factor ρ is close to 0 when L/µ /similarequal 1, for well conditioned problems. As outlined above, we may consider ˜ F θ = F K θ in which case Corollary 1 applies with a smaller value of ρ , 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 is obviously not an exact estimator in the sense that one does not necessarily have J OS x k ( θ ) → J θ ¯ x ( θ ) as k →∞ . However, it is easy to see that this estimator is exact in the case of exact single-step algorithms, meaning F satisfies F ( x, θ ) = ¯ x ( θ ) for all x, θ . Indeed, in this case, one has J x F ( x, θ ) = 0 and J θ F ( x, θ ) = J θ ¯ x ( θ ) for all x, θ . Such a situation occurs, for example, when applying Newton'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 : R n × R m → R n is C 1 and that the recursion x k +1 = F θ ( x k ) converges globally, locally uniformly in θ to the unique fixed point ¯ x ( θ ) of F θ such that J x F (¯ x ( θ ) , θ ) = 0.

Note that under Assumption 2, it is always possible to find a small neighborhood of ¯ x such that ‖ J x F θ ‖ op remains small, that is, Assumption 1 holds locally and Lemma 1 applies. Furthermore, it is possible to show that the derivative estimate is asymptotically correct as follows.

Corollary 2 (Jacobian convergence) Let F : R n × R m → R n be as in Assumption 2. Then J OS x k ( θ ) → J θ ¯ x ( θ ) as k →∞ , and J OS ¯ x ( θ ) = J θ ¯ x ( θ ) .

Proof : Since J x F (¯ x ( θ ) , θ ) = 0, implicit differentiation of the fixed point equation reduces to J θ ¯ x ( θ ) = J θ F (¯ x ( θ ) , θ ), and the result follows by continuity of the derivatives. □

Example 2 (Superlinearly convergent algorithm) Assume that F is C 1 and for each ρ > 0, there is R > 0 such that ‖ F θ ( x ) -¯ x ( θ ) ‖ ≤ ρ ‖ x -¯ x ( θ ) ‖ for all x, θ such that ‖ x -¯ x ( θ ) ‖ ≤ R . Then Corollary 2 applies as for any v

2.3.2 Quadratically convergent algorithms

Under additional quantitative assumptions, it is possible to obtain more precise convergence estimates similar to those obtained for implicit differentiation, see Remark 2.

Corollary 3 Let F be as in Assumption 2 such that x ↦→ J ( x,θ ) F ( x, θ ) (joint jacobian in ( x, θ ) ) is L J Lipschitz (in operator norm). Then, the recursion is asymptotically quadratically convergent and for each k ≥ 1 ,

0

2

4

6

8

Iteration

10

9

10

6

10

3

10 0

Error

One-step diff

Forward diff

Dist optimum

Implicit diff

Proof : Following the same argument as in the proof of Corollary 2, we have

As for the quadratic convergence, we may assume that ¯ x ( θ ) = 0 and drop the θ variable to simplify notations. We have F (0) = 0 and for all x ,

Thus L J / 2 ‖ x k +1 ‖ ≤ [ L J / 2 ‖ x k ‖ ] 2 , and asymptotic quadratic convergence follows. □

Example 3 (Newton's algorithm) Assume that f : R n × R m → R is C 3 with Lipschitz derivatives, and for each θ , x ↦→ f ( x, θ ) is µ -strongly convex. Then Newton's algorithm with backtracking line search satisfies the hypotheses of Corollary 3, see [15, Sec. 9.5.3]. Indeed, it takes unit steps after a finite number of iterations, denoting by ¯ x ( θ ) the unique solution to ∇ x f ( x, θ ) = 0, for all x locally around ¯ x ( θ )

We have, ∇ 2 xx f ( x, θ ) F ( x, θ ) = ∇ 2 xx f ( x, θ ) x -∇ x f ( x, θ ), differentiating using tensor notations,

so that J x F (¯ x ( θ ) , θ ) = 0 and Lipschitz continuity of the derivatives of f implies Lipschitz continuity of J x F ( x, θ ) using the fact that ∇ 2 xx f ( x, θ ) /followsequal µ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

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, θ ) where F is as in Assumption 1. The problem may then be rewritten as

see illustrations in Remark 1. Gradient descent (or hyper-gradient) on (9) using our one-step estimator in (4) consists in the following recursion

where α > 0 is a step size parameter. Note that the quantity J OS x k ( θ ) T ∇ 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 ω > 0 be the multiplicative overhead of gradient evaluation, in typical applications, ω ≤ 5 (cheap gradient principle [6]). The time cost of evaluating the Jacobian of F is nω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ωC F m (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 reverse-mode of AD has time complexity kωC F (cheap gradient principle on F k θ ) and requires to store k vectors of size n . Implicit differentiation requires one full Jacobian ( ωC F n ) and solution of one linear system of size n × n , that is roughly n 3 . Finally, one-step differentiation is given by only differentiating a single step of the algorithm at cost ω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 ≥ n 2 , then reverse AD has a computational advantage if k ≤ 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 = -[ ∇ 2 f ] -1 ∇ f , implicit differentiation should be applied to the gradient mapping ∇ f = 0, not F . In typical application C ∇ 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 ∇ 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 ρ , the contraction factor, and distance to the solution for the inner algorithm. The proof is given in Appendix A.

Corollary 4 Let F : R n × R m → 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 ∇ Lipschitz gradient. Then,

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 : R m → R is C 1 with L Lipschitz gradient and lower bounded by h ∗ . Assume that for some /epsilon1 > 0 , for all l ∈ N , ∥ ∥ θ l +1 -θ l + 1 L ∇ h ( θ l ) ∥ ∥ ≤ /epsilon1 L . Then, for all K ∈ N , K ≥ 1 , we have

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 ≥ 2 and assume the following

sup θ ‖ x 0 ( θ ) -F θ ( x 0 ( θ )) ‖ ≤ M , for some M > 0 .

F is L F Lipschitz and J ( x,θ ) F is L J Lipschitz jointly in operator norm.

g is C 1 , l g Lipschitz with l ∇ Lipschitz gradient.

g ◦ ¯ x is lower bounded by g ∗ .

Then setting /epsilon1 = ρ 1 -ρ ( L F l g +( L J l g + L F l ∇ ) Mρ k -2 ) , for all K ∈ N ,

The level of approximate criticality is the sum of a term proportional to ρ 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 ρL F l g / (1 -ρ ) which is small if ρ is close to 0 ( e.g. superlinear algorithms).

4 Numerical experiments

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 ∈ R N × n be a design matrix, the first column being made of 1s to model an intercept. Rows are denoted by ( a 1 , . . . , a N ). Let x ∈ R n and y ∈ {-1 , 1 } N . We consider the regularized logistic regression problem

Interior point for QP

Newton for logistic regression

0

500

1000

1500

0

500

1000

1500

2000

2

4

6

0.0

2.5

5.0

7.5

10.0

12.5

Total number of parameters (1e3)

Time (s)

Gradient

Autodiff

Implicit

None

One step

where /lscript is the logistic loss, /lscript : t ↦→ log(1 + exp( -t )), λ > 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 ( θ ) the solution to problem (11) and try to evaluate the gradient of θ ↦→ ‖ x ( θ ) ‖ 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):

where Q ∈ R n × n is positive definite, A ∈ R m × n and G ∈ R p × n are matrices, q ∈ R n , θ ∈ R m and h ∈ R p are vectors. We consider x ( θ ) the solution of problem (12) as a function of θ , 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 θ x ( θ ) using implicit differentiation. More details on this experiment are given in Appendix B.

We consider evaluating the gradient of the function θ ↦→‖ x ( θ ) ‖ 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

autodiff (Jacobian)

autodiff (iter.)

1-step (Jacobian)

1-step (iter.)

implicit diff (Jacobian)

implicit diff (iter.)

100

suboptimality

101

10-1

￥

10-2

10-1

Relatives

10-3

水

10-4

10-3

10-5

×

0

50

100

150

200

250

300

350

400

0

50

100

150

200

250

300

350

400

Time (s)

Time (s)

100

Relative suboptimality

10-1

10-2

10-3

X

10-4

10-5

￥

10-6

10-7

10-8

X

0

10

20

30

40

50

60

70

0

10

20

30

40

50

60

70

Time (s)

Time (s)

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 ∈ R N × n , y ∈ R N , λ > 0 and a vector of weights θ ∈ R N :

We solve this problem using gradient descent with adequate step-size F ( x, θ ) = x -α ∇ f θ ( x ) with x 0 ( θ ) = 0, and we consider the K -step truncated Jacobian propagation ˜ F = F K θ with K = 1 /κ where κ 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 ( θ ) -¯ x ( θ ) ‖ 2 for the iterates, and ‖ J θ x k ( θ ) -J θ ¯ x ( θ ) ‖ 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 allows 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].

5 Conclusion

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.

References

Martín 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.

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.

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.

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.

Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems , 32, 2019.

Walter Baur and Volker Strassen. The complexity of partial derivatives. Theoretical computer science , 22(3):317-330, 1983.

Amir Beck. First-order methods in optimization . SIAM, 2017.

Thomas Beck. Automatic differentiation of iterative processes. Journal of Computational and Applied Mathematics , 50(1-3):109-118, 1994.

Yoshua Bengio. Gradient-Based Optimization of Hyperparameters. Neural Computation , 12(8):1889-1900, 2000.

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.

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.

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.

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.

Jerome Bolte, Edouard Pauwels, and Samuel Vaiter. Automatic differentiation of nonsmooth iterative algorithms. In Advances in Neural Information Processing Systems , 2022.

Stephen Boyd and Lieven Vandenberghe. Convex Optimization . Cambridge University Press, 2004.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye WandermanMilne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

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.

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.

Matthias J Ehrhardt and Lindon Roberts. Analyzing inexact hypergradients for bilevel learning. arXiv preprint arXiv:2301.04764 , 2023.

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.

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.

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.

Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is attention better than matrix decomposition? In ICLR , 2021.

Zhengyang Geng, Xin-Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. Advances in Neural Information Processing Systems , 34:2424724260, 2021.

Jean Charles Gilbert. Automatic differentiation and iterative processes. Optimization methods and software , 1(1):13-21, 1992.

Andreas Griewank and Christele Faure. Piggyback differentiation and optimization. In Large-scale PDE-constrained optimization , pages 148-164. Springer, 2003.

Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

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.

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.

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.

Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. In ICLR , 2019.

Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In AISTATS , pages 1540-1552, 2020.

Giulia Luise, Alessandro Rudi, Massimiliano Pontil, and Carlo Ciliberto. Differential properties of sinkhorn approximation for learning with wasserstein distance. In NeurIPS , volume 31, 2018.

Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani Raiko. Scalable gradientbased tuning of continuous regularization hyperparameters. In ICML , pages 2952-2960, 2016.

Sheheryar Mehmood and Peter Ochs. Automatic differentiation of some first-order methods in parametric optimization. In International Conference on Artificial Intelligence and Statistics , pages 1584-1594. PMLR, 2020.

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 in Neural Information Processing Systems 32 , pages 8024-8035, 2019.

Edouard Pauwels and Samuel Vaiter. The derivatives of sinkhorn-knopp converge. SIAM J Optim , 33(3):1494-1517, 2023.

Boris Polyak. Introduction to optimization. In Optimization Software, Publications Division . Citeseer, 1987.

Zaccharie Ramzi, Florian Mannel, Shaojie Bai, Jean-Luc 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.

Subham Sekhar Sahoo, Anselm Paulus, Marin Vlastelica, Vít Musil, Volodymyr Kuleshov, and Georg Martius. Backpropagation through combinatorial algorithms: Identity with projection works. In The Eleventh International Conference on Learning Representations , 2023.

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.

Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated back-propagation for bilevel optimization. In AISTATS , pages 1723-1732. PMLR, 2019.

Lieven Vandenberghe. The cvxopt linear and quadratic cone program solvers. Online: http://cvxopt. org/documentation/coneprog. pdf , 2010.

R. E. Wengert. A simple automatic derivative evaluation program. Communications of the ACM , 7(8):463-464, August 1964.

Matt Johnson Zico Kolter, David Duvenaud. Deep implicit layers - neural odes, deep equilibirum models, and beyond, 2020. NeurIPS Tutorial.

A Proof of Section 3

Proof of Corollary 4: We have ∇ θ ( g ◦ ¯ x )( θ ) = J θ ¯ x ( θ ) T ∇ g (¯ x ( θ )), and

The result follows from the triangular inequality combined with Corollary 1 and the following

and

Proof of Lemma 3: We have, using the 'descent lemma' (see [7]), for all l ∈ 0 . . . K ,

Summing for l = 0 , . . . , K , we have

and we deduce, by using concavity of the square root, that

Lemma 4 Let A : R d → R n × m and B : R d → R m × p , be M A and M B bounded and L A and L B Lipschitz respectively, in operator norm, then A · B is M A L B + L A M B Lipschitz in operator norm.

Proof : We have for all x, y ∈ R d

□

Lemma 5 The function M ↦→ ( I -M ) -1 is (1 -ρ ) -2 Lipschitz in operator norm on the set of square matrices M such that ‖ M ‖ op ≤ ρ < 1 .

Proof : For any two square matrices A,B with operator norm bounded by ρ < 1, I -A and I -B are invertible and the operator norm of the inverse is at most (1 -ρ ) -1 (see proof of

□

Proof of Corollary 5: Let us evaluate a Lipschitz constant of ∇ θ ( g ◦ ¯ x ). We have

This can be seen as a composition of a matrix function with ¯ x . Combining Lemma 4 and Lemma 5 with the fact that J x F and J θ F are both L J Lipschitz and L F bounded, the function ( x, θ ) ↦→ ( I -J x F ( x, θ )) -1 J θ f ( x, θ ) is L J L F (1 -ρ ) -2 + L J 1 -ρ = L J 1 -ρ ( L F 1 -ρ +1 ) Lipschitz.

Invoking Lemma 4 and using the fact that the operator norm of a vector is the euclidean norm, we have that the function ( x, θ ) ↦→ ( I -J x F ( x, θ )) -1 J θ f ( x, θ ) ∇ g ( x ) is L J 1 -ρ ( L F 1 -ρ +1 ) l g + l ∇ L F 1 -ρ Lipschitz.

From the implicit function theorem, ‖ J θ ¯ x ( θ ) ‖ op ≤ L F / (1 -ρ ) so that ¯ x is L F / (1 -ρ ) Lipschitz and ∇ θ ( g ◦ ¯ x ) is ( L J 1 -ρ ( L F 1 -ρ +1 ) l g + l ∇ L F 1 -ρ ) L F 1 -ρ Lipschitz so the choice of α is such that 1 /α a Lipschitz constant for the gradient.

Using Corollary 4, we are in the setting of Lemma 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 a 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 linear 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 a left incoming vector in equation (3). In this sense, we did implement VJP for implicit differentiation although we explicitly formed Jacobians because this was the most efficient.

Generation of random QPs: To generate the QP instances, the quantities n, m, p are varied on a logarithmic scale. We set Q = M T M where M is of size n × 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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).\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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\nQUESTIONS:\nIn [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)\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n-> 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.\n\nWEAKNESSES:\n-> 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\nQUESTIONS:\n-> 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.\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nAutomatic 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.\n\nWEAKNESSES:\nAs 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.\n\nQUESTIONS:\nline 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.\n\nLIMITATIONS:\nThe paper lacks a discussion of potential limitations or scenarios where the one-step differentiation method may not be suitable or may exhibit suboptimal performance.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\nL20: 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)\n\nLIMITATIONS:\nIn 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)", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\nPlease see above.\n\nLIMITATIONS:\nNo limitations are included.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-15"}