Paper: DeepEMD: A Transformer-based Fast Estimation of the Earth Mover’s Distance (ICLR-2024)

Reviewer comment (Novelty):
Lack of Novelty: The problem of point-cloud registration is very well-studied, with numerous learning-based methods proposed over the years. Some of these also include attention mechanisms to do the matching. See the methods listed here: https://github.com/XuyangBai/awesome-point-cloud-registration

Author response:
W1,W2 : We are not solving a point cloud registration task. While we evaluate on 3D point clouds,
DeepEMD does not assume anything about the input data and does not process 3D point clouds in
any special way, in particular does not try to “match” according to a local appearance similarity.
The main goal is to approximate EMD between two distributions which is ’similar’ to point cloud
matching when working with point clouds as uniform distributions. In principle, DeepEMD can be
utilized as a fast surrogate for EMD wherever EMD is applicable, including tasks such as single-view
point cloud generation or point cloud completion.

Reviewer comment (Experiments):
Comparative Evaluation: The paper does not compare its method to state-of-the-art registration approaches. Given that point-cloud registration is a long-studied problem, there are likely several strong baseline methods against which the proposed model should be evaluated.

Author response:
W1,W2 : We are not solving a point cloud registration task. While we evaluate on 3D point clouds,
DeepEMD does not assume anything about the input data and does not process 3D point clouds in
any special way, in particular does not try to “match” according to a local appearance similarity.
The main goal is to approximate EMD between two distributions which is ’similar’ to point cloud
matching when working with point clouds as uniform distributions. In principle, DeepEMD can be
utilized as a fast surrogate for EMD wherever EMD is applicable, including tasks such as single-view
point cloud generation or point cloud completion.

Reviewer comment (Experiments):
Generalization may be an issue: While the paper mentions that the model generalizes well to larger point clouds during inference than those seen during training, most of the experiments are carried out on datasets consistent of synthetic 3D models. What happens when this is applied to real LiDAR scans or depth images?

Author response:
W3. As discussed above, there is no inherent limitation in the DeepEMD architecture and we
only choose to evaluate on point clouds as they represent uniform probability distribution and
shape/object based point clouds datasets are easily available (number of samples and reasonable
cardinality). We believe that the model could generalize well to these other datasets (LiDAR, etc.)
but we would need it to scale to larger point clouds (or support sizes/ number of histogram bins).
Further, test time generalization for a model trained on shapes might not be accurate enough
for very different datasets but we can always train on those new datasets. The novelty in our
method is to regress for the matching/coupling rather than the distance. We believe that this
ground truth supervision with the matching allows the excellent generalisation capabilities to the
model as compared to methods which regress for the distance itself. In particular, in the SetVAE
application, the distribution of the generative model changes over iterations, hence, the population
EMD also changes since one marginal (decoder output) is changing. We observe that the behaviour
of DeepEMD is excellent during the whole process. This may be due in part to the noise we inject
to the data samples during training of DeepEMD (explained in Sec. 4.1 and Sec. B.1). Note that
we designed this data augmentation scheme before running any SetVAE related application, and
did not revisit it, which makes us confident it is actually extremely robust.

Reviewer comment (Experiments):
Scalability Issues: Attention mechanisms, particularly in the context of a large number of outputs, can be memory-intensive. Given that point clouds can be large (millions of points), it might be difficult to scale the proposed approach to such cases.

Author response:
W4. The same limitations also arise with other methods like Hungarian, Sinkhorn, EMD, etc.
We agree that investigating this line of work will be very useful. Concurrent research on attention
models for increasing context length, reducing attention complexity, etc. can be applied to our
method and bring in the gains seamlessly.

Reviewer comment (Experiments):
While the paper compares DeepEMD with other algorithms, it fails to compare against other state of the art works that compute EMD in terms of speed and accuracy.

Author response:
We experimented with several variants of Sinkhorn, including sinkhorn log, sinkhorn stabilized,
greenkhorn, etc. However, we discovered that they either exhibited significantly slower performance
compared to the vanilla variant, frequently failed to converge to a satisfactory optimal transport
matrix within a finite timeframe, or sometimes exhibited both issues.

Reviewer comment (Experiments):
The authors talk about several works like CD, DPDist etc. in related works section which are faster than the proposed work. However, authors fail to compare their work with these works.

Author response:
We do provide comparisons with CD and for the other related works we provide qualitative
reasoning behind why DeepEMD is superior.

Reviewer comment (Evaluation):
Figure 8 shows DeepEMD to be performing 100x faster than Hungarian algorithm but has almost half the accuracy against the same algorithm. How much does DeepEMD sacrifice accuracy to achieve speed? Can you provide a more detailed speed vs accuracy tradeoff?

Author response:
Plese refer to the answer to W2.

Reviewer comment (Experiments):
Is there a specific reason to use only 100 iterations for comparison? Can a better graph be provided that compares the methods over iterations?

