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

TL;DR

Rex introduces a family of reversible exponential (stochastic) Runge-Kutta solvers that enable exact inversion of neural differential equation models, improving stability and accuracy in generative tasks.

Contribution

The paper proposes Rex, a novel reversible exponential (stochastic) Runge-Kutta solver family, extending explicit schemes with Lawson methods for better inversion in neural differential models.

Findings

Enhanced sampling of Boltzmann distributions

Improved image generation and editing

Theoretically rigorous with empirical validation

Abstract

Deep generative models based on neural differential equations have quickly become the state-of-the-art for numerous generation tasks across many different applications. These models rely on ODE/SDE solvers which integrate from a prior distribution to the data distribution. In many applications it is highly desirable to then integrate in the other direction. The standard solvers, however, accumulate discretization errors which don't align with the forward trajectory, thereby prohibiting an exact inversion. In applications where the precision of the generative model is paramount this inaccuracy in inversion is often unacceptable. Current approaches to solving the inversion of these models results in significant downstream issues with poor stability and low-order of convergence; moreover, they are strictly limited to the ODE domain. In this work, we propose a new family of reversible exponential (stochastic) Runge-Kutta solvers which we refer to as Rex developed by an application of Lawson methods to convert any explicit (stochastic) Runge-Kutta scheme into a reversible one. In addition to a rigorous theoretical analysis of the proposed solvers, we also empirically demonstrate the utility of Rex on improving the sampling of Boltzmann distributions with flow models, and improving image generation and editing capabilities with diffusion models.