Paper: A Computational Framework for Solving Wasserstein Lagrangian Flows (ICLR-2024)

Reviewer comment (Theory):
The authors propose a method claimed to be effective and efficient across a wide variety of domains, including transport with constraints and unbalanced transport, but the specifics sometimes contradict the claimed generality. For example, the constraints they enforce are on the marginal values themselves, but it is unclear how the method would handle constraints in other forms, such as symmetries in the problem or distributions, which are common in physics.

Author response:
> they say (method) is effective and efficient in a wide variety of domains including transport with constraints and unbalanced transport... For example, the constraints.. are on the marginal values themselves. However, what if the constraints come in a different form, e.g. a symmetry in the problem or distributions? 
 
We have been careful to only refer to 'marginal’ constraints ('intermediate’ or 'endpoint’) throughout the paper.   Symmetry constraints are fascinating, but beyond the scope of this work.

Reviewer comment (Experiments):
The authors motivate their problem using an RNA sequencing task, where population statistics are observed and the goal is to infer individual statistics. However, none of the experiments test this specific idea — all experiments predict population statistics. This gap leaves the reviewer unconvinced that the approach achieves the high-level goal.

Author response:
> we can only observe population statistics and we want to infer individual statistics from this… But looking at the experiments, none of the experiments test this specific point. All the experiments are predicting population statistics.

We agree that it would be interesting to test against ground truth individual trajectories.  Since this is not possible on real data, we constructed the leave-one-out experiments to test the quality of our marginal interpolations.

Reviewer comment (Reproducibility):
Below equation 27, the authors say they perform optimization with respect to $\eta$ using a reparameterization trick, but the reviewer could not find a reference or explanation of this trick.

Author response:
> Below eq 27, the authors say that they perform optimization with respect to $\eta$ using a reparameterization trick. What is this trick? I couldn’t find a reference to it.

The reparameterization trick refers to optimizing an expectation $E_{\rho}[f]$ under $\rho_t(x_t)$ by sampling some relevant stochasticity (here, $(x_0, x_1)$), parameterizing a transformation $x_t = g_\eta(x_0,x_1,t)$, and taking gradients using the chain rule, we have

$$\frac{d}{d\eta}E_{\rho_t (x_{t})}\left[f\left(x_{t}\right)\right]=E_{\rho_0\left(x_{0}\right)\rho_1\left(x_{1}\right)}\left[\frac{\partial x_{t}}{\partial \eta}\frac{\partial f}{\partial x_{t}}\right]$$

For more information, see: https://en.wikipedia.org/wiki/Variational_autoencoder#Reparameterization

Reviewer comment (Experiments):
In the experiments, it is unclear which time points are the left-out marginals. It is not specified whether they are in the future (requiring extrapolation) or in between known points (requiring interpolation). The reviewer considers the extrapolation task significantly harder.

Author response:
> Are the left-out marginals in the future and need to be extrapolated? Or is it interpolated as in-between points?

We leave out one of the intermediate marginals each time we run our algorithm, and report the average W1 distance over runs.

Reviewer comment (Theory):
The reviewer wishes the authors had expanded more on the potential extension to the quantum Schrödinger equation. In the appendix, the Madelung transformation appears to be only for one particle, which is not of practical relevance. The authors do not explicitly state this, and more precision is desired. Extending to multiple particles in the quantum setting is challenging due to issues such as entanglement, as addressed in approaches like density functional theory.

Author response:
>The Madelung transformation in the appendix appears to only be for one particle… which is not of practical relevance.   The authors do not state this exactly

The Madelung transformation in the appendix can be extended to $n$ particles by concatenating their coordinates into the vector $x \in \mathbb{R}^{dn}$, where $d$ is the dimensionality of the space, and introducing the corresponding symmetries (bosonic, fermionic) into the density and the vector field. The notorious challenges with the introduction of the symmetries, together with the fact that the Lagrangian becomes non-linear, suggest leaving the study of quantum systems beyond the scope of our paper.

Reviewer comment (Experiments):
The experiment task is difficult to follow without consulting the appendix and details of the Kaggle competition. The reviewer wishes the authors had described more clearly what the raw data is, noting that the resulting task is the PCA of this raw data.

Author response:
> The particular task in the experiment is tough to follow… e.g. simple things like what the raw data is in the form of. The resulting task is the PCA of this raw data.

Thank you for the suggestion.   However, due to space constraints, we deferred details to Appendix and referenced related literature.

Reviewer comment (Novelty):
The reviewer does not believe this paper unifies the field, as the authors themselves note that their method may not apply well in some situations (e.g., Schrödinger equation). The reviewer advises caution in the authors’ language regarding this point.

Author response:
>  I do not think this paper unifies the field. In fact, there were many situations where the authors stated to their credit that their method may not apply very well (e.g. Schrodinger equation). I am sure people would have to get a different perspective for these problems. So I would caution the authors in their language surrounding this point.

We have tried to limit our use of 'unified’ to refer to our computational approach for the 'linearizable’ objectives in Sec. 3.2 and Eq. 17.   This level of generality is still notable, as many previous works have derived problem-specific methods which rely on additional assumptions or known structure in each setting.

We feel that the perspective of ‘Lagrangian optimization on the density manifold’ and our novel dual objective in Thm 1 can be insightful to understand (and make connections between) a diverse range of problem settings, even where our computational approach does not directly apply.   For example, solving problems with nonlinear Hamiltonians, our dual objective suggests maintaining both a sampler and a density model to parameterize $\rho_t$.

