arXiv · 2502.08834
Rex: A Family of Reversible Exponential (Stochastic) Runge-Kutta Solvers
Abstract
Deep generative models based on neural differential equations have become state-of-the-art for many generation tasks. These models rely on ODE/SDE solvers that integrate from a prior distribution to the data distribution; in many applications it is also highly desirable to integrate in the inverse direction. Standard solvers, however, accumulate discretization errors that prohibit exact inversion, an inaccuracy that is unacceptable in precision-critical applications. Existing inversion methods suffer from poor stability and low order of convergence, and are strictly limited to the ODE setting. In this work, we propose Rex, a family of reversible exponential (stochastic) Runge-Kutta solvers obtained by applying Lawson methods to convert any explicit (stochastic) Runge-Kutta scheme into an algebraically reversible one for both diffusion ODEs and SDEs. Beyond a rigorous theoretical analysis -- establishing arbitrary-order convergence and a non-zero region of linear stability -- we empirically demonstrate that Rex achieves near-machine-precision reconstruction and improves Boltzmann sampling with flow models as well as image generation and editing with diffusion models.
Explore related subjects
Keep this discovery
Zander W. Blasingame, Chen Liu. 2025-02-12. Rex: A Family of Reversible Exponential (Stochastic) Runge-Kutta Solvers. https://arxiv.org/abs/2502.08834
Cite the original work for its findings. Save a collection to share your selection of sources.