Rex: Reversible Solvers for Diffusion Models

ICLR 2026 Conference Submission22418 Authors

Keywords: numerical methods, neural differential equations, stochastic processes, stochastic differential equations, runge-kutta methods, diffusion models

TL;DR: We propose a family of algebraically reversible solvers for diffusion models including SDEs

Abstract:

Diffusion models have quickly become the state-of-the-art for numerous generation tasks across many different applications.
Encoding samples from the data distribution back into the models underlying prior distribution is an important task that arises in many downstream applications.
This task is often called the *inversion* of diffusion models.
Prior approaches for solving this task, however, are often simple heuristic solvers that come with several drawbacks in practice.
In this work, we propose a new family of solvers for diffusion models by exploiting the connection between this task and the broader study of *algebraically reversible* solvers for differential equations.
In particular, we construct a family of reversible solvers using an application of Lawson methods to construct exponential Runge-Kutta methods for the diffusion models.
We call this family of reversible exponential solvers *Rex*.
In addition to a rigorous theoretical analysis of the proposed solvers we also emonstrate the utility of the methods through a variety of empirical illustrations.

Primary Area: generative models

Submission Number: 22418