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 proposes an approximate gradient, called phantom gradient, for training implicit models. They highlight that the method reduces the cost of gradient computation and it provides a descent direction suitable to train models. Their gradients are based on two methods, fixed-point unrolling and the Neumann series. The paper proves that under some network conditions on the eigenvalues of the derivative of the implicit module, the phantom gradient has a descent direction positively corelated to the exact gradient. The paper also provide almost surely convergence of the phantom gradient to the exact gradient under some conditions.