Learning Differential Equations that are Easy to Solve
This addresses a computational bottleneck for researchers and practitioners using neural differential equations, offering an incremental improvement in efficiency.
The paper tackles the problem of expensive numerical solving of neural network-parameterized differential equations during training by introducing a differentiable surrogate for solver time cost, enabling a trade-off between model performance and solving speed. The result is substantially faster models with nearly as good accuracy across classification, density estimation, and time-series tasks.
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.