Differentiable Generalized Sliced Wasserstein Plans
Authors: Laetitia Chapel
Romain Tavenard
Samuel Vaiter
Venue: NeurIPS 2025 poster ()
Keywords: Optimal Transport plans
Generalized Sliced Wasserstein
Primary area: general_machine_learning
TL;DR: We propose a Generalized Sliced Wasserstein method that provides an approximated transport plan and is differentiable.

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.

Differentiable
Generalized Sliced Wasserstein Plans
Laetitia Chapel
Romain Tavenard"
IRISA I, I''nntitut Agro Rennes-Angers Rennes, France.
Université de Rennes 2 IRISA
Rennes, France.
laetitia. cha.etiaacccael@riis...f
romain tavenardduui---rreeee.....
Samuel Vaiter
CNRS & Université Côte d' Azur
Laboratoire J. A. Dicudonné
Nice, France.
samuel, Sraitterdnnss.c
Abstract
Optimal Transport (OT) has attracted significant interest in the machine learning community, not only for its ability define meaningful distances between proba- bility distributions such the Wasserstein distance but also for its formulation of OT plans. Its computational complexity remains 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 plan. Despite its computational and theoretical advantages, min-SWGG inherits typical limitations slicing methods: (i) the number of required slices grows exponentially with the data dimension, and (ii) it is constrained to linear projections. Here, reformulate min-SWGG as a bilevel optimization problem and propose a differentiable approx- imation scheme efficiently identify the optimal slice, even in high-dimensional settings. We furthermore define its generalized extension for accommodating data living manifolds. Finally, demonstrate the practical value of our approach in various applications, including gradient flows on manifolds and high- dimensional spaces, as well novel sliced OT-based conditional flow matching for image generation where fast computation of transport plans is essential.
1 Introduction
Optimal Transport (OT) has emerged as foundational tool in modern machine learning, primarily due to its capacity to provide meaningful comparisons between probability distributions. Rooted in the seminal works of Monge (38) and Kantorovich (27). OT introduces mathematically rigorous frame- work that defines distances, such as the Wasserstein distance, that have demonstrated high performance in various learning tasks. Used as loss function, it is now the workhorse of learning problems ranging from classification, transfer learning or generative modelling, see [39] for a review. One of the key ad- vantages of OT lies in its dual nature: it also constructs optimal coupling transport plan between distributions. This coupling reveals explicit correspondences between samples, enabling a wide range of applications. For example, OT has proven valuable in shape matching (10), color transfer (47] 49), domain adaptation [4], and, more recently, in generative modeling through conditional flow match- where it provides alignment between the data distribution and samples from prior.
ing 45
"qual contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

