Tree-Sliced Wasserstein Distance with Nonlinear Projection

We direct the Reviewer to Tables R1-2 and Figure R1, available at https://sites.google.com/view/nonlinear-tsw-4.

Q1. It seems that any continuous injective function $h$ can be used to define the generalized Radon transform. However, the criteria for selecting a suitable $h$ for different datasets or tasks remain unclear.

Answer Q1. We kindly refer the Reviewer to our response to Q1+W1 in Reviewer m8Ge's review.

Q2. In Section 5.2, the authors explain that $h(x)$ can be determined by a neural network, but the structure and size of the neural network are not specified.

When $h$ is implemented as a neural network, the cost of optimizing its parameters is not included in the computational complexity analysis in Section 5.3. This suggests that the actual complexity should be higher than the proposed value.

Answer Q2. A potential neural network for implementing $h(x)$ can be inspired by [6], where the neural network is a single feedforward layer that maps $\mathbb{R}^d \to \mathbb{R}^d$. Such a layer would have $d \times d$ weights and $d$ biases, introducing $O(d^2)$ parameters. If $h$ is implemented as a neural network, the cost of optimizing its parameters would scale the overall complexity linearly by the number of optimization steps. Incorporating learnable parameters into $h$ could enhance performance by better adapting to the data distribution, but it substantially increases the computational cost.

In our work, efficiency is a major concern, so we opted for a simple mapping to avoid this overhead. Specifically, for SpatialTSW, we use the mapping $h(x) = (f_1(x), \ldots, f_d(x))$, where $f_i(x) = x_i + x_i^3$. This mapping costs $O(n \cdot d_{\theta})$, which is trivial compared to other steps like projection or sorting.

Q3. What is the reasoning behind the significant reduction in computational cost when using the Circular Radon Transformed Sliced-Wasserstein distance?

Additionally, in Figure 2, the computational cost of circular TSW appears to increase significantly when $n \geq 400$, but no explanation for this behavior is provided.

Answer Q3. The significant reduction in computational cost when using the Circular Sliced Wasserstein variants arises from the efficiency of computing $L_2$ norms compared to inner products. This advantage is illustrated in Figure R1.

As shown in Figure 2, the computational cost of $CircularTSW_{r=0}$ increases when $n \geq 400$. This behavior is attributed to the cost of computing the splitting map $\alpha$, which remains $\mathcal{O}(L k n d_{\theta})$ — the same complexity as other tree-sliced variants. While $CircularTSW_{r=0}$ is theoretically and empirically more efficient in the projection and sorting steps, the cost of computing $\alpha$ becomes dominant as the number of support points $n$ grows.

In practice, $CircularTSW_{r=0}$ is slower than standard SW when the number of supports increases significantly since it involves the additional computation of the splitting map $\alpha$, whereas SW does not. However, in a large-scale experiment such as a diffusion model, $CircularTSW_{r=0}$ increases training time by only $12$% compared to SW while substantially improving the FID from $3.64$ to $2.48$, highlighting a favorable practical trade-off. Additionally, $CircularTSW_{r=0}$ outperforms the best existing tree-sliced method, Db-TSW, in training time and FID.

Q4. The Circular Radon Transform (Equation (9)) is defined in $\mathbb{R}^d$ rather than in a circular space. Could you clarify why it is referred to as the Circular Radon Transform?

Answer Q4. The term circular refers to the circular defining function used in the Radon Transform, rather than the domain of the space itself. This naming convention is consistent with several prior works (see [4], [5]).

In contrast, for distances involving measures defined on the sphere, the term spherical is typically used (see [1], [2], [3]).

---

We thank the Reviewer for the constructive feedback, as well as for pointing out typos and missing references, which we will address. If the Reviewer finds our clarifications satisfactory, we kindly ask you to consider raising the score. We would be happy to address any further concerns during the next stage of the discussion.

---

References.

[1] Hoang Tran et al., Spherical Tree-Sliced Wasserstein Distance, ICLR 2025.

[2] Clément Bonet et al., Spherical Sliced-Wasserstein, ICLR 2023.

[3] Huy Tran et al., Stereographic Spherical Sliced Wasserstein Distances, ICML 2024.

