Paper: Stochastic interpolants with data-dependent couplings (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
***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.

Reviewer comment (Theory):
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 justification and the definition of conditional expectation. You kinda have to use that the expected value only depends on (the time derivative) of the stochastic coupling.

Author response:
***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)\cdot \nabla \phi(x_t)|x_t]]
= \mathbb{E}[\mathbb{E}[(\dot\alpha x_0+ \dot\beta x_1 + \dot\gamma z)|x_t]\cdot \nabla \phi(x_t)]$$

Reviewer comment (Theory):
Is there any intuitive interpretation for the losses like in the diffusion case?

Author response:
***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".

Reviewer comment (Theory):
In the paper it is discussed, that one needs $\sigma>0$ in the superresolution experiment. If not one would try to establish a normalizing flow between a lower and a higher dimensional manifold. However when one thinks about the inpainting example filling all the boxes with random (Gaussian) noise could lead to an overestimation of the dimension of the target data, therefore prohibiting "a true transport equation" to hold as this defines a normalizing flow.

Author response:
***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.

Reviewer comment (Theory):
Is it possible to derive the losses from the forward/reverse SDEs so one does not necessarily enforce invertibility?

Author response:
To clarify the theoretical question: The current formulation of our Theorem 1 and Corollary 2 does require that the base and the target distributions have densities, in which case the flow is invertible. These results can be generalized to situations where we go from a higher dimensional manifold to a lower dimensional one, but in this case the drift must become singular at the end point (to focalize the process) whether we use an ODE or and SDE as generative process — our loss will  in principle learn this singular drift in this setup, rendering the flow noninvertible. This is similar to what happens with score based diffusion models since in the same setup the score also becomes singular at the end point (and plays the same focalizing role).  However, score-based diffusion models can also go from a lower dimensional manifold to a higher dimensional one, with the help of the diffusion term at initial time. In our approach, if we use a generative model based on an ODE, this cannot be achieved since the velocity would then have to be singular at initial time, and the flow cannot be started. We note that we could generalize our construction by using the SDE instead of the ODE as generative model  but this is not done here.  In practice, it seems to us that having the capability to focalize the process at final time (i.e. on the target distribution, since we always go forward in time from base to target) is what matters most.

Reviewer comment (Writing):
A small nitpick: I think the formulation "reverse" SDE is more appropriate since backward has a different meaning in the probability theory context.

Author response:
**Reversed-time vs backward SDE:** We agree that the first denomination is better and we have adopted it.

Reviewer comment (Novelty):
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 methodologies. However, the proposed framework does generalize beyond this to other base distributions (for example), so I would expect (or rather, like to see) more empirical emphasis to be placed on this setting.

Author response:
***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.

Reviewer comment (Experiments):
In a similar vein, for the inpainting experiments, there is a big issue in that existing score based methods (e.g. ScoreSDE) can inpaint (up to some approximation error + some necessary hacks) without having to retrain, while the current results come about through retraining.

Author response:
***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.

Reviewer comment (Novelty):
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 as being widely used in diffusion models (e.g., [4]), and since they can be thought of as particular cases of stochastic interpolants, the addition of conditioning is straightforward.

Author response:
***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.

Reviewer comment (Evaluation):
**Empirical evaluation** - the empirical evaluation is solely qualitative, which makes it impossible to assess whether there is a benefit in using conditional and data dependent couplings in the stochastic interpolant framework.

Author response:
***Empirical evaluation:*** 
- As mentioned in our general reply, we have added FID benchmarks to demonstrate the utility of our data-dependent coupling.

Reviewer comment (Novelty):
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?

Author response:
> *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.

Reviewer comment (Novelty):
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).

Author response:
> *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.