Distance-Based Tree-Sliced Wasserstein Distance

Q2. In Section 5.2, when suggesting to sample orthogonal directions, it is stated that it is "to ensure that the sampled tree systems do not include similar directions". While I understand what is meant here; the goal is to have different directions to better capture the topology of the data; the sentence is not very clear as in practice, we would have $\theta_i \not= \theta_j$ almost surely. Note also that it was proposed e.g. in [1] to sample orthogonal directions for Sliced-Wasserstein and in [2] for a max-Sliced-Wasserstein variant, but these references are missing.

Answer. We recall the sentence mentioned by the reviewer in Section 5.2, located on line 382.

The intuitive motivation for this choice is to ensure that the sampled tree systems do not include lines with similar directions.

Here, "similar" refers to "almost identical." Specifically, in a tree system $\{(x,\theta_1), \ldots, (x,\theta_k)\}$, we require $| \langle \theta_i, \theta_j \rangle | \ll 1$ to ensure that no two lines have nearly the same direction. This requirement motivates the introduction of the orthogonality condition in the space $\mathbb{T}^\perp$.**

Q3. I do not think I saw the step size used for each variant of SW in the gradient flow experiment. I believe that using the same step size might not always be fair as they do not necessarily have the same scale/smoothness properties (although I would agree that it is complicated to use a different step size for each method).

Answer. In the Gradient Flow experiment, we adopt the training settings, including the step size, from [1] and [2] for the existing methods and tune the step size for our approach. Since the step sizes for the existing methods were already optimized in the referenced papers, and we did not modify the settings or training process, we believe this ensures a fair comparison.

Reference.

[1] Guillaume Mahey et al., Fast optimal transport through sliced generalized wasserstein geodesics. In Advances in Neural Information Processing Systems 36 (2024).

[2] Khai Nguyen et al., Sliced wasserstein estimation with control variates. In The Twelfth International Conference on Learning Representations, 2024.

We sincerely thank the reviewer for the valuable feedback. If our responses adequately address all the concerns raised, we kindly hope the reviewer will consider raising the score of our paper.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1+W2. There is no statistical analysis of the new method. One could wonder how the distance evaluates w.r.t. the number of samples or the number of projections.

Another weakness is the lack of an ablation study, which would help us to see what really make it better than TSW-SL. Is it the new way of sampling trees which makes a big difference? Or is it because the distance is E(d)-equivariant? Or is it because the splitting map depends now on the distance to the lines?

Answer. We present an ablation study on the number of tree systems $L$ and the number of lines $k$ within each tree system.

In Figure 8 in Appendix C.8, we include the ablation study for $L$ and $k$ on the Gradient flow task. We chose the Gaussian 100d dataset. When not chosen to vary, we used the default values: number of projections $n=100$, learning rate $lr=0.005$, number of iterations $n_{iter}=5000$, number of trees $L=500$, number of lines per tree $k=2$, and coefficient $\delta=10$. We fixed all other values while varying $k\in \{2,4,16,64,128\}$ and $L \in \{100, 500, 1000, 1500, 2000, 3000\}$.

From Figure 8a in Appendix C.8, the ablation analysis of the number of lines $k$ reveals an interesting relationship between the number of lines per tree ($k$) and the algorithm's performance. When fixing the number of trees ($L$), configurations with fewer lines per tree demonstrate faster convergence to lower Wasserstein distances. This behavior is reasonable because increasing $k$ expands the space of tree systems, which in turn requires more samples to attain similar performance levels. For instance, $k=2$ and $k=4$ configurations reach lower Wasserstein distances compared to $k=64$ and $k=128$. Additionally, as demonstrated in the runtime and memory analysis (Appendix C.7), configurations with higher number of lines per tree incur greater computational and memory costs while yielding suboptimal performance.

From Figure 8b in Appendix C.8, the convergence analysis demonstrates a clear relationship between the number of trees ($L$) and the algorithm's performance. When fixing the number of lines per tree ($k$), configurations with more trees achieve significantly better convergence, reaching much lower Wasserstein distances. This behavior is expected because increased sampling in the Monte Carlo method used in Db-TSW leads to a more accurate approximation of the distance. Specifically, comparing $L=3000$ and $L=100$, we observe that $L=3000$ achieves a final Wasserstein distance of $10^{-5}$, while $L=100$ only reaches $10^2$. However, there is a computation-accuracy trade-off since more trees involves more computation and memory.

