GeONet: a neural operator for learning the Wasserstein geodesic

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

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

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.

(1) Theoretical analysis. We acknowledge that this paper does not deal with the theoretical aspect of statistical guarantees. We expect that some theoretical work can be done under certain assumptions given the well-known regularity of the (static and dynamic) OT problems. For instance, if the initial and terminal distributions are smooth, then Caffarelli's global regularity ensures that the static OT map between the two distributions is also smooth (to a less degree). Then with reasonable assumptions on the convexity and smoothness of the optimal potential for pushing mass along the Wasserstein geodesic, we expect that there is a stability between the GeONet loss function and the underlying Wasserstein geodesic, leading to a reasonable generalization error bound that is useful to control the predictive risk of geodesic in the test time. Given the experimental nature of this paper introducing a new concept of Wasserstein geodesic from an operator learning perspective, we leave the theoretical analysis to the future work.

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

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

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

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

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

[1] Benamou, Carlier, Cuturi, Nenna, Peyré. (2015) Iterative Bregman Projections for Regularized Transportation Problems. $*SIAM Journal on Scientific Computing.*$

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

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

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

(3) Our method only suffers from $**input**$ curse-of-dimensionality in the branch network, and not $**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.

[2] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. $*ICLR*$ 2023.

[3] Alexander Tong, Nikoly Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, Yoshua Bengio. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. 2023.

[4] Solomon et al. (2015) Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. $*ACM Trans. Graph.*$

[5] Julien Lacombe, Julie Digne, Nicolas Courty, and Nicolas Bonneel. Learning to generate Wasserstein barycenters, $*Journal of Mathematical Imaging and Vision*$ 2021.

[6] Brandon Amos, Giulia Luise, Samuel Cohen, and Ievgen Redko. Meta Optimal Transport, $*ICML*$ 2023.

(1) Comparisons and related work.

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

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

(2) More on experiments and comparisons.

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

(b) Convolutional Wasserstein Barycenters (CWB) is most suitable for computing the OT maps and dynamics on meshes over large domains in the continuous (non-point cloud) setting, and favorable timing and numerics over other algorithms (linear programs, Sinkhorn, iterative projections) have been reported in [4]. Given the discretization size in our problems, we used off-the-shelf CWB solver in the standard POT Python library to compute the Wasserstein geodesic (since the barycenter is a special of $t = 0.5$). There may be numerical issue when the regularization parameter is smaller than the resolution of the domain discretization because the convolution kernel is ill-conditioned; but similar issues would also occur for other entropic regularized algorithms with low regularization values. In our experiments, it appears that the ``ground truth" geodesics computed by CWB is entirely reasonable (cf. Figures 3, 4, 7 in the revised version). We presume you are referring to the errors between GeONet and the CWB framework. Since CWB acts as the best 2D geodesic calculator in the continuous setting, it is reasonable for comparison to GeONet, yielding such errors.

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

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