Rex: Algebraically reversible solvers for diffusion models

Diffusion sampling is a flow; existing solvers cannot reverse it exactly

Diffusion models  have quickly become the state of the art across many generation tasks; sampling from them amounts to integrating an ODE or SDE from a simple prior to a complex data distribution . For diffusion models the reverse-time generative process is itself an SDE , and many downstream applications require solving the same equation in the encoding direction. Encoding samples from the data distribution back into the model's underlying prior is referred to as the inversion of the diffusion model, and it requires a bijective map between the two distributions.

Inversion underpins a range of tasks: image editing encodes a real image back to its latent noise and re-decodes under a new prompt ; Boltzmann sampling uses the change-of-variables formula to assign exact likelihoods ; gradient-based fine-tuning and reward-guided sampling require accurate adjoints . In each of these settings the round trip $\boldsymbol x_0 \!\to\! \boldsymbol x_T \!\to\! \boldsymbol x_0$ must close exactly; an approximate inverse is not sufficient.

The continuous flow itself is a bijection; truncation error from discretization is what breaks the round trip. Several prior works have proposed reversible solvers for the probability flow ODE, namely EDICT , BDIA , and BELM . These schemes are, however, plagued by issues of low order of convergence and a lack of linear stability, amongst other undesirable properties. To the best of our knowledge, no prior scheme achieves exact inversion for diffusion SDEs without storing the entire Brownian trajectory in memory, which precludes adaptive step sizes and is impractical at scale.

Moreover, reversibility has so far been constructed one solver at a time; no general procedure exists for converting an explicit (S)RK scheme — Euler, midpoint, RK4, Dormand–Prince, Euler–Maruyama — into a reversible solver tailored to the semi-linear structure of the diffusion ODE/SDE. This is the gap Rex closes.

Two largely independent research traditions have circled this problem from opposite directions. The first grew up inside the diffusion community itself, retrofitting bijective structure onto specific samplers: EDICT’s coupled latent trick, BDIA’s bidirectional integration, and BELM’s linear-multistep construction each reverse-engineer reversibility into one particular solver. The second comes from the neural differential-equation community, where reversibility has been pursued as a memory-efficiency device for training — from the original Neural ODE adjoint to MALI, reversible Heun, and the recent McCallum–Foster framework. Rex sits at the confluence: it borrows the McCallum–Foster coupling from the neural-DE lineage and adapts it to the exponential-integrator structure that the diffusion lineage relies on, yielding a single construction that subsumes both.

The construction itself is the subject of the next section. We first present the explicit (S)RK family that Rex operates on, then the exponential–integrator step that absorbs the semi-linear drift, and finally the McCallum–Foster coupling that closes the round trip in closed form.

Constructing reversible solvers from arbitrary explicit (S)RK schemes

We propose Rex, a Reversible Exponential (Stochastic) Runge–Kutta family of solvers for diffusion models. Given any explicit (S)RK scheme $\boldsymbol\Phi$, Rex constructs an algebraically reversible solver tailored to the semi-linear structure of the diffusion ODE/SDE. The construction proceeds in two stages: an exponential-integrator step which absorbs the linear drift, followed by a McCallum–Foster-style coupling which yields a closed-form inverse.

Algebraic reversibility

Every numerical scheme is reversible in the analytic sense: one can recover the previous step by fixed-point iteration on the update equation.Provided the step size $h$ is small enough for the fixed-point iteration to converge. This is, however, both expensive and only approximate. A scheme is algebraically reversible if the inverse map admits a closed-form expression. Rex maintains an auxiliary state $\hat{\boldsymbol x}_n$ alongside $\boldsymbol x_n$ precisely so that the inverse map can be written in closed form.

