Paper: Computing high-dimensional optimal transport by flow neural networks (ICLR-2024)

Reviewer comment (Presentation):
The Figures are not all of good quality. Notably, Figure 2 and 4 are a bit too small and we cannot really distinguish the results of Figure 2,b.

Author response:
1. **(Presentation) The Figures are not all of good quality. Notably, Figure 2 and 4 are a bit too small and we cannot really distinguish the results of Figure 2,b.**  

Thank you. We will re-make the figures in future revisions.

Reviewer comment (Novelty):
I think that some related works are not cited. For instance, [1] parameterize Normalizing Flows (NFs) with Monge maps, [2] train NFs using the JKO scheme and the dynamic formulation of OT and [3] improves the OT cost of Normalizing Flows. Also, [4] proposes a way to find a Normalizing Flow between two arbitrary distributions.

Author response:
2. **(Citation) I think that some related works are not cited. For instance, [1] parameterize Normalizing Flows (NFs) with Monge maps, [2] train NFs using the JKO scheme and the dynamic formulation of OT and [3] improves the OT cost of Normalizing Flows. Also, [4] proposes a way to find a Normalizing Flow between two arbitrary distributions.**  

    [1] Huang, Chin-Wei, Ricky TQ Chen, Christos Tsirigotis, and Aaron Courville. "Convex potential flows: Universal probability distributions with optimal transport and convex optimization." arXiv preprint arXiv:2012.05942 (2020).  

    [2] Vidal, Alexander, Samy Wu Fung, Luis Tenorio, Stanley Osher, and Levon Nurbekyan. "Taming hyperparameter tuning in continuous normalizing flows using the JKO scheme." Scientific Reports 13, no. 1 (2023): 4501.  

    [3] Morel, Guillaume, Lucas Drumetz, Simon Benaïchouche, Nicolas Courty, and François Rousseau. "Turning Normalizing Flows into Monge Maps with Geodesic Gaussian Preserving Flows." arXiv preprint arXiv:2209.10873 (2022).  

    [4] Panda, Nishant, Natalie Klein, Dominic Yang, Patrick Gasda, and Diane Oyen. "Semi-supervised Learning of Pushforwards For Domain Translation & Adaptation." arXiv preprint arXiv:2304.08673 (2023).  

Thank you. We will add the citations in future revisions. Meanwhile, please see our response to the common questions above regarding the difference of our approach with theirs.

Reviewer comment (Writing):
Typo above equation (5): "The inner-loop training of $r_1$ is by"

Author response:
3. **Typos: Above equation (5): "The inner-loop training of r_1 is by”**  

Thank you. We will fix this in the revised draft in future revisions.

Reviewer comment (Theory):
How the proposed two loss function satisfies the condition $\partial_t\rho+\nabla\cdot(\rho v)=0$ during training.

Author response:
The condition $\partial_t \rho+\nabla \cdot (\rho v)=0$ is the continuity equation which relates the density $\rho(x,t)$ to the velocity field $v(x,t)$. Specifically, if $x(t)\sim \rho(x,t)$ and $\partial_t x(t) = v(x,t)$, then $\rho(x,t)$ *will* satisfy the continuity equation. Therefore, we do not need to enforce this condition in practice; during training, we only need to focus on the boundary conditions (i.e., $\rho(x,0)=p(x), \rho(x,1)=q(x)$).

Reviewer comment (Experiments):
How does the KL loss impact the final training results (i.e., if the terminal condition is not considered, how does the final result become)?

Author response:
If the terminal condition is not considered, the learnt dynamic trajectory would not be between $P_0=P$ and $P_1=Q$, but between $P$ and $\hat{P}_1$, where $\hat{P}_1$ can differ from $Q$. In that case, this is different from our original goal of learning the dynamic OT between $P$ and $Q$ based on the BB-formula (see Eq (2)). In practice, as in the case of handbag -> shoes shown in Section 4.5, this would imply a poor match between the true shoes and generated ones, which are translated from handbags by the trained flow.

Reviewer comment (Experiments):
The authors are encouraged to compare their proposal with recent state-of-the-art diffusion based generation methods.

Author response:
Thank you for the suggestion. We want to emphasize that our main goal is not to train a generative model, but to approximate the dynamic OT between arbitrary data distributions $P$ and $Q$. Doing so can be helpful in statistical inference tasks such as DRE. In comparison, diffusion models have reached SOTA performance on generative modeling (where $Q$ is typically the standard Gaussian), but these models are not learning the dynamic OT.

Reviewer comment (Theory):
At the first glance, the authors position their main optimization objective as a minimization problem. However, having understood the paper better, one may realize that their objective actually constitutes the min-max optimization because it uses the variational (discriminator-based) estimation of KL divergence. It seems like the classification net can be viewed as a discriminator (in accordance with the formula (5) and Table 2 of the paper [9]). The trained flow-based network is used as a generator according to expressions (6) and (3). Unfortunately, the authors did not mention this important fact over the paper, which seems a little bit unfair with respect to the reader.

