DISE: Dynamic Integrator Selection to Minimize Forward Pass Time in Neural ODEs

1 INTRODUCTION . Neural ordinary differential equations ( Neural ODEs ) are to learn time-dependent physical dynamics describing continuous residual networks ( Chen et al. , 2018 ) . It is well known that residual connections are numerically similar to the explicit Euler method , the simplest integrator to solve ODEs . In this regard , Neural ODEs are considered as a generalization of residual networks . In general , it is agreed by many researchers that Neural ODEs have two advantages and one disadvantage : i ) Neural ODEs can sometimes reduce the required number of neural network parameters , e.g. , ( Pinckaers & Litjens , 2019 ) , ii ) Neural ODEs can interpret the neural network layer ( or time ) as a continuous variable and a hidden vector at an arbitrary layer can be calculated , iii ) however , Neural ODEs ’ s forward-pass inference can sometimes be numerically unstable ( i.e. , the underflow error of DOPRI ’ s adaptive step size ) and/or slow to solve an integral problem ( i.e. , too many steps in DOPRI ) ( Zhuang et al. , 2020b ; Finlay et al. , 2020 ; Daulbaev et al. , 2020 ; Quaglino et al. , 2020 ) . Much work has been actively devoted to address the numerically unstable nature of solving integral problems . In this work , however , we are interested in addressing the problem of long forward-pass inference time . To overcome the challenge , we i ) directly regularize the numerical errors of the Dormand–Prince ( DOPRI ) method ( Dormand & Prince , 1980 ) , which means we try to learn an ODE that can be quickly solved by DOPRI , and ii ) dynamically select an appropriate integrator for each sample rather than relying on only one integrator . In many cases , Neural ODEs use DOPRI , one of the most advanced adaptive step integrator , for its best accuracy . However , our method allows that we rely on simpler integrators , such as the Euler method or the fourth-order Runge–Kutta ( RK4 ) method ( Ixaru & Vanden Berghe , 2004 ) , for carefully selected inputs . Table 1 shows an experimental result that our proposed regularization not only reduces the number of function evaluations ( NFE ) — the inference time is linearly proportional to the number of function evaluations in Neural ODEs — but also increases the inference accuracy in the MNIST classification task . We can reduce the inference time by reducing the average number of steps ( and thus , the average NFE ) of DOPRI , which can be obtained when the learned ODE is trained to be in a suitable form to solve with DOPRI with a proper regularization . However , the NFE of DOPRI in a step is 6 whereas RK4 has 4 and the Euler method has 1 . So , the Euler method is six times faster than DOPRI even when their step sizes are identical . Therefore , the automatic step size adjustment of DOPRI is not enough to minimize the NFE of forward-pass inference ( see Section B in Appendix for more detailed descriptions with a concrete example ) . To this end , we design an auxiliary network that chooses an appropriate integrator for each sample . The combination of our regularization and the proposed Dynamic Integrator SElection ( DISE ) shows the best performance in the table . We conduct experiments for three different tasks and datasets : MNIST image classification , PhysioNet mortality prediction , and continuous normalizing flows . Our method shows the best ( or close to the best ) accuracy with a much smaller NFE than state-of-the-art methods . Our contributions can be summarized as follows : 1 . We design an effective regularization to reduce the number of function evaluations ( NFE ) of Neural ODEs . 2 . We design a sample-wise dynamic integrator selection ( DISE ) method to further accelerate Neural ODEs without significantly sacrificing model accuracy . 3 . We conduct in-depth analyses with three popular tasks of Neural ODEs . 2 RELATED WORK . In this section , we review the literature on Neural ODEs . In particular , we review recent regularization designs for Neuarl ODEs and numerical methods to solve ODEs . 2.1 NEURAL ODES . It had been attempted by several researchers to model neural networks as differential equations ( Weinan , 2017 ; Ruthotto & Haber , 2019 ; Lu et al. , 2018 ; Ciccone et al. , 2018 ; Chen et al. , 2018 ; Gholami et al. , 2019 ) . Among them , the seminal neural ordinary differential equations ( Neural ODEs ) , as shown in Fig . 1 , consist of three parts in general : a feature extractor , an ODE , and a classifier ( Chen et al. , 2018 ; Zhuang et al. , 2020a ) . Given an input x , the feature extractor produces an input to the ODE , denoted h ( 0 ) . Let h ( t ) be a hidden vector at layer ( or time ) t in the ODE part . In Neural ODEs , a neural network f with a set of parameters , denoted θ , approximates ∂h ( t ) ∂t and h ( t1 ) becomes h ( 0 ) + ∫ t1 t0 f ( h ( t ) , t ; θ ) dt , where f ( h ( t ) , t ; θ ) = ∂h ( t ) ∂t . In other words , the internal dynamics of the hidden vector evolution is described by an ODE . One key advantage of Neural ODEs is that we can reduce the number of parameters without sacrificing model accuracy . For instance , one recent work based on a Neural ODE marked the best accuracy for medical image segmentation with an order of magnitude smaller parameter numbers ( Pinckaers & Litjens , 2019 ) . In general , we calculate h ( 1 ) 1 and feed it into the next classifier and its final prediction is made . One can accordingly modify the architecture in Fig . 1 for other types of tasks . For simplicity but without loss of generality , we assume the architecture in our discussion . Neural ODEs have been used in many tasks , ranging from classification and regression to time series forecasting and generative models ( Yildiz et al. , 2019 ; Grathwohl et al. , 2019 ; Rubanova et al. , 2019 ) . 2.2 ODE SOLVERS . DOPRI is one of the most powerful integrators ( Hairer et al. , 1993 ) and widely used in Neural ODEs . It is a member of the Runge–Kutta family of ODE solvers . DOPRI dynamically controls the step size while solving an integral problem . It is now the default method for MATLAB , GNU Octave , and Simulink . It internally estimates an error by using a heuristic method and the step size is determined by a function inversely proportional to the estimated error — the larger the error , the shorter the step size . The error at i-th step of DOPRI for an integral problem x , denoted errx , i , is estimated by the difference between the fourth-order and the fifth-order Runge–Kutta methods at the moment . The intuition behind the heuristic error estimation is simple yet effective . Among simpler methods , we consider the Euler method , and the fourth-order Runge–Kutta ( RK4 ) method . The Euler method is the simplest method to solve ODEs and both the Euler method and RK4 use a fixed step size . Therefore , their solving time is deterministic . One step of DOPRI involves six function evaluations , i.e. , six function calls of f . The Euler method calls the network f only once in a step and RK4 calls four times . Therefore , the Euler method is six times faster than DOPRI for a step . The term ‘ NFE ’ refers to the number of function evaluations to solve an integral problem . For the Euler method and RK4 , NFE is deterministic and does not vary . In DOPRI , however , NFE varies from one sample to another , depending on the estimated error and the number of steps . We refer readers to Section B in Appendix for more detailed descriptions with a concrete example . 2.3 REGULARIZATIONS IN NEURAL ODES . To make Neural ODEs faster , one possible way is regularizing the ODE function f . Two naı̈ve methods are regularizing θ with the L1 or L2 regularizers ( Ng , 2004 ) . Strictly speaking , these two regularizers are to prevent overfitting . Therefore , preventing overfitting does not necessarily mean quick forward-pass inference . To this end , Dupont et al . showed that by augmenting h ( t ) with additional zeros , i.e. , augmenting the dimensionality of h ( t ) , one can achieve similar effects ( Dupont et al. , 2019 ) . However , this method is meaningful when we can not freely control the dimensionality of h ( t ) , which is not our setting . Recently , a kinetic regularization concept has been proposed by Finlay et al . ( Finlay et al. , 2020 ) , which is written as follows : Rk def = ∫ t1 t0 ‖f ( h ( t ) , t ; θ ) ‖22 dt . ( 1 ) Among all regularization terms designed so far , this kinetic regularization ’ s goal is the closest to ours . It can encourage Neural ODEs to learn straight-line paths from h ( t0 ) to h ( t1 ) . 3 PROPOSED METHOD . While enabling the design of compact models , Neural ODEs have one critical drawback that they require solving integral problems , for which many approximation methods have been proposed : the Euler method , RK4 , and DOPRI , to name a few . Almost all of them are based on discretizing t and converting an integral into a series of additions . In many cases , therefore , it requires a dense discretization , resulting in a long forward-pass inference time . 1For simplicity but without loss of generality , the time duration can be fixed into t ∈ [ 0 , 1 ] . Any arbitrary length ODEs can be compressed into a unit time interval ODE . In some time series datasets , however , the final integral time t1 is given in a sample . In such a case , t1 is set to the sample time . In this paper , we tackle the problem of minimizing the number of function evaluations ( and thereby , the forward-pass inference time ) of Neural ODEs without significantly sacrificing model accuracy . Our proposed method consists of two parts : i ) using the DOPRI ’ s error estimation as a regularizer and ii ) using an auxiliary network to select an appropriate integrator for each input sample . 3.1 DOPRI ’ S ERROR ESTIMATION AS A REGULARIZER . We re-implement the DOPRI method in PyTorch and make it return the estimated error terms . Let { errx,1 , errx,2 , · · · , errx , N } , where N is the number of steps of DOPRI , be an array of errors estimated by DOPRI while solving an integral problem for an input x . Note that the adaptive step size is an inverse function of the error at each step . We use the following regularizer while training Neural ODEs : Rerr def = ∑ x∈T N∑ i=1 errx , i , ( 2 ) where x is an input sample for which we have to solve an integral problem , and T is a training set . For instance , x can be an image sample to classify . If we train a Neural ODE to classify images with the cross-entropy loss in conjunction with the regularizer , the trained Neural ODE will learn an ODE that can correctly classify images while reducing the forward-pass time of DOPRI . The backward-pass calculation of our proposed regularizer can be done in O ( 1savg ) by maintaining the forward-pass computation graph , where savg is the average step size of DOPRI . However , this complexity will decrease as training goes on with our regularizer because the average step size will increase .

Summary: This paper addresses the complexity of the forward pass inference in neural ODEs. The paper proposes to augment training of the neural ODE  with an auxiliary neural network that dynamically selects the best numerical integrator for a given input sample. Furthermore, the paper also proposes a regularizer that uses the errors of the numerical integrator to reduce the number of function evaluations, without sacrificing accuracy.