[4] Soheil Kolouri et al., Generalized Sliced Wasserstein Distances, NeurIPS 2019.

[5] Gaik Ambartsoumian et al., On the injectivity of the circular Radon transform, Inverse Problems 21 (2005).

[6] Xiongjie Chen et al., Augmented Sliced Wasserstein Distances, ICLR 2022.

We direct the Reviewer to Table R1-2 at https://sites.google.com/view/nonlinear-tsw.

Q1. The paper feels a bit incremental, but there are lots of results, which compensate.

It is not really stated when one would prefer one type of non linear projection compared to another.

Answer Q1. We kindly refer the Reviewer to our response to Q1+W1 in Reviewer m8Ge's review.

Q2. You are proposing new choices for the polynomials $h$: how does it compare to the construction proposed in [1]?

Answer Q2. In [1], the author proposed using homogeneous polynomials of odd degree for the mapping $h: \mathbb{R}^d \to \mathbb{R}^{d_\theta}$, where the output dimension $d_\theta$ grows exponentially with $d$. For large $d$, such as in our Diffusion Experiment where $d \approx 3000$, this results in $d_\theta \approx 4.5 \times 10^9$ (see line 263-270), making the approach computationally impossible. In contrast, our proposed mapping, defined as $h(x) = (f_1(x), \ldots, f_d(x))$ with $f_i(x) = x_i + x_i^3$, maintains the same output dimension as the input ($d_\theta = d$) and introduces a trivial computational cost while still ensuring a non-linear projection.

Q3. How the trees are constructed? I am not sure it is precised (but I may have missed it)

Answer Q3. The tree structure used in our paper is the concurrent-lines structure introduced in [2]. We will make this explicit in the revised version of the paper.

Q4. The background is not very clear, and lot of important things to understand the paper are in Appendix. Also, more figures could help understanding the constructions.

Answer Q4.

We acknowledge that the background for the tree-sliced approach can be dense, as it builds upon the foundations laid in [2] and [4]. We thank the Reviewer for pointing this out and will revise the paper to include additional figures to improve readability.

Q5. In Figure 2, the result are given for $n$ up to $n=500$, which is rather small for sliced settings.

In Figure 2, we use $n$ up to $500$ because Diffusion Model experiments typically train with batch sizes $n < 500$. Additionally, we set $d = 3000$, $L = 2500$, and $k = 4$, aligning with the practical settings of the Diffusion Model experiment. We also provide runtime and memory analysis for our Tree-Sliced Wasserstein variants in Appendix D.3, with $n$ up to $50000$.

Q6. In Table 2 and 3, this is basically the same experiment, but the results are not provided in the same way ($W_2 \text{ versus } \log_2 W_2$).

Table 2 is about Gradient Flow on Euclidean space $\mathbb{R}^d$, while Table 3 is about Gradient Flow on the sphere $\mathbb{S}^d$. Therefore, the choice of experimental setting and baselines are different. Note that Table 2 follows the setting of [2], while Table 3 follows the setting of [4].

---

We thank the Reviewer for the constructive feedback, as well as for pointing out typos and missing references. We will address them accordingly. If the Reviewer finds our clarifications satisfactory, we kindly ask you to consider raising the score. We would be happy to address any further concerns during the next stage of the discussion.

---

References.

[1] Soheil Kolouri et al., Generalized Sliced Wasserstein Distances. NeurIPS 2019.

[2] Hoang Tran et al., Distance-Based Tree-Sliced Wasserstein Distance, ICLR 2025.

[3] Hoang Tran et al., Tree-Sliced Wasserstein Distance on a System of Lines.

[4] Hoang Tran et al., Spherical Tree-Sliced Wasserstein Distance, ICLR 2025.

[5] Tam Le et al., Tree-Sliced Variants of Wasserstein Distances, NeurIPS 2019.

[6] Tam Le et al., Sobolev Transport: A Scalable Metric for Probability Measures with Graph Metrics, AISTATS 2022.

[7] Makoto Yamada et al., Approximating 1-Wasserstein Distance with Trees, TMLR 2022.

We direct the Reviewer to Table R1-4 and Figure R1 at https://sites.google.com/view/nonlinear-tsw-2.

Q1. Explicitly highlight ... to prior work.