Author response:
1. **(Implicit variational learning) At the first glance, the authors position their main optimization objective as a minimization problem. However, having understood the paper better, one may realize that their objective actually constitutes the min-max optimization because it uses the variational (discriminator-based) estimation of KL divergence. It seems like the classification net can be viewed as a discriminator (in accordance with the formula (5) and Table 2 of the paper [9]). The trained flow-based network is used as a generator according to expressions (6) and (3). Unfortunately, the authors did not mention this important fact over the paper, which seems a little bit unfair with respect to the reader…[Meanwhile,] existing flow-based methods [1],[2],[3],[4]...have simpler non-adversarial optimization.**

    [1] - “Action matching: Learning Stochastic Dynamics from Samples”, Neklyudov et al., 2022 

    [2] - “Building normalizing flows with stochastic interpolants”, Michael S. Albergo et al., 2023 

    [3] - “Flow matching for generative modeling”, Lipman et al., 2022 

    [4] - “Rectified flow: A marginal preserving approach to Optimal transport”, Liu Qiang, 2022

Thank you for pointing this out. We will make this implicit variational learning clear in future revisions. We agree that compared to [1-4], we have to use a “discriminator” to enforce the boundary condition. However, there are essentially differences in terms of the objective and what we can achieve. First, [1-3] learn interpolations between $P$ and $Q$, but such interpolations are not necessarily the (dynamic) OT between $P$ and $Q$. In fact, such an interpolation can be used as the initialized flow for our OT refinement; we have considered using [2] as an initialization strategy in the handbag -> shoe example shown in Section 4.5. Regarding [4], the method can be viewed as an alternative direction descent of the BB formula (see Section 5.4 therein). However, the linear interpolation process in the alternative scheme “is *not* deterministic or ODE-inducible unless the fixed point is achieved”, making it essentially different from the BB approach that focuses on deterministic ODE solutions. In addition, there lacks empirical evaluation of the proposed rectified flow approach.

Reviewer comment (Presentation):
There are lot of metrics used in the experiment section: mutual information, FID, and BPD. If you can group them in one table or plot, it would be cleaner to compare the methods with all three metrics.

Author response:
Thank you for the suggestion. We would do so in the future revisions, which will include additional experiments to verify the method.

Reviewer comment (Novelty):
We recommend the authors cite the following two recent works on MMD and gradient flow:  
Fan, J. and Alvarez-Melis, D., 2023. Generating synthetic datasets by interpolating along generalized geodesics. arXiv preprint arXiv:2306.06866.  
Hua, X., Nguyen, T., Le, T., Blanchet, J. and Nguyen, V.A., 2023. Dynamic Flows on Curved Space Generated by Labeled Data. arXiv preprint arXiv:2302.00061.

Author response:
Thank you, we will cite these works in the future.

Reviewer comment (Experiments):
Have you done an ablation study on different loss functions or their weights?

Author response:
We performed an ablation study on the weight parameter $\gamma$ in Table A.2, where the DRE performance barely varies across $\gamma$. We will consider different loss functions (e.g., the MMD loss) in the future.

Reviewer comment (Presentation):
the title of the paper is "Computing high-dimensional optimal transport by flow neural networks". However, a significant part of the paper is devoted to benchmarking of the capability of the algorithm to perform density ratio estimation. So it is not clear what is the main aim of the paper - to compute OT, or to estimate log density ratio.

Author response:
Thank you for the question. We intend to present methods that can be useful for both the OT problem and the DRE task. Specifically, we aim to approximate the dynamic OT solution in BB formula (see Eq. (2)) by a flow model, where the learnt continuous OT trajectory enables more accurate DRE. In future revision, we will make our contribution more clear with better experimental design.

Reviewer comment (Reproducibility):
it is not clear how to tune a value of gamma in (3). Any recipes for automatic tuning?

Author response:
Thank you for the question. At present, we don’t have a recipe for automatic tuning, but find that using $\gamma=0.5$ is typically a good starting point. We will explore such recipes in the future. Nevertheless, In Table A.2, we conducted an ablation study on the MNIST example with respect to $\gamma$, where we verified that the performance is not sensitive to the choice of $\gamma$.

Reviewer comment (Theory):
page 18 (the second line after the displayed formula 15): Why Q = N(0,I_d) if P = N(0,Sigma)?

Author response:
Thank you for the question. This example is taken from [Rhodes et al., 2020, Appendix E]. Specifically, we note that $\Sigma$ is a block-diagonal matrix with 2x2 blocks. As a result, for $X\sim N(0,\Sigma)$, $U=(x_1,x_3,\ldots,x_{d-1})$ and $V=(x_2,x_d,\ldots,x_d)$ are random vectors drawn from $N(0,I_{d/2})$. As a result, $p(U)p(V)=Q(x)$ for $Q=N(0,I_d)$.  

Reference:  

[Rhodes et al., 2020] Telescoping Density-Ratio Estimation

Reviewer comment (Experiments):
page 19: it is not clear how gamma = 0.5 was selected. Why not 0.6?

Author response:
Thank you for the question. At present, we don’t have a recipe for automatic tuning, but find that using $\gamma=0.5$ is typically a good starting point. We will explore such recipes in the future. Nevertheless, In Table A.2, we conducted an ablation study on the MNIST example with respect to $\gamma$, where we verified that the performance is not sensitive to the choice of $\gamma$.