Paper: GeONet: a neural operator for learning the Wasserstein geodesic (ICLR-2024)

Reviewer comment (Experiments):
The application of neural networks to compute Wasserstein geodesics is interesting. The accuracy is still a big issue compared to classical mesh-based approaches. See Fu, et.al., High order computation of optimal transport, mean field planning, and potential mean field games. Journal of computational physics, 2023.

Author response:
(1) We agree with you that mesh-based methods are more accurate than learning-based methods. In general, amortized inference are based on reasonable assumptions on data and model (such as no distribution shift in the test time and the model capacity). Mesh-based approaches can accurately solve the Wasserstein geodesic (and other mean-field control/game/planning) down to machine precision (i.e., optimization error), while learning-based methods can only solve up to statistical (coming from data) and approximation (coming from model) errors. On the other hand, learning-based methods are often faster (and sometimes much faster) than mesh-based methods on higher-dimensional and larger-size problems.

Reviewer comment (Experiments):
Some more interesting computational examples in mean field control problems can be considered in future work. See Lin, et.al. Alternating the Population and Control Neural Networks to Solve High-Dimensional Stochastic Mean-Field Games, PNAS.

Author response:
(2) We thank you for pointing the interesting direction to make connection to mean-field control problems. The mean-field control problems are very much related to the dynamical OT formulation, both of which can be formulated as a convex-concave primal-dual saddle-point learning (and optimization) problem. On the other hand, there are two main differences. First, initial density and terminal state are required boundary conditions for mean-field control, while Wasserstein geodesic requires the boundary condition as fixed two densities. Second, the dynamical OT problem lacks of the interaction term, which is important in the mean-field control. Thus, it seems that mean-field planning is an intermediate problem which we believe can be studied as an operator learning problem in a similar manner as learning the Wasserstein geodesic. However, learning the dynamics of a general mean-field control is likely harder than dynamical OT problem and we leave it to the future work. We include this point in the *LIMITATIONS* section in the revision.

Reviewer comment (Experiments):
Some more numerical examples are needed to compare with the current method and the ones studied in the above literature.

Author response:
We remark we have added an additional experiment comparing our methodology to those of alternative literature, on empirical densities constructed from point clouds of Gaussian mixtures. We refer you to section 4.2.

Reviewer comment (Experiments):
It lacks more examples in application.

Author response:
(2) We have since remedied this by the introduction of a new experiment on empirical densities constructed from point clouds sampled from Gaussian mixtures. The primary purpose of this experiment is to present a comparison to other methodology, notably rectified flow (RF) [2], and conditional flow matching (CFM) [3]; however, we also remark that this experiment presents GeONet as suitable for a point cloud setting, if point clouds are made into empirical densities. The process can be reversed, by sampling points according to the generated geodesic densities. While GeONet is not a framework for learning the direct translocation of points, we hope this experiment is satisfactory is illustrating effectiveness in the discrete setting.

Reviewer comment (Reproducibility):
It is not clear how the endpoint densities are defined in the MNIST dataset experiment. The reviewer requests the code for the experiment on Gaussian Mixture densities.

Author response:
(3) We thank you for your appreciation of our code. We use an autoencoder framework to encode the entire MNIST data, as a $28 \times 28$ grid made into a $32$-dimensional vector, in which a decoder subsequently maps the encoded representation back to its original image. There are no restrictions on the encoding, and so the encoded representation naturally converges to highly irregular data. These encoded representations are shifted by a constant to ensure nonnegativity, and normalized to be made into densities. The geodesic is then learned between such encoded densities. We revisited this experiment this rebuttal, and ensured the encoded representations were made into densities, which lowered error by several percentage points for all times, in addition to including the full MNIST dataset is used for training. In the revision, we also included the source code for the full MNIST experiments in the supplementary material to reproduce the discretization of the endpoint densities in MNIST data experiment.

Reviewer comment (Novelty):
Lack of comparisons and related work: [1,2] both amortize Wasserstein geodesic learning. These should be at least cited and discussed.

