[Paper Review] Learning Differential Equations that are Easy to Solve

Jacob Kelly, Jesse Bettencourt+2 more|arXiv (Cornell University)|Jul 9, 2020

Model Reduction and Neural Networks38 references47 citations

TL;DR

The paper introduces a differentiable speed regularizer that minimizes higher-order time derivatives of neural ODE trajectories to reduce the number of function evaluations in adaptive solvers, trading off speed for accuracy.

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.

Motivation & Objective

• Motivate reducing numerical solving cost in neural ODEs by controlling trajectory complexity.
• Propose a differentiable surrogate for time cost based on higher-order derivatives of z(t).
• Show that speed-regularized models can be substantially faster to evaluate with limited accuracy loss across tasks.

Proposed method

• Define R_K(theta) as the integral of the squared norm of the K-th time derivative of the state along the trajectory.
• Add lambda * R_K(theta) to the training objective to trade accuracy for solver speed.
• Use Taylor-mode automatic differentiation to compute total derivatives efficiently with O(K^2) cost.
• Augment the ODE system to include the integrand of R_K so it can be computed in a single solver pass.
• Empirically evaluate speed-regularized models on supervised MNIST, PhysioNet time-series, and FFJORD density estimation.

Experimental results

Research questions

• RQ1 Does regularizing higher-order derivatives of neural ODE trajectories reduce the required number of function evaluations (NFE) of adaptive solvers?
• RQ2 What is the impact of matching the regularization order K to the solver order on speed vs. accuracy?
• RQ3 How does speed regularization affect performance across classification, time-series modeling, and density estimation tasks?
• RQ4 Is Taylor-mode AD essential for practical computation of the regularizer?
• RQ5 Do speed-regularized models exhibit stable training and generalization behavior?

Key findings

• Regularizing the K-th derivative can reduce solver evaluations by approximately a factor of 2 in several tasks with modest loss increase.
• Matching the regularization order to the solver order generally yields the best speed-accuracy trade-offs, with diminishing returns for higher orders beyond a point.
• Speed regularization substantially lowers NFE during training and evaluation, while test-time evaluation becomes faster due to simpler dynamics.
• Taylor-mode automatic differentiation provides efficient computation of total derivatives, enabling practical use of R_K during training.
• Compared to prior regularizers, the approach can maintain comparable accuracy while achieving faster inference in neural ODEs.