Efficient, Accurate and Stable Gradients for Neural ODEs

Sam McCallum, James Foster

10/15/2024

cs.LGstat.ML

No official code found. Best implementation: sammccallum/reversible-solvers (5 stars, last commit 2025-01-29). 1 more below.

Abstract

Training Neural ODEs requires backpropagating through an ODE solve. The state-of-the-art backpropagation method is recursive checkpointing that balances recomputation with memory cost. Here, we introduce a class of algebraically reversible ODE solvers that significantly improve upon both the time and memory cost of recursive checkpointing. The reversible solvers presented calculate exact gradients, are high-order and numerically stable -- strictly improving on previous reversible architectures.

View on arXiv View PDF

Code Implementations(2)

sammccallum/reversible-solversCommunity100%

Algebraically reversible differential equation solvers.

50Shell, PythonNov 27, 20241 years ago

sammccallum/reversibleCommunity100%

JAX implementation of the Reversible Solver method.

23Shell, PythonOct 14, 20241 years ago

Similar Papers

Explicit and Effectively Symmetric Schemes for Neural SDEs

88% match1 repo42025

Mixed Precision Training of Neural ODEs

85% match2025

Adaptive Feedforward Gradient Estimation in Neural ODEs

84% match2024

PETRA: Parallel End-to-end Training with Reversible Architectures

83% match1 repo2024

MinBackProp -- Backpropagating through Minimal Solvers

82% match1 repo92024

Stabilizing Backpropagation Through Time to Learn Complex Physics

81% match1 repo112024

Cite this paper

@article{mccallum2024efficient,
  title  = {Efficient, Accurate and Stable Gradients for Neural ODEs},
  author = {Sam McCallum and James Foster},
  year   = {2024},
  eprint = {2410.11648},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2410.11648v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

[![Code of Paper](https://codeofpaper.com/badge/2410.11648.svg)](https://codeofpaper.com/paper/2410.11648)