📊 statistics

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

This paper introduces Rex, a family of reversible exponential (stochastic) Runge-Kutta solvers that overcome the exact inversion limitations of standard methods by converting explicit schemes into algebraically reversible ones, thereby enabling near-machine-precision reconstruction and improved performance in diffusion-based generative models.

Original authors: Zander W. Blasingame, Chen Liu

Published 2026-08-05

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Zander W. Blasingame, Chen Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake a perfect cake, but instead of a recipe, you have a magical, self-mixing bowl that slowly transforms a pile of flour and eggs into a delicious dessert. This is how modern AI creates images, music, or even protein structures today. It starts with pure chaos (random noise) and slowly, step-by-step, guides it into something structured and beautiful. This process is like a river flowing downstream. But what if you wanted to go upstream? What if you wanted to take a finished cake and perfectly reverse the process to see exactly what the original pile of flour looked like?

In the world of computer science, this "upstream" journey is called inversion. It's incredibly useful for editing images (changing a cat into a dog while keeping the background perfect) or for scientific simulations where you need to know the exact starting point of a system. However, there's a catch. The mathematical tools (solvers) we use to guide the AI usually accumulate tiny, invisible crumbs of error with every step they take. When you try to go backward, these crumbs don't disappear; they pile up, and you end up with a slightly different pile of flour than you started with. It's like trying to un-bake a cake and finding that a few crumbs of sugar are missing, leaving you with a slightly different recipe. For years, scientists have struggled to build a "perfect" backward path that leaves zero crumbs behind, especially when the process involves random noise.

This is where a new family of tools called Rex comes in. The researchers behind this work, Zander W. Blasingame and Chen Liu, have invented a clever new way to solve these equations that guarantees you can go forward and backward without losing a single crumb. They didn't just fix the math; they built a whole new set of "reversible" solvers that work for both smooth, predictable paths and chaotic, noisy ones. Their method, which they call Rex (Reversible Exponential), allows computers to travel back in time through an AI's generation process with near-perfect precision. This means we can now edit images with surgical accuracy, train AI models more efficiently, and even simulate complex molecules like tri-alanine with a level of trust that was previously impossible. It's like finally finding a time machine that doesn't just take you back, but ensures you arrive at the exact same spot you left, down to the last atom.

The Problem: The "Un-baking" Problem

To understand why Rex is a big deal, we have to look at how these AI models work. They use something called Neural Differential Equations. Think of these as a set of instructions that tell the AI how to change a state (like a pixel in an image) over time. The AI starts with random noise and integrates these instructions forward to create an image.

The problem arises when we want to go backward. Standard solvers are like a hiker walking through a forest. If they take a step forward, then immediately try to step back, they might not land exactly where they started because they misjudged the terrain or slipped on a rock. In math terms, this is called discretization error. When you try to reverse the process, these tiny errors accumulate. If you are just generating a picture, a tiny error might not matter. But if you are trying to edit a photo (removing a person but keeping the background) or calculating the exact probability of a molecule's shape, those errors ruin the result. You might end up with a blurry background or a molecule that doesn't physically make sense.

Previous attempts to fix this were like trying to walk backward in the dark. Some methods were unstable (they would crash if you took a big step), others were slow, and some only worked for smooth paths, not the noisy, random paths that many modern AI models use. One method, called EDICT, was reversible but produced low-quality results. Another, BDIA, was unstable and often failed to reconstruct the original image.

The Solution: The "Perfect Echo" Solver

The authors propose Rex, a family of solvers that are algebraically reversible. This is a fancy way of saying that the math is designed so that if you take a step forward and then immediately take the corresponding step backward, you land exactly on your starting point. There is no "drift."

How did they do it? They took inspiration from a method called Lawson methods, which are used to handle equations with exponential growth or decay. They combined this with a technique called the McCallum-Foster method, which is known for being stable and reversible.

Here is the magic trick:

The Setup: They take a standard, high-speed solver (like the ones used in popular tools like DDIM or DPM-Solver).
The Twist: They add a "shadow" state variable. Imagine you are walking forward, but you also keep a perfect, synchronized copy of your path in a parallel universe.
The Reversal: When they need to go backward, they don't just reverse the steps. They use the "shadow" copy to calculate the exact reverse step. Because of the specific way they constructed the math, the forward and backward steps cancel each other out perfectly.

