Stochastic interpolants with data-dependent couplings
This work addresses the need for more flexible conditional generative modeling in tasks such as image restoration, though it is incremental as it builds on existing stochastic interpolant methods.
The paper tackles the problem of constructing conditional generative models by formalizing data-dependent couplings between base and target densities within the stochastic interpolants framework, enabling applications like super-resolution and in-painting through a simple square loss regression.
Generative models inspired by dynamical transport of measure -- such as flows and diffusions -- construct a continuous-time map between two probability densities. Conventionally, one of these is the target density, only accessible through samples, while the other is taken as a simple base density that is data-agnostic. In this work, using the framework of stochastic interpolants, we formalize how to \textit{couple} the base and the target densities, whereby samples from the base are computed conditionally given samples from the target in a way that is different from (but does preclude) incorporating information about class labels or continuous embeddings. This enables us to construct dynamical transport maps that serve as conditional generative models. We show that these transport maps can be learned by solving a simple square loss regression problem analogous to the standard independent setting. We demonstrate the usefulness of constructing dependent couplings in practice through experiments in super-resolution and in-painting.