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

Summary

The paper introduces Rex, a family of algebraically reversible exponential (stochastic) Runge–Kutta solvers for diffusion ODEs and SDEs. The conceptual move is a three-step recipe: (i) rewrite the semi-linear diffusion (S)ODE by combining exponential integrators with a signal-to-noise-ratio time change (Proposition 3.1 / Proposition C.1), yielding a "Princeps" scheme $\Psi$ whose linear drift is absorbed into an integrating factor; (ii) apply an arbitrary explicit (S)RK $\Phi$ to that reparameterized equation; (iii) close the loop with the McCallum–Foster coupling to obtain an algebraically reversible solver $\Upsilon$. This is a genuinely useful synthesis. On the theoretical side, the paper inherits arbitrary-order convergence and a non-trivial linear stability region from McCallum–Foster (Theorem 3.3), extends the strong-convergence analysis to the SDE side via a Dambis–Dubins–Schwarz time change (Theorem 3.4, Lemma B.3), and shows that Princeps subsumes DDIM, the DPM-Solver family, SEEDS-1, and gDDIM as non-reversible special cases (Theorem 3.5), so Rex furnishes reversible versions of these solvers "for free." On the empirical side, Rex (Euler) attains near-machine-precision round-trip inversion on Stable Diffusion v1.5, Rex (Dopri5) roughly halves editing LPIPS versus the strongest reversible baseline, and Rex (Dopri5) improves the energy 2-Wasserstein on tri-alanine Boltzmann sampling. Executed replication confirms the reconstruction story that anchors the paper — the fp32 and fp16 reconstruction tables reproduce within an order of magnitude and reversibility holds in the operating regime — while showing that the strongest phrasings are somewhat optimistic and that roughly half the empirical claims could not be reproduced from the released artifact. The main conceptual limitations are that (a) the convergence theorems are proved only for variance-preserving schedules and hinge on assumptions ($\alpha_T > 0$, $b(t)/a(t) \to \infty$) whose relation to the finite $[0,T]$ setting is not spelled out; (b) the "exact" SDE inversion is coupled to an oracle reconstructing the Brownian realization, and the double role of $\zeta$ is under-analyzed; and (c) the arbitrary-order theory does not activate empirically — higher-order Rex is beaten by lower-order Rex on the image benchmarks.

Strengths

• Conceptual contribution. The paper isolates why previous reversible diffusion solvers were unsatisfying (EDICT / BDIA / O-BELM lack a non-trivial linear stability region, and BELM must store $k-1$ prior states) and fixes all three issues in one move by grafting the McCallum–Foster coupling onto a Lawson-transformed (S)RK backbone. The survey of prior reversible solvers in Appendix A and Table 2 is unusually thorough.
• Unifying view of existing diffusion solvers. Theorem 3.5 and Appendix E show the Princeps construction subsumes DDIM, DPM-Solver-1/2/12, DPM-Solver++1/(2S), SDE-DPM-Solver-1/++1, SEEDS-1, and gDDIM, turning each into a reversible variant with essentially no design effort.
• First reversible diffusion SDE solver without full-trajectory storage. Prior SDE inversion (CycleDiffusion) stored the entire Brownian trajectory; Rex (Euler-Maruyama / ShARK) sidesteps this via splittable-PRNG reconstruction, and Rex (ShARK) is a novel strong-order-1.5 reversible diffusion SDE solver.
• Near-machine-precision inversion, independently reproduced. The fp32 reconstruction table reproduces within an order of magnitude on every baseline cell, Rex sits at the finite-precision floor ($\sim 10^{-8}$), and O-BELM's error correctly grows with step count while Rex's does not — matching the stability argument.
• Clean coupling ablation. The three-ingredient decomposition (reversible coupling, exponential transform, time reparameterization) cleanly identifies the reversible coupling as the dominant contributor to inversion accuracy, and this is corroborated by the reconstruction sweep (removing the coupling inflates latent MSE by $\sim 6.5$–$6.9$ orders).

Weaknesses