We believe that a comprehensive investigation into the analytical and statistical properties of Db-TSW is essential to theoretically explain these findings. However, as this paper primarily focuses on analyzing distance-based designs for splitting maps and the corresponding Radon Transform, and given the already broad scope of the work, we have opted to leave these aspects of Db-TSW for future research.

Q1. In equation (2), $Rf_{\mu}$ is not a probability. So I think there is an abuse of notation. In equation (8), $R_{\mathcal{L}}^{\alpha}\mu$ is used, but $R_{\mathcal{L}}^{\alpha}$ is only defined on $f \in L^2$ in equation (7).

Answer. We thank the reviewer for the correction. There is an abuse of notation between probability distributions and measures in the paper. We have made necessary adjustments in the revision.

Thank you for your response and additional suggestion. We appreciate your support and will include an ablation study exploring the effects of changing the sampling tree and modifying the splitting maps soon.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Q6. In the conclusion, it is mentioned that improved performance is expected by adapting recent advanced sampling schemes or techniques from SW to TSW-SL. However, it seems that the extension is not straightforward: can you provide some examples that could be adapted to Dd-TSW?

Answer. The tree-sliced framework in our paper has the following local perspective: Each line in a tree system is treated similarly to a line in the Sliced Wasserstein (SW) framework. Splitting maps determine how the mass at each point is distributed across the lines, and then the projection of these mass portions onto the lines is processed in the same way as in SW. A significant challenge with this approach is to verify whether the injectivity of the corresponding Radon Transform is preserved, as this determines whether the proposed metric qualifies as a true metric or merely a pseudo-metric. However, we addressed this concern by providing the proof in Appendix B.3of the paper. Considering, for example, the Generalized Sliced Wasserstein (GSW) distance [1], in GSW, the SW framework is retained, but the projection mechanism is altered. Specifically, GSW generalizes the integration level set in the Radon Transform, replacing the level set defined by the inner product (representing orthogonal hyperplanes) with one defined by an arbitrary function. Similarly, in TSW-SL, which currently relies on the inner product, a framework could be developed that generalizes this by using an arbitrary function, offering new flexibility and applications.

Reference.

[1] Soheil Kolouri et al., Generalized Sliced Wasserstein Distances. Advances in neural information processing systems, 32, 2019

[2] Khai Nguyen et al., Sliced wasserstein with random-path projecting directions. In Forty-first International Conference on Machine Learning, 2024

We sincerely thank the reviewer for the valuable feedback. If our responses adequately address all the concerns raised, we kindly hope the reviewer will consider raising the score of our paper.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1. Unsufficient analysis of its various parameters. The settings of some parameters are not discussed (e.g., the number of lines $k$ and the value of the inverse temperature parameter $\delta$).

Answer. We appreciate the reviewer's suggestion about the analysis for hyperparameters of Db-TSW. In Figure 8 in Appendix C.8, we include the ablation study for $L$, $k$, and $\delta$ on the Gradient Flow task. We chose the Gaussian 100d dataset. When not chosen to vary, we used the default values: number of projections $n=100$, learning rate $lr=0.005$, number of iterations $n_{iter}=5000$, number of trees $L=500$, number of lines per tree $k=2$, and coefficient $\delta=10$. We fixed all other values while varying $k\in \{2,4,16,64,128\}$, $L \in \{100, 500, 1000, 1500, 2000, 3000\}$, and $\delta \in \{1, 2, 5, 10, 20, 50, 100\}$.

From Figure 8a in Appendix C.8, the ablation analysis of the number of lines $k$ reveals an interesting relationship between the number of lines per tree ($k$) and the algorithm's performance. When fixing the number of trees ($L$), configurations with fewer lines per tree demonstrate faster convergence to lower Wasserstein distances. This behavior is reasonable because increasing $k$ expands the space of tree systems, which in turn requires more samples to attain similar performance levels. For instance, $k=2$ and $k=4$ configurations reach lower Wasserstein distances compared to $k=64$ and $k=128$. Additionally, as demonstrated in the runtime and memory analysis (Appendix C.7), configurations with higher number of lines per tree incur greater computational and memory costs while yielding suboptimal performance.