The paper shows that this works for both ODEs (smooth, deterministic paths) and SDEs (paths with random noise). This is a major breakthrough because most previous reversible methods only worked for smooth paths. By handling the noise correctly, Rex opens the door for exact inversion in the most complex, realistic AI models.

What They Found: Precision and Power

The researchers didn't just write down the math; they tested it extensively.

Perfect Reconstruction: In their tests, when they used Rex to go forward and then backward, the error was so small it was essentially zero (near-machine precision). In contrast, other methods like BDIA or O-BELM showed visible errors that grew larger as they took more steps. For example, in a test with 50 steps, Rex's error was orders of magnitude smaller than its competitors.
Better Image Editing: When they used Rex to edit images (changing a prompt like "a man on a horse" to "a man on a dragon"), the results were much cleaner. The background stayed true to the original image, whereas other methods introduced weird artifacts or blurred the scene. Rex reduced the visual distortion by about half compared to the best existing reversible methods.
Scientific Sampling: They also tested Rex on a molecule called tri-alanine. In this field, scientists need to sample from a "Boltzmann distribution" (a way of finding the most likely shapes a molecule can take). Using Rex, they were able to sample these shapes more accurately than previous methods, improving the accuracy of the energy calculations.

What Rex is NOT

It is important to note what this paper does not claim.

It does not say that Rex is the only way to do this, but it suggests it is currently the most robust and accurate method for both ODEs and SDEs.
It does not claim that higher-order methods (like using a 4th-order solver) are always better. In fact, their experiments showed that for some tasks, a simpler, lower-order version of Rex (Euler) actually performed better than the more complex versions.
It does not solve the problem of "hallucinations" in AI (where the AI makes things up). It solves the problem of precision in the math that moves the AI from noise to data.

Why This Matters

The ability to perfectly reverse a process is like having a "Ctrl+Z" for the entire universe of AI generation.

For Artists: It means you can edit an AI-generated image with total confidence, knowing the parts you didn't touch will remain exactly as they were.
For Scientists: It allows for more accurate simulations of physical systems, like how proteins fold or how molecules interact, which is crucial for drug discovery.
For AI Researchers: It enables better training of models by allowing them to calculate exact probabilities, which helps in understanding how the AI "thinks."

The authors have made their code available, inviting others to use this "perfect echo" solver. While the math behind Rex is complex, the result is simple: a tool that lets us walk forward and backward through the AI's creative process without ever losing our way.

Technical Summary: Rex: A Family of Reversible Exponential (Stochastic) Runge-Kutta Solvers

Problem Statement

Deep generative models based on neural differential equations (ODEs and SDEs) have become state-of-the-art for tasks ranging from image generation to protein design. These models typically integrate from a prior distribution to a data distribution. In many critical applications—such as gradient-based training, image editing, and calculating exact likelihoods for Boltzmann sampling—it is essential to integrate in the inverse direction (from data back to the prior).

However, standard numerical solvers accumulate discretization errors that prevent exact inversion. While "analytic reversibility" (solving implicit equations via fixed-point iteration) exists, it is computationally expensive and approximate. Existing algebraically reversible solvers (those with closed-form inverse steps) suffer from significant drawbacks:

Poor Stability: Most previous reversible methods (e.g., Asynchronous Leapfrog, Reversible Heun) lack a non-zero region of linear stability, making them prone to divergence.
Low Order of Convergence: Many existing schemes are limited to low orders (e.g., zero or first order), leading to high truncation errors.
SDE Limitations: Existing exact inversion methods are largely restricted to ODEs (probability flow ODEs). Methods for SDEs often require storing the entire Brownian motion trajectory in memory, which is impractical for adaptive step sizes and high-dimensional problems.

Methodology

The authors propose Rex, a family of reversible exponential (stochastic) Runge-Kutta solvers designed for both diffusion ODEs and SDEs. The construction follows a three-step recipe:

Selection of Base Scheme ( $\Phi$ ): The user selects an explicit (stochastic) Runge-Kutta (SRK) scheme, such as Euler, Dormand-Prince, or ShARK.
Construction of Princeps ( $\Psi$ ): The authors apply Lawson methods to the semi-linear structure of diffusion equations. By using an integrating factor (exponential integrator) and a change of variables, the original SDE/ODE is transformed into a reparameterized form where the drift is simplified. An explicit RK scheme is applied to this transformed system to create the "Princeps" scheme. This step effectively handles the semi-linear drift and additive noise structure common in diffusion models.
Construction of Rex ( $\Upsilon$ ): The Princeps scheme is converted into an algebraically reversible solver using the McCallum-Foster method. This involves introducing an auxiliary state variable $\hat{x}_n$ and a coupling parameter $\zeta \in (0, 1]$ . The forward and backward steps are defined algebraically to ensure that applying the backward step immediately after the forward step returns the exact initial state (in exact arithmetic).

