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
把任意显式(随机)Runge-Kutta经Lawson指数变换与McCallum-Foster耦合,做成同时覆盖扩散ODE与SDE的代数可逆求解器。Stable Diffusion v1.5上FP32潜空间往返MSE低至8.85×10⁻¹⁰,图像编辑LPIPS从DDIM的0.214降到0.107。
扩散模型和连续归一化流把噪声积到数据。真实图像编辑、沿求解器反传梯度、用精确似然做 Boltzmann 抽样,都还要把数据积回去。常规步进每步留下离散化误差,正向再反向走不回原点,步数越多漂移越大。
已有「精确反演」方法卡在三件事上。EDICT 后来被证明是零阶方法,局部截断误差是 O(h)。BDIA 和 O-BELM 在线性测试方程上稳定域是空集;O-BELM 的重建误差甚至随步数变差。这些方法都只覆盖概率流 ODE。SDE 侧先前所谓可逆,是把整条布朗轨迹存下来,内存随步数线性涨,自适应步长也用不了。
Rex 的配方分三步,底层求解器可以换。
SDE 版本把正向和反向绑在同一条布朗路径上。路径不存盘,用可分裂伪随机数发生器加 Brownian Interval,从一个种子按需重建增量和空间-时间 Lévy 面积。这样才谈得上自适应步长。图像实验默认 ζ=0.999;三丙氨酸 Boltzmann 抽样改成 0.001,那边更在意线性稳定。
McCallum-Foster 每步要调两次底层格式,一次 Φh、一次 Φ{-h},状态也多一份。换来的是:ODE 情形继承底层 Runge-Kutta 的任意阶收敛,并有一块非空线性稳定域。可逆 Heun 和 leapfrog 的稳定域是虚轴上的 [-i, i],对 Re(λ)<0 的测试方程无处稳定。Rex 继承 McCallum-Foster 那块非空域,面积仍小于原 RK。
重建是这篇最硬的数字。100 张真实图,Stable Diffusion v1.5,512×512,CFG=1.0,测 FP32 潜空间往返均方误差,把 VAE 误差排除在外:
| 求解器 | 10 步 | 20 步 | 50 步 |
| 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 平均比 O-BELM 低约三个数量级,比不可逆的 DDIM 低约八个。O-BELM 步数加、误差涨,和「无处线性稳定」对得上。像素空间误差被 VAE 盖住,可逆求解器之间几乎一样,都在 1.86×10⁻³ 附近。
无条件生成用 CelebA-HQ 256 的预训练 DDPM,抽 10⁴ 张,DINOv2 特征上算 Fréchet 距离。50 步时 Rex (Euler-Maruyama) 的 FD 是 391.93,DDIM 是 490.88,O-BELM 是 476.29。10 步时 O-BELM 的 605.52 略好于 Rex (Midpoint) 的 607.20。高阶 Rex (RK4) 在这个低步数图像设定里不如一阶 SDE 变体,和扩散社区的老观察一致;流匹配和科学计算里,自适应高阶才更稳。
条件生成用 SD v1.5 和 1000 条 COCO 标题。10 步 Image Reward:Rex (ShARK) 0.239,Rex (Euler-Maruyama) 0.222,DDIM 0.033,O-BELM 0.051。各 Rex 变体的 CLIP 相对 DDIM 差距不超过 0.23。
图像编辑走 pix2pix 往返:源字幕反演到 t=0.6,再按编辑字幕采样回来。LPIPS 越低,未改区域保得越好。Rex (Dopri5) 0.107,Rex (Euler) 0.109,O-BELM 0.140,DDIM 0.214,BDIA 0.885。EDICT 在这个设定塌成近似恒等映射,表里直接拿掉。Rex (Dopri5) 是这篇给出的第一个用在扩散编辑上的自适应步长可逆求解器。
三丙氨酸 Boltzmann 抽样,10⁴ 个提案。同一套 3.1M 参数 DiT,Dopri5 换成 Rex (Dopri5) 之后,能量 2-Wasserstein 从 0.737 降到 0.495,表内最好;有效样本量 ESS 从 0.140 降到 0.104,二面角 T-W2 从 0.468 升到 0.497。能量分布对上了,有效样本量没有一起涨。
做扩散反演、真实图编辑、需要精确变量替换的流模型,这是可以把 DDIM 或 DPM-Solver 直接换成可逆版本的求解器,不是新的生成模型。SDE 精确反演以前要么存整条噪声,要么没有闭式。Rex 用种子重建布朗路径,论文把它写成目前不存整条轨迹的做法。
代价清楚:多一份辅助状态,每步两次底层格式调用。图像生成质量在可逆求解器里经常最好或并列最好,相对不可逆 DDIM 是渐进改进。高阶格式别默认上。图像低步数用 Euler 或 Euler-Maruyama 更稳,科学计算和似然场景再用 Dopri5。超参没为 CelebA 榜单单独调过,和 EDICT、BDIA、O-BELM 的对比对 Rex 并不吃亏。
任意阶收敛写在 ODE 上。SDE 侧定理证明的是 Princeps 继承底层随机 RK 的强收敛阶,可逆封装后的 Rex-SDE 没有同等陈述。适用范围是能写成半线性加性噪声的方程,状态依赖的扩散项不在里面。
ζ 很敏感。图像实验默认 0.999,Boltzmann 改成 0.001;ζ 太小会让反向里的 ζ⁻¹ 放大浮点误差。稳定域非空,仍小于原 RK。附录给出 ShARK 在 5 步上的失败轨迹,随机高阶格式在极少步时并不自动好用。
Boltzmann 只报了三丙氨酸,ESS 还降了。图像实验停在 SD v1.5 和 CelebA DDPM,没有和现在的流匹配模型对同等函数评价次数。无条件生成表里没放 Rex (Euler),和重建表不是同一组变体,跨表对比要小心。