Reviewer comment (Writing):
I may have missed it, but I think the notation $\phi_{\overset{.}{\mu_t}}$, for instance in Equation (7), is not defined.

Author response:
$\phi_{\dot{\mu}_t}$ indicates the cotangent vector corresponding to $\dot{\mu}_t$ as in Eq. 6 (which is written in terms of $\rho_t$ instead of $\mu_t$).

Reviewer comment (Theory):
If I understand correctly, the representational capacity of equation (20) is not clear or discussed. E.g., can it universally approximate distribution paths, and is it obvious that it can capture an optimal transport path? Experimentally it seems that it is the case, but it would be good to present or discuss more intuition behind this family of paths.

Author response:
Thank you for the suggestion!  In the general response, we have argued that our parameterization in Eq. 26 can universally approximate any absolutely-continuous distributional path (e.g. OT).   We will include a rigorous proof in the camera-ready.

Reviewer comment (Evaluation):
Appendix C.2 has some image generation experiments, but they seem incomplete and there are no quantitative evaluations of the method.

Author response:
This is not a main focus of the work.   However, we note that exact OT requires the velocity vector field to be a gradient field, which is a more strict condition than alternative 1-step generation methods such as rectified flows.   In fact, exact OT methods have rarely been applied to high-dimensional settings such as images.    The generated images show that multi-step generation and one-step generations are nearly identical, which indicates satisfying the zero-acceleration condition.

Reviewer comment (Novelty):
From a methodological standpoint, could the authors highlight the conceptual differences between your method and Chow et al. (2020) as well as Carrillo et al. (2021)? They seem to capture several of the frameworks discussed in this work and it is unclear in what sense you are providing a "unifying framework" here.

Author response:
Comparison with Chow et al. (2020) was discussed in detail at the top of page 9 in the original submission.  Carrillo et al. (2021) is concerned with JKO-style discretizations of W2 gradient flows, which does not address endpoint-constrained problems such as optimal transport.

Reviewer comment (Experiments):
I am surprised by the choices of baselines. Could the authors comment on why they only compared against Tong et al., 2023a and Tong et al., 2023b? There are several approaches that propose Schrödinger bridges for modeling single-cell dynamics and beyond, in particular Shi et al. (2023), Somnath et al. (2023), Liu et al. (2023), and Chen et al. (2022) to name a few. Vargas et al. (2021) and Somnath et al. (2023) in particular contain single-cell experiments.

Author response:
In the general response, we have added comparison with additional SB baselines (Vargas et. al 2021, Chen et. al 2022 (FB-SDE), de Bortoli et. al 2022, along with Koshizuka et. al 2023 (Neural Lagrangian SB) and others).

Reviewer comment (Experiments):
What is the motivation of running the method on the first 5 principal components (Table 1) besides running on 50 and 100 PCs? Does the method not scale to larger problems? Usually, single-cell data is reduced to 50 to 100 PCs to capture most of the variance, 5 PCs are not enough.

Author response:
We demonstrate the scalability of our method in Table 2 in the original submission, showing results on 50 and 100 principal components for the Cite and Multi datasets.   As in the 5-d cases, our proposed approach achieves improvements over the baselines.

Reviewer comment (Evaluation):
Why does the method barely outperform the exact OT method? For Table 2, why is exact OT even better on 50 and 100 PCs than the proposed method WLF-OT?

Author response:
WLF-OT and OT-CFM approximate the straight-path exact OT solution (which is not differentiable when transitioning between transportation plans) with a differentiable function.  This may lead to better performance in some cases, since the exact OT solution does not reflect the ground-truth dynamics of the cells in either Table 1 or 2. 

We also note that the evaluation metric in this setting is the W1 distance rather than W2, in keeping with previous works on the subject (Tong et. al 2023a,b, Koshizuka et. al 2023).

Reviewer comment (Experiments):
Is there any reasoning when to use which version of the method? If you compare the three methods using the same number of neural network parameters, do the results change?

Author response:
Our goal was to emphasize throughout the paper (Introduction, first sentence of each Example in Sec. 3.3, Conclusion) that prior knowledge or inductive bias of the ground truth dynamics of the task should inform the choice of Lagrangian.

The results for our proposed WLF methods use MLPs with the same number of layers and neurons.
For baseline methods, the differences in architectures or number of parameters are insignificant from the perspective of the runtime.

Reviewer comment (Experiments):
In your conclusion you state: "we studied the problem of trajectory inference in biological systems, and showed that we can integrate various priors into trajectory analysis while respecting the marginal constraints on the observed data, resulting in significant improvement in several benchmarks". This is not what the results suggest in my opinion. The chosen datasets do not allow for a thorough study of cell death and birth and thus do not provide a framework that allows studying unbalanced formulations of OT. However, there are multiple datasets that contain screens of cancer drugs that induce cell death (see Srivatsan et al., 2020) as well as the datasets studied in Pariset et al. (2023), Lübeck et al. (2022), or Eyring et al. (2022) that would be more suitable.

Author response:
We believe this statement is supported by our exposition and results, and does not make specific claims about unbalanced OT or birth/death dynamics.   As requested, we have included an additional unbalanced OT experiment capturing birth/death dynamics in the general response.