Computing high-dimensional optimal transport by flow neural networks

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.

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.

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.

1. (Continuity equation) How the proposed two loss function satisfies the condition $\partial_t \rho+\nabla \cdot (\rho v)=0$ during training.

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)$).

2. (Loss to final result) 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)?

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.

3. (High-dim images) No large-scale/high-resolution image generation experiments.

Thank you. We will consider such examples in future revisions, besides the handbags -> shoes example in Section 4.5 on 64x64 images. We also want to emphasize that the main goal of the work is not to perform image generation, but to learn the dynamic OT leveraging the BB formula in Eq (2). Leveraging the learnt OT, we can thus effectively perform many other tasks such as DRE.

4. (Against diffusion models) The authors are encouraged to compare their proposal with recent state-of-the-art diffusion based generation methods.

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.

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.

2. (Simulation-based approach) The computation of integral of Neural-ODE along learned trajectories lies at the heart of computation inefficiency of the proposed algorithm in accordance with section 3.2. That is, the method is simulation-based

Thank you for pointing this out. We agree that the method is simulation-based. However, this is because we are directly trying to learn the dynamic OT based on the Benamou-Brenier formula (see Eq (2)). Prior works [Liu et al., 2022, Tong et al., 2023] have tried to learn the dynamic OT in a simulation-free manner, but as explained in the answers to the first common question, these approaches either fail to recover the exact deterministic ODE solution [Liu et al., 2022], or the method only works in theory with infinitely large batches [Tong et al., 2023].

References:

[Liu 2022] Rectified flow: A marginal preserving approach to Optimal transport]

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

3. (W2 objective) The formula (7) seems to only be an upper-bound for the true Wasserstein-2 distance but not the exact distance, right?

Thank you for the question. As explained in the paper, “the population form of (7) in minimization can be interpreted as the discrete-time summed (square) Wasserstein-2 distance”, which is between densities $\rho(\cdot, t_{k-1})$ and $\rho(\cdot, t_k)$.

1. (Presentation) There are a 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.

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

2. (Cite) 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.

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

3. (Different objective) In section 3.1, is it possible to use MMD in the loss instead of KL divergence?

Thank you for the question. We agree that the MMD objective can be an alternative relaxation of the terminal condition $\rho(\cdot, 1) =q$. We will explore the use of this objective in future works.

4. (Ablation) Have you done an ablation study on different loss functions or their weights?

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.

5. (Additional evaluation) Similar to question 2, is it possible to compute the KL divergence with the images themselves or embeddings of the images?

Thank you for the question. We can estimate the KL divergence between generated images and true images using Eq (6). Specifically, assume we have trained a classification network $\hat r_1$ that is close to the optimal classifier $r_1^*$ between $P_1$ and $Q$. Then, the KL divergence between $P_1$ and $Q$ can be computed as $-\mathbb{E}_{x\sim P_1} \hat{r}_1(x)$, where the expectation is approximated by finite samples.

1. (contribution, OT or DRE) 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 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

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.

2. (Not showing OT through FID) experimental results to verify efficiency of computed W2 high-dimensional optimal transport are not enough to claim accuracy and efficiency of the proposed approach. E.g. The authors consider some image translation tasks, but the FID score used to characterize accuracy does not guarantee that the computed W2 high-dimensional optimal transport map is accurate.

Thank you for pointing this out. We agree that a low FID score does not indicate an accurate learnt OT trajectory; we considered this example mainly to illustrate the scalability of our approach to non-trivial high-dimensional image datasets. In the future, we will more quantitatively evaluate our learnt OT trajectory, such as through this benchmark [Korotin et al., 2021]

Reference:

[Korotin et al., 2021] Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark

3. ($\gamma$ choice) it is not clear how to tune a value of gamma in (3). Any recipes for automatic tuning? On page 19: it is not clear how gamma = 0.5 was selected. Why not 0.6?

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$.

4. (Clarification) page 18 (the second line after the displayed formula 15): Why Q = N(0,I_d) if P = N(0,Sigma)?

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

2. Method

• (R2, R3, R4) Why do we have bi-directional design? Cannot one just train along one direction?
• (R1, R4) Why does the design of loss objective learn the dynamic OT? Any guarantees?

(Response to the first bullet) In practice, introducing bidirectional training enhances numerical accuracy and stability. We acknowledge the added computational cost of bidirectional training compared to training in one direction. In the future, we'll explore methods to enhance training efficiency and potentially simplify the design.

(Response to the second bullet) We based our loss objective on the Benamou-Brenier formula (see Eq (2)), which comprises the KL term in (6) and the transport cost term in (7). It's worth noting that (6) relaxes the terminal condition by employing a classifier, while (7) can be seen as the discrete-time sum of the (square) Wasserstein-2 distance.

We intuitively explain why minimizing (3) leads to learning dynamic OT. Firstly, when we train the classifier $r_1$ near optimality, minimizing the expected value of (6) enforces the terminal condition. Secondly, with a sufficiently fine time grid $0=t_0<t_1 < \ldots < t_K=1$, the expected value of (7) approximates the continuous-time transport cost well. We plan to formalize this intuition in future revisions.

3. Experiment

• (R1, R3, R5) More comparison with DRE baselines and/or examples. MNIST is not convincing (R3), and wants larger datasets (if DRE on images).

  Paper list:

  [Feng et al., 2021] Relative entropy gradient sampler for unnormalized distributions

  [Miller et al., 2022] Contrastive Neural Ratio Estimation

  [Heng et al., 2023] Generative Modeling with Flow-Guided Density Ratio Learning

  [Givens et al., 2023] Density Ratio Estimation and Neyman Pearson Classification with Missing Data

We thank the reviewers for the suggestion. We will compare our proposed OT+DRE method against these baselines on larger datasets in future revisions. Below, we summarize the key ideas of these methods and highlight the difference between what we proposed and theirs.

Some of these DRE baselines are based on Bregman divergence [Feng et al., 2021, Heng et al., 2023] while others are based on cross entropy loss [Miller et al., 2022]. The work by [Givens et al., 2023] further extends DRE under the Kullback-Leibler importance estimation procedure when missing data are present, but the DRE estimator is not based on neural networks.

We discovered that existing baselines only focus on learning a single-step estimator that optimally recovers $\log(p/q)$. However, previous studies [Rhode et al. 2020, Choi et al. 2022] have shown that DRE's empirical performance can significantly degrade when the KL-divergence between $p$ and $q$ exceeds tens of nats. This has prompted us to adopt a different approach: constructing an OT trajectory between $p$ and $q$ before applying DRE, which is continuous in time (see Appendix A). Empirically, this approach has resulted in improved performance compared to both the single estimator and telescopic estimators using non-OT trajectories.

References:

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

[Choi et al., 2022] Density Ratio Estimation via Infinitesimal Classification

• (R1, R3, R5) More comparison with OT baselines. For example, this benchmark https://github.com/iamalexkorotin/Wasserstein2Benchmark.

We thank the reviewers for the suggestion. We will compare our proposed method against these OT baselines in future revisions, especially on the high-dimensional examples of Celeb-A dataset.