Answer Q1. The key technical challenge in developing this approach lies in proving the injectivity of the proposed Radon Transforms, which ensures that the resulting distances qualify as proper metrics. This differs from prior work on TSW, such as [2]. Compared to SW approaches—including those that also employ nonlinear projection techniques like [1]—the difference is even more pronounced due to the introduction of splitting maps $\alpha$.

Given these factors, we believe our approach to addressing these challenges is non-trivial.

Q2. The notions of well-definedness ... Radon Transform variants".

Answer Q2. The two mentioned properties above hold under certain assumptions on the functions $g$ and $h$, such as continuity, smoothness, and injectivity. These conditions have been discussed in previous works (see lines 134–144).

In our work, we do not elaborate on these general assumptions; instead, we provide specific types of functions that satisfy the necessary conditions for our method. A detailed discussion of why these functions meet the required properties are presented Section 4.2 and Appendix B.

Injectivity is typically required for variants of the Radon Transform because it ensures that the derived metrics are proper metrics rather than pseudo-metrics.

Q3. In what sense ... Radon Transform?

Q4. Why does $\alpha$ ... being referred to

Answer Q3+Q4. Previous works on TSW [2, 3, 4] show that replacing lines in the SW framework with tree structures via the splitting mechanism $\alpha$ consistently enhances performance, even with the same number of projections.

Notably, $\alpha$ is independent of SW components, suggesting that other SW improvements could be integrated into the tree-sliced framework. However, verifying well-definedness and injectivity in this new setting requires novel analytical approaches.

Q5. Design a simpler ... clearly demonstrated.

Answer Q5. The non-linearity in SpatialTSW complicates the design of settings where it outperforms other variants. Empirically, SpatialTSW matches Db-TSW (see Tables 1, 2, and R1), suggesting it as a drop-in replacement for Db-TSW with potential performance gains.

We provide intuition for when CircularTSW and CircularTSW$_{r=0}$ may outperform other variants. Since these distances rely on the $L_2$ norm for projection, they are likely to excel when the $L_2$ norms of the data are diversely distributed. We validate this advantage over Db-TSW and SpatialTSW in Table R2 on a synthetic dataset. The improvement is more pronounced for high-dimensional data (large $d$), indicating that variants using circular defining functions perform well while others struggle in such settings.

Q6. Across all experiments, ... image generation quality.

Answer Q6. In our paper, uncertainty quantification for Table 2 is provided in the Appendix (see Table 7). We also include uncertainty quantification for the Diffusion Model and SWAE experiments in Tables R3 and R4, respectively.

We provide a new qualitative result for Point Cloud Gradient Flow, visualizing the faster convergence of SpatialTSW in Figure R1.

Q7. The authors focus on the first-order ... and well-definedness.

Answer Q7. For $p>1$, the proposed approach can be extended. However, the Tree-Wasserstein distance with $p>1$ lacks a closed-form solution (see [5]). A meaningful alternative is provided by Sobolev Transport (ST) [6], which offers a closed-form solution and has been applied in the tree-sliced framework, as discussed in [3, Eq. (15)].

Although TSW works such as [2] and [4] do not explicitly address this aspect, their implementations support arbitrary $p>1$. We omit the complex ST literature to reduce presentation complexity while keeping implementation flexible, and thus chose not to include it.

Due to space constraints, we strongly encourage the Reviewer to refer to the ST literature [6]. Based on its properties, our extension to the $p>1$ case satisfies the theoretical guarantees previously discussed.

Q8. The results in Tables 1 and 2 ... is negligible

Answer Q8. We acknowledge that the reduction in computation time in Table 2 is marginal, as this experiment uses a toy synthetic dataset. However, in Table 1, when applied to real-world Diffusion Models, CircularTSW$_{r=0}$ reduces the total training time to 105.5 hours compared to 131 hours for Db-TSW over 1800 epochs, saving 25.5 hours of computation time.

---

We thank the Reviewer for the constructive feedback, as well as for pointing out typos and missing references, which we will address. If our clarifications are satisfactory, we kindly ask you to consider raising the score. We are happy to address any further concerns in the next discussion stage.

---

References. Kindly refer to References in our response to Reviewer jNth.

We direct the Reviewer to Table R1-2 at https://sites.google.com/view/nonlinear-tsw.

