SNODE: Spectral Discretization of Neural ODEs for System Identification

1 INTRODUCTION . Neural Ordinary Differential Equations ( ODE-Nets ; Chen et al. , 2018 ) can learn latent models from observations that are sparse in time . This property has the potential to enhance the performance of neural network predictive models in applications where information is sparse in time and it is important to account for exact arrival times and delays . In complex control systems and model-based reinforcement learning , planning over a long horizon is often needed , while high frequency feedback is necessary for maintaining stability ( Franklin et al. , 2014 ) . Discrete-time models , including RNNs ( Jain & Medsker , 1999 ) , often struggle to fully meet the needs of such applications due to the fixed time resolution . ODE-Nets have been shown to provide superior performance with respect to classic RNNs on time series forecasting with sparse training data . However , learning their parameters can be computationally intensive . In particular , ODE-Nets are memory efficient but time inefficient . In this paper , we address this bottleneck and propose a novel alternative strategy for system identification . Summary of contributions . We propose SNODE , a compact representation of ODE-Nets for system identification with full state information that makes use of a higher-order approximation of its states by means of Legendre polynomials . This is outlined in Section 4 . In order to find the optimal polynomial coefficients and network parameters , we develop a novel optimization scheme , which does not require to solve an ODE at each iteration . The resulting algorithm is detailed in Section 3 and is based on backpropagation ( Linnainmaa , 1970 ; Werbos , 1981 ; Lecun , 1988 ) and automatic differentiation ( Paszke et al. , 2017 ) . The proposed method is fully parallel with respect to time and its approximation error reduces exponentially with the Legendre polynomial order ( Canuto et al. , 1988 ) . Summary of numerical experiments . In Section 5 , our method is tested on a 6-state vehicle problem , where it is at least one order or magnitude faster in each optimizer iteration than explicit and adjoint methods , while convergence is achieved in a third of the iterations . At test time , the MSE is reduced by one order of magnitude . In Section 6 , the method is used for a 30-state system consisting ∗Corresponding author . of identical vehicles , coupled via a known collision avoidance policy . Again , our method converges in a third of the iterations required by backpropagation thourgh a solver and each iteration is 50x faster than the fastest explicit scheme . 2 NEURAL ORDINARY DIFFERENTIAL EQUATIONS . The minimization of a scalar-valued loss function that depends on the output of an ODE-Net can be formulated as a general constrained optimization problem : min θ∈Rm ∫ t1 t0 L ( t , x ( t ) ) dt , ( 1 ) s.t . ẋ ( t ) = f ( t , x ( t ) , u ( t ) ; θ ) , x ( t0 ) = x0 , where x ( t ) ∈ X is the state , u ( t ) ∈ U is the input , the loss and ODE functions L and f are given , and the parameters θ have to be learned . The spaces X and U are typically Sobolev ( e.g . Hilbert ) spaces expressing the smoothness of x ( t ) and u ( t ) ( see Section 8 ) . Equation ( 1 ) can be used to represent several inverse problems , for instance in machine learning , estimation , and optimal control ( Stengel , 1994 ; Law et al. , 2015 ; Ross , 2009 ) . Problem ( 1 ) can be solved using gradient-based optimization through several time-stepping schemes for solving the ODE . ( Chen et al. , 2018 ; Gholami et al. , 2019 ) have proposed to use the adjoint method when f is a neural network . These methods are typically relying on explicit time-stepping schemes ( Butcher & Wanner , 1996 ) . Limitations of these approaches are briefly summarized : Limitations of backpropagation through an ODE solver . The standard approach for solving this problem is to compute the gradients ∂L/∂θ using backpropagation through a discrete approximation of the constraints , such as Runge-Kutta methods ( Runge , 1895 ; Butcher & Wanner , 1996 ) or multistep solvers ( Raissi et al. , 2018 ) . This ensures that the solution remains feasible ( within a numerical tolerance ) at each iteration of a gradient descent method . However , it has several drawbacks : 1 ) the memory cost of storing intermediate quantities during backpropagation can be significant , 2 ) the application of implicit methods would require solving a nonlinear equation at each step , 3 ) the numerical error can significantly affect the solution , and 4 ) the problem topology can be unsuitable for optimization ( Petersen et al. , 2018 ) . Limitations of adjoint methods . ODE-Nets ( Chen et al. , 2018 ) solve ( 1 ) using the adjoint method , which consists of simulating a dynamical system defined by an appropriate augmented Hamiltonian ( Ross , 2009 ) , with an additional state referred to as the adjoint . In the backward pass the adjoint ODE is solved numerically to provide the gradients of the loss function . This means that intermediate states of the forward pass do not need to be stored . An additional step of the ODE solver is needed for the backward pass . This suffers from a few drawbacks : 1 ) the dynamics of either the hidden state or the adjoint might be unstable , due to the symplectic structure of the underlying Hamiltonian system , referred to as the curse of sensitivity in ( Ross , 2009 ) ; 2 ) the procedure requires solving a differential algebraic equation and a boundary value problem which is complex , time consuming , and might not have a solution ( Ross & Karpenko , 2012 ) . Limitations of hybrid methods . ANODE ( Gholami et al. , 2019 ) splits the problem into time batches , where the adjoint is used , storing in memory only few intermediate states from the forward pass . This allows to improve the robustness and generalization of the adjoint method . A similar improvement could be obtained using reversible integrators . However , its computational cost is of the same order of the adjoint method and it does not offer further opportunities for parallelization . 3 RELAXATION OF THE SUPERVISED LEARNING PROBLEM . Our algorithm is based on two ingredients : i ) the discretization of the problem using spectral elements leading to SNODE , detailed in Section 4 , and ii ) the relaxation of the ODE constraint from ( 1 ) , enabling efficient training through backpropagation . The latter can be applied directly at the continuous level and significantly reduces the difficulty of the optimization , as shown in our examples . The problem in ( 1 ) is split into two smaller subproblems : one finds the trajectory x ( t ) that minimizes an unconstrained relaxation of ( 1 ) . The other trains the network weights θ such that the trajectory becomes a solution of the ODE . Both are addressed using standard gradient descent and backpropagation . In particular , a fixed number of ADAM or SGD steps is performed for each problem in an alternate fashion , until convergence . In the following , the details of each subproblem are discussed . Step 0 : Initial trajectory . The initial trajectory x ( t ) is chosen by solving the problem min x ( t ) ∈X ∫ t1 t0 L ( t , x ( t ) ) dt , ( 2 ) s.t . x ( t0 ) = x0 , If this problem does not have a unique solution , a regularization term is added . For a quadratic loss , a closed-form solution is readily available . Otherwise , a prescribed number of SGD iterations is used . Step 1 : Coordinate descent on residual . Once the initial trajectory x ? ( t ) is found , θ is computed by solving the unconstrained problem : min θ∈Rm ∫ t1 t0 R ( t , x ? ( t ) , θ ) dt , R ( t , x ( t ) , θ ) : = ‖ẋ ( t ) − f ( t , x ( t ) , u ( t ) ; θ ) ‖2 ( 3 ) If the value of the residual at the optimum θ∗ is smaller than a prescribed tolerance , then the algorithms stops . Otherwise , steps 1 and 2 are iterated until convergence . Step 2 : Coordinate descent on relaxation . Once the candidate parameters θ ? are found , the trajectory is updated by minimizing the relaxed objective : min x ( t ) ∈X ∫ t1 t0 γL ( t , x ( t ) ) +R ( t , x ( t ) , θ ? ) dt , ( 4 ) s.t . x ( t0 ) = x0 . Discussion . The proposed algorithm can be seen as an alternating coordinate gradient descent on the relaxed functional used in problem ( 4 ) , i.e. , by alternating a minimization with respect to x ( t ) and θ . If γ = 0 , multiple minima can exist , since each choice of the parameters θ would induce a different dynamics x ( t ) , solution of the original constraint . For γ 6= 0 , the loss function in ( 4 ) trades-off the ODE solution residual for the data fitting , providing a unique solution . The choice of γ implicitly introduces a satisfaction tolerance ( γ ) , i.e. , similar to regularized regression ( Hastie et al. , 2001 ) , implying that ‖ẋ ( t ) − f ( t , x ( t ) ; θ ) ‖ ≤ ( γ ) . Concurrently , problem ( 3 ) reduces the residual . 4 SNODE – HIGH-ORDER DISCRETIZATION OF THE RELAXED PROBLEM In order to numerically solve the problems presented in the previous section , a discretization of x ( t ) is needed . Rather than updating the values at time points ti from the past to the future , we introduce a compact representation of the complete discrete trajectory by means of the spectral element method . Spectral approximation . We start by representing the scalar unknown trajectory , x ( t ) , and the known input , u ( t ) , as truncated series : x ( t ) = p∑ i=0 xiψi ( t ) , u ( t ) = z∑ i=0 uiζi ( t ) , ( 5 ) where xi , ui ∈ R and ψi ( t ) , ζi ( t ) are sets of given basis functions that span the spaces Xh ⊂ X and Uh ⊂ U . In this work , we use orthogonal Legendre polynomials of order p ( Canuto et al. , 1988 ) for ψi ( t ) , where p is a hyperparameter , and the cosine Fourier basis for ζi ( t ) , where z is fixed . Collocation and quadrature . In order to compute the coefficients xi of ( 5 ) , we enforce the equation at a discrete set Q of collocation points tq . Here , we choose p+ 1 Gauss-Lobatto nodes , which include t = t0 . This directly enforces the initial condition . Other choices are also possible ( Canuto et al. , 1988 ) . Introducing the vectors of series coefficients xI = { xi } pi=0 and of evaluations at quadrature points x ( tQ ) = { x ( tq ) } q∈Q , the collocation problem can be solved in matrix form as xI = M −1x ( tQ ) , Mqi : = ψi ( tq ) . ( 6 ) We approximate the integral ( 3 ) as a sum of residual evaluations over Q . Assuming that x ( 0 ) = x0 , the integrand at all quadrature points tQ can be computed as a component-wise norm R ( tQ , x ( tQ ) , θ ) = ‖DM−1x ( tQ ) − f ( tQ , x ( tQ ) , u ( tQ ) ; θ ) ‖2 , Dqi : = ψ̇i ( tq ) . ( 7 ) Fitting the input data . For the case when problem ( 2 ) admits a unique a solution , we propose a new direct training scheme , δ-SNODE , which is summarized in Algorithm 1 . In general , a least-squares approach must be used instead . This entails computing the integral in ( 2 ) , which can be done by evaluating the loss function L at quadrature points tq . If the input data is not available at tq , we approximate the integral by evaluating L at the available time points . The corresponding alternating coordinate descent scheme α-SNODE is presented in Algorithm 2 . In the next sections , we study the consequences of a low-data scenario on this approach . We use fixed numbers Nt and Nx of updates for , respectively , θ and x ( t ) . Both are performed with standard routines , such as SGD . In our experiments , we use ADAM to optimize the parameters and an interpolation order p = 14 , but any other orders and solvers are possible . Algorithm 1 δ-SNODE training Input : M , D from ( 6 ) - ( 7 ) . x∗I ← argminx L ( Mx ) while R ( Mx∗I , θ∗ ) > δ do θ∗ ← ADAMθ [ R ( Mx∗I , θ ) ] end while Output : θ∗ Algorithm 2 α-SNODE training Input : M , D , θ∗ ∈ Rm , x∗I ∈ Rp , γ > 0. while γL ( Mx ? I ) +R ( Mx∗I , θ∗ ) > δ do for i = 0 . . . Nx do x∗I ← SGDx [ γL ( Mx ) +R ( Mx , θ∗ ) ] end for for i = 0 . . . Nt do θ∗ ← ADAMθ [ R ( Mx∗I , θ ) ] end for end while Output : θ∗ Ease of time parallelization . If R ( tq ) = 0 is enforced explicitly at q ∈ Q , then the resulting discrete system can be seen as an implicit time-stepping method of order p. However , while ODE integrators can only be made parallel across the different components of x ( t ) , the assembly of the residual can be done in parallel also across time . This massively increases the parallelization capabilities of the proposed schemes compared to standard training routines . Memory cost . If an ODE admits a regular solution , with regularity r > p , in the sense of Hilbert spaces , i.e. , of the number of square-integrable derivatives , then the approximation error of the SNODE converges exponentially with p ( Canuto et al. , 1988 ) . Hence , it produces a very compact representation of an ODE-Net . Thanks to this property , p is typically much lower than the equivalent number of time steps of explicit or implicit schemes with a fixed order . This greatly reduces the complexity and the memory requirements of the proposed method , which can be evaluated at any t via ( 5 ) by only storing few xi coefficients . Stability and vanishing gradients . The forward Euler method is known to have a small region of convergence . In other words , integrating very fast dynamics requires a very small time step , dt , in order to provide accurate results . In particular , for the solver error to be bounded , the eigenvalues of the state Jacobian of the ODE need to lie into the circle of the complex plane centered at ( −1 , 0 ) with radius 1/dt ( Ciccone et al. , 2018 ; Isermann , 1989 ) . Higher-order explicit methods , such as Runge-Kutta ( Runge , 1895 ) , have larger but still limited convergence regions . Our algorithms on the other hand are implicit methods , which have a larger region of convergence than recursive ( explicit ) methods ( Hairer et al. , 1993 ) . We claim that this results in a more stable and robust training . This claim is supported by our experiments . Reducing the time step can improve the Euler accuracy but it can still lead to vanishing or exploding gradients ( Zilly et al. , 2016 ; Goodfellow et al. , 2016 ) . In Appendix C , we show that our methods do not suffer from this problem . Experiments setup and hyperparameters . For all experiments , a common setup was employed and no optimization of hyperparameters was performed . Time horizon T = 10s and batch size of 100 were used . Learning rates were set to 10−2 for ADAM ( for all methods ) and 10−3 for SGD ( for α-SNODE ) . For the α-SNODE method , γ = 3 and 10 iterations were used for the SGD and ADAM algorithms at each epoch , as outlined in Algorithm 2 . The initial trajectory was perturbed as x0 = x0 + ξ , ξ ∼ U ( −0.1 , 0.1 ) . This perturbation prevents the exact convergence of Algorithm 1 during initialization , allowing to perform the alternating coordinate descent algorithm .

Summary: This work proposes a new approach for the evolution of Neural ODEs for particle systems. The authors suggest to replace the traditional backpropagation through ODEs or the recent adjoint method for backpropagation and instead solve the problem as an alternating optimization scheme. In particular it is suggested that using spectral methods (Legendre's polynomials) first compute a minimizer of the trajectory (optimizing a trajectory x(t)) given the Loss/Langrangian (that is based on the data times of training). After an initial trajectory is computed, follow an alternating minimization, where in the first step, minimize the discrepancy between the network (that describes the time change of the ODE) and the time derivative of the current trajectory. In the second step, re-compute the trajectory with the new updated network and modified lagrangian to update the network parameters again. This two step optimization is applied back and forth until a residual condition is reached, i.e. the loss is small. The network's parameters in the two stages are optimized via SGD and ADAM respectively. Further, to perform the required numerical integration in each step the authors apply Gaussian quadrature and turn the overall optimization scheme into the repeating application of a finite number of gradient updates in an alternating fashion (as has become popular in recent Deep Learning approaches e.g. GANs). The authors test the new method on different particle systems' trajectories and observe a numerical speed-up and improved accuracy as compared with previous approaches.