BézierFlow: Learning Bézier Stochastic Interpolant Schedulers for Few-Step Generation
This work addresses the challenge of efficient sampling in generative models for AI practitioners, offering a lightweight training approach that is incremental over prior timestep-learning methods.
The paper tackles the problem of improving few-step generation with pretrained diffusion and flow models by learning optimal stochastic interpolant schedulers, achieving a 2-3x performance improvement for sampling with ≤10 NFEs while requiring only 15 minutes of training.
We introduce BézierFlow, a lightweight training approach for few-step generation with pretrained diffusion and flow models. BézierFlow achieves a 2-3x performance improvement for sampling with $\leq$ 10 NFEs while requiring only 15 minutes of training. Recent lightweight training approaches have shown promise by learning optimal timesteps, but their scope remains restricted to ODE discretizations. To broaden this scope, we propose learning the optimal transformation of the sampling trajectory by parameterizing stochastic interpolant (SI) schedulers. The main challenge lies in designing a parameterization that satisfies critical desiderata, including boundary conditions, differentiability, and monotonicity of the SNR. To effectively meet these requirements, we represent scheduler functions as Bézier functions, where control points naturally enforce these properties. This reduces the problem to learning an ordered set of points in the time range, while the interpretation of the points changes from ODE timesteps to Bézier control points. Across a range of pretrained diffusion and flow models, BézierFlow consistently outperforms prior timestep-learning methods, demonstrating the effectiveness of expanding the search space from discrete timesteps to Bézier-based trajectory transformations.