Rex: A Family of Reversible Exponential (Stochastic) Runge-Kutta Solvers

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.

Problem. Modern AI image generators, protein designers, and molecular simulators work by gradually transforming random noise into a meaningful sample, one small step at a time. For many important uses — editing a real photograph, training the model from its outputs, or accurately measuring how likely a molecular shape is — we need to run this process in reverse, recovering the original noise from a finished sample. In practice, however, tiny rounding errors pile up at every step, so the reverse trip never lands exactly where it started. Existing fixes are unstable, only mildly accurate, and almost none of them work when the generation process involves randomness. Solution. We introduce Rex, a new family of numerical "step-takers" that are algebraically reversible: the backward step is the exact mathematical undo of the forward step, so the reconstruction is perfect by design — down to the precision of the computer's arithmetic. Rex works for both non-random and random generative models, and it is the reversible counterpart of many widely used solvers. Impact. Rex enables faithful image editing and more accurate sampling of molecular configurations — unlocking precision-critical applications that were previously out of reach.