Paper: Lagrangian Flow Networks for Conservation Laws (ICLR-2024)

Reviewer comment (Novelty):
To my knowledge, the idea of using the Lagrangian view of the continuity equation to build a neural parameterization that gives both velocity and density access is not new. In Section 3.1 and 3.2 of [1], the TIPF parameterization appears to be identical to LFlow proposed in the present paper, with the only difference being that in (13) of the present paper, $\Phi_t$ appears instead of the inverse, which is the case of TIPF. This is a minor difference for TIPF, since in [1] the conditional normalizing flow used is RealNVP [2], which has ready access to the inverse flow map.

Author response:
We are thankful to the reviewer for pointing out a related work, which we will include in the revised version of the paper. However, we respectfully disagree that this constitutes lack of novelty. 

The proposed approach shares only a high-level similarity with TIPF. Both papers use known results from classical theory, namely that conditional diffeomorphisms define probability densities that satisfy the continuity equation by construction. However, (i) the class of problems addressed by the two papers is fundamentally different and (ii) the proposed architectures are consequently designed for different purposes.
 
&nbsp;

### (i) Scope of the papers:
- TIPF consider well-posed forward problems for probability flows, where initial conditions, boundary conditions, and the full system of equations are known. Initial conditions are evolved forward in time with conditional bijections that fulfill the continuity equation by construction, providing an advantage compared to numerical methods.\

&nbsp;

- In contrast, we consider inverse problems for real physical phenomena. We rely only on sparse (and noisy) observations, with unknown initial conditions, velocity and total mass. This is a challenging ill-posed setting and cannot be directly approached with classical numerical forward solvers.
 
&nbsp;
### (ii) Proposed architecture:
- TIPF evolve a known initial density $p_{t_0}$ over time by directly parameterizing the Lagrangian flow map, which we denote as $X_t$ in Section 3.1. We shortly discuss this parameterization in Section 3.2. As $X_t$ is directly parameterized, $X_{t_0}$ must be enforced to be the identity map with an additional loss. Since there is no base density, observations are always mapped back to the known initial density $p_{t_0}$. One important limitation is that the bijection has to be inverted, which is computationally expensive. Therefore, TIPF rely on bijections with analytical inverses, which however severely limit the expressivity of the model.

&nbsp;
- In contrast, LFlows do **not** directly parameterize $X_t$ and $\rho_{t_0}$. Instead, only a single time-conditioned bijection $\Phi_t$ is learned, such that $X_t(x) = \Phi_t^{-1}(\Phi_{t_0}(x))$ (see Section 3.2). By mapping the time dependent density to a common base density, we do not even need an explicit $\rho_{t_0}$. Importantly, our approach enables us to avoid computing the inverse explicitly. We can compute the velocity by simply inverting the Jacobian of $\Phi_t$ instead of explicitly inverting $\Phi_t$. We would like to highlight that this is a significant contribution, since unlike TIPS we are free to choose highly flexible bijections with unknown or expensive inverses.

Reviewer comment (Theory):
Could the authors elaborate on the advantages of avoiding the inverse here?

Author response:
>Could the authors elaborate on the advantages of avoiding the inverse here?

Highly expressive bijective layers usually result in the loss of an efficient (analytical) inverse. Lipschitz constrained residual networks for example require a fixed-point iteration for obtaining the inverse. Backpropagating through such an operation (for minimizing a velocity loss) quickly turns infeasible for moderately sized networks. By completely avoiding the inverse we do not share the same limitations for the choice of the bijective layers as TIPF.

Reviewer comment (Experiments):
Could the authors also comment on the choice of the Lipschitz-constrained invertible densenets as the backbone architecture? What is the advantage/disadvantage compared to e.g. RealNVP [2]?

Author response:
>Could the authors also comment on the choice of the Lipschitz-constrained invertible densenets as the backbone architecture? What is the advantage/disadvantage compared to e.g. RealNVP [2]?

Invertible densenets have free-form Jacobians and showed high performance on generative modeling tasks compared to other models [3], including RealNVP. We further found them easy to finetune and less prone to training instabilities and overfitting compared competing bijective layers.

[3] Perugachi-Diaz, et al. Invertible densenets with concatenated LipSwish. NeurIPS 2021.

Reviewer comment (Experiments):
The experiments section is weak in my opinion, and weaker than that in [Richter-Powell 2022]. If I understand correctly, all experiments are in 2D. How would the parameterization fare in higher dimensions? Higher dimensional experiments are done in [Richter-Powell 2022]. In [1] it is found that TIPF does not work well in dimension >= 10. I suspect the same behavior occurs here, due to the reason that invertible pushforward networks might not be expressive (or might be hard to train).