OT (essut), (C 13.00
minGSWGG (C.s 1222222
DGSWP (NN), (C 13
souren Lurget
Figure 1: 8Gaussians (source) to Two Moons (target) distributions and associated OT plans Grey: (Left) exact solution (Middle) min-SWGG, that projects samples on an optimal line determined by random sampling (Right) Differentiable Generalized SW plan, that relies on a neural network to get non linear-based ordering of the samples. Gradient of colors represent the ordering of the samples.
Despite the many successes optimal transport in machine learning, computing OT plans remains computationally challenging problem. The most used exact algorithms are drawn from linear programming, typically resulting in a (nn²) complexity with respect to the number of samples n. To alleviate this issue, several strategies have been developed in the last decade, that can be roughly classified three families: i) regularization based methods, such as entropic optimal transport [16], ii) minibatch-basse problems and iii) approximation based through closed-form formulas such as projection- based OT. This paper is concerned with this third class of methods, with the underlying goal to obtain transport plan. We quickly review them below.
methods [23 20). that average the outputs of several smaller optimal transport
Approximation of OT loss with sliced-OT. The Sliced-Wart Wasserstein Distance (SWD) 48
approximates the Wasserstein distance by projecting onto one-dimensional subspaces, where OT has closed-form solution obtained in On n). While SWD averages over multiple projections, max- SWD [18] selects the most informative one, cnabling efficient computations on large-scale problems while preserving key properties. Generalized SWD alleviates the inefficiencies caused by the linear projections using polynomial projections or neural networks [30]. They build on generalized Radon transforms use polynomial projections neural networks, and provide conditions for which it remains a valid distance: the generalized Radon transform must be injective. Rather than considering non-linear projections, Chen et al. 12 aim at better capturing non-linearities by augmenting the input space, using injective neural networks, such that linear projections could better capture them. Generalized SWD has also been defined in the case of tree metrics (55).
Approximation of OT plan with sliced-OT. Addressing the challenge that none of the previous works provide an approximated transport plan, (37] proposes min- SWGG slicing scheme that lifts one-dimensional OT plans to the original space: [35] follow the same line but rather define expected plans computed over all the lines instead of retaining the best one.
Differentiable OT. As is discrete problem, OT is not differentiable per se. Several formulations, that mostly rely on smoothing the value function by adding a regularization term, have then been defined, the examples. It can be efficiently solved via Sinkhorn's 's matrix scaling algorithm, providing OT approximation that differentiable. Blondel et al. 7 use a l2 regularization term to define smooth and sparse approximation. Regarding the unregularized OT formulation, an approximated gradient can be computed using sub-gradient on the dual formulation (see (2 for instance).
version of the OT problem [16] being one of the most striking
Optimizing on discrete problems. The OT problem is a linear program and the optimal plan belong to the (potentially rescaled) Birkhoff polytope, i.e. the set of doubly-stochastic matrices, which discrete. It is also the case for wide range of problems for which the solution is prescribed to belong to such discrete set. For instance, one can cite the sorting operator, the shortest path or the top-k operator. They break back proppaaraagttion along the computational graph, hence preventing their use in deep learning pipelines. Numerous works have been proposed to provide differentiable proxies, based on smoothing the operators: dataSP
32 define a differentiable all-to-all shortest path algorithm, smoothed sorting and ranking [8 can be defined by projection onto the permutahedron, perturbed optimizers [4] have also been defined, relying on perturbing the input data, to name few.

---

Contributions. In this paper, we aim at approximating the OT distance and plan. We rely on slicing scheme, extending the framework of [37], highlighting that it is an instance of a bilevel optimisation problem. We provide two main contributions: a generalized provides approximated OT plans relying on non-linear projections (see Fig. for an illustration), then GPU-friendly optimization algorithm to find the best projector. We showcase the benefits brought by these contributions in three different experimental contexts: on a 2 dimensional example where non-linear projection is sought; conducting gradient flow experiment in high dimensional space; introducing a novel sliced OT-based conditional flow matching for image generation.
Waseed--sssrsttenn that
Preliminaries on Optimal Transport
We now give the necessary background on Optimal Transport and sliced- based approximations. For a reference on computational OT, the reader can refer to [4]
2.1 Discrete Optimal Transport
We consider two points clouds {z; [=)) E 2 and { E x where X' C R2d is discrete subset. We denote their associated empirical distributions y = Li= a.Sx, and v = > bggg.-
Formulation and Wasserstein distance. We consider a cost function e that will be in our setting cuu, v) ||u - =Ilß for p > 1 and Cyu- = (c(æ,, M/ p. ( --- Traditional Kantorovich formulation of optimal transport is given by the Wasserstein metric on ((TR)) defined as
WF(K.)) min
Tig ||,
yo lp subject to ((,
where b) = for E R250 -im a, π T 1n = b) is the set of couplings between and v, Eq. (1) is a convex optimization problem and an optimal transport plan oor is a solution. Note that, a priori, there is no reason for this minimizer be uniquely defined.
One-dimensional OT. When d = 1, and μ, P are empirical distributions with a; = b, = 1/n, the optimal transport problem is equivalent to the assignment problem. In this case, the Wasserstein distance can be computed by sorting the empirical samples, resulting in an overall complexity of On n). Let a and r be permutations such that (1) S FT(2))
Za(a) and yr(1) S 9r(2) S S Wr(n). The Wasserstein distance then given by WP( v) = - (Ta()) - V-(i) l.. The
optimal transport is thus monotone and its plan Tor has the form of a permutation matrix. Note that this approach can be easily extended # m and arbitrary marginals b 44
2.2 Sliced Wasserstein Distance
Formulation. Sliced- Wasserstein (SWD) relies on a simple idea: disintegrate the original problem onto unidimensional ones, and average the different solutions. More precisely, SWD 48 approximates the Wasserstein distance by averaging along projection directions E sa as
J&L. 8g
SW SWD,(u,
Wy(Psu.PPw)aNN
where p0 Rd IR is the ID projection onto the unit vector 8, P*(x) (I. 8), and À is the uniform distribution on the unit sphere gd-1, Typically, SWD is computed thanks to Monte-Carlo approximation in which directions are drawn independently. leading to a computational complexity of O(dLn + Ln log n)). One of the main drawback of SWD is that is requires a high number of random projections, which leads to intractability for high dimensional problems. Then, there have
been works to perform selective sampling, e.g. [4I) or to optimize over the directions [8] 37]. The
two later works rely on non-convex formulations that can be optimized.
Generalized Sliced WD. The idea of non- linear slicing has been explored in several works, with the aim to improve the projection efficiency, e.g. when the data live on non-linear manifolds. 30) the generalized SWD uses nonlinear projections such neural network-based ones; the conditions on which it yields a valid metric are also stated: the projection map must injective, In [41, a gen-
eralized projection is also proposed in addition to selective sampling. Augmented Sliced Wasserstein
3

---

(AWD) [12 pursued the same goal but, rather than considering nonlinear projections, they use neu- ral network to augment the imput space, which leads to a space on which linear projection better fits the data. All these works lead to significant improvements in a wide set of machine learning scenarios.
2.3 Sliced Wasserstein Plan
As it is defined as an average over ID OT distances, SWD does not provide inherently transport plar
Recently, some works have tackled this limitation by lifting one or several one-dimensional
plans back to the original multidimensional space 35]. We now give more details 37 as 35
extend the latter framework to by averaging several plans.
Sliced Wasserstein Generalized Geodesies. Mahey et al. [37 introduced an OT surrogate that lifts the plan from the ID projection onto the original space. more details, when n = m, it defined as:
min SWGG u, = 1 Σ |Fao(()) 1r DE24-1
min-SWGG/(a.v)
(3)
where oy and та are the permutations obtained by sorting 1. and PL Note that it extends naturally to the case where n & m which we omit here for brevity. By setting p = 2, it hinges on the notion of Wasserstein generalized geodesics 2 with pivot measure supported on line. This alternative formulation allows deriving optimization scheme to find the optimal 0 that relies on multiple copies of the projected samples, which holds only for p = 2. For a fixed direction 0, provided that families (xr)), and (PE (u:)); are injective, that is to say there is no ambiguity on the orderings dg and Tp- SWGG(,, 0) is itself metric [36.. Min-SWGG has appealing properties: it yields an upper bound of the Wasserstein distance that still provides an explicit transport map between the input measures. The authors show that MiG-SWGG metrises weak convergence and is translation-equivariae
Limitations of min-SWGG. While min-SWGG provides an upper bound on the Wasserstein distance, its tightness is not guaranteed in general. However, two settings are known where the bound is exact: (i) when one of the distributions supported on a one-dimensional subspace ; and (ii) when the ambi- ent dimension satisfies d > 2n (37). More generally, the number of reachable permutations increases with the ambient dimension [15, making data dimensionality a critical factor in the quality of the ap- proximation. These insights motivate first contribution: the definition of Generalized Wasserstein plans, which aim to extend the set of reachable permutations and better capture non-linear structures.
Generalized Sliced Wasserstein Plan
Generalized Sliced Wasserstein Plan (GSWP) is built upon the idea of generalizing SWGG to an arbitrary scalar field parameterized by some R2. We consider a map : Rd x IR - IR typically a one-dimensional projection 6(x, 0) = (x, 0) (in which case q = d), or a neural network, as in Fig. for example. Generalized Sliced Wasserstein Plan distance defined as the dieddimeniooall Wasserstein distance between the point clouds through the image of de == o(- f) for given 0 E R
R99 and p,v E P(2). The G-Gnerrllzeed Sliced Wasserstein Plan distance between j and is defined as
Definition 1. Let p
(p, = ( (0l (fV)
Any element n (u, E Uaa b) that achieves d (n, v) is called a a d-Generalized Sliced Wasserstein Plan (0-GSWP).
Denoting 0 Co (0) the cost matrix between ns, v through the image of Φ. 0-GSWP reads
x² E arg min g( 0) = -C ( (0), T) (a.(a.
A o(y,,0)''
(4)
The g(x, function is continuous but not convex, as illustrated in Fig. 2 Note that, i) for p > the map C has the same regularity as o and ii) a solution n² is suboptimal point of the standard (I) problem. Indeed, since the optimization occurs on the same coupling space U (a,b),, T0 is an admissible point for the original problem. GSWP defines a distance on the space of measures P(2)).
Note that a transport plan can also be inferred when performing SW gradient flows by putting into corre- spondence the original and final samples when the algorithm has converged.

---

Proposition 1. Let A E R and assume &" is an injective map on X. Then d' is distance on P(x))
For clarity, all our proofs are presented in Sec A.I Note that the injectivity condition is akin to designing sufficient and necessary conditions for the injectivity of generalized Radon transform [S]
and relates to the problem of reconstructing measures from discrete measurements [52]. In practice,
using general & increases expressivity. enabling wider range of permutations compared to
linear map. By an application of Rockafellar's 's enveloppe theorem, we also have a characterization of the gradient of d as function of 0. Let us assume that p > 1, Φ is jointly C' and (4) has
unique solution at 0 E R9. Then we have Vod -00 (0) T x Note that similar result on the subdifferential of d is true if Φ is not differentiable but convex thanks to Danskin's theorem.
ac*a
Let p 2 0 and μ4, V E P(Z)) The minimal Generalized Sliced Wasserstein Plan semimetric min-GSWP between μ and v is defined as
Definition 2.
min h(0) = [=((T) )+²)= = (Cues T²) BER
(5)
subject to
T E arg min g(T,
(EU(a.
Equation (5) defines a bilevel optimization problem. In the case where 6(x 0) = (r. and we further constrain
live on the unit sphere. (5) the min-SWGG (37 approximation of OT. Observe that since x is an admissible coupling for the original OT problem ① min-GSWP(Ru, v) an upper bound for W, (P,,,)) v). Unfortunately, the value function 0 h(O) does not, in general, possess desirable regularity properties, even for the simple choice of a one- me-dimensional projection ((z, = (ar, f): it is discontinuous, as illustrated in Fig. 2 Moreover, as soon as the maps (2z, 0) and (vy, f) give rise the same permutations dg and To. the value of x² remains the same. As consequence, one cannot use (stochastic) gradient- gradient-based based bilevel methods such as [3
on (5)
It turns out that even if min-SWGG was introduced the minimization over the sphere Sd-1, it possible to see it as an unconstrained optimization problem thanks to the following lemma.
Lemma 1. Assume that f ++ o(x, o) is I-homogeneous for all E Rdd. Then, 0 d (u,, is I-homogeneous, and h is invariant by scaling: h(cff) = h(0) for all c <>>00 0.
In particular, for every open where h is diffcrentiable, if is I-homogencous, the gradient flow Ô (() has orthogonal level lines (6. f) O, hence the dynamics occurs the sphere of radius equal to the norm of the initialization. In practice. it means and we observed - that we do not have to care about the normalization of & during gradient descent.
When (xx 0) = (ar, f)), min-SWGG can be smoothed by making perturbed copies of the projections olx, 0); heuristics for determining the number of copies and the scale of the noise are given in (37] Nevertheless, the continuity of the obtained smooth surrogate cannot be guaranteed, and the formulation is only valid for p 2. The following section proposes differentiable approximation of min-GSWP rooted probability theory which valid for any > 1. The main difference with the aforementioned scheme is that we here perturb the parameters (or direction) 0 rather than the samples.
4 Differentiable Approximation of min-GSWP
Smoothing estimators by averaging a popular way to tackle nonsmooth problems. We mention three families of strategies: i) smoothing by (infimal) convolution, ii) smoothing by using Gumbel-like trick I 4, and tii) smoothing by reparameterization, and in particular by using Stein's 's lemma. Strategy i) would be computationally intractable in high dimension and strategy ii) doesn't 't yield consistent transport plan due to perturbations in the cost matrix. We then focus on the third option.
We begin by recalling this classical result due to Stein [11, which plays a central role in our analysis. The lemma, restated below in our specific setting, provides an identity for computing gradients through expectations involving Gaussian perturbations. Typically, Stein's 's lemma is stated for "almost differentiable" function, we require here slightly less regularity
Assumption 1.
There exists open set C C Es with Hausdorff dimension H \ C* 0 (in particular Lebesgue-negliging such that the mapping 0 + h(0) continuously differentiable on C.
Under this assumption, Stein's 's lemma remains true.
5

---

9
Figure 2: Example y and h (seen as func- tion of 0 on the sphere $t for a 2D OT problem. he, is Monte-Carlo estimate of lie with gradient VN. Note that is continuous whereas h is piecewise constant, hence the need for a smoothing mechanism, that results he.v-
h
by N
J5
Lemma 2 (Stein's lemma). Suppose Assumption 7 holds and that for all indices j. the partial derivatives satisfy the integrability condition h[]0,h(Z)) < +c. Let Z (0, Id,) be standard multivariate Gaussian random variable, and > 0. Then, the following identity holds:
Ez(Vh(O Z) - - Ez =Z)Z).
Note that this version of Stein's lemma is slightly more general than the original of Stein 51 in term of regularity asked to the value function h (in particular, we do not require it to be absolutely continuous).
We define the following smoothed value function hg, which corresponds to Gaussian smoothing of the original non-differentiable outer optimization problem:
(0) Ez [h(0 + +2)] = (Cnu, =
where the D-Differentiable Generalized Sliced Wasserstein Plan (9-DGSWP) at smoothing level reads
T Ez
arg min 9( 0 +=))+ rEt((.))
(6)
Due to the nice regularity properties of Stein's approximation, we get the following proposition regarding our approximation of 0-GSWP. Proposition 2. Suppose Assumption 7 holds. pet P(X). The following statements are true:
1. (Admissibiliity) For any E > O, xh is admissible (i.e., Th (a,,)) b)). Hence, h. hz(0) gives an upper-bound of the Wasserstein metric h0) 2 WP (y,(u,)) v).
2. (Differentiabilltyy) For any > O, the map 8 + (+) is differentiable. Moreover, we have Voh (0) = E (6 EZ)Z].
3. (Consistency.) For almost all 9 E R7, lim - h (0) h(0), and if Vh(() exists, then lim;-+0 (0) <h(())
4. (Distance.) Let 1 R9, of is injective on X then the map (p,v) + (h(0)(H, v) is distance on P(x). Assume that for almost all 0 E RR, 8² is injective. Then, (p,v)
(O)(,, is also a distance on P(X).
While this result allows for the unbiased estimation of gradients, it is known that naive Monte Carlo approximation of the right-hand side tends suffer from high variance. To address this, one may consider an alternative formulation that often results in a reduced variance estimator. Specifically, using a control variate approach 66. one observes that Ez(Vh( + = =Z)
EE2 + =Z) h(0))2).
which holds due to the zero mean of the Gaussian distribution and the linearity of the expectation. Equation leads to a Monte-Carlo estimator of the gradient Vha defined as
N
he EN
Σ (l:0 + ezx) h(0)) k=1
Σ (Cav: 9teek
n )2
(8)
EN k=1
where the vectors z; ~ NO0 Id,,) are independent standard Gaussian samples. Hence, estimating the Monte-Carke Carlo gradient (T, 0) requires to solve N + ID-optimal transport problems for an overall cost of O(N(n + m) log(n + m)). Algorithm (in Supplementary material) describes gradient descent method to perform the minimization of h2 using this Monte-Carhe approximation.

---

5 Experiments
We evaluate the performance of our Differentiable Generalized Sliced Wasserstein Plans, coined DGSWP, by assessing its ability to provide meaningful approximated OT plan in several contexts. First, we consider a toy example where a non-linear projection must be considered; we then perform gradient flow experiments on Euclidean and hyperbolic spaces, demonstrating the versatilliyy of our approach. Finally. we integrate sliced-OT plans in an OT-based conditional flow matching in 0.05 and N 20 DGSWP-specinee
lieu of mini-batch OT. In all the experiments, we use hyperparameters. Full experimental setups and additional results are provided App. A.3 Implementation is available online we also use POT toolbox (1
5.1 DGSWP as an OT approximation
We begin by examining the illustrative scenario shown in Fig. where the task to compute an optimal transport (OT) plan between a mixture of eight Gaussians (source) and the Two Moons dataset (target). While some of the Gaussian modes are properly matched by min-SWGG, others are matched to the more distant moon, due to information loss along the direction orthogonal to the projection. To address this limitation, we apply DGSWP with a neural network parameterizing the projection function Φ, enabling more expressive projections. This improvement is reflected in the quantitative results: the transport cost associated to the DGSWP plan is notably closer to the squared Wasserstein distance than that of the min-SWGG method.
This transport plan above is obtained using the approach outlined in Sec. 4 To further evaluate the impact of the variance reduction technique introduced in Eq. (8). we now repeat the same experiment across I0 different random initial- izations of the neural network. The average learning curves for the variants with and with-
50
Without Variance Reduction
With Variance Reduction
40
(Caoa
30
out variance reduction are shown in Fig. E] The results clearly indicate that using variance redue- tion improves both the final transport cost and the stability of the learning process. Based on this evidence, we adopt the variance reduction strategy in all subsequent experiments.
20
0
200
40D
600
800
1000
Iseratton
Figure 3: Impact of the variance reduction scheme (first 1,000 iterations).
5.2 Gradient flows
In these gradient flow experiments, our objective is to iteratively transport the particles of a source distribution toward target distribution by progressively minimizing the DGSWP objective.
Without manifold assumption. We compare DGSWP with a linear and NN-based & mapping against several baseline methods: Sliced Wasserstein distance (SWD), Augmented Sliced Wasserstein Distance, that has been shown in [12] to outperform GSW methods, and min- SWGG, evaluated in both its random search and optimization- based forms. Fig. presents results across range of target distributions designed to capture diverse structural and dimensional characteristics, In all experiments, the source distribution is initialized as a uniform distribution within a hypereube, We set the number of samples at = 50, and repeat each experiment over 10 random cost along with the first and third quartiles. We use the same learning rate for all experiments.
the median transport
In two-dimensionnal settings, DGSWP consistently converges to a meaningful solution, whereas methods based on the min-SWGG objective sometimes fail to do so, even when using their original optimization scheme. More notably, in the high-dimensional regime, DGSWP stands out as the only slicing-bar based method capable of producing satisfactory transport plans, underscoring its robustness and scalability in challenging settings. This holds true even though min-SWGG is theoretically equivalent to Wasserstein when the dimension is high enough (here we have 2n); practice. its optimization often struggles to find informative directions. In contrast, our proposed optimization strategy succeeds even in the linear case (where the formulation effectively reduces to min-SWGG) highlighting the practical benefit of our approach.
httpp://gittub..co http:///ggh.aa.oooo.mtttww/////t
7

---

SWD
min-SWGG (random search)
DGSWP (dineac)
ASWD
SWGG (oplim)
DGSWP (NN)
logs ((
log
Gauzeian 500d
Number inevations
Number iterations
lognn biiiiaga
No
do 100
Number of Recations
Number of Beeations
Figure 4: Log of the Wasserstein Distance as function of the number of iterations of the gradient flow, considering several target distributions. The source distribution is uniform in all cases.
SWD
InSS
100
150
300
300
100
200
300
Iterations
Iscrations
Figure 5: Log of the WD (second and fourth panels) for two different targets (first and third ones) as wrapped normal distributions for HHSW, SWD and DGWSP.
On hyperbolic spaces. Hyperbolic spaces Riemannian manifolds of negative constant curva- ture (33] They manifest in several representations and we consider here the Poincaré ball of dimension d, B We aim to construct manifold feature map for hyperbolic space B x BB - R. typically relying on an analog of hyperplanes. Horospheres (26 generalize hyperplanes in the Poincaré ball and
parametrized by base point located the boundary of the space 0B gd-1 Akin 28
we consider a map that corresponds to horosphere projection o(x, 0) = log [I 0 ( log
Similary to [9]. that live in the Poincaré disk. We compare with the Horospherical sliced-Was Wasserstein discrepancy (HHSW, (9]) and Sliced Wasserstein computed on the Poincaré ball. For the gradient estimation, rely on the von Mises-Fisher which is generalization of a Gaussian distribution from IR to Sd- and perform Riemannian Gradient Descent using Python toolbox geoopt [29 to optimize on 0 and the source data (that live on Bd). Figure plots the evolution of the exact log 10- Wasserstein distance between the learned distribution and the target, using the geodesic distance as ground cost. We wrapped normal distributions as source and set the same learning rate for all methods. DGWSP enables fast convergence towards the target, demonstrating its versatility for hyperbolic manifolds.
assess the ability of generalized sliced Wasserstein plans learn distributions
5.3 Sliced-OT based Conditional Flow Matching
We now investigate the use of DGSWP in the context of generative modeling. In the flow matching (FM) framework, time-dependent velocity field (x) learned such that it defines continuous transformation from a prior to the data distribution. In practice, d4 is trained by minimizing regression loss on synthetic trajectories sampled from known coupling between source and target distributions that determines the starting and ending points (zu, X1) of the trajectory. Several variants of flow matching have been proposed depending on how these pairs are sampled. In Independent CFM (I-CFM, (44), pairs are sampled independently from the prior and data distributions respectively. However, this approach ignores any explicit alignment between source and target samples. To address this, OT-CFM [54] proposes to deterministically couple source and target
8

---

samples using mini-batch OT, so that (zo, 21) ~ т where т is the OT plan computed between batch of prior samples and batch of data samples. This coupling tends to straighten the diffusion trajectories, which leads to improved generation quality in few-step sampling regimes.
Despite its advantages, OT-CFM relies on trade-off: since computing exact OT is infeasible for large datasets, mini-batch OT is used, which leads to imperfect matchings, especially in high-dimensional spaces involving complex distributions, for which mini-batch unlikely to be representative of the whole distribution. This motivates the use of DGSWP as an alternative. By estimating transport plans using DGSWP, we can leverage significantly larger batches, resulting better couplings, while maintaining low computational cost. In our experiments, we aggregate samples from 10 minibatches to compute the transport plan, and find that the additional computational cost remains negligible. Our approach is also supported by the findings of Cheng and Schwing [33] who show that increasing the batch size improves performance for OT-CFM in the few-sampling-stepy steps regime.
We conduct experiments on CIFAR- 10 [3]], reporting FID scores for DGSWP-CFM, OT-CFM, and I-CFM across varying numbers of sampling steps. We use the experimental setup and hyperparameters from Tong et al. [54]. For DGSWP, we evaluate both a linear projector and more expressive non-linear embedding implemented by a neural network. We compare a fixed-step Euler solver with the adaptive Dornnan-Prrinc method [9] for trajectory integration. The results in Table[] the importance of learning straight transport trajectories, which translates more efficient generation: models that induce straighter flows require significantly fewer function evaluations (NFE) to achieve high-quality samples. Notably, even a simple 100-step Euler scheme proves highly effective when used with DGSWP-CFM, demonstrating the practicality of the method for fast generation.
Integration method-
Euler
DoPri5
Algorithm
FID
NFE
FID
NFE
I-CFM
4.63
100
3.65
138.16
OT-CFM
4.82
100
3.86
132.90
DGSWP-CER CFM (linear)
4.17
100
4.47
110.37
DGSWP-CFM (NN)
3.56
100
3.87
120.04
Table I: FID score and average number of function evalua- tions (NFE) per batch.
A closer analysis of the results leads to three key observations: (i) Among the projection choices, linear DGSWP underperforms compared to its neural network-based counterpart, highlighting the benefit of extending min-SWGG with non- linear projections; (ii) When using the adaptive Dormand- Prince solver, DGSWP achieves per- formance comparable to OT-CFM in terms of FID, but with lower computational cost indi- cated by the reduced number of function eval- uations; (ii) In the fixed-step regime, DGSWP clearly outperforms all baselines under the Fuler solver, offering the best trade-off between sam- ple quality and efficiency, as illustrated in Fig. 6
LCFM
DGSWP-CFM (linea))
OT-CFM
DGSWP-CFM (NN)
7
5
100k
200k
300k
400k
Training Step
Figure 6: FID function of training iterations for various algorithms using 100-step Euler sampling.
Conclusion
This paper presents novel differentiable approach to approximate sliced Wasserstein plans, incorpo- rating non-lincar projections. Its differentiable and GPU-efficient formulation enables the definition of optimal projections. The proposed method offers several key advantages: i) efficient computation com- parable to that of SWD, ii) ability to provide an approximated transport plan, akin to min-SWGG, iii) improved performance in high-dimensional settings thanks to its improved optimization strategy iv) the capacity to handle data supported on manifolds through a generalized formulation. To the best of our knowledge, also provide the first empirical evidence that slicing techniques be effectively used in conditional flow matching, resulting in better performance and fewer function evaluations.
A key strength of slicing-based approaches for approximating OT lies in their favorable statistical prop- erties, such as improved sample complexity [4]] These advantages also extend to projection- based methods on k-dimensional subspaces 42. as well as to Min-SWGG (36). Future works will explore

---

the conditions that the projection map & must satisfy to preserve these statistical guarantees. Addition- ally, ensuring the injectivity of ø is essential for DGSWP define proper distance. Investigating injective neural architectures, such as those proposed in 46 is promising research direction.
References
[1] J. Abernethy, C. Lee, A. Sinha, and A. Tewari. Online linear optimization via smoothing. In Conference on Learning Theory, pages 807-823, 2014.
[2] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
[3] Michael Arbel and Julien Mairal. Amortized implicit differentiation for stochastic bilevel optimization. In ICLR, 2022.
[4] Quentin Berthet, Mathieu Blondel, Olivier Teboul, Marco Cuturi, Jean- Yean-Philippe Vert, and Francis Bach. Learning with differentiable pertubed optimizers. In Advances in Neural Information Processing Systems, volume 33, pages 9508-9519, 2020.
[5] Gregory Beylkin. The inversion problem and applications of the generalized radon transform. Communications on Pure and Applied Mathematics, 37(5):579-555 -999, 1984.
[6] Mathieu Blondel and Vincent Roulet. The elements of differentiable programming, 2024. URL attts://arxiv.oom abs 2403. 14606
[7] Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In International conference on artificial intelligence and statistics, pages 880-889. PMLR, 2018.
[8] Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. Fast differentiable sorting and ranking. In International Conference Machine Learning, pages 950-959. PMLR, 2020.
[9] Clément Bonet, Laetitia Chapel, Lucas Drumetz, and Nicolas Courty. Hyperbolic sliced- wasserstein via geodesic and horospherical projections. In Proceedings 2nd Annual Workshop on Topology, Algebra, and Geometry in Machine Learning (TAG-ML.), volume 221 of Proceed- ings of Machine Learning Research, pages 334-370. PMLR, 2023.
[10] Nicolas Bonneel and David Coeurjolly. Spot: sliced partial optimal transport. ACM Transactions on Graphics (TOG), 38(4):1-13, 2019.
[1]] Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasser- stein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22-45, 2015.
[12] Xiongjie Chen, Yongxin Yang, and Yunpeng Li. Augmented sliced wasserstein distances. In International Conference on Learning Representations, 2022.
[13] Ho Kei Cheng and Alexander Schwing. The curse of conditions: Analyzing and improving optimal transport for conditional flow-based generation. arXiv preprint arXiv:2503. 2025.
[14] Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9): 1853-1865, 2016.
[15] Thomas M Cover. The number of linearly inducible orderings of points in d-space. SIAM Journal Applied Mathematics, 15(2):434-439 1967.
[16] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
[17] Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. NeurIPS, 35: 267100, 2022.
10

---

[18] Ishan Deshpande, Yuan- Ting Hu, Ruoyu Sun, Ayis Pyrros. Nasir Siddiqui. Sanmi Koyejo. Zhizhen Zhao, David Forsyth, and Alexander G Sehwing. Max-sliced wasserstein distance and its use for gans. In Proceedings of the IEEF/CVE conference on computer vision and pattern recognition, pages 10648-10656, 2019.
[19] John Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
[20] Kilian Fatras, Younes Zine, Rémi Flamary, Remi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. In International Conference Artificial Intelligence and Statistics, pages 2131-2141. PMLR, 2020.
[21] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, al. Pot: Python optimal transport. Journal Machine Learning Research, 22(78):1-8, 2021.
[22] Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016.
[23] Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608- 1617. PMLR, 2018.
[24] Niels Richard Hansen. On stein's ' unbiased risk estimate for reduced rank estimators. Statistics & Probability Letters, 135:76 82, 2018.
[25] Niels Richard Hansen and Alexander Sokol. Degrees of freedom for nonlinear least squares estimation. arXiv preprint arXiv: (rXXiv:402.2997, 2014.
[26] Ernst Heintze and Hans-Christoph Im Hof. Geometry of horospheres. Journal of Differential geometry, 12(4):481 1-491, 1977.
[27] Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS), volume 37, pages 199-201, 1942.
(28) Isay Katsman, Eric Chen, Sidhanth Holalkere, Anna Aseh, Aaron Lou, Ser Nam Lim, and Christopher M De Sa. Riemannian residual neural networks. Advances Neural Information Processing Systems, 36:63502-63514, 2023.
[29] Max Kochurov, Rasul Karimov, and Serge Kozlukov. Geoopt: Riemannian optimization in pytorch, 2020.
[30] Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Gener- alized sliced wasserstein distances. Advances in neural information processing systems, 32, 2019.
[31] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers features from tiny images. 2009.
[32] Alan Lahoud, Erik Schaffernicht, and Johannes Stork. Datasp: differential all-to-all shortest path algorithm for learning costs and predicting paths with context. In Unceriainty in Artificial Intelligence, pages 2094-211 22122 PMLR, 2024.
[33] John M Lee. Riemannian manifolds: an introduction to curvature, volume 176. Springer Science Business Media, 2006.
[34] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In ICLR, 2023. URL /RLptttp////ooeeeeviee http:/pparreviiieeeeefrrrrrrr net/forum? id=PqvMRDCJTTE
[35] Xinran Liu, Rocío Díaz Martín, Yikun Bai, Ashkan Shahbazi, Matthew Thorpe, Akram Aldroubi, and Soheil Kolouri. Expected sliced transport plans, 2024. URIpsttps: : //arxiv. org/ abs/
2410. 12176]

---

[36] Guillaume Mahey. Unbalanced and linear optimal transport for reliable estimation of the Wasserstein distance. PhD thesis, INSA Rouen Normandie, 2024.
[37] Guillaume Mahey, Laetitia Chapel, Gilles Gasso, Clément Bonet, and Nicolas Courty. Fast opti- mal transport through sliced generalized wasserstein geodesics. Advances in Neural Information Processing Systems, 36, 2024.
[38] Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pages 666-704, 1781.
[39] Eduardo Fernandes Montesuma, Fred Maurice Ngole Mboula, and Antoine Souloumiac. Recent advances in optimal transport for machine learning, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
[40] Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli. Statistical and topological properties of sliced probability divergences. Advances in Neural Information Processing Systems, 33:20802 12, 2020.
[41] Khai Nguyen. Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced-wasseree wasserstein and applications to generative modeling. In International Conference Learning Representations, 2021. URL StRRl.ttps:////tt http:::s///ppeeeeieieeee http://wwwwweettteeeeteeeee mett//frrumidddy(0oo0ooooo
[42] François- Pierre Paty and Marco Cuturi. Subspace robust wasserstein distances. In International conference on machine learning, pages 5072-5081. PMLR, 2019.
[43] Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In ICML, pages 737-746, 2016.
[44] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 11(5-6)355-607, 2019.
[45] Aram-Alexandre Pooladian, Heli Ben- Hamu, Carles Domingo-Earich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning, pages 28100-28127. PMLR, 2023.
[46] Michael Puthawala, Konik Kothari, Matti Lassas, Ivan Dokmanié, and Maarten De Hoop. Globally injective relu networks. Journal of Machine Learning Research, 23(105): 2022.
[47] Julien Rabin, Julie Delon, and Yann Gousseau. Regularization of transportation maps for color and contrast transfer. In 2010 IEEE International Conference on Image Processing, pages 1933-1936. IEEE, 2010.
[48] Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernt. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 201 I, Ein-Gedi, Israel, May 29-June 2, 201 Revised Selected Papers 3, pages 435-446. Springer, 2012.
[49] Julien Rabin, Sira Ferradans, and Nicolas Papadakis. Adaptive color transfer with relaxed optimal transport. In 2014 IEEE international conference on image processing (ICIP), pages 1852-4856. IEFE, 2014.
[50] Filippo Santambrogio. Optimal transport for applied mathematicians, volume 87. Springer, 2015.
[51] Charles Stein. Estimation of the mean of a multivariate normal distribution. The Annals of Statistics, 96): 160-1135-1111 1151, 1981.
[52] Eloi Tanguy, Rémi Flamary, and Julie Delon. Reconstructing discrete measures from projections. Consequences on the empirical Sliced Wasserstein Distance. Comptes Rendus. Mathématique, 362: 162:1121-1129 2024.
[53] Ryan Tibshirani. Degrees of freedom and model search. Statistica Sinica, pages 1265-1296. 2015.
12

