Stochastic interpolants with data-dependent couplings

We thank the reviewer for their input and address their main concerns one-by-one below. We also refer them to our general reply for some more information.

Lack of quantitative experiments: We thank the reviewer for raising this point. We are currently running additional quantitative tests by comparing FID scores for our approach to existing baselines in the literature on super-resolution and in-painting. These results are shown in the table in our summary reply, where we achieve improved scores relative to earlier methods. We are continuing to train our models and expect the metrics to further improve for when we post our revision.

Proof of Theorem 1: The penultimate step in Eq. (23) uses the the tower property of the conditional expectation. That is, $\mathbb{E}[(\dot\alpha x_0+ \dot\beta x_1 + \dot\gamma z)\cdot \nabla \phi(x_t)] = \mathbb{E}[\mathbb{E}[(\dot\alpha x_0+ \dot\beta x_1 + \dot\gamma z)\cdot \nabla \phi(x_t)|x_t]].$ The last step uses the fact that the conditioning of $\phi(x_t)$ on $x_t$ is trivial, i.e.

= \mathbb{E}[\mathbb{E}[(\dot\alpha x_0+ \dot\beta x_1 + \dot\gamma z)|x_t]\cdot \nabla \phi(x_t)]$$ ***Role of $\alpha(t)$, $\beta(t)$, and $\gamma(t)$:*** In the present paper, these functions of time are prescribed and are not optimized upon. The optimization is only performed over the parameters in the velocity field. ***Invertibility:*** The flow from base to target that we construct does not need to be invertible. In particular, it can map samples from a higher-dimensional manifold to samples on a lower-dimensonal one (though this comes at the expense of requiring a velocity that is singular at $t=1$). In contrast, traversing from a lower-dimensional manifold to a higher-dimensional one is more problematic, since it requires the velocity to be singular at $t=0$, and therefore there is ambiguity on how to start the flow. To avoid this latter scenario, we add a bit of noise to the sample from the base density that we use. ***Interpretation of the loss:*** One loss gives the conditional velocity appearing in the transport equation. The second loss gives the denoiser $\mathbb{E}[z|x_t=x]$, which allows us to estimate the score since $\nabla \log \rho(t,x) = - \gamma^{-1}(t) \mathbb{E}[z|x_t=x]$. That is, the loss for the denoiser $\mathbb{E}[z|x_t=x]$ can be viewed as a reweighted version of the loss for the score $\nabla \log \rho(t,x)$. This trick is exploited in diffusion models when learning a "noise model". **Reversed-time vs backward SDE:** We agree that the first denomination is better and we have adopted it.

We thank the reviewer for their input and we address their main concerns one-by-one below. We also refer them to our general reply for further information.

Data-dependent coupling vs. class-conditioning: We stress that there is a distinct difference between putting conditional information in the velocity/score field and learning a map between coupled densities. For a visual explanation, see the figure posted here. The point of our paper is that data-dependent couplings can offer advantages that are orthogonal to (but can be combined with) conditional information placed in the velocity field. To emphasize this point, we have removed the conditional variable $\xi$ from the main text, and have clarified how it can be used with our coupling framework in the appendix.

Originality: Our method with data-dependent coupling is general, as the base density $\rho(x_0|x_1)$ can be designed in many different ways. This offers a degree of design flexibility that is not available to existing approches.

Need to retrain in the in-painting experiments: We would kindly like to request some clarification on the precise meaning of "retraining". In the response below, we have assumed that this refers to "classifier" or "guidance"-based methods for diffusion models, which can leverage a pre-trained model for conditioning. If this is incorrect, please let us know and we will provide an additional reply.

To this end, we would first like to emphasize that we focus here on data-dependent couplings, as emphasized in our summary reply above. Nevertheless, to improve performance in super-resolution and in-painting, we also leverage conditional information. We would like to point out that conditioning can be used to improve performance in our method, but is not strictly necessary. For example, for in-painting, our coupling pairs a masked image with its unmasked counterpart. By contrast, a guidance-free diffusion model would likely perform much worse because it would have no information about the mask.

We would also like to stress that our conditioning is over high-dimensional images, rather than the more standard scalar class label leveraged in guidance-based approaches. This means that the corresponding guidance model used in score-based methods is a second high-dimensional generative model, and its training is comparably expensive to the problem we solve here. In this sense, we would like to emphasize that our method does not require "retraining", as guidance-based methods would require an equivalent training effort.

As stated above, please let us know if this is not what was meant by "retraining", and we will provide further explanation.

Confidence in the experiments: We have added the requested FID benchmarks and see that they outperform the existing approaches.

Thank you for the feedback!

For super-resolution, [Ho et al 2021] (Cascading Diffusion Models, CDM) actually compares against the [Nichol and Dhariwal, 2021] (Improved DDPM) and [Dhariwal and Nichol, 2021] (ADM/guided-diffusion) models mentioned. For 64x64->256x256, [Ho et a 2021] reports FID scores of 12.26 for improved DDPM; 7.49 for ADM; and 4.88 for CDM. The method we propose gives an FID score of 2.13. While the numbers here are reported from the tables in [Ho et al 2021], we will be glad to additionally run versions of them in our codebase over the coming weeks and report both the original and reproduced numbers.

