We sincerely thank the reviewer for acknowledging the merit of our idea, the contribution of OT in comparing and embedding distributions within the tree structure, the solid analysis of extending the ℓ2-CPCC framework to OT, and the strong experimental improvements demonstrated under various Wasserstein distance computations. Let us answer your questions and concerns below:

In Table 1 (Sec 2.2), I do not think that these complexities are straightforward to derive. Can the authors point to the reasoning to their derivation? For instance, EMD's complexity is tipically $O(n^3\log n)$ instead of $O(n(d + n^2\log n))$.

Sorry for confusion. We fixed some minor mistakes in Table 1, and let us walk you through all time complexities listed in the updated Table 1.

For simplicity, in the time complexity analysis, we assume each class conditioned cluster has the same size $n$, which implies batch size $b = nk$.

For $\ell_2$, given one batch of batch size $b$, computing Euclidean centroids for all classes requires scanning through all $d$-dimensional data points once, which requires $O(bd) = O(nkd)$. We have at most $k$ classes in one batch, so computing the pairwise distances between class centroids will cost $O(k^2d)$. The total cost for $\ell_2$ is $O((nk+k^2)d)$.

For OT methods, we have $k$ class conditioned clusters, so there are $O(k^2)$ pairs of matrices as input for all OT methods. For each pair, OT methods computation can take from linear time to cubic time w.r.t $n$. For example, each FastFT computation is $O(nd)$ as we shown in Theorem 3.2, so the total computation time is $O(k^2nd)$. In general, the computation time for all OT methods will be $O(k^2 \cdot \text{OT-time})$.

Let us elaborate the “OT-time” part.

For EMD, EMD’s complexity is typically $O(n^3\log n)$ in the past literature, mainly focusing on the influence of the number of samples $n$. However, in our setting, because $d >> n$, we want to emphasize the effect of $d$. In the expression of the time complexity for EMD, $O(n^3\log n)$ comes from solving the linear programming problem, and the $O(nd)$ term comes from the computation of the objective function $\left<\mathbf{P}, \mathbf{D}\right>$. Due to a sparse solution of linear programming, there are only $O(n)$ non-zero values in $\mathbf{P}$, we need to compute $O(n)$ $d$-dimensional pairwise distances in $D$, so the dot product takes $O(nd)$ in total. The total complexity of EMD is $O(n(d + n^2\log n))$.

To get the approximated flow matrix, Sinkhorn mainly depends on the iterative matrix scaling algorithm. For each $I$ iteration, we normalize the row and column of a matrix of size $O(n^2)$, so the complexity for running this iterative algorithm is $O(n^2I)$. Unlike in EMD, the approximated flow matrix $P$ derived from Sinkhorn is not sparse. Therefore, the dot product of the objective function will now take $O(n^2d)$ in total. The total complexity of Sinkhorn is $O(n^2(d+I))$.

For SWD, the first step is to map $d$-dimensional points to 1D by computing dot products of the data points with random vectors for $p$ repetitions, which involves $O(npd)$. For each of the $p$ projection, we need to sort the projected $n$ values, and compute the cumulative sum of the sorted values. Due to sorting, this step takes $O(np\log n)$. The total complexity of SWD is $O(np(d + \log n))$.

For FT, the most expensive part is the tree construction using QuadTree algorithm. [1] shows that we can build this tree in $O(nd\log (d\Phi))$. After the tree is constructed, by lemma 3.1 in [1], FT requires $O(n(d + \log (d\Phi)))$ for running the bottom-up algorithm and computing the objective function, which is less expensive than the tree construction. The total complexity of FT is therefore $O(nd\log (d\Phi))$.

For FastFT, please refer to the proof of Theorem 3.2.

For TWD, we follow the Algorithm 1 in [2]. The first step is to construct a tree. Instead of using QuadTree, to save some computation, we can leverage the same technique in FastFT to use our augmented tree, so this step is $O(1)$. Now for each of $I$ iteration, for supervised learning, we need to compute all of the entries in the pairwise distance matrix $D$, which takes $O(n^2d)$. Therefore, even when tree structure is known, TWD will take $O(n^2dI)$.