---

[54] Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024.
[55] Thanh Tran, Viet-Hoang Tran, Thanh Chu, Trang Pham, Laurent El Ghaoui, Tam Le, and Tan M Nguyen. Tree-sliced wasserstein distance with nonlinear projection. arXiv preprint arXiv:2505.00968, 2025.
13

---

A Appendix supplemental material
A.1 Proofs
We start by the proof of Lemmal concerned with 0-homogeneity of our target 0 + h(0).
Proof of Lemmal / We prove in fact that for all > 0, n7 x - This is a consequence the fact that if 0 (xx, 8) is 1-homegennous, then does 0 (f). Thus, the level set of π g(T, cB) the level-set of + g 0) dilated by factor Hence, they share the same minimizers, and in consequence h(cf) = h(0).
For the sake of completeness. we also prove the comment stating that gradient flow on h preserves the norm of the initialization with the following lemma
Lemma 3. Let U IR' an open set, assume that U differentiable and h is O-homegeneous, i.e., h(c0) = h00). Consider the gradient flow dynamics
0(0) =0 EU ((t) -
(9)
Then, there exists an interval S Rzo and a unique solution t O(t) such that for all t e (0(t), ((t)) = 0 and [[())]] = 0g-
Proof. Orthogonal gradient of O-homogeneous function. Consider the real function ui R + R
defined by (c) (c0). By 0-homogeneity, (c()) h(0) (1).. Hence, s is constant R" thus (c) = 0 for all + 0. Using the chain rule for real function, we have that for all * 0
* (c) = ((c (c)), Vh(co)) = (0, Vh(co)) =0 O.
Using this fact for 1, we conclude that
v0 EEU, U, (0, Vh(0)) 0.
(10)
Orthogonal dynamics. The existence and uniqueness of the Cauchy problem (9 comes from the Cauchy-Lipschitz theorem. Consider this solution t + 0(t) defined over /. Then,
(é(1), 0(1) = (V)(0()), = 0,
Hence, (tt) (tt for all t using 10
Conservation of the norm. Consider r(t) = [|0(t)]1. The chain rule tells us that for all t 1,
r (t) 2(0()),(()),
hence r' 0 and thus is constant function.
The proof of Proposition D relies on the fact that the p- Wasserstein distance metric on ID measures, and that ø" is conveniently supposed to be injective over the reference set X.
Proof of Proposition Let 1, R?, and assume that 0² is injective map from X X2? to IR Let μ, & three discrete distributions in P(x)). Recall that
d (p, v) Wp Wp(dfu, v5v).
Using the fact that Wp is a metric on (R)), we also obtain that Wp metric on P(X)) as restriction.
Well-posed Since d (t, p) W,(0gu, Lfn), and that W, is metrie, have that d8 (f,, 11)
0.
Symmetry. The symmetry comes directly from the of W
14

---

Positivity. Suppose that f * v. Since Wp is metric, we only need to prove that 0ff1 # D.V. Assuming that, where # By for all i o i' and yj # yy for all # j'.
m
1
azS2; and
j=1
we have that
a:d (x))) and 0. = Σ by (500"(())
1=1
Using the injectivity of o., we have that @ (x:) # (xx) for all i + 1' and (v;)) * a (0 for all
j + j'. Hence, of = of if, and only if, n = m and there exists permutation a : {1,.... {1...., such that
a: = bi)) and d' = (yo(()) for all i. But then, using the injectivity of d² have {a:li= (uieer Hence, which is contradiction.
Triangle inequality. We have that d (H, = Wp 0.v). Using the triangle inequality on Wp. we have that d (u, )) S Wp(ofu, 2f5) + W Wo, olv) d du, (), (E,
We now turn to the Stein's ss lemma. The proof of the Stein's lemma under a (weak) differentiability cri- terion [5T] is classic, and relies on an integration by part and the properties of the normal distribution. Nevertheless, we concerned with function 0 - (0) that typically will have discontinuities, breaking the classical proof. Note that one cannot expect the Stein's ' lemma to hold true for any kind of discontinuities,
with almost everywhere differentiabilitty The celebrated example is the Heaviside function h(0) = 1820 in ID where the Stein's 's lemma needs a correction term if there is a non-negligible number of them in sense of the H Hausdorff dimension. This setting was studied for various applications in statistics and signal processing, in particular for Stein's 's unbiased risk estimation.
by (25 53 24
Proof of Lemma 2 The proof of this result is mostly contained in [24] Proposition 1]. We outline the strategy. Assumption requires to have H (R2 C*) 0. Hence, for all {l,...,q] and Lebesgue almost all (fI..... 0,-1, 01+1. 0g) E IR- the map
1 (0 O, 0.,
absolutely continuous on every compact interval of IR. So, in turns, h E Wilel (RR) which in turns show that h is almost differentiable in the sense of Stein [5]] and we can thus apply [5]] Lemma 1].
We now turn to the proof of Proposition 2 regarding the properties of the smoothed version he of h.
Proof of Proposition
(Admissibility)) Let >> 0 and R2, Recall that
: Ez-N(O,L.)) arg min g( 0 + +z2) (a,b)
Denote by u(z) arg (() ( 0 ez), hence o. E BZz For all z E R9, u(2) (a, b) by definition of the minimization problem. Hence,
JR u(=)p(z)dz,
where p(z) = -22))=/2eepp(
is the probability density function of the multivariate normal law and dz is the Lebesgue measure on Rs. Since U(a, b) is convex and u(z) E U(a,b) b). then Sa u(z)p(2)dz E U(a,b) b)
also. In turn, since Th U(a, b). then given the true solution पहं of (T). we have E EØ= Tor lx, uyllg < E Σ -., - 13
(Differentiabiliyyy) This a direct consequence of Lemma and Lebesgue dominated convergence theorem to invert expectation and derivative.
15

