SHINE: SH ARING THE IN VERSE E STIMATE FROM THE FORWARD PASS FOR BI - LEVEL OPTIMIZATION AND IMPLICIT MODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
In recent years, implicit deep learning has emerged as a method to increase the effective depth of deep neural networks.
While their training is memory-efﬁcient, they are still signiﬁcantly slower to train than their explicit counterparts.
In Deep
Equilibrium Models (DEQs), the training is performed as a bi-level problem, and its computational complexity is partially driven by the iterative inversion of a huge Jacobian matrix.
In this paper, we propose a novel strategy to tackle this computational bottleneck from which many bi-level problems suffer.
The main idea is to use the quasi-Newton matrices from the forward pass to efﬁciently approximate the inverse Jacobian matrix in the direction needed for the gradient computation.
We provide a theorem that motivates using our method with the original forward algorithms.
In addition, by modifying these forward algorithms, we further provide theoretical guarantees that our method asymptotically estimates the true implicit gradient.
We empirically study this approach in many settings, ranging from hyperparameter optimization to large Multiscale DEQs applied to CIFAR and
ImageNet.
We show that it reduces the computational cost of the backward pass by up to two orders of magnitude.
All this is achieved while retaining the excellent performance of the original models in hyperparameter optimization and on CIFAR, and giving encouraging and competitive results on ImageNet.
I NTRODUCTION
Implicit deep learning models such as Neural ODEs (Chen et al., 2018), OptNets (Amos and
Zico Kolter, 2017) or Deep Equilibrium models (DEQs) (Bai et al., 2019; 2020) have recently emerged as a way to train deep models with inﬁnite effective depth without the associated memory cost.
Indeed, while it has been observed that the performance of deep learning models increases with their depth (Telgarsky, 2016), an increase in depth also translates into an increase in the memory footprint required for training, which is hardware-constrained.
While other works such as invertible neural networks (Gomez et al., 2017; Sander et al., 2021) or gradient checkpointing (Chen et al., 2016) also tackle this issue, implicit models do so using an O (1) memory cost and with constraints on the architecture that are usually not detrimental to the model performance (Bai et al., 2019).
These models have been successfully applied to large-scale tasks such as language modeling (Bai et al., 2019), computer vision tasks (Bai et al., 2020) and inverse problems (Gilton et al., 2021; Heaton et al., 2021).
In general, the formulation of DEQs can be cast as a bi-level problem of the following form:
[Equation]
We will refer to the root ﬁnding problem g θ ( z (cid:63) )
=
as the inner problem , and call its resolution the forward pass .
On the other hand, we will refer to arg min θ L ( z (cid:63) ) as the outer problem , and call the computation of the gradient of L ( z (cid:63) ) w.r.t.
θ the backward pass .
The core idea for DEQs is that their output z (cid:63) is expressed as a ﬁxed point of a parametric function f θ from R d to R d ,
i.e., g θ
( z (cid:63) )
=
z (cid:63) − f θ ( z (cid:63) )
= 0 .
This model is said to have inﬁnitely many weight-tied layers as z (cid:63) can be obtained by successively applying the layer f θ inﬁnitely many times, provided f θ is contractive.
In practice, DEQs’ forward pass is not computed by applying successively the function but usually relies on quasi-Newton (qN) algorithms, such as Broyden’s method (Broyden, 1965), which approximates efﬁciently the Jacobian matrix ∂g θ ∂z and its inverse for root-ﬁnding.
To compute DEQs’ gradient efﬁciently and avoid high memory cost, one does not rely on backpropagation but uses the implicit function theorem (Krantz and Parks, 2013) which gives an analytical expression of the Jacobian of z (cid:63) with respect to θ , ∂z (cid:63) ∂θ .
While this method is memory efﬁcient, it requires the computation of matrix-vector products involving the inverse of a large Jacobian matrix, which is computationally demanding.
To make this computation tractable, one needs to rely on an iterative algorithm based on vector-Jacobian products, which renders the training particularly slow, as highlighted by the original authors (Bai et al., 2020)
(see also the break down of the computational effort in Section E.3).
Moreover, the formulation (1) allows us to also consider general bi-level problems such as hyperparameter optimization under the same framework.
For instance, hyperparameter optimization for
Logistic Regression (LR) can be written as
[Equation]
to
[Equation]
where L train and L val correspond to the training and validation losses from the LR problem (Pedregosa, 2016).
Here, z corresponds to the weights of the LR model while θ is the regularisation parameter.
As the training loss is smooth and convex, the inner problem can be written as in (1) with g θ = ∇ z r θ to ﬁt (1).
Similarly to DEQ, the inner problem is often solved using qN methods, which approximate the inverse of the Hessian in the direction of the steps, such as the LBFGS algorithm (Liu and Nocedal, 1989), and the gradient computation suffers from the same drawback as it is also obtained using the implicit function theorem.
Lorraine et al.
(2020) review the different hypergradient approximations for bi-level optimization and evaluate them on multiple tasks.
With the increasing popularity of DEQs and the ubiquity of bi-level problems in machine learning, a core question is how to reduce the computational cost of the resolution of (1).
This would make these methods more accessible for practitioners and reduce the associated energy cost.
In this work, we propose to exploit the estimates of the (inverse of the) Jacobian/Hessian produced by qN methods in the hypergradient computation.
Moreover, we also propose extra updates of the qN matrices which maintain the approximation property in the direction of the steps, and ensure that the inverse Jacobian is approximated in an additional direction.
In effect, we can compute the gradient using the inverse of the ﬁnal qN matrix instead of an iterative algorithm to invert the Jacobian in the gradient’s direction, while stressing that the inverse of a qN matrix, and thus the multiplication with it, can be computed very efﬁciently.
We emphasize that the goal of this paper is neither to improve the algorithms used to compute z (cid:63) , nor is it to demonstrate how to perform the inversion of a matrix in a certain direction as a stand-alone task.
Rather, we are describing an approach that combines the resolution of the inner problem with the computation of the hypergradient to accelerate the overall process.
Our work is the ﬁrst to consider modifying the inner problem resolution in order to account for the bi-level structure of the optimization The idea to use additional updates of the qN matrices to ensure additional approximation properties is not new, and it is also known that a full matrix inversion can be accomplished in this way.
For instance, Gower and Richtárik (2017) used sketching to design appropriate extra secant conditions in order to obtain guarantees of uniform convergence towards the inverse of the Jacobian.
The novelty in our work is that we integrate additional update to yield the inverse in a speciﬁc direction, which is substantially cheaper than computing the inverse.
A concurrent work by Fung et al.
(2021) is also concerned with the acceleration of DEQs’ training, where the inverse Jacobian is approximated with the identity.
Under strong contractivity and conditioning assumptions, it is proven that the resulting approximation is a descent direction and the authors show good empirical performances for small scale problems.
The contributions of our paper are the following:
• We introduce a new method to greatly accelerate the backward pass of DEQs (and generally, the differentiation of bi-level problems) using qN matrices that are available as a by-product of the forward computations.
We call this method SHINE ( SH aring the IN verse E stimate).
• We enhance this method by incorporating knowledge from the outer problem into the inner problem resolution.
This allows us to provide strong theoretical guarantees for this approach in various settings.
• We additionally showcase its use in hyperparameter optimization.
Here, we demonstrate that it provides a gain in computation time compared to state-of-the-art methods.
• We test it for DEQs for the classiﬁcation task on two datasets, CIFAR and ImageNet.
Here, we show that it decreases the training time while remaining competitive in terms of performance.
• We extend the empirical evaluation of the Jacobian-Free method to large scale multiscale DEQs and show that it performs well in this setting.
We also show that it is not suitable for more general bi-level problems.
• We propose and evaluate a natural reﬁnement strategy for approximate Jacobian inversion methods (both SHINE and Jacobian-Free) that allows a trade-off between computational cost and performances.
H YPERGRADIENT O PTIMIZATION WITH A PPROXIMATE J ACOBIAN I NVERSE 2. SHINE: H YPERGRADIENT D ESCENT WITH A PPROXIMATE J ACOBIAN I NVERSE
Hypergradient Optimization Hypergradient optimization is a ﬁrst-order method used to solve (1).
We recall that in the case of smooth convex optimization, ∂g θ ∂z is the Hessian of the inner optimization problem, while for deep equilibrium models, it is the Jacobian of the root equation.
In the rest of this paper, with a slight abuse of notation, we will refer to both these matrices with J g θ whenever the results can be applied to both contexts.
To enable Hypergradient Optimization, i.e. gradient descent on L with respect to θ , Bai et al.
(2019, Theorem 1) show the following theorem, which is based on implicit differentiation (Krantz and Parks, 2013):
Theorem 1 (Hypergradient (Bai et al., 2019; Krantz and Parks, 2013)) .
Let θ ∈ R p be a set of parameters, let L : R d → R be a loss function and g θ : R d → R d be a root-deﬁning function.
Let z (cid:63) ∈ R
d
such that g θ ( z (cid:63) )
=
and J g θ ( z (cid:63) ) is invertible, then the gradient of the loss L wrt.
θ , called
Hypergradient, is given by
[Equation] [Equation]
Algorithm g θ ( z (cid:63) ) =
Result:
Root z (cid:63) , qN matrix B b = true if using Broyden’s method, b = false if using BFGS n = 0 , z
[Equation] = 0 , B
= I while not converged do [Equation]
[Equation] n can be 1 or determined by line-search
[Equation]
if b then
[Equation]
else
[Equation]
The norm used in BFGS is a weighted Frobenius norm end
[Equation]
In practice, we use an algorithm to approximate z (cid:63) , and Theorem 1 gives a plug-in formula for the backward pass.
Note that this formula is independent of the algorithm chosen to compute z (cid:63) .
Moreover, as opposed to explicit networks, we do not need to store intermediate activations, resulting in the aforementioned training time memory gain for DEQs.
Once z (cid:63) has been obtained, one of the major bottlenecks in the computation of the Hypergradient is the inversion of
[Equation]
( z (cid:63) ) in the directions
∂g
θ ∂θ (cid:12) (cid:12)(cid:12)
z (cid:63) or ∇ z L ( z (cid:63) ) .
Quasi-Newton methods In practice, the forward pass is often carried out with qN methods.
For instance, in the case of bi-level optimization for Logistic Regression, Pedregosa (2016) used L-BFGS (Liu and Nocedal, 1989), while for Deep Equilibrium Models, Bai et al.
(2019) used Broyden’s method (Broyden, 1965), later adapted to the multi-scale case in a limitedmemory version (Bai et al., 2020).
These quasi-Newton methods were ﬁrst inspired by Newton’s method, which ﬁnds the root of g θ via the recurrent Jacobian-based updates z
[Equation]
Speciﬁcally, they replace the Jacobian J g θ ( z n ) by an approximation B n that is based on available values of the iterates z n and g θ rather than its derivative.
These B n , called qN matrices, are deﬁned recursively via an optimization problem with constraints called secant conditions.
Solving this problem leads to expressing B n as a rank-one or rank-two update of B n − 1 , so that B n is the sum of the initial guess
B 0 (in our settings, the identity) and n low-rank matrices (less than n in limited memory settings).
This low rank structure allows efﬁcient multiplication by B n and B − 1 n .
We now explain how the use of qN methods as inner solver can be exploited to resolve this computational bottleneck.
SHINE Roughly speaking, our proposition is to use B − 1 = lim n →∞ B − 1 n as a replacement for
J g θ
( z (cid:63) ) − 1 in (3), i.e. to share the inverse estimate between the forward and the backward passes.
This gives the approximate Hypergradient
[Equation]
In practice we will consider the nonasymptotical direction p
[Equation]
Thanks to the Sherman-Morrison formula (Sherman and Morrison, 1950), the inversion of B n can be done very efﬁciently (using scalar products) compared to the iterative methods needed to invert the true
Jacobian J g θ
( z (cid:63) ) .
In turn, this signiﬁcantly reduces the computational cost of the Hypergradient computation.
Relationship to the Jacobian-Free method Because B
= I in our setting, we may regard B as an identity matrix perturbed by a few rank-one updates.
In the directions that are used for updates, B is going to be different from the identity, and hopefully closer to the true Jacobian in those directions.
However, in all orthogonal directions we fall exactly into the setting of the Jacobian-Free method introduced by Fung et al.
In that work, J g θ ( z (cid:63) ) − 1 is approximated by I , and the authors highlight that this is equivalent to using a preconditioner on the gradient.
Under strong assumptions on g θ they show that this preconditioned gradient is still a descent direction.
Transition to the exact Jacobian Inverse.
The approximate gradient p ( n ) θ can also be used as the initialization of an iterative algorithm for inverting J g θ
( z (cid:63) ) in the direction ∇ z L ( z (cid:63) ) .
With a good initialization, faster convergence can be expected.
Moreover, if the iterative algorithm is also a qN method, which is the case in practice in the Multiscale DEQ implementation, we can use the qN matrix B from the forward pass to initialize the qN matrix of this algorithm.
We refer to this strategy as the reﬁne strategy .
Because the reﬁne strategy is essentially a smart initialization scheme, it recovers all the theoretical guarantees of the original method (Bai et al., 2019; 2020; Pedregosa, 2016).
2. C ONVERGENCE TO THE TRUE GRADIENT
To further justify and formalize the idea of SHINE, we show that the direction p ( n ) θ converges to the
Hypergradient ∂ L ∂θ (cid:12)(cid:12)(cid:12) z (cid:63) .
We now collect the assumptions that will be used for this purpose.
Assumption 1 (Uniform Linear Independence (ULI) (Li et al., 1998)) .
There exist a positive constant ρ > 0 and natural numbers n 0 ≥ 0 and m ≥ d with the following property:
For any n ≥ n 0 we can ﬁnd indices n ≤ n 1 ≤ . . .
≤ n d ≤ n + m such that, for p n deﬁned in Algorithm 1, the smallest singular value of the d
[Equation]
is no smaller than ρ .
Assumption 2 (Smoothness and convergence to the ﬁxed point) .
(i) (cid:80) ∞ n =0 (cid:107) z n − z (cid:63) (cid:107) <
∞ for some z (cid:63) with g θ
( z (cid:63) )
= 0 ; (ii) g θ is C 1 , J g θ is Lipschitz continuous near z (cid:63) , and J g θ
( z (cid:63) ) is invertible; (iii) ∇ z L is continuous, and ∀ θ , ∂g θ
∂θ is continuous.
Remark.
The Assumption 2 (i) implies lim n →∞ z n = z (cid:63) .
The existence of the Jacobian and its inverse are assumptions that are already made in the regular DEQ setting just to train the model.
Theorem 2 (Convergence of SHINE to the Hypergradient using ULI) .
Let us denote p ( n ) θ , the SHINE direction for iterate n in Algorithm 1 with b = true .
Under Assumptions 1 and 2, for a given parameter θ , ( z n ) converges q-superlinearly to z (cid:63) and
[Equation]
Proof.
From More and Trangenstein (1976, Theorem 5.7) we obtain that lim n →∞ B n =
J g θ
( z (cid:63) ) .
We can then conclude using the continuity of the inversion operator on the space of invertible matrices and of the right and left matrix vector multiplications.
A complete proof is given in Section B.1.
Theorem 2 establishes convergence of the SHINE direction to the true Hypergradient, but relies on
Assumption 1 (ULI).
While ULI is often used to prove convergence results for qN matrices, e.g. in (Conn et al., 1991; Li et al., 1998; Nocedal and Wright, 2006), it is a strong assumption whose satisfaction in practice is debatable, cf., e.g., (Fayez Khalfan et al., 1993).
For Broyden’s method,
ULI is violated in all numerical experiments in (Mannel, 2020; 2021a;b), and those works also prove that ULI is necessarily violated in certain settings (but the setting of this work is not covered).
In the following we therefore derive results that do not involve ULI.
2. O UTER P ROBLEM A WARENESS
The ULI assumption guarantees convergence of B − 1 n to J g θ ( z (cid:63) ) − 1 .
However, (3) only requires the multiplication of J g θ
( z (cid:63) ) − 1 with ∂g θ
∂θ
| z (cid:63) from the right and ∇ z L ( z (cid:63) ) from the left.
BFGS with OPA In order to strengthen Theorem 2, let us consider the setting of bi-level optimization with a single regularizing hyperparameter θ .
There, the partial derivative ∂g
θ
∂θ | z (cid:63) is a d -dimensional vector and it is possible to compute its approximation
∂g
θ ∂θ | z n at a reasonable cost.
We propose to incorporate additional updates of the quasi-Newton matrix B n into Algorithmthat improve the approximation quality of B − 1 n in the direction ∂g
θ ∂θ | z n (thus asymptotically in the direction ∂g
θ ∂θ | z (cid:63) ).
Given a current iterate pair ( z n , B n ) , these additional updates only change B n , but not z n .
We will demonstrate that a suitable update direction e n ∈ R d is given by
[Equation]
where ( t n ) ⊂ [0 , ∞ ) satisﬁes (cid:80) n t n < ∞ .
This update direction will be used to create an extra secant condition X − 1 ( g θ ( z n + e n ) − g θ ( z n ))
=
e n for the additional update of B n .
Since this extra update is based on the outer problem, we refer to this technique as Outer-Problem Awareness (OPA).
The complete pseudo code of the OPA method in the LBFGS algorithm (Liu and Nocedal, 1989) is given in Appendix A.
We now prove that if extra updates are applied at a ﬁxed frequency, then fast (q-superlinear) convergence of ( z n ) to z (cid:63) is retained, while convergence of the SHINE direction to the true Hypergradient is also ensured.
To show this, we use the following assumption.
Assumption 3 (Assumptions for BFGS) .
Let g θ ( z ) =
∇ z r θ ( z ) for some C 2 function r θ : R d → R .
Consider Algorithm 1 with b = false .
We assume some regularity on r and that an appropriate line search is used.
An extended version of this assumption is given in Section B.2 (Assumption 5).
Theorem 3 (Convergence of SHINE to the Hypergradient for BFGS with OPA) .
Let us consider p ( n ) θ , the SHINE direction for iterate n in Algorithm 1 that is enriched by extra updates in the direction e n deﬁned in (5).
Under Assumptions 2 (ii-iii) and 3, for a given parameter θ , we have the following:
Algorithm 1, for any symmetric and positive deﬁnite matrix B 0 , generates a sequence ( z n ) that converges q-superlinearly to z (cid:63) , and there holds
[Equation] [Equation]
Proof.
It follows from known results that the extra updates do not destroy the q-superlinear convergence of ( z n ) .
The proof of (6) relies ﬁrstly on the fact that by continuity of the derivative of g θ , we have lim n →∞ ∂g θ ∂θ | z n = ∂g θ
∂θ | z (cid:63) .
Due to the extra updates we can show convergence of the qN matrices to the true Hessian in the direction of the extra steps e n , from which (6) follows.
A full proof is provided in Section B.2.
Remark.
Theorem 3 also holds without line searches (i.e., α n = 1 for all n ) and any C 2 function r θ (such that g θ ( z ) =
∇ z r θ ( z ) ) with locally Lipschitz continuous Hessian if z 0 is close enough to some z (cid:63) with ∇ z r θ ( z (cid:63) )
=
and ∇ 2 zz r θ ( z (cid:63) ) positive deﬁnite.
We note that Theorem 3 guarantees fast convergence of the iterates ( z n ) and that z 0 does not have to be close to z (cid:63) for that guarantee.
Also, there is no restriction on B 0 other than being symmetric and positive deﬁnite (which is satisﬁed for our choice B
= I ).
Finally, Theorem 3 does not rely on ULI.
From a practical standpoint we thus regard Theorem 3 as a much stronger result than Theorem 2.
Adjoint Broyden with OPA It is not practical to use the partial derivative ∂g
θ ∂θ in the DEQ setting because it is a huge Jacobian that we do not have access to in practice.
In order to still leverage the core idea of OPA, we propose to use extra updates that ensure that B − 1 n approximates J g θ ( z (cid:63) ) −in the direction ∇ z L ( z (cid:63) ) applied by left-multiplication, as required by (3).
An appropriate secant condition is given by
[Equation]
where
[Equation]
To incorporate the secant condition (7), we use the Adjoint Broyden’s method (Schlenkrich et al., 2010), a qN method relying on the efﬁcient vector-Jacobian multiplication by J g θ using autodifferentiation tools.
To prove convergence of the SHINE direction for this method, we need the following assumption.
Assumption 4 (Uniform boundedness of the inverse qN matrices) .
The sequence ( B n ) generated by
Algorithm 1 satisﬁes
[Equation]
Remark.
Convergence results for quasi-Newton methods usually include showing that Assumptionholds, cf.
Broyden et al.
(1973, Theorem 3.2) for Broyden’s method and the BFGS method, respectively,
Schlenkrich et al.
(2010, Theorem 1) for the Adjoint Broyden’s method.
It can also be proved that
Assumption 4 holds for globalized variants of these methods, e.g., for the line-search globalizations of Broyden’s method proposed by Li and Fukushima (2000).
We point out that Assumption 1 entails lim B n =
J g θ
( z (cid:63) ) and thus lim B − 1 n =
J g θ
( z (cid:63) ) − 1 , so it is clearly stronger than Assumption 4.
Theorem 4 (Convergence of SHINE to the Hypergradient for Adjoint Broyden with OPA) .
Let us consider p ( n ) θ , the SHINE direction for iterate n in Algorithm 1 with the Adjoint Broyden secant condition (7) and extra update in the direction v n deﬁned in (8).
Under Assumptions 2 and 4, for a given parameter θ , we have q-superlinear convergence of ( z n ) to z (cid:63) and
[Equation]
Proof.
The q-superlinear convergence of ( z n ) follows from Schlenkrich et al.
(2010, Theorem 2).
To establish convergence of the SHINE direction, we proceed in three steps.
First, it is shown that for ∇ z L ( z (cid:63) ) =
the claim holds due to continuity and Assumption 4.
Then ∇ z L ( z (cid:63) ) (cid:54) = 0 is considered and it is proved that the desired convergence holds on the subsequence that corresponds to the additional updates.
Lastly, this result is transferred to the entire sequence by involving the ﬁxed frequency of the additional updates.
The complete proof is provided in Section B.3.
Using the Adjoint Broyden’s method comes at a computational cost.
Indeed, because we now rely on J g θ , we have to store the activations of g θ ( z ) (which has a computational cost in addition to a memory cost), but also perform the vector-Jacobian product in addition to the function evaluation.
R ESULTS
We test our method in 3 different setups and compare it to the original iterative inversion and its closest competitor, the Jacobian-Free method (Fung et al., 2021).
We draw the reader’s attention to the fact that although the Jacobian-Free method (Fung et al., 2021) is used outside the assumptions needed to have theoretical guarantees 2 of descent, it still performs relatively well in the Deep Equilibrium setting.
The same is true for SHINE: While the ULI assumption is not met (and we are in practice far from the ﬁxed point convergence), it performs well in practice.
Implementations.
All the bi-level optimization experiments were done using the HOAG code (Pedregosa, 2016) 3 , which is based on the Python scientiﬁc ecosystem (Harris et al., 2020; Pedregosa et al., 2011; Virtanen et al., 2020).
Deep Equilibrium experiments were done using the PyTorch (Paszke et al., 2019) code for Multiscale DEQ (Bai et al., 2020) 4 , which was distributed under [Figure] Time (s)
[Figure]
T e s t
L o ss Sub o p t i m a li t y 20news [Figure] Time (s)
[Figure]
real-sim SHINE (ours) [Figure] SHINE
[Figure] Jacobian-Free Grid search the MIT license. Plots were done using Matplotlib (Hunter, 2007), with Science Plots style (Garrett and Peng, 2021).
DEQ trainings were done in a publicly funded HPC, using nodes
In practice, we never reach convergence of ( z n ) , hence the approximate gradient might be far from the true gradient.
To improve the approximation quality, we now propose a variant of our method.
Fallback in the case of wrong inversion.
Empirically, we noticed that using B can sometimes produce bad approximations, although with very low probability.
We propose to detect this with by monitoring a telltale sign based on the norm of the approximation, as we veriﬁed on several examples that cases with a huge norm compared to the correct inversion also had a very bad correlation with the correct inversion.
In these cases, we can simply fallback onto another inversion method.
For the
Deep Equilibrium experiments, when the norm of the inversion using SHINE is 1.3 times above the norm of the inversion using the Jacobian-Free method (which is available at no extra computational cost), we use the Jacobian-Free inversion.
We refer to this strategy as the fallback strategy .
3. B I - LEVEL OPTIMIZATION – H YPERPARAMETER OPTIMIZATION IN L OGISTIC R EGRESSION
We ﬁrst test SHINE in the simple setting of bi-level optimization for (cid:96) 2 -regularized LR, using the code from Pedregosa (2016) and the same datasets.
Convergence on unseen data is illustrated in Figure 1.
An acceptable level of performance is reached twice faster for the SHINE method compared to any other competitor.
Another ﬁnding is that the reﬁne strategy does not provide a deﬁnitive improvement over the vanilla version of SHINE.
In order to verify that the performance gain of SHINE is not simply driven by truncated inversion, we also run HOAG with limited number of inversion iteration and showed that this degrades its performances (see HOAG limited backward in Section E.1).
We also tested our implementation of OPA on the 20news dataset and present the results in Figure 2.
In order to get a fair comparison, we implemented both SHINE, SHINE-OPA and HOAG using the same full Python code instead of relying on the original code which relied on the Fortran implementation of L-BFGS from (Virtanen et al., 2020). While SHINE with OPA does not outperform the vanilla SHINE, it reaches similar performances, outperforming HOAG, and comes with strong theoretical grounding. [Figure]
t
i
[Figure]
a li
[Figure]
We also showed on a smaller dataset, the breast cancer dataset (Dua and Graff, 2017), that OPA indeed ensures a better approximation of the inverse in the prescribed direction.
For a given split of the data, we compared the quality of the approximation of the inversion in three different directions: a prescribed direction chosen randomly but used for the OPA update, the Krylov direction
[Figure] [Equation] z n − z n − 1 ) and a random direction not used in the qN algorithm. The results for 100 runs with different random seeds are depicted in Figure 2, where we can observe that OPA indeed ensures a better inversion in the prescribed direction compared to a random direction.
We also notice that a poor direction for the inversion seems correlated with a small magnitude.
Next, we tested SHINE on the more challenging DEQ setup. Two experiments illustrate the performance of SHINE on the image classiﬁcation task on two datasets. For both datasets, we used the same model conﬁguration as in the original Multiscale DEQ paper (Bai et al., 2020) and did not ﬁne tune any hyperparameter.
For the different DEQ training methods, models for a given seed share the same unrolled-pretraining steps. We do not include OPA in the DEQ results because while the gradients are well correlated with the true ones (see Figure E.2), we observe a sharp initial performance drop that reduces its performance on Imagenet. We provide partial results in Section E.4. CIFAR-10.
The ﬁrst dataset is CIFAR-10 (Krizhevsky, 2009) which features 60,000 32 × 32 images representing 10 classes.
For this dataset, the size of the multi-scale ﬁxed point is d = 50 k. We train the models for ﬁve different random seeds. The results in Figure 3 show that for the vanilla version, SHINE slightly outperforms the JacobianFree method (Fung et al., 2021). Additionally, our results suggest that SHINE (in its vanilla version) is able to reduce the time taken for the backward pass almost 10-fold compared to the original method while retaining a competitive performance (on par with Res-Net-18 (He et al., 2016) at 92.9%). Finally, we do highlight that the Jacobian-Free method (Fung et al., 2021) is able to perform well outside the scope of its theoretical assumptions, albeit with slightly worse performance than SHINE.
We conjecture that the batched stochastic gradient descent helps accelerated methods by averaging out the errors made in the approximation. ImageNet. The second dataset is the ImageNet dataset (Deng et al., 2009) which features 1. million images cropped to 224 × 224 , representing 1000 classes. This dataset is recognized as a large-scale computer vision problem and the dimension of the ﬁxed point to ﬁnd is d = 190 k.
For this challenging task, we noticed that the vanilla version of SHINE was suffering a big drop just after the transition from unrolled pre-training to actual equilibrium training. To remedy partly this problem, we introduced the fallback to Jacobian-Free inversion. The results for a single random seed presented in Figure 3 for the ImageNet dataset are given for SHINE with fallback. We veriﬁed that
[Figure]
the fallback is barely used : in 1000 batches of size 32, only 2 samples used fallback, a proportion of 6 .
25 × 10 − 5 . Despite the drop suffered at the beginning of the equilibrium training, SHINE in its reﬁned version is able to perform on par with the Jacobian-Free method (Fung et al., 2021).
We also conﬁrm the importance of choosing the right initialization to perform accelerated backpropagation, by showing that with a limited iterative inversion, the performance of the original method deteriorates. Finally, while the drop in performance for the accelerated methods is signiﬁcant when applied in their vanilla version, we remind the reader that no ﬁne-tuning was performed on the training hyperparameters, making those results encouraging (on par with architectures like ResNet-18 (He et al., 2016)). The key take-away from Figure 3 is that both SHINE and Jacobian-Free approximation methods allow to accelerate the DEQ’s backward pass at a relatively low accuracy cost. Moreover, using the proposed reﬁned versions of these methods, the performance drop can be reduced by reducing the acceleration.
C ONCLUSION AND D ISCUSSION
We introduced SHINE, a method that leverages the qN matrices from the forward pass to obtain an approximation of the gradient of the loss function, thereby reducing the time needed to compute this gradient.
We showed that this method can be used on a wide range of applications going from bi-level optimization to small and large scale computer vision tasks.
We found that both SHINE and the
Jacobian-Free method reduce the required amount of time for the backward pass of implicit models, potentially lowering the barriers for training implicit models.
As those methods still suffer from a small performance drop, there is room for further improvement.
In particular, a potential experimentation avenue would be to understand how to balance the efforts of the Adjoint Broyden method in order to come closer to guaranteeing the asymptotical correctness of the approximate inversion.
On the theoretical side, this may involve the rate of convergence of the approximated gradient.
It also seems desirable to develop a version of Theorem 4 in which convergence of ( z n ) to z (cid:63) is not an assumption but rather follows from the assumptions, as achieved in Theorem 3.
We have no doubt that the contraction assumption used for the Jacobian-Free method would allow to prove such a result, but expect that a signiﬁcantly weaker assumption will sufﬁce.
R EPRODUCIBILITY S TATEMENT
We provide with the submission of this paper the full code necessary to reproduce the ﬁgures and the other quantitative results of the paper, from the training, to the evaluation and the actual ﬁgure drawing.
We made sure to use seeds and veriﬁed that the seeding was indeed allowing reproducible results.
We also provide time estimates for the reproduction of the ﬁgures.
We made sure to provide the full proofs for our theorems in the supplementary material of this manuscript.
The core concepts used in the proofs, and their sketches are also laid out in the main text.
A OPA ALGORITHM
Algorithm LBFGS: (Limited memory) BFGS method with OPA
Input: initial guess ( z 0 , B − 1 0 ) , where B − 1 0 is symmetric and positive deﬁnite, tolerance (cid:15) > 0 , frequency of additional updates M ∈ N , memory limit L ∈ N ∪ {∞} , ( t n ) a null sequence of positive numbers with (cid:80) n t n < ∞ Let F := ∇ z g θ for n = 0 , 1 , 2 , . . .
do if (cid:107) F ( z n ) (cid:107) ≤ (cid:15) then let z (cid:63) :
=
z n and let B := B n ; STOP
Let ˆ B − 1 n := B − 1 n if ( n mod M ) = 0 then let e n :=
t n B − 1 n ∂g θ ∂θ (cid:12)(cid:12)(cid:12) z n , ˆ y n := F ( z n + e n ) − F ( z n ) and ˆ r n := ( e n ) T ˆ y n if ˆ r n > 0 then let ˆ a n :=
e n − B − 1 n ˆ y n and let
[Equation]
Let B − 1 n := ˆ B − 1 n if n ≥ L then remove update n − L from B − 1 n
Let p n := − B − 1 n F ( z n )
Obtain α n via line-search and let s n := α n p n
Let z n +1 :=
z n + s n , y n := F ( z n +1 ) − F ( z n ) and r n :=
( s n ) T y n if r n > 0 then let a n :=
s n − B − 1 n y n and let
B − 1 n +1 :=
B − 1 n + a n ( s n ) T + s n ( a n ) T r n − ( a n ) T y n ( r n ) 2 s n ( s n ) T else let B − 1 n +1 :
=
B − 1 n if n ≥ L then remove update n − L from B − 1 n +
Output: z (cid:63) , B
Remark.
A possible choice for ( t n ) is to use an arbitrary t 0 > 0 and t n := (cid:107) s n − 1 (cid:107) for n ≥ 1 .
B P ROOFS OF SHINE CONVERGENCE
To facilitate reading, we restate the results before proving them.
B. C ONVERGENCE USING ULI
Theorem 2 (Convergence of SHINE to the Hypergradient using ULI) .
Let us denote p ( n ) θ , the SHINE direction for iterate n in Algorithm 1 with b = true .
Under Assumptions 1 and 2, for a given parameter θ , ( z n ) converges q-superlinearly to z (cid:63) and
[Equation]
Proof.
Under Assumptions 1 and 2, More and Trangenstein (1976, Theorem 5.7) shows that B n satisﬁes
[Equation]
The inversion operator is continuous in the space of invertible matrices, so we have:
[Equation]
Because ∇ z L and ∂g θ
∂θ are continuous at z (cid:63) by Assumption 2 (iii), we also have thanks to Assumption 2 (i):
[Equation]
and lim
[Equation]
By continuity we then deduce that, as claimed,
[Equation]
B. C ONVERGENCE FOR BFGS WITH OPA
Assumption 5 (Extended Assumptions for BFGS) .
Let g θ ( z ) =
∇ z r θ ( z ) for some C 2 function r θ : R d → R .
Consider Algorithm 1 with b = false and suppose that
the set Ω := { z ∈ R d : r θ ( z ) ≤ r θ ( z 0 ) } is convex; 2.
r θ is strongly convex in an open superset of Ω (this implies that r θ has a unique global minimizer z (cid:63) ) and has a Lipschitz continuous Hessian near z (cid:63) ; 3.
there are positive constants η 1 , η 2 such
that the
line search
used in the algorithm ensures that for each n ≥ 0 either
[Equation]
θ ( z n ) + η 2 ∇ r θ ( z n ) T p n is satisﬁed; 4.
the line search has the property that α n = 1 will be used if both (cid:107) ( B n − J g θ ( z n ))
s n (cid:107) (cid:107) s
n (cid:107) and (cid:107) z n − z (cid:63) (cid:107) are sufﬁciently small.
Remark.
The requirements 3.
and 4.
on the line search are, for instance, satisﬁed under the well-known Wolfe conditions, see Byrd et al.
(1988, section 3) for further comments.
Theorem 3 (Convergence of SHINE to the Hypergradient for BFGS with OPA) .
Let us consider p ( n ) θ , the SHINE direction for iterate n in Algorithm 1 that is enriched by extra updates in the direction e n deﬁned in (5).
Under Assumptions 2 (ii-iii) and 3, for a given parameter θ , we have the following:
Algorithm 1, for any symmetric and positive deﬁnite matrix B 0 , generates a sequence ( z n ) that converges q-superlinearly to z (cid:63) , and there holds
[Equation]
Proof.
The proof is divided into four steps.
The ﬁrst step is to establish the q-superlinear convergence of ( z n ) to z (cid:63) .
Denoting by N e ⊂ { 0 , M, 2 M, . . .
} the set of indices of extra updates that are actually applied, the second step consists of showing
[Equation]
where, in this proof, B n always represents the matrix from Algorithm LBFGS before the update in the direction e n is applied, i.e., the matrix whose inverse appears in the deﬁnition of e n , while ˆ B n always represents the matrix from Algorithm LBFGS after the update in the direction e n has been applied; if the update in the direction e n is not applied, then B n = ˆ B n .
The third step is to prove that (9) implies the desired convergence (6) of the SHINE direction if the limit n → ∞ is replaced by
N e (cid:51) n → ∞ , i.e., the limit is taken on the subsequence corresponding to N e .
The fourth step is then to transfer the convergence to the entire sequence.
It is easy to check that instead of updating B − 1 n , respectively, ˆ B − 1 n , we can also obtain the sequences ( B n ) and ( ˆ B n ) by updating according to
[Equation]
for the usual update (skipping the update if y Tn s n ≤ 0 ), respectively,
[Equation]
for the extra update (skipping the update if ˆ y Tn e n ≤ 0 ).
Here, the quantities y n , ˆ y n and e n are deﬁned as in Algorithm LBFGS.
We can now argue essentially as in the proof of Byrd et al.
Theorem 3.1) to show that ( z n ) converges q-superlinearly to z (cid:63) .
As part of that proof we obtain that ˆ B n (cid:54) =
B n for at least (cid:100) 0 .
5 Q (cid:101) of the indices n = 0 , M, 2 M, . . .
, QM for any Q ∈ N (namely for all n ∈ N e satisfying n ≤ QM ) and that we can apply Byrd and Nocedal (1989, Theorem 3.2), which yields
[Equation]
For the third step, we abbreviate v n :=
∂g θ
∂θ | z n .
From the deﬁnition of e n and (10) we infer that
[Equation]
After multiplication with J g θ ( z (cid:63) ) − 1 this entails lim N
[Equation]
which shows that
[Equation]
by Assumption 2 (iii).
Using Assumption 2 (iii) again it follows that
[Equation]
concluding the third step.
To infer that (6) holds, it sufﬁces to show that lim
N e
[Equation]
n →∞
[Equation]
B n − B j n
[Equation]
= 0 for any sequence ( j n ) n ∈ N e ⊂ N such that { j n , j n +1 , . .
, n − 1 }∩ N e = ∅ for all n ∈ N e sufﬁciently large.
Indeed, since for C := max { sup n (cid:107) B n (cid:107) , sup n (cid:107) B − 1 n (cid:107)} , which is ﬁnite by Byrd and Nocedal (1989, Theorem 3.2), there holds
[Equation]
A
[Equation]
≤ C, (cid:107)
[Equation]
C
[Equation]
and the set on the right-hand side of the inclusion is compact by the Banach lemma, inversion is a uniformly continuous operation on this set, hence
lim N e
(cid:51)
n →∞ (cid:107) B − 1 n − B − 1 j n (cid:107) = 0 ,
so lim N e (cid:51) n →∞ (cid:107) p ( n )
[Equation]
− p
[Equation]
j n )
[Equation]
by continuity, and therefore
[Equation]
by the third step, establishing the claim.
It remains to show the validity of lim N e (cid:51) n →∞ (cid:107) B n − B j n (cid:107)
{+= 
+}
for any sequence ( j n ) n ∈ N e such that { j n , j n + 1 , . .
. , n − 1 } ∩ N e = ∅ for all n ∈ N e sufﬁciently large.
Since at least every second extra update is actually carried out, the condition on the intersection implies n − j n ≤ 2 M − 1 for all these n .
Now let ( j n ) n ∈ N e be any such sequence.
Then B n − B j n = (cid:80) n − 1 m =
j n B m +1 − B m is a sum of at most 2 M − 1 BFGS updates in search directions, but contains no extra updates.
Hence, the secant conditions B n − l s n − 1 − l = y n − 1 − l , l ∈ { 0 , 1 , . . . , n − j n } , are satisﬁed, allowing us to deduce
[Equation]
for all l ∈ { 0 , 1 , . . . , n − j n − 1 } .
{+For each of these l , both terms on the right-hand side tend to zero for N e (cid:51) n → ∞ (for the second term this follows from the ﬁrst identity in (10) due to B n − l − 
+}
= ˆ
B n − l − 1 ).
Recalling that B n − B j n = (cid:80) n − 1 m = j n B m +1 − B m we ﬁnd lim N e
[Equation]
n →∞ (cid:107) B n − B j n (cid:107) = 0 , which ﬁnishes the fourth step and thus concludes the proof.
B. C ONVERGENCE FOR A DJOINT B ROYDEN WITH OPA
Theorem 4 (Convergence of SHINE to the Hypergradient for Adjoint Broyden with OPA) .
Let us consider p ( n ) θ , the SHINE direction for iterate n in Algorithm 1 with the Adjoint Broyden secant condition (7) and extra update in the direction v n deﬁned in (8).
Under Assumptions 2 and 4, for a given parameter θ , we have q-superlinear convergence of ( z n ) to z (cid:63) and
[Equation]
Proof.
Due to Assumption 2, the superlinear convergence of ( z n ) follows from Schlenkrich et al.
(2010, Theorem 2).
The proof of the remaining claim is divided into two cases.
Case 1: Suppose that ∇ z L ( z (cid:63) )
= 0 .
By continuity this implies lim n →∞ ∇ z L ( z n ) = 0 .
Since the sequence ( B − 1 n ∂g θ ∂θ | z n ) is bounded by Assumption 4, it follows that
[Equation]
as claimed.
Case 2: Suppose that ∇ z L ( z (cid:63) ) (cid:54) = 0 .
{+By continuity this implies ∇ z L ( z n ) (cid:54) = 
+}
for all sufﬁciently large n ∈ N .
Let us denote by N e ⊂ N the set of indices of extra updates.
We stress that this set is inﬁnite since, by construction, every M -th update is an extra update.
{+We have v n (cid:54) = 
+}
for all sufﬁciently large n ∈ N e , hence Schlenkrich et al.
(2010, Lemma 3) yields
[Equation]
This implies
[Equation]
thus necessarily
[Equation]
Since lim N e
[Equation]
L
[Equation]
by continuity, we ﬁnd
[Equation]
whence lim
[Equation]
where we have used continuity again.
To prove that these limits hold not only for N e (cid:51) n → ∞ but in fact for all N (cid:51) n → ∞ , we establish, as intermediate claim, that for any ﬁxed m ∈ N we have lim n →∞ (cid:107) B n + m − B n (cid:107) = 0 .
Note that this claim is equivalent to lim n →∞ (cid:107) B n +1 − B n (cid:107) = 0 .
{+Denoting by L ≥ 
+}
the Lipschitz constant of J g θ near z (cid:63) , we ﬁnd
[Equation]
Both terms on the right-hand side go to zero as n goes to inﬁnity: the ﬁrst one due to lim n →∞ z n = z (cid:63) and the second one since lim n →∞ (cid:107) E Tn v n (cid:107) (cid:107) v n (cid:107)
{+= 
+}
by Schlenkrich et al.
(2010, Lemma 3).
This shows that lim n →∞ (cid:107) B n +1 − B n (cid:107) = 0 , which concludes the proof of the intermediate claim.
From lim n →∞ (cid:107) B n + m − B n (cid:107) = 0 for any ﬁxed m ∈ N it follows that for any sequence ( j n ) ⊂ N with sup n | j n − n | < ∞ there holds lim n →∞ (cid:107) B j n − B n (cid:107)
= 0 .
This implies for any such sequence ( j n ) the limit lim n →∞ (cid:107) B − 1 j n − B − 1 n (cid:107) = 0 .
To establish this, note that for
C := max { sup n (cid:107) B n (cid:107) , sup n (cid:107) B − 1 n (cid:107)} , which is ﬁnite by Assumption 4 and the combination of the bounded deterioration principle (Schlenkrich et al., 2010, Lemma 2) with Assumption 2 (i), the set (cid:110)
A ∈ R d × d : A − 1 exists , (cid:107) A (cid:107) ≤ C, (cid:107)
A − 1 (cid:107) ≤ C (cid:111) includes the sequence ( B n ) and is compact by the Banach lemma, so inversion is a uniformly continuous operation on this set.
Now let us construct a sequence ( j n ) ⊂
N e by deﬁning, for every n ∈ N , j n := arg min m ∈
N e | n − m | .
That is, for every n , j n denotes the member of N e with the smallest distance to n .
It is clear that | n − j n | ≤ M − 1 for all n , hence lim n →∞ (cid:107) B − 1 j n − B − 1 n (cid:107) = 0 .
Using this and, again, continuity it is easy to see that
[Equation]
which implies by (11) that
[Equation]
thereby establishing the claim.
Remark.
An inspection of the proof reveals that if B n is never updated in the direction z n , but only updated in the direction v n deﬁned in (8), then Assumption 4 can be replaced by the signiﬁcantly weaker assumption that the sequence ( B − 1 n ∂g θ ∂θ | z n ) is bounded.
The price to pay is that the convergence rate of ( z n ) to z (cid:63) will be slower (q-linear instead of q-superlinear) since the updates in the direction z n are critical for ensuring fast convergence of ( z n ) to z (cid:63) .
C L OGISTIC R EGRESSION H YPERPARAMETERS
For both datasets we split the data randomly (with a different seed for each run) between trainingvalidation-test, with the following proportions: 90%-5%-5%.
The hyperparameters are the same as in the original HOAG work (Pedregosa, 2016), except:
• We use a memory limitation of 30 updates (not grid-searched) for accelerated methods (Jacobian-Free and SHINE), compared to 10 for the original method.
This is because the approximation should be better using more updates.
We veriﬁed that using 30 updates for the original method does not improve the convergence speed.
That number is 60 for OPA. • We use a smaller exponential decrease of 0.78 (not grid-searched) for the accelerated methods, compared to 0.99 for the original method. This is because in the very long run, the approximation can cause oscillations. We also use the same setting as Pedregosa (2016) for the Grid and Random Search.
Finally, we highlight that warm restart is used for both the inner problem and the Hessian inversion in the direction of the gradient. OPA inversion experiments For the OPA experiments, we used a memory limitation of 60, and a tolerance of 10 − 6 .
The OPA update is done every 5 regular updates. D DEQ TRAINING DETAILS
The training details are the same as the original Multiscale DEQ paper (Bai et al., 2020): all the hyperparameters are kept the same and not ﬁne-tuned, and the data split is the same. We recall here some important aspects.
For both datasets, the network is ﬁrst trained in an unrolled weight-tied fashion for a few epochs in order to stabilize the training. We also underline that the DEQ models, in addition to having a ﬁxed-point-deﬁning sub-network, also have a classiﬁcation and a projection head. Finally, for Figure 3, the median backward pass is computed with 100 samples on a single V
GPU for a batch size of 32.
D. CIFAR
Adam optimizer (Kingma and Ba, 2015) is used with a 10 − 3 start learning rate, and a cosine annealing schedule. D. I MAGE N ET
The Stochastic Gradient Descent optimizer is used with a 5 × 10 − 2 start learning rate, and a cosine annealing schedule. The images are downsampled 2 times before being fed to the ﬁxed-point deﬁning sub-network. [Figure]
HOAG - lim
[Figure]
backward SHINE (ours) SHINE reﬁne ( [Figure]
Jacobian-
[Figure] E A DDITIONAL RESULTS E. B I - LEVEL OPTIMIZATION EXTENDED In order to make sure that SHINE was indeed improving over HOAG (Pedregosa, 2016), we also looked at the results obtained when performing an inversion with a precision lower than that prescribed by Pedregosa (2016) originally (i.e. truncating the iterative inversion).
These results, also complemented with Random Search (Bergstra and Bengio, 2012), can be seen in Figure E.1. They conﬁrm that the advantage provided by SHINE cannot be retrieved with a looser tolerance on the inversion.
E. C ONTRACTIVITY ASSUMPTION
One of the main limiting assumptions in the original Jacobian-Free method work (Fung et al., 2021), is the contractivity assumption.
We showed here that it was not important to enforce this in order to achieve excellent results, but one can wonder whether this assumption is not met in practice thanks to the unrolled pretraining of DEQs.
We looked at the contractivity of the ﬁxed-point deﬁning sub-network empirically by using the power-method applied to a nonlinear function, in the CIFAR setting. The results, summarized in Table E.1, show that the ﬁxed-point deﬁning sub-network is not contractive at all.
E. T IME GAINS
Because the total training time is not only driven by backward pass but also by the forward pass and the evaluation, we show for completeness in Table E.2 the time gains for the different acceleration methods for the overall epoch. We do not report in this table the time taken for pre-training which is [Table] [Table]
equivalent across all methods, and is not something on which SHINE has an impact. It is clear in
Table E.2 that accelerated methods can have a signiﬁcant impact on the training of DEQs because we see that half the time of the total pass is spent on the backward pass (more on ImageNet (Deng et al., 2009)).
We also notice that while SHINE has a slightly slower backward pass than the Jacobian-Free method (Fung et al., 2021), the difference is negligible when compared to the total pass computational cost. E. DEQ OPA RESULTS We can clearly see in Figure E.2 that in the case of DEQs, OPA also signiﬁcantly improves the inversion over the other accelerated methods. We also see that the improvements of SHINE over the
Jacobian-Free method without OPA are marginal.
Because the inversion is so good, we would expect that the performance of SHINE with OPA would be on par with the original method’s.
However, this is not what we see in the results presented in Table E.3.
Indeed, OPA does improve on SHINE with only Adjoint Broyden, but it does not outperform SHINE done with Broyden.
[Table]
[Table]
(Broyden)
[Table]
SHINE (
[Table]
. Broyden)
[Table]
SHINE (Adj. Broyden/OPA) [Table] [Figure] c o ss i m ( a , b )
Method
Jacobian-Free
SHINE w. Broyden
SHINE w. Adj.
Broyden
SHINE w. Adj.
Broyden / OPA