Year: 2022

Abstract:
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-efficient, they are still significantly 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 efficiently 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 and the recent Jacobian-Free method in different settings, ranging from hyperparameter optimization to large Multiscale DEQs (MDEQs) applied to CIFAR and ImageNet. Both methods reduce significantly the computational cost of the backward pass. While SHINE has a clear advantage on hyperparameter optimization problems, both methods attain similar computational performances for larger scale problems such as MDEQs at the cost of a limited performance drop compared to the original models.

heading: 1 INTRODUCTION
text: 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 infinite 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 bear an O(1) memory cost and with constraints on the architecture that are usually not detrimental to the 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 (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:
arg min θ L(z?) subject to gθ(z?) = 0. (1)
We will refer to the root finding problem gθ(z?) = 0 as the inner problem, and call its resolution the forward pass. On the other hand, we will refer to arg minθ L(z?) as the outer problem, and
call the computation of the gradient of L(z?) w.r.t. θ the backward pass. The core idea for DEQs is that their output z? is expressed as a fixed point of a parametric function fθ from Rd to Rd, i.e., gθ(z
?) = z?−fθ(z?) = 0.1 This model is said to have infinitely many weight-tied layers as z? can be obtained by successively applying the layer fθ infinitely 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 efficiently the Jacobian matrix ∂gθ∂z and its inverse for root-finding. To compute DEQs’ gradient efficiently 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? with respect to θ, ∂z ?
∂θ . While this method is memory efficient, 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.4). 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
min θ Lval(z∗) subject to z∗ = min z rθ(z) , Ltrain(z) + θ‖z‖22, (2)
where Ltrain and Lval 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θ = ∇zrθ to fit (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 final 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 efficiently. We emphasize that the goal of this paper is neither to improve the algorithms used to compute z?, 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 first 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 specific 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:
1Here, we do not explicitly write the dependence of fθ on the input x of the DEQ, usually referred to as the injection.
• 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 (SHaring the INverse Estimate).
• 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 classification 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 refinement strategy for approximate Jacobian inversion methods (both SHINE and Jacobian-Free) that allows a trade-off between computational cost and performances.

heading: 2 HYPERGRADIENT OPTIMIZATION WITH APPROXIMATE JACOBIAN INVERSE

heading: 2.1 SHINE: HYPERGRADIENT DESCENT WITH APPROXIMATE JACOBIAN INVERSE
text: Algorithm 1: qN method to solve gθ(z?) = 0 Result: Root z?, qN matrix B b = true if using Broyden’s method, b = false if using BFGS n = 0, z0 = 0, B0 = I while not converged do
pn = −B−1n gθ(zn), zn+1 = zn + αnpn // αn can be 1 or determined by line-search yn = gθ(zn+1)− gθ(zn) sn = zn+1 − zn if b then
Bn+1 = arg min X: Xsn=yn ‖X −Bn‖F else
Bn+1 = arg min
X: X=XT ∧Xsn=yn ‖X−1 −B−1n ‖ // The norm used in BFGS is a weighted Frobenius norm
end n← n+ 1
end z? = zn, B = Bn
Hypergradient Optimization Hypergradient optimization is a first-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 Jgθ 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 θ ∈ Rp be a set of parameters, let L : Rd → R be a loss function and gθ : Rd → Rd be a root-defining function. Let z? ∈ Rd such that gθ(z?) = 0 and Jgθ (z
?) is invertible, then the gradient of the loss L wrt. θ, called Hypergradient, is given by ∂L ∂θ ∣∣∣ z? = ∇zL(z?)>Jgθ (z?)−1 ∂gθ ∂θ ∣∣∣ z? . (3)
In practice, we use an algorithm to approximate z?, 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?. 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? has been obtained, one of the major bottlenecks in the computation of the Hypergradient is the inversion of Jgθ (z ?) in the
directions ∂gθ∂θ ∣∣∣ z? or∇zL(z?).
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 LBFGS (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 limited-memory version (Bai et al., 2020). These quasi-Newton methods were first inspired by Newton’s method, which finds the root of gθ via the recurrent Jacobian-based updates zn+1 = zn − Jgθ (zn)−1gθ(zn). Specifically, they replace the Jacobian Jgθ (zn) by an approximation Bn that is based on available values of the iterates zn and gθ rather than its derivative. These Bn, called qN matrices, are defined recursively via an optimization problem with constraints called secant conditions. Solving this problem leads to expressing Bn as a rank-one or rank-two update of Bn−1, so that Bn is the sum of the initial guess B0 (in our settings, the identity) and n low-rank matrices (less than n in limited memory settings). This low rank structure allows efficient multiplication by Bn and B−1n . 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 = limn→∞B−1n as a replacement for Jgθ (z
?)−1 in (3), i.e. to share the inverse estimate between the forward and the backward passes. This gives the approximate Hypergradient
pθ = ∇zL(z?)B−1 ∂gθ ∂θ ∣∣∣ z? . (4)
In practice we will consider the nonasymptotical direction p(n)θ = ∇zL(zn)B−1n ∂gθ∂θ ∣∣∣ zn . Thanks to the Sherman-Morrison formula (Sherman and Morrison, 1950), the inversion of Bn can be done very efficiently (using scalar products) compared to the iterative methods needed to invert the true Jacobian Jgθ (z
?). In turn, this significantly reduces the computational cost of the Hypergradient computation.
Relationship to the Jacobian-Free method Because B0 = 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. (2021). In that work, Jgθ (z
?)−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 Jgθ (z
?) in the direction∇zL(z?). 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 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 refine strategy. Because the refine strategy is essentially a smart initialization scheme, it recovers all the theoretical guarantees of the original method (Bai et al., 2019; 2020; Pedregosa, 2016).

heading: 2.2 CONVERGENCE TO THE TRUE GRADIENT
text: To further justify and formalize the idea of SHINE, we show that the direction p(n)θ converges to the
Hypergradient ∂L∂θ ∣∣∣ z? . 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 n0 ≥ 0 and m ≥ d with the following property: For any n ≥ n0 we can find indices n ≤ n1 ≤ . . . ≤ nd ≤ n+m such that, for pn defined in Algorithm 1, the smallest singular value of the d× d matrix (
pn1 ‖pn1‖ , pn2 ‖pn2‖ , . . . , pnd ‖pnd‖ ) is no smaller than ρ.
Assumption 2 (Smoothness and convergence to the fixed point). (i) ∑∞ n=0 ‖zn− z?‖ <∞ for some z? with gθ(z?) = 0; (ii) gθ is C1, Jgθ is Lipschitz continuous near z ?, and Jgθ (z
?) is invertible; (iii) ∇zL is continuous, and ∀θ, ∂gθ∂θ is continuous. Remark. The Assumption 2 (i) implies limn→∞ zn = z?. 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 θ, (zn) converges q-superlinearly to z? and
lim n→∞
p (n) θ = ∂L ∂θ ∣∣∣ z? .
Proof. From More and Trangenstein (1976, Theorem 5.7) we obtain that limn→∞Bn = Jgθ (z ?). 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.

heading: 2.3 OUTER PROBLEM AWARENESS
text: The ULI assumption guarantees convergence of B−1n to Jgθ (z
?)−1. However, (3) only requires the multiplication of Jgθ (z
?)−1 with ∂gθ∂θ |z? from the right and∇zL(z?) 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? is a d-dimensional vector and it is possible to compute its approximation ∂gθ∂θ |zn at a reasonable cost. We propose to incorporate additional updates of the quasi-Newton matrix Bn into Algorithm 1 that improve the approximation quality of B−1n in the direction ∂gθ ∂θ |zn (thus asymptotically in the direction ∂gθ∂θ |z? ). Given a current iterate pair (zn, Bn), these additional updates only change Bn, but not zn. We will demonstrate that a suitable update direction en ∈ Rd is given by
en = tnB −1 n ∂gθ ∂θ ∣∣∣ zn , (5)
where (tn) ⊂ [0,∞) satisfies ∑ n tn < ∞. This update direction will be used to create an extra secant condition X−1(gθ(zn + en)− gθ(zn)) = en for the additional update of Bn. 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 fixed frequency, then fast (q-superlinear) convergence of (zn) to z? 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) = ∇zrθ(z) for some C2 function rθ : Rd → 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 en defined in (5). Under Assumptions 2 (ii-iii) and 3, for a given parameter θ, we have the following: Algorithm 1, for any symmetric and positive definite matrix B0, generates a sequence (zn) that converges q-superlinearly to z?, and there holds
lim n→∞
p (n) θ = ∂L ∂θ ∣∣∣ z? . (6)
Proof. It follows from known results that the extra updates do not destroy the q-superlinear convergence of (zn). The proof of (6) relies firstly on the fact that by continuity of the derivative of gθ, we have limn→∞ ∂gθ∂θ |zn = ∂gθ ∂θ |z? . Due to the extra updates we can show convergence of the qN matrices to the true Hessian in the direction of the extra steps en, 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 C2 function rθ (such that gθ(z) = ∇zrθ(z)) with locally Lipschitz continuous Hessian if z0 is close enough to some z? with∇zrθ(z?) = 0 and ∇2zzrθ(z?) positive definite. We note that Theorem 3 guarantees fast convergence of the iterates (zn) and that z0 does not have to be close to z? for that guarantee. Also, there is no restriction on B0 other than being symmetric and positive definite (which is satisfied for our choice B0 = 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−1n approximates Jgθ (z ?)−1 in the direction ∇zL(z?) applied by left-multiplication, as required by (3). An appropriate secant condition is given by
vTnBn+1 = v T n Jgθ (zn+1), (7)
where vTn = ∇zL(zn)B−1n . (8)
To incorporate the secant condition (7), we use the Adjoint Broyden’s method (Schlenkrich et al., 2010), a qN method relying on the efficient vector-Jacobian multiplication by Jgθ 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 (Bn) generated by Algorithm 1 satisfies
sup n∈N ‖B−1n ‖ <∞.
Remark. Convergence results for quasi-Newton methods usually include showing that Assumption 4 holds, 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 limBn = Jgθ (z ?) and thus limB−1n = Jgθ (z ?)−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 vn defined in (8). Under Assumptions 2 and 4, for a given parameter θ, we have q-superlinear convergence of (zn) to z? and
lim n→∞
p (n) θ = ∂L ∂θ ∣∣∣ z? .
Proof. The q-superlinear convergence of (zn) 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 ∇zL(z?) = 0 the claim holds due to continuity and Assumption 4. Then ∇zL(z?) 6= 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 fixed 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 Jgθ , 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.

heading: 3 RESULTS
text: 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 guarantees2 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 fixed point convergence), it performs well in practice.
2See the results on contractivity in Section E.3.
Implementations. All the bi-level optimization experiments were done using the HOAG code (Pedregosa, 2016)3, which is based on the Python scientific 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 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, on nodes with 4 V100 GPUs. In practice, we never reach convergence of (zn), 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 verified 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.

heading: 3.1 BI-LEVEL OPTIMIZATION – HYPERPARAMETER OPTIMIZATION IN LOGISTIC REGRESSION
text: We first test SHINE in the simple setting of bi-level optimization for `2-regularized LR, using the code from Pedregosa (2016) and the same datasets. Convergence on unseen data is illustrated in Figure 1.5 An acceptable level of performance is reached twice faster for the SHINE method compared to any other competitor. Another finding is that the refine strategy does not provide a definitive 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
3https://github.com/fabianp/hoag 4https://github.com/locuslab/mdeq 5To facilitate the reader’s understanding of the figures, we plot the empirical suboptimality, but we do remind them that there is no guarantee of convergence on held-out test data ; the kink present in the case of the real-sim dataset is an example of that.
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. Additional results on hyperparameter optimization for the regularized nonlinear least squares problem are available in Section E.2. 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 ∂gθ ∂z ∣∣∣ z?
(zn − zn−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.

heading: 3.2 DEEP EQUILIBRIUM MODELS
text: Next, we tested SHINE on the more challenging DEQ setup. Two experiments illustrate the performance of SHINE on the image classification task on two datasets. For both datasets, we used the same model configuration as in the original Multiscale DEQ paper (Bai et al., 2020) and did not fine 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.3), we observe a sharp initial performance drop that reduces its performance on Imagenet. We provide partial results in Section E.5.
CIFAR-10. The first 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 fixed point is d = 50k. We train the models for five 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.2 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 fixed point to find is d = 190k.
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. 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 refined version is able to perform on par with the Jacobian-Free method (Fung et al., 2021). We also confirm 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 significant when applied in their vanilla version, we remind the reader that no fine-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.6 Moreover, using the proposed refined versions of these methods, the performance drop can be traded-off for acceleration.

heading: 4 CONCLUSION AND DISCUSSION
text: 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 (zn) to z? 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 significantly weaker assumption will suffice.
6More on the overall computational effort can be found in Table E.2

heading: REPRODUCIBILITY STATEMENT
text: We provide with the submission of this paper the full code necessary to reproduce the figures and the other quantitative results of the paper, from the training, to the evaluation and the actual figure drawing. We made sure to use seeds and verified that the seeding was indeed allowing reproducible results. We also provide time estimates for the reproduction of the figures. 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.

heading: A OPA ALGORITHM
text: Algorithm LBFGS: (Limited memory) BFGS method with OPA Input: initial guess (z0, B−10 ), where B −1 0 is symmetric and positive definite, tolerance > 0,
frequency of additional updates M ∈ N, memory limit L ∈ N ∪ {∞}, (tn) a null sequence of positive numbers with ∑ n tn <∞

heading: Let F := ∇zgθ
text: for n = 0, 1, 2, . . . do
if ‖F (zn)‖ ≤ then let z? := zn and let B := Bn; STOP Let B̂−1n := B −1 n if (nmodM) = 0 then let en := tnB−1n ∂gθ ∂θ ∣∣∣ zn
, ŷn := F (zn + en)− F (zn) and r̂n := (en)T ŷn if r̂n > 0 then
let ân := en −B−1n ŷn and let
B̂−1n := B −1 n +
ân(en) T + en(ân) T
r̂n − (ân) T ŷn (r̂n)2 en(en) T
Let B−1n := B̂ −1 n if n ≥ L then remove update n− L from B−1n Let pn := −B−1n F (zn) Obtain αn via line-search and let sn := αnpn Let zn+1 := zn + sn, yn := F (zn+1)− F (zn) and rn := (sn)T yn if rn > 0 then
let an := sn −B−1n yn and let
B−1n+1 := B −1 n +
an(sn) T + sn(an) T
rn − (an) T yn (rn)2 sn(sn) T
else let B−1n+1 := B−1n if n ≥ L then remove update n− L from B−1n+1
Output: z?, B
Remark. A possible choice for (tn) is to use an arbitrary t0 > 0 and tn := ‖sn−1‖ for n ≥ 1.

heading: B PROOFS OF SHINE CONVERGENCE
text: To facilitate reading, we restate the results before proving them.

heading: B.1 CONVERGENCE USING ULI
text: 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 θ, (zn) converges q-superlinearly to z? and
lim n→∞
p (n) θ = ∂L ∂θ ∣∣∣ z? .
Proof. Under Assumptions 1 and 2, More and Trangenstein (1976, Theorem 5.7) shows that Bn satisfies
lim n→∞
Bn = Jgθ (z ?)
The inversion operator is continuous in the space of invertible matrices, so we have: lim n→∞ B−1n = Jgθ (z ?)−1
Because ∇zL and ∂gθ∂θ are continuous at z? by Assumption 2 (iii), we also have thanks to Assumption 2 (i):
lim n→∞ ∇zL(zn) = ∇zL(z?) and lim n→∞ ∂gθ ∂θ ∣∣∣ zn = ∂gθ ∂θ ∣∣∣ z?
By continuity we then deduce that, as claimed,
lim n→∞ p (n) θ = limn→∞ ∇zL(zn)B−1n ∂gθ ∂θ (zn) = ∇zL(z?)Jgθ (z?)−1 ∂gθ ∂θ ∣∣∣ z? = ∂L ∂θ ∣∣∣ z?

heading: B.2 CONVERGENCE FOR BFGS WITH OPA
text: Assumption 5 (Extended Assumptions for BFGS). Let gθ(z) = ∇zrθ(z) for some C2 function rθ : Rd → R. Consider Algorithm 1 with b = false and suppose that
1. the set Ω := {z ∈ Rd : rθ(z) ≤ rθ(z0)} is convex; 2. rθ is strongly convex in an open superset of Ω (this implies that rθ has a unique global
minimizer z?) and has a Lipschitz continuous Hessian near z?; 3. there are positive constants η1, η2 such that the line search used in the algorithm ensures
that for each n ≥ 0 either rθ(zn+1) ≤ rθ(zn)− η1 [∇rθ(zn)T pn
‖pn‖
]2 or rθ(zn+1) ≤ rθ(zn) + η2∇rθ(zn)T pn
is satisfied; 4. the line search has the property that αn = 1 will be used if both
‖(Bn − Jgθ (zn))sn‖ ‖sn‖
and ‖zn − z?‖
are sufficiently small. Remark. The requirements 3. and 4. on the line search are, for instance, satisfied 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 en defined in (5). Under Assumptions 2 (ii-iii) and 3, for a given parameter θ, we have the following: Algorithm 1, for any symmetric and positive definite matrix B0, generates a sequence (zn) that converges q-superlinearly to z?, and there holds
lim n→∞
p (n) θ = ∂L ∂θ ∣∣∣ z? . (6)
Proof. The proof is divided into four steps. The first step is to establish the q-superlinear convergence of (zn) to z?. Denoting by Ne ⊂ {0,M, 2M, . . .} the set of indices of extra updates that are actually applied, the second step consists of showing
lim Ne3n→∞ (Bn − Jgθ (z?)) en ‖en‖ = 0, (9)
where, in this proof, Bn always represents the matrix from Algorithm LBFGS before the update in the direction en is applied, i.e., the matrix whose inverse appears in the definition of en, while B̂n always represents the matrix from Algorithm LBFGS after the update in the direction en has been applied; if the update in the direction en is not applied, then Bn = 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 Ne 3 n → ∞, i.e., the limit is taken on the subsequence corresponding to Ne. The fourth step is then to transfer the convergence to the entire sequence. It is easy to check that instead of updating B−1n , respectively, B̂ −1 n , we can also obtain the sequences (Bn) and (B̂n) by updating according to
Bn+1 = Bn + yny
T n
yTn sn − Bnsn(Bnsn)
T
sTnBnsn
for the usual update (skipping the update if yTn sn ≤ 0), respectively,
B̂n = Bn + ŷnŷ
T n
ŷTn en − Bnen(Bnen)
T
eTnBnen
for the extra update (skipping the update if ŷTn en ≤ 0). Here, the quantities yn, ŷn and en are defined as in Algorithm LBFGS. We can now argue essentially as in the proof of Byrd et al. (1988,
Theorem 3.1) to show that (zn) converges q-superlinearly to z?. As part of that proof we obtain that B̂n 6= Bn for at least d0.5Qe of the indices n = 0,M, 2M, . . . , QM for any Q ∈ N (namely for all n ∈ Ne satisfying n ≤ QM ) and that we can apply Byrd and Nocedal (1989, Theorem 3.2), which yields
lim n→∞
( B̂n − Jgθ (z?) ) sn ‖sn‖ = 0 and lim Ne3n→∞ (Bn − Jgθ (z?)) en ‖en‖ = 0. (10)
For the third step, we abbreviate vn := ∂gθ∂θ |zn . From the definition of en and (10) we infer that
0 = lim Ne3n→∞ (Bn − Jgθ (z?)) en ‖en‖ = lim Ne3n→∞
( I − Jgθ (z?)B−1n ) vn ‖B−1n vn‖ .
After multiplication with Jgθ (z ?)−1 this entails
lim Ne3n→∞
( Jgθ (z ?)−1 −B−1n ) vn ‖B−1n vn‖ = 0,
which shows that
lim Ne3n→∞ B−1n vn = lim Ne3n→∞ Jgθ (z ?)−1vn = Jgθ (z ?)−1 ∂gθ ∂θ |z?
by Assumption 2 (iii). Using Assumption 2 (iii) again it follows that
lim Ne3n→∞ p (n) θ = limNe3n→∞ ∇zL(zn)B−1n ∂gθ ∂θ ∣∣∣ zn = ∇zL(z?)Jgθ (z?)−1 ∂gθ ∂θ ∣∣∣ z? = ∂L ∂θ ∣∣∣ z? ,
concluding the third step. To infer that (6) holds, it suffices to show that limNe3n→∞ ‖Bn−Bjn‖ = 0 for any sequence (jn)n∈Ne ⊂ N such that {jn, jn+1, . . . , n−1}∩Ne = ∅ for all n ∈ Ne sufficiently large. Indeed, since for C := max{supn ‖Bn‖, supn ‖B−1n ‖}, which is finite by Byrd and Nocedal (1989, Theorem 3.2), there holds
(Bn) ⊂ { A ∈ Rd×d : A−1 exists , ‖A‖ ≤ C, ‖A−1‖ ≤ C } 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 limNe3n→∞ ‖B−1n −B−1jn ‖ = 0, so
lim Ne3n→∞
‖p(n)θ − p (jn) θ ‖ = 0
by continuity, and therefore
lim Ne3n→∞ p (jn) θ = limNe3n→∞ p (n) θ = ∂L ∂θ ∣∣∣ z?
by the third step, establishing the claim. It remains to show the validity of limNe3n→∞ ‖Bn − Bjn‖ = 0 for any sequence (jn)n∈Ne such that {jn, jn + 1, . . . , n− 1} ∩Ne = ∅ for all n ∈ Ne sufficiently large. Since at least every second extra update is actually carried out, the condition on the intersection implies n− jn ≤ 2M − 1 for all these n. Now let (jn)n∈Ne be any such sequence. Then Bn −Bjn = ∑n−1 m=jn
Bm+1 −Bm is a sum of at most 2M − 1 BFGS updates in search directions, but contains no extra updates. Hence, the secant conditions Bn−lsn−1−l = yn−1−l, l ∈ {0, 1, . . . , n − jn}, are satisfied, allowing us to deduce
‖Bn−l −Bn−l−1‖ = ‖(Bn−l −Bn−l−1)sn−l−1‖
‖sn−l−1‖
≤ ‖yn−l−1 − Jgθ (z ?)sn−l−1‖ ‖sn−l−1‖ + ‖(Bn−l−1 − Jgθ (z?))sn−l−1‖
‖sn−l−1‖ for all l ∈ {0, 1, . . . , n− jn − 1}. For each of these l, both terms on the right-hand side tend to zero for Ne 3 n→∞ (for the second term this follows from the first identity in (10) due to Bn−l−1 = B̂n−l−1). Recalling that Bn −Bjn = ∑n−1 m=jn
Bm+1 −Bm we find limNe3n→∞ ‖Bn −Bjn‖ = 0, which finishes the fourth step and thus concludes the proof.

heading: B.3 CONVERGENCE FOR ADJOINT BROYDEN WITH OPA
text: 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 vn defined in (8). Under Assumptions 2 and 4, for a given parameter θ, we have q-superlinear convergence of (zn) to z? and
lim n→∞
p (n) θ = ∂L ∂θ ∣∣∣ z? .
Proof. Due to Assumption 2, the superlinear convergence of (zn) follows from Schlenkrich et al. (2010, Theorem 2). The proof of the remaining claim is divided into two cases. Case 1: Suppose that ∇zL(z?) = 0. By continuity this implies limn→∞∇zL(zn) = 0. Since the sequence (B−1n ∂gθ ∂θ |zn) is bounded by Assumption 4, it follows that
lim n→∞ p (n) θ = limn→∞ ∇zL(zn)B−1n ∂gθ ∂θ ∣∣∣ zn = 0 = ∂L ∂θ ∣∣∣ z? ,
as claimed. Case 2: Suppose that ∇zL(z?) 6= 0. By continuity this implies ∇zL(zn) 6= 0 for all sufficiently large n ∈ N. Let us denote by Ne ⊂ N the set of indices of extra updates. We stress that this set is infinite since, by construction, every M -th update is an extra update. We have vn 6= 0 for all sufficiently large n ∈ Ne, hence Schlenkrich et al. (2010, Lemma 3) yields
lim Ne3n→∞ ‖∇zL(zn)(I −B−1n Jgθ (z?))‖ ‖(∇zL(zn)B−1n )T ‖ = lim Ne3n→∞ ‖(vn)T (Bn − Jgθ (z?))‖ ‖vn‖ = 0.
This implies
lim Ne3n→∞ ‖∇zL(zn)(Jgθ (z?)−1 −B−1n )‖ ‖∇zL(zn)B−1n ‖ = 0,
thus necessarily lim
Ne3n→∞ ‖∇zL(zn)(Jgθ (z?)−1 −B−1n )‖ = 0.
Since limNe3n→∞∇zL(zn)Jgθ (z?)−1 = ∇zL(z?)Jgθ (z?)−1 by continuity, we find lim
Ne3n→∞ ∇zL(zn)B−1n = ∇zL(z?)Jgθ (z?)−1,
whence
lim Ne3n→∞ p (n) θ = limNe3n→∞ ∇zL(zn)B−1n ∂gθ ∂θ ∣∣∣ zn = ∇zL(z?)Jgθ (z?)−1 ∂gθ ∂θ ∣∣∣ z? = ∂L ∂θ ∣∣∣ z? , (11)
where we have used continuity again. To prove that these limits hold not only for Ne 3 n → ∞ but in fact for all N 3 n → ∞, we establish, as intermediate claim, that for any fixed m ∈ N we have limn→∞ ‖Bn+m −Bn‖ = 0. Note that this claim is equivalent to limn→∞ ‖Bn+1 −Bn‖ = 0. Denoting by L ≥ 0 the Lipschitz constant of Jgθ near z?, we find
‖Bn+1 −Bn‖ = ‖vnvTn [Jgθ (zn+1)−Bn] ‖
‖vn‖2 ≤ ‖Jgθ (zn+1)− Jgθ (z?)‖+ ‖ [Jgθ (z?)−Bn]T vn‖ ‖vn‖
≤ L‖zn+1 − z?‖+ ‖ETn vn‖ ‖vn‖ .
Both terms on the right-hand side go to zero as n goes to infinity: the first one due to limn→∞ zn = z? and the second one since limn→∞ ‖ETn vn‖ ‖vn‖ = 0 by Schlenkrich et al. (2010, Lemma 3). This shows that limn→∞ ‖Bn+1 −Bn‖ = 0, which concludes the proof of the intermediate claim. From limn→∞ ‖Bn+m − Bn‖ = 0 for any fixed m ∈ N it follows that for any sequence (jn) ⊂ N with supn |jn − n| < ∞ there holds limn→∞ ‖Bjn − Bn‖ = 0. This implies for any such sequence (jn) the limit limn→∞ ‖B−1jn − B−1n ‖ = 0. To establish this, note that for C := max{supn ‖Bn‖, supn ‖B−1n ‖}, which is finite by Assumption 4 and the combination of the bounded deterioration principle (Schlenkrich et al., 2010, Lemma 2) with Assumption 2 (i), the set{
A ∈ Rd×d : A−1 exists , ‖A‖ ≤ C, ‖A−1‖ ≤ C }
includes the sequence (Bn) and is compact by the Banach lemma, so inversion is a uniformly continuous operation on this set. Now let us construct a sequence (jn) ⊂ Ne by defining, for every n ∈ N, jn := arg minm∈Ne |n−m|. That is, for every n, jn denotes the member of Ne with the smallest distance to n. It is clear that |n− jn| ≤M − 1 for all n, hence limn→∞ ‖B−1jn −B−1n ‖ = 0. Using this and, again, continuity it is easy to see that
lim n→∞
‖p(n)θ − p (jn) θ ‖ = 0,
which implies by (11) that
lim n→∞ p (n) θ = limn→∞ p (jn) θ = limNe3n→∞ p (n) θ = ∂L ∂θ ∣∣∣ z? ,
thereby establishing the claim.
Remark. An inspection of the proof reveals that if Bn is never updated in the direction zn, but only updated in the direction vn defined in (8), then Assumption 4 can be replaced by the significantly weaker assumption that the sequence (B−1n ∂gθ ∂θ |zn) is bounded. The price to pay is that the convergence rate of (zn) to z? will be slower (q-linear instead of q-superlinear) since the updates in the direction zn are critical for ensuring fast convergence of (zn) to z?.

heading: C LOGISTIC REGRESSION HYPERPARAMETERS
text: 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 verified 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.

heading: D DEQ TRAINING DETAILS
text: 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 fine-tuned, and the data split is the same. We recall here some important aspects. For both datasets, the network is first 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 fixed-point-defining sub-network, also have a classification and a projection head. Finally, for Figure 3, the median backward pass is computed with 100 samples on a single V100 GPU for a batch size of 32.

heading: D.1 CIFAR
text: The Adam optimizer (Kingma and Ba, 2015) is used with a 10−3 start learning rate, and a cosine annealing schedule.
D.2 IMAGENET 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 fixed-point defining sub-network.
0 10 20 30 40 50 60
Time (s)
10−1
100
101
102
103 20news
0 5 10 15 20 25 30
Time (s)
10−1
100
101
102
103
real-sim
HOAG
HOAG - lim. backward
SHINE (ours)
SHINE refine (ours)
Grid search
Random search
Jacobian-Free
T es
t L
o ss
S u
b o p
ti m
al it
y
Figure E.1: Bi-level optimization: Convergence of different hyperparameter optimization methods on the `2-regularized logistic regression problem for two datasets (20news (Lang, 1995) and realsim (lib)) on held-out test data.

heading: E ADDITIONAL RESULTS

heading: E.1 BI-LEVEL OPTIMIZATION EXTENDED
text: 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 confirm that the advantage provided by SHINE cannot be retrieved with a looser tolerance on the inversion.

heading: E.2 REGULARIZED NONLINEAR LEAST SQUARES
text: In order to further validate the efficiency of SHINE compared to competing methods, we also benchmarked it on the regularized nonlinear least squares task. For a training set (xtrain,i, ytrain,i)Ni=1 and a test set (xtest,i, ytest,i)Mi=1, this problem reads
min θ
1
2 M∑ i=1 ‖ytest,i − σ((z∗)>xtest,i)‖22
z∗ = arg min z
1
2 N∑ j=1 ‖ytrain,j − σ(z>xtrain,j)‖22 + θ 2 ‖z‖22
(12)
where σ denotes the sigmoid function σ(x) = 11+e−x . For a fixed hyper-parameter θ, this task is typically solved using L-BFGS (Berahas et al., 2021; Xu et al., 2020).
0 100 200 300 400 500 600
Time (s)
10−1
100
101
102
20news
SHINE (ours)
HOAG
Jacobian-Free SHINE refine (ours) SHINE - OPA (ours)
T es
t L
os s
S u
b op
ti m
al it
y
Figure E.2: Bi-level optimization on regularized nonlinear least squares: Convergence of different hyperparameter optimization methods on the `2-regularized nonlinear least squares for the 20news (Lang, 1995) dataset on held-out test data.
Table E.1: Nonlinear spectral radius obtained by the power method for the fixed-point defining sub-network for the 3 different methods.
Method Nonlinear spectral radius Original 230.5 Jacobian-Free 193.7 SHINE 234.2
We can see in Figure E.2 that SHINE clearly outperforms the Jacobian-Free method and it is also quicker to converge compared to HOAG. We can also notice the benefit of OPA compared to the vanilla SHINE method is more pronounced. We hypothesize that this is due to the nonconvex nature of the inner problem making the Hessian inverse approximation more difficult, as was noted by Berahas et al. (2021).

heading: E.3 CONTRACTIVITY ASSUMPTION
text: 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 fixed-point defining 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 fixed-point defining sub-network is not contractive at all.

heading: E.4 TIME GAINS
text: 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 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 significant 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.
Table E.2: The time required for each method on the different datasets during the equilibrium training. For the forward and backward passes, the time is measured offline, for a single batch of 32 samples, with a single GPU, using the median to avoid outliers. This time is given in milliseconds. For the epochs, the time is measured by taking an average of the 6 first epochs, and given in hours-minutes for Imagenet and minutes-seconds for CIFAR. The epoch time for SHINE without improvement on Imagenet is not given because it never reaches the 26 forward steps: the implicit depth is too short. Fallback is not used for CIFAR. Numbers in parenthesis indicate the number of inversion steps for the refined versions.
Dataset Name CIFAR (Krizhevsky, 2009) ImageNet (Deng et al., 2009) Method Name Forward Backward Epoch Forward Backward Epoch Original (Bai et al., 2020) 256 210 4min40 644 798 3h38 Jacobian-Free (Fung et al., 2021) 249 12.9 3min10 621 13.5 2h02 SHINE Fallback (ours) 218 16.0 3min20 622 35.3 2h13 SHINE Fallback refine (5, ours) 272 96.6 3min50 622 212 2h44 Jacobian-Free refine (5) 260 86.5 3min40 620 186 2h43 Original limited backprop 281 86.4 3min50 653 187 2h40
1.2 1.4 1.6 1.8 2.0
‖a‖/‖b‖
0.5
0.6
0.7
0.8
0.9
co ss
im (a ,b
) Method
Jacobian-Free
SHINE w. Broyden
SHINE w. Adj. Broyden
SHINE w. Adj. Broyden / OPA
Figure E.3: Quality of the inversion using OPA in DEQs : Ratio of the inverse approximation over the exact inverse function of the cosine similarity between the inverse approximation b = ∇zL(z?)B−1n and the exact inverse a = ∇zL(z?)Jgθ (z?)−1 for different methods. For OPA, the extra update frequency is 5. 100 runs were performed with different batches.

heading: E.5 DEQ OPA RESULTS
text: We can clearly see in Figure E.3 that in the case of DEQs, OPA also significantly 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 E.3: CIFAR DEQ OPA results : Top-1 accuracy of different methods on the CIFAR dataset, and epoch mean time.
Methode name Top-1 Accuracy (%) Epoch mean time Original 93.51 4min40 Jacobian-Free 93.09 3min10 SHINE (Broyden) 93.14 3min20 SHINE (Adj. Broyden) 92.89 4min SHINE (Adj. Broyden/OPA) 93.04 4min40

Review summaries:
In implicit deep learning such as deep equilibrium models, computing the inverse Jacobian for the forward pass is computationally expensive. This paper propose an interesting approach to combine the information from the forward and backward pass to make an efficient estimate of the Jacobian inverse. In one approach, they propose to replace the Jacobian in the backward update with the quasi-Newton matrix, which is being already used/estimated in the forward pass solved by quasi-Newton method. Additionally, they propose an iterative update to the quasi-Newton matrix such that to helps its estimate toward the direction useful in the backward pass (they call this outer problem awareness).  They provide theoretical analysis of their proposed method and show that under certain conditions/assumptions, the forward pass still converges to the desired solution and the sequences of backward estimates converges to the loss gradient of needed to parameter updates. They provide numerical results in bi-level optimization (regularized logistic regression) and training DEQ for classification. In certain settings (bi-level optimization), they show that they outperform Jacobian-free and have similar performance to the state-of-the-arts but it is faster than all. For DEQ, they show similar performance to Jacobian-free.

In various machine learning problems which can be formulated as a bilevel optimization problem and solved using gradient-based methods, the computation of hypergradients is necessary. However, the involved inverse Jacobian matrix in the hypergradient has been a computational bottleneck in high-dimensional settings. This paper proposes to use quasi-Newton matrices from the forward pass to approximate this inverse Jacobian matrix in the direction needed for the gradient computation which appears in the computation of hypergradients. The proposed algorithm is applied to both hyperparameter optimization and deep equilibrium models for CIFAR-10 and ImageNet, showing that it reduces the computational cost of the backward pass by up to two orders of magnitude.