The Curse of Unrolling: Rate of Differentiating Through Optimization
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
[Figure]
J a c o b i a n s ub o p t i m a li t y
# iterations
[Figure]
I t e r a t e s ub o p t
i m a li t y [Figure] L √ 2 L 2 L + ‘ s t e p s i
z e
Introduction
Let x ⋆ ( θ ) be a function defined implicitly as the solution to an optimization problem,
[Equation]
Implicitly defined functions of this form appear in different areas of machine learning, such as reinforcement learning (Pfau et al., 2016; Du et al., 2017), generative adversarial networks (Metz et al., 2016), hyper-parameter optimization (Bengio, 2000; Pedregosa, 2016; Franceschi et al., 2017a;
Lorraine et al., 2020;
Bertrand et al., 2020), meta-learning (Franceschi et al., 2018; Rajeswaran et al., 2019), deep equilibrium models, (Bai et al., 2019) or optimization as a layer (Kim et al., 2017; Amos et al., 2017; Wang et al., 2019), to name a few.
The main computational burden of using implicit functions in a machine learning pipeline is that the Jacobian computation ∂ θ x ⋆ ( θ ) is challenging: since the implicit function x ⋆ ( θ ) 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 differentiation and unrolled differentiation.
This paper focuses on unrolled differentiation while recent surveys on implicit differentiation are (Duvenaud et al., 2020; Blondel et al., 2021).
Unrolled differentiation, also known as iterative differentiation, starts by approximating the implicit function x ⋆ ( · ) by the output of an iterative algorithm, which we denote x t ( · ) , and then differentiates through the algorithm’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 ∥ ∂ x t ( θ ) − ∂ x ⋆ ( θ ) ∥ for different methods, more precisely:
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. 2.
We identify the “curse of unrolling” as a consequence of 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 a slower asymptotic rate (Figure 4).
Finally, we describe a 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 of 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 et al.
(2020), where they develop a complexity analysis for non-quadratic functions.
The rate they obtain are valid only for monotone optimization algorithms, such as gradient descent with small step size.
In our case, we instead focus on the more restrictive quadratic optimization setting but obtain rates for a more general class of algorithms, including non-monotone and accelerated algorithms such as GD with large step size, or the Chebyshev method.
Thanks to the more restrictive quadratic setting, we’re able not only to provide a tighter analysis, but also derive accelerated variants.
Finally, none of the discussed works observes or provides explanation for the burn-in phase that we discuss in Section 3.
Preliminaries and Notations
In this paper, we consider an objective function f parametrized by two variables x ∈ R d and θ ∈ R k .
We are interested in the derivative of optimization problem solutions:
[Equation]
We also assume x ⋆ ( θ ) ∈ R
d is the unique minimizer of f ( x , θ ) , for some fixed value of θ .
In particular, we will describe the rate of convergence of ∥ ∂ x t ( θ ) − ∂ x ⋆ ( θ ) ∥ in the specific case where f is a quadratic function in its first argument, and x t 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 , b ), and bold upper-case for matrices ( H ).
We write P t the set of polynomials of degree at least t .
We distinguish ∇ f ( x , θ ) , that refers to the gradient of the function f in its first argument, and ∂ θ f ( x , θ ) is the partial derivative of f w.r.t.
its second argument, evaluated at ( x , θ ) (if there is no ambiguity, we write ∂ instead of ∂ θ ).
Similarly, ∂ x ( θ ) is the Jacobian of the vector-valued function x ( · ) evaluated at θ and P ′ ( · ) the derivative of the polynomial P .
The Jacobian ∂ H ( θ ) a tensor of size k × p × p .
We’ll denote its tensor multiplication by a matrix Q ∈ R p × q by ∂ H ( θ ) Q , with the understanding that this denotes the multiplication along the first axis, that is, the resulting tensor is characterized by [ ∂ H ( θ ) Q ] i =
∂
H ( θ ) i Q for 1 ≤ i ≤ k .
Finally, we denote by ℓ the strong convexity constant of the objective function f , by L its smoothness constant, and by κ = ℓ/L 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 1 (Quadratic objective) .
The function f is a quadratic function in its first argument,
[Equation]
where ℓ I ⪯ H ( θ )
⪯
L I for 0 < ℓ < L .
We write x ⋆ ( θ ) the minimizer of f w.r.t.
the first argument.
Assumption 2 (Commutativity of Jacobian) .
We assume that H ( θ ) commutes with its Jacobian, in the sense that
[Equation]
In the case in which θ is a scalar ( k = 1 ), this condition amounts to the commutativity between matrices ∂ H ( θ ) H ( θ ) =
H ( θ ) ∂
H ( θ )
Importance.
The previous assumption allows to have simpler expression for the Jacobian of H t .
Notably, with this assumption the Jacobian of H can be expressed as
[Equation]
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 ∈ R n × d , ¯ x ∈ R d , b ∈ R n , and let H ( θ ) =
A ⊤ A + θ I , where θ in this case is a scalar.
The ridge regression problem
[Equation]
satisfies Assumptions 1, 2, as f ( x , θ ) is quadratic in x , and ∂ H ( θ ) H ( θ )
=
H ( θ ) ∂
H ( θ ) =
H ( θ ) .
In our last assumption we restrict ourselves to first-order method , widely used in large-scale optimization.
This includes methods like gradient descent or Polyak’s heavy-ball.
Assumption 3 (First-order method) .
The iterates { x t } t =0 ... are generated from a first-order method:
[Equation]
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 et 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 as polynomial bounding problem.
2.2.1 Connection with Residual Polynomials
When minimizing quadratics, after t iterations, one can associate to any optimization method polynomial P t of degree at most t such that P t (0)
= 1 , i.e.,
[Equation]
In such a case, the error at iteration t then can be expressed as
[Equation]
This polynomial P t ( H ( θ )) is called the residual polynomial and represents the output of evaluating the originally real-valued polynomial P t ( · ) at the matrix H .
Example 2.
In the case of gradient descent, the update reads
[Equation]
which yields the residual polynomial P t
[Equation]
2.2.2 Worst-Case Convergence Bound
From the above identity, one can quickly compute a worst-case bound on the associated optimization method.
Using the Cauchy-Schwartz inequality on (4), we obtain
[Equation]
We are interested in the performance of the first-order method on a class of quadratic functions (see
Assumption 1), whose Hessian has bounded eigenvalues.
Using the fact that the ℓ 2 -norm of a matrix is equal to its largest singular value, the worst-case performance of the algorithm then reads
[Equation]
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 ( λ ∈ [ ℓ, L ] ).
2.2.3 Expected Spectral Density and Average-Case Complexity
We recall the average-case complexity framework (Pedregosa et 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 an 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 H ( θ ) , the vector b ( θ ) and the initialization x 0 ( θ ) in Assumption 3 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 µ ( λ ) , defined as
[Equation]
In Equation 6, λ i ( H ( θ )) is the i -th eigenvalue of H ( θ ) , δ ( · ) is the Dirac’s delta and
µ H ( θ ) ( λ ) is the empirical spectral density (i.e., µ H ( θ ) ( λ )
= 1 /d if λ is an eigenvalue of H ( θ ) and 0 otherwise).
Assuming x 0 ( θ ) − x ⋆ ( θ ) is independent of H ( θ ) , 1 the average-case complexity of the first-order method associated to the polynomial P t as
[Equation]
Here, the term in P t is algorithm-related, while the term in d µ is related to the difficulty of the (distribution over the) problem class.
As 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 t 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.
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 a link between how well a first-order methods estimates ∂ x ⋆ ( θ ) , its associated residual polynomial P t , and its derivative P ′ t .
Theorem 1 (Master identity) .
Under Assumptions 1, 2, 3, let x t ( θ ) be the t th iterate of a first-order method associated to the residual polynomial P t .
Then the Jabobian error can be written as
[Equation]
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
[Equation]
3.1 Worst-case Rates for Gradient Descent
We consider the fixed-step gradient descent algorithm,
[Equation]
As mentioned in Example 2, the associated polynomial reads P t =
(1 − hλ ) t .
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 t ( θ ) be the t th iterate of gradient descent scheme with step size h > 0 .
Then,
[Equation]
[Table]
[Equation]
Discussion.
The bound above is a product of two terms: the first term, (1 − hλ ) t − 1 , is the convergence rate of gradient descent and decreases exponentially for h ≤ 2 L + ℓ , while the second term is increasing in t .
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 = 1 L (short steps) and h = 2 L + ℓ (large steps, maximize asymptotic rate).
Corollary 1.
For the step size h = 1 /L , the rate of Theorem 2 reads
[Equation]
Assuming G = 0 , the above bound is monotonically decreasing.
If instead we take the worst-case optimal (for minimization) step size h = 2 / ( L + ℓ ) , then we have
[Equation]
Moreover, assuming G = 0 , the maximum of the upper bound over t can go up to
[Equation]
[Figure]
J a c o b i a n s ub o p t i m a li t y
[Figure]
u r n - i n P h a s e
L i n e a r
C o n v e r g e n ce
Gradient Descent
[Figure]
L i n e a r
C
[Figure]
n v
[Figure]
ce
Chebyshev method
In this Corollary, we see a trade-off between the linear convergence rate (exponential in t ) and the 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 / ( ℓ + L ) .
However, the term in t is way smaller for the small step size.
This makes a big difference in the convergence: for h = 1 /L , there is no local increase, which is not the case for h = 2 ℓ + L .
In the next Corollary we provide a bound on the step size h to guarantee a monotone convergence of the Jacobian.
Corollary 2.
Assuming G = 0 , the bound of Theorem 2 is monotonically decreasing for t ≥ 1 if the step size h from Theorem 2 satisfies 0 < h < √ 2 /L .
This bound contrasts with the condition on the step size of gradient descent, which is h ≤ 2 /L , with an optimal value of h = 2 / ( ℓ + L ) (Nesterov, 2004).
This trade-off between asymptotic rate and length of the burn-in phase leads us to formulate:
To ensure convergence of the Jacobian with gradient descent, we must either 1) accept that the algorithm has a burn-in period proportional to the condition number 1 /κ , or 2) choose a small step size that will slow down the algorithm’s asymptotic convergence.
The curse of unrolling 3.2 Worst-case Rates for the Chebyshev method
We now derive a convergence-rate analysis for the Chebyshev method, which achieves the best worst-case convergence rate for the minimization of a quadratic function with a bounded spectrum.
Chebyshev method and Chebyshev polynomials.
We recall the properties of the Chebyshev method (see e.g. (d’Aspremont et al., 2021, Section 2) for a survey).
As mentioned in §2.2.2, the rate of convergence of a first-order method associated with the residual polynomial P t can be upper bounded by max λ ∈ [ ℓ,L ] | P t ( λ ) | .
Let ˜ C t be the Chebyshev polynomial of the first kind of degree t , and define
[Equation]
A known property of Chebyshev polynomials is that the shifted and normalized Chebyshev polynomial
C t is the residual polynomial with smallest maximum value in the [ ℓ, L ] 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,
[Equation]
where h t is the step size and m t the momentum.
Those parameters are time-varying and depend only on ℓ and L .
The following Proposition shows the rate of convergence of the Chebyshev method.
Theorem 3 (Jacobian Suboptimality Rate for Chebyshev Method) .
Under Assumptions 1,2, let ξ def
=
(1 − √ κ ) / (1 + √ κ ) , and x t ( θ ) denote the t th iterate of the Chebyshev method.
Then, we have the following convergence rate
[Equation]
In short, the rate of the Chebyshev algorithm for unrolling in O ( t 2 ξ t ) .
Moreover, assuming G = 0 , the maximum of the upper bound over t can go up to
[Equation]
Discussion.
Despite being optimal for minimization, the rate of the Chebyshev method has an additional O ( t 2 ) factor.
Due to this term, the bound diverges at first, similarly to gradient descent with the optimal step size h = 2 ℓ + L , 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 a lower bound on the Jacobian suboptimality and then propose a 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 x t be the t -th iterate of a first-order method.
Then, for all iterations t and for all θ , there exists a quadratic function f that verifies Assumption 1 such that G = 0 , and
[Equation]
This result tells us that unrolling is at least as 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 are not sure if the lower bound is tight for all t , we have that when t → ∞ , the rate of Chebyshev method matches the above rate.
4.2 Average-Case Accelerated Unrolling with Sobolev Polynomials
We now describe an accelerated method for unrolling based on Sobolev polynomials.
We first introduce the definition of the Sobolev scalar product for polynomials.
Definition 1.
The Sobolev scalar product (and its norm) for two polynomials P, Q and a density function µ is defined as
[Equation]
In the following we’ll assume µ 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 ∥ ∂ x t ( θ ) − ∂ x ⋆ ( θ ) ∥ F and in Prop. 3, a polynomial minimizing this bound.
Proposition 2.
Assume that ∥ ∂ H ( θ )( x 0 ( θ ) − x ⋆ ( θ ))
∥ F ≤ η ∥ ∂ x 0 ( θ ) − ∂ x ⋆ ( θ ) ∥ F .
Then, under Assumption 1, 2 and 3, we have the following bound for the average-case rate
[Equation]
Proposition 3.
Let { S t } be a sequence of orthogonal Sobolev polynomials, i.e., ⟨ S i , S j ⟩ > [-
-]0 if i = j and 0 otherwise, normalized such that S i (0)
= 1 .
Then, the residual polynomial that minimizes the Sobolev norm can be constructed as
[Equation]
where a i = ∥ S t ∥ 2 η and A t = t (cid:88)
i = a i .
Moreover, we have that ∥ P ⋆t ∥ 2 η = 1 /A t .
Limited burn-in phase.
Using the algorithm associated with P ⋆ with parameters η and µ , we have
[Equation]
This inequality follows directly from Proposition 2 and the optimality of P t for ∥ · ∥ η : we have that ∥ P t ∥ η ≤ ∥ P t − 1 ∥ η (because P t − 1 is a feasible solution for P t ) and ∥ P t ∥ η ≤ 1 (because P t = 1 is a feasible solution for any t > 1 ).
That is much better than the maximum bump of ( O (1 /κ )
from gradient descend (Theorem 1) and from Chebyshev (Theorem 3).
4.3 Gegenbaueur-Sobolev Algorithm
In most practical scenarios one does not have access to the expected spectral density µ .
Furthermore, the rates of average-case accelerated algorithms have been shown to be robust with respect to distribution mismatch (Cunha et al., 2022).
In these cases, we can approximate the expected spectral density by some distribution that has the same support.
A classical choice is the Gegenbauer parametric family indexed by α ∈ R , which encompasses important distributions such as the density associated with Chebyshev’s polynomials or the uniform distribution:
[Equation]
We’ll call the sequence of Sobolev orthogonal polynomials for this distribution Gegenbaueur-Sobolev polynomials.
Although in general Sobolev orthogonal polynomials don’t enjoy a three-term recurrence as classical orthogonal polynomials do, for this class of polynomials it’s possible to build a recurrence for S t involving only S t − 2 , Q t and Q t − 2 , where { Q t } is sequence of Gegenbaueur polynomials (Marcellán et al., 1994).
Unfortunately, existing work on Gegenbaueur and GegenbaueurSobolev polynomials considers the un-shifted distribution ˜ µ , but not µ , and also doesn’t consider the residual normalization Q t (0)
= 1 or S t (0)
= 1 .
After (painful) changes to shift and normalize the polynomials, we obtain a three-stages algorithm, summarized in the next Theorem.
Theorem 4 (Accelerated Unrolling) .
Let P ⋆t be defined in Proposition 3, where the Sobolev product is defined with the density function µ (12) .
Then, the optimization algorithm associated to P ⋆t reads
[Equation]
[Equation] [Equation] (15) for some step size h t , momentum m t , parameters c (1) t , c (2) t , c (3) t that depend on α , ℓ , L and η , and
A t , a t are defined in Proposition 3, whose recurrence are detailed in Appendix D. Moreover, when t → ∞ , the recurrence simplifies into
[Equation]
where m = (cid:16) 1 −√ κ 1+ √ κ (cid:17) 2 and h = (cid:16) 2 √ ℓ + √ L (cid:17) 2 are the momentum and step size of Polyak’s Heavy Ball.
Moreover, as t → ∞ , we have the same asymptotic linear convergence as the Chebyshev method,
[Equation]
[Figure]
J a c o b i a n s ub o p t i m a li t y
Gradient descent.
, h = 1 /L
Gradient descent, h = 2 / ( ‘ + L )
Chebyshev method
Sobolev method
Heavy-ball
Sobolev (asymptotic)
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 to an algorithm whose associated polynomials are Gegenbaueur polynomials.
This is expected, as Pedregosa et 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 t .
Finally, (15) is the weighted average of Sobolev polynomials that builds P ⋆ in Proposition 2.
The non-asymptotic algorithm is rather complicated to implement; see Appendix D. Moreover, it requires a bound on the spectrum of H ( θ ) , namely [ ℓ, L ] , and one also has to choose an associated expected spectral density µ ( λ ) (parametrized by α ) and the parameter η .
Nevertheless, this is the method that achieves the best performance for problems that satisfy our assumptions.
Surprisingly, the asymptotic version is extremely simple, as it corresponds to a weighed average of Heavy-Ball iterates : the only required parameters are ℓ and L .
It means that asymptotically, the algorithm is universally optimal, i.e., it achieves the best performing rate as long as we can identify the bounds on the spectrum of H ( θ ) .
Such universal properties have been identified previously in (Scieur et al., 2020b), who showed that all average-case optimal algorithms converge to Polyak’s momentum independently of the expected spectral density µ (up to mild assumptions).
We have the same phenomenon here, but with the additional surprising (and counter-intuitive) feature that the asymptotic algorithm is also independent of η .
Experiments and Discussion 5.1 Experiments on 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 θ = 10 − 3 .
Figure 1 shows the objective and
Jacobian suboptimality on a ridge regression problem with the breast-cancer 2 as underlying dataset.
Figure 4 shows the Jacobian suboptimality as a function of the number of iterations, on both the breast-cancer and bodyfat 3 dataset, and for a synthetic dataset (where H ( θ ) is generated as A ⊤ A , where each entry in A is generated from a 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 2 and Theorem 3.
Compared to Figure 3, that showed the theoretical rates, we 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. [Figure] i [Figure]
i
[Figure]
5.2 Experiments on logistic regression objective
In this section we provide some extra experiments on a non-quadratic objective.
We choose the following regularized logistic regression objective
[Equation]
18) where φ is the binary logistic loss
[Figure]
A
[Figure]
b is the data
[Figure]
which we generated both from a synthetic dataset and the breast-cancer dataset as described in §
[Figure]
significant difference with the least squares loss is the range step-size values initial burn quadratic loss
-sizes close to 2 / ( L + µ ) , in the case of logistic regression, L is a crude upper bound and so this step-size is not necessarily the one that achieves the fastest convergence rate.
The featured two-phase curve was computed using the step-size with a fastest asymptotic rate, computed through a grid-search on the step-size values.
Limitations.
Our theoretical results are limited to first-order methods applied to quadratic functions.
Many applications use first-order methods, the quadratic Assumption 1, as well as the commutativity
Assumption 2, are somewhat restrictive.
However, experiments on objectives violating the nonquadratic 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.
Acknowledgements.
The authors would like to thank Pierre Ablin, 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.
A On the Commutativity Assumption
We consider the problem
[Equation]
, which is a generalization of Example 1 for the matrix norm ∥ x ∥ 2 D with a diagonal matrix D .
Contrary to Example 1, the matrix D is not an identity matrix, but instead a diagonal matrix where the diagonal entries are generated from a Chi-squared distribution.
In this case, Assumption 2 is no longer verified.
To investigate whether the two phases dynamics appear also on this class of problems, we repeat the same experiment as in Figure 2 with the above objective.
We plot the result here below, confirming the same dynamics of an initial Burn-in-Phase followed by a linear convergence phase observed in the initial experiment.
[Figure]
# iterations
[Figure]
J a c o b i a n s ub o p t i m a li t y
B u r n - i n
P h
a s e
L i n e a r
C o n v e r g e n ce
Gradient
[Figure]
Chebyshev method
We also reproduced the same setup as in Figure 4 with this matrix norm, obtaining again comparable results as in the commutative case.
This suggest that results regarding the two-phase dynamics could potentially be developed without Assumption 2, as we observe similar results as in Figure 4.
[Figure]
J a c o b i a n s ub o p t i m a li t y
[Figure]
B Experiments B.1 Further experimental details
Hyperparameters.
Initialization is always zero, x 0 = 0 , the regularization parameter θ in the ridge regression problem is always set to λ = 10 − 3 ∥ A ∥ 2 . κ
B REAST C ANCER  7.2 ×
BODYFAT  0.021 S YNTHETIC  0.
Train-test split.
For every dataset, we only use the train set, where the split is given by the libsvmtools 4 project.
Run-time.
Given the reduced size of these datasets, the script to compare all methods, which does a full unrolling for each iteration, runs in under 5 minutes running on CPU.
C Proofs C.1 Proof of Theorem
Theorem 1 (Master identity) .
Under Assumptions 1, 2, 3, let x t ( θ ) be the t th iterate of a first-order method associated to the residual polynomial P t .
Then the Jabobian error can be written as
[Equation]
Proof.
We differentiate both sides of (4) and use Assumption 2:
[Equation]
We now differentiate the equation b ( θ )
=
H ( θ ) x ⋆ ( θ ) w.r.t.
θ , ∂ b
[Equation]
) = ∂ H (
[Equation]
) x ⋆ ( θ ) + H ( θ ) ∂ x ⋆ ( θ ) .
We first substitute ∂ H ( θ ) x
⋆ ( θ ) by ∂ b ( θ ) − H ( θ ) ∂
x ⋆ ( θ ) .
After rearrangement, we finally get
[Equation]
It suffices to notice that the terms inside the square brackets are the cross-derivative of f :
[Equation]
C.2 Proof of Corollary
Corollary 2.
Assuming G = 0 , the bound of Theorem 2 is monotonically decreasing for t ≥ 1 if the step size h from Theorem 2 satisfies 0 < h < √ 2 /L .
Proof.
In this proof, we assume that t ≥ 1 .
Indeed, when t = 0 and t = 1 , the worst-case bound do not guarantee any progress over ∥ ∂ x 1 ( θ ) − ∂ x ⋆ ( θ ) ∥ F .
First, we notice that when hλ ≤ 1 (i.e., h ≤ 1 /L ) , we have that the rate from Theorem 2 is monotonically decreasing.
Indeed, the derivative over t gives
[Equation]
If the following condition is satisfied for all t ≥ 1 , the derivative is negative, and therefore the bound is monotonically decreasing:
[Equation]
This is always true since the right-hand side is negative, because hλ < 1 , and the left-hand side is always positive since t ≥ 1 .
We now assume that there exist some values of λ such that hλ > 1 .
For those values of hλ , the expression in Theorem 2 becomes
[Equation]
We now compute its maximum value.
First, we compute its derivative over t and solve d · d t = 0 .
We obtain the unique solution
[Equation]
This means there is only one maximum in the expression.
We now seek a value of hλ where the bound decrease monotonically for t > 1 , i.e.,
[Equation]
Since we know there is only one maximum, we compute hλ such that, in the worst case, ∥
[Equation]
x
[Equation]
We therefore have to solve
[Equation]
In particular, this means that if hλ < √ 2 , the bound decreases monotonically for t = 1 , 2 , . . . .
C.3 Proof of Theorem
Theorem 3 (Jacobian Suboptimality Rate for Chebyshev Method) .
Under Assumptions 1,2, let ξ def
=
(1 − √ κ ) / (1 + √ κ ) , and x t ( θ ) denote the t th iterate of the Chebyshev method.
Then, we have the following convergence rate
[Equation]
[Table]
[Equation]
[Table]
[Equation]
In short, the rate of the Chebyshev algorithm for unrolling in O ( t 2 ξ t ) .
Moreover, assuming G = 0 , the maximum of the upper bound over t can go up to
[Equation]
Proof.
First, we recall that the derivative of the Chebyshev polynomial of the first kind can be expressed as a function of the Chebyshev polynomial of the second kind (written ˜
U t ):
[Equation]
Therefore, we replace the polynomial P in Theorem 1 by C t , and evaluate
[Equation]
This polynomial achieves its maximum in absolute value at the end of the interval [ ℓ, L ] .
Therefore, after replacement, and using the fact that m ( L ) = 1 , ˜ C (1) = 1 , and ˜ U t (1) =
t , we obtain
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
Similarly, for the second term, we have
[Equation]
It suffices now to evaluate 1 | ˜ C
( m (0)) | .
Using (for example) (d’Aspremont et al., 2021, Theorem 2.1), we finally have
[Equation]
C.4 Proof of Proposition
Proposition 1.
Let x t be the t -th iterate of a first-order method.
Then, for all iterations t and for all θ , there exists a quadratic function f that verifies Assumption 1 such that G = 0 , and
[Equation]
Proof.
The proof is based on a reduction to the optimization case.
Indeed, consider the specific case of ridge regression, with a free scaling parameter α > 0 ,
[Equation]
In such a case, for all x 0 , we have ∥ ∂ θ ∇ f ( x 0 ( θ ) , θ ) ∥ F = 0 .
Moreover, this function is [ σ 2min ( A ) + α θ ] strongly convex and [ σ 2max ( A )+ α θ ] -smooth, where σ min and σ max are respectively the smallest and largest singular value of a matrix.
Let us write H =
A ⊤ A + α θI and x ⋆ =
H − 1 ( θ ) A T b .
Now, consider any quadratic function ˜ f of the form
[Equation]
Using the notation ¯ θ to be a fixed value of theta θ (i.e., ¯ θ = θ but ∂ θ ¯ θ = 0 ), it is possible to write f such that it matches ˜ f , by setting
[Equation]
This is possible only if ˜ H −
α ¯ θ ≻ 0 , or equivalently, if ℓ > α ¯ θ .
It suffices to set ℓ ¯ θ > α to ensure that condition.
This means we can cast any quadratic function that does not depends on θ into one that depends on θ , such that ∥ ∂ θ ∇ f ( x 0 ( θ ) , θ ) ∥ F = 0 .
In such a case, the master identity from Theorem 1 reads
[Equation]
where H ( θ ) =
A ⊤ A + θI .
Now, write Q t ( λ ) =
P t ( λ ) − λP ′ t ( λ ) .
We now have the following identity,
[Equation]
This identity is similar to the one we have in optimization:
[Equation]
and for this identity, we have the lower bound (Nemirovski, 1995, Proposition 12.3.2)
[Equation]
However, in the case of unrolling, we have different constraints on Q t , which are the following:
[Equation]
Therefore, we have more constraints on Q (i.e., on how fast we can decrease the accuracy bound).
Since we have seen that the functional class we work on is at least as large as the one of quadratic optimization, the lower bound can only be worse than the one for minimizing quadratic function with a bounded spectrum.
C.5 Proof of Proposition
Proposition 2.
Assume that ∥ ∂ H ( θ )( x 0 ( θ ) − x ⋆ ( θ ))
∥ F ≤ η ∥ ∂ x 0 ( θ ) − ∂ x ⋆ ( θ ) ∥ F .
Then, under Assumption 1, 2 and 3, we have the following bound for the average-case rate
[Equation]
Proof.
We first derive both sides of (4) and use Assumption 2, then we use Cauchy-Schwartz and ( a + b ) 2 ≤ [Equation] 2 +
[Equation]
[Equation]
Since the trace of a symmetric matrix is the sum of its eigenvalues [Equation] after taking the expectation on both sides [Equation]
we obtain the desired result:
[Equation]
C.6 Proof of Proposition
Proposition 3.
Let { S t } be a sequence of orthogonal Sobolev polynomials, i.e., ⟨ S i , S j ⟩ > [-
-]0 if i = j and 0 otherwise, normalized such that S i (0)
= 1 .
Then, the residual polynomial that minimizes the Sobolev norm can be constructed as
[Equation]
where a i = ∥ S
[Equation]
∥ 2 η and A t = t (cid:88) i = a i .
Moreover, we have that ∥ P ⋆t ∥ 2 η = 1 /A t .
Proof.
We have that the sequence { S i } i =0 ...t is a orthogonal basis for P t .
Therefore, we can write any polynomials as a weighted sum of S i .
Also, since P t (0) = 1 and S i (0)
= 1 , we have to enforce that the linear combination sums to one.
This means that
P t =
[Equation]
i = a i S i , t (cid:88) i
= a i = 1 .
We now minimize over α .
[Equation]
The Lagrangian of the optimization problem reads
[Equation]
Taking its derivative to zero gives the desired result:
[Equation]
Injecting the optimal solution into ∥ P ∥ 2 η gives
[Equation]
D Optimal Sobolev algorithm
We recall the Sobolev algorithm:
[Equation]
parametrized by:
• [ ℓ, L ] , lower and upper bound on the eigenvalues of H ( θ ) , • α , parameter of the Gegenbaueur distribution (12), supposed to be the expected spectral density (6).
Note: α = 0 leads to a sequence of Chebyshev polynomials for y t .
• η , assumed to satisfy the inequality ∥
∂ H ( θ )( x 0 ( θ ) − x ⋆ ( θ )) ∥ F ≤ η ∥ ∂ x 0 ( θ ) − ∂ x ⋆ ( θ ) ∥ F .
Intuitively, this parameter is the balance between ∥ P ∥ and ∥ P ′ ∥ .
D.1 Initialization (required for t = 0 and t = 1 ) D.1.1 Side parameters
[Equation]
D.1.2 Main parameters
[Equation]
D.2 Recurrence (for t ≥ 2 ) D.2.1 Side parameters
[Equation]
D.2.2 Main parameters
[Equation]
E Derivation of the Sobolev algorithm E.1 Notations
In this section, we use the following notations.
We denote by µ the Gegenbaueur density 12 defined in [ ℓ, L ] , ˜ µ the Gegenbaueur density defined in [ − 1 , 1] : [Equation] [Equation] and m : [
[Equation]
, L
[Equation]
[Equation]
where
[Equation]
We also denote by G t and ˜ G t the sequence of Gegenbaueur polynomials that are orthogonal respectively w.r.t.
the measure µ and ˜ µ , that it, for all i, j ≥ 0 , we have
[Equation]
L
[Equation]
In terms of normalization, we have that G t is a residual polynomial, and ˜ G t is a monic polynomials.
In other terms,
[Equation]
In such a case, by using the linear mapping m ( λ ) from [ ℓ, L ] to [ − 1 , 1] , see (19), we have the following relation:
[Equation]
Similarly, we define S t and ˜ S t the sequence of orthogonal Sobolev polynomials w.r.t.
the Sobolev product involving the Gegenbaueur density, i.e.,
[Equation]
and
[Equation]
Originally, they are called Gegenbaueur-Sobolev polynomials (Marcellán et al., 1994) because µ is a Gegenbaueur density, but for conciseness, we simply call them Sobolev polynomials.
As for the
Gegenbaueur polynomials, S t is a residual polynomial while ˜ S t is a monic polynomial.
Finally, we have that
[Equation]
and only if
[Equation]
Note that we make a distinction between plain symbols and tilde ˜ symbols, where the tilde ˜ notation is used for polynomials that are defined on [ − 1 , 1] , while the plain notation is the counterpart defined on [ ℓ, L ] .
E.2 Monic Sobolev polynomial
We now describe the construction of ˜ S , detailed in (Marcellán et al., 1994).
The monic Gegenbaueur polynomial is constructed as
[Equation]
Then, the Sobolev polynomials are defined as a simple recurrence involving [Equation] t and ˜ G t
[Equation]
where
[Equation]
Note that the following property will be important later:
[Equation]
where
[Equation]
E.3 Shifted, normalized Sobolev polynomials
We now shift and normalize the Sobolev polynomials, that it, instead of being defined in [0 , 1] and being monic, we make them defined in [ ℓ, L ] (evaluate the polynomial at x = m ( λ ) ) and residual (divide the polynomial by ˜ S t ( m (0)) ).
We begin by doing it to the Gegenbaueur polynomials.
By applying the technique from (Pedregosa et al., 2020, Proposition 18) on the polynomial ˜ G t ( m ( λ )) ,
[Equation]
We obtain the recurrence
[Equation]
) where
[Equation]
This expression can be cast into a recurrence that involves a step size and a momentum,
[Equation]
where
[Equation]
We now show how to shift and normalize the Sobolev polynomial.
The shifting operation is not complicated, as it suffice to evaluate the polynomial ˜ S t at x = m ( λ ) .
The difficult part is the normalization.
Using the relations (20), (21) and (23), we obtain
[Equation]
Therefore, we have to compute those quantities that involves ratio of polynomials evaluated at λ = 0 , whose recurrence is detailed in the next Proposition.
Proposition 4.
Let
[Equation]
Then,
[Equation]
Proof.
We now show, one by one, each terms of the recurrence.
We begin by ∆ Pt .
Indeed,
[Equation]
Therefore, using (27), we obtain
[Equation]
After comparing this expression with (26)
[Equation]
we deduce that
[Equation]
In other words,
[Equation]
To show the other recurrences
[Equation]
we will often use the fact that
[Equation]
We now show how to form τ t .
Indeed, using (32),
[Equation]
Using the same technique
[Equation]
we have for κ t :
[Equation]
However,
[Equation]
Therefore,
[Equation]
Finally
[Equation]
it remains to show the recurrence for ∆ St .
As usual,
[Equation]
We have seen before that
[Equation]
which finally gives
[Equation]
− 2 .
E.4 Norm of Sobolev Polynomials
Now that we can build the shifted, normalized Gegenbaueur and Sobolev polynomials, we still need to compute the norm of the Sobolev polynomial to compute P ⋆t .
First, for simplicity, we write
[Equation]
Indeed, to obtain the optimal method
[Equation]
we need to compute the coefficients
[Equation]
To do so [Equation] we will use the property (25):
[Equation]
We begin by the explicit expression of the norm of the shifted, normalized Sobolev polynomials
[Equation]
and express it as a function of the norm of the plain, monic Sobolev polynomial.
Indeed,
[Equation]
Since m ′ ( λ ) =
σ
[Equation]
and since ˜ η
= σ 1 η
[Equation]
we have
[Equation]
Note that, by definition of ∆ St [Equation] we have the recursion [Equation]
[Equation]
Let ¯ G t be defined as
[Equation]
i.e., ¯ G t is a scaled version of G t , which is the classical definition of Gegenbaueur polynomials.
Then
[Equation]
Since Γ( x + 1) = x Γ( x ) , we can deduce a recurrence equation.
Indeed,
[Equation]
with the initial condition
[Equation]
However, there is a factor between ¯ G t and the monic polynomial ˜ G t .
Indeed,
[Equation]
This factor can be computed recursively.
Let k t =
[Equation]
Then,
[Equation] [Equation]
Therefore, using successively (35)
[Equation] [Equation]
then (36)
[Equation]
we have
[Equation]
with the same initial condition
[Equation]
We now compute the recursion for ∥ S t ∥ 2 η .
Indeed, by using successively (33), (34), (25), (38) × 2 , (25) then (33),
[Equation]
We finally have the desired recurrence for the a t ’s since
[Equation]
where ¯ a is a nonzero multiplicative constant.
We can arbitrarily decide that ¯ a = [Equation] which gives us a 0 = 1 .
Given that S 1 = G [Equation] and after using (38) [Equation]
(25) and (33)
[Equation]
we have
[Equation]
F Asymptotic algorithm F.1 Asymptotics of Sobolev-Gegenbaeur polynomials
From (Scieur et al., 2020b), we know that the parameters converges asymptotically to
[Equation]
(cid:18)
[Equation]
In addition, it is easy to see that
[Equation]
Therefore,
[Equation]
Thus [Equation] the recurrence simplifies into (after replacing d ∞ by 0 )
[Equation]
This means that the asymptotic recurrence for S reads
[Equation]
Moreover, we have
[Equation]
When t → ∞ , we have that K t → 1 / 4 , ξ t → 1 / 4 , ∆ t → 4 m .
Therefore,
[Equation]
Moreover, d i d i − 2 → 1 .
So, we have in the end that
[Equation]
or more simply,
[Equation]
Therefore, when t → ∞ , we have
[Equation]
This means that the asymptotic dynamic for P ⋆ reads
[Equation]
G Asymptotic algorithm and asymptotic rate
The asymptotic recurrence of the polynomials reads
[Equation]
This can be simplified into
[Equation]
Translated into an algorithm, we finally have a weighted average of HB iterates:
[Equation]
Note that the asymptotic rate reads
[Equation]
Therefore
[Equation]
when t → ∞ ,
[Equation]