The coupling parameter $\zeta \in (0,1]$ controls a trade-off between linear stability and exact inversion.$\zeta = 1$ recovers exact algebraic reversibility; $\zeta < 1$ introduces a small mixing of the auxiliary state which enlarges the linear stability region at the cost of a residual inversion error. We use $\zeta = 1$ throughout the empirical results. The weights $\{\kappa_n\}$ and reparameterized time variable $\varsigma$ are determined by the noise schedule $(\alpha_t, \sigma_t)$; specializing these recovers the data-prediction, noise-prediction, ODE, and SDE cases under a single scheme.

Reversibility for SDEs poses an additional challenge: a naïve construction stores the entire realization of the Brownian motion in memory to replay it in reverse,This is the approach taken by CycleDiffusion and similar reversible-by-storage methods. which prohibits adaptive step sizes. Rex avoids this by fixing a single seed $\omega \in \Omega$ and reconstructing any sub-interval of the Brownian trajectory on demand through a splittable PRNG, following the Brownian interval construction of Kidger et al. . The resulting scheme supports adaptive step-size SDE solvers while remaining exactly reversible.

Rex inherits convergence and stability from $\boldsymbol\Phi$, and subsumes prior diffusion solvers

The numerical properties of Rex are inherited from the underlying scheme $\boldsymbol\Phi$: the order of convergence, the region of linear stability, and (in the SDE case) the strong order of convergence are all preserved by the construction. To the best of our knowledge, Rex is also the first reversible solver framework that achieves exact inversion for diffusion SDEs without storing the entire realization of the Brownian motion .

Princeps subsumes prior diffusion solvers

A consequence of the Lawson-based construction is that Princeps $\boldsymbol\Psi$, instantiated with appropriate choices of the base scheme $\boldsymbol\Phi$, reduces exactly to many of the most popular numerical solvers used in the diffusion literature.A full statement appears in the paper as Theorem 3.10 — Rex subsumes previous solvers; the proofs are deferred to the appendix. Each row of the table below identifies a choice of $\boldsymbol\Phi$, the existing solver it recovers, and the corresponding algebraically reversible Rex variant.

Comparison with prior reversible solvers

Empirical results across inversion, generation, editing, and Boltzmann sampling

We evaluate Rex in five regimes: (a) the empirical inversion error of a Stable Diffusion v1.5 round trip; (b) unconditional image generation on CelebA-HQ; (c) text-to-image generation on Stable Diffusion v1.5 with COCO captions; (d) round-trip image editing , for which exact inversion is essential; and (e) Boltzmann sampling of molecular conformations , for which an approximate inverse invalidates the change-of-variables likelihood and breaks self-normalized importance sampling.

Inversion error · Stable Diffusion v1.5

For an algebraically reversible solver, encoding a Stable Diffusion  latent and then decoding it should recover the original latent up to floating-point rounding error. We measure the latent-space MSE of this round trip on Stable Diffusion v1.5 in fp32 at 10, 20, and 50 sampling steps; the results are shown in Figure 3.

Several observations follow. DDIM is not reversible, so its error simply reflects the truncation error of the underlying solver. EDICT and BDIA both improve with the number of steps but plateau in the $10^{-7}$–$10^{-9}$ range due to numerical instability. O-BELM exhibits a more striking failure mode: the error grows with the number of steps, since the iteration is nowhere linearly stable and round-off error accumulates.This is consistent with the theoretical claim in Section 03: a scheme with no non-trivial region of linear stability cannot maintain bounded inversion error as the number of steps grows. Rex sits at the floating-point round-off floor and improves monotonically with step count.

Unconditional generation · CelebA-HQ 256×256 , DDPM

We evaluate $10^4$ samples under identical seeds, using the DINOv2  Fréchet distance , FD$_\infty$ , precision/recall , density, and coverage . Rex with Euler–Maruyama (the SDE variant) outperforms every prior reversible solver at 20 and 50 sampling steps; the ODE variants (Midpoint, RK4) match O-BELM at lower step counts and exceed it at higher ones. Figure 4 and Table 3 summarize the results.

Same initial noise · 10 discretization steps · DDPM on CelebA-HQ 256×256.

Text-to-image · Stable Diffusion v1.5, COCO captions

We sample under 1000 COCO captions with shared seeds, scored by CLIP score, ImageReward, and PickScore. The stochastic variants of Rex (Euler–Maruyama and ShARK) achieve higher aesthetic and prompt-alignment scores than every reversible baseline; both also exceed the non-reversible DDIM baseline. Figure 5 and Table 4 summarize the metrics across step counts.

"White plate with fried fish and lemons sitting on top of it." — 10 steps

"A young boy riding skis with ski poles." — 10 steps

"A lady enjoying a meal of some sort." — 10 steps

Round-trip image editing · pix2pix, Stable Diffusion v1.5

Round-trip editing is a canonical task for which exact inversion is required: a real image is encoded into noise under its original caption, then decoded under a modified edit caption. We evaluate on the pix2pix dataset, measuring prompt alignment with CLIP, ImageReward, and PickScore, and faithfulness to the source image with LPIPS.

Rex with the Dopri5 base scheme (an adaptive 5(4) Dormand–Prince method) is, to our knowledge, the first adaptive-step-size reversible solver applied to diffusion-based editing.Adaptive step-size SDE solvers, prior to Rex, required storing the entire Brownian trajectory in memory. The Brownian-interval construction discussed in Section 02 is what makes the adaptive variant reversible. EDICT failed entirely on this benchmark, collapsing to the identity map; BDIA also failed, with an LPIPS of $0.885$ — roughly an order of magnitude worse than the next-best reversible method. Rex preserves source-image faithfulness while still following the edit instruction, as shown in Figure 6 and Table 5.

Boltzmann sampling of molecular conformations · tri-alanine with DiT

Sampling from a Boltzmann distribution $p_{\text{target}}(\boldsymbol x) \propto \exp(-\mathcal E(\boldsymbol x))$ with a neural ODE  requires exact likelihoods, computed via the change-of-variables formula. If the discretized solver is not a bijection the resulting likelihoods are biased, and self-normalized importance sampling then produces estimators that are inconsistent in expectation .The bias does not, in general, decrease with more samples; it reflects a systematic mismatch between the integrated density and the true density.

We train a Diffusion Transformer on tri-alanine (Ala–Ala–Ala) and evaluate sampling quality with the effective sample size (ESS) and the 2-Wasserstein distance to the molecular-dynamics ground truth, taken with respect to the energy distribution ($\mathcal E\text{-}\mathcal W_2$) and the dihedral-angle distribution ($\mathbb T\text{-}\mathcal W_2$).

Substituting Rex (Dopri5) for the non-reversible Dopri5 baseline within the same DiT pipeline reduces the energy-distribution Wasserstein error from $0.737$ to $\mathbf{0.495}$, a $33\%$ reduction, while leaving the dihedral-angle metric essentially unchanged. This is consistent with the inversion-error analysis above: exact reversibility ensures that the change-of-variables likelihoods used in importance sampling are unbiased, which the non-reversible Dopri5 baseline cannot guarantee.

Conclusion

We have proposed Rex, a family of algebraically reversible solvers for diffusion models, constructed by combining Lawson exponential integration with a McCallum–Foster-style coupling. Rex inherits an arbitrarily high order of convergence from the underlying explicit RK scheme in the ODE case, and to the best of our knowledge furnishes the first method for exact inversion of diffusion SDEs without storing the entire realization of the Brownian motion. The intermediate Princeps family subsumes several widely used diffusion solvers — DDIM, DPM-Solver, gDDIM, and SEEDS-1 — as special cases, so that Rex yields algebraically reversible analogues of each. Empirically, Rex functions as a capable numerical scheme across unconditional and conditional image generation, round-trip image editing, and Boltzmann sampling of molecular conformations; in the Boltzmann setting, exact reversibility yields an unbiased change-of-variables likelihood and a $33\%$ reduction in the energy-distribution Wasserstein metric over the non-reversible baseline.