Rex: A Family of Reversible Exponential (Stochastic) Runge-Kutta Solvers
Zander W. Blasingame, Chen Liu
ICML 2026
cs.LG, cs.AI, stat.ML
2025-02-13
Rex converts any explicit (S)RK scheme into a reversible diffusion ODE/SDE solver. SD v1.5 FP32 latent MSE hits 8.85e-10; editing LPIPS drops from DDIM's 0.214 to 0.107.
Diffusion models and continuous normalizing flows integrate noise into data. Editing a real image, backpropagating through the solver, and computing exact likelihoods for Boltzmann sampling all need the reverse trip. A standard stepper leaves truncation error on every step, so integrating forward then backward misses the start, and the miss grows with the number of steps.
Prior exact-inversion solvers stall on three fronts. EDICT was later shown to be order zero, with local truncation error O(h). BDIA and O-BELM have an empty linear stability region; O-BELM's reconstruction error gets worse as you take more steps. All of them cover only the probability-flow ODE. On the SDE side, earlier reversible schemes cache the whole Brownian path, so memory scales with steps and adaptive stepping is off the table.
Rex is a three-step recipe. The base stepper is interchangeable.
The SDE variant pins the forward and backward passes to the same Brownian path. The path is not stored. A splittable PRNG plus Kidger's Brownian Interval rebuilds increments and space-time Lévy areas from a single seed, which is what makes adaptive stepping possible. Image experiments use ζ=0.999; the tri-alanine Boltzmann run switches to 0.001, where linear stability matters more.
Each McCallum-Foster step calls the base scheme twice, Φh and Φ{-h}, and carries the extra state. In the ODE case Rex inherits the arbitrary convergence order of the underlying Runge-Kutta method and a non-empty linear stability region. Reversible Heun and leapfrog are stable only on the imaginary interval [-i, i], nowhere for test equations with Re(λ)<0. Rex inherits McCallum-Foster's non-empty region, still smaller than that of the original RK method.
Reconstruction is the hardest number in the paper. One hundred real images, Stable Diffusion v1.5 at 512×512, CFG=1.0, FP32 latent-space round-trip MSE, VAE error excluded:
| Solver | 10 steps | 20 steps | 50 steps |
| DDIM | 3.57×10⁻¹ | 9.89×10⁻² | 1.66×10⁻² |
| EDICT | 1.59×10⁻⁶ | 1.98×10⁻⁷ | 2.23×10⁻⁹ |
| O-BELM | 1.05×10⁻⁷ | 2.24×10⁻⁷ | 3.35×10⁻⁷ |
| Rex (Euler) | 3.77×10⁻⁹ | 1.98×10⁻⁹ | 8.85×10⁻¹⁰ |
Rex sits about three orders of magnitude below O-BELM on average, and about eight below non-reversible DDIM. O-BELM's error grows with step count, matching the empty stability region. Pixel-space MSE is dominated by the VAE and sits near 1.86×10⁻³ for every reversible solver.
Unconditional generation uses a pretrained DDPM on CelebA-HQ 256, 10⁴ samples, Fréchet distance on DINOv2 features. At 50 steps, Rex (Euler-Maruyama) scores FD 391.93 against DDIM 490.88 and O-BELM 476.29. At 10 steps O-BELM's 605.52 edges Rex (Midpoint) at 607.20. Higher-order Rex (RK4) loses to the first-order SDE variant in this low-step image setting, a known pattern in the diffusion literature. Adaptive high-order schemes pay off more in flow matching and scientific computing.
Conditional generation uses SD v1.5 on 1000 COCO captions. Image Reward at 10 steps: Rex (ShARK) 0.239, Rex (Euler-Maruyama) 0.222, DDIM 0.033, O-BELM 0.051. CLIP scores of the Rex variants stay within 0.23 of DDIM.
Image editing is a pix2pix round trip: invert under the source caption to t=0.6, then resample under the edit caption. Lower LPIPS means unedited regions stay closer. Rex (Dopri5) 0.107, Rex (Euler) 0.109, O-BELM 0.140, DDIM 0.214, BDIA 0.885. EDICT collapsed to an approximate identity map on this benchmark and was dropped from the table. Rex (Dopri5) is the first adaptive-step reversible solver the paper applies to diffusion editing.
Boltzmann sampling of tri-alanine uses 10⁴ proposals. On the same 3.1M-parameter DiT, swapping Dopri5 for Rex (Dopri5) cuts energy 2-Wasserstein from 0.737 to 0.495, the best in the table. ESS falls from 0.140 to 0.104, and dihedral T-W2 rises from 0.468 to 0.497. The energy histogram lines up; effective sample size does not rise with it.
For diffusion inversion, real-image editing, and flow models that need a valid change of variables, Rex is a drop-in reversible version of DDIM or DPM-Solver, not a new generative model. Exact SDE inversion previously meant storing the full noise path or giving up a closed form. Reconstructing Brownian motion from a seed is the route this paper actually takes.
The cost is explicit: an auxiliary state and two base-scheme calls per step. Among reversible solvers, image quality is often best or tied for best; against non-reversible DDIM the gain is incremental. Do not default to high order. Euler or Euler-Maruyama is the safer pick in few-step image sampling; save Dopri5 for scientific and likelihood settings. Hyperparameters were not tuned for the CelebA leaderboard, so the comparison with EDICT, BDIA, and O-BELM is not stacked against Rex.
Arbitrary-order convergence is proven for the ODE case. On the SDE side the theorem shows that Princeps inherits the strong order of the underlying stochastic RK method; Rex-SDE after the reversible wrap has no matching statement. The construction covers semi-linear additive-noise equations, not state-dependent diffusion.
ζ is sensitive. Image runs use 0.999, the Boltzmann run uses 0.001, and a small ζ lets ζ⁻¹ amplify floating-point error in the inverse update. The stability region is non-empty and still smaller than the original RK region. The appendix shows a failed ShARK trajectory at 5 steps; a higher-order stochastic scheme does not automatically invert well at tiny step counts.
Boltzmann sampling is reported only on tri-alanine, and ESS dropped. Image experiments stop at SD v1.5 and a CelebA DDPM, with no matched-NFE comparison against current flow-matching models. The unconditional table omits Rex (Euler), so it is not the same variant set as the reconstruction table.