Fast Sampling for Flows and Diffusions with Lazy and Point Mass Stochastic Interpolants

TL;DR

This work develops a unified stochastic interpolant framework that bridges flows and diffusions and introduces pathwise conversion results between SDEs with arbitrary schedules and diffusion scales. It extends to point-mass schedules and defines lazy schedules under a Gaussian-data assumption, revealing variance-preserving ODE schedules and point-mass SDE schedules, along with simple space reparameterizations to adapt pretrained flow models. Theoretical results hinge on the optimal diffusion scale $\varepsilon^*$ and the ability to reparameterize to lazy or SDE schedules, supported by intra- and inter-interpolant conversion formulas. Empirically, converting a state-of-the-art flow model to lazy schedules enables generating high-quality images with fewer steps, especially for SDE sampling, demonstrating practical gains in sampling efficiency without retraining.

Abstract

Stochastic interpolants unify flows and diffusions, popular generative modeling frameworks. A primary hyperparameter in these methods is the interpolation schedule that determines how to bridge a standard Gaussian base measure to an arbitrary target measure. We prove how to convert a sample path of a stochastic differential equation (SDE) with arbitrary diffusion coefficient under any schedule into the unique sample path under another arbitrary schedule and diffusion coefficient. We then extend the stochastic interpolant framework to admit a larger class of point mass schedules in which the Gaussian base measure collapses to a point mass measure. Under the assumption of Gaussian data, we identify lazy schedule families that make the drift identically zero and show that with deterministic sampling one gets a variance-preserving schedule commonly used in diffusion models, whereas with statistically optimal SDE sampling one gets our point mass schedule. Finally, to demonstrate the usefulness of our theoretical results on realistic highly non-Gaussian data, we apply our lazy schedule conversion to a state-of-the-art pretrained flow model and show that this allows for generating images in fewer steps without retraining the model.

Figures (9)

• Figure 1: ODE ($\varepsilon \equiv 0$) and statistically optimal SDE ($\varepsilon = \varepsilon^*$) sample path under three different interpolation schedules for $\rho_X$ a Gaussian mixture density for which the dynamics are analytically known stochastic_interpolants. All paths start in the same initial condition which sits in the initial position (cyan circle) for the leftmost two interpolants with density-admitting schedule and in the initial drift for the rightmost point mass schedule. The same Wiener process realization is shared between all three SDE solutions. Leftmost is the linear schedule from \ref{['def:linear_interpolant']}, middle and rightmost is the lazy interpolant for the ODE and statistically optimal SDE, respectively, as in \ref{['ex:u_t_eq_t']}. Since $u_t = t$ for all three schedules, all sample paths are equivalent up to a reparameterization of space. Note that all sample paths end in the same position for the ODE and SDE, respectively. See \ref{['fig:schedules']} for a plot of the three schedules and their statistically optimal diffusion scales.
• Figure 2: Sample images from the PRX flow model with a varying number of predictor-corrector steps for the ODE and statistically optimal SDE, respectively, sampled using the original linear flow model schedule versus converting to the lazy schedule. The text prompt "A car is stopped at a red light" was used. As expected, the image quality improves as the number of solver steps increases and seems to converge to the "ground truth" reference image which depends on whether an ODE or SDE is used but less so on which schedule is used. See \zenodourl for an animation.
• Figure 3: Left: Pixel-wise RMSE predictor-corrector convergence plot for ODE (circle markers and solid curves) and statistically optimal SDE (square markers and dashed curves) generation with the original linear flow model schedule versus the converted lazy schedule as in \ref{['fig:convergence_samples']}. Average and 95% confidence interval (shaded bands) is over $100$ prompts, initial conditions and Wiener process realizations. One sees that convergence is (almost) monotone as a function of the number of solver steps for all four configurations. Compare with \ref{['fig:convergence_plot_within_step']}. Right: Visual answer to the question "Using $n_i$ numerical solver steps with the lazy schedule, how many solver steps would I on average need to use with the linear flow model schedule to achieve the same RMSE to the reference image?". The answer for each $n_i$ on the $x$-axis is plotted on the $y$-axis (cyan labels indicate nearest rounded integer). Shaded bands indicate $95\%$ confidence intervals. The lazy schedule almost always performs statistically significantly better (and never worse) for both ODE and SDE generation, but the improvement is much bigger for SDE generation. All confidence intervals are calculated by bootstrapping with $10000.0$ samples.
• Figure 4: Schedules $(\alpha,\beta)$ (top row) and corresponding statistically optimal diffusion scale $\varepsilon^*$ (bottom row). These are the schedules and diffusion scales used in \ref{['fig:typst_interpolant_paths']}.
• Figure 5: Sample images from the PRX flow model with varying number of steps using different integration schemes for the ODE with original linear flow model schedule. A guidance strength of $5$ with the text prompt "Two husky's hanging out of the car windows" was used.

Theorems & Definitions (57)

• Definition 1.1: stochastic interpolant
• Definition 1.2: schedule
• Definition 1.3
• Theorem 1.4: SDE
• Definition 4.1
• Proposition 4.1: intra-interpolant conversion formulas
• Theorem 4.2
• Definition 5.1: linear interpolant
• Proposition 5.1
• Definition 5.1