Regarding inpainting benchmarks: after reading through the works you suggested, we found that [Song et al 2021] do not train any Imagenet models and that [Song, Durkan, et al 2021] only train on Imagenet 32x32. We have not trained any models at the lower resolutions of 32x32 and 64x64, so we will need time to do this.

We will work on these additional benchmarks. In the meantime, we hope that our internal benchmark, in which we compare the results of our approach with data-coupling to those obtained using the interpolant framework with conditioning only and no data-coupling (similar to what score-based diffusion models would use), is informative: using only conditioning performs well visually, but with worse FID score (1.35) than what we get with data-conditioned (1.13): see the table in the main reply above).

[Ho et a 2021] Cascaded Diffusion Models for High Fidelity Image Generation

[Nichol and Dhariwal, 2021] Improved Denoising Diffusion Probabilistic Models

[Dhariwal and Nichol, 2021] Diffusion Models Beat GANs on Image Synthesis

[Song et al 2021] Score-based Generative Modeling through Stochastic Differential Equations

[Song, Durkan, et al 2021] Maximum Likelihood Training of Score-Based Diffusion Models

We thank the reviewer for their input and we address their main concerns one-by-one below. We also refer them to our general reply for further information.

Derivation:

• While class-conditioning is indeed easily incorporated into the interpolant framework, coupling the base density to the data is a more significant change compared to the work in [3]. The resulting derivation of the transport equation follows similar steps, but the way to construct the coupled density $\rho_0(x_0,x_1) = \rho_0(x_0|x_1) \rho_1(x_1)$ via proper definition of $\rho_0(x_0|x_1)$ is non-trivial, as emphasizd in the next point.

Data-dependent coupling vs mini-batch OT:

• The data-dependent coupling introduced here is general and can be used with any suitable $\rho(x_0,x_1) = \rho_0(x_0|x_1) \rho_1(x_1)$, which can be tailored to the application domain. In the experiments, we focus on examples in which the base $\rho(x_0 | x_1)$ is Gaussian conditional on the data $x_1$, but we stress that both the coupled density $\rho(x_0, x_1)$ and the marginal density $\rho_0(x_0) = \int \rho(x_0,x_1) dx_1$ are not Gaussian even in this case.
• While related work by [1] and [2] showed how to construct minibatch couplings that try to approximate the optimal coupling, their perspective does not give tools to perform the more general types of data-dependent modeling we consider here. Moreover, the couplings we introduce are exact, and do not require (but also do not preclude) the use of additional algorithms such as Sinkhorn.

We believe that these features make the approach we propose more flexible and performant than existing methods. In the revision, we will rework our presentation to clarify these points. We will also notify the reviewer as soon as the revised version becomes available.

Empirical evaluation:

• As mentioned in our general reply, we have added FID benchmarks to demonstrate the utility of our data-dependent coupling.

Questions:

Can the authors emphasize the analogies to existing works and highlight the benefits of formulating conditioning and data-dependent coupling concepts in the stochastic interpolants framework?

• To re-emphasize our previous points, the main novelty of our work is to propose a generic way to perform data-dependent coupling. To do so, we exploit the flexibility in the base density afforded by the stochastic interpolant framework. This is different from class-conditioning (though the two approaches can be used in conjunction) and it offers a degree of design flexibility that is not available to existing methods.