---

(Consistency.) The first fact is a consequence of Lebesgue dominated convergence theorem. The second one use the expression the gradient through the variance reduction expression:
h(6 EZ) h(0) Z
Vgh+(0) = EE Ex + h(0))
Ez
E
Hence, again using Lebesgue dominated convergence theorem, we have
hhe =Z) h(0) 7 =Ez
h(0 + (10+=Z)= h(0))
lim V Vahg (0) lim Ez c0
lim
c-0
E
E
Recognizing the directional derivative of h(0) h is differentiable at 0, we get that lim (f) = Ez (Vh(()) Z)Z] Vh(0).
(Distance.) We assume here that Φ is injective on X. We split the proof for h (1.) and hgg (2.).
1. Proof that (H, v) + (u, v) is distance over P(X). The positivity comes from the subop- timality of 7 v), that is (10) v) 2 Wg(u, v) > 0 if + (as W, is metric itself). The symmetry comes from the fact that Cyu is symmetric and that T (H, v) (z²(1, the well-posedness, since @º is injective, then W Wo dg = 0 and T (m, p1) is the identity matrix. Hence,
Regarding
h(0) = (Caps x" (p, = ΣΣ (C (m
(( =0,
2=1
since for all i, (Cma)is = |:: - xi = Concerning the triangle inequality, let μis U2, Mz E P(X). Let us denote
π 12 = T T (p1, 12) E R R2ixx
n 13 = R (M1, E R
T 23 T (p2, E R
Using the specific structure of the ID optimal transport [50]. there exists a tensor II E R X of order 3 such that admits 12 18 and 28 as marginals, that is
Vi, j.
12
>>
T ji
Vi, k, The Vj. k, T5 -=l
12
Lig.
I.j..b
Since this structure provides us "gluing lemma", we continue the proof similarly the standard proof of the triangular inequality of the Wasserstein distance.
(h(0)(1,,s)) = ((Cpyngs π 101) 1/p
1/p
13 Tik
10 - 2k 2x
by definition
1=1
1/2
Tigxlla; - -2x1BB
as glue.
Using that ||:; - 11 < llac - ys 1i + 1185
II we get that
1/p
(MO))HI)aa)) <
ΣΣΣ
I1 [1
wslls + 11 =x|1p)
Applying now the Minkowski inequality, we obtain that
1/P
1/2
ns
(4)) <
ΣΣΣ
Hy 12
2oll
+
Σ
2x
16

---

Using the fact that II has marginals π and 2² we get that
1/p
1/p
R1
F2
ng
(M)))
12 17 20
Vsl1
+
Jk
Hence,
(h (h(0) (a)) (h(0)(m),
(M,0)(1,12))44
2. Proof that (p. 2) + he (2) is distance over P(2'). The well- posednesss comes from the fact The symmetry also a direct consequence of the symmetry of (1,) v) and the positivity comes from the fact that the expectation of a positive quantity positive (understood almost surely on IR9). For the triangle inequality, we use the linearity of the expectation: let M2, 12,148 13 E P(2). Then, for all 0 E IR, and for all E IRy using the fact that h00 + æz) is distance
that he (0) p)) Ez = p)] [z[0]
h(0 + (2), pa) & h(0 + ez) p2) + 1:0 (2), PP)
Hence, taking the expectation and using linearity gives that
Ez|h(0 + =Z) (B /r) - Ez [h(h(0 + =Z) E2lh(0 =Z) us)
A.2 Algorithm
Algorithm D
describes gradient descent method perform the minimization of hig using the Monte-Carlo approximation from Eq. (8)
Algorithm Monte-Carlo gradient descent (0)
Require: 8g E IR,, step size policty 7). smoothing parameter
0, number Monte Carlo samples
N, number of iterations T
1: for /= 0 to T
do Sample i.i.d. perturbation vectors ;IN (O, Idg))
2:
3:
for = to N do Solve OT problems obtain getez and x
4:
using ID OT solver
5:
gk - (Car, π tezz - x()
6:
Vhe,N (0.) EN Ek=1 Ik
N
approximate gradient
7:
Bet1 t- 0₁ Tne
update parameter
8: return Or
A.3 Additional results and experiment details
All experiments except the Conditional Flow Matching (CFM) were run on MacBook Pro M2 Max
with 32 GB of RAM. On this machine, Fig. took approximately 3 minutes per run (10.000 iterations).
Fig. (3 about 6 minutes for 10 runs (with two models trained sequentially. 1,000 iterations). Fig. 4 required roughly 30 minutes, and Fig. 5 took around 10 minutes total (all models considered. 10 repetitions). The CFM experiments were dispatched over GPU cluster composed of (PPUAA100 80G, GPU-A6000 48 Go, with total runtime of 130h for training and inference all presented models.
We estimate that the total compute time over the course of the project -ineetiinciinngg experimentation. debugging, and hyperparameter tuning-is approximately two orders of magnitude larger than the reported runtimes for CPU-based experiments, and one order of magnitude larger for the GPU-based experiments.
A.3.1 Hyperparameter settings
We report here the hyperparameter configurations used across the main experiments. Figures and correspond to the same experiment-Figg highhlights early training dynamics, while Fig. [ results at convergence. The projection network used is 3-layer MLP with ReLU activations: (with
17

---

Varianoo Recluction With Vartence Reduction
Withoat
9
(C
dus
200
Figure 7: Impact of the variance reduction scheme from Eq. Here, the same learning rate used for both variants, in which case the variant without variance reduction does not even converge.
dimensions 2 - 64 + 16 + 1). Optimization is done using SGD with a learning rate 0.2; for the variant without variance reduction, a lower learning rate of 0.0002 is used ensure convergence (cf. Fig. 7 in which the same learning rate is used for both variants). In Figure 4 (gradient flow experiments),
perform 2000 outer flow steps using SGD with learning rate of 0.01. At each flow step, we execute 20 projection steps (or inner optimization updates when using learnable projectors). For the latter, we use Adam with learning rate of 0.01. The neural projector for our method is hayde--aadde MLP with ReLU activations and He initialization. In Figure which investigates gradient flows hyperbolic manifolds, vary the outer learning rate across methods to account for differences in convergence speed: the base learning rate is 2.5, used for HHSW; SW uses scaled learning rate of 17.5, and DGSWP uses a reduced rate of 0.83. Each flow step is composed of 100 projection or inner optimization steps. For the Conditional Flow Matching (CFM) experiment shown Figures 6, g and Table [ we adopt the same training hyperparameters as Tong al. [54] For our method specifically, the projection model is 3-layer fully connected network with SELU activations:
32 x 32 + 256 256 > Its parameters are optimized using Adam with learning rate of 0.01. We perform 1000 optimization steps for the projection model at initialization, followed by step per CFM training iteration.
A.3.2
Additional results
Fig g presents a set of generated images using I-CFM, OT-CFM and DGSWP (NN) after 400k iterations using the 100-step Fuler integrator.
(a) 1-CFM
(b) OT-CFM
(c) DGSWP-CFM
Figure 8: Example of generated images after 400k steps.
Fig. presents the evolution the image generation quality during training when using the adaptive- step Dormand Normand-Prince strategy for the integration.
18

---

1-CMM
min-DGSWP-CFM (linear)
CFM
min DGSWP CFM (NN)
FID
5
4
100k
200k
300k
400k
Training Step
Figure 9: FID as function of training iterations for various algorithms using adaptive-step DoPri5 sampling.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide theoretical but also experimental results that support the claims made in the abstract and introduction. That is to say we "reformulate min-SWGG as bilevel optimization problem and propose differentiable approximation scheme efficiently identify the optimal slice, even in 3high-dimmnnsonnll 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". Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We address the main limitations of our work as future works the conclusion: "Future works will explore the conditions that the projection map ó must satisfy to preserve [...] statistical guarantecs. Additionally, ensuring the injectivity of ó essential for DGSWP define proper distance. Investigating injective neural architectures [...] is promising direction for future research." We also discuss the computational efficiency into the appendix and deseribe the full experimental setup. Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
19

---

The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-texl system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
The proofs are given in the supplementary material.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide the entire implementation, that allows running the all set of experiments and reproduce the illustrative figures of the paper, in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
20

---

If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide the entire implementation, that allows running the all set of experiments and reproduce the illustrative figures of the paper, in the supplementary material. We provide the entire commands or notebooks for ease reproductibility. Datasets used for the experiments are public datasets; for toy examples, code for producing the dataset is also provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips.cc//ubllc//uuieeee http://oooSSuuiiissee wwi.bmmiiisiooooiiii.icc.c for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
21

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting//etails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We present the experimental setting in the core of the paper, and describe thoroughly the parameter values in the supplementary material. The parameters are also given into the code.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For most of results, we provide several runs our results. For the OT-CFM experiments, we run the experiments only once, as in (54), to avoid extensive use of computational resources. Note that, in this case, do not claim any statistical difference but rather propose qualitative evaluation of the results.
the variability of
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The type of compute workers that have been used is described into the supplementary material.
22

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compule required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't I make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
httc:ppubbic//tiiiccsuddeeeeeeid ines?
Answer: [Yes]
Justification: We have conducted the research conform in every respect with the NeurIPS code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no clear societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
23

---

Answer:
[NA]
We rely on CIFARIO and synthetic datasets, which limits the risk of misuse. Guidelines:
Justification:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-us should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We mainly use the POT library (21) (under MIT licence) and torchCFM [54) (MIT licence). We also use the code of [9] for which no information of licence is provided. All these resources are cited in the paper. We also use the CIFARIO dataset [31], for which the official page does not mention any licence.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, haperswithcode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide the implementation of the method.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.

---

14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
25


Reviews:
[{"review_id": "Z48Qd4TLay", "reviewer": "Reviewer_kzcn", "summary": "The 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.", "questions": "Questions:\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", "limitations": "see weakness and questions", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \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.", "was_revised": true, "final_justification": "authors have addressed my concerns in a sufficient way, I raise the scaore to 5", "rebuttal": "We would like to thanks the reviewer for these positive comments on our work, and the careful reading. We will improve the manuscrit based on _all_ your minor comments, in particular regarding the presentation of the different objects introduced in our paper. We address in the following some specific comments that may merit discussions or clarifications.\n\n**Regarding injectivity**: We recall that there is no injective map from $\\mathbb{R}^d \\to \\mathbb{R}$ for $d > 1$, and we emphasize that require only injectivity from the (potential) support to $\\mathbb{R}$. It is true that asking for injectivity of a fine grid would have numerical consequences, but here we merely need it on the real support of the transport problem. We would like to add that this condition is satisfied almost surely in the linear case, and could be also satisfied with injective neural network [1]. Note that we require this property to prove that we obtain a metric, the object $d^\\theta$ _still_ make sense without this hypothesis. In particular, we did not observe performance collapse due to noninjective $\\phi$ map. From a numerical point of view, we simply rely on the order given by the sort algorithm used by Python that can be interpreted as a selection map from the set of minimizers.\n\n[1] Puthawala, Michael, et al. \"Globally injective relu networks.\" _Journal of Machine Learning Research_ 23.105 (2022): 1-55.\n\n**Regarding the differentiability**: Thanks for this comment. You're absolutely right. Formally, we should work in the context of Bolte & Pauwels [2] with conservative Jacobians. Since this kind of notions is far away from the interest of the OT community of NeurIPS, we believe that it is better to keep a differentiable hypothesis to simplify the exposition. We will add a comment in the revision in the direction of a formal proof with conservative Jacobians (that are exactly design to tackle for instance ReLU functions).\n\n[2] Bolte, Jérôme, and Edouard Pauwels. \"Conservative set valued fields, automatic differentiation, stochastic gradient methods and deep learning.\" _Mathematical Programming_ 188.1 (2021): 19-51.\n\n**Regarding practical benefits over SWGG**. You are right, we erroneously introduce ASWD as AWD in Sec. 2.2 which impairs understanding. This will be fixed in the next iteration of the paper.\nRegarding comparison between min-SWGG (optim) and DGSWP (linear) at equal timings (sorry, figures are not allowed), we get (note that the iteration for DGSWP (linear) is chosen here so as to lead to a timing lower or equal to that of the corresponding min-SWGG run): \n\n| Dataset | Method | Iteration | $\\log_{10} W$ | Timing |\n|---|---|---|---|---|\n| Swiss Roll | min-SWGG (optim) | 500 | **-0.81** | 3.34 |\n| Swiss Roll | DGSWP (linear) | 173 | -0.49 | 3.34 |\n| Swiss Roll | min-SWGG (optim) | 1000 | **-0.89** | 6.67 |\n| Swiss Roll | DGSWP (linear) | 345 | -0.84 | 6.67|\n| Swiss Roll | min-SWGG (optim) | 1500 | -0.80 | 10.00 |\n| Swiss Roll | DGSWP (linear) | 516 | **-1.07** | 9.99 |\n| Swiss Roll | min-SWGG (optim) | 2000 | -0.79 | 13.33 |\n| Swiss Roll | DGSWP (linear) | 687 | **-1.28** | 13.32 |\n| Gaussian 2d | min-SWGG (optim) | 500 | **-1.42** | 3.40 |\n| Gaussian 2d | DGSWP (linear) | 172 |-0.40| 3.38 |\n| Gaussian 2d | min-SWGG (optim) | 1000 |**-2.01**| 6.81 |\n| Gaussian 2d | DGSWP (linear) | 342 |-1.12| 6.80 |\n| Gaussian 2d | min-SWGG (optim) | 1500 |**-1.96**| 10.25 |\n| Gaussian 2d | DGSWP (linear) | 514 |-1.75| 10.24 |\n| Gaussian 2d | min-SWGG (optim) | 2000 |-2.28| 13.67|\n| Gaussian 2d | DGSWP (linear) | 686 |**-2.46**| 13.66 |\n\nWe will add such numbers in the next iteration of the paper and discuss the fact that min-SWGG can be more efficient early in the process and DGSWP (linear) becomes more accurate as running time grows.\n\n**Regarding the comparaison with Expected Sliced Wasserstein [35].** When considering ESW as an additional competitor, we get the results from Tables S.1 and S.2 below:\nFor the dataset in Fig. 1, varying the temperature parameter $\\tau$ for Expected Sliced Wasserstein, we get the following results in terms of transportation cost:\n\n_Table S.1_\n\n| Method | $\\langle \\pi, C\\rangle$ cost |\n|---|---|\n| Exact OT | 13.50 |\n| min-SWGG | 22.22 |\n| Expected Sliced Wasserstein ($\\tau=0$) | 32.28 |\n| Expected Sliced Wasserstein ($\\tau=1$) | 21.30 |\n| Expected Sliced Wasserstein ($\\tau=10$) | 20.09 |\n| DGSWP (NN) | 13.80 |\n\nFor the gradient flow experiment in Fig. 4, we get the following performance after 2,000 iterations for the Swiss Roll dataset (note that results are consistent across all 4 datasets).\n\n_Table S.2_\n\n| Method | $\\log_{10} W$ (at convergence) |\n|---|---|\n| min-SWGG (optim) | -0.79 |\n| ESW ($\\tau=0$) | -0.07 |\n| ESW ($\\tau=1$) | -0.08 |\n| ESW ($\\tau=10$) | -0.13 |\n| DGSWP (NN) | **-5.83** |\n\nIn other words, in the experimental setup of Fig. 1 and Fig. 4, ESW is consistently outperformed by both min-SWGG and DGSWP in terms of transportation cost and convergence of the corresponding gradient flows (varying the hyper-parameter from $\\tau=0$ to $\\tau=10$).\nBased on these initial comparisons on toy datasets, we have decided not to include ESW in the comparison for larger-scale problems such as the CFM experiment.\n\n**Regarding function evaluations.** In the CFM context, at generation time, given a model for $v(x, t)$, different solvers can be used to integrate the path from $x_0$ to $x_1$. It is common in the CFM litterature to evaluate the complexity of the generation steps as the number of calls to $v$ required for the integration [13, 34, 54]. \nNote that in all our experiments, $v$ is implemented using a UNet architecture, which is the same across methods. In practice, between I-CFM, OT-CFM and DGSWP-CFM, the only difference lies in the training procedure that is used to train $v$.\nAs a consequence, NFE, which is the number of times the model $v$ is called is the adequate measure of complexity for the generation step.\n\n**Regarding local minima w.r.t $\\theta$.** You are absolutely right, DGSWP is non-convex and it converges towards a local minima. Even with simple scenario with 4 samples can lead to non-convex objective function. To assess the impact of the initialization, we perform several runs for each experiments. We notice that, in practice, results are similar across runs, even for the CFM experiment with is in dimension $d> 3000$.\n\n**Regarding positive curvature.** It is indeed a very interesting question. DGSWP necessitates to define a $\\phi$ function that projects the sample $\\in \\mathbb{R}^d \\to \\mathbb{R}$, and uses a sliced mechanism to obtain the permutations (i.e. it suffices to sort the samples according to their projection). It is not straightforward to adapt DGSWP to the context of spherical slicing: for instance, [4] projects on a circle (and not on a line), which necessitates to employ the SW distance on the circle; [5] uses stereographic projections that could be used but is not naturally rotationaly invariant. We believe that some adaptations of DGSWP should allows it to be adapted in this context but we leave it for future work.\nRegarding positive curvature spaces, and again relying the formulation of [6] to define manifold feature maps $\\mathcal{M} \\to \\mathbb{R}$, one could naturally use DGSWP to compare SPD matrices (see [7] for an example with sliced-Wasserstein).\n\n[4] Bonet, C., Berg, P., Courty, N., Septier, F., Drumetz, L., & Pham, M. T. Spherical Sliced-Wasserstein. ICLR.\n\n[5] Tran, H., Bai, Y., Kothapalli, A., Shahbazi, A., Liu, X., Diaz_Martin, R. P., & Kolouri, S. (2024). Stereographic Spherical Sliced Wasserstein Distances. ICML.\n\n[6] Katsman, I., Chen, E., Holalkere, S., Asch, A., Lou, A., Lim, S. N., & De Sa, C. M. (2023). Riemannian residual neural networks. NeurIPS.\n\n[7] Bonet, C., Malézieux, B., Rakotomamonjy, A., Drumetz, L., Moreau, T., Kowalski, M., & Courty, N. (2023). Sliced-Wasserstein on symmetric positive definite matrices for M/EEG signals. ICML"}, {"review_id": "VoWDfBBRqJ", "reviewer": "Reviewer_Q6uB", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# 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", "was_revised": true, "final_justification": "The authors addressed all my questions. I believe the paper addresses an important question in the literature of sliced optimal transport.", "rebuttal": "We would like first to thanks the reviewer for their positive comments, acknowledging the quality of the writing and the novelty. We would like to address your main comments below.\n\n**Regarding statistical significance.** We have provided gradient flows results with several repetitions. For CFM, the paper contains results only for one run. We performed additional runs to better assess the variability of the results. Below, we report the average +- std for 3 runs, which keeps the conclusion unchanged. \n\n| |  Euler  (FID) | DoPri5 FID (NFE) | \n|:--------------- |:--------:| ---------------:| \n| I-CFM | 4.61 +- 0.06 | 3.74 +- 0.04 (6896 +- 47) | \n| OT-CFM | 4.75 +-0.06     | 3.81 +- 0.04 (6490 +- 36) | \n| DGSWP-CFM (linear) |  4.16 +- 0.13  | 4.45 +- 0.06 (5556 +- 108) | \n| DGSWP-CFM (NN)               |   **3.67 +- 0.03**      |   3.95 +- 0.06 (5696 +- 64)              | \n\n**Regarding Expected Sliced Optimal Transport.** When considering ESW as an additional competitor, we get the results from Tables S.1 and S.2 below:\nFor the dataset in Fig. 1, varying the temperature parameter $\\tau$ for Expected Sliced Wasserstein, we get the following results in terms of transportation cost:\n_Table S.1_\n| Method | $\\langle \\pi, C\\rangle$ cost |\n|----|----|\n| Exact OT | 13.50 |\n| min-SWGG | 22.22 |\n| Expected Sliced Wasserstein ($\\tau=0$) | 32.28 |\n| Expected Sliced Wasserstein ($\\tau=1$) | 21.30 |\n| Expected Sliced Wasserstein ($\\tau=10$) | 20.09 |\n| DGSWP (NN) | 13.80 |\n\nFor the gradient flow experiment in Fig. 4, we get the following performance after 2,000 iterations for the Swiss Roll dataset (note that results are consistent across all 4 datasets).\n\n_Table S.2_\n| Method | $\\log_{10} W$ (at convergence) |\n|---|---|\n| min-SWGG (optim) | -0.79 |\n| ESW ($\\tau=0$) | -0.07 |\n| ESW ($\\tau=1$) | -0.08 |\n| ESW ($\\tau=10$) | -0.13 |\n| DGSWP (NN) | **-5.83** |\n\nIn other words, in the experimental setup of Fig. 1 and Fig. 4, ESW is consistently outperformed by both min-SWGG and DGSWP in terms of transportation cost and convergence of the corresponding gradient flows (varying the hyper-parameter from $\\tau=0$ to $\\tau=10$).\nBased on these initial comparisons on toy datasets, we have decided not to include ESW in the comparison for larger-scale problems such as the CFM experiment.\n\n**Regarding control variate for Monte Carlo estimation of gradient.** Thanks for this comment. We believe that the line of works that you mention is very interesting. Overall, we kept a discussion on the variance reduction aspect quite minimalistic, both from a theoretical and practical aspects. We believe that finding better, or even optimal, variance reduction method working alongside with Stein's approximation is an interesting topic of future works.\n\n**Regarding assessing the quality of a transportation plan.** Assessing the quality of a transportation plan is hard. In this work, we do so by computing the cost associated to a plan, ie. $\\langle \\pi, C\\rangle$ (cf. Fig 1). When doing so for the dataset in Fig. 1, varying the temperature parameter $\\tau$ for Expected Sliced Wasserstein, we get the  results obtained above in Table S.2:\nNotably, we observe that the cost for ESW is consistently larger than that for min-SWGG (which is expected, since ESW averages plans over drawn directions while min-SWGG retains the direction of minimal cost), and in practice larger than that of DGSWP (NN).\nNote that one could argue that assessing the similarity of a transport plan with respect to the ground truth one could be done by computing the matrix norm of the differences between plans, but we observe in practice that such a distance tends to favor uninformative uniform plans over more meaningful sharp ones, which is why we avoid such metrics in our work.\nIf the reviewer has suggestions for an alternative way to assess the quality of an estimated transport plan, we would be happy to provide such metrics for the compared methods during the author-reviewer discussion period."}, {"review_id": "0tRrjId1dd", "reviewer": "Reviewer_WETw", "summary": "The paper introduces a differentiable approximation of OT distance and plan using generalized Sliced-Wasserstein projections and Sliced-Wasserstein Generalized Geodesics framework.", "questions": "Regarding 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.", "limitations": ".", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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", "was_revised": true, "final_justification": "I thank the authors for their response. I believe all of my concerns have been adequately addressed. I am now satisfied with the clarifications and revisions provided, and I have updated my score in favor of acceptance.", "rebuttal": "We thank the reviewer for their comments. We remain available throughout the discussion period for any further clarifications, and we hope you will consider increasing your overall evaluation in the light of these clarifications.\n\n**Regarding the difficulty of our contribution.** We thank the reviewer for highlighting the connection with GSW  [37] and minSWGG  [30].  While our approach builds on these ideas, and we obviously do not deny it, it is not a simple juxtaposition. The key contribution is to view this generalization of minSWGG as a bilevel optimization problem, and to carefully design a gradient-like method to solve it. It is well known that bilevel optimization in the context of non-strongly convex inner problem is difficult (see e.g., [1]) and we believe that our approximation results of GSWP are non-trivial. We agree that the proofs of Lemma 1 and Proposition 1 do not present technical difficulties, but we believe it is still important to formally state them, as they are necessary first steps before stating results on DGSPW. Despite these simple proofs, the approximation results that follow in Section 4 are less classical, and we believe that the proof of Proposition 2 contains interesting arguments (especially the fact that we obtain a gluing lemma \"for free\"). \n\n[1] Bolte, J., Lê, Q. T., Pauwels, E., & Vaiter, S. (2025). Geometric and computational hardness of bilevel programming. Mathematical Programming, 1-36.\n\n**Regarding the injectivity condition.** We would like to recall that there is no injective map from $\\mathbb{R}^d \\to \\mathbb{R}$ for $d > 1$, and we emphasize that we require only injectivity from the (potential) support to $\\mathbb{R}$. This condition is satisfied almost surely in the linear case (all points in general configuration), and could be also satisfied with injective neural network [2]. Note that we require this property to prove that we obtain a metric, the object $d^\\theta$ _still_ make sense without this hypothesis. In particular, we did not observe performance collapse due to noninjective $\\phi$ map, and we believe that it due that it is unlikely that the training procedure select a MLP that is degenerate with regard to this property. We do not know however to prove such result, since it involves structural properties of the training dynamics of neural networks. From a numerical point of view, we simply rely on the order given by the sort algorithm used by Python that can be interpreted as a selection map from the set of minimizers.\n\n[2] Puthawala, Michael, et al. \"Globally injective relu networks.\" _Journal of Machine Learning Research_ 23.105 (2022): 1-55.\n\n**Regarding the choice of Conditional Flow Matching.** The goal of the experimental section is to assess the performance of DGSWP in various contexts. The first experiment aims to provide a qualitative evaluation of the plan, while the second examines the convergence properties of DGSWP. Lastly, we seek to evaluate our method in a setting where a transport plan is required, and the number of samples is large. In this context, mini-batch OT is the main competitor, as, except for the expected sliced Wasserstein (see General comment), no sliced methods provide a transport plan. For these reasons, and to challenge DGSWP in a context where slicing methods have not been considered before, we focus on CFM, as [3] is well-positioned in the literature. Therefore, the primary motivation was not to deliver state-of-the-art performance on a generative modeling task but rather to investigate the performance of DGSWP across different settings. We propose to clarify this in the final version of the paper.\n\n[3] Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., ... & Bengio, Y. (2024). Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 1-34.\n\n**Regarding runtime comparison with OT solvers.** You are perfectly right, we have not explicitly stated the complexity for our DGSWP method in the paper. We will fix that in the final version of the manuscript upon acceptance.\nLet $n$ be the number of samples in each distribution, $d$ be the dimensionality of the input space, $m$ be the number of Monte Carlo samples used for the gradient estimation in Eq. (8) and $p$ the number of gradient steps required to reach convergence, then the time complexity for the linear DGSWP is $O(p m n \\log n + p m n d))$.\nUpon acceptance of the manuscript, we will add a visualization of the corresponding running times for both linear and NN variants of DGSWP.\n\n**Regarding CIFAR-10 evaluation.** Results were originally reported for one single run; we performed additional experiments to report results computed on 3 runs:\n    \n|                    |   Euler  (FID)   |           DoPri5 FID (NFE) |\n|:------------------ |:----------------:| --------------------------:|\n| I-CFM              |   4.61 +- 0.06   |  3.74 +- 0.04 (6896 +- 47) |\n| OT-CFM             |   4.75 +-0.06    |  3.81 +- 0.04 (6490 +- 36) |\n| DGSWP-CFM (linear) |   4.16 +- 0.13   | 4.45 +- 0.06 (5556 +- 108) |\n| DGSWP-CFM (NN)     | **3.67 +- 0.03** |  3.95 +- 0.06 (5696 +- 64) |\n\nTo achieve these results, we performed the computations on a cluster consisting of multiple GPUs (utilizing only 2 GPUs for each experiment) with different GPU types, using a job scheduler based on resource reservation and availability. As a result, it is challenging to provide a definitive answer to the question about the required runtime. However, we can assert that the CFM experiment with mini-batch OT and DGSWP has roughly the same runtime (as mentioned in the paper: \"aggregate samples from 10 minibatches to compute the transport plan, and find that the additional computational cost remains negligible\"). As for I-CFM, the running time is approximately 25% lower than that of OT-CFM."}, {"review_id": "HY3UUTw2Be", "reviewer": "Reviewer_8PTq", "summary": "The 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.", "questions": "- 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", "limitations": "The 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$.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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}$).", "was_revised": true, "final_justification": "I believe the paper is worthy of acceptance and further the authors have address the points raised in the review. I will maintain my rating at 5 for acceptance.", "rebuttal": "We are grateful for the reviewer's very positive feedback and are pleased that our work was well received. We would like to take this opportunity to address your comments and improve the manuscript based on your thoughtful suggestions.\n\n**Regarding a better exposition of our contributions.** We will clarify our exact contribution more clearly. We want to emphasize that Definitions 1 and 2 are **new** objects that were never introduced in previous works. More precisely, [37] introduced SWGG that only tackles linear projection $\\phi(x,\\theta) = \\langle x, \\theta \\rangle$, whereas [30] proposed _generalized sliced Wasserstein_ that is a distinct metric from the one proposed in our submission. We will rephrase the first paragraph of Section 3 to better clarify this point.\n\n**Regarding comparison with other slicing methods.** Except expected sliced Wasserstein (which we did not consider in the CFM experiment due to performance issues), sliced-based methods do not provide a transport plan. For CFM, getting such a plan is mandatory to sample trajectories. Then, they cannot be considered in this context. \nOne major strength of DGSWP is its ability to provide good results in high dimensional problems, you are absolutely right. In the CFM experiments, the samples are set of images, one sample being one image seen as a vector of dimension 32 pixels $\\times$ 32 pixels $\\times$ 3 colors = 3072, which remains a challenging setting. Because we do not have access to extensive computation resources, we did not consider higher dimensional settings.\n\n**Regarding NFE values.** We run additional experiments for a range of NFE values, which confirms the good performance of DGSWP. The following table gives the average values of FID over 3 runs +- std:\n\n| Euler  (FID) for -> |     NFE = 20     |         NFE = 50 |        NFE = 100 |\n|:------------------- |:----------------:| ----------------:| ----------------:|\n| I-CFM               |   8.33 +- 0.09   |     5.49 +- 0.02 |     4.61 +- 0.06 |\n| OT-CFM              |   8.34 +- 0.07   |     5.78 +- 0.03 |     4.75 +- 0.06 |\n| DGSWP-CFM (linear)  |   8.57 +- 0.13   |     5.49 +- 0.01 |     4.16 +- 0.13 |\n| DGSWP-CFM (NN)      | **7.94 +- 0.17** | **4.72 +- 0.04** | **3.67 +- 0.03** |\n\nRegarding the ability of DGSWP to provide straight plans, it is indeed a desiderable behavior that is worth being investigating. Nevertheless, we believe that this question deserves a deeper investigation before raising complete conclusions, and we stick to the goal of this paper, which is to assess the ability of DGSWP to provide a meaningful plan. \n\n**Regarding evolutionary strategies.** One of the points of our submission is to show that a typical gradient descent-type algorithm can be used to compute DGSWP. We fully agree that different methods could be considered, for instance evolutionary strategies or reinforcement learning methods. We want to point out that [37] used a 0-order method to compute minSWGG (that boils down to GSWP with a linear projection). We will add a comment that mention the work you indicated as a potential alternative to design an algorithm to solve the bilevel problem behind GSWP.\n\n**Regarding scaling.** This is indeed an interesting question, thank you. It has been shown that, provided that i) we are able to sample the direction $\\theta$ that minimizes SWGG ii) the samples are in general position (see [2]), min-SWGG is equal to Wasserstein when the dimension $d\\geq 2n$. As DGSWP allows optimizing over $\\theta$ (or a non-linear version of it), it is therefore natural to wonder under which conditions one could recover this result. However, it is hard to give a definite answer to it, for several reasons:\n- as we face a non-convex problem, it is hard to guarantee that we will converge to the global optima\n- to our knowledge, there is no network that guarantees that the embeddings will be in general position\n        \nNevertheless, one can rely on this result to design a neural network architecture that is likely to provide good results, that is to say expressive enough to provide an embedding space in a large enough dimension. We propose to add a discussion about this on the paper.\n   \n\n[2] Cover, T. M. (1967). The number of linearly inducible orderings of points in d-space. SIAM Journal on Applied Mathematics, 15(2), 434-439."}, {"review_id": "SV7r9xhSuO", "reviewer": "Reviewer_enMR", "summary": "This 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.", "questions": "- What is the sample complexity for the gradient estimation of h, is this better than the previous work min-SWGG?", "limitations": "The 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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "- 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.", "was_revised": false, "final_justification": "", "rebuttal": "We would like to first thanks you for your praises of our work such as finding our bilevel formulation \"_highly novel_\" and our numerical experiments \"_convincing evidence of the method’s effectiveness\"_. We discuss below the main point that you addressed in your rebuttal regarding sample complexity.\n\n**Regarding sample complexity.** We are not entirely sure we fully understand your question. We provide a tentative answer addressing two different points, regarding two different sample complexities, and we hope this will offer enough clarity. We remain available during the discussion period to better develop one of these two points.\n\n- _Regarding the sample complexity of the Monte Carlo approximation of the gradient_.\nThe Stein's approximation through DGSWP is done through a standard Monte Carlo approximation of the expectation. In this sense, the convergence rate of $\\hat \\nabla h_{\\varepsilon, N}$ towards $\\nabla h_\\varepsilon$ has a rate of $O(1/\\sqrt{N})$ where $N$ is the number of samples used (and crucially does _not_ depends on the dimension of the features). Note that more refined use of Monte Carlo methods, for instance by using MCMC methods could lead to a better rate such as $O(1/N)$. We leave this issue to future works.\n\n- _Regarding the sample complexity of (sliced-)Wasserstein and DGSWP_. \nThe study of the statistical properties of (sliced-)Wasserstein has been the focus of considerable research in recent years. One significant result is that most sliced-based methods (including min-SWGG, see [1]) have a dimension-free sample complexity, which is not true for Wasserstein. This raises the question of the sample complexity of DGSWP. At this point, we do not have a definitive answer, as it depends on the function $\\phi$ used to project the samples. If we consider a straight line, the dimension-free sample complexity of min-SWGG holds. Otherwise, it depends on the expressiveness of the neural network. For example, if the network can project samples onto a space of dimension $d \\geq 2n$ where the samples are in general position, we would recover the sample complexity of Wasserstein. Then, we hypothesize that the resulting sample complexity will depend on the choice of the $\\phi$ function. This is a very interesting question that we plan to explore in the near future.\n\n[1] Mahey, G. (2024). Unbalanced and linear optimal transport for reliable estimation of the Wasserstein distance (Doctoral dissertation, INSA Rouen Normandie)."}]

Author rebuttal:

Meta-review: The authors proposed an approach to calculate a sliced wasserstein plan. They generalized a standard approach based on a linear projection to more general functions including neural networks. Also, for efficient usage of SGD they proposed a trick to calculate gradients using Stein’s lema.

As a strengths of the paper I would highlight 
- The approach works well compared to similar approaches of such kind,
- The approach is generalized  to non-linear projections via neural networks and to manifold-valued data 
- The paper is well written

As a weaknesses I would highlight
- To some extent the approach is a combination of existing ideas (GSW and min-GSW)

The main reasons for the accept decision are
- Sufficiently novel generalization of existing algorithm
- Good experimental verification
- Sufficiently clearly written text

During the discussion with the reviewers several issues were raised that should be addressed in the final version of the paper. Here is the list of the main issues
- Improve the text and clarify the main contributions, provide information about computational performance
 (reviewer kzcn, reviewer WERw)
- More extensive experiments with error bars and comparison with expected sliced OT (reviewer Q6uB)
- Compare the performance of other slicing methods on the flow matching experiment (reviewer 8PTq)
Decision: Accept (poster)