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.

翻译：基于神经微分方程的深度生成模型已成为许多生成任务中的最先进方法。这些模型依赖于从先验分布积分至数据分布的ODE/SDE求解器；在许多应用中，实现逆方向积分也具有高度需求。然而，标准求解器会累积离散化误差，导致无法精确反转，这种不精确性在精度关键的应用中不可接受。现有逆方法存在稳定性差、收敛阶数低的问题，且严格局限于ODE场景。本文提出Rex，一类可逆指数（随机）龙格-库塔求解器家族，通过应用Lawson方法将任意显式（随机）龙格-库塔格式转换为适用于扩散ODE和SDE的代数可逆格式。除了严格的理论分析（建立任意阶收敛性和非零线性稳定区域），我们通过实验证明，Rex能够实现接近机器精度的重建，并改进基于流模型的玻尔兹曼采样以及基于扩散模型的图像生成与编辑。