🤖 AI Summary
Neural ODE training suffers from high computational cost, excessive memory consumption, and numerical instability during backpropagation. To address these challenges, this paper introduces the Algebraically Invertible ODE Solver family, grounded in algebraically invertible numerical integration. The method integrates high-order implicit/explicit reversible schemes, adjoint-state techniques, and memory–computation co-optimization to achieve, for the first time, high-order accuracy, strict numerical stability, and exact gradient computation in backpropagation. Unlike recursive checkpointing, our approach achieves strictly superior time and memory complexity bounds. Extensive evaluation on multiple benchmark ODE tasks demonstrates a 2.1× reduction in training latency and a 68% decrease in GPU memory usage, while preserving gradient precision and numerical robustness.
📝 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.