Disintegrated Stochastic Interpolants
Abstract
Diffusion 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. In scientific applications, known equality constraints can confine the target distribution to a lower-dimensional set, making it singular with respect to the ambient Lebesgue measure. An ambient model must learn both how to drive samples onto this set and how to recover the nontrivial distribution along it. We introduce disintegrated stochastic interpolants, which prescribe the collapse onto the constraint set and learn the remaining distributional transport. Our construction decomposes the state space into level sets of a known constraint map and separates transport between these sets from transport within them. The constraint values follow an analytically prescribed evolution towards zero, while the model learns only transport within the moving level sets. We establish conditions under which these components recover the target distribution and derive corresponding training objectives and sampling algorithms for flow and diffusion models. Beyond singular targets, the proposed framework can also exploit conservation laws and symmetries. For symmetric distributions, we recover the quotient-space diffusion loss xu2026quotient as a special case of our training objective. Across geometric and physics-based problems, our disintegrated models learn target distributions more accurately at matched training budgets while satisfying the prescribed constraints substantially more closely than ambient alternatives.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.