SHINE: SHaring the INverse Estimate from the forward pass for bi-level optimization and implicit models

Introduction

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 $\mathcal{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).s

In general, the formulation of DEQs can be cast as a bi-level problem of the following form:

To compute DEQs’ gradient efficiently and avoid high memory cost, one does not rely on back-propagation but uses the implicit function theorem (Krantz and Parks, 2013) which gives an analytical expression of the Jacobian of $z^{\star}$ with respect to $\theta$, $\frac{\partial z^{\star}}{\partial\theta}$. 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 (missing) 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

where $\mathcal{L}_{\text{train}}$ and $\mathcal{L}_{\text{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 $\theta$ is the regularisation parameter. As the training loss is smooth and convex, the inner problem can be written as in (missing) 1 with $g_{\theta}=\nabla_{z}r_{\theta}$ to fit (missing) 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 (missing) 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^{\star}$, 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:

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.

Hypergradient Optimization with Approximate Jacobian Inverse

Hypergradient optimization is a first-order method used to solve (missing) 1. We recall that in the case of smooth convex optimization, $\frac{\partial g_{\theta}}{\partial 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_{\theta}}$ whenever the results can be applied to both contexts. To enable Hypergradient Optimization, i.e. gradient descent on $\mathcal{L}$ with respect to $\theta$, Bai et al. (2019, Theorem 1) show the following theorem, which is based on implicit differentiation (Krantz and Parks, 2013):

In practice, we use an algorithm to approximate $z^{\star}$, and 1 gives a plug-in formula for the backward pass. Note that this formula is independent of the algorithm chosen to compute $z^{\star}$. 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^{\star}$ has been obtained, one of the major bottlenecks in the computation of the Hypergradient is the inversion of $J_{g_{\theta}}(z^{\star})$ in the directions \frac{\partial g_{\theta}}{\partial\theta}\Bigr{|}_{z^{\star}} or $\nabla_{z}\mathcal{L}(z^{\star})$.

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 limited-memory version (Bai et al., 2020).

These quasi-Newton methods were first inspired by Newton’s method, which finds the root of $g_{\theta}$ via the recurrent Jacobian-based updates $z_{n+1}=z_{n}-J_{g_{\theta}}(z_{n})^{-1}g_{\theta}(z_{n})$. Specifically, they replace the Jacobian $J_{g_{\theta}}(z_{n})$ by an approximation $B_{n}$ that is based on available values of the iterates $z_{n}$ and $g_{\theta}$ rather than its derivative. These $B_{n}$, called qN matrices, are defined 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 efficient multiplication by $B_{n}$ and $B_{n}^{-1}$. 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\to\infty}B_{n}^{-1}$ as a replacement for $J_{g_{\theta}}(z^{\star})^{-1}$ in (missing) 3, i.e. to share the inverse estimate between the forward and the backward passes. This gives the approximate Hypergradient

In practice we will consider the nonasymptotical direction p_{\theta}^{(n)}=\nabla_{z}\mathcal{L}(z_{n})B_{n}^{-1}\frac{\partial g_{\theta}}{\partial\theta}\Bigr{|}_{z_{n}} . Thanks to the Sherman-Morrison formula (Sherman and Morrison, 1950), the inversion of $B_{n}$ can be done very efficiently (using scalar products) compared to the iterative methods needed to invert the true Jacobian $J_{g_{\theta}}(z^{\star})$. In turn, this significantly reduces the computational cost of the Hypergradient computation.

Relationship to the Jacobian-Free method

Because $B_{0}=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, $J_{g_{\theta}}(z^{\star})^{-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_{\theta}$ they show that this preconditioned gradient is still a descent direction.

Transition to the exact Jacobian Inverse.

The approximate gradient $p_{\theta}^{(n)}$ can also be used as the initialization of an iterative algorithm for inverting $J_{g_{\theta}}(z^{\star})$ in the direction $\nabla_{z}\mathcal{L}(z^{\star})$. 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 (Pedregosa, 2016, Bai et al., 2019; 2020).

2 Convergence to the true gradient

To further justify and formalize the idea of SHINE, we show that the direction $p_{\theta}^{(n)}$ converges to the Hypergradient \frac{\partial\mathcal{L}}{\partial\theta}\Bigr{|}_{z^{\star}}. We now collect the assumptions that will be used for this purpose.

There exist a positive constant $\rho>0$ and natural numbers $n_{0}\geq 0$ and $m\geq d$ with the following property: For any $n\geq n_{0}$ we can find indices $n\leq n_{1}\leq\ldots\leq n_{d}\leq n+m$ such that, for $p_{n}$ defined in Algorithm 1, the smallest singular value of the $d\times d$ matrix

(i) $\sum_{n=0}^{\infty}\|z_{n}-z^{\star}\|<\infty$ for some $z^{\star}$ with $g_{\theta}(z^{\star})=0$; (ii) $g_{\theta}$ is $C^{1}$, $J_{g_{\theta}}$ is Lipschitz continuous near $z^{\star}$, and $J_{g_{\theta}}(z^{\star})$ is invertible; (iii) $\nabla_{z}\mathcal{L}$ is continuous, and $\forall\theta$, $\frac{\partial g_{\theta}}{\partial\theta}$ is continuous.

The 2 (i) implies $\lim_{n\to\infty}z_{n}=z^{\star}$. The existence of the Jacobian and its inverse are assumptions that are already made in the regular DEQ setting just to train the model.

Let us denote $p_{\theta}^{(n)}$, the SHINE direction for iterate $n$ in Algorithm 1 with $b=\texttt{true}$. Under Assumptions 1 and 2, for a given parameter $\theta$, $(z_{n})$ converges q-superlinearly to $z^{\star}$ and

From More and Trangenstein (1976, Theorem 5.7) we obtain that $\lim_{n\to\infty}B_{n}=J_{g_{\theta}}(z^{\star})$. 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. ∎

2 establishes convergence of the SHINE direction to the true Hypergradient, but relies on 1 (ULI). While ULI is often used to prove convergence results for qN matrices, e.g. in (Li et al., 1998, Nocedal and Wright, 2006, Conn et al., 1991), 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, 2021a; b; 2020), 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.

3 Outer Problem Awareness

The ULI assumption guarantees convergence of $B_{n}^{-1}$ to $J_{g_{\theta}}(z^{\star})^{-1}$. However, (missing) 3 only requires the multiplication of $J_{g_{\theta}}(z^{\star})^{-1}$ with $\frac{\partial g_{\theta}}{\partial\theta}|_{z^{\star}}$ from the right and $\nabla_{z}\mathcal{L}(z^{\star})$ from the left.

where $(t_{n})\subset[0,\infty)$ satisfies $\sum_{n}t_{n}<\infty$. This update direction will be used to create an extra secant condition $X^{-1}(g_{\theta}(z_{n}+e_{n})-g_{\theta}(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 fixed frequency, then fast (q-superlinear) convergence of $(z_{n})$ to $z^{\star}$ is retained, while convergence of the SHINE direction to the true Hypergradient is also ensured. To show this, we use the following assumption.

Let us consider $p_{\theta}^{(n)}$, the SHINE direction for iterate $n$ in Algorithm 1 that is enriched by extra updates in the direction $e_{n}$ defined in (missing) 5. Under Assumptions 2 (ii-iii) and 3, for a given parameter $\theta$, we have the following: Algorithm 1, for any symmetric and positive definite matrix $B_{0}$, generates a sequence $(z_{n})$ that converges q-superlinearly to $z^{\star}$, and there holds

It follows from known results that the extra updates do not destroy the q-superlinear convergence of $(z_{n})$. The proof of (missing) 6 relies firstly on the fact that by continuity of the derivative of $g_{\theta}$, we have $\lim_{n\to\infty}\frac{\partial g_{\theta}}{\partial\theta}|_{z_{n}}=\frac{\partial g_{\theta}}{\partial\theta}|_{z^{\star}}$. 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 (missing) 6 follows. A full proof is provided in Section B.2. ∎

3 also holds without line searches (i.e., $\alpha_{n}=1$ for all $n$) and any $C^{2}$ function $r_{\theta}$ (such that $g_{\theta}(z)=\nabla_{z}r_{\theta}(z)$) with locally Lipschitz continuous Hessian if $z_{0}$ is close enough to some $z^{\star}$ with $\nabla_{z}r_{\theta}(z^{\star})=0$ and $\nabla_{zz}^{2}r_{\theta}(z^{\star})$ positive definite.

We note that 3 guarantees fast convergence of the iterates $(z_{n})$ and that $z_{0}$ does not have to be close to $z^{\star}$ for that guarantee. Also, there is no restriction on $B_{0}$ other than being symmetric and positive definite (which is satisfied for our choice $B_{0}=I$). Finally, 3 does not rely on ULI. From a practical standpoint we thus regard 3 as a much stronger result than 2.

Adjoint Broyden with OPA

It is not practical to use the partial derivative $\frac{\partial g_{\theta}}{\partial\theta}$ 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_{n}^{-1}$ approximates $J_{g_{\theta}}(z^{\star})^{-1}$ in the direction $\nabla_{z}\mathcal{L}(z^{\star})$ applied by left-multiplication, as required by (missing) 3. An appropriate secant condition is given by

To incorporate the secant condition (missing) 7, we use the Adjoint Broyden’s method (Schlenkrich et al., 2010), a qN method relying on the efficient vector-Jacobian multiplication by $J_{g_{\theta}}$ using auto-differentiation tools. To prove convergence of the SHINE direction for this method, we need the following assumption.

The sequence $(B_{n})$ generated by Algorithm 1 satisfies

Convergence results for quasi-Newton methods usually include showing that 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 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 1 entails $\lim B_{n}=J_{g_{\theta}}(z^{\star})$ and thus $\lim B_{n}^{-1}=J_{g_{\theta}}(z^{\star})^{-1}$, so it is clearly stronger than 4.

Let us consider $p_{\theta}^{(n)}$, the SHINE direction for iterate $n$ in Algorithm 1 with the Adjoint Broyden secant condition (missing) 7 and extra update in the direction $v_{n}$ defined in (missing) 8. Under Assumptions 2 and 4, for a given parameter $\theta$, we have q-superlinear convergence of $(z_{n})$ to $z^{\star}$ and

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 $\nabla_{z}{\mathcal{L}}(z^{\star})=0$ the claim holds due to continuity and 4. Then $\nabla_{z}{\mathcal{L}}(z^{\star})\neq 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 $J_{g_{\theta}}$, we have to store the activations of $g_{\theta}(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.

Results

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 guaranteesSee the results on contractivity in Section E.3. 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.

All the bi-level optimization experiments were done using the HOAG code (Pedregosa, 2016)https://github.com/fabianp/hoag, which is based on the Python scientific ecosystem (Harris et al., 2020, Virtanen et al., 2020, Pedregosa et al., 2011). Deep Equilibrium experiments were done using the PyTorch (Paszke et al., 2019) code for Multiscale DEQ (Bai et al., 2020)https://github.com/locuslab/mdeq, 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 $(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 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.

1 Bi-level optimization – Hyperparameter optimization in Logistic Regression

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. 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 {\frac{\partial g_{\theta}}{\partial z}\Bigr{|}_{z^{\star}}(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.

2 Deep Equilibrium Models

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.

The first dataset is CIFAR-10 (Krizhevsky, 2009) which features 60,000 $32\!\times\!32$ images representing 10 classes. For this dataset, the size of the multi-scale fixed point is $d=50$k. We train the models for five different random seeds.

The results in Figure 3 show that for the vanilla version, SHINE slightly outperforms the Jacobian-Free 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\!\times\!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=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. The fallback is barely used : in 1000 batches of size 32, only 2 samples used fallback, a proportion of $6.25\times 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.More on the overall computational effort can be found in Table E.2 Moreover, using the proposed refined versions of these methods, the performance drop can be traded-off for acceleration.

Conclusion and Discussion

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 4 in which convergence of $(z_{n})$ to $z^{\star}$ is not an assumption but rather follows from the assumptions, as achieved in 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.

Reproducibility Statement

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.

Acknowledgements

This work was performed using HPC resources from GENCI-IDRIS (Grant AD011011153R2).

References

Appendix A OPA algorithm

A possible choice for $(t_{n})$ is to use an arbitrary $t_{0}>0$ and $t_{n}:=\|s_{n-1}\|$ for $n\geq 1$.

Appendix B Proofs of SHINE convergence

To facilitate reading, we restate the results before proving them.

Under Assumptions 1 and 2, More and Trangenstein (1976, Theorem 5.7) shows that $B_{n}$ satisfies

The inversion operator is continuous in the space of invertible matrices, so we have:

Because $\nabla_{z}\mathcal{L}$ and $\frac{\partial g_{\theta}}{\partial\theta}$ are continuous at $z^{\star}$ by 2 (iii), we also have thanks to 2 (i):

By continuity we then deduce that, as claimed,

B.2 Convergence for BFGS with OPA

$r_{\theta}$ is strongly convex in an open superset of $\Omega$ (this implies that $r_{\theta}$ has a unique global minimizer $z^{\star}$) and has a Lipschitz continuous Hessian near $z^{\star}$;

there are positive constants $\eta_{1},\eta_{2}$ such that the line search used in the algorithm ensures that for each $n\geq 0$ either

the line search has the property that $\alpha_{n}=1$ will be used if both

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.

The proof is divided into four steps. The first step is to establish the q-superlinear convergence of $(z_{n})$ to $z^{\star}$. Denoting by $N_{e}\subset\{0,M,2M,\ldots\}$ the set of indices of extra updates that are actually applied, the second step consists of showing

where, in this proof, $B_{n}$ always represents the matrix from Algorithm 2 before the update in the direction $e_{n}$ is applied, i.e., the matrix whose inverse appears in the definition of $e_{n}$, while $\hat{B}_{n}$ always represents the matrix from Algorithm 2 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}=\hat{B}_{n}$. The third step is to prove that (missing) 9 implies the desired convergence (missing) 6 of the SHINE direction if the limit $n\to\infty$ is replaced by $N_{e}\ni n\to\infty$, 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_{n}^{-1}$, respectively, $\hat{B}_{n}^{-1}$, we can also obtain the sequences $(B_{n})$ and $(\hat{B}_{n})$ by updating according to

for the usual update (skipping the update if $y_{n}^{T}s_{n}\leq 0$), respectively,

For the third step, we abbreviate $v_{n}:=\frac{\partial g_{\theta}}{\partial\theta}|_{z_{n}}$. From the definition of $e_{n}$ and (missing) 10 we infer that

After multiplication with $J_{g_{\theta}}(z^{\star})^{-1}$ this entails

by 2 (iii). Using 2 (iii) again it follows that

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}\ni n\to\infty}\|B_{n}^{-1}-B_{j_{n}}^{-1}\|=0$, so

by the third step, establishing the claim.

It remains to show the validity of $\lim_{N_{e}\ni n\to\infty}\|B_{n}-B_{j_{n}}\|=0$ for any sequence $(j_{n})_{n\in N_{e}}$ such that $\{j_{n},j_{n}+1,\ldots,n-1\}\cap N_{e}=\emptyset$ for all $n\in N_{e}$ sufficiently large. Since at least every second extra update is actually carried out, the condition on the intersection implies $n-j_{n}\leq 2M-1$ for all these $n$. Now let $(j_{n})_{n\in N_{e}}$ be any such sequence. Then $B_{n}-B_{j_{n}}=\sum_{m=j_{n}}^{n-1}B_{m+1}-B_{m}$ is a sum of at most $2M-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\in\{0,1,\ldots,n-j_{n}\}$, are satisfied, allowing us to deduce

for all $l\in\{0,1,\ldots,n-j_{n}-1\}$. For each of these $l$, both terms on the right-hand side tend to zero for $N_{e}\ni n\to\infty$ (for the second term this follows from the first identity in (missing) 10 due to $B_{n-l-1}=\hat{B}_{n-l-1}$). Recalling that $B_{n}-B_{j_{n}}=\sum_{m=j_{n}}^{n-1}B_{m+1}-B_{m}$ we find $\lim_{N_{e}\ni n\to\infty}\|B_{n}-B_{j_{n}}\|=0$, which finishes the fourth step and thus concludes the proof. ∎

B.3 Convergence for Adjoint Broyden with OPA

Due to 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 $\nabla_{z}\mathcal{L}(z^{\star})=0$. By continuity this implies $\lim_{n\to\infty}\nabla_{z}\mathcal{L}(z_{n})=0$. Since the sequence $(B_{n}^{-1}\frac{\partial g_{\theta}}{\partial\theta}|_{z_{n}})$ is bounded by 4, it follows that

Since $\lim_{N_{e}\ni n\to\infty}\nabla_{z}\mathcal{L}(z_{n})J_{g_{\theta}}(z^{\star})^{-1}=\nabla_{z}\mathcal{L}(z^{\star})J_{g_{\theta}}(z^{\star})^{-1}$ by continuity, we find

Both terms on the right-hand side go to zero as $n$ goes to infinity: the first one due to $\lim_{n\to\infty}z_{n}=z^{\star}$ and the second one since $\lim_{n\to\infty}\frac{\|E_{n}^{T}v_{n}\|}{\|v_{n}\|}=0$ by Schlenkrich et al. (2010, Lemma 3). This shows that ${\lim_{n\to\infty}\|B_{n+1}-B_{n}\|=0}$, which concludes the proof of the intermediate claim.

includes the sequence $(B_{n})$ and is compact by the Banach lemma, so inversion is a uniformly continuous operation on this set.

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}$ defined in (missing) 8, then 4 can be replaced by the significantly weaker assumption that the sequence $(B_{n}^{-1}\frac{\partial g_{\theta}}{\partial\theta}|_{z_{n}})$ is bounded. The price to pay is that the convergence rate of $(z_{n})$ to $z^{\star}$ 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^{\star}$.

Appendix C Logistic Regression Hyperparameters

For both datasets we split the data randomly (with a different seed for each run) between training-validation-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.

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.

Appendix 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 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.

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\times 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.

Appendix E Additional results

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.

E.2 Regularized Nonlinear Least Squares

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 $(x_{train,i},y_{train,i})_{i=1}^{N}$ and a test set $(x_{test,i},y_{test,i})_{i=1}^{M}$, this problem reads

where $\sigma$ denotes the sigmoid function $\sigma(x)=\frac{1}{1+e^{-x}}$. For a fixed hyper-parameter $\theta$, this task is typically solved using L-BFGS (Xu et al., 2020, Berahas et al., 2021).

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).

E.3 Contractivity 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 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.

E.4 Time 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 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.

E.5 DEQ OPA results

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.