Author response:
We can observe that Sinkhorn with 5 iterations is more expensive than DeepEMD and
performs worse than DeepEMD. While further increasing the number of iteration might lead to a
better convergence but the cost-performance trade-off in this regime is not very relaevant for our
comparisons.

Reviewer comment (Experiments):
The authors claim that their proposed approach to compute EMD approximately is significantly faster than existing techniques. They primarily consider two baselines, the well-known Hungarian method which computes EMD exactly and the well-known method that used entropic regularization to compute approximate EMD (Cuturi et al 2013) using matrix scaling / Sinkhorn Knopp algorithm. It is true that the proposed method is significantly faster than the Hungarian method and the Sinkhorn implementation tested here. However, there are a number of works in the literature which are either variants of the Sinkhorn method that have a better convergence than the original Sinkhorn method [C] or have investigated faster algorithms for computing EMD [D, E] that do not require any learning and therefore do not have the issue of overfitting/failure to generalize. A comparison with methods such as [C, D, E] is needed (perhaps there are a few other works in the literature) to make a stronger case for the proposed method.

Author response:
We experimented with several variants of Sinkhorn, including sinkhorn log, sinkhorn stabilized,
greenkhorn (ref [C], POT provides an implmentation), etc. However, we discovered that they ei-
ther exhibited significantly slower performance compared to the vanilla variant, frequently failed
to converge to a satisfactory optimal transport matrix within a finite timeframe, or sometimes
exhibited both issues. Further [D] has a similar complexity as Sinkhorn. Also note that [D] demonstrates that ”regularized
variations with large regularization parameter will degrade the performance in several important
machine learning applications, and small regularization parameter will fail due to numerical sta-
bility issues with existing algorithms.” which applies to our case for using Sinkhorn and variants
for benchmarking. The methods in [E] is limited to the Wasserstein-1 metric and doesn’t generalize for Wasserstein-p.

Reviewer comment (Theory):
The title/abstract is not particularly appropriate. The proposed approach is aimed at computing EMD-based distance between 3D point clouds of geometric shapes. However, the title of the paper makes it sound like a generic method to compute EMD between arbitrary distributions has been proposed. Moreover, the approximation factor in the proposed method is somewhat adhoc and not well analyzed, as some of the existing methods in the literature.

Author response:
Although we choose to evaluate on 3D point clouds, note that DeepEMD does not assume
anything about the input data and does not process 3D point clouds in any special way but works
directly with the raw features. Thus the limitation to 3D objects is not significant.The main goal
is to approximate EMD between two distributions which is ’similar’ to point cloud matching when
working withpoint clouds as uniform distributions. In principle, DeepEMD can be utilized as a
fast surrogatefor EMD wherever EMD is applicable.

Reviewer comment (Presentation):
More 3D visualization can improve the representation. The paper do not contain any 3D visualization for reconstruction, generation, etc. This makes the paper lack intuitive qualitative comparison.

Author response:
W2. This is indeed a very good point, such representation will be added to the final version.

Reviewer comment (Experiments):
It will be much more convincing if the authors adopt DeepEMD as a loss to previous generative models (e.g. single-view point cloud generation, point cloud completion) and report the performance compared to the baseline models.

Author response:
W3. We do show that DeepEMD (see Sec 4.4) as a loss works well for training an  auto-encoder as well as a variational auto-encoder. In principle, DeepEMD can be utilized as a fast surrogate for EMD wherever EMD is applicable, including tasks such as single-view point cloud generation or point cloud completion. In addition to good approximation to the EMD and its gradient, we show the most straight-forward evaluation by training  auto-encoders which are completely based on the distance metric and thus appropriate in assessing the quality of DeepEMD as a surrogate metric.

Reviewer comment (Theory):
I think that the argmax in Eq.(10) is not differentiable, how did you solve that?

Author response:
Q1. We only use eq. 10 during inference and thus it doesn't need to be differentiable.

Reviewer comment (Experiments):
Will DeepEMD still be efficient with a large point number?

Author response:
Q2. We do show in Sec 4.3 the test time performance with point clouds as large as 8196 and observe that it performs excellently. With transformer architectures which are either faster or scale up to larger context sizes, we should be able to scale up DeepEMD.

Reviewer comment (Reproducibility):
How long dose it take to train DeepEMD, and how much GPU do you use?

Author response:
Q3. In general, we typically need to train a single DeepEMD model for a given dataset (note that this step may not be needed given the generalization capabilities of the model). Further, training DeepEMD takes about 4 mins for an epoch of training on a single NVidia RTX 3090 and 16 GB memory with a batch size of 16 (we train for 75 epochs). During inference, it needs 0.8 GB GPU memory. Also note that the inference and training times can be improved by 5-8x with efficient attention implementations, eg. flash attention. When used as a loss, the parameters of DeepEMD are fixed and gradients for the same are not needed, thus its memory footprint is very limited.