Fast Sampling for Flows and Diffusions with Lazy and Point Mass Stochastic Interpolants
This work addresses the computational inefficiency in generative modeling for AI applications, offering an incremental improvement by optimizing existing schedules for faster sampling.
The paper tackled the problem of slow sampling in flow and diffusion models by introducing a method to convert sample paths between different interpolation schedules and diffusion coefficients, enabling faster image generation without retraining. They demonstrated this by applying a lazy schedule conversion to a pretrained flow model, reducing the number of steps needed for image synthesis.
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.