[论文解读] Learning Differential Equations that are Easy to Solve

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

Model Reduction and Neural Networks参考文献 38被引用 47

一句话总结

本文提出一种可微分的速度正则化器，通过最小化神经 ODE 路径的高阶时间导数来减少自适应求解器的函数评估次数，在速度与准确性之间进行权衡。

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.

研究动机与目标

• 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.

提出的方法

• 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.

实验结果

研究问题

• 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?

主要发现

• 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.