Author response:
(a) Thanks for pointing out the references. In the revised version, we added these two references and made a brief discussion.

Reviewer comment (Novelty):
Lack of comparisons and related work: There are many methods to compute Wasserstein geodesics relatively quickly, although without amortization. I would be curious how the quality of interpolation compares to these more recent methods, and would amend this statement for more recent work [e.g. 3,4,5]. Recently, a machine learning method to compute the Wasserstein geodesic for a *given* input pair of probability measures has been considered in (Liu et al., 2021).

Author response:
(b) In the revision, we added comparison with two more suggested learning-based methods: the rectified flow (RF) [2] and conditional flow matching (CFM) [3]. We run a similar 2D Gaussian mixture simulation setup in the discrete setting, constructing empirical distributions from sampled point clouds belonging from fixed densities. We use POT as the ground truth for comparing GeONet, RF and CFM. The $L^1$ estimation error for geodesic at time point $t = 0, 0.25, 0.5, 0.75, 1$ are reported in Table 3 and an estimated geodesic example is shown in Figure 4. There are a few observations we can draw from this new experiments. First, RF and CFM have 3-4 times comparably larger estimation errors than GeONet, except for the initial time $t = 0$, only because this initial data is given and learned directly for RF and CFM. GeONet is the only framework among the comparison which encapsulates the geodesic behavior to a considerable degree. While the other methods are suitable for point cloud representations, the geodesic behavior is entirely lost and highly inaccurate to ground truth. Second, RF and CFM have the same fixed resolution as the input probability distribution pairing, while GeONet can be smoothed out for estimating the density flows on higher resolution than the input pairing (cf. the third row in Figure 4).

Reviewer comment (Experiments):
The experiments show how well GeONet works in toy settings, although even this is not very clear given the lack of comparisons.

Author response:
(a) Regarding the MNIST experiment, in the revision, we retrained GeONet on the full MNIST data, splitting with 30,000 training data for $\mu_0$ and 30,000 training data for $\mu_1$ and randomly pairing them together in training. New testing error in the $L^1$ metric is reported in Table 4, consisting of both the error in the encoded space and ambient pixel space. Comparing the encoded-space error in the initial submission (Table 3 therein), we see that increasing the MNIST data indeed decreases the testing errors. Furthermore, the encoded data was also normalized prior to computing the geodesics with GeONet, which likely also contributed to the decrease in error. Moreover, as expected, the ambient-space error is much larger than the encoded-space error, meaning that the geodesics in the encoded space and ambient image space do not coincide. We mentioned this in the *LIMITATIONS* subsection 4.5 (initial submission) that encoding strategy seems to be necessary with such a dataset. In addition, we included the source code for the full MNIST experiments in the supplementary material for our revised submission.

Reviewer comment (Experiments):
The experiments are all extremely toy with limited examples and dimensionality, with the largest experiment being on a small subset (5000) of MNIST on a 30 dimensional embedded space.

Author response:
(a) Regarding the MNIST experiment, in the revision, we retrained GeONet on the full MNIST data, splitting with 30,000 training data for $\mu_0$ and 30,000 training data for $\mu_1$ and randomly pairing them together in training. New testing error in the $L^1$ metric is reported in Table 4, consisting of both the error in the encoded space and ambient pixel space. Comparing the encoded-space error in the initial submission (Table 3 therein), we see that increasing the MNIST data indeed decreases the testing errors. Furthermore, the encoded data was also normalized prior to computing the geodesics with GeONet, which likely also contributed to the decrease in error. Moreover, as expected, the ambient-space error is much larger than the encoded-space error, meaning that the geodesics in the encoded space and ambient image space do not coincide. We mentioned this in the *LIMITATIONS* subsection 4.5 (initial submission) that encoding strategy seems to be necessary with such a dataset. In addition, we included the source code for the full MNIST experiments in the supplementary material for our revised submission.