From Figure 8b in Appendix C.8, the convergence analysis demonstrates a clear relationship between the number of trees ($L$) and the algorithm's performance. When fixing the number of lines per tree ($k$), configurations with more trees achieve significantly better convergence, reaching much lower Wasserstein distances. This behavior is expected because increased sampling in the Monte Carlo method used in Db-TSW leads to a more accurate approximation of the distance. Specifically, comparing $L=3000$ and $L=100$, we observe that $L=3000$ achieves a final Wasserstein distance of $10^{-5}$, while $L=100$ only reaches $10^2$. However, there is a computation-accuracy trade-off since more trees involves more computation and memory.

From Figure 8c in Appendix C.8, the analysis of different $\delta$ values indicates variations in the algorithm's performance. Moderate values of delta ($\delta=1-10$) enhance the convergence speed compared to $\delta = 1$, as they effectively accelerate the optimization process. However, when delta becomes too large ($\delta=20-100$), the converged speed decreased.

W2. In the gradient flow experiment, differences in ground costs (e.g., $L_2$ and tree metric) between methods make it challenging to compare results at a fixed number of iterations for a fixed LR.

Answer. As described in Section 6.2 of the paper (Lines 459–466), the objective of this task is to update the source distribution to make it as close to the target distribution as possible. The empirical advantage of Db-TSW over other SW-variants and over the original TSW-SL is demonstrated by measuring the $2$-Wasserstein distance between the source and target distributions at steps 500, 1000, 1500, 2000, and 2500. In this task:

• Db-TSW and baselines are used as the loss function.
• The $2$-Wasserstein distance serves as the evaluation metric.

It is crucial to note that the roles of the $2$-Wasserstein distance (evaluation metric) and the other distances, such as Db-TSW, TSW-SL, and SW-variants (loss functions), are separate. Hence, there is no difference when evaluating the distance between the updated source distribution and the target distribution.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Thank you for your response, and we appreciate your endorsement. For high-dimensional cases, advanced methods such as SWGG may require a more suitable number of Monte Carlo samples to fully realize their advantages over the original SW. The behaviors of SW and TSW variants in high dimensions are indeed intriguing problems, and we hope that more techniques related to SW and TSW will be discovered in the future.

W4. The paper would be strengthened if explicitly discussing sample/projection complexity, and also error bounds, convergence rate for MC approximation.

Answer. This paper emphasizes the analysis of both traditional and newly introduced splitting maps, particularly the $E(d)$-invariant splitting maps, in relation to the Radon Transform. Additionally, we propose the novel Db-TSW approach. Due to the extensive content presented, we have deferred the analytical and statistical examination of STSW to future work.

It is important to note that the analysis of Db-TSW poses unique challenges arising from the inclusion of splitting maps—a distinctive feature of Tree-Sliced Wasserstein (TSW) variants that differentiates them from Sliced Wasserstein (SW) variants.

We are actively investigating the properties of splitting maps, which appear to be a highly promising research direction.

W5. What are the authors' rationales for not using a consistent set of baseline methods to evaluate relative performance across tasks? Moreover, is it a good idea to compare against non-linear slicing methods like GSW, ASW? Why or why not?

Answer. In each experiment, we compare our methods with the baselines provided in the most recent papers that include the corresponding task. Empirical evidence demonstrates that Db-TSW consistently outperforms state-of-the-art SW variants across various tasks. For instance, in the Denoising Diffusion Model task, we include all baselines from the latest relevant paper [1]. As shown in Table 1, [1] reports the best performance with an FID of 2.70, while our method achieves FIDs of 2.60 and 2.525. We believe this effectively highlights the advantages of Db-TSW.

Q1. Can any component of this framework be made learnable? That would probably enable various interesting extensions.

Answer. We appreciate the reviewer’s suggestion. There are two straightforward approaches to incorporate a learnable perspective: (1) tuning the parameter $\delta$ and (2) modifying the splitting map $\alpha$.

For $\delta$, it can be parameterized as a single learnable variable. For splitting maps, recall that the splitting maps used in our method are defined in Eq. (25):

$\alpha(x,\mathcal{L})\_l = softmax \Bigl( \{\delta \cdot d(x,\mathcal{L})\_l\}\_{l \in \mathcal{L}} \Bigr)$.

More generally, one can design splitting maps in the form shown in Eq. (24):

$\alpha(x,\mathcal{L})\_l = \beta \Bigl( \{d(x,\mathcal{L})\_l\}\_{l \in \mathcal{L}} \Bigr)$,

