A Computational Framework for Solving Wasserstein Lagrangian Flows

Thank you for your significant engagement with our paper. Despite a stated lack of expertise, we appreciated your perspective.

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.

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.

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

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.

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.

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.

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

Thank you for the positive review! We are glad that you felt we explained technical content in a comprehensive and digestible way.

We have included further experiments comparing with related SB methods which, in light of other reviewers' comments, we felt was higher priority than synthetic experiments.

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

Thank you for your review. Below, we emphasize the `distinguishing characteristics’ of our approach and address other concerns.

it is not clear what is the distinguishing characteristic of the proposed approach in relation to other works.

Our approach provides a novel method for solving optimal transport and its variants.

Using a common parameterization and the objective in Eq. 17, our computational method applies to a diverse collection of ‘linearizable’ Lagrangians (Def 3.1). Existing works often solve a single one of these problems (see Sec. 3.3) by leveraging known structure or particular assumptions (e.g. W2 OT solvers using ICNNs, Brenier Thm).

Our main technical contribution is Proposition 1, which uses duality of tangent and cotangent space on the Wasserstein manifold to suggest computational objectives for solving Wasserstein Lagrangian Flows.

The paper does not provide sufficient evidence for the validity of their approach to solve all Lagranian type problems

We agree that our computational approach does not apply for all Lagrangians, but applies for `linearizable’ Lagrangians (defined in Sec. 3.2). This includes W2 OT, physically constrained OT, OT with general Lagrangian costs, unbalanced OT, and entropic OT.

We provide evidence for the effectiveness of these methods for solving trajectory inference problems in Sec. 4 and the experiments in the general response.

Minimizing the Lagrangian may not be the appropriate objective in general, because Lagrangian flows are not the minimizer, but the extreme points of the action integral.

Optimal transport is explicitly defined via cost- or kinetic-energy minimization, and the existence of minimizers is guaranteed, for example, in the Benamou-Brenier formulation of W2 OT. While the action functional could have multiple critical points (some of which are not minima), we know that all minimizers are critical points (i.e. Lagrangian flows).

There is not enough motivation for parametrizing the flow according to (26). For a standard OT problem, what is the optimal function that replaces the neural net in this equation?

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. We have also defined the optimal function which the neural network approximates.

Most of the background material can be moved to appendix as they are not critical to the proposed algorithm

Thank you for the feedback.
In Sec. 2.2 and 2.3, we hope to introduce important concepts (kinetic energy, (co)tangent space) using familiar examples, before the more abstract and general presentation in Sec. 3.1-3.2.
We will consider moving Sec. 2.1 to the appendix, where we give detailed consideration to OT with state-space Lagrangian costs (App. A.2.2).

I appreciate the author's response to my questions.

Optimal transport is explicitly defined via cost- or kinetic-energy minimization, and the existence of minimizers is guaranteed, for example, in the Benamou-Brenier formulation of W2 OT. While the action functional could have multiple critical points (some of which are not minima), we know that all minimizers are critical points (i.e. Lagrangian flows).

You are right about optimal transport, because it only contains the Kinetic energy. But in the paper, you present the Lagranigian as difference of "Kinetic" - "potential" energies, which I assume it is inspired by classical Lagrangian mechanics. My point was to say that in Lagrangian mechanics, we are interested in the stationary points of the action integral, not minimizers.

https://en.wikipedia.org/wiki/Stationary-action_principle

we have argued that our parameterization in Eq. 26 can universally approximate any absolutely-continuous distributional path.

Thanks for your answer. I can now see how one can parameterize the optimal path in the way you did. But this representation does not seem very efficient as it requires a single neural net to represent a discontinuous function of x, and learning each of the potential functions \phi_0 or \phi_1 is difficult on their own.

Thank you for your review. In particular, we feel the suggestions regarding Lagrangian SB experiments and expressivity of parameterization will improve the paper.

1. My main concern with the paper (and perhaps the research area) is that the quality of the solution to the OT/UOT/SB problems is not directly measured

Note that in Table 1-2, we compare our WLF-OT to the exact OT solution. However, we agree that evaluation metrics are difficult in this research area. Among scRNA setups, our experimental setting was motivated by (i) testing our ability to achieve accurate interpolation between given marginals and (ii) demonstrating how one might explore various Lagrangians or inductive biases for approximating ground-truth dynamics for a given problem of interest.

