Rex: Reversible Solvers for Diffusion Models

TL;DR

Rex introduces a reversible solver family for diffusion models that achieves arbitrarily high-order convergence in the probability-flow ODE and exact inversion for reverse-time SDEs without storing Brownian paths. By combining McCallum-Foster algebraic reversibility with Lawson exponential integrators and a data/noise-prediction reparameterization, Rex delivers a unified, exact-reversibility framework for diffusion models. Theoretical results establish convergence and stability properties, while extensive experiments on unconditional and conditional image generation, as well as interpolation, demonstrate practical gains over prior reversible methods. These contributions enable reliable, high-fidelity inversion and sampling in diffusion-based generative modeling, with broad implications for editing, inversion tasks, and downstream applications.

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.

Figures (14)

• Figure 1: The computation graph of the Rex solver. Here $\bm \Psi_h$ denotes an exponentially weighted Runge-Kutta scheme (cf.\ref{['sec:rex_ode']}) or exponential stochastic Runge-Kutta scheme (cf.\ref{['sec:rex_sde']}), $\zeta \in (0, 1)$ is a coupling parameter, and $\{w_n\}_{n=1}^N$ denotes the set of weighting variables derived from the exponential schemes. For ODEs we have $w_n = \sigma_n$ and for SDEs we have $w_n = \frac{\sigma_n^2}{\alpha_n}$. The visualization of the computation graph is inspired by mccallum2024efficient.
• Figure 2: Overview of the construction of $\bm \Psi$ for the probability flow ODE from an underlying RK scheme $\bm \Phi$ for the reparameterized ODE. This graph holds for the SDE and noise prediction cases mutatis mutandis.
• Figure 3: Qualitative comparison of unconditional sampling with different reversible solvers with a pre-trained DDPM model on CelebA-HQ ($256 \times 256$) with the non-reversible DDIM as a baseline. Each method used 10 discretization steps.
• Figure 4: Qualitative comparison of text-to-image conditional sampling with different reversible solvers with Stable Diffusion v1.5 ($512 \times 512$) and 10 discretization steps. Prompts from top to bottom are: "White plate with fried fish and lemons sitting on top of it.", "A lady enjoying a meal of some sort.", and "A young boy riding skis with ski poles.".
• Figure 5: Unconditional interpolation between two real images from FRLL frll with a DDPM model trained on CelebA-HQ. Top row is BELM, middle is Rex (Euler), and bottom is Rex (ShARK). 50 steps used for each method.

Theorems & Definitions (99)

• Definition 1: McCallum-Foster method
• Remark 2
• Proposition 0: Time reparameterization of the probability flow ODE
• Proposition 0: Time reparameterization of the reverse-time diffusion SDE
• Definition 3: Space-time Lévy area
• Proposition 0: Rex
• Theorem 1: Rex is a $k$-th order solver
• Theorem 2: Convergence order for stochastic $\bm \Psi$
• Definition 4: Asynchronous leapfrog method
• Remark 5