On Training Implicit Models

1 Introduction . Conventional neural networks are typically constructed by explicitly stacking multiple linear and non-linear operators in a feed-forward manner . Recently , the implicitly-defined models [ 1 , 2 , 3 , 4 , 5 ] have attracted increasing attentions and are able to match the state-of-the-art results by explicit models on several vision [ 3 , 6 ] , language [ 2 ] and graph [ 4 ] tasks . These works treat the evolution of the intermediate hidden states as a certain form of dynamics , such as fixed-point equations [ 2 , 3 ] or ordinary differential equations ( ODEs ) [ 1 , 7 ] , which represents infinite latent states . The forward passes of implicit models are therefore formulated as solving the underlying dynamics , by either black-box ODE solvers [ 1 , 7 ] or root-finding algorithms [ 2 , 3 ] . As for the backward passes , however , directly differentiating through the forward pass trajectories could induce a heavy memory overhead [ 8 , 9 ] . To this end , researchers have developed memory-efficient backpropagation via implicit differentiation , such as solving a Jacobian-based linear fixed-point equation for the backward pass of deep equilibrium models ( DEQs ) [ 2 ] , which eventually makes the backpropagation trajectories independent of the forward passes . This technique allows one to train implicit models with essentially constant memory consumption , as we only need to store the final output and the layer itself without saving any intermediate states . However , in order to estimate the exact gradient by implicit differentiation , implicit models have to rely on expensive black-box solvers for backward passes , e.g. , ODE solvers or root-solving algorithms . These black-box solver usually makes the gradient computation very costly in practice , even taking weeks to train state-of-the-art implicit models on ImageNet [ 10 ] with 8 GPUs . This work investigates fast approximate gradients for training implicit models . We found that a firstorder oracle that produces good gradient estimates is enough to efficiently and effectively train implicit * Equal contribution †Corresponding author , zlin @ pku.edu.cn 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . models , circumventing laboriously computing the exact gradient as in prior arts [ 2 , 3 , 4 , 11 , 12 ] . We develop a framework in which a balanced trade-off is made between the precision and conditioning of the gradient estimate . Specifically , we provide the general condition under which the phantom gradient can provide an ascent direction of the loss landscape . We further propose two instantiations of phantom gradients in the context of DEQ models , which are based on the the damped fixed-point unrolling and the Neumann series , respectively . Importantly , we show that our proposed instantiations satisfy the theoretical condition , and that the stochastic gradient descent ( SGD ) algorithm based on the phantom gradient enjoys a sound convergence property as long as the relevant hyperparameters , e.g. , the damping factor , are wisely selected . Note that our method only affects , and thus accelerates , the backward formulation of the implicit models , leaving the forward pass formulation ( i.e. , the root-solving process ) and the inference behavior unchanged so that our method is applicable to a wide range of implicit models , forward solvers , and inference strategies . We conduct an extensive set of synthetic , ablation , and large-scale experiments to both analyze the theoretical properties of the phantom gradient and validate its speedup and performances on various tasks , such as ImageNet [ 10 ] classification and Wikitext-103 [ 13 ] language modeling . Overall , our results suggest that : 1 ) the phantom gradient estimates an ascent direction ; 2 ) it is applicable to large-scale tasks and is capable of achieving a strong performance which is comparable with or even better than that of the exact gradient ; and 3 ) it significantly shortens the total training time needed for implicit models roughly by a factor of 1.4 ∼ 1.7× , and even accelerates the backward passes by astonishingly 12× on ImageNet . We believe that our results provide strong evidence for effectively training implicit models with the lightweight phantom gradient . 2 Method . 2.1 Inspection of Implicit Differentiation . In this work , we primarily focus on the formulation of implicit models based on root-solving , represented by the DEQ models [ 2 ] . The table of notations is arranged in Appendix A . Specifically , given an equilibrium module F , the output of the implicit model is characterized by the solution h∗ to the following fixed-point equation : h∗ = F ( h∗ , z ) , ( 1 ) where z ∈ Rdu+dθ is the union of the module ’ s input u ∈ Rdu and parameters θ ∈ Rdθ , i.e. , z > = [ u > , θ > ] . Here , u is usually a projection of the original data point x ∈ Rdx , e.g. , u =M ( x ) . In this section , we assume F is a contraction mapping w.r.t . h so that its Lipschitz constant Lh w.r.t . h is less than one , i.e. , Lh < 1 , a setting that has been analyzed in recent works [ 14 , 15 ] 1 . To differentiate through the fixed point by Eq . ( 1 ) , we need to calculate the gradient of h∗ w.r.t . the input z . By Implicit Function Theorem ( IFT ) , we have ∂h∗ ∂z = ∂F ∂z ∣∣∣∣ h∗ ( I − ∂F ∂h ∣∣∣∣ h∗ ) −1 . ( 2 ) Here , ( ∂a/∂b ) ij = ∂aj/∂bi . The equilibrium point h ∗ of Eq . ( 1 ) is then passed to a post-processing function G to obtain a prediction ŷ = G ( h∗ ) . In the generic learning scenario , the training objective is the following expected loss : R ( θ ) = E ( x , y ) ∼P [ L ( ŷ ( x ; θ ) , y ) ] , ( 3 ) where y is the groundtruth corresponding to the training example x , and P is the data distribution . Here , we omit the parameters of G , because given the output h∗ of the implicit module F , training the post-processing part G is the same as training explicit neural networks . The most crucial component is the gradient of the loss function L w.r.t . the input vector z > = [ u > , θ > ] , which is used to train both the implicit module F and the input projection moduleM . Using Eq . ( 2 ) with the condition h = h∗ , we have ∂L ∂u = ∂F ∂u ( I − ∂F ∂h ) −1 ∂L ∂h , ∂L ∂θ = ∂F ∂θ ( I − ∂F ∂h ) −1 ∂L ∂h . ( 4 ) The gradients in Eq . ( 4 ) are in the same form w.r.t . u and θ . Without loss of generality , we only discuss the gradient w.r.t . θ in the following sections . 1Note that the contraction condition could be largely relaxed to the one that the spectral radius of ∂F/∂h∗ on the given data is less than 1 , i.e. , ρ ( ∂F/∂h∗ ) < 1 , as indicated by the well-posedness condition in [ 5 ] . 2.2 Motivation . The most intriguing part lies in the Jacobian-inverse term , i.e. , ( I − ∂F/∂h ) −1 . Computing the inverse term by brute force is intractable due to the O ( n3 ) complexity . Previous implicit models [ 2 ] approach this by solving a linear system involving a Jacobian-vector product iteratively via a gradient solver , introducing over 30 Broyden [ 16 ] iterations in the backward pass . However , the scale of the Jacobian matrix can exceed 106 × 106 in the real scenarios , leading to a prohibitive cost in computing the exact gradient . For example , training a small-scale state-of-the-art implicit model on ImageNet can consume weeks using 8 GPUs while training explicit models usually takes days , demonstrating that pursuing the exact gradient severely slows down the training process of implicit models compared with explicit models . Secondly , because of the inversion operation , we cast doubt on the conditioning of the gradient and the stability of training process from the numerical aspect . The Jacobian-inverse can be numerically unstable when encountering the ill-conditioning issue . The conditioning problem might further undermine the training stability , as studied in the recent work [ 17 ] . Plus , the inexact gradient [ 9 , 18 , 19 , 20 , 21 , 22 ] is widely applied in the previous learning protocol , like linear propagation [ 23 ] and synthetic gradient [ 24 ] . Here , the Jacobian-inverse is used to calculate the exact gradient which is not always optimal for model training . Moreover , previous research has used a moderate gradient noise as a regularization approach [ 25 ] , which has been shown to play a central role in escaping poor local minima and improving generalization ability [ 26 , 27 , 28 ] . The concerns and observations motivate us to rethink the possibility of replacing the Jacobian-inverse term in the standard implicit differentiation with a cheaper and more stable counterpart . We believe that an exact gradient estimate is not always required , especially for a black-box layer like those in the implicit models . Hence this work designs an inexact , theoretically sound , and practically efficient gradient for training implicit models under various settings . We name the proposed gradient estimate as the phantom gradient . Suppose the Jacobian ∂h∗/∂θ is replaced with a matrixA , and the corresponding phantom gradient is defined as ∂̂L ∂θ : = A ∂L ∂h . ( 5 ) Next , we give the general condition on A so that the phantom gradient can be guaranteed valid for optimization ( Sec . 2.3 ) , and provide two concrete instantiations of A based on either damped fixed-point unrolling or the Neumann series ( Sec . 2.4 ) . The proofs of all our theoretical results are presented in Appendix C . 2.3 General Condition on the Phantom Gradient . Previous research on theoretical properties for the inexact gradient include several aspects , such as the gradient direction [ 22 ] , the unbiasedness of the estimator [ 29 ] , and the convergence theory of the stochastic algorithm [ 19 , 30 ] . The following theorem formulates a sufficient condition that the phantom gradient gives an ascent direction of the loss landscape . Theorem 1 . Suppose the exact gradient and the phantom gradient are given by Eq . ( 4 ) and ( 5 ) , respectively . Let σmax and σmin be the maximal and minimal singular value of ∂F/∂θ . If∥∥∥∥A ( I − ∂F∂h ) − ∂F ∂θ ∥∥∥∥ < σ2minσmax , ( 6 ) then the phantom gradient provides an ascent direction of the function L , i.e. , 〈 ∂̂L ∂θ , ∂L ∂θ 〉 > 0 . ( 7 ) Remark 1 . Suppose only the ( I − ∂F/∂h ) −1 term is replaced with a matrix D , namely , A = ( ∂F/∂θ ) D. Then , the condition in ( 6 ) can be reduced into∥∥∥∥D ( I − ∂F∂h ) − I ∥∥∥∥ < 1κ2 , ( 8 ) where κ is the condition number of ∂F/∂θ . ( See Appendix C.1 for the derivation . )

Summary: The paper presents a method called phantom gradient to increase the efficiency in training implicit deep models by approximating the exact gradient. The authors first exploit Neumann series for the damped inverse matrix in implicit gradient calculation. They then truncate the series to k terms to arrive at finite term approximation of the exact implicit gradient. They finally provide theoretical justification for when phantom gradients are descent directions for implicit models. The authors evaluate their methods by conducting experiments on thoroughly analyzing their descent direction and demonstrating acceleration over a current implicit training method.