• Convergence guarantees are narrower than the framing. The abstract and contribution bullets advertise "arbitrary order of convergence" without qualification, but Theorem 3.3 restricts this to variance-preserving schedules and Theorem 3.4 additionally requires $\alpha_T > 0$. The Conclusion's generalization to "any additive-noise SDE that can be written in the semi-linear form" is asserted without proof — and, pointedly, the headline Boltzmann experiment applies Rex (Dopri5) to an OT flow-matching CNF that is not variance-preserving, i.e. outside the proven regime.
• $\zeta$ plays a double role that is under-analyzed. Image tasks use $\zeta = 0.999$ and Boltzmann uses $\zeta = 0.001$, justified by "we do not need to invert the solver exactly, only to guarantee that the scheme is invertible." If reversibility is algebraically exact for all $\zeta \in (0,1)$, this distinction should not exist; if small $\zeta$ erodes exactness, the Boltzmann likelihoods use a subtly incorrect change of variables. No $\zeta$ sweep is reported, and replication found $\zeta = 0.5$ numerically unstable.
• Higher-order theory does not translate into higher-order performance. Rex (RK4) is repeatedly beaten by Rex (Euler) on CelebA-HQ, and the rationalization is a precedent, not a mechanism. No controlled convergence-order study confirms the predicted $O(h^k)$ slopes are attained at all.
• "Exact SDE inversion" depends on an external oracle. Exactness in the SDE case is a property of $\Upsilon$ composed with a splittable-PRNG Brownian reconstruction; the failure mode when that oracle is approximate, or when adaptive step control queries the path at different times on the two passes, is not stated. Relatedly, the adaptive-step reversibility claim never explains how the forward and backward passes are forced onto the same step grid.
• Empirical claims exceed the reported numbers. The editing "$2\times$ improvement over the strongest reversible baseline" is really $\sim 1.31\times$ against O-BELM ($0.140 \to 0.107$); the "$2\times$" matches the non-reversible DDIM. The "every Rex variant stays within a tenth of a point on CLIP of DDIM" claim is contradicted by several Table 6 entries with gaps of 0.16–0.23.
• No variance / seed information; opaque bolding rule. Tables 1, 4, 5, 6, 7 report a single point per cell, and the bolding rule is a range-normalized rank criterion, not a statistical test, so borderline margins cannot be interpreted.
• Thin evaluation splits. The editing benchmark is a single 100-triple subset of InstructPix2Pix with no stated selection protocol or error bars, and the Boltzmann evaluation uses one small peptide and one architecture (alanine dipeptide is mentioned but not reported).
• Derivational and interpretive slips. The noise-prediction proof writes the off-diagonal quadratic variation vanishes "for all $i = j$" (should be $i \neq j$), inverting the intended statement; the noise-prediction convergence constant is labelled $C_1 = \sigma_T C$ when it should involve $\alpha_T$; the $(\kappa_n, \varsigma_t)$ summary for the four cases uses an undefined $\gamma_n$; the SDE strong-order proof does not track the state-scaling Jacobian or the time-change factor; and Theorem 3.4's stability inheritance is argued for McCallum–Foster rather than for the reparameterized Rex. The $W_2$ metric is also mislabelled an integral probability metric, and Remark E.1's "wholly unique / no existing counterpart" overreaches given Appendix A.4's own SEEDS/SRKL discussion.

Reproducibility & code

• Reconstruction (fp32/fp16) reproduces; the tightest Rex numbers do not. Running the shipped reconstruction pipeline ($n=100$) reproduced the full fp32 table within an order of magnitude on every baseline cell, and the fp16 half matched. But the Rex (Euler) fp32 row reproduced at $9.26\times10^{-9}$ / $1.67\times10^{-8}$ / $7.87\times10^{-9}$ versus the paper's $3.77\times10^{-9}$ / $1.98\times10^{-9}$ / $8.85\times10^{-10}$ (roughly $2.5$–$9\times$ higher), the paper's smooth monotone Rex decrease was not reproduced, and at 50 steps EDICT ($2.21\times10^{-9}$) edged Rex — so "most accurate at every step count / orders of magnitude below all baselines" is stronger than the executed evidence.
• Algebraic reversibility holds only in the stable regime. With $\zeta = 0.999$ the executed round-trip returns the latent to fp32-$\epsilon$ relative error, but the released quick-start config ($\zeta = 0.5$, RK4) diverges (round-trip error $\sim 2.5\times10^{11}$, encode $|x_T| \sim 1.7\times10^{18}$), and the README snippet reuses the same time span for the decode until it is reversed. The literal "$\epsilon = 0$ regardless of step size or precision" does not hold for the shipped default.
• Boltzmann pipeline absent. The release contains no boltzmann-sampling module (README: "coming soon"), no tri-alanine loader, no DiT/Rex integration, and no ESS/E-W2/T-W2 evaluator; the checkpoint was private and the MD data unpinned. Table 1 could not be reproduced at all.
• DINOv2 / dgm-eval metrics not shipped. Tables 5–6 rely on DINOv2 FD/FD$_\infty$/Precision/Recall/Density/Coverage, but only Inception clean-FID ships; the conditional-generation metric aggregation had to be reimplemented during replication. Neither table could be reproduced.
• Code–paper mismatches. The editing evaluator loads ViT-B/16 while Appendix G.2.2 specifies ViT-L/14; the reconstruction driver defaults $\zeta = 0.5$ and $\epsilon = 0.02$ against the paper's $\zeta = 0.999$, $\epsilon = 0.0002$; the editing step count is 100+100 in Appendix G.3 but "50+50" in the Table 7 caption; and the Table 8 profiler harness is not shipped, so the negligible-overhead claim is unverified.
• What could not be attempted, and why. Image editing (Table 7), unconditional generation (Table 5), conditional generation (Table 6), Boltzmann sampling (Table 1), and timing (Table 8) produced no comparable numbers — partly a single-CPU-core environment that throttled the run, partly genuine code/data-availability gaps (Boltzmann code, dgm-eval suite, reference sets, and metric-aggregation scripts). One positive note: the reconstruction driver already exposes a --dtype flag, so the fp16 rows do not require source edits.