Reviewer comment (Evaluation):
Why is error calculated in the encoded space for MNIST? It would be much more meaningful to calculate error in the ambient space. The error in the encoded space is difficult to understand and an unreproducible metric, particularly given the lack of code.

Author response:
(a) Regarding the MNIST experiment, in the revision, we retrained GeONet on the full MNIST data, splitting with 30,000 training data for $\mu_0$ and 30,000 training data for $\mu_1$ and randomly pairing them together in training. New testing error in the $L^1$ metric is reported in Table 4, consisting of both the error in the encoded space and ambient pixel space. Comparing the encoded-space error in the initial submission (Table 3 therein), we see that increasing the MNIST data indeed decreases the testing errors. Furthermore, the encoded data was also normalized prior to computing the geodesics with GeONet, which likely also contributed to the decrease in error. Moreover, as expected, the ambient-space error is much larger than the encoded-space error, meaning that the geodesics in the encoded space and ambient image space do not coincide. We mentioned this in the *LIMITATIONS* subsection 4.5 (initial submission) that encoding strategy seems to be necessary with such a dataset. In addition, we included the source code for the full MNIST experiments in the supplementary material for our revised submission.

Reviewer comment (Evaluation):
The L_1 error is also twice the Total variation distance. I find it somewhat strange to use TV here, usually Wasserstein or MMD are used, but I guess this is okay for toy problems.

Author response:
(c) We choose the $L^1$ error as our primary performance measure because it is a bounded metric. This way, we can interpret this error for different probability distributions across simulation setups. To accommodate your concern, we also include the errors (for the same simulation setups) in terms of $L^2$ distance and Wasserstein distance in Appendix J, Table 6 in the revised version. We see that the $L^2$ and Wasserstein errors exhibits similar patterns as the $L^1$ error, except that it is difficult to interpret the magnitude of the $L^2$ and Wasserstein errors.

Reviewer comment (Experiments):
Zero shot super resolution can be done by many modern methods. I’m not sure this is a useful experiment without significantly more experimentation and comparison. Perhaps on benchmarks that are constructed with a known map? (see [6]).

Author response:
(d) We highlight the zero-shot super resolution is consequence of the operator learning (from functions to functions) nature of our method and a feature that is not present in many other existing learned-based [2, 3, 5, 6] and (more traditional) optimization-based Wasserstein geodesic methods. As such, we do not claim the state-of-the-art of super-resolution problems in compute graphics. Given the limited time period for rebuttal, we do not intend to pursue in this direction.

Reviewer comment (Presentation):
“x-axis is the log of grid length in one dimension. This is somewhat confusing (also which log base?). Can this be replaced by the actual grid length?

Author response:
(e) The log base is the natural logarithm with based $e$. The reason we chose a log-log plot for our runtime comparison is that a log-log plot is capable of displaying a trend when there is a difference in order of magnitude, as exhibited by the linear patterns. According to your suggestion, we added a Figure 7 in Section 4.5 in revision (to replace Figure 7 in Section 4.4 in initial submission) to include plots for runtime vs. discretization length on both non-log and log-log scales.

Reviewer comment (Experiments):
While I can see the usefulness of amortizing W2 computation for faster inference, I do not think the comparison in 4.4 is fair. I would be interested to know for a comparable accuracy, how fast POT and GeONet are, as I assume the POT solver is extremely accurate. Or, similar to shown in MetaOT, a comparison showing GeONet provides a better initialization and speeds up convergence of Sinkhorn-based solvers.

Author response:
(f) Thanks for your suggestion! In the revision, we replaced Figure 7 in Section 4.4 by Figure 6 in Section 4.5, which compares the GeONet inference time and two versions of POT: one solves the Wasserstein geodesic problem to the machine precision, and the other solves the same problem with early stopping to the GeONet precision level. The left two panels in Figure 6 (revised version) display the runtime and discretization length (not on the log-scale), and the right panel displays the log-runtime and log-discretization-length to show the order of magnitude differences. The added experiment suggests that: the computational gain of GeONet over both versions of POT is more sizable for higher-dimensional and larger-size problems. The only exception for the reduced accuracy of POT beating GeONet is 1D Gaussian mixture with coarse discretization, a scenario that is not realistically interesting.

