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

1 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 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 . 2 HYPERGRADIENT OPTIMIZATION WITH APPROXIMATE JACOBIAN INVERSE . 2.1 SHINE : HYPERGRADIENT DESCENT WITH APPROXIMATE JACOBIAN INVERSE . 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 ) .

Summary: 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.