where $\beta \colon \mathbb{R}^k \to \Delta_{k-1}$ is an arbitrary map. As discussed in Section 5.2, this design ensures that the splitting maps remain $E(d)$-invariant. To make $\beta$ learnable, it can be parameterized as an MLP with input and output dimensions equal to $k$, using a softmax layer at the end to ensure the output lies on the standard simplex $\Delta_{k-1}$.

Due to the limited time available during the discussion phase, we are unable to provide empirical results for this approach. However, we have conducted an ablation study on $\delta$. In Figure 8c in Appendix C.8, the analysis of different $\delta$ values indicates variations in the algorithm's performance. Moderate values of delta ($\delta=1-10$) enhance the convergence speed compared to $\delta = 1$, as they effectively accelerate the optimization process. However, when delta becomes too large ($\delta=20-100$), the converged speed decreased.

Q2. Are there instances where prior knowledge on the data allows a good/optimal tree design?

Answer. Currently, the only prior knowledge about the data used for tree design is that the root of concurrent-line tree systems is sampled near the data's mean. Specifically, if $m$ represents the mean position of all data points, the root $r$ is sampled uniformly from a small closed ball centered at $m$. This was mentioned in line 421 of our paper.

This allows for more effective sampling, specifically by sampling trees located around the support of target distributions (ideally near their means).

Through empirical observations, we discover that initializing the tree root near the mean of the target distribution results in more consistent performance for our models. We believe that a comprehensive investigation into the analytical and statistical properties of Db-TSW is essential to theoretically explain these findings. However, as this paper primarily focuses on analyzing distance-based designs for splitting maps and the corresponding Radon Transform, and given the already broad scope of the work, we have opted to leave these aspects of Db-TSW for future research.

References

[1] Khai Nguyen et al., Sliced wasserstein with random-path projecting directions. ICML 2024

We sincerely thank the reviewer for the valuable feedback. If our responses adequately address all the concerns raised, we kindly hope the reviewer will consider raising the score of our paper.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1. It would be nice to have an explicit discussion of computational/memory complexity.

Answer. Since GPUs are now standard in machine learning workflows due to their parallel processing capabilities, and Db-TSW is primarily intended use for deep learning tasks, like generative modeling and optimal transport problems, which are typically GPU-accelerated, we provide both complexity analysis and empirical runtime/memory of Db-TSW with GPU settings. We have also included these information in Appendix C.6 and C.7 of our revised manuscript.

Assume $n \ge m$, the computational and memory complexity of Db-TSW are $O(Lknd + Lkn\log n)$ and $O(Lkn + Lkd + nd)$, respectively. In details, the Db-TSW algorithm will have three most costly operations as described in the table below (also Table 5 in Appendix C.6):

Table: Complexity Analysis of Db-TSW

The kernel fusion trick: In the table above, the distance-based weight splitting operation involves: (1) finding the distance vector from each point to each line, (2) calculating the distance vectors' norms, and (3) applying softmax over all lines in each tree. The first step costs $O(Lknd)$ computation and $O(Lknd)$ memory. Similarly, the second step costs $O(Lkdn)$ computation and $O(Lknd)$ memory. Lastly, the final step costs $O(Lkn)$ computation and $O(Lkn)$ memory. Therefore, this operation theoretically costs $O(Lknd)$ in terms of both computation and memory. However, we leverage the automatic kernel fusion technique provided by torch.compile (torch document) to "fuse" all these three steps into one single big step, thus contextualizing the distance vectors of size $Lkn \times d$ in a shared GPU memory instead of global memory. In the end, we only need to store two input matrices of size $Lk\times d$ (a line matrix) and $n \times d$ (a support matrix), along with one output matrix of size $Ln \times k$ (a split weight), which helps reduce the memory stored in GPU global memory to only $O(Lkn + Lkd + nd)$.

In Appendix C.7, we have conducted a runtime and memory comparison of Db-TSW with respect to the number of supports and the support's dimension in a single Nvidia A100 GPU. We fix $L=100$ and $k=10$ for all settings and varies $N \in \{500, 1000, 5000, 10000, 50000\}$ and $d \in \{10, 50, 100, 500, 1000\}$.

In Figure 7a (Appendix C.7), the relationship between runtime and number of supports demonstrates predominantly linear scaling, particularly beyond $10000$ supports, though there is subtle non-linear behavior in the early portions of the curves for higher dimensions. The performance gap between high and low-dimensional cases widens as the number of supports increases, with $d=1000$ taking approximately twice the runtime of $d=500$, suggesting a linear relationship between dimension and computational time.