Does this formulation add flexibility compared to [1,2] or [6] which uses both conditioning and data dependent coupling (sec 5.3). (I'm aware [6] is not to be considered as previous work, but I want to understand the superiority of this work over applications already present in other frameworks).

• For our reply about references [1] and [2], please see the second bullet under Data-dependent coupling vs mini-batch OT.
• In the recent reference [6], to the best of our knowledge, the coupling proposed is another OT coupling akin to that of [1] and [2]. While the iterative algorithm presented there is interesting, and while we will surely cite this work, it remains a somewhat orthogonal to the general method we propose here.

[1] Pooladian et. al., Multisample Flow Matching: Straightening Flows with Minibatch Couplings (2023)

[2]Tong et. al., Improving and generalizing flow-based generative models with minibatch optimal transport (2023)

[3] Albergo et. al., Stochastic Interpolants: A Unifying Framework for Flows and Diffusions (2023)

[4] Saharia et. al., Palette: Image-to-Image Diffusion Models (2022)

[5] Lipman et. al., Flow Matching for Generative Modeling (2022)

[6] Song et. al., Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation

I thank the authors for their response.

Data-dependent coupling vs mini-batch OT: I better understand the subtleties now. [1,2] focus on constructing the coupling satisfying both marginal constraints, i.e., define $\rho(x_0,x_1)$ such that $\rho_0(x) = \int \rho(x_0,x_1) dx_1$, $\rho_1(x) = \int \rho(x_0,x_1) dx_0$ given $\rho_0$ and $\rho_1$. In this work, the coupling is assumed to be a given $\rho(x_0,x_1)$ and the paper shows how to construct $\rho_0$ implicitly by defining $\rho(x_0|x_1)$ given the coupling.

The above perspective drives me to the realms of image to image translation applications (as also shown in the paper) but I believe there is not enough discussion on the relations to that. A very closely related paper with a great overlap in the basic ideas is [7]. Although framed under the Schrödinger Bridge formulation, to my understanding, the concepts are very similar except for this work smoothing with a gaussian the base $\rho(x_0|x_1)$. Furthermore, [7] has already shown the benefits of changing the base distribution over conditioning (like you showed in the additional experiments now). Can the authors comment on the relation to this work?

In that case, I stand by my statement that I believe the contribution of this paper is limited.

I do appreciate the formulation in another framework and I think the stochastic interpolants framework is simple and elegant, but this paper either needs to be reframed as a "formulation paper" and not as one introducing new concepts and adding a thorougher discussion on existing works or show strong experimental evaluation and ablations for the benefits of the subtle differences proposed in this work.

[7] Liu et. al., $I^2SB$: Image-to-Image Schrödinger Bridge (ICML 2023)

Clarifying our contributions compared to prior and concurrent works:

Clarifying our theoretical contributions:

• The data-dependent coupling introduced here is general and can be used with any suitable $\rho(x_0,x_1) = \rho_0(x_0|x_1) \rho_1(x_1)$, which can be tailored to the application domain. In the experiments, we focus on examples in which the base $\rho(x_0 | x_1)$ is Gaussian conditional on the data $x_1$, but we stress that both the coupled density $\rho(x_0, x_1)$ and the marginal density $\rho_0(x_0) = \int \rho(x_0,x_1) dx_1$ are not Gaussian even in this case.
• While related work by [1] and [2] showed how to construct minibatch couplings that try to approximate the optimal coupling, their perspective does not give tools to perform the more general types of data-dependent modeling we consider here. Moreover, the couplings we introduce are exact, and do not require (but also do not preclude) the use of additional algorithms such as Sinkhorn.
• This work extends previous results on stochastic interpolants to expound on the utility of using an arbitary base distribution.

Clarifying our experimental contributions: Here we delineate the differences between the experiments we have presented and prior work on in-painting and super-resolution. Our primary point is that the use of a data-dependent base density has significant numerical advantages over uncoupled approaches, even if they make use of conditional information.

• Super-resolution: Reviewer v6gN pointed to [3] and [4] as examples of work that have already implemented super-resolution. We emphasize that these are suitable baselines for our method. In both papers, a flow or a diffusion is used to push a Gaussian sample to a super-resolution image, where the vector field is conditioned on a low-resolution image. This is standard conditional generation. Instead, we propose to use a density in which each low resolution image is coupled to its high-resolution counterpart, rather than just to use a velocity field that is conditional on the low-resolution image. The advantage of this second approach is apparent from the FID table above.
• In-painting: Here, we again use a dependent density as our base distribution, where we couple each masked image to its unmasked counterpart. This has the advantage that the velocity field can be taken as zero outside the mask, which allows for controlled generation. We show the benefit of this in compared to an uncoupled, standard normal base density in our numerics, where we see an improved FID.

Adjustments made to the exposition

Thanks to the feedback from the reviewers, we have adjusted the exposition in the text to highlight the subtleties in the differences of our method compared to prior works (as detailed above). Moreover, we have adjusted the numerical experiments, some preliminary results of which are discussed below and summarized in the table above. We will update the paper with the final numbers at the end of the rebuttal period.

New Numerical Results: FIDs

As per the requests of the reviewers, we have included Frechet Inception Distance (FID) benchmarks of our model vs others.

In-painting: We compare FIDs for in-painting when generating from the coupled density $\rho(x_0 | x_1)$ against the baseline of generating from a standard normal on ImageNet $256\times256$. With only 200,000 steps of training, we achieve and FID of 1.13, compared to the baseline of 1.38.

Superresolution: With only 300k steps of training, we achieve an FID for super-resolution of ImageNet 64x64 to ImageNet 256x256 of 2.13, which surpasses other known methods such as those mentioned by the reviewers.

References

[1] Pooladian et. al., Multisample Flow Matching: Straightening Flows with Minibatch Couplings (2023)

[2]Tong et. al., Improving and generalizing flow-based generative models with minibatch optimal transport (2023)

[3] Saharia et al: Image Super-Resolution via Iterative Refinement https://arxiv.org/abs/2104.07636 (2021)

[4] Ho et al: Cascaded Diffusion Models for High Fidelity Image Generation, https://arxiv.org/abs/2106.15282 (2021)

[5] Lipman et al: Flow matching for generative modeling, https://arxiv.org/abs/2210.02747 (2023)