Author response:
>If I understand correctly, all experiments are in 2D. How would the parameterization fare in higher dimensions?

We would like to correct the misunderstanding regarding the dimension. Both the  experiments in section 5.1 and the real-world bird migration experiment in 5.3 are in 3D. In the simulated fluid experiment we also show that the proposed LFlow outperforms competing methods by an even greater margin when going from 2D to 3D (see Figure 2, top right)

>The experiments section is weak in my opinion, and weaker than that in [Richter-Powell 2022].

Concerning the experiments, we respectfully disagree that our section is weaker compared to [Richter-Powell 2022]. From our point of view, the main purpose of PDE-constrained neural networks is the seamless incorporation of measurements for solving ill-posed inverse problems. Different to [Richter-Powell 2022], we actually evaluate such a setting with the bird migration experiment. 

> How would the parameterization fare in higher dimensions? [...]  In [1] it is found that TIPF does not work well in dimension >= 10. I suspect the same behavior occurs here, due to the reason that invertible pushforward networks might not be expressive (or might be hard to train)


We assume "working well in dimensions >= 10" refers to (i) computational scalability, (ii) generalization capability, (iii) the enforcement of additional PDEs, and (iv) expressivity of the bijections.

(i) In our 2D and 3D experiments we showed that the proposed LFlows scale better than competing methods in terms of memory and/or training time.

(ii) While the generalization capability in higher dimensions might be interesting to study, we are not aware of any real world settings where a density or velocity of a fluid in high (>3) dimensional spaces is observed. The work of [1] considers probability flows in high dimensions rather than physical observations of densities. 

(iii) Settings like the ones discussed in [1] require enforcement of additional equations in high-dimensional spaces, similar to our optimal transport experiment. As in PINNs, a PDE-Loss has to be minimized on the support of the density by sampling collocation points in high dimensions. We agree that the proposed LFLows would suffer from the same limitations in high dimensions. The fulfillment of the continuity equation however remains unaffected by higher dimensions.

(iv) Regarding expressivity, Normalizing Flows (and i-DenseNets) have already been succesfully employed in very high-dimensional settings (32x32 images) [3], leading to our assumption that the expressivity should not be the limiting factor.


> [...] Higher dimensional experiments are done in [Richter-Powell 2022]. [...]

We note that the high-dimension (>>3) experiments of [Richter-Powell 2022] are about decomposing vector fields into irrotational and divergence-free components. This is specific to divergence-free networks and unrelated to our work.

Reviewer comment (Theory):
At the bottom of the first page, it reads "in setting (ii) we measure only the density of the fluid but we know additional equations defining the velocity." If density is already known for all time, then isn't it already enough to reconstruct the velocity without needing to know any additional equations, by simply solving the continuity equation?

Author response:
> 1. [...] it reads "in setting (ii) we measure only the density of the fluid but we know additional equations defining the velocity." If density is already known for all time, then isn't it already enough to reconstruct the velocity without needing to know any additional equations, by simply solving the continuity equation?

&nbsp;

We never assume dense observations at all times, but we agree that in such a setting one could and should directly solve the continuity equation numerically. In contrast, for our setting we always assume sparse measurements in spacetime. We see how this is not apparent from the quoted sentence. We will state it more clearly in the revised version of the paper.

Reviewer comment (Evaluation):
In the evaluation in Section 5.2, the gold standard of $W_2^2$ is computed using discrete estimates (and then averaged). However, the discrete estimate of W2 is known to be biased. How can one be sure that the red bars on the right of Figure 3 are accurate? Would the bar shift if more samples were used?

Author response:
> 2. In the evaluation in Section 5.2, the gold standard of 
 is computed using discrete estimates (and then averaged). However, the discrete estimate of W2 is known to be biased. How can one be sure that the red bars on the right of Figure 3 are accurate? Would the bar shift if more samples were used?

We thank the reviewer for highlighting that, as we were unaware of this issue. 
We will rerun the discrete solvers with up to 50.000 samples (previously 20.000) to explore the behaviour and add the results to the appendix in a revised version. For the Circle-Pinwheel example, we can already state that the estimate remains rather stable, going from 7,29E-02 (+- 2,10E-03) to 7,26E-02 (+- 6,66E-04), where we report AVG (+- STDEV).

