Stochastic Interpolants A Unifying Framework for Flows and Diffusions

Abstract

A class of generative models that unifies flow-based and diffusion-based methods is introduced. These models extend the framework proposed in Albergo & Vanden-Eijnden (2023), enabling the use of a broad class of continuous-time stochastic processes called ‘stochastic interpolants’ to bridge any two arbitrary probability density functions exactly in finite time. These interpolants are built by combining data from the two prescribed densities with an additional latent variable that shapes the bridge in a flexible way.

Summary

The stochastic interpolant framework provides a mathematically rigorous unification of flow-based and diffusion-based generative models. A stochastic interpolant is defined as $x_t=I(t,x_0,x_1)+\gamma (t)z$ where $I$ interpolates between endpoints ($I(0,x_0,x_1)=x_0$, $I(1,x_0,x_1)=x_1$), $\gamma (t)$ controls latent noise ($\gamma (0)=\gamma (1)=0$, $\gamma (t)>0$ for $t\in (0,1)$), $(x_0,x_1)\sim \nu$ with marginals ${\rho }_0,{\rho }_1$, and $z\sim \mathcal{N}(0,I_d)$ independent of the data. The density $\rho (t)$ of $x_t$ is proved to be absolutely continuous and strictly positive, satisfying both a transport equation ${\partial }_t\rho +\nabla \cdot (b\rho )=0$ and forward/backward Fokker-Planck equations with tunable diffusion coefficient $ϵ(t)$.

The velocity $b(t,x)=\mathbb{E}[{\dot{x}}_t|x_t=x]$ and score $s(t,x)=-{\gamma }^{-1}(t)\mathbb{E}[z|x_t=x]$ are characterized as unique minimizers of simple quadratic objectives (Theorems 7, 8). This yields three equivalent generative models: a probability flow ODE ${\dot{X}}_t=b(t,X_t)$, a forward SDE $d{X}_t^F=b_{F}dt+\sqrt{2ϵ}d{W}_t$ with $b_F=b+ϵs$, and a backward SDE. Crucially, the density $\rho (t)$ is the same for all three — the interpolant construction is decoupled from the sampling process.

The framework’s key advantages over your random bridge approach: (1) the latent variable $\gamma (t)z$ smooths the intermediate density, eliminating spurious modes that appear in pure linear interpolation; (2) the coupling $\nu$ between $(x_0,x_1)$ can be arbitrary (independent, OT, or data-adapted); (3) the Schrödinger bridge is recovered as the solution to a max-min problem over interpolants (Theorem 41, Section 3.4).

Key Contributions

• General stochastic interpolant $x_t=I(t,x_0,x_1)+\gamma (t)z$ with rigorous density theory (Theorem 6)
• Velocity and score as unique minimizers of quadratic objectives (Theorems 7, 8)
• Simultaneous derivation of ODE and SDE generative models with tunable $ϵ(t)$
• Likelihood control: SDE models bound KL divergence via velocity+score loss (Theorem 23); ODE models additionally require Fisher divergence control
• Recovery of score-based diffusion as one-sided interpolant with time reparameterization (Section 5.1)
• Recovery of rectified flow as $\gamma =0$ linear interpolant (Section 5.3)
• Recovery of Schrödinger bridge via optimization over interpolant function $I$ (Section 3.4, Theorem 41)
• Remark 48 explicitly states: “straight line solutions is a necessary condition for optimal transport, but it is not sufficient” — confirming the rectified-flow-doesn’t-solve-OT result

Methodology

The interpolant $x_t=I(t,x_0,x_1)+\gamma (t)z$ bridges ${\rho }_0$ and ${\rho }_1$ by construction. The velocity $b$ is learned by minimizing ${\mathcal{L}}_b={\int }_0^1\mathbb{E}[\frac{1}{2}|\hat{b}{|}^2-({\partial }_{t}I+\dot{\gamma }z)\cdot \hat{b}]dt$. The score is learned via ${\mathcal{L}}_s={\int }_0^1\mathbb{E}[\frac{1}{2}|\hat{s}{|}^2+{\gamma }^{-1}z\cdot \hat{s}]dt$. For the spatially linear case $x_t^{\text{lin}}=\alpha (t)x_0+\beta (t)x_1+\gamma (t)z$, the velocity factorizes as $b=\dot{\alpha }{\eta }_0+\dot{\beta }{\eta }_1+\dot{\gamma }{\eta }_z$ where ${\eta }_0=\mathbb{E}[x_0|x_t]$, ${\eta }_1=\mathbb{E}[x_1|x_t]$, ${\eta }_z=\mathbb{E}[z|x_t]$.

Key Findings

• The latent variable $\gamma (t)z$ smooths intermediate densities, suppressing spurious modes (Figure 4)
• The diffusion coefficient $ϵ(t)$ affects sample trajectories but NOT the density $\rho (t)$ — a decoupling not present in standard diffusion models
• For ODE-based generative models: learning $b$ alone is insufficient for likelihood control; one must also control the Fisher divergence (unlike SDE models)
• Rectification (Section 5.3) produces straight-line flows but does NOT guarantee optimal transport (Remark 48): “straight line solutions is a necessary condition for optimal transport, but it is not sufficient”
• Optimizing over gradient velocity fields and iterating rectification converges to Brenier’s polar decomposition (Remark 49, citing Liu 2022)

Important References

1. Flow Matching for Generative Modeling — Concurrent work arriving at similar conditional objectives via a different derivation
2. Score-Based Generative Modeling through Stochastic Differential Equations — Score-based diffusion recovered as one-sided interpolant
3. Flow Straight and Fast — Rectified flow recovered as γ=0 linear interpolant; rectification analyzed in Section 5.3

Atomic Notes

• stochastic interpolant
• conditional flow matching

---

paper