2. representational capacity of equation (20) not clear or discussed. E.g., can it universally approximate distribution paths, and it is obvious that it can capture an optimal transport path?

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.

3. seems like the experimental settings considered by the other Lagrangian SB and generalized SB papers would be reasonable to compare against

In the general response, we have included results in the experimental setting of the Lagrangian SB paper (which is also scRNA, but is their only non-toy experiment). All of our proposed methods perform comparably or better than the best baseline method.

4. Appendix C.2 has some image generation experiments, but they seem incomplete

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.

We appreciate the knowledgeable and thorough review! We have run two sets of new experiments to address the concern about additional baselines, and summarized our results in the general response (and soon in the new revision). We address the remaining concerns below.

highlight the conceptual differences between your method and Chow et al. (2020) as well as Carrillo et al. (2021)?

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.

1. several approaches that propose Schrödinger bridges for modeling single-cell dynamics and beyond

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

2. It seems a bit odd not to compare against methods that also incorporate unbalanced solutions… the most obvious comparison would be with Pariset et al. (2023) and Pooladian et al. (2023b), which you cite but do not compare against. Is there a reason why?

Thank you for the suggestion. We have included the comparison with Pariset et al. (2023) in our general response.

Note, Pooladian et. al 2023b is not an unbalanced transport method. Their Sec. 5.1 considers physically-constrained OT, but only in two-dimensional experiments.

3. First 5 principal components (Table 1). Does the method not scale to larger problems?

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.

4. 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?

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

5. Is there any reasoning when to use which version of the method?

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.

5. If you compare the three methods using the same number of neural network parameters, do the results change?

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.

6. "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 / unbalanced OT.

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.

Parameterization of Distributional Path $\rho_t$

Reviewers uk1R (5,3) and zR9h (8,3) asked about the expressivity of our parameterization, while E9Se (3,3) asked about its motivation and optimality conditions. We argue below that our parameterization can universally express absolutely-continuous distributional paths $\rho_t$, and agree with reviewers that this result will improve our revised submission.

First, consider the W2 optimal transport case, where (given the optimal $\phi^*$) one can sample $x_t \sim \rho_t^*$ from the optimal marginals using $x_0 \sim \mu_0$ and $x_t =x_0 + t \nabla \phi_0^*(x_0)$, or $x_1 \sim \mu_1$ and $x_t =x_1 - (1-t) \nabla \phi_1^*(x_1)$ by the Brenier Thm. Consider a function $f: (x_0, x_1, t, \mathbb{I}[t < .5]) \mapsto \begin{cases} x_0 + t \nabla \phi_0^*(x_0) \qquad \text{ if } t<0.5 \newline x_1 - (1-t) \nabla \phi_1^*(x_1) \text{ if } t \geq 0.5 \end{cases}$.

The NN in our parameterization in Eq. 26 has the same signature as $f$ above. Since we also pass an indicator variable $\mathbb{I}[t < .5]$ to our neural network (as described below Eq. 28), the NN is able to learn a discontinuous function. By the universal expressivity of neural networks, our parameterization is thus expressive enough to learn the optimal $f$ ( after accounting for the $(1-t)x_0 + t x_1$ term), which yields samples $x_t \sim \rho_t^*$ from the optimal marginals for the W2 OT problem.

More generally, any absolutely-continuous distributional path $\rho_t$ on the W2 manifold can be traced using an ODE with a vector field $v_t^* = \nabla \phi^*_t$ (Ambrosio et. al 2006 Thm 8.3.1 [1]). Following similar reasoning as above, we define the function $f$ by running the ODE in the appropriate direction, e.g. $x_t = ODE(x_0; \phi^*_t)$ if $\mathbb{I}[t < .5]$.

For evaluation, held-out test samples are propagated and compared to the test samples using the W1 distance.

Table 1 of the Lagrangian SB paper [2] uses the W2 distance, not the W1 distance. Is it possible that the baseline methods in the table you sent use the W2 distance as they were reported in [2], but the WLF results are measured using the W1 distance?

(Also, it would be useful to include the NLSB results in the final version of the table as well)