The Curse of Unrolling: Rate of Differentiating Through Optimization
Authors: Damien Scieur
Gauthier Gidel
Quentin Bertrand
Fabian Pedregosa
Venue: NeurIPS 2022 Accept (2022)
Keywords: implicit differentiation
unrolling
optimization
bi-level
meta-learning
sobolev

Abstract: Computing the Jacobian of the solution of an optimization problem is a central problem in machine learning, with applications in hyperparameter optimization, meta-learning, optimization as a layer, and dataset distillation, to name a few. Unrolled differentiation is a popular heuristic that approximates the solution using an iterative solver and differentiates it through the computational path. This work provides a non-asymptotic convergence-rate analysis of this approach on quadratic objectives for gradient descent and the Chebyshev method. We show that to ensure convergence of the Jacobian, we can either 1) choose a large learning rate leading to a fast asymptotic convergence but accept that the algorithm may have an arbitrarily long burn-in phase or 2) choose a smaller learning rate leading to an immediate but slower convergence. We refer to this phenomenon as the curse of unrolling.
Finally, we discuss open problems relative to this approach, such as deriving a practical update rule for the optimal unrolling strategy and making novel connections with the field of Sobolev orthogonal polynomials.

The Curse of Unrolling: Rate of Differentiating Through Optimization
Damien Scieur Samsung SAIL Montreal damien wea.ien.sceerrgmaaiiee com
Quentin Bertrand
Mila & Universtié de Montréal quentin. bertranddmil quebec
Gauthier Gidel
Fabian Pedregosa
Mila & Université de Montréal
Google Research
Canada CIFAR AI Chair
pedregosa@ggogle.. com
gidelgau@mila, www.uueemmaauuueeecccceo.o.
Abstract
Computing the Jacobian of the solution of an optimization problem central problem in machine learning, with applications in hyperparameter optimization, meta-learning, optimization as layer, and dataset distillation, name a few. Unrolled differentiation is popular heuristic that approximates the solution using iterative solver and differentiates it through the computational path. This work provides non-asymptotic convergence-rate analysis of this approach quadratic objectives for gradient descent and the Chebyshev method. We show that to ensure convergence of the Jacobian, we can either 1) choose large learning rate leading to a fast asymptotic convergence but accept that the algorithm may have an arbitrarily long burn-in phase or 2) choose smaller learning rate leading to immediate but slower convergence. We refer to this phenomenon as the curse of unrolling. Finally, discuss open problems relative this approach, such as deriving a practical update rule for the optimal unrolling strategy and making novel connections with the field of Sobolev orthogonal polynomials.
10
T.4
10
Jacobian suboptimality
suboptimality
Iterate
step size
10a
10-1-
100
150
200
50
100
150
200
# iterations
# iterations
Figure 1: The Curse Unrolling: Better function suboptimality does not imply better Jacobian
suboptimality. The speed of the function suboptimality (right) (0), 0) f(x. (0) 0) of gradient descent is not representative of the speed convergence of the Jacobian [dgx; ( (0) (left) early iterations. The Jacobian suboptimality exhibits burn-in phase where the suboptimality initially increases. The length of this burn-in phase depends on the step-size, with larger step-sizes having smaller function suboptimality but larger burn-in phase.
36th Conference Neural Information Processing Systems (NeurIPS 2022).

---

Introduction
Let (0) be a function defined implicitly as the solution to an optimization problem,
2z (0) arg min f(x 0). æER4
Implicitly defined functions of this form appear in different areas of machine learning, such as reinforcement learning (Pfau al., 2016; Du et al., 2017), generative adversarial networks (Metz et al., 201 hyper-parameter optimization (Bengio, 2000; Pedregosa, 2016; Franceschi et al., 2017 Lorraine al., 2020; Bertrand et al., 2020), meta-learning (Franceschi al., 2018; Rajeswaran et al., 2019), deep equilibrium models, (Bai et al., 2019) or optimization as a layer (Kim et al., 2017: Amos al., 2017; Wang al., 2019), name few. The main computational burden of using implicit functions in a machine learning pipeline is that the Jacobian computation dgx,(0) is challenging: since the implicit function x. (0) does not usually admit an explicit formula, classical automatic differentiation techniques cannot be applied directly.
Two main approaches have emerged to compute the Jacobian of implicit functions: implicit differenti- ation and unrolled differentiation. This paper focuses on unrolled differentiation while recent surveys on implicit differentiation (Duvenaud al., 2020; Blondel al., 2021).
Unrolled differentiation, also known as iterative differentiation, starts by approximating the implicit function (.).)) (+) by the output of an iterative algorithm, which we denote æ; (-), and then differentiates through the algorithm's 's computational path (Wengert, 1964; Domke, 2012; Deledalle et al., 2014;
Franceschi et al., 2017b; Shaban et al., 2019).
Contributions. We analyze the convergence of the unrolled Jacobian by establishing worst-case bounds on the Jacobian suboptimality |0xz (0) - (0)( for different methods, more precisely:
1. We provide a general framework for analyzing unrolled differentiation on quadratic objectives for any gradient-based method (Theorem 1). For gradient descent and the Chebyshev iterative method, we derive closed-form worst-case convergence rates in Corollary 1 and Theorem 3.
We identify the "curse of unrolling" as a consequence this analysis: A fast asymptotic rate inevitably leads to a condition number-long burn-in phase where the Jacobian suboptimality increases. While it is possible to reduce the length, or the peak, of this burn-in phase, this comes at the cost of slower asymptotic rate (Figure 4).
3. Finally, we describe novel approach to mitigate the curse of unrolling, motivated by the theory of Sobolev orthogonal polynomials (Theorem 4).
Related work The analysis of unrolling was pioneered in the work Gilbert (1992), who showed the asymptotic convergence of this procedure for a class of optimization methods that includes gradient descent and Newton's method. These results have been recently extended by Ablin al. (2020) and Grazzi et al. (2020), where they develop a complexity analysis for non non- non-quudratic functions. The rate they obtain are valid only for monotone optimization algorithms. such as gradient descent with small step size. We note that Grazzi et al. (2020) developed descent that matches our Theorem 2, and provided plots where one can appreciate the burn-in phase, although they did not discuss this behaviour nor the trade-off between the length of this phase and the step-size. Compared to these two papers, we instead focus on the more restrictive quadratic optimization setting. Thanks to this, obtain tight rates for larger class of functions, including non-monotone algorithms such as gradient descent with a large learning rate and the Chebyshev method. Furthermore, we also derive novel accelerated variants for unrolling (84.3).
asymsttmt rate for gradient
2 Preliminaries and Notations
In this paper, we consider an objective function f parametrized by two variables æ Rd and O E R We are interested the derivative of optimization problem solutions:
Goal: approximate Jacobian dx, (0), where (0) = arg min ((x 0).
(OPT)
wERd

---

We also assume 2. (0) E R is the unique minimizer of f(z, 0). for some fixed value of 0. In particular, we will describe the rate of convergence of [0)xx(() 0x, (0) in the specific case where I is quadratic function in its first argument, and æ; is generated by a first-order method.
Notations. In this paper, we use upper-case letter for polynomials (P; Q). bold lower-case for vectors (x, 6), and bold upper-case for matrices (H). We write Pt the set of polynomials of degree at least We distinguish V fxx 0). that refers to the gradient of the function f in its first argument, and fa 0) is the partial derivative of w.r.t. its second argument, evaluated at (æe, 0) (if there is no ambiguity, we write 0 instead of 0g). Similarly, 0z the Jacobian of the vector- valuud function zx() evaluated at t and (-) the derivative of the polynomial P. The Jacobian HH (0) tensor of size l x p p. We'll denote its tensor multiplication by matrix E Rxx by H (0)Q. with the understanding that this denotes the multiplication along the first axis, that is, the resulting tensor characterized by (OH(0)Q)
(0) for 1 S S k. Finally, we denote by the strong convexity constant of the objective function f. by I its smoothness constant, and by e/ its inverse condition number.
2.1 Problem Setting and Main Assumptions
Throughout the paper, we make the following three assumptions. The first one assumes the problem is quadratic. The second one is more technical and assumes the Hessian commutes with its derivative respectively, which simplifies considerably the formulas. As we'll discuss in the Experiments section, we believe that some of these assumptions could potentially be relaxed. The third assumption restricts the class of algorithms to first-order methods.
Assumption (Quadratic objective). The function f a quadratic functtion in its first argument,
[xx 0) def y H x + b(0)T== x,
where fI 3 H()) = LI for < < L. We write x. (0) the minimizer of f w.r.t the first argument. Assumption (Commutativity of Jacobian). We assume that H(0) commutes with its Jacobian, in the sense that
OH (()) = (O) for S iish k. In the case in which 0 is a scalar (k- = 1), this condition amounts to the commutativity between matrices ÖH (H) = (O)HH
(2)
Importance. The previous assumption allows to have simpler expression for the Jacobian of H Notably, with this assumption the Jacobian of H can be expressed as
i [()) = 10H 0)
(3)
This assumption is verified for example for Ridge regression (see below). Although quite restrictive, empirical evidence (see Appendix A) suggest that this assumption could potentially be relaxed or even lifted entirely.
Example 1 (Ridge regression). Let us fix A E Rand & E Rd, RR", and let H()) = A A 0I, where 0 in this case is scalar. The ridge regression problem
f(x, 0) 2 (I1Ax 2112 0|| x 112)
satisfies Assumptions 1, 2, f(x, 0) is quadratic in x, and H (()) = (O))HH) (()..
In our last assumption we restrict ourselves to first-order method, widely used large-scale opti- mization. This includes methods like gradient descent or Polyak's heavy-ball.
Assumption 3 (First-order method). The iterates { are generated from a first-order method:
(0) E (0) + spann V f 0), ..., (0).)-(0))0)
2.2 Polynomials and First-Order Methods on Quadratics
The polynomial formalism has seen a revival in recent years, thanks to its simple and constructive analysis (Scieur al., 2020a; Pedregosa et al., 2020; Agarwal et al., 2021). It starts from a connection between optimization methods and polynomials that allows to cast the complexity analysis of optimization methods polynomial bounding problem.
3