W1. [...] (many variants of TSW) outperform SW or TSW with linear projection

Q1. [...] (intuition) one (TSW) over another

Answer. Our motivation for proposing the non-linear projection framework is inspired by Generalized Sliced-Wasserstein (GSW) [1], which also includes both Circular and Spatial variants. It is underexplored in prior studies that among the three versions—original SW, SpatialSW, and CircularSW, which variant is most suitable for a given task.

This suggests that among the corresponding TSW variants—Db-TSW [2], CircularTSW, and Spatial TSW—there is no guarantee that the versions with non-linear projections will consistently outperform the linear-projection TSW. However, the two new distance variants each offer distinct advantages over standard TSW, as outlined below:

• The definition of SpatialTSW subsumes Db-TSW as a special case when the function $h$ is the identity map (see lines 122–124). This implies that models leveraging SpatialTSW have, in theory, greater representational capacity than those using Db-TSW. A similar relationship holds between the corresponding SW variants.
• The definition of CircularTSW is theoretically non-comparable to Db-TSW due to their fundamentally different constructions. However, $CircularTSW_{r=0}$ offers improved runtime efficiency. This benefit does not hold in the SW context, where $CircularSW_{r=0}$ performs poorly (see lines 252–261). One reason is that $CircularSW_{r=0}$ defines only a pseudo-metric, while $CircularTSW_{r=0}$ is a true metric.

Our framework offers greater flexibility by enabling a broader selection of distance functions. However, in Machine Learning, predicting the best variant for a task often requires empirical experimentation. Table R1 shows that both Db-TSW and SpatialTSW perform well, but the non-linearity in SpatialTSW makes it hard to determine in advance which variant is better suited for a given task.

We offer intuition for selecting CircularTSW and $CircularTSW_{r=0}$. Since these distances rely on the $L_2$ norm for the projection step, they are likely to perform well when the $L_2$ norms of the data are diversely distributed. We validate this advantage over Db-TSW and SpatialTSW in Table R2, where the distribution of $L_2$ norms is uniform. We speculate that this property explains why CircularTSW performs effectively for the Diffusion experiment (Table 1).

To the best of our knowledge, Db-TSW [2] is the only tree-sliced distance effectively suited for large-scale generative tasks involving transport from a training measure to a target measure in Euclidean space. Previously, [3] presents a basic and limited version of [2], primarily emphasizing the constructive aspects of the tree-sliced approach, which serve as foundational groundwork. Meanwhile, [4] explores the method in a spherical setting. Other works on Tree-Sliced Wasserstein (TSW), such as [5], [7], and others, are mainly designed for classification tasks and are not applicable to generative settings. This limitation arises because these methods rely on a clustering-based framework for computing slices—which is theoretically unsuitable (as the clustering must be recomputed each time the training measure is updated, rendering previous clustering results irrelevant) and empirically inefficient (since clustering is significantly more computationally expensive than linear or non-linear projection methods).

Q2. [...] (select $r$) affect performance empirically?

Answer Q2. Selecting the optimal hyperparameter, such as $r$ for CircularTSW, is challenging and often requires empirical tuning. Intuitively, $r$ should be large enough to ensure diverse projections onto the lines but should not exceed the data's magnitude. For normalized data, we suggest starting with $r = \frac{1}{\sqrt{d}}$ and tuning from there.

Q3. [...] (include) CRT and SRT?

Answer Q3. The TSW-SL [3] offers more general tree structure than the concurrent-lines tree structure used in [2]. However, this generality comes at the cost of runtime efficiency, as the concurrent-lines structure allows a GPU-friendly implementation. Since efficiency is a priority in our work, we adopt the concurrent-lines structure.

When extending TSW-SL to non-linear projections, we initially believed SRT could apply to general tree structures, as key properties like injectivity are preserved. However, this does not seem to hold for CRT, since parts of the proof that CircularTSW defines a valid metric rely on the specific concurrent-lines structure.

---

We thank the Reviewer for the constructive feedback and for pointing out the typos, which we will address. If our clarifications are satisfactory, we kindly ask you to consider raising the score. We are happy to address any further concerns in the next discussion stage.

---

References. Kindly refer to References in our response to Reviewer jNth.