SNODE: Spectral Discretization of Neural ODEs for System Identification

TL;DR

This paper introduces a spectral element method for Neural ODEs that accelerates training and improves accuracy in system identification by using Legendre polynomial series and a parallelized optimization scheme.

Contribution

It presents a novel spectral discretization approach for Neural ODEs, enabling faster training and better generalization compared to standard methods.

Findings

At least ten times faster training speed.

One order of magnitude lower testing MSE.

Enhanced generalization capabilities.

Abstract

This paper proposes the use of spectral element methods \citep{canuto_spectral_1988} for fast and accurate training of Neural Ordinary Differential Equations (ODE-Nets; \citealp{Chen2018NeuralOD}) for system identification. This is achieved by expressing their dynamics as a truncated series of Legendre polynomials. The series coefficients, as well as the network weights, are computed by minimizing the weighted sum of the loss function and the violation of the ODE-Net dynamics. The problem is solved by coordinate descent that alternately minimizes, with respect to the coefficients and the weights, two unconstrained sub-problems using standard backpropagation and gradient methods. The resulting optimization scheme is fully time-parallel and results in a low memory footprint. Experimental comparison to standard methods, such as backpropagation through explicit solvers and the adjoint technique \citep{Chen2018NeuralOD}, on training surrogate models of small and medium-scale dynamical systems shows that it is at least one order of magnitude faster at reaching a comparable value of the loss function. The corresponding testing MSE is one order of magnitude smaller as well, suggesting generalization capabilities increase.