Reviewer comment (Experiments):
It could be helpful to include a comparison to non-amortized methods to make clear under what circumstances amortization becomes beneficial.

Author response:
(g) We addressed the problem in the reply to a previous question on comparison to non-amortization methods. Specifically, we added comparison with two more suggested learning-based methods: the rectified flow (RF) [2] and conditional flow matching (CFM) [3] in the revision.

Reviewer comment (Theory):
The method suffers from the curse of dimensionality as it currently requires fixed sized grids as input. It would be interesting to consider more general input forms.

Author response:
(3) Our method only suffers from $\textbf{input}$ curse-of-dimensionality in the branch network, and not $\textbf{output}$ in the trunk network. This suggests that there is no effect on inference time for fine grids of which the output is to be evaluated. Moreover, the input suffers the curse-of-dimensionality only partially: this could be mitigated by using alternative neural network frameworks, such as convolutional neural networks, for the branch networks.

Reviewer comment (Experiments):
Thank you for the additional comparisons. However, I believe it is much more meaningful to compare against the OT versions of both papers. Specifically, using the re-flow strategy of [3] or the minibatch strategy of [4,5].

Author response:
(b) In the revision, we added comparison with two more suggested learning-based methods: the rectified flow (RF) [2] and conditional flow matching (CFM) [3]. We run a similar 2D Gaussian mixture simulation setup in the discrete setting, constructing empirical distributions from sampled point clouds belonging from fixed densities. We use POT as the ground truth for comparing GeONet, RF and CFM. The $L^1$ estimation error for geodesic at time point $t = 0, 0.25, 0.5, 0.75, 1$ are reported in Table 3 and an estimated geodesic example is shown in Figure 4. There are a few observations we can draw from this new experiments. First, RF and CFM have 3-4 times comparably larger estimation errors than GeONet, except for the initial time $t = 0$, only because this initial data is given and learned directly for RF and CFM. GeONet is the only framework among the comparison which encapsulates the geodesic behavior to a considerable degree. While the other methods are suitable for point cloud representations, the geodesic behavior is entirely lost and highly inaccurate to ground truth. Second, RF and CFM have the same fixed resolution as the input probability distribution pairing, while GeONet can be smoothed out for estimating the density flows on higher resolution than the input pairing (cf. the third row in Figure 4).

Reviewer comment (Experiments):
Thank you for including error on the ambient space. The ambient error seems extremely high, even at times zero and one. Perhaps this is not 0-1normalized? I’m not sure how much we can say on this experiment with such a low-quality pre-trained autoencoder even on this small dataset. (Note: I did not have time to read the updated source code).

Author response:
(a) Regarding the MNIST experiment, in the revision, we retrained GeONet on the full MNIST data, splitting with 30,000 training data for $\mu_0$ and 30,000 training data for $\mu_1$ and randomly pairing them together in training. New testing error in the $L^1$ metric is reported in Table 4, consisting of both the error in the encoded space and ambient pixel space. Comparing the encoded-space error in the initial submission (Table 3 therein), we see that increasing the MNIST data indeed decreases the testing errors. Furthermore, the encoded data was also normalized prior to computing the geodesics with GeONet, which likely also contributed to the decrease in error. Moreover, as expected, the ambient-space error is much larger than the encoded-space error, meaning that the geodesics in the encoded space and ambient image space do not coincide. We mentioned this in the *LIMITATIONS* subsection 4.5 (initial submission) that encoding strategy seems to be necessary with such a dataset. In addition, we included the source code for the full MNIST experiments in the supplementary material for our revised submission.