In Figure 7b (Appendix C.7), the memory consumption analysis of Db-TWD reveals several key patterns which align with the theoretical complexity analysis and suggest predictable scaling behavior. First, there is a clear linear relationship between memory usage and the number of samples across all dimensions. Second, in higher-dimensional cases ($d = 100, 500, 1000$), the memory requirements exhibit a proportional relationship with the dimension, as evidenced by the approximately equal spacing between these curves. Finally, for lower dimensions ($d = 10, 50, 100$), the memory curves nearly overlap due to the dominance of the $Lkn$ term in the memory complexity formula $O(Lkn + Lkd + nd)$, where the number of supports components ($Lkn$) has a greater impact than the dimensional components ($Lkd$ and $nd$) when $L=100$ and $k=10$.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1. Although Db-TSW is optimized for GPU implementation, the method’s reliance on high-dimensional Euclidean distances for each point-line relationship could still pose challenges for scalability in very high-dimensional settings. This could result in increased computation times or memory usage for large datasets, particularly when fine-grained positional information is necessary. I would like to see a thorough runtime comparison of Db-TSW with respect to the number of samples and dimensionality of the data.

Answer. We agree with the reviewer that the point-line relationship theoretically introduces an additional memory and computation usage. Particularly, assume that $n \ge m$, the weight splitting operator will involve $O(Lkdn)$ memory and $O(Lkdn)$ computation to calculate distance vectors from all points to all lines. We overcome the $O(Lkdn)$ memory issue by using automatic kernel fusion feature of torch.compile. This technique allows us to fuse the distance vector calculating function and its norm calculating function into one function, thus only contextualize the distance vectors in the shared memory instead of in the global memory of a GPU. Meanwhile, the $O(Lkdn)$ computation of the weight splitting operator does not affect the overall computation complexity since it has the same computation complexity with the projection operator used widely in SW and its variants implementation. Both our weight splitting operator and the projection operator require a matrix multiplication for two matrixes of size $n \times d$ and $Lk\times d$. We have also included the detail of complexity analysis and runtime/memory evolution in Appendix C.6 and C.7 of the paper.

In Appendix C.7 of our revision, we have conducted a runtime comparison of Db-TSW with respect to the number of supports and the support's dimension in a single Nvidia A100 GPU. We fix $L=100$ and $k=10$ for all settings and varies $N \in \{500, 1000, 5000, 10000, 50000\}$ and $d \in \{10, 50, 100, 500, 1000\}$.

In Figure 7a (Appendix C.7), the relationship between runtime and number of supports demonstrates predominantly linear scaling, particularly beyond $10000$ supports, though there is subtle non-linear behavior in the early portions of the curves for higher dimensions. The performance gap between high and low-dimensional cases widens as the number of supports increases, with $d=1000$ taking approximately twice the runtime of $d=500$, suggesting a linear relationship between dimension and computational time.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Thank you for your response, and we deeply appreciate your thoughtful endorsement.

Q2. The effect of $L$, $k$ and $\delta$ to the method's performance

Answer. In Figure 8 in Appendix C.8, we include the ablation study for $L$, $k$, and $\delta$ on the Gradient flow task. We chose the Gaussian 100d dataset. When not chosen to vary, we used the default values: number of projections $n=100$, learning rate $lr=0.005$, number of iterations $n_{iter}=5000$, number of trees $L=500$, number of lines per tree $k=2$, and coefficient $\delta=10$. We fixed all other values while varying $k\in \{2,4,16,64,128\}$, $L \in \{100, 500, 1000, 1500, 2000, 3000\}$, and $\delta \in \{1, 2, 5, 10, 20, 50, 100\}$.

From Figure 8a in Appendix C.8, the ablation analysis of the number of lines $k$ reveals an interesting relationship between the number of lines per tree ($k$) and the algorithm's performance. When fixing the number of trees ($L$), configurations with fewer lines per tree demonstrate faster convergence to lower Wasserstein distances. This behavior is reasonable because increasing $k$ expands the space of tree systems, which in turn requires more samples to attain similar performance levels. For instance, $k=2$ and $k=4$ configurations reach lower Wasserstein distances compared to $k=64$ and $k=128$. Additionally, as demonstrated in the runtime and memory analysis (Appendix C.7), configurations with higher number of lines per tree incur greater computational and memory costs while yielding suboptimal performance.

