Explicit and Effectively Symmetric Schemes for Neural SDEs
This addresses a bottleneck in scalable and accurate training of neural SDEs for researchers and practitioners in machine learning, offering a practical improvement over existing reversible methods.
The paper tackles the trade-off between memory efficiency and gradient accuracy in training neural SDEs by introducing stable, near-reversible Runge-Kutta schemes, which enable memory-efficient training without severe step size or model restrictions, as demonstrated through numerical experiments.
Backpropagation through (neural) SDE solvers is traditionally approached in two ways: discretise-then-optimise, which offers accurate gradients but incurs prohibitive memory costs due to storing the full computational graph (even when mitigated by checkpointing); and optimise-then-discretise, which achieves constant memory cost by solving an auxiliary backward SDE, but suffers from slower evaluation and gradient approximation errors. Algebraically reversible solvers promise both memory efficiency and gradient accuracy, yet existing methods such as the Reversible Heun scheme are often unstable under complex models and large step sizes. We address these limitations by introducing a novel class of stable, near-reversible Runge--Kutta schemes for neural SDEs. These Explicit and Effectively Symmetric (EES) schemes retain the benefits of reversible solvers while overcoming their instability, enabling memory-efficient training without severe restrictions on step size or model complexity. Through numerical experiments, we demonstrate the superior stability and reliability of our schemes, establishing them as a practical foundation for scalable and accurate training of neural SDEs.