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

TL;DR

This paper introduces a method to convert stochastic interpolant sample paths between different schedules and diffusion coefficients, enabling more efficient image generation with fewer steps in flow models.

Contribution

It provides a theoretical framework for schedule conversion in stochastic interpolants and extends the framework to include point mass schedules, improving sampling efficiency.

Findings

Converted sample paths between schedules and diffusion coefficients.

Identified lazy schedules that make the drift zero, enabling variance-preserving sampling.

Applied schedule conversion to a pretrained flow model to reduce image generation steps.

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.