---

2.2.1 Connection with Residual Polynomials
When minimizing quadratics, after t iterations, one can associate to any optimization method polyno- mial Pt of degree at most such that P.(O) = 1, i.e.,
P.(A) a.
In such case, the error at iteration t then can be expressed
( (0) = P. (HH aa0.--.,,.0 (0))
(4)
This polynomial (0) is called the residual polynomial and represents the output of evaluating the originally vellev polynomial P((.) at the matrix H.
Example the case of gradient descent, the update reads Bt+1 x²' = -I x²),
which yields the residual polynomial P.(H(O) = (I --H(()))'
2.2.2 Worst-Case Convergence Bound
From the above identity, one can quickly compute worst-case bound on the associated optimization method. Using the Cauchy-Schward inequality on (4), we obtain
| x. (0)1 S (H(()) 2
(0)ll-
We are interested in the performance of the first-order method class of quadratic functions (see Assumption 1). whose Hessian has bounded eigenvalues. Using the fact that the --rnm of matrix is equal to its largest singular value, the worst-case performance of the algorithm then reads
x. (0) (0)l < max (E(E,L)
[|
(0)]
(5)
Therefore, the worst-case convergence bound is a function of the polynomial associated with the first-order method, and depends on the bound over the eigenvalue of the Hessian (A [e, L)).
2.2.3 Expected Spectral Density and Average-Case Complexity
We recall the average-cast complexity framework (Pedregosa al., 2020; Paquette et al., 2022; Cunha et al., 2022), which provides a finer-grained convergence analysis than the worst-case. This framework is crucial in developing accelerated method for unrolled differentiation (Section 4).
Instead of considering the worst instance from a class of quadratic functions, average-case analysis considers that functions are drawn at random from the class. This means that, in Assumption 1, the matrix (0), the vector b(0) and the initialization (0) in Assumption are sampled from some (potentially unknown) probability distributions. Surprisingly, we do not require the knowledge of these distributions, instead, the quantity of interest is the expected spectral density p(A), defined as
(A) def EH(0) [( (A)],
(A)
def a I - (H)))
(6)
In Equation 6, A (0)) is the i-th eigenvalue of H (0), 8(.) is the Dirac's delta and ( (A) is the empirical spectral density (i.e., (H(() (A) = 1/d À is an eigenvalue of H()) and 0 otherwise).
Assuming (0 (0) x. (0) is independent of H (0), the average-case complexity of the first-order method associated to the polynomial P. as
E (0) x. (0)) E[|| (0) (0)1 / (2)) du(A)-
Here, the term in Pe is seoorith--reltted while the term in dy is related to the difficulty of the (distribution over the) problem class. opposed to the worst-case analysis, where only the worst value of the polynomial impacts the convergence rate, the average-case rate depends on the expected value of the squared polynomial P; over the whole distribution. Note that in the previous equation, the expectation is taken over problem instances and not over any stochasticity of the algorithm.
'This assumption can be removed in (Cunha al., 2022) the price of more complicated expected spectral density.

---

The Convergence Rate of differentiating through optimization
We now analyze the rate of convergence of gradient descent and Chebyshev algorithm (optimal on quadratics). We first introduce the master identity (Theorem 1), which draws link between how well a first-order methods estimates dx, (0), its associated residual polynomial Pt, and derivative P/.
Theorem (Master identity). Under Assumptions 1, 2, let æ (0) be the th iterate of a first-order method associated to the residual polynomial Pt. Then the Jabobian enror can be written as
dz. (0) - - = (P: - ( H()) (oxo( - -dx -0)((0))
+ (H (0) (Zn
(8)
The above identity involves the derivative of the residual polynomials and not only the residual polynomial, as was the case for minimization (4). This difference is crucial and will result in different rates for the Jacobian suboptimality than classical ones for objective or iterate suboptimality.
For conciseness, our bounds make use of the following shorthand notation
G def (( Fl
3.1 Worst-case Rates for Gradient Descent
We consider the fixed-step gradient descent algorithm,
( (0) (f())-1(0).0).
As mentioned in Example 2, the associated polynomial reads P: = (1 hA)². We can deduce its convergence rate after injecting the polynomial in the master identity from Theorem 1.
Theorem 2 (Jacobian Suboptimality Rate for Gradient Descent). Under Assumptions 1, 2, let x, (0) be the t iterate of gradient descent scheme with step size h 0. Then,
max XE[F.I]
(1- (1- exponential decrease
[1 (t 1)hA)
linear increuse
Discussion. The bound above is product of two terms: gence rate of gradient descent and decreases exponentially for h S T,=z- while the second term is increasing in l. This results in two distinct phases in training: an initial burn-in phase, where the second term dominates and the Jacobian suboptimality might be increasing, followed by an linear convergence phase where the exponential term dominates. This phenomenon can be seen empirically, see Figure 1. As an illustration, the next corollary exhibits explicit rates for gradient descent in the two special cases where h (short steps) and L+t (large steps, maximize asymptotic rate).
the first term, (1 - -_)-\)"-1 is the conver-
Corollary 1. For the step size h L,L, the rate of Theorem 2 reads
|(xx (x)(())] <
(1-k) (--=)²1
n(( 1)) (0) (0)(() +
G
L
exponential decrease
Assuming G O, the above bound is monotonically decreasing.
If instead we take the worst-case optimal (for minimization) step size h= 2/((LLL E), then we huve
21
[0xx((0) (0)]l <
2t 1
|00z (0) (0))(
G
L
exponential decrease
finear increase
Moreover, assuming G = 0, the maximum of the upper bound over can go up
[[0x;(0)= (0)]l 5 Ox-10 ( (H110xz((0)) dx. 0x. (0)]l)) at t t
5

---

Gradient Descent
Chebyshev method
10
10
x 1P-
x 10P
Jacobian suboptimality
Burn-in Phase
2 10P-
x 10
Burn-in Phase
Linear Convergence
Convergence
Lincar
1001
10-1
250
500
750
1000
1250
iterations
iterations
Figure 2: The Phases of Unrolling, Gradient descent and the Chebyshev method exhibit two distinct phases during unrolled differentiation: initial burn-in phase, where the Jacobian suboptimality increases, followed by a convergent phase with an asymptotic linear convergence. The maximum of the suboptimality is similar for both algorithms, but Chebyshev peaks sooner than gradient descent, after a number of iterations equal to (roughly) the square root of the one required by gradient descent. The distinct phases, as well their relative duration, are predicted by the theoretical worst-case bounds of Corollary and Theorem 3. Both plots are run on the same problem, a ridge regression loss on the breast-cancer dataset.
In this Corollary, linear growth in t. When the step size is small (h = 1/L), the linear rate is slightly slower than the rate of the larger step size (h 2/(e T.). However, the term in / is way smaller for the small step size. This makes big difference in the convergence: for h = 1/L, there no local increase, which is not the for E In the next Corollary we provide bound the step size h to guarantee monotone convergence of the Jacobian.
trade-off between the linear convergence rate (exponential in t) and the
Corollary 2. Assuming G O, the bound of Theorem 2 is monotonically decreasing for 2 if the step size h from Theorem 2 satisfies 0 < h < L
This bound contrasts with the condition on the step size gradient descent, which S 2/L, with an optimal value of h 2/(e + L) (Nesterov, 2004). This trade-off between asymptotic rate and length of the burn-in phase leads us to formulate:
The curse of unrolling
To ensure convergence of the Jacobian with gradient descent, we must either 1) accept that the algorithm has burn-in period proportional to the condition number 1/x, or choose a small step size that will slow down the algorithm's 's asymptotic convergence.
3.2 Worst-case Rates for the Chebyshev method
We now derive a convergence-rail analysis for the Chebyshev method, which achieves the best worst-case convergence rate for the minimization of quadratic function with bounded spectrum.
Chebyshev method and Chebyshev polynomials. We recall the properties of the Chebyshev method (see e.g. (d' Aspremont al., 2021, Section 2) for survey). As mentioned in $2.2.2, the rate of convergence of first-order method associated with the residual polynomial P. can be upper bounded by max r. [P))) Let G2 be the Chebyshev polynomial of the first kind of degree t, and define
C₂(A) (m((A)
m: Le, [0, m((A
2M-L-1 L-
(10)
(.((O)))
known property of Chebyshev polynomials is that the shifted and normalized Chebyshev polynomial C is the residual polynomial with smallest maximum value in the [E, T.) interval. This implies that the

---

Chebyshev method, which is the method associated with this polynomial, enjoys the best worst-case convergence bound on quadratic functions. Algorithmically speaking, the Chebyshev method reads,
( 2:-1 (0) h4 ( + (x (0. (0) 20)),
where h4 is the step size and my the momentum. Those parameters are time- time-varying and depend only on and L. The following Proposition shows the rate of convergence the Chebyshev method.
Theorem 3 (Jacobian Suboptimality Rate for Chebyshev Method). Under Assumptions 1,2, let def /R) RR), and (0) denote the thh iterate of the Chebyshev method. Then, we have
the following convergence rate
2t
(0)x,(0) - 0x)(())lF W
2r
[10zo(0) 2x. (0) +
G
L-
exponeential
quadratic
In short, the rate of the Chebyshev algorithm for unrolling in O(P2gt). Moreover, assuming G the maximum of the upper bound over t can go up to
[1xx; (0))(())= 3kx=0 ( (m) (0) dv. (0)l
1 2
Discussion. Despite being optimal for minimization, the rate of the Chebyshev method has an additional Ot²²) factor. Due to this term, the bound diverges at first, similarly to gradient descent with the optimal step size h = E4L. but sooner. This behavior is visible on Figure 2.
Accelerated Unrolling: How fast can we differentiate through optimization?
We now show to accelerate unrolled differentiation. We first derive lower bound on the Jacobian suboptimality and then propose method based on Sobolev orthogonal polynomials (Marcellán et al., 2015), which are extremal polynomials for a norm involving both the polynomial and its derivative.
4.1 Unrolling is at least as hard as optimization
Proposition 1. Let a, be the I-th iterate of a first-order method. Then, for all iterations and for all 0, there exists a quadratic function f that verifies Assumption such that G = and
2 (()x;(0)= - Ox. - 110xo(0)) (x)(
1 Vk & = 1+ VR
(11)
This result tells us that unrolling is least difficult as optimization. Indeed, the rate in (11) is known to be the lower bound on the accuracy for minimizing smooth and strongly convex function (Nemirovski, 1995). Moreover, although we sure if the lower bound is tight for all I, we have that when - xc, the rate of Chebyshev method matches the above rate.
4.2 Average-Case Accelerated Unrolling with Sobolev Polynomials
We describe an accelerated method for unrolling based on Sobolev polynomials. We first introduce the definition of the Sobolev scalar product for polynomials.
Definition The Sobolev scalar product (and its norm) for two polynomials P, Q and a density function μ is defined as
P(N)((A) du +u P'(N) ||PI| del (P.
(P. Q)
In the following we'll assume j is the expected spectral density associated with the current problem class and discuss in the next section some practical choices. Using this scalar product, we can compute a (loose) upper-bound for (0) Üx (f)ll)) and Prop. polynomial minimizing this bound.

---

Proposition 2. Assume that ||0
(0) (0)
(0) (.)ol I nl(āxu(0) - 0x. (0)lE.. Then, under
Assumption 1, 2 and we have the following bound for the average-case rate
E O)) 10x,
æ. (0)l S P H (0xa(0) - dr. (0)
Proposition 3. Let } be a sequence of orthogonal Sobolev polynomials, i.e., (S,. S2) 0 if i == and otherwise, normalized such that S,(0) = Then, the residual polynomial that minimizes the Sobolev norm can be constructed
P
arg min P0)=:((0)=1
(P, P)v
Σ o;Si,
where 04
and A.
Σ ai. 1=0
A, 1=0
|
Moreover, we have that
|| 1/A,.
Limited burn-in phase.
Using the algorithm associated with P with parameters and J4, we have Ox, - 0x, (0)l s 2 2-EH( (0x 2.)0l
This inequality follows directly from Proposition 2 and the optimality of P. for 11: we have that [e S. || (because P:-1 1 is a feasible solution for P:) and ||P:ll S 1 (because Pt = is feasible solution for any > 1). That is much better than the maximum bump of (0(1/k) from gradient descend (Theorem 1) and from Chebyshev (Theorem 3).
4.3 Gegenbaueur-Soboley Algorithm
In most practical scenarios does not have access to the expected spectral density μ.. Furthermore, the rates average-case accelerated algorithms have been shown to be robust with respect to distribution mismatch (Cunha et al., 2022). In these cases, can approximate the expected spectral density by some distribution that has the same support. A classical choice is the Gegenbauer parametric family indexed by I, which encompasses important distributions such the density associated with Chebyshev's polynomials or the uniform distribution:
2X L-
(a()= = (mm())), ((xx) = - 2²)0=2 and m [e, I.] - [0. 1). m(A)
(12)
--
We'll call the sequence of Soholev orthogonal polynomials for this distribution Gegenbaueur-Sooole polynomials. Although in general Sobolev orthogonal polynomials don't enjoy a three-term recur- rence as classical orthogonal polynomials do, for this class of polynomials it's possible to build a recurrence for S, involving only S.-2, Qt and Qr-2, where {Q:} is sequence of Gegenbaueur polynomials (Marcellán et al., 1994). Unfortunately, existing work on Gegenbaueur and Gegenbaueur- Sobolev polynomials considers the un-shifted distribution j., but not μt, and also doesn't consider the S,(0) = 1. After (painful) changes to shift and normalize the
residual normalization (0) polynomials, we obtain a three-stages algorithm, summarized in the next Theorem.
Theorem 4 (Accelerated Unrolling). Let P. be defined in Proposition where the Sobolev product is defined with the density function (12). Then, the optimization algorithm associated to P; reads
V1-1 h4 f +mmyyy11 91-2)
(13)
(1)
(2)
(3)
21 = #1-2 + ci Y/
y:-2
(14)
Ar-1
Dt A, 11
A, zc:
(15)
for some step size hy, momentum My parameters (1) (2) (3) that depend on a, L and n, and At, at are defined in Proposition 3, whose recurrence are detailed in Appendix D. Moreover when t + co, the recurrence simplifies into
y: U hV (Vy +m)m))+++... m(y (1-2)
(16)
ye + m( 91-2).
(17)
where 1-1x and h ( the momentum and step size of Polyak's Heavy Ball. [+v/
Moreover, as + X0, we have the same asymptotic linear convergence as the Chebyshev method, 00) 000 (
lim
< Vm.
12 ( S2.
8

---

10a
Gradient descent., h 1/L
Jacobian suboptimality
109
Gradient descent, = 2/ L)
Chebyshev method
10-2
Sobolev method
10-1
Heavy-ball
Sobolev (asymptotic)
10-6
50
100
# iterations
Figure 3: (Theoretical) Worst-case convergence rates for different algorithms, with / 0.5, L 10. = 1 and G = O. The upper bound of a method associated with the polynomials {Pt} is defined as maxxele.4] [A)], where is the iteration counter. The plot compares Gradient descent with large ( (14n)) ) and small ( step size, Chebyshev, Sobolev (with = I/(),, and their asymptotic variants. We recognize in those curves the peaks of gradient descent and Chebyshev from Figure 2.
The accelerated algorithm for unrolling is divided into three parts, First, (13) corresponds an algo- rithm whose associated polynomials are Gegenbaueur polynomials. This is expected, as Pedregosa al. (2020) identified that all average-case optimal methods take the form of gradient descent with momentum. Second. (14) builds the Sobolev polynomial that corresponds to a weighted average of y, Finally, (15) is the weighted average of Sobolev polynomials that builds 1 in Proposition 2.
The non-asymptotic algorithm is rather complicated to implement; see Appendix D. Moreover, requires bound the spectrum of H namely [e, L], and one also has choose associated expected spectral density (*A) (parametrized by a)) and the parameter Nevertheless, this is the method that achieves the best performance for problems that satisfy assumptions.
Surprisingly, the asymptotic version is extremely simple, as corresponds to a weighed average of Heavy-Ball iterates: the only required parameters are l and L. It means that asymptotically, the algorithm is universally optimal, i.c., it achieves the best performing rate as long as we can identify the bounds on the spectrum of H (0). Such universal properties have been identified previously in (Scieur et al., 2020b), who showed that all average-cast optimal algorithms converge to Polyak's ss momentum independently of the expected spectral density μ (up to mild assumptions). We have the same phenomenon here, but with the additional surprising (and counter-intuine feature that the asymptotic algorithm is also independent of
Experiments and Discussion
5.1 Experiments least squares objective
We compare multiple algorithms for estimating the Jacobian (OPT) of the solution of a ridge regression problem (Example (1)) for a fixed value of 0 10-3, Figure shows the objective and Jacobian suboptimality on a ridge regression problem with the breast-cancer" as underlying dataset. Figure 4 shows the Jacobian suboptimality as function of the number of iterations, on both the breast-cancer and bodyfat dataset, and for synthetic dataset (where H (0) is generated as A A. where each entry in A is generated from standard Gaussian distribution). Appendix B contains further details and experiments on a logistic regression objective.
We observe the early suboptimality increase of Gradient descent and Chebyshev algorithm as predicted by Theorem and Theorem 3. Compared to Figure 3, that showed the theoretical rates, see that there's a remarkable agreement between theory and practice, as both the early increase, the asymptotic rate and the ordering of the methods matches the theoretical prediction. We also see that Sobolev is the best performing algorithm in practice, as it avoids the early increase while matching the accelerated asymptotic rate of Chebyshev.
"nctpe://archive. ics. www. htt.:///w...laeeeeeeeeeeeeeeeeea (Diagoooo))) "http://iibstaa" cmu, edu/dataseta/

---

breast-cances dataset
bodyfat dataset
synthetic dataset
101
10
101
suboptimality
Jacobian
10""
10"1-
10
10-1-
10
10
10-14
10-10
10
0
100
2.0
306
400
500
0
100
2001
300
400
500
100
2001
300
400
500
# iterations
# iterations
#
iterations
Sobolev
Chebyshev
Gradient Descent
Figure 4: Empirical comparison of the Sobolev method introduced in $4.3 (with a and rj 1), the Chebyshev method and Gradient descent on different datasets. The Sobolev algorithm has the shortest burn-in phase, does not locally diverge and has an accelerated asymptotic rate of convergence.
5.2 Experiments on logistic regression objective
In this section provide some extra experiments non-quadratic objective. We choose the following regularized logistic regression objective
12
f(x,0) = Σ (AA æ. sign(b))
(18)
where yp is the binary logistic loss, A, b is the data, which generated both from synthetic dataset and the breast-cancer dataset as described in $5.
synthelic dalasel
breast-cancer dalaset
Jacobian suboptimality
10'
10
10
10"
10"
10"
60
100
150
200
0
50
100
160
200
# iterations
# iterations
Figure 5: Two-phase dynamics in logistic regression. The two-phase dynamics predicted by Corollary and Theorem empirically hold for logistic regression objective. This objective not covered by our theory since it would violate the quadratic assumption (Assumption 1).
The only significant difference with the least squares loss is the range of step-size values that exhibit the initial burn-in phase. While for the quadratic loss, these are step-sizes close to 2/(T. na), in the case of logistic regression, that achieves the fastest convergence rate. The featured two-phase curve was computed using the step-size with a fastest asymptotic rate, computed through grid-search on the step-size values.
a crude upper bound and so this step-size is not necessarily the one
Limitations. Our theoretical results are limited to first-order methods applied to quadratic functions. Many applications use first-order methods, the quadratic Assumption 1, well the commutativity Assumption 2, are somewhat restrictive. However, experiments on objectives violating the non- quadratic and non-commutative assumption (Appendix B and A) show that the two-phase dynamics empirically translate to more general objectives. The Sobolev algorithm developed in this paper, however, might not generalize well outside the scope of quadratics. Nevertheless, the development of this accelerated method for unrolling highlights that we can adapt the design of current optimization algorithms so that they might perform better for automatic differentiation.
10

---

The authors would like to thank Pierre Ablin, Riccardo Grazzi, Paul Vicol, Mathieu Blondel and the anonymous reviewers for feedback on this manuscript. QB would like to thank Samsung Electronics Co., Ldt. for funding this research.
Acknowledgements.
References
Ablin, Pierre, Gabriel Peyré, and Thomas Moreau (2020). "Super--fficiency of automatic differentia- tion for functions defined as minimum". . In: Proceedings of the 37 th International Conference on Machine Learning.
Agarwal, Naman, Surbhi Goel, and Cyril Zhang (2021). "Acceleration via Fractal Learning Rate Schedules". In: Proceedings of the 38th International Conference on Machine Learning. Ed. by Marina Meila and Tong Zhang. Proceedings of Machine Learning Research. PMLR.
Amos, Brandon and J Zico Kolter (2017). "Optnet: Differentiable optimization as a layer in neural networks". In: International Conference on Machine Learning. Bai, Shaojie, J Zico Kolter, and Vladlen Koltun (2019). "Deep equilibrium models". In: Advances in
Neural Information Processing Systems.
Bengio, Yoshua (2000). "Gradient-basee optimization of hyperparameters" In: Neural computation. Bertrand, Quentin et al. (2020). "Implicit differentiation of Lasso-type models for hyperparameter optimization". In: International Conference on Machine Learning. PMLR.
Blondel, Mathieu et al. (2021). "Efficient and Modular Implicit Differentiation". In: arXiv preprint
arXiv:2105. 183.
Cunha, Leonardo al. (2022). "Only tails matter: Average-Case Universality and Robustness in the
Convex Regime". In. d' Aspremont, Alexandre, Damien Scieur, and Adrien Taylor (2021). "Acceleration methods". In:
arXiv preprint arXiv:2 105.09555
Deledalle, Charles-Alban al. (2014). "Stein Unbiased GrAdient estimator of the Risk (SUGAR) for multiple parameter selection". In: SIAM Journal on Imaging Sciences. Domke, Justin (2012). "Generic methods for optimization- based modeling". In: Artificial Intelligence and Statistics.
Du, Simon S et al. (2017). "Stochastic variance reduction methods for valuation". In: International Conference on Machine Learning.
Duvenaud, David, Zico Kolter, and Matthew Johnson (2020). "Deep Implicit Layers Tutorial
Neural ODEs, Deep Equilibirum Models, and Beyond".. In: Neural Information Processing Systems Tutorial.
Franceschi, Luca et al. (2017a). "Forward and reverse gradient-based hyperparameter optimization".
In: International Conference on Machine Learning.
(2017b). "imp reverse tradient-based hyperparameter optimization". In: Proceedings of the 34th International Conference on Machine Learning- Volume 70. JMLR. org.
Franceschi, Luca et al. (2018). "Bilevel programming for hyperparameter optimization and meta- learning". In: arXiv preprint arXiv: 04910.
Gilbert, Jean Charles (1992). "Automatic differentiation and iterative processes". In: Optimization methods and software.
Grazzi, Riccardo al. (2020). "On the iteration complexity of hypergradient computation". In: International Conference on Machine Learning. PMLR.
Kim, Yoon al. (2017). "Structured attention networks". In: arXiv preprint arXiv: 1oXiv:/702... 7702.008877
Lorraine, Jonathan, Paul Vicol, and David Duvenaud (2020). "Optimizing millions of hyperparameters by implicit differentiation".. In: International Conference on Artificial Intelligence and Statistics. Marcellán, Francisco, Teresa E Pérez, and Miguel A Piñar (1994). Siggenbauer-Sobolle orthogonal polynomials" In: Nonlinear Numerical Methods and Rational Approximation II. Springer. Marcellán, Francisco and Yuan Xu (2015). "On Sobolev orthogonal polynomials". In: Expositiones Mathematicae.
Metz, Luke et al. (2016). "Unrolled generative adversarial networks". In: arXiv preprint /.ki://611.01622
Nemirovski, Arkadi (1995). "Information-baee" complexity of convex programming". In: Lecture Notes.
Nesterov, Yurii (2004). Introductory lectures convex optimization. Springer.
Paquette, Courtney et al. (2022). "Halting Time is predictable for large models: universality property and average-case analysis". In: Foundations of Computational Mathematics.
11

---

Pedregosa, Fabian (2016). "Hyperparameter optimization with approximate gradient". In: Interna- tional Conference on Machine Learning.
Pedregosa, Fabian and Damien Scieur (2020). "Average-case Acceleration Through Spectral Density
Estimation". In: Proceedings of the 37th International Conference on Machine Learning (ICML). Pfau, David and Oriol Vinyals (2016). "Connecting generative adversarial networks and actor-critic methods". In: arXiv preprint arXiv:1610.0194)
Rajeswaran, Aravind et al. (2019).** "Meta-learning with implicit gradients". In: Advances in Neural Information Processing Systems.
Scieur, Damien, Alexandre d' A'Apremont, and Francis Bach (2020a). "Regularized Nonlinear acceler- ation". In: Mathematical Progranming.
Scieur, Damien and Fabian Pedregosa (2020b). "Rage-Case Optimality of Polyak Momentum". In: Proceedings of the 37th International Conference on Machine Learning (ICML). Shaban, Amirreza et al. (2019). "TrMiack-pree ppropaattoootttttionn for bilevel optimization". In: The 22nd International Conference on Artificial Intelligence and Statistics.
Wang, Po- Wei et al. (2019). "Satnet: Bridging deep learning and logical reasoning using differen- tiable satisfiability solver". In: International Conference on Machine Learning.
Wengert, Robert Edwin (1964). "A simple automatic derivative evaluation program", In: Communica- tions of the ACM.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [No]
(b) Did you specify all the training details (e.g., data splits. hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [No]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] (Datasets are not sensitive)
5. If you used crowdsourcing conducted research with human subjects... [N/A]
13


Reviews:
[{"review_id": "V8D18OvoKGP", "reviewer": "Reviewer_9FSi", "summary": "This paper pointed out an interesting phenomenon in bi-level optimization: the tradeoff between the convergence rate and the length of \"burn-in\" phase. Mathematical descriptions and analysis are provided.\n\n----\nThe authors addressed my concerns and I updated my score.", "questions": "1. At the beginning of this paper, Fig. 1 shows the difference between the function suboptimality and Jacobian suboptimality. This phenomenon is really interesting. However, in the main text, there is only the upper bound for Jacobian suboptimality. Why is there no bounds for the function gap? I believe that this paper will be better if the function gap bound can be provided and it matches Fig. 1.\n2. Through this paper, only quadratic functions are considered both in the theoretical and experimental parts. Is it possible to provide some experiments to show that the phenomenon \"curse of unrolling\" exists not only with quadratic functions but also with more general functions? At least some references on this phenomenon should be provided. If such a phenomenon only happens on quadratic functions, the significance of this paper will be weakened.\n3. Line 182: typo: \"Figure ??\" Such typo also occurs in the caption of Figure 3.\n4. Figure 4: Why does gradient descent on \"synthetic dataset\" diverge?", "limitations": "My main concerns are Questions 1 & 2 in the previous section, some potential suggestions are provided following the questions.\nThis paper is purely theoretical and I think it has no potential negative societal impact.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The \"curse of unrolling\" and its mathematical descriptions are really interesting. Thus, I think this paper is very insightful. But the completeness and presentation have room to improve. I listed some questions below and I'm open to updating the grade if the authors are able to fix my concerns.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "NfdFc6yQ-nM", "reviewer": "Reviewer_eLXs", "summary": "This paper studies the use of unrolling as a way to differentiate through an optimization problem. Specifically, it focuses on the estimation of the Jacobian (of the solution w.r.t. the trainable parameters of the network) when the problem is a quadratic optimization and the algorithm is a first-order method. The main contribution is the analysis which shows ``the curse of unrolling'' and on top of it, this paper proposes a scheme to accelerate the unrolled differentiation based on Sobolev orthogonal polynomials. The experiments show curves that match well with the proposed analysis.", "questions": "Please see weaknesses.", "limitations": "The authors have discussed the limitations in detail and claim to have no potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The topic is both theoretically and practically important considering that the unrolling scheme has been utilized by a large family of deep learning works.\n2. The results (theorem1-3) are novel and can explain the practical observation. Discussions are given to explain how the results can lead to certain pattern in the plot.\n3. The proposed accelerated scheme for unrolling has a limited ``burn-in'' phase which is a useful property.\n\nWeaknesses:\n1. My largest concern is that for deep pipelines, both the Jacobian $\\partial x_t - \\partial x_*$ and the final solution $||x_t-x_*||$ matter at the same time. With only the bound for the former one, the results cannot be directly used to guide the practical use of unrolling. Especially for the proposed accelerated scheme, it is not clear how this modified scheme will influence $||x_t-x_*||$ .\n2. The datasets used in the experiment are kind of toy-size by modern standards. Since the topic is the use of unrolling in deep learning pipelines, it would be more convincing if more serious experiments were included.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Ic-Iy7qoM_", "reviewer": "Reviewer_ZGFf", "summary": "The authors propose a framework for convergence analysis of unrolled differentiation of quadratic objectives. To that end, they use a connection between optimization methods and residual polynomials which allows one to establish convergence bounds by evaluating appropriate polynomials on the spectrum of the hessian. From their theoretical results it can be seen that the convergence behaviour of the Jacobian of the argmin and the optimization objective behave differently during unrolling. A larger step size that yields fastest convergence of the optimization objective induces in turn a longer burn-in phase of the Jacobian sequence. Smaller step sizes can mitigate this burn-in phase while inducing slower asymptotic convergence. This tradeoff is called the curse of unrolling by the authors. Finally the authors propose an accelerated unrolling procedure for quadratic objectives that is motivated by the theory of Sobolev Polynomials and is shown to outperform reference methods (GD and Chebyshew) on different datasets.", "questions": "Some typos I found:\n\n17: function(s)\n\n27: differentiation(m)\n\n30: of (an) iterative algorithm\n\n32: What does \"mode differentiation\" stand for here?\n\n49: algorithm(s)\n\n189: iteration(s)\n\n194: (that) that\n\n217: (we) don't have\n\nTop of page 8: Figure ??\n\n260: there is no Figure 5\n\n280: should it be \"formalism does (not) apply\"?\n\nOne thing I was wondering about is whether the analysis still works when $n > 1$, i.e., when $\\theta$ is not a scalar but a vector. How do I need to read for example (2) in that case? What is the dimension of $\\partial H(\\theta)$ then?", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "From my perspective, the presented results are novel and technically correct. Apart from a few typos and missing references that can likely be fixed, I also find that the paper is well-structured, clearly written, and has some comprehensive figures. I personally think that the theoretical underpinning of distinct convergence phases during unrolling is a very interesting result. Also I would say that the methodology and the derived accelerated method are strengths of the paper. Clear limitations that the authors also mention in their discussion are that the results only apply to first-order methods and quadratic objectives. With that said, the presented results are in my opinion still interesting, although maybe less relevant to a broader audience.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept