0% | | 0/0
ICML 2026

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

Zander W. Blasingame1,2 · Chen Liu2
1 AITHYRA, 2 Clarkson University
Aithyra Aithyra / Clarkson University Clarkson University
introduction

An Illustration of Reversible Solvers

$ sde_sampler.py — non-reversible
$ rex_solver.py — reversible
$$\mathrm{d}\boldsymbol X_t = \big[a(t)\,\boldsymbol X_t + b(t)\,\boldsymbol f_\theta(t,\boldsymbol X_t)\big]\,\mathrm{d}t + g(t)\,\mathrm{d}\boldsymbol W_t$$
data distribution — tri-alanine conformation forward SDE — diffuse data $\boldsymbol X_0 \to$ noise $\boldsymbol X_T$ reverse solve — non-reversible drifts, Rex retraces exactly
overview

Reversible Solvers via Coupling

To make a numerical scheme reversible, we need at least one auxiliary state — the same structural trick Normalizing Flows use in coupling layers.

forward · standard $\Phi$ · not reversible
$x_n$
$\Phi$
$x_{n+1}$
cannot invert
backward · standard $\Phi$ · not reversible
$x_{n+1}$
$\Phi$
$x_n$
cannot invert
forward · reversible
$x_n$
 + +
$x_{n+1}$
$\Psi_h$
$\Psi_{-h}$
$\hat x_n$
 −
$\hat x_{n+1}$
backward · reversible
$x_n$
 − −
$x_{n+1}$
$\Psi_h$
$\Psi_{-h}$
$\hat x_n$
 +
$\hat x_{n+1}$
e.g. McCallum–Foster coupling

Reversibility requires ≥ 1 auxiliary state — the same structural trick as coupling layers in Normalizing Flows.

motivation

Reversible Diffusion Solvers

2021 2022 2023 2024 2025 2026 DDIM Inversion non-exact ICLR 2021 EDICT coupled latents CVPR 2023 BDIA bidirectional ECCV 2024 O-BELM linear multistep NeurIPS 2024 MALI async leapfrog ICLR 2021 Reversible Heun reversible SDE solver NeurIPS 2021 McCallum–Foster arbitrary order / non-trivial stability Rex ICML 2026

Current diffusion inversion methods are bespoke schemes for diffusion ODEs which suffer from:

  • Low order of convergence
  • Poor linear stability
  • No support for the SDE formulation
  • No adaptive step sizing
  • Constructed one solver at a time — no general procedure

Can we design a general reversible solver — high-order, stable, and works for both ODEs and SDEs?

We propose

Rex: a family of Reversible Exponential (Stochastic) Runge–Kutta solvers.

Three main building blocks

What this gives us

  • High-order convergence & linear stability (ODEs)
  • Reversible adaptive step sizing
  • One framework for both diffuions ODEs and SDEs
  • Reversible versions of popular diffusion solvers
building blocks

Explicit Stochastic Runge-Kutta Schemes

ODE Runge-Kutta extended to SDEs via iterated stochastic integrals which are generally intractable. We use the space-time Lévy area and the Foster-Reis-Strange (S)RK scheme  to develop nice SDE solvers.

$s$ $t$ Brownian motion $W_t$ $W_s$ $W_t$ $H_{s,t}$

$$\class{srk-wcolor}{\boldsymbol W_{s,t}} \sim \mathcal N(\boldsymbol 0,\, h\,\boldsymbol I)$$

$$\class{srk-hcolor}{\boldsymbol H_{s,t}} \sim \mathcal N\!\left(\boldsymbol 0,\, \tfrac{h}{12}\,\boldsymbol I\right), \quad \class{srk-wcolor}{\boldsymbol W_{s,t}} \perp\!\!\!\perp \class{srk-hcolor}{\boldsymbol H_{s,t}}$$

"where" you end up | "how curvy" the path was 

$$\renewcommand\arraystretch{1.2}\begin{array}{c|c|c|c} \boldsymbol c & \boldsymbol a & \class{srk-wcolor}{\boldsymbol a^W} & \class{srk-hcolor}{\boldsymbol a^H} \\ \hline & \boldsymbol b & \class{srk-wcolor}{\boldsymbol b^W} & \class{srk-hcolor}{\boldsymbol b^H} \end{array}$$

Two new columns for the Brownian increment $\boldsymbol W$ and Lévy area $\boldsymbol H$.

$$\begin{array}{c|c} 0 & \\ \hline & 1 \end{array}$$

Euler — ODE, order 1

$$\renewcommand\arraystretch{1.5}\begin{array}{c|cccc} 0 & \\ \tfrac12 & \tfrac12 \\ \tfrac12 & 0 & \tfrac12 \\ 1 & 0 & 0 & 1 \\ \hline & \tfrac16 & \tfrac13 & \tfrac13 & \tfrac16 \end{array}$$

RK4 — ODE, order 4

$$\renewcommand\arraystretch{1.3}\begin{array}{c|cccccc} 0 & \\ \tfrac15 & \tfrac15 \\ \tfrac3{10} & \tfrac3{40} & \tfrac9{40} \\ \tfrac45 & \tfrac{44}{45} & \tfrac{-56}{15} & \tfrac{32}9 \\ \tfrac85 & \tfrac{19372}{6561} & \tfrac{-25360}{2187} & \tfrac{64448}{6561} & \tfrac{-212}{729} \\ 1 & \tfrac{9017}{3168} & \tfrac{-355}{33} & \tfrac{46732}{5247} & \tfrac{49}{176} & \tfrac{-5103}{18656} \\ 1 & \tfrac{35}{384} & 0 & \tfrac{500}{1113} & \tfrac{125}{192} & \tfrac{-2187}{6784} & \tfrac{11}{84} \\ \hline & \tfrac{35}{384} & 0 & \tfrac{500}{1113} & \tfrac{125}{192} & \tfrac{-2187}{6784} & \tfrac{11}{84} & 0 \\ & \tfrac{5179}{57600} & 0 & \tfrac{7571}{16695} & \tfrac{393}{640} & \tfrac{-92097}{339200} & \tfrac{187}{2100} & \tfrac1{40} \end{array}$$

Dormand–Prince 4(5) — ODE, adaptive (rows 7–8: order 4 / order 5)

$$\renewcommand\arraystretch{1.5}\begin{array}{c|c|c|c} 0 & 0 & 0 & 0 \\ \hline & 1 & \class{srk-wcolor}{1} & \class{srk-hcolor}{0} \end{array}$$

Euler–Maruyama — SDE, strong order 0.5

$$\renewcommand\arraystretch{1.5}\begin{array}{c|cc|c|c} 0 & & & \class{srk-wcolor}{0} & \class{srk-hcolor}{1} \\ \tfrac56 & \tfrac56 & & \class{srk-wcolor}{\tfrac56} & \class{srk-hcolor}{1} \\ \hline & 0.4 & 0.6 & \class{srk-wcolor}{1} & \class{srk-hcolor}{0} \\ & -0.6 & 0.6 & & \end{array}$$

ShARK — SDE, strong order 1.5, adaptive (row 3: error estimator)

  • Strong convergence — pathwise accuracy (order 1.5 via ShARK); SEEDS only achieves weak order 1
  • Markov for free — $\boldsymbol W, \boldsymbol H$ independent on non-overlapping intervals, no special decomposition needed
  • Trivial to sample — just two independent Gaussians extend the ODE Butcher tableau
Building Blocks

Exponential Integrators

The diffusion SDE has a semi-linear structure — a stiff linear term and a noise term that can both be solved in closed form, leaving only the neural network term for numerical integration.

Diffusion SDE

$$\underbrace{\class{lawson-stiff}{a(t)\,\boldsymbol X_t\,\mathrm dt}}_{\text{stiff linear}} + \underbrace{\class{lawson-nn}{b(t)\,\boldsymbol f_\theta(t,\boldsymbol X_t)\,\mathrm dt}}_{\text{neural network}} + \underbrace{\class{lawson-noise}{g(t)\,\mathrm d\boldsymbol W_t}}_{\text{noise}}$$

→
After change of variables

$$\class{lawson-nn}{\boldsymbol f_\theta(\varsigma,\,\Xi(\varsigma)\,\boldsymbol Y_\varsigma)\,\mathrm d\varsigma} + \mathrm d\boldsymbol W_\varsigma$$

  • Linear term solved exactly — $\Xi(t) = \exp\!\bigl(\int_0^t a(\tau)\,\mathrm d\tau\bigr)$ absorbs $a(t)\,\boldsymbol X_t$ via $\boldsymbol Y_t = \Xi^{-1}(t)\,\boldsymbol X_t$
  • Noise term solved exactly — time change $\varsigma_t$ converts $g(t)\,\mathrm d\boldsymbol W_t$ into standard Brownian motion $\mathrm d\boldsymbol W_\varsigma$
  • Only the NN term remains — any standard (S)RK scheme now applies to a clean, non-stiff equation
data pred. $\mathrm d\boldsymbol X_t = \bigl[\bigl(\tfrac{\dot\alpha_t}{\alpha_t} + \tfrac{g^2(t)}{\sigma_t^2}\bigr)\boldsymbol X_t - \tfrac{\alpha_t g^2(t)}{\sigma_t^2}\,\boldsymbol x_{0|t}(\boldsymbol X_t)\bigr]\,\mathrm dt + g(t)\,\mathrm d\boldsymbol W_t$
→
$\mathrm d\boldsymbol Y_\varsigma = \tfrac{\sigma_T}{\gamma_T}\,\boldsymbol x_{0|\varsigma}\!\bigl(\tfrac{\gamma_T\sigma_\varsigma}{\sigma_T\gamma_\varsigma}\boldsymbol Y_\varsigma\bigr)\,\mathrm d\varsigma + \tfrac{\sigma_T}{\gamma_T}\,\mathrm d\boldsymbol W_\varsigma$
noise pred. $\mathrm d\boldsymbol X_t = \bigl[\tfrac{\dot\alpha_t}{\alpha_t}\,\boldsymbol X_t + \tfrac{g^2(t)}{\sigma_t}\,\boldsymbol x_{T|t}(\boldsymbol X_t)\bigr]\,\mathrm dt + g(t)\,\mathrm d\boldsymbol W_t$
→
$\mathrm d\boldsymbol Y_\chi = 2\alpha_T\,\boldsymbol x_{T|\chi}\!\bigl(\tfrac{\alpha_\chi}{\alpha_T}\boldsymbol Y_\chi\bigr)\,\mathrm d\chi + \alpha_T\,\mathrm d\boldsymbol W_{\chi^2}$
where $g^2(t) = \dot\sigma_t^2 - 2\,\tfrac{\dot\alpha_t}{\alpha_t}\,\sigma_t^2$,   $\gamma_t = \tfrac{\alpha_t}{\sigma_t}$,   $\varsigma_t = \tfrac{\alpha_t^2}{\sigma_t^2}$,   $\chi_t = \tfrac{\sigma_t}{\alpha_t}$
building blocks

McCallum–Foster Coupling

For a differential equation $\tfrac{\mathrm d\boldsymbol x}{\mathrm dt} = \boldsymbol f_\theta(t, \boldsymbol x_t)$, $\boldsymbol\Phi$ is a numerical scheme such that $\boldsymbol x_{n+1} = \boldsymbol x_n + \boldsymbol\Phi_{t_n,t_{n+1}}(\boldsymbol x_n)$.

We use the McCallum–Foster reversible coupling  to convert arbitrary explicit RK schemes into reversible RK schemes:

$\boldsymbol\Phi$
 explicit RK scheme McCallum–Foster reversible scheme
$\boldsymbol\Upsilon$
  • inherits convergence order
  • non-zero region of stability
a family of diffusion solvers

Princeps $\boldsymbol\Psi$

Combine the exponential integrator with an explicit (S)RK scheme $\boldsymbol\Phi$ — the result is a large family of exponential (stochastic) Runge–Kutta solvers we call Princeps, $\boldsymbol\Psi$.

$\begin{array}{c|c}0&\\\hline&1\end{array}$
Euler
$\begin{array}{c|cc}0&\\\tfrac12&\tfrac12\\\hline&0&1\end{array}$
Midpoint
$\begin{array}{c|c|c|c}0&0&0&0\\\hline&1&1&0\end{array}$
Euler–Maruyama
$\boldsymbol\Phi$
 explicit (S)RK scheme Princeps exponential (S)RK
$\boldsymbol\Psi$
DDIM / DPM-Solver-1
DPM-Solver-2
SEEDS-1

Easy to build new solvers! E.g., adaptive step-sizing via Dormand–Prince 4(5) , or a state-of-the-art SDE solver via ShARK .

making this family reversible

Rex

When building Princeps $\boldsymbol\Psi$ we instead first convert $\boldsymbol \Phi$ into a stable reversible solver via the McCallum–Foster coupling with an auxiliary state $\hat{\boldsymbol Y}_t$ in the exponetially lifted domain, before building Rex in the $\boldsymbol X_t$ space.

$\boldsymbol X$ space
$\boldsymbol Y$ space
continuous
discretized
SDE
exponential SDE
reversible scheme
Rex
exponentially lift
apply McCallum–Foster coupled $\boldsymbol\Phi$
transform back to $\boldsymbol X$ space
image inversion

Round-Trip Latent MSE

Encode a Stable Diffusion v1.5 latent → decode → measure $\|\boldsymbol X_0-\hat{\boldsymbol X}_0\|^2$.  Lower ↔ closer to exact.

