"Hey, that's not an ODE'": Faster ODE Adjoints with 12 Lines of Code

1 INTRODUCTION . We begin by recalling the usual set-up for neural differential equations . 1.1 NEURAL ORDINARY DIFFERENTIAL EQUATIONS . The general approach of neural ordinary differential equations ( E , 2017 ; Chen et al. , 2018 ) is to use ODEs as a learnable component of a differentiable framework . Typically the goal is to approximate a map x 7→ y by learning functions ` 1 ( · , φ ) , ` 2 ( · , ψ ) and f ( · , · , θ ) , which are composed such that z ( τ ) = ` 1 ( x , φ ) , z ( t ) = z ( τ ) + ∫ t τ f ( s , z ( s ) , θ ) ds and y ≈ ` 2 ( z ( T ) , ψ ) . ( 1 ) The variables φ , θ , ψ denote learnable parameters and the ODE is solved over the interval [ τ , T ] . We include the ( often linear ) maps ` 1 ( · , φ ) , ` 2 ( · , ψ ) for generality , as in many contexts they are important for the expressiveness of the model ( Dupont et al. , 2019 ; Zhang et al. , 2020 ) , though our contributions will be focused around the ODE component and will not depend on these maps . Here we will consider neural differential equations that may be interpreted as a neural ODE . 1.2 APPLICATIONS . Neural differential equations have to the best our knowledge three main applications : 1 . Time series modeling . Rubanova et al . ( 2019 ) interleave Neural ODEs with RNNs to produce ODEs with jumps . Kidger et al . ( 2020 ) take f ( t , z , θ ) = g ( z , θ ) dXdt ( t ) , dependent on some time-varying input X , to produce a neural controlled differential equation . 2 . Continuous Normalising Flows as in Chen et al . ( 2018 ) ; Grathwohl et al . ( 2019 ) , in which the overall model acts as coupling or transformation between probability distributions , 3 . Modeling or controlling physical environments , for which a differential equation based model may be explicitly desired , see for example Zhong et al . ( 2020 ) . 1.3 ADJOINT EQUATIONS . The integral in equation ( 1 ) may be backpropagated through either by backpropagating through the internal operations of a numerical solver , or by solving the backwards-in-time adjoint equations with respect to some ( scalar ) loss L. az ( T ) = dL dz ( T ) , az ( t ) = az ( T ) − ∫ t T az ( s ) · ∂f ∂z ( s , z ( s ) , θ ) ds and dL dz ( τ ) = az ( τ ) , aθ ( T ) = 0 , aθ ( t ) = aθ ( T ) − ∫ t T az ( s ) · ∂f ∂θ ( s , z ( s ) , θ ) ds and dL dθ = aθ ( τ ) , at ( T ) = dL dT , at ( t ) = at ( T ) − ∫ t T az ( s ) · ∂f ∂s ( s , z ( s ) , θ ) ds and dL dτ = at ( τ ) , ( 2 ) These equations are typically solved together as a joint system a ( t ) = [ az ( t ) , aθ ( t ) , at ( t ) ] . ( They are already coupled ; the latter two equations depend on az . ) As additionally their integrands require z ( s ) , and as the results of the forward computation of equation ( 1 ) are usually not stored , then the adjoint equations are typically additionally augmented by recovering z by solving backwards-intime z ( t ) = z ( T ) + ∫ t T f ( s , z ( s ) , θ ) ds . ( 3 ) 1.4 CONTRIBUTIONS . We demonstrate that the particular structure of the adjoint equations implies that numerical equation solvers will typically take too many steps , that are too small , wasting time during backpropagation . Specifically , the accept/reject step of adaptive-step-size solvers is too stringent . By applying a correction to account for this , we demonstrate that the number of steps needed to solve the adjoint equations may be reduced by typically about 40 % . We observe improvements on some problems by as much as 62 % . Factoring in the forward pass ( which is unchanged ) , the overall training time is roughly halved . Our method is hyperparameter-free and requires no tuning . We do not observe any change in model performance , and at least with the torchdiffeq package ( our chosen differential equation package ) , this correction may be applied with only 12 lines of code . 2 METHOD . 2.1 NUMERICAL SOLVERS . Both the forward pass given by equation ( 1 ) , and the backward pass given by equations ( 2 ) and ( 3 ) , are solved by invoking a numerical differential equation solver . Our interest here is in adaptive-stepsize solvers . Indeed the default choice for solving many equations is the adaptive-step-size Runge– Kutta 5 ( 4 ) scheme of Dormand–Prince ( Dormand & Prince , 1980 ) , for example as implemented by dopri5 in the torchdiffeq package or ode45 in MATLAB . A full discussion of the internal operations of these solvers is beyond our scope here ; the part of interest to us is the accept/reject scheme . Consider the case of solving the general ODE y ( t ) = y ( τ ) + ∫ t τ f ( s , y ( s ) ) ds , with y ( t ) ∈ Rd . Suppose for some fixed t the solver has computed some estimate ŷ ( t ) ≈ y ( t ) , and it now seeks to take a step ∆ > 0 to compute ŷ ( t + ∆ ) ≈ y ( t + ∆ ) . A step is made , and some candidate ŷcandidate ( t + ∆ ) is generated . The solver additionally produces yerr ∈ Rd representing an estimate of the numerical error made in each channel during that step . Given some prespecified absolute tolerance ATOL ( for example 10−9 ) , relative tolerance RTOL ( for example 10−6 ) , and ( semi ) norm ‖ · ‖ : Rd → [ 0 , ∞ ) ( for example ‖y‖ = √ 1 d ∑d i=1 y 2 i the RMS norm ) , then an estimate of the size of the equation is given by SCALE = ATOL+RTOL ·max ( ŷ ( t ) , ŷcandidate ( t+ ∆ ) ) ∈ Rd , ( 4 ) where the maximum is taken channel-wise , and the error ratio r = ∥∥∥ yerr SCALE ∥∥∥ ∈ R ( 5 ) is then computed . If r ≤ 1 then the error is deemed acceptable , the step is accepted and we take ŷ ( t+∆ ) = ŷcandidate ( t+∆ ) . If r > 1 then the error is deemed too large , the candidate ŷcandidate ( t+∆ ) is rejected , and the procedure is repeated with a smaller ∆ . Note the dependence on the choice of norm ‖ · ‖ : in particular this determines the relative importance of each channel towards the accept/reject criterion . 2.2 ADJOINT SEMINORMS . Not an ODE The key observation is that aθ ( and in fact also at ) does not appear anywhere in the vector fields of equation ( 2 ) . This means that ( conditioned on knowing z and az ) , the integral corresponding to aθ is just an integral—not an ODE . As such , it is arguably inappropriate to solve it with an ODE solver , which makes the implicit assumption that small errors now may propagate to create large errors later . Accept/reject This is made manifest in the accept/reject step of equation ( 5 ) . Typical choices of norm ‖ · ‖ , such as L2 , will usually weight each channel equally . But we have just established that to solve the adjoint equations accurately , it is far more important that z and az be accurate than it is that aθ be accurate . Seminorms Thus , when solving the adjoint equations equation ( 2 ) , we propose to use a ‖ · ‖ that scales down the effect in those channels corresponding to aθ . In practice , in our experiments , we scale ‖ · ‖ all the way down by applying zero weight to the offending channels , so that ‖ · ‖ is in fact a seminorm . This means that the integration steps are chosen solely for the computation of az and z , and the values of aθ are computed just by integrating with respect to those steps . Example As an explicit example , note that aθ ( T ) = 0 . When solving the adjoint equation numerically , this means for t close to T that the second term in equation ( 4 ) is small . As ATOL is typically also small , then SCALE is additionally small , and the error ratio r in equation ( 5 ) is large . This implies that it becomes easy for the error ratio r to violate r ≤ 1 , and it is easy for the step to be rejected . Now there is nothing intrinsically bad about a step being rejected—we would like to solve the ODE accurately , after all—the problem is that this is a spurious rejection , as the rejection occurred to ensure the accuracy of aθ , which is as already established unnecessary . In practice , we observe that spurious rejections may occur for any t , not just those near T . Other channels In fact , essentially the same argument applies to at as well : this does not affect the value of the vector field either . In a continuous normalising flow , the log-probability channel is also only an integral , rather than an ODE , and again the same argument may be applied . Does this reduce the accuracy of parameter gradients ? One obvious concern is that we are typically ultimately interested in the parameter gradients aθ , in order to train a model ; with respect to this our approach seems counter-intuitive . However , we verify empirically that models still train without a reduction in performance . We explain this by noting that the z , az channels truly are ODEs , so that small errors now do propagate to create larger errors later . Thus these are likely the dominant source of error overall . 2.3 CODE . Depending on the software package , the code for making this change can be trivial . For example , using PyTorch ( Paszke et al. , 2019 ) and torchdiffeq ( Chen et al. , 2018 ) , the standard set-up requires only a few additional lines of code . The additional 12 lines are marked . 1 import t o r c h d i f f e q 2 3 def rms norm ( t e n s o r ) : # 4 re turn t e n s o r . pow ( 2 ) . mean ( ) . s q r t ( ) # 5 # 6 def make norm ( s t a t e ) : # 7 s t a t e s i z e = s t a t e . numel ( ) # 8 def norm ( a u g s t a t e ) : # 9 y = a u g s t a t e [ 1 : 1 + s t a t e s i z e ] # 10 a d j y = a u g s t a t e [ 1 + s t a t e s i z e : 1 + 2 ∗ s t a t e s i z e ] # 11 re turn max ( rms norm ( y ) , rms norm ( a d j y ) ) # 12 re turn norm # 13 # 14 t o r c h d i f f e q . o d e i n t a d j o i n t ( func = . . . , y0 = . . . , t = . . . , 15 a d j o i n t o p t i o n s = d i c t ( norm=make norm ( y0 ) ) ) # This amounts to the extra 12 lines of code stated in the title—a number that even includes the additional whitespace and visual indents . To keep the remainder of this discussion software-agnostic , we defer further explanation of this specific code to Appendix A .

Summary: The paper proposes a modification for the adjoint method, such that to improve the training efficiency of neural ODEs. The proposed idea is that the solution of some terms in the adjoint method can be less accurate, because these are not ODEs but simple integrals, and hence, the error does not propagate. Thus, the solver can utilize bigger steps, and in total to perform less steps. In the experiments the efficiency is demonstrated under different scenarios where neural ODEs are used.