Learning Differential Equations that are Easy to Solve

TL;DR

This paper introduces a method to train neural differential equations that are easier to solve numerically by optimizing a differentiable surrogate for solver time cost, leading to faster models without significant accuracy loss.

Contribution

The authors propose a novel differentiable surrogate for numerical solver time cost, enabling the training of neural differential equations that are computationally more efficient to solve.

Findings

Models trained with the surrogate are substantially faster to solve.

Achieves near state-of-the-art accuracy in various tasks.

Reduces computational cost without sacrificing performance.

Abstract

Differential equations parameterized by neural networks become expensive to solve numerically as training progresses. We propose a remedy that encourages learned dynamics to be easier to solve. Specifically, we introduce a differentiable surrogate for the time cost of standard numerical solvers, using higher-order derivatives of solution trajectories. These derivatives are efficient to compute with Taylor-mode automatic differentiation. Optimizing this additional objective trades model performance against the time cost of solving the learned dynamics. We demonstrate our approach by training substantially faster, while nearly as accurate, models in supervised classification, density estimation, and time-series modelling tasks.