Beyond the PF-ODE: Learn the Reverse Dynamics for Few-Step Diffusion Sampling

Diffusion solvers accelerate sampling by numerically approximating the reverse process with only a few function evaluations. This reverse process is commonly formulated as either a deterministic probability-flow ODE (PF-ODE) or a stochastic reverse SDE. Reverse SDEs can dissipate distribution mismatch accumulated during sampling, but stronger stochasticity also amplifies the effect of score error. PF-ODEs avoid this additional stochasticity, yet preserve inherited distribution mismatch under exact scores. We show that this creates a step-dependent trade-off between distribution correction and score reliability, suggesting that the preferable reverse dynamics need not remain fixed throughout sampling. Motivated by this observation, we propose RD-Solver, which jointly learns the reverse dynamics and their discretization. Specifically, RD-Solver learns an interval-wise stochasticity coefficient together with multistep predictor–corrector weights and calibrated score-evaluation times. We optimize these parameters with a frozen network using an endpoint loss that upper-bounds the squared 2-Wasserstein distance to a high-accuracy teacher in the supervision space. Experiments across diffusion and flow models demonstrate improved few-step generation, achieving FID scores of 2.71 on CIFAR-10 and 4.62 on ImageNet-64 with five function evaluations (5 NFEs).