Reviewer comment (Reproducibility):
In (15), how is the double integral computed in practice? It looks like it requires uniform samples from the domain $\Omega$, which could be very ineffective if the flow only concentrates on a small portion of the domain.

Author response:
> 3. In (15), how is the double integral computed in practice? It looks like it requires uniform samples from the domain, which could be very ineffective if the flow only concentrates on a small portion of the domain. 
 
For the synthetic experiment we adopted the procedure of [Richter-Powell 2022] to be consistent with their setting. They rely on importance sampling for evaluating the integral (15), with a sampling distribution $q(t,x)$. Samples in time are drawn uniformly, i.e. $q(t)\sim U(t_0, t_1)$. Samples in space are drawn from a uniform mixture: $q(x) = \frac{1}{3} p_{t_0}(x) + \frac{1}{3} p_{t_1}(x) + \frac{1}{3}U_\Omega$ with $U_\Omega$ being the uniform distribution on the domain.

Different to [Richter-Powell 2022] we could also directly sample from the density, reformulating the integral as an expectation. To avoid backpropagating through the numerical inverse, one might also detach the density samples from the computational graph and use them as a proposal distribution for importance sampling.

Reviewer comment (Presentation):
About the presentation. For me, the writings in Section 3.1 and 4 are confusing and hard-to-read. In Section 3.1, the authors use several theorems in an informal way and sometimes not give the exact formalization of theorem. I prefer to place the definitions of theorems in the main text and leave the proof in appendix, such as Theorem 1, 2 and the existence and uniqueness of Lagrangian flows for smooth vector fields. In Section 4, I don’t think the implementation of LFlows are well described. The model architecture in Appendix A.3 should be placed in main text. Besides, the implementation for baselines should be deferred into the Appendix.

Author response:
We thank the reviewer for the valuable feedback on the presentation of our work. We are eager to make the paper clearer and more accessible.

--------
&nbsp;

> For me, the writings in Section 3.1 [...] are confusing and hard-to-read.
> In Section 3.1, the authors use several theorems in an informal way and sometimes not give the exact formalization of theorem. I prefer to place the definitions of theorems in the main text and leave the proof in appendix, such as Theorem 1,2 and the existence and uniqueness of Lagrangian flows for smooth vector fields.

We understand the argument with respect to the exact formulation of the theorems, and would generally agree with this preference. 

However, in a previous submission, we did provide exact formalizations and statements of the theorems in 3.1, which we now moved to A.1.1. This led to large misunderstandings among the reviews. We received the unanimous feedback that including the formal and large theorems in the main text decreased the clarity and potentially led the reader to believe that it was the main constribution of the work, instead of the starting point.

We would thus lean towards keeping the current version of Section 3.1, with the full theorems provided in the Appendix A.1.


-----------------
&nbsp;


> [...] the writings in Section [...] 4 are confusing and hard-to-read.
> In Section 4, I don’t think the implementation of LFlows are well described. The model architecture in Appendix A.3 should be placed in main text. Besides, the implementation for baselines should be deferred into the Appendix.

We thank the reviewer for highlighting these clarity issues, and agree that Section 4 can be improved by following these recommendations.

We rewrote the implementation section accordingly. That is, Section 4 now only discusses the implementation of LFlows, and does so in a clearer way. The implementation of the baselines were deferred to the Appendix A.3.

Reviewer comment (Evaluation):
About the noise observations in the real-world dataset. Since LFlows enable the strict continuity equation, does it come across problems when the input data is noisy even wrong? Some discussions are expected.

Author response:
> About the noise observations in the real-world dataset. Since LFlows enable the strict continuity equation, does it come across problems when the input data is noisy even wrong? Some discussions are expected.

We do not observe issues with respect to noisy observations. We think such noise should be addressed in the loss function / likelihood. For example, by minimizing the MSE on the log1p transformed density we assume log(1p)-Gaussian noise. The problem of designing a problem specific noise model is however not unique to our method, and is an essential part of many classical machine learning and data assimilation problems.

Another point is that we are in highly ill-posed settings. From the data alone it can not be determined whether mass is strictly conserved and comes from an unobserved region, or whether there is a sink or source. Thus, we argue that the motivation for enforcing a physical constraint in such a setting must come from domain knowledge. We assume that for few consecutive nights during migration, there is a negligible change in bird-mass compared to the resolution and accuracy of our sensors. If this assumption is wrong, then of course the strict enforcement of the constraint will be unhelpful and even harming for the prediction.