Lipschitz-Guided Design of Interpolation Schedules in Generative Models

TMLR Paper9074 Authors

Abstract:

We study the design of interpolation schedules in flow and diffusion-based generative models from both statistical and numerical perspectives.
Within the stochastic interpolants framework, we first show that scalar interpolation schedules are statistically equivalent under the Kullback--Leibler divergence in path space, after optimal a posteriori tuning of the diffusion coefficient.
This equivalence motivates focusing on numerical properties of the drift field rather than purely statistical criteria.
We propose minimizing the averaged squared Lipschitzness of the drift as a principled criterion for schedule design, in contrast with kinetic-energy minimization in optimal transport.
A simple transfer formula expresses the drift of one schedule in terms of the drift of another, allowing the designed schedule to be used at inference time with a model trained under a different (e.g., linear) schedule, without retraining.
We work out the optimal schedules analytically for Gaussian and Gaussian-mixture targets: for Gaussians, we obtain exponential improvements in the Lipschitz constant over linear schedules; for Gaussian mixtures, we obtain schedules that mitigate mode collapse in few-step sampling.
We then validate the approach on high-dimensional invariant measures of stochastic Allen--Cahn and Navier--Stokes equations, where the designed schedule yields markedly more accurate fine-scale statistics at fixed integrator budget.

Submission Type: Long submission (more than 12 pages of main content)

Assigned Action Editor: ~Julius_Berner1

Submission Number: 9074