Fast Neural ODE / UDE: Improved Parallelism and Memory Performance Differentiating Stiff ODEs

From ICIAM 2023 August 24th
iciam2023.org/registered_data...
[03338] Improved Parallelism and Memory Performance for Differentiating Stiff Differential Equations
Christopher Vincent Rackauckas (Julia Hub, Pumas-AI, MIT)
Abstract : Previous work demonstrated trade-offs in performance, numerical stability, and memory usage for ODE solving and differentiation of solutions. Our new time stepping methods expose more parallelism is shown to accelerate small ODE solves, while new GPU-based ODE solvers demonstrate a 10x performance improvement over Jax and PyTorch-based solvers. New adjoint methods achieve linear cost scaling with respect to parameters in stiff ODEs, as opposed to the cubic of Jax/PyTorch, while limiting the memory scaling.
I was on vacation so I had fun making that thumbnail. Isn't it awesome? I just CZcams'd the snizz out of that algorithm.
More information on the nonlinear solvers: sciml.ai/news/2024/01/23/nlso...