latent MSE 10⁰ 10⁻¹ 10⁻² 10⁻³ 10⁻⁴ 10⁻⁵ 10⁻⁶ 10⁻⁷ 10⁻⁸ 10⁻⁹ 10⁻¹⁰ 10⁻¹¹ 10 steps 20 steps 50 steps DDIM EDICT BDIA O-BELM Rex (Euler) 1.7×10⁻² 3.4×10⁻⁷ 2.3×10⁻⁸ 2.2×10⁻⁹ 8.9×10⁻¹⁰ ↗ unstable
unconditional generation

CelebA-HQ 256×256 · DDPM · 10 / 20 / 50 steps

steps10
FD FD∞ Prec. Recall Density Coverage
steps20
FD FD∞ Prec. Recall Density Coverage
steps50
FD FD∞ Prec. Recall Density Coverage
EDICT DDIM BDIA O-BELM Rex (Mid) Rex (RK4) Rex (E–M, SDE)

10⁴ samples · larger polygon = better. Rex (E–M) sweeps five of six metrics at 20 & 50 steps.

conditional generation & editing

Stable Diffusion v1.5 · COCO & pix2pix

text-to-image · 1000 COCO captions · 10/20/50 steps
CLIP 10 Image Reward 10 PickScore 10 CLIP 20 Image Reward 20 PickScore 20 CLIP 50 Image Reward 50 PickScore 50
DDIM EDICT BDIA O-BELM Rex Mid Rex RK4 Rex E–M Rex ShARK

3 metrics × 3 step counts, min‑max per metric. Rex SDE variants lead Image Reward & PickScore at every step budget.

image editing · pix2pix · 50+50 steps
Image Reward CLIP PickScore LPIPS
DDIM BDIA O-BELM Rex Euler Rex Dopri5

LPIPS inverted (larger = better). Rex (Dopri5) on the ring; EDICT failed entirely.

boltzmann sampling · what it looks like

From Noise to an Unbiased Sample

flow matching · tri-alanine noise
①flow matching → biased sample (mode A)
②reweight via exact log p(x) — needs reversibility
③unbiased sample — reaches mode B
density over reaction coord. biased q(x)
mode A mode B
reweight by w(x) = p(x)/q(x): the biased proposal q becomes the true target p — mode B appears.

Noise → biased sample → reweight via exact log p(x) → unbiased sample. Reversibility gives unbiased likelihoods.

boltzmann sampling

Impact of Reversibility on Resampling

𝐸-W₂ ↓ lower
Rex Dopri5 0.495 SBG (SMC) 0.598 DiT Dopri5 0.737 SBG (IS) 0.758 RegFlow 1.051 ECNF++ 2.206
𝕋-W₂ ↓ lower
DiT Dopri5 0.468 Rex Dopri5 0.497 SBG (IS) 0.502 SBG (SMC) 0.503 ECNF++ 0.962 RegFlow 1.612
ESS ↑ higher
DiT Dopri5 0.140 Rex Dopri5 0.104 SBG (IS) 0.052 RegFlow 0.029 ECNF++ 0.003 SBG (SMC) —

Same DiT, same pipeline — swap Dopri5 for Rex. Energy‑W₂ 0.737 → 0.495 (33% reduction).

summary

What Rex Gives You

Rex is a family of algebraically reversible diffusion solvers.

[ ⇆ ]
Reversible family
A whole family of solvers whose numerical scheme has an exact inverse.
[ ⊂ ]
Subsumes popular schemes
Reversible variants of DDIM, DPM-Solver, gDDIM & SEEDS-1.
[ dx/dt ]
ODE & SDE
Handles both the deterministic and stochastic settings.
[ Δt ]
Adaptive step-sizing
Compatible with adaptive solvers such as Dopri5.
[ O(hⁿ) ]
High-order convergence
Inherits the arbitrarily high order of the base RK scheme.
[ ≈ ]
Non-trivial stability
Stable reverse dynamics via the McCallum–Foster coupling.
future work

The Stability ↔ Reversibility Trade-off

  • EDICT / BDIA / O-BELM are algebraically reversible — but sit right on the axis with zero linear stability.
  • Rex keeps that perfect reversibility while recovering some limited stability from the underlying explicit scheme.
  • DDIM has good stability, but is not practically reversible.
  • Recent EES schemes (Shmelev & Salvi) are only nearly reversible, but far more stable — likely landing in the trade-off region between DDIM and Rex. A compelling option for image editing, where an approximate inverse can suffice.
stability → linear stability of the reverse dynamics reversibility → exactness of the inverse EDICT / BDIA / O-BELM reversible · zero stability Rex (ours) perfect reversibility · some stability EES nearly reversible · better stability DDIM good stability · not reversible
?
rex-solver.github.io
Poster Session July 8 · 10:30–12:15 Hall A · #2402
selected citations

References