From Figure 8b in Appendix C.8, the convergence analysis demonstrates a clear relationship between the number of trees ($L$) and the algorithm's performance. When fixing the number of lines per tree ($k$), configurations with more trees achieve significantly better convergence, reaching much lower Wasserstein distances. This behavior is expected because increased sampling in the Monte Carlo method used in Db-TSW leads to a more accurate approximation of the distance. Specifically, comparing $L=3000$ and $L=100$, we observe that $L=3000$ achieves a final Wasserstein distance of $10^{-5}$, while $L=100$ only reaches $10^2$. However, there is a computation-accuracy trade-off since more trees involves more computation and memory.

From Figure 8c in Appendix C.8, the analysis of different $\delta$ values indicates variations in the algorithm's performance. Moderate values of delta ($\delta=1-10$) enhance the convergence speed compared to $\delta = 1$, as they effectively accelerate the optimization process. However, when delta becomes too large ($\delta=20-100$), the converged speed decreased.

We are glad to answer any further questions you have on our submission.

Q1. Computation and memory complexity of Db-TSW

Answer. Since GPUs are now standard in machine learning workflows due to their parallel processing capabilities, and Db-TSW is primarily intended use for deep learning tasks like generative modeling and optimal transport problems which are typically GPU-accelerated, we provide both complexity analysis and empirical runtime/memory of Db-TSW with GPU settings. We have also included these information in Appendix C.6 and C.7 of the paper.

Assume $n \ge m$, the computational complexity and memory complexity of Db-TSW are $O(Lknd + Lkn\log n)$ and $O(Lkn + Lkd + nd)$, respectively. In details, the Db-TSW algorithm will have there most costly operations as described in the table below (also Table 5 in Appendix C.6):

Table: Complexity Analysis of Db-TSW

The kernel fusion trick: In the table, the distance-based weight splitting operation involves: (1) finding the distance vector from each point to each line, (2) calculating the distance vectors' norms, and (3) applying softmax over all lines in each tree. The first step costs $O(Lknd)$ computation and $O(Lknd)$ memory. Similarly, the second step costs $O(Lkdn)$ computation and $O(Lknd)$ memory. Lastly, the final step costs $O(Lkn)$ computation and $O(Lkn)$ memory. Therefore, this operation theoretically costs $O(Lknd)$ in terms of both computation and memory. However, we leverage the automatic kernel fusion technique provided by torch.compile (torch document) to "fuse" all these three steps into one single big step, thus contextualizing the distance vectors of size $Lkn \times d$ in a shared GPU memory instead of global memory. In the end, we only need to store two input matrices of size $Lk\times d$ (a line matrix) and $n \times d$ (a support matrix), along with one output matrix of size $Ln \times k$ (a split weight), which helps reduce the memory stored in GPU global memory to only $O(Lkn + Lkd + nd)$.

In Appendix C.7, we have conducted a runtime and memory comparison of Db-TSW with respect to the number of supports and the support's dimension in a single Nvidia A100 GPU. We fix $L=100$ and $k=10$ for all settings and varies $N \in \{500, 1000, 5000, 10000, 50000\}$ and $d \in \{10, 50, 100, 500, 1000\}$.

In Figure 7a (Appendix C.7), the relationship between runtime and number of supports demonstrates predominantly linear scaling, particularly beyond $10000$ supports, though there's subtle non-linear behavior in the early portions of the curves for higher dimensions. The performance gap between high and low-dimensional cases widens as the number of supports increases, with $d=1000$ taking approximately twice the runtime of $d=500$, suggesting a linear relationship between dimension and computational time.

In Figure 7b (Appendix C.7), the memory consumption analysis of Db-TWD reveals several key patterns which align with the theoretical complexity analysis and suggest predictable scaling behavior. First, there is a clear linear relationship between memory usage and the number of samples across all dimensions. Second, in higher-dimensional cases ($d = 100, 500, 1000$), the memory requirements exhibit a proportional relationship with the dimension, as evidenced by the approximately equal spacing between these curves. Finally, for lower dimensions ($d = 10, 50, 100$), the memory curves nearly overlap due to the dominance of the $Lkn$ term in the memory complexity formula $O(Lkn + Lkd + nd)$, where the number of supports components ($Lkn$) has a greater impact than the dimensional components ($Lkd$ and $nd$) when $L=100$ and $k=10$.