Faster Training of Neural ODEs Using Gau{\ss}-Legendre Quadrature

TL;DR

This paper introduces a novel approach to accelerate neural ODE training by employing Gau{e2}ss-Legendre quadrature for faster integral computation, improving efficiency especially for large models, and extends this method to SDE training via the Wong-Zakai theorem.

Contribution

The paper proposes using Gau{e2}ss-Legendre quadrature to speed up neural ODE training and extends this technique to SDEs through the Wong-Zakai theorem, offering a more efficient training method.

Findings

Faster training of neural ODEs demonstrated with large models.

Memory-efficient integral computation using Gau{e2}ss-Legendre quadrature.

Extension of the method to SDE training via the Wong-Zakai theorem.

Abstract

Neural ODEs demonstrate strong performance in generative and time-series modelling. However, training them via the adjoint method is slow compared to discrete models due to the requirement of numerically solving ODEs. To speed neural ODEs up, a common approach is to regularise the solutions. However, this approach may affect the expressivity of the model; when the trajectory itself matters, this is particularly important. In this paper, we propose an alternative way to speed up the training of neural ODEs. The key idea is to speed up the adjoint method by using Gau{\ss}-Legendre quadrature to solve integrals faster than ODE-based methods while remaining memory efficient. We also extend the idea to training SDEs using the Wong-Zakai theorem, by training a corresponding ODE and transferring the parameters. Our approach leads to faster training of neural ODEs, especially for large models. It also presents a new way to train SDE-based models.