Lipschitz-Guided Design of Interpolation Schedules in Generative Models
This work addresses numerical efficiency issues in generative model training, offering a method to improve sampling quality without retraining, which is incremental but practically useful for researchers and practitioners in machine learning.
The paper tackles the problem of designing interpolation schedules in generative models by focusing on numerical efficiency rather than statistical criteria, proposing Lipschitzness minimization as an optimization criterion. This approach achieves exponential improvements in Lipschitz constants for Gaussian distributions and reduces mode collapse in Gaussian mixtures, with validation on high-dimensional distributions from stochastic Allen-Cahn and Navier-Stokes equations.
We study the design of interpolation schedules in the stochastic interpolants framework for flow and diffusion-based generative models. We show that while all scalar interpolation schedules achieve identical statistical efficiency under Kullback-Leibler divergence in path space after optimal diffusion coefficient tuning, their numerical efficiency can differ substantially. This observation motivates focusing on numerical properties of the resulting drift fields rather than statistical criteria for schedule design. We propose averaged squared Lipschitzness minimization as a principled criterion for numerical optimization, providing an alternative to kinetic energy minimization used in optimal transport approaches. A transfer formula is derived that enables conversion between different schedules at inference time without retraining neural networks. For Gaussian distributions, our optimized schedules achieve exponential improvements in Lipschitz constants over standard linear schedules, while for Gaussian mixtures, they reduce mode collapse in few-step sampling. We also validate our approach on high-dimensional invariant distributions from stochastic Allen-Cahn equations and Navier-Stokes equations, demonstrating robust performance improvements across resolutions.