Stochastic interpolants with data-dependent couplings
Michael S. Albergo, Mark Goldstein, Nicholas M. Boffi, Rajesh, Ranganath, Eric Vanden-Eijnden
TL;DR
This paper introduces a new framework for constructing data-dependent couplings in stochastic interpolants, enabling conditional generative models that improve tasks like super-resolution and in-painting.
Contribution
It formalizes how to couple base and target densities conditionally, allowing the creation of data-dependent transport maps for generative modeling.
Findings
Transport maps can be learned via simple regression.
Dependent couplings improve super-resolution performance.
The method enables conditional in-painting tasks.
Abstract
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…
Peer Reviews
Decision·ICML 2024 Spotlight
- The paper formalizes two important notions in generative modeling, conditional and data dependent coupling, in the stochastic interpolates framework. - The authors show how to construct both conditional and data-dependent coupling.
1. **Limited contribution** - the work does not introduce a new concept and is a formulation of existing concepts into an existing framework. 1. The derivation of the transport equations in section 2.1, which takes a great portion of the paper, was already done in section 4 of [3] for the unconditional case, where the addition of the conditioning repeats the same derivation with marginalization over the condition. Furthermore, conditioning for super-resolution has been shown in [5] as well
* The paper is theoretically sound, as the derivations follow directly from the continuity equation. * Additionally, some experiments show the method's viability for common image generation tasks.
* I'm not sure the method is that original in practice. In particular, the paper notes that much of the construction can be connected with existing SDE and ODE formulations, all of which depend on the score function ([1] for the straight path ODE that is described in the paper, otherwise the standard OU process). In that case, the conditional methodology would follow from the score function argument as well, implying there would be little difference on the empirical side with existing methodolog
The idea is very neat and the theory seems well-executed. Qualitatively the experiments look really nice. Furthermore, the approach presents a nice unifying framework for many papers attempting to handcraft the couplings.
1) The biggest glaring weakness is the lack of quantitative experiments. While I am a huge fan of the idea and the developed theory, I think quantitative experimental evaluation is necessary for acceptance. An appropriate baseline could be the OT coupling flow from Tong et al. 2) In the proof of Theorem 1, I do not fully understand one step. In equation (23) for the second equality apparently the definition of conditional expectation is used. Please clarify this via some additional justificat
Code & Models
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsGenerative Adversarial Networks and Image Synthesis · Cell Image Analysis Techniques · Computer Graphics and Visualization Techniques
MethodsBalanced Selection