Recommended Changes

Essential

• Scope the convergence claims. Add the variance-preserving assumption (and $\alpha_T > 0$) to the abstract/contribution bullet on convergence order, and either prove the Conclusion's "any additive-noise SDE in semi-linear form" generalization or downgrade it — explicitly noting that the flow-matching Boltzmann experiment lies outside the proven regime. Maps to the convergence-scope and generalization weaknesses.
• Resolve the $\zeta$ tension and add a sweep. State whether $\zeta \in (0,1)$ preserves algebraic exactness (making $\zeta$ purely a stability knob) or trades off exactness (with consequences for the SNIS estimator), and report a $\zeta$ sweep on the reconstruction and Boltzmann tasks. Maps to the $\zeta$ double-role weakness.
• Fix the two overstated empirical sentences so the "$2\times$" editing figure names the correct baseline and the CLIP gap-to-DDIM matches Table 6. Maps to the "empirical claims exceed reported numbers" weakness.
• Report per-seed variance (mean $\pm$ std over $\geq 3$ seeds) for the headline cells of Tables 1, 4, 5, 6, 7, and give the explicit bolding-rule formula. Maps to the "no variance" weakness.
• Ship the Boltzmann code and data (DiT training, Rex/DiT sampling, SNIS reweighting, ESS/E-W2/T-W2 evaluators, and a pinned MD reference) or mark Table 1 as not reproducible from the release. Maps to the Boltzmann reproducibility bullet.
• Ship or pin the DINOv2/dgm-eval evaluators for Tables 5–6, including the metric-aggregation script and exact command lines. Maps to the DINOv2/dgm-eval bullet.
• Reconcile code with the manuscript: default $\zeta = 0.999$ and $\epsilon = 0.0002$ in the driver, load ViT-L/14 in the editing evaluator, choose one editing step count and update caption/appendix/script consistently, and fix the quick-start round-trip (reverse the decode span, use a stable $\zeta$). Maps to the code–paper-mismatch and reversibility replication bullets.
• Correct the load-bearing derivation slips: the "for all $i = j$" quadratic-variation condition, the $C_1 = \sigma_T C$ constant in the noise-prediction bound, and the undefined $\gamma_n$ in the $(\kappa_n, \varsigma_t)$ specification. Maps to the "derivational slips" weakness.

Suggested

• Add an empirical convergence-order study on a controlled ODE showing log-log slopes matching the predicted order for Euler / Midpoint / RK4 / Dopri5, and discuss mechanically why RK4 underperforms Euler on the image benchmark. Maps to the "higher-order theory does not translate" weakness.
• Specify the SDE-inversion oracle and adaptive-step reversibility. State which Brownian mechanism delivers exactness and its failure mode, and explain how the forward/backward step grids are kept identical under adaptive control. Maps to the "exact SDE inversion depends on an oracle" weakness.
• Supply the missing SDE-side bound. Add the explicit inequality tracking the state-scaling Jacobian and time-change factor in Theorem 3.4, and argue that a non-trivial stability region survives the reparameterization (not just for McCallum–Foster). Maps to the strong-order-proof and stability weaknesses.
• Broaden the evaluation. State the 100-triple editing subset-selection procedure and add bootstrap intervals; add a second Boltzmann system (e.g. alanine dipeptide). Maps to the "thin evaluation splits" weakness.
• Add a $(\kappa_n, \varsigma_t)$ summary table for the four data/noise $\times$ ODE/SDE cases. Maps to the kappa/varsigma weakness.
• Ship the Table 8 profiler harness and define the VRAM column convention. Maps to the profiler reproducibility bullet.
• Report BDIA $\gamma$ per table row and hedge the Remark E.1 novelty claim relative to SEEDS/SRKL. Maps to the BDIA-$\gamma$ and "wholly unique" points.