Key Technical Innovations:

Algebraic Reversibility for SDEs: Unlike prior SDE inversion methods that cache the entire trajectory, Rex uses a single seed for a splittable pseudo-random number generator (PRNG) to recalculate Brownian motion increments deterministically during the backward pass. This enables adaptive step sizes without storing the full path.
Exponential Integration: By leveraging the semi-linear structure of diffusion models, Rex incorporates exponential integrators, allowing it to generalize popular solvers like DDIM and DPM-Solvers into their reversible counterparts.
Stability Mechanism: The use of the McCallum-Foster coupling parameter $\zeta$ ensures that the resulting solver possesses a non-zero region of linear stability, a property previously absent in reversible solvers.

Key Contributions

Rex Solver Family: The construction of an algebraically reversible solver family for both diffusion ODEs and SDEs. The ODE variant inherits arbitrary-order convergence and a non-zero stability region from the underlying McCallum-Foster method.
Theoretical Generalization: The authors demonstrate that the intermediate "Princeps" scheme subsumes many popular existing solvers (DDIM, DPM-Solver-1/2/12, DPM-Solver++, SEEDS-1, gDDIM). Consequently, Rex provides the reversible versions of these widely used methods.
Convergence and Stability Proofs: The paper establishes that Rex inherits the convergence order of the base explicit scheme (Theorem 3.3) and proves that the ODE variant has a non-zero region of linear stability.
Exact SDE Inversion: The proposal of the first method for exact inversion of diffusion SDEs that does not require storing the entire Brownian motion trajectory, enabling practical adaptive step-size integration.

Empirical Results

The authors evaluate Rex across four major contexts:

Reconstruction Error: On Stable Diffusion v1.5, Rex (Euler) achieves near-machine-precision reconstruction (latent-space MSE) over 10, 20, and 50 steps, outperforming baselines like EDICT, BDIA, and O-BELM by orders of magnitude. Notably, O-BELM's error grows with step count due to instability, whereas Rex remains stable.
Unconditional Image Generation: On CelebA-HQ, Rex variants (Midpoint, RK4, Euler-Maruyama) outperform reversible baselines (EDICT, BDIA, O-BELM) and often surpass the non-reversible DDIM baseline in metrics including Fréchet Distance (FD), Precision, Recall, Density, and Coverage.
Conditional Image Generation & Editing: Using Stable Diffusion v1.5 on COCO captions and the pix2pix dataset, Rex variants achieve superior CLIP Scores, Image Rewards, and PickScores. In image editing (round-trip inversion), Rex (Dopri5) significantly reduces perceptual loss (LPIPS) compared to DDIM and O-BELM, preserving non-edited content better. BDIA fails catastrophically in this task due to instability.
Boltzmann Sampling: In sampling from the Boltzmann distribution of tri-alanine, applying Rex (Dopri5) to a Diffusion Transformer (DiT) improves the 2-Wasserstein distance of energy distributions compared to standard non-reversible baselines, enabling accurate likelihood-based sampling.

Significance and Claims

The paper claims that Rex addresses the critical trade-off between precision (exact inversion) and stability in neural differential equation solvers.

Precision: Rex enables "near-machine-precision" reconstruction, which is vital for applications requiring exact likelihoods or differentiable rewards through the generative process.
Stability: By inheriting the stability properties of the McCallum-Foster method, Rex overcomes the historical weakness of reversible solvers, allowing for larger step sizes and robust performance in SDE settings.
Versatility: The framework is general, supporting adaptive step sizes and both ODE and SDE formulations, covering a broad class of generative models including affine probability path flow matching.
Practicality: The ability to invert diffusion SDEs without memory-heavy trajectory storage makes the method suitable for real-world applications like image editing and scientific sampling (e.g., molecular conformations).

The authors position Rex not just as a theoretical improvement but as a robust numerical scheme ready for integration into existing AI4Science and generative modeling pipelines where bijectivity and precision are paramount.