Disintegrated Stochastic Interpolants: Constrained Training for Flow Matching and Diffusion Models

Diffusion-based and flow-based models are effective across a wide range of generative tasks, from images and text to physical systems governed by partial differential equations. Both are formulated through ambient densities, and therefore assume that the data distribution has a strictly positive density everywhere. Scientific data are usually known to follow specific equality constraints. A single neural network must therefore learn both how to drive samples towards the manifold of the data support and how to model the nontrivial distribution within it. We show that these two tasks can be separated and doing so can give us significant performance gains. We introduce \emph{disintegrated stochastic interpolants}, which is an extension of stochastic interpolants that unifies both flow-based and diffusion-based models into a single theoretical framework. Within this framework, we first decompose the ambient space into level sets of a known constraint function; \emph{then the transport of measures can also be separated into transport within and between level sets}. Using this, we can derive novel training objectives for target measures that lie on some known submanifold of the Euclidean space. Across geometric and physics-based problems, disintegrated models learn the target distribution more accurately at matched training budgets while satisfying the prescribed constraints substantially more closely than ambient alternatives. Experiments span generation on the Stiefel manifold, MNIST inpainting, and sparse conditional generation of Burgers trajectories.