Title: Differentiable Generalized Sliced Wasserstein Plans
Venue: NeurIPS 2025
Accepted: yes
Keywords: Optimal Transport plans

Generalized Sliced Wasserstein
Abstract: Optimal Transport (OT) has attracted significant interest in the machine learning community, not only for its ability to define meaningful distances between probability distributions -- such as the Wasserstein distance -- but also for its formulation of OT plans. 
Its computational complexity remains a bottleneck, though, and slicing techniques have been developed to scale OT to large datasets. Recently, a novel slicing scheme, dubbed min-SWGG, lifts a single one-dimensional plan back to the original multidimensional space, finally selecting the slice that yields the lowest Wasserstein distance as an approximation of the full OT plan. Despite its computational and theoretical advantages, min-SWGG inherits typical limitations of slicing methods: (i) the number of required slices grows exponentially with the data dimension, and (ii) it is constrained to linear projections. Here, we reformulate  min-SWGG as a bilevel optimization problem and propose a differentiable approximation scheme to efficiently identify the optimal slice, even in high-dimensional settings. We furthermore define its generalized extension for accommodating data living on manifolds. Finally, we demonstrate the practical value of our approach in various applications, including gradient flows on manifolds and high-dimensional spaces, as well as a novel sliced OT-based conditional flow matching for image generation -- where fast computation of transport plans is essential.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the generalized sliced Wasserstein distance using arbitrary projection functions, including neural networks, thereby overcoming the limitations of previous methods that restricted the choice of projections. The core approach is based on a bilevel optimization framework combined with the application of Stein’s lemma for gradient estimation. Specifically, the authors formulate the minimization of the projection parameters θ as the upper-level problem, while the lower-level problem involves solving a sliced optimal transport (OT) problem with fixed θ. Due to the structure of sliced OT, the lower-level problem can be solved very efficiently using permutations. The gradient of the smoothed upper-level loss function is estimated via Stein’s lemma. To further accelerate convergence, the authors introduce a variance reduction technique leveraging the zero-mean property of Gaussian distributions. Numerical experiments on toy examples and flow matching demonstrate the improvements achieved by the proposed method.\n\nSTRENGTHS AND WEAKNESSES:\n- The proposed bilevel formulation combined with the application of Stein’s lemma for gradient estimation is highly novel. Although bilevel optimization problems are typically challenging to solve, the authors leverage the closed-form solution of sliced OT, enabling efficient gradient estimation at the upper level.\n- The introduced variance reduction technique substantially enhances computational efficiency.\n- The numerical experiments provide convincing evidence of the method’s effectiveness.\n\nQUESTIONS:\n- What is the sample complexity for the gradient estimation of h, is this better than the previous work min-SWGG?\n\nLIMITATIONS:\nThe authors claim that a major drawback of min-SWGG is that the number of slices grows exponentially with the data dimension. However, the paper does not provide any analysis of the sample complexity for the proposed method, leaving it unclear whether the dimensionality issue has been addressed.\nAlthough a variance reduction technique is proposed, the paper offers limited explanation or insight into its underlying mechanism.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors reformulate Generalized Sliced Wasserstein Plans as a bilevel optimisation problem, generalising minimum Sliced Wasserstein Generalized Geodesics plans, to allow for non-linear slices parametrised by neural networks, in order to increase the expressiveness of their approximation to the true OT plan. The authors then propose a tractable optimisation scheme by Gaussian smoothing the outer objective of the bilevel problem, for which they provide a practical gradient estimator. They then validate their approach in variety of downstream applications.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- The paper is well written and does a good job at explaining the necessary background and their contribution.\n- The proposed method is simple and appears easy to implement.\n- The proposed method addresses an important problem in machine learning which is validated as an improved alternative to the commonly-used mini-batch OT in flow matching applications.\n\nWeaknesses:\n\n- The exact contribution of the authors could been signposted better - i.e. in section 3, the authors' reformulation in Definition 2 comes straight after the discussion of Generalized Sliced Wasserstein Plan which comes from previous work. \n- It would have been nice to compare the performance of other slicing methods on the flow matching experiment. From looking at the previous experiments, it seems like DGSWP is especially strong in high-dimensional problems, so checking this on more realistic data would have been beneficial.\n- In addition, for the flow matching experiments, it would be illuminating to report FID results over a range of NFE values for the Euler solver to demonstrate how each method improves the straightness of paths. This could also be achieved by looking at the path straightness metric from [1]\n\n[1] Liu et al. (2022) Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow\n\nTypos/Grammer:\n\nLine 870: There is $\\mathbb{X} \\hspace{1mm} \\mathcal{X}?$ which appears to be a comment that has been left in the appendix (I think the correct notation here should just be $\\mathcal{X}$).\n\nQUESTIONS:\n- I am not sure why there is a need to use (and to extend to the non-absolutely continuous) Stein’s lemma in section 4. To me, this appears to just be an application of REINFORCE/Evolutionary Strategies [1] which naturally handles the discontinuous setting. This citation should be added to direct readers to this literature, which could be used for improving the optimisation scheme proposed in the current paper.\n\n- Are the authors able to provide any results/intuition for when DGSWP provides a tight approximation to the true OT distance/plan, akin to the results with min-SWGG? For instance, what performance can we expect to be able to achieve by scaling the neural network $\\phi^\\theta$? \n\n[1] Salimans et al. (2017) Evolution Strategies as a Scalable Alternative to Reinforcement Learning\n\nLIMITATIONS:\nThe authors could have conduced more ablations for the stability of their optimisation scheme, for example, with respect to the hyper-parameters of $\\epsilon$ and the number of Monte Carlo particles $N$.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a differentiable approximation of OT distance and plan using generalized Sliced-Wasserstein projections and Sliced-Wasserstein Generalized Geodesics framework.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths.**\n\n- The paper is largely self-contained. The first two sections provide a clear and comprehensive introduction to related work on Optimal Transport (OT), Sliced OT (or Sliced Wasserstein, SW), and its extensions—namely Generalized SW (GSW) [1] and SW Generalized Geodesics (min-SWGG) [2]. These prior formulations are thoroughly recalled and serve as a strong foundation for the developments in later sections.\n\n- The proposed method in Sections 3 and 4 appears to be novel to the best of my knowledge. It is accompanied by theoretical guarantees, with detailed proofs provided in Appendix A.1. The overall methodology is clearly described, including a step-by-step algorithm in Appendix A.2, which I believe is sufficient for readers to understand and implement the proposed framework.\n\n- The experimental results presented in Section 5 support the claims made, and the authors provide code in the supplementary material (though I have not verified it).\n\n**Weaknesses.**\n\n- The method proposed in Section 3 appears to be a straightforward combination of two prior works: Generalized Sliced Wasserstein (GSW) for projection, and min-SWGG for selecting the transport plan with minimal cost. Proposition 1 and Lemma 1 can be viewed as direct consequences of these two existing methods.\n\n- Section 4 introduces a differentiable formulation for transport plan estimation. However, the construction of the injective projection map $\\phi^\\theta$, which is essential to both the theoretical results (e.g., Proposition 2.4) and the experiments, is not explicitly defined. While injective maps can be constructed (e.g., via $f:x\\mapsto(x,f(x))$), the paper does not provide details or evidence that the maps used in experiments satisfy the injectivity conditions required for the proposed distance to be valid.\n\n- Regarding the experimental evaluation: I appreciate the first experiment for its clarity. The second experiment follows a common setup seen in many related works and, while somewhat standard, still adds value. However, my concerns lie mainly with the third experiment involving generative modeling on CIFAR-10. Several recent works in OT and SW-based generative modeling have achieved substantially better FID scores on this dataset using alternative frameworks (e.g., [3]). The authors’ decision to constrain their method to Conditional Flow Matching is not well-justified. \n\n- Additionally, while the paper reports the total runtime for the proposed method, it does not provide any runtime comparisons with the baselines. For reference, the CIFAR-10 experiment reportedly takes over 130 hours (>5 days) on a cluster of A100s and A6000s—though the exact number of GPUs is not disclosed—which raises concerns about efficiency. Lastly, the paper does not report standard deviations or the number of runs used to compute the reported results in this task, limiting reproducibility and statistical confidence.\n\n---\n**Reference.**\n\n[1] Soheil Kolouri et al., Generalized Sliced Wasserstein Distances\n\n[2] G Mahey et al., Fast optimal transport through sliced generalized wasserstein geodesics\n\n[3] Khai Nguyen et al., Sliced Wasserstein with Random-Path Projecting Directions\n\nQUESTIONS:\nRegarding the weaknesses raised, I request clarification on the following points:\n\n1. Injective Map: Could the authors specify the construction of the injective map $\\phi^\\theta$ used in the experiments? Additionally, please indicate where this component is implemented in the code supplement. This would help verify whether the assumption required by Proposition 2.4 is satisfied in practice.\n\n2. Choice of Conditional Flow Matching (CFM): While we agree that novel methods need not outperform the state of the art, the rationale for selecting Conditional Flow Matching as the sole generative modeling framework is currently underexplained. Could the authors elaborate on the motivation and intuition behind this design choice? A clearer justification would help contextualize its relevance and potential benefits.\n\n3. Runtime Comparison with Standard OT Solvers: Since a central objective of the paper is to approximate both the OT distance and its corresponding transport plan, please provide a runtime and memory usage comparison between the proposed method and a standard OT solver (e.g., the network simplex or Sinkhorn algorithm). These two factors are critical for evaluating the method’s practical efficiency. Such a comparison—conducted under controlled settings with varying numbers of support points and data dimensionality—would provide valuable insight into the computational trade-offs involved in adopting the proposed approach over classical solvers.\n\n4. CIFAR-10 Evaluation Details: For the CIFAR-10 generative modeling task, please clarify (i) the number of independent runs used to compute the reported performance metrics, and (ii) the runtime required for the baseline methods used in comparison. This information is important for assessing the statistical reliability and computational efficiency of the proposed method.\n\n---\n\nI would be willing to consider increasing my overall evaluation if the authors are able to clarify the points raised above.\n\nLIMITATIONS:\n.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new algorithm, Differentiable Generalized Sliced Wasserstein Plans (DGSWP), an extension of sliced-based approximations of Optimal Transport (OT). The conventional computation of OT is costly. Slicing algorithms like min-SWGG approximation address the problem by using one-dimensional optimal transport plan, which can be computed efficiently, as an approximation. However, current variants suffer from drawbacks such as linear projection dependence and inability to scale to large dimensions. The paper reframe min-SWGG into a bilevel optimization problem and present a differentiable approximation scheme using Gaussian smoothing and Stein's lemma for gradient-based optimization. The generalized formulation also supports non-linear projections (e.g., neural networks) and can be applied to manifold-valued data, like hyperbolic spaces. They empirically validate the method on three environments: non-linear toy problems for OT approximation, gradient flows in high-dimensional Euclidean and hyperbolic spaces, and a novel application to conditional flow matching for image synthesis, showcasing better sample efficiency and quality.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\nThe paper  addresses an important question of obtaining transportation plans from sliced optimal transport, which is a valuable contribution to the field of computational optimal transport. Differentiable Generalized Sliced Wasserstein Plans (DGSWP) is  a method enabling efficient, differentiable OT plan approximation through a bilevel optimization reformulation of the min-SWGG problem. The paper addresses several major limitations of existing sliced OT methods, most notably linear projections and the lack of differentiability. By taking advantage of Gaussian smoothing and Stein's lemma, the authors construct a differentiable surrogate objective that can be optimized via stochastic gradient descent, even in high dimensions. The generalization to non-linear projections via neural networks and to manifold-valued data (e.g., hyperbolic spaces via horospherical slicing) demonstrates excellent flexibility. Experiments are well thought out and range across a range of settings, from toy transport problems, high-dimensional gradient flows, to a novel application of OT plans in conditional flow matching for generative models (which is new itsself). The empirical results improve over previous sliced OT methods across the board, especially in structured and high-dimensional settings. \n\n# Weaknesses\n\nWhile the proposed smoothed optimization method with Stein's lemma is elegant, the outer optimization problem remains non-convex and initialization-sensitive. An issue particularly in high-dimensional regimes where it is difficult to find informative directions. Experimentally, although the method shows empirical gains on conditional flow matching on CIFAR-10, the results are for single runs and without error bars or statistical significance tests, so it is hard to assess robustness. In addition, the paper misses comparison with a recent competitor,  the expected sliced optimal transport plan [1, 2], which averages sliced optimal transport plans with respect to a slicing distribution over projection parameter (which can be a direction or a neural network) instead of finding the minimum projection parameter. In addition, comparisons against more established differentiable OT solvers such as entropic Sinkhorn-based methods would provide a clearer picture of the scalability, accuracy, and plan quality trade-offs. Finally, while the theoretical formulation is comprehensive, some parts, especially the connection between smoothing methods and their variance, would benefit from more empirical exploration or ablation studies.\n\n[1] Expected Sliced Transport Plans, Liu et al\n[2] Energy-Based Sliced Wasserstein Distance, Nguyen et al\n\nQUESTIONS:\n1. Since the authors use control variate for Monte Carlo estimation of gradient, I wonder if control variates techniques in sliced Wasserstein cases be applicable here for the linear case?  [2] [3] It is just a minor aspect but it is worth to discuss.\n\n[3] Sliced Wasserstein Estimation with Control Variates, Nguyen et al\n[4] Sliced-Wasserstein Estimation with Spherical Harmonics as Control Variates, Leluc et al\n\n2. Can the author verify the accuracy of the obtained plans with the true OT in some cases where the ground metric is known (Eculidean, manifold, non-manifold)? It is worth to add the comparison with expected sliced optimal transport plan. I will definitely raise my score if this question is addressed.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper is an extension of the min-SWGG\ntowards two direction: 1) generalization of the linear projections of the measure supports to more general function $\\phi$, in particular neural networks, and 2) smoothing of the objective wrt $\\pi$ based on (a generalization of) Stein's lemma. Finally DGSWP comes out.\nNumerical examples on the approximation of the OT plans themselves (also on hyperbolic spaces)\nand on their use in flow matching (also on hyperbolic spaces) are provided.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \nIt is a solid paper on a generalization of min-SWGG (min-SWGG does not work so well in our experience). The smoothing idea based on Stein's lemma is original and the performance appears to improve previous methods.\n\nWeaknesses:\nThe writing is sometimes suboptimal, e.g. the key formula is (5). Here it would be good to see the definition of $g$ again or a link to formula (4), see also minor items below.\nIn the numerical part, I am missing some experiment that compares the computed approximate planes to the original ones and on the reduction of computation time achieved by the new method, see also below.\n\nQUESTIONS:\nQuestions:\n- Proposition 1 seems of limited practical use, as it requires an injective function from $\\mathbb R^d$ \nto $\\mathbb R$, which is numerically super-unstable.\n- When $\\phi$ is a NN with a ReLU activation function as in the experiments (see also Lemma 1 - smoother activation functions are rarely positive one-homogeneous), it is not differentiable \nat hyperplanes in $\\R^q$. Therefore Assumption 1 is likely not fulfilled.\n- It is not clear if (linear) DGSW has practical benefits to SWGG. \nFig. 1 does not show the former, and Fig. 4 gives the error depending on the iteration number. \nIt would be better to have the computation time, as one interaction of DGSW needs the additional sampling of the Gaussian. In Fig. 4, ASWD appears without explanation, and seems to be the best algorithm in most cases.\n- A comparison with expected plans from [35] might be interesting.\n- Which function evaluations does Table 1 mean? There (and in Fig. 6), computation time would be more relevant than number of iterations/function evaluations.\n- What about local minima in the minimization with respect to theta; there should my many high dimensions?\n- Could also e.g. $d$-spheres be considered or is the positive curvature a problem?\n(since I stepped in just as a emergency reviewer I had not enough time to go through the whole supplement)\n\nI have some minor items:\n- Until definition 2 you deal with $p>1$, why you move to $p \\ge 0$ (and waht sense make $p=0$ at all),\nmoreover the setting for $p=\\infty$ should be explained or excluded.\n- line 121: ''extend the latter framework to by ...'' something is missing here\n- line 132: ''The authors ...'' in the context this refers to 836] which is a PhD by one person\n- line 151: why $g$ is continuous without any assumptions on $\\phi$?\n- line 152: waht is a suboptimal point?\n- The formulation of Lemma 2 is not nice. The definition of $Z$ comes after its use \nand one should write ''for any $\\varepsilon >0$.\n- caption Fig 2: ''hence there is the need ...''\n- Lemma 3 (supplement): make comment in the main paper better visible (e.g. as a remark with a number to refer to or just as lemma)\n- I do not like the notation $C_{\\nu,\\mu}$ since the cost function does not depend on the measures, but only\non their support, i.e. the $x_i$, $y_j$\n- Please start personal names with capitals (Radon, Wasserstein, Sinkhorn, ... ) in the text and the references; same with names of journals\n\nLIMITATIONS:\nsee weakness and questions", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-09"}