Reviewer comment (Evaluation):
Thank you for the revision of Figure 6. I am more convinced by this, however I would like to ideally see the empirical error against time for these methods, not just a few different settings of the numerical error since a method like GeONet does not really have such a setting. I think this would add context to when GeONet is useful and can provide meaningful amortization.

Author response:
(f) Thanks for your suggestion! In the revision, we replaced Figure 7 in Section 4.4 by Figure 6 in Section 4.5, which compares the GeONet inference time and two versions of POT: one solves the Wasserstein geodesic problem to the machine precision, and the other solves the same problem with early stopping to the GeONet precision level. The left two panels in Figure 6 (revised version) display the runtime and discretization length (not on the log-scale), and the right panel displays the log-runtime and log-discretization-length to show the order of magnitude differences. The added experiment suggests that: the computational gain of GeONet over both versions of POT is more sizable for higher-dimensional and larger-size problems. The only exception for the reduced accuracy of POT beating GeONet is 1D Gaussian mixture with coarse discretization, a scenario that is not realistically interesting.

Reviewer comment (Experiments):
I am not sure about whether this net can be generalized to unseen distribution Wasserstein distance calculations, especially the ones that are very different from the training data. Since many times we don't know whether the distributions in the application field are similar to the ones in training, and if we just calculate the OT question, we can get the accurate results although it will be slow. But I think it is hard to ensure such a learned neural operator is generally applicable.

Author response:
(1) We thank you for pointing the OOD retraining issue. To clarify, our OOD experiments on 1D and 2D Gaussian mixtures were tested upon data that was not seen in training, but still belonging to the same family of Gaussian mixture distributions with different set of variance parameters from training. Recall that the training dataset size for 1D Gaussian mixtures is 20,000 pairs and for 2D Gaussian mixtures is 5,000 pairs. In the current Gaussian mixture setting, if we were to retrain the experiments on new data with different variance parameters of mixture Gaussian components for GeONet to achieve similar $L^1$ errors, we empirically observe around 500 new pairings are needed.

Reviewer comment (Writing):
The readability is too bad. Even though I am familiar with many of the maths inside, the symbols are not defined precisely, making it very difficult to understand. Much more detail of the proposed algorithms and backgrounds needs to be added.

Author response:
(3) Thanks for your suggestion. We revised the paper to address questions raised by all reviewers. We hope the rebuttal is further clarifying and that the revised version of the paper is satisfactory.

Reviewer comment (Experiments):
How much data is needed for retraining for OOD cases?

Author response:
(1) We thank you for pointing the OOD retraining issue. To clarify, our OOD experiments on 1D and 2D Gaussian mixtures were tested upon data that was not seen in training, but still belonging to the same family of Gaussian mixture distributions with different set of variance parameters from training. Recall that the training dataset size for 1D Gaussian mixtures is 20,000 pairs and for 2D Gaussian mixtures is 5,000 pairs. In the current Gaussian mixture setting, if we were to retrain the experiments on new data with different variance parameters of mixture Gaussian components for GeONet to achieve similar $L^1$ errors, we empirically observe around 500 new pairings are needed.

Reviewer comment (Theory):
I don't quite understand the justification of "no need for retraining for new input"; I think this is not applicable for all usages.

Author response:
(2) We apologize for the confusion. This sentence was meant to compare with traditional OT solvers based on optimization (such as iterative Bregman projections [1]) and PINNs that learns the solution of a given (i.e., single) PDE (not the solution operator as a mapping from the boundary conditions to PDE solutions). GeONet is an operator learning method that we do not need to retrain the OT dynamics for new boundary conditions. To accommodate your concern, we changed "no need for retraining for new input" to "operator learning" in Table 1 in the revision.

Reviewer comment (Writing):
The current presentation is too rough, making it very difficult for readers to understand all the details. The paper should be revised with clearer explanations and better writing.

Author response:
(3) Thanks for your suggestion. We revised the paper to address questions raised by all reviewers. We hope the rebuttal is further clarifying and that the revised version of the paper is satisfactory.