Paper: Diffeomorphic Mesh Deformation via Efficient Optimal Transport for Cortical Surface Reconstruction (ICLR-2024)

Reviewer comment (Novelty):
The novelty of the idea is limited -- optimizing meshes with CD and/or earth mover distance (EMD) between sampled point clouds is not new (for instance [1]). Replacing EMD with another distributional distance seems more or less some pure trial-and-error endeavor. Indeed, the paper [2] has shown that optimization with SWD on point clouds is beneficial. It is therefore not surprising that SWD can be used for mesh optimization.

Author response:
We respectfully disagree with Reviewer in two aspects: 

**First of all**,  while acknowledging the prior exploration of optimizing meshes with Chamfer Distance (CD)/Earth Mover's Distance (EMD) between sampled point clouds, we argue that the computation of EMD as an optimization metric is not practical for high-resolution meshes like the cortical surface. We refer the reviewer to the issue: Earth Mover's Distance (EMD) raised on pythorch3D GitHub and its answer by Georgia Gkioxari ([link](https://github.com/facebookresearch/pytorch3d/issues/211)). To the best of our knowledge, the computation of the exact transport map between two point clouds with $n$ points has a polynomial complexity and is not affordable for $n > 20,000$ [5]. The reason [1] successfully employs EMD is due to their experiments on the ShapeNet dataset, which is representable by a few thousand points. In contrast, the complex surface topology of the cortical surface requires mesh approximations with several hundred thousand points. 

**Secondly**, we do not replace EMD with another distributional distance due to merely a result of trial-and-error. Specifically, we aim to derive a valid, fast, and scalable metric for high-resolution mesh deformation. While EMD is a valid metric, its applicability to high-resolution meshes is limited, as previously mentioned. On the other hand, Chamfer distance (CD) is faster, but it is not a valid distance since it does not satisfy the identity of indiscernibles, i.e., when CD is 0, two point clouds are not necessarily the same, thus leading to the local minimum issue presented in our paper. Furthermore, as shown in Section 4.2.2 and Figures 3 in our paper, CD (and its variants) is not as scalable as our proposed metric. To the best of our knowledge, we are the first to propose a metric that is valid, fast, and scalable in learning high resolution shape deformation.

Reviewer comment (Writing):
The presentation of the paper needs to be improved. There are many notations not explained (e.g., the # operator in Eqn. 5).

Author response:
We respectfully disagree with reviewer for two reasons. 

**First**, our probabilistic interpretation, especially our theorem 1, is necessary in the context of this paper. This theorem lays a robust foundation, motivating the representation of meshes as probability measures and justifying the use of Sliced Wasserstein Distance (SWD) as our optimization metric. In fact, from independent point of view, Reviewer **4B9Y** also acknowledges our theorem as "novel and promising" and "suitable for practical applications". 

**Secondly**, although EMD has already used for mesh optimization, to the best of our knowledge, there is no theoretical justification for its usage. Furthermore, as presented in Q1, EMD is not a suitable metric for our use case with high-resolution mesh.

Reviewer comment (Experiments):
While the proposed method seems general, the scope of this paper is quite limited to cortical structures. The authors may consider trying the proposed method on general shapes and compare the results on common benchmarks (if any).

Author response:
The operator $\sharp$ used in Equation 5 has been previously defined in Section 1 under 'Notation.' Allow us to restate the notation here: *We denote $\theta \sharp \mu$ as the push-forward measures of $\mu$ through the function $f:\mathbb{R}^{d} \to \mathbb{R}$ that is $f(x) = \theta^\top x$*. For further notation, we encourage the reviewer to refer back to this section in the main text. 

We apologize for any confusion regarding notation that may have arisen during the review process. We would greatly appreciate it if the reviewer could point out any notations that are not clearly explained, enabling us to enhance the clarity and comprehensibility of our final draft.

Reviewer comment (Evaluation):
I am not sure if the computation of chamfer distance (CD) should be slower than that of SWD. Some paper argues that the complexity of SWD is similar to CD [2]. And by using KD-tree for nearest neighbor retrieval or some recent methods [3], one should be able to compute CD much faster. I would like the authors explain the speed difference between CD and SWD shown in Fig. 3.

Author response:
We chose to demonstrate our method using cortical surface reconstruction due to its relevance in neuroscience and neuroimaging. The high resolution of the reconstructed mesh also provides an excellent opportunity to highlight the efficiency of the sliced Wasserstein distance in terms of both scalability and precision. As for its applicability to other datasets, our method is versatile and can be extended to the reconstruction of other organs, including heart [8] and lung [9] reconstruction. Due to the time constrain of discussion process, we are unable to give a comprehensive comparisons on other benchmark. We plan to address this in our future work, expanding the scope of our method's applications and evaluations.

Reviewer comment (Experiments):
There are some improved variants of CD, such as [4]. Just out of curiosity: how do they perform compared to SWD?

Author response:
In Fig. 3,  we compare the computation of SWD with the implementation of Chamfer distance implemented in Pytorch3D [6], which is employed by all of our competing methods. It is important to note that the Pytorch3D implementation accelerates nearest neighbor (NN) retrieval using custom CUDA kernels, as detailed in Figure 1 of [6]. This implementation takes advantage of parallel computation capabilities on GPUs, a key factor in its efficiency. 

Regarding KD-tree, the time and complexity of CD using KD-tree is theoretically fast, i.e. $\mathcal{O}(n \log n)$, quite faster than $\mathcal{O}(n^2)$ complexity of the vanilla CD. However, current implementations of KD-trees do not support parallel processing, which limits their scalability, particularly with a large number of points. Additionally, the efficiency of KD-trees diminishes with increasing dimensionality, making them less suitable for GPU computation of high-dimensional data. 

To validate our claims, we conducted a comparative time analysis of SWD, CD using Pytorch3D, and CD with KD-trees. Since PyTorch doesn't natively support KD-Tree, so we use an external library like [scipy](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.html) for this purpose. In low dimensional points, CD-KDtree's speed is comparable to CD-Pytorch3D, but it is still less scalable than SWD. In high dimensional points, e.g. oriented varifold, the speed of CD-KDtree is much slower than both CD-Pytorch3D and SWD. We provide the comparison visualization and our implementation of KD-Tree in the [anonymously link](https://drive.google.com/drive/folders/143lmPXeByrvgsalWxg-r_suOoWagA0Bn).

Reviewer comment (Theory):
The proposed method represents a mesh as a distribution in the position-orientation space, and the optimal transportation map is carried out in this space. It should be explained what the cost function is in this context.

Author response:
The cost function defined on the position-orientation space is delienated at Appendix D. Specifically, the cost is computed as the Monte Carlo sampling of sliced Wasserstein distance between the predicted mesh $\mathcal{\hat{M}}$ and the target mesh ${\mathcal{M}^\ast}$ represented as probability measures $\mu^{\mathcal{\hat{M}}}$ and $\mu^{\mathcal{M}^\ast}$, respectively: $$\mathcal{L}(\mathcal{\hat{M}}, \mathcal{M}^\ast) = \widehat{SW}\_p^p (\tilde{\mu}^{\mathcal{\hat{M}}},\tilde{\mu}^{\mathcal{M}^\ast}) =  \frac{1}{L}\sum\_{l=1}^L \text{W}\_p^p (\theta\_l \sharp \tilde{\mu}^{\mathcal{\hat{M}}},\theta\_l \sharp \tilde{\mu}^{\mathcal{M}^\ast}),$$

where $\text{W}_p^p (\theta_l \sharp \tilde{\mu}^{\mathcal{\hat{M}}},\theta_l \sharp \tilde{\mu}^{\mathcal{M}^\ast})$ is the Wasserstein-$p$ [1] distance between $\tilde{\mu}^{\mathcal{\hat{M}}}$ and $\tilde{\mu}^{\mathcal{M}^\ast}$, and $\theta\_1,\ldots,\theta\_L$ ($L$ is the number of projections) are drawn i.i.d from $\mathcal{U}(\mathbb{S}^{d-1})$. We empirically set the hyperparameter by tuning on a small subset of OASIS dataset. To be specific, we set $L=100, p=2$. In terms of oriented varifold, for each support, we concatenate the barycenter of vertices of the face and the unit normal vector as a single vector in $\mathbb{R}^6$.

Reviewer comment (Theory):
The representation of the mesh depends on the position and the orientation of the mesh, which means the Wasserstein distance varies if one mesh is transformed by a rigid motion.

Author response:
It is possible that the Wasserstein distance changes when one mesh is transformed under rigid motion. Therefore, during training we do not apply any transformation to keep Wasserstein distance stable.

Reviewer comment (Theory):
What is the cost function defined on the position-orientation space?

Author response:
The cost function defined on the position-orientation space is delienated at Appendix D. Specifically, the cost is computed as the Monte Carlo sampling of sliced Wasserstein distance between the predicted mesh $\mathcal{\hat{M}}$ and the target mesh ${\mathcal{M}^\ast}$ represented as probability measures $\mu^{\mathcal{\hat{M}}}$ and $\mu^{\mathcal{M}^\ast}$, respectively: $$\mathcal{L}(\mathcal{\hat{M}}, \mathcal{M}^\ast) = \widehat{SW}\_p^p (\tilde{\mu}^{\mathcal{\hat{M}}},\tilde{\mu}^{\mathcal{M}^\ast}) =  \frac{1}{L}\sum\_{l=1}^L \text{W}\_p^p (\theta\_l \sharp \tilde{\mu}^{\mathcal{\hat{M}}},\theta\_l \sharp \tilde{\mu}^{\mathcal{M}^\ast}),$$

where $\text{W}_p^p (\theta_l \sharp \tilde{\mu}^{\mathcal{\hat{M}}},\theta_l \sharp \tilde{\mu}^{\mathcal{M}^\ast})$ is the Wasserstein-$p$ [1] distance between $\tilde{\mu}^{\mathcal{\hat{M}}}$ and $\tilde{\mu}^{\mathcal{M}^\ast}$, and $\theta\_1,\ldots,\theta\_L$ ($L$ is the number of projections) are drawn i.i.d from $\mathcal{U}(\mathbb{S}^{d-1})$. We empirically set the hyperparameter by tuning on a small subset of OASIS dataset. To be specific, we set $L=100, p=2$. In terms of oriented varifold, for each support, we concatenate the barycenter of vertices of the face and the unit normal vector as a single vector in $\mathbb{R}^6$.

Reviewer comment (Theory):
Why use the position-orientation space to represent the mesh measure? Why not just use position space?

Author response:
By relying on position-orientation space, we can employ oriented varifold theory, which gives us a theoretical bound [2] for approximating discrete mesh. Furthermore, our experimental results demonstrate that using position-orientation space yields better performance compared to position space only (see Sec. 4.3. Ablation). Therefore, in this paper, we rely on position-orientation space to represent the mesh measure.

Reviewer comment (Theory):
Can we say something about the regularity of the optimal transportation map? Is it diffeomorphic?

Author response:
We are assuming that you are asking whether the optimal transport map is a diffeomorphism. It is important to clarify that our approach does not seek to establish an optimal transport map between two manifolds. Instead, we utilize the sliced Wasserstein distance to determine the OT cost. This method primarily serves as a loss function, aiding in quantifying the discrepancy between the predicted and target mesh. Consequently, given the focus of our methodology on sliced Wasserstein distance for loss computation, we cannot definitively ascertain whether the OT map resulting from this process is a diffeomorphism. We are happy to discuss more if Reviewer is still not satisfied with our answer.

Reviewer comment (Writing):
The authors depict a mesh as an oriented varifold, a pivotal concept in this paper. However, for enhanced clarity and accessibility, it's better to provide a concise introduction to oriented varifolds, even though the foundational idea is rooted in earlier works. This brief overview will promote a fundamental understanding of oriented varifolds and their importance in our study, improving the paper's readability.

Author response:
We thank Reviewer for the valuable suggestion. In response, we have included a concise overview of oriented varifolds in Appendix A. Additionally, a detailed derivation of mesh as an oriented varifold can be found in Section 3.1. We believe that our current information about oriented varifold sufficiently conveys our concepts and remain easily comprehensible for scholars familiar with the field.

Reviewer comment (Experiments):
Is the initial surface in your method the same as the initial surfaces used in competing methods, which are also extracted from the white matter segmentation mask of the brain MRI image?

Author response:
In the process of generating initial surfaces from MRI images, each of our competing methods uses a different initial surface. The initial surface in our approach, as well as in CortexODE, is derived from the segmentation of MRI images. Conversely, CFPP and Vox2Cortex commence with a pre-defined mean shape template. DeepCSR, on the other hand, employs an implicit-based methodology and thus does not rely on a template for initiation.

Reviewer comment (Theory):
How about the topological information of the initial surface and the ground truth surfaces? Are they all genus-0 surfaces? Does the DDOT framework maintain the topological information during the deformation process?

Author response:
Yes, the topological information of initial surface and the groundtruth surface is the same, which is genus-0 surfaces. The utilization of the diffeomorphic deformation framework, DDOT, ensures the preservation of this topological information throughout the deformation process.

Reviewer comment (Writing):
It is better to give a brief introduction of oriented varifold for paper's readability.

Author response:
We thank Reviewer for the valuable suggestion. In response, we have included a concise overview of oriented varifolds in Appendix A. Additionally, a detailed derivation of mesh as an oriented varifold can be found in Section 3.1. We believe that our current information about oriented varifold sufficiently conveys our concepts and remain easily comprehensible for scholars familiar with the field.