Lipschitz-Guided Design of Interpolation Schedules in Generative Models

TL;DR

The paper tackles interpolation schedule design in unit-time stochastic interpolants for flow- and diffusion-based generative models, showing that scalar schedules are statistically equivalent under the KL divergence in path space after optimal diffusion-coefficient tuning. It then introduces a numerical criterion, averaged squared Lipschitzness $A_2$, and a transfer formula to switch schedules at inference without retraining, demonstrating substantial numerical gains. Analytically, Gaussian targets achieve exponential Lipschitz improvements and Gaussian mixtures show reduced mode collapse; these insights extend to high-dimensional invariant measures from stochastic PDEs like Allen–Cahn and Navier–Stokes, delivering robust improvements across resolutions. The work highlights the limits of scalar-schedule statistics and motivates exploring richer interpolation mechanisms to further enhance numerical stability and sampling efficiency in complex scientific applications.

Abstract

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 motivates focusing on numerical properties of the resulting drift fields rather than purely 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, the 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.

Figures (5)

• Figure 1: Comparison of different interpolation schedules $\beta_t$. Left: $M=5$. Right: $M=20$. We set $p=0.3$. For the dilated schedule, we take $\kappa = 1$.
• Figure 2: Left: $128\times 128$ Gaussian fields generated by using linear schedules with $20$ steps of the RK4 integrator. Middle: $128\times 128$ Gaussian fields generated by using the designed schedules with $20$ steps of the RK4 integrator. Right: linear and designed schedules.
• Figure 3: Energy spectra of Gaussian fields: comparison between truth, generated via designed schedules or standard linear schedules, with $20, 40$ or $80$ RK4 steps. The three figures correspond to different resolutions. Left: $32\times 32$; middle: $64\times 64$; right: $128\times 128$.
• Figure 4: Energy spectra of invariant distributions of stochastic Allen-Cahn: comparison between truth, generated via designed schedules or standard linear schedules, with $10, 20$ or $40$ RK4 steps. The three figures correspond to different resolutions. Left: $32$; middle: $64$; right: $128$.
• Figure 5: Left: generated $128\times 128$ sample using linear schedule and $10$ steps of RK4; middle: generated $128\times 128$ sample using designed schedule and $10$ steps of RK4; enstrophy spectra of samples using different schedules.

Theorems & Definitions (30)

• Definition 2.1
• Proposition 2.2
• Proposition 2.3
• Remark 2.4
• Proposition 2.5
• proof
• Remark 2.6
• Proposition 3.1
• proof
• Definition 3.2