Reproduce “Stiff Neural Ordinary Differential Equations” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Suyong Kim, Weiqi Ji, Sili Deng, Yingbo Ma, Christopher Rackauckas

Researchers at MIT and Julia Computing developed a framework that enables Neural Ordinary Differential Equations (NODEs) to robustly model stiff dynamical systems. The approach incorporates stable adjoint gradient computations and equation scaling techniques, successfully learning the dynamics of benchmark problems like Robertson's problem and a 20-species air pollution model.

Key insights

  • Traditional continuous adjoint methods are inherently unstable for stiff systems because reverse integration of a stiff ODE can lead to exponential error amplification.
  • Explicitly scaling neural network outputs and normalizing loss function components is crucial for overcoming gradient pathologies caused by large magnitude differences in state variables within stiff systems.
  • Despite using stiff ODE solvers, efficient training of NODEs for stiff systems requires advanced adjoint techniques (e.g., parameter splitting) to manage the cubic computational complexity of Jacobian factorizations with respect to parameters.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.