We included this detailed discussion of time complexities in App.I in the latest version.

[1] Backurs, Arturs, et al. "Scalable nearest neighbor search for optimal transport." International Conference on machine learning. PMLR, 2020.

[2] Yamada, Makoto, et al. "Approximating 1-wasserstein distance with trees." arXiv preprint arXiv:2206.12116 (2022).

Run OT-CPCC with Bures-Wasserstein or Mixture-Wasserstein

Yes, this is a natural extension of the $\ell_2$ CPCC method. We first restate the expression of both metrics. Bures-Wasserstein is defined as $\mathcal{W}_{2}(P,Q)^{2} = \|\mu_P - \mu_Q\|_2^2 + \text{Tr}(\Sigma_P + \Sigma_Q - 2(\Sigma_P^{1/2}\Sigma_Q\Sigma_P^{1/2})^{1/2}),$

and the objective function of Mixture-Wasserstein [1] is

$\min \sum w_{PQ} \mathcal{W}_{2}(P,Q)^{2}.$

We can observe that both metrics depend on the covariance term which requires at least $O(nd^2)$ computation, and Mixture-Wasserstein requires additional time for fitting GMM. Taking the mean of a class conditioned cluster requires $O(nd)$. All other OT methods only have a linear $d$ term in the big-O term too. In practice, $n$ is the size of a class conditioned cluster. In CIFAR100, if batch size = 1024, $n$ is approximately $1024/100 \approx 10$ if all classes have the same size. If we use ResNet18 as the backbone, $d = 512$ which is typically much larger than $n$. Therefore, the factor of $d$ from computing the covariance matrix can be very expensive in our case. To verify this statement, we computed the running time for CIFAR100 (batch size = 128) for different methods below, for one batch update. The results are averaged over three random seeds.

We can observe that Bures-Wasserstein is much slower than both $\ell_2$ and EMD. Therefore, for scalability concern, all covariance based metrics do not fit our problem setup very well.

[1] Delon, Julie, and Agnes Desolneux. "A Wasserstein-type distance in the space of Gaussian mixture models." SIAM Journal on Imaging Sciences 13.2 (2020): 936-970.

While I got what the authors are trying to express with prop. 3.1, and the result is true if stated correctly, I think the authors could rephrase it for better clarity.

Thanks for pointing out this writing issue. We will take your suggestions and revise the statement as “EMD reduces to $\ell_2$ between means of Gaussian with the same covariance.” The change is reflected in line 181-183 of the latest version.

In table 4, the objective "Flat" is not explained. From the overall context, I guess this corresponds to learning with no structure tree, but this is not clear to me.

Yes, “Flat” corresponds to standard cross entropy learning without using any hierarchical information. We now added one extra line for clarification in line 332.

For better uniformity and clarity, the authors might want to change references to the EMD as a distance in favor of 2−Wasserstein distance, which is what is actually used in the paper. For instance, in most tables the authors use EMD and terms involving the Wasserstein distance (TWD or SWD).

Thanks for your suggestion! We understand your concern about consistency of naming. As we mentioned in App.B in the proof of Prop.3.1, EMD is equivalent to 1-Wasserstein distance, using $\ell_2$ as the ground metric, for discrete measures. The reason why we use the name of EMD is because in our paper, we restrict our discussion only on discrete measures where the distance metric can be derived from a linear programming problem in Eq.3, while the name of 1-Wasserstein distance can be applied to a more general case where $\mu$ and $\nu$ can be continuous. With the linear programming setup in Eq.3, we can explain gradient analysis in Sec 3.3 easier.

In section 3.3, when authors are comparing the gradients of ℓ2−CPCC and OT-CPCC, I think the authors can make 2 further links, if they agree with the following discussion.

We appreciate your insightful comments! We added the discussion of the relationship between MMD and $\ell_2$ in line 43-44.

However, the equivalence between Sinkhorn with $\epsilon \to \infty$ and $\ell_2$ does not hold in general. Assume we have two matrices $X \in \mathbb{R}^{n \times d}$ and $Y \in \mathbb{R}^{m \times d}$ representing two class conditioned features. When $\epsilon \to \infty$, Sinkhorn’s flow matrix are uniformly distributed with each entry to be $1/mn$, so the dot product of flow matrix $\mathbf{P}$ and pairwise distance matrix $\mathbf{D}$ is the average of all sample pairwise distances $\frac{1}{nm} \sum_{i,j} \mathbf{D}(X_i, Y_j)$. On the other hand, $\ell_2$ is the distance between two class means, which is $\|\frac{1}{n}\sum_{i}X_i - \frac{1}{m}\sum_{j}Y_j\|$. These two expressions are only equivalent when all points in each matrix are equidistant from each other, which does not hold in general.

Thanks for carefully reading our rebuttal and we appreciate your insightful feedback. We just updated the submission and included the discussion about the possibility of using Bures-Wasserstein-CPCC in our rebuttal into App.K (line 1613-1634). Because of the scalability issue for Bures-Wasserstein, it is unfortunately difficult to train with full epochs to evaluate other generalization/distortion metrics specifically for Bures-Wasserstein. Feel free to let us know if you have any remaining suggestions, questions, or concerns. We would be happy to continue the discussion.

We sincerely thank the reviewer for appreciating the clarity of our writing and recognizing the theoretical extension of CPCC with EMD, which enables handling complex, non-unimodal distributions and effectively captures nuanced hierarchical structures in data. Let us answer your questions and concerns below:

The improvement of FastFT and EMD, which appears to be the paper’s key contribution, is modest on empirical datasets, showing less than a 1% gain on average across three random seeds.

We understand your concern. We want to clarify that the 1% gain is only for the accuracy metrics in Table 4 and 5. However, on retrieval metrics in Table 4 and 5 (MAP), we can observe that the gain is significant, and typically much larger than 1%. Besides, when the multi-mode distribution assumption is valid (including CIFAR10, CIFAR100, INAT, based on GMM test in Table 5), we can see a huge improvement in the preciseness of hierarchical information in Table 3.

The computation time for empirical datasets is not reported.

We include the wall clock run time of different CPCC methods training for one epoch. The result is averaged for three seeds, and we use two very different settings, CIFAR10 with batch size = 128 and model pretraining from scratch, and INAT with batch size = 1024 and model fine-tuned from ImageNet1K checkpoint. The table below summarizes the time comparison:

Although two datasets have very different settings, we can observe that FastFT consistently runs much faster than exact EMD computation, better than other two approximation methods SWD and Sinkhorn, which matches our theoretical analysis in Table 1. For the comparison between FastFT and $\ell_2$, FastFT isn’t much slower than $\ell_2$, and in INAT, FastFT is even faster than $\ell_2$. We can therefore safely conclude that our proposed method, especially FastFT, is scalable.

We include this table in App.K of the paper.

The paper is a follow-up to the CPCC work by Zeng et al. (2022), with the addition of EMD. The novelty of introducing Fast FlowTree is somewhat limited.

We want to first emphasize that from our motivating example (Fig. 1), OT-CPCC is well-inspired by very relevant practical multi-mode scenarios and generalizes the previous $\ell_2$-CPCC algorithm as a special case, thereby offering novel, stronger and more flexible solutions overall.

In addition to proposing the novel OT-CPCC method for learning structured representation, we also

• Proposed the novel FastFT algorithm which is a very high-quality approximation of exact EMD distance. To our best knowledge, we do not aware of the FastFT algorithm being proposed elsewhere under the context of EMD approximation.
• Discovered the surprisingly close connection between FastFT, FlowTree, and the 1d-OT transportation algorithm.
• Added a novel comprehensive forward/backward theoretical analysis of $\ell_2$/OT-CPCC to prove our method is a generalization of previous $\ell_2$-CPCC.

We therefore believe we have made several original contributions that span both the research areas of structured representation learning and optimal transport. We expect our novel FastFT algorithm to have a much broader impact in many scenarios that require Wasserstein computation.

We sincerely thank the reviewer for highlighting the novelty and scalability of our algorithms, the performance improvements across multiple datasets, and the impact of our work on applications requiring hierarchical class embedding. Let us answer your questions and concerns below:

I feel section 3.3 can be simplified by moving some derivations and technical discussion to the appendix.

Thank you for your suggestion to simplify Section 3.3 by moving some content to the appendix. We feel the current level of detail is important for maintaining the paper’s flow and clarity. However, we will revisit this section if space adjustments are needed to address other reviewers’ feedback.

Label embedding methods have been employed in extreme classification (see datasets: http://manikvarma.org/downloads/XC/XMLRepository.html). The authors may want to experiment on those datasets to verify scalability and compare to existing embedding methods.

Thanks for suggesting these datasets! Extreme classification sounds like a very interesting direction to apply our method. However, we want to show that it is nontrivial to apply CPCC in the extreme classification setting:

Extreme classification deals with multi-label classification problems involving an exceptionally large number of labels. Instead of a multi-label classification setting, our hierarchical classification in Sec.4.3 is a multi-class setting where each data point only has one fine or one coarse label for fine and coarse level evaluations. To apply CPCC for multi-label classification, since each data point can have multiple labels, for a pair of data point, the “tree distance” between two data points cannot be the shortest distance between only two class nodes on the label hierarchy anymore. Some modifications of CPCC must be made, and we would like to leave it as future work.

Besides, we want to show the difference between label embedding methods and OT-CPCC methods. For label embedding methods, each label is mapped to some high dimensional vector for downstream tasks: for example in the standard multi-class classification, we typically use one-hot vector embeddings. While $\ell_2$-CPCC falls into this category, OT-CPCC methods are not label embedding methods. Instead of representing one class with a Euclidean centroid, OT-CPCC methods intentionally avoided this step by considering a more nuanced relationship between labels. OT-CPCC depends on pairwise relationships between data points from two classes, to include a more complicated distributional geometry during learning. Since we rely on data point relations, there is no one-to-one correspondence between a label and a representation.

We include the discussion about the difference between label embedding methods/learning label representation vs. structured representation learning in App.B with related works.

In Table 3, the preciseness of hierachical information listed show that in most cases, the presented FastFT based method cannot beat the ℓ2 norm based method. How to interprete the results in Table 3?

The preciseness of hierarchical information depends on the multi-modality of the class conditioned features. For all datasets, when we see a “FastFT CPCC lower than $\ell_2$ CPCC,” it only happens for BREEDS settings. From Figure 5 right, the “multimodality-ness” of 4 BREEDS dataset is 1.25, which is almost close to Gaussian. When the features are Gaussian, from Proposition 3.1 we know that $\ell_2$ and EMD are equivalent if all class conditioned features share the same covariance. The minor performance drop can be either explained by the approximation error of FastFT (w.r.t EMD), or by the violation of the shared covariance assumption. For other datasets with high modality, we can observe FastFT embeds the tree information very well.

We included this explanation in line 487-489 of the paper.

Using AIC to estimate the number of modality of the classes sounds interesting. Is there any more intuitive evidence to support the existence of multi-mode in feature distribution? Or, whether the identified multi-mode classes could be visualized by e.g., t-SNE or UMAP?

Thanks for suggesting this! We included UMAP in Fig.6, yet only for visualizing the difference between different learned representations qualitatively. We thought about using t-SNE/UMAP for multi-modality measurement, but we have several concerns, such as:

• t-SNE/UMAP are not quantitative, because it is difficult to numerically determine the exact number of modes directly from visualization. Our current GMM based method, on the other hand, tells us the approximate number of modes.
• t-SNE/UMAP are sensitive to parameters.
• Since t-SNE/UMAP can be only visualized in 2D/3D, we may create some noisy artifacts in the process of dimensionality reduction. We directly learn GMM on high-dimensional data.

Therefore, we still think our current AIC based method might be a better choice, and we can use the result to get some insights to explain some phenomenon, such as the performance drop of preciseness of hierarchical information as you mentioned above.

We sincerely thank the reviewer for appreciating our extension of CPCC to OT-CPCC using EMD, the proposed linear-time algorithm for solving the OP problem, and the extensive experiments validating our approach. Let us answer your questions and concerns below:

To replace the $\ell_2$ distance between two class-centroids, the EMD is introduced based on representing class label by all samples in the class. If the number of samples in each class is large, then solving a single EMD is expensive in time cost. Notice that the distances are computed for pairwise classes, thus the times to compute EMD is quadratic to the number of classes. If the number of classes is large, then it is extremely expensive to compute the EMD based hierarchy information among embeddings. While different approximation methods can be used, it is not clear about the computation cost, compared to the baseline methods.

Yes indeed, all EMD and all its approximation OT methods will be more expensive than the $\ell_2$ method. We made some minor mistakes in computing the time complexity of $\ell_2$ in the first submission. Now we updated the corrected results in Tab.1. Let us elaborate:

For simplicity, in the time complexity analysis, we assume each class conditioned cluster has the same size $n$, which implies batch size $b = nk$.

For $\ell_2$, given one batch of batch size $b$, computing Euclidean centroids for all classes requires scanning through all $d$-dimensional data points once, which requires $O(bd) = O(nkd)$. We have at most $k$ classes in one batch, so computing the pairwise distances between class centroids will cost $O(k^2d)$. The total cost for $\ell_2$ is $O((nk+k^2)d)$.

For OT methods, we have $k$ class conditioned clusters, so there are $O(k^2)$ pairs of matrices as input for all OT methods. For each pair, OT methods computation can take from linear time to cubic time w.r.t $n$. For example, each FastFT computation is $O(nd)$ as we shown in Theorem 3.2, so the total computation time is $O(k^2nd)$. In general, the computation time for all OT methods will be $O(k^2 \cdot \text{OT-time})$.

Since all OT methods are at least $O(k^2nd)$, and $O(k^2nd) \geq O((nk+k^2)d)$, it can be expensive to use exact EMD, and thus we propose using approximation methods to close this gap.

In practice, however, for the standard batch size like 128/1024, in Fig.3 left, we can see when $k$ is fixed ($k$ = 2 in Fig.3), the difference between our proposed method, FastFT, and $\ell_2$ is much smaller than the gap between $\ell_2$ and other methods. Note that in Fig. 3, $\ell_2$ looks like a constant time w.r.t. $n$. This might be because the mean operation in Pytorch is well-vectorized, so the effect of $n$ is improved due to parallelism. Besides, we include the wall clock time of all CPCC methods for running one epoch below:

We can observe that: first, the gap between FastFT and $\ell_2$ isn’t too large, due to the different settings of batch size ($b$) , number of classes per batch ($k$), and the size of each class-conditioned cluster ($n$); second, we can also observe the difference between OT approximation methods approximately matches our time complexity analysis in Tab.1 and Fig.3 left.

We included the details of time complexity analysis in App.I, and this empirical run time table in App.K.

Finally, a literature review on learning label representations, accompanied by illustrative examples, would provide valuable context and motivation for the study. This would help readers appreciate the importance of this approach within the broader scope of label representation learning.

Thanks for mentioning label representation learning! However, we believe that OT-CPCC methods are not label representation learning methods. In label representation learning, each label is mapped to some high dimensional vector for downstream tasks: for example in the standard multi-class classification, we typically use one-hot vector embeddings. Alternative choices include using a soft label to assign some probability values in (0,1) on each class [1]. While $\ell_2$-CPCC falls into this category, OT-CPCC methods are not label embedding methods. Instead of representing one class with an Euclidean centroid in $\ell_2$, OT-CPCC methods intentionally avoided this step by considering a more nuanced relationship between labels. OT-CPCC depends on pairwise relationships between data points from two classes, to include a more complicated distributional geometry during learning. Since we rely on data point relations, there is no one-to-one correspondence between a label and a representation.

We include the difference between label representation/embedding learning and structured representation learning in App.B of the current version.

[1] Nguyen, Quang, Hamed Valizadegan, and Milos Hauskrecht. "Learning classification models with soft-label information." Journal of the American Medical Informatics Association 21.3 (2014): 501-508.

Another issue arises with the reliance on an assumed ground-truth tree structure. In most cases, this ideal tree is unknown or impractical to obtain, even if the underlying assumptions hold. As a result, multiple approximated trees may be used instead. The paper does not adequately address how gradient computations would handle these varied approximations, which is an important practical consideration.

We understand your concern about using ideal trees. We want to emphasize that our method only relies on the existence of some tree, not necessarily ideal.

• First, there is no way to verify the tree given in the dataset is ideal. For example, in CIFAR100, although we know fish and mammals are two different species biologically, we cannot notice a huge structural difference visually on fish and dolphins where the latter is a type of mammal. The given CIFAR100 hierarchy relies on some factual knowledge about nature, but it is hard to verify if this is ideal for the structured representation learning problem.
• We can use some approximated versions of the trees in OT-CPCC. We included the Effect of Tree Depth in App.H.1 and Effect of Tree Weight in App.H.2.
  • In the setting of Effect of Tree Depth, we started from tree depth = 3, which is the ground truth root-coarse-fine hierarchy. We create synthetic mid/coarser levels for tree depth = 4,5, and this modification can be seen as an approximation of the original depth = 3 ground truth tree. From Fig.10, we can observe that OT-CPCC successfully embedded the approximated, or synthetic information well. We also attached the several performance metrics in Tab.13 and at the very end of this reply. We can observe that all CPCC values are high, which matches our visualized pattern in Fig.10. For the fine and coarse accuracy, we didn’t observe a huge disadvantage of using approximated trees.
  • In the setting of Effect of Tree Weights, we started from weight x1, which is the ground truth hierarchy. To create approximated hierarchical trees, we adjusted the distance within the fine classes under the coarse label of aquatic mammals by multiplying the original weights by x2, and x4. We include the visualization in Fig.11, and performance metrics in Tab.14. The conclusions are similar to the experimental results for different tree depths.
• Since all trees are treated as a constant in our learning pipeline, i.e, in CPCC computation, only the feature distances are parametrized, and tree distances do not depend on the currently learned feature, we do not expect any problems of using approximated trees for back propagation. We can construct a tree based on learned features: for example, we can apply QuadTree to the learned feature to make tree weights depend on the parameters. However, this tree will not be as informative as injecting external structural knowledge to constrain the representation geometry, which is the major focus of structured representation learning.
• Besides, there are many ways to construct trees. Even when the hierarchy does not exist in the original dataset, we can extract information from other knowledge bases such as WordNet [1], or using LLMs to generate label hierarchies [2].
• We can also use “multiple” approximated trees in parallel for OT-CPCC learning, which is reasonable if we want to have a more accurate approximation of the ideal tree. There are two scenarios. First, if all approximated trees have the same structure, we can merge all trees into one by using a new tree with the same structure, but using the edge weights averaged from all trees. Then this problem reduces to the standard OT-CPCC problem. Second, if approximated trees have different structures, we can leverage algorithms used in Federated Learning. For example, if we use FederatedSGD, for each tree, we can get a copy of gradient w.r.t. the model parameters. Because these gradients have the same dimension, the final aggregated gradient update can be set as the average of the gradient for each tree. In summary, OT-CPCC methods do not rely on the existence of an ideal tree. Any tree can be used in the learning pipeline optimized by backpropagation without any issue.

[1] Miller, George A. "WordNet: a lexical database for English." Communications of the ACM 38.11 (1995): 39-41.

[2] Wan, Mengting, et al. "Tnt-llm: Text mining at scale with large language models." Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2024.

We sincerely thank the reviewer for recognizing our use of OT distance to capture distributional differences, the efficient tree-based approximation method addressing OT’s computational complexity, and the experimental validation showing improved performance in classification and class representation learning. Let us answer your questions and concerns below:

Theorem 3.2 appears to be a key contribution, as it establishes that Algorithms 1 and 2 can approximate the Earth Mover's Distance (EMD). However, the proof is vague and leaves several critical points open to interpretation. I found it challenging to follow the complete argument leading to the stated approximation result. Including a detailed proof sketch could greatly enhance clarity and strengthen the manuscript.

Sorry for confusion. Let us clarify that the purpose of Theorem 3.2 is to show the equivalence between Alg.1 and Alg.2 under the construction of the augmented tree (Fig. 2), rather than showing any of these algorithms can approximate the Earth Mover’s Distance. The key idea is as follows: we can observe that the blue code blocks, the update rule of the flow matrix, between Alg.1 and Alg.2 are almost identical. By [1], Alg.2 provides the optimal flow matrix for the optimal transport problem using the tree metric as the ground metric in $\mathbf{D}$. From Fig.2, the tree distance between two samples in the augmented tree are always the same, so the order of seeing any vertex within source or target distribution does not matter. This reduces a FlowTree problem in Alg.2 to a flat 1d OT problem in Alg.1, where the major difference with Alg.2 is we do not need to pass the flow from the leaf to the root node in Alg.1.

Due to the limitation of space, we included a more detailed proof sketch in line 854-861 in the current version, and made a clarification about the statement of this theorem in line 210-214.

Additionally, the approximation error of FT and FastFT is an open question, and we want to leave it for further theoretical analysis in the future. In FlowTree [2], the authors analyze the approximation factor of FastFT only in the nearest neighbor setup, where they only rely on the fact that $\text{EMD} = \left<\mathbf{P}^*_{\ell_2}, \mathbf{D}_{\ell_2}\right> \leq \left<\mathbf{P}_{\text{FT}}, \mathbf{D}_{\ell_2}\right> = \text{FT}$ by the definition of the linear programming expression of EMD. However, the exact approximation factor/error or FT is unknown.

[1] Kalantari, Bahman, and Iraj Kalantari. "A linear-time algorithm for minimum cost flow on undirected one-trees." Combinatorics Advances. Boston, MA: Springer US, 1995. 217-223.

[2] Backurs, Arturs, et al. "Scalable nearest neighbor search for optimal transport." International Conference on machine learning. PMLR, 2020.

Similarly, while Lemma 3.4 provides the derivative of the EMD, the result itself is relatively straightforward. More importantly, it is unclear whether the same derivative result directly applies to FastTree. A formal proof confirming that this result holds for FastTree would significantly enhance the paper's contribution and rigor.

Sorry for being vague about the differentiability of FastFT. For FastFT, the flow matrix $\mathbf{P}$ does not depend on the model parameters $\theta$. Instead, from Alg.1 and Alg.2, we can observe that it only depends on the structure of the input (augmented) label tree. Therefore, in the computation of the objective function $\left<\mathbf{P}, \mathbf{D}(\theta)\right>$, $\mathbf{P}$ is treated as a constant with respect to the model parameters $\theta$, and the backward gradient flow does not pass through $\mathbf{P}$. So, we can still use $\mathbf{P}\cdot\frac{\partial\mathbf{D}}{\partial \theta}$ for the back propagation for Fast FlowTree. This expression has a format similar to EMD’s partial derivative.

We include this explanation in line 295-297 of the current version.