Convergence of SGD for Training Neural Networks with Sliced Wasserstein Losses
Authors: Eloi Tanguy
Venue: Accepted by TMLR ()

Abstract: Optimal Transport has sparked vivid interest in recent years, in particular thanks to the Wasserstein distance, which provides a geometrically sensible and intuitive way of comparing probability measures. For computational reasons, the Sliced Wasserstein (SW) distance was introduced as an alternative to the Wasserstein distance, and has seen uses for training generative Neural Networks (NNs). While convergence of Stochastic Gradient Descent (SGD) has been observed practically in such a setting, there is to our knowledge no theoretical guarantee for this observation. Leveraging recent works on convergence of SGD on non-smooth and non-convex functions by Bianchi et al. (2022), we aim to bridge that knowledge gap, and provide a realistic context under which fixed-step SGD trajectories for the SW loss on NN parameters converge. More precisely, we show that the trajectories approach the set of (sub)-gradient flow equations as the step decreases. Under stricter assumptions, we show a much stronger convergence result for noised and projected SGD schemes, namely that the long-run limits of the trajectories approach a set of generalised critical points of the loss function.

Published in Transactions on Machine Learning Research (11/2023)
Convergence of SGD for Training Neural Networks with Sliced Wasserstein Losses
Eloi Tanguy
eloi. tonguy@u-paris
Université Paris Cité, CNRS, MAP5, 7505006 Paris, France
Reviewed on OpenReview: https: openreview. net/ forun? id= aqgfB3p9ZA
Abstract
Optimal Transport has sparked vivid interest in recent years, in particular thanks to the Wasserstein distance, which provides a geomelrically sensible and intuitive way of comparing probability measures. For computational reasons, the Sliced Wasserstein (SW) distance was introduced as an alternative to the Wasserstein distance, and has seen uses for training gen- erative Neural Networks (NNs). While convergence of Stochastic Gradient Descent (SGD) has been observed practically such setting, there is to our knowledge no theoretical guarantee for this observation. Leveraging recent works on convergence of SGD on non- smooth and non-convex functions by BBiancii al. (2022), gap; and provide realistic context under which lixed-slep SGD trajectories for the SW loss on NN parameters converge. More precisely, we show that the trajectories approach the set of (sub)-gradient flow equations as the step decreases. Under stricter assumptions, we show a much stronger convergence result for noised and projected SGD schemes, namely that the long-run limits of the trajectories approach set of generalised critical points of the loss function.
aim to bridge that knowledge
1 Introduction
1.1 Optimal Transport in Machine Learning
Optimal Transport (OT) allows the comparison of measures metric space by generalising the use of the ground metric. Typical applications use the so-called 2- 2-Wasserstein distance, defined as
Vx,y P2(TR), W2(x,y) inf xEll(xy) Rd
|| y yl2dn(z,
(W2)
where (IRd)) is the set of probability measures on Rd admitting a second-order moment and where 1l(x, y) is the set of measures of P2(Rd x RR)) first marginal x and second marginal y. One may find a thorough
presentation of properties in classical monographs such as Peyré & Cuturi (2019); Santambrogio (2015)
Villani 2009g
The ability to compare probability measures is useful in probability density fitting problems, which are sub-genre of generation tasks. In this formalism, considers probability measure parametrised by vector u which is designed to approach a target data distribution (typically the real-world dataset). In order to determine suitable parameters, one may choose any probability discrepancy (Kullback-Leibles Ciszar divergences, f-divergences or Maximum Mean Discrepancy (Gretton et al., 2006)), or in our case, the Wasserstein distance. In the case of Generative Adversarial Networks, the optimisation problem which trains the 'Wasserstein GAN' (Arjovsky al. 2017) stems from the Kantorovitch-Rubiin Rubinstein dual expression of the 1-Wasserstein distance.

---

Published in Transactions on Machine Learning Research (11/2023)
1.2 The Sliced Wasserstein Distance Alternative
The Wasserstein distance sulfers from the curse of dimensionality, in the sense that the sample complexity for n samples in dimension d is of the order ((n'/d) (Dudley 1969). Due to this practical limitation and to the computational cost of the Wasserstein distance, the study of cheaper alternatives has bocome a prominent field of research. A prominent example is Entropic OT introduced by Cuturi (2013), which adds an entropic regularisation term, advantageously making the problem strongly convex. Sample complexity bounds have been derived by Genevay et al. (2019), showing a convergence in ((vn) with constant depending on the regularisation factor.
Another alternative is the Sliced Wasserstein (SW) Distance introduced by Rabin et al. 2012), which consists in computing the 1D Wasserstein distances between projections of input measures, and averaging over the projections. The aforementioned projection of a measure x on Id is done by the push-forward operation by the map Po x + zzz. Formally, # is the measure on R such that for any Borel set B C R, PB = x (B)) Once the mcasures are projected onto line RA, the computation of the Wasscrstein distance becomes substantially simpler numerically. We illustrate this fact in the discrete case, which arises in practical optimisation settings. Let two discrete measures on Fd ūx := 1 Ex Ozg TY == - E supports X = (I1, Inn and Y= (V1, Un) E Raxd Their push-forwards by Pe are simply computed by the formula Po#ox = 1 Σx (pw)))) and the 2- Wasserrstninn distance between their projections can be computed by sorting their supports: let permutation sorting (0 "0",,"...,,0) 0TTx),,, and permutation sorting (#Ty,. ",.....8 _ Vn). one has the simple expression
E 6yg with
ToT Wa = -(OT [o(k) - 0T v=()))²
The SW distance is the expectation of this quantity with respect to 0 o, i.c. uniform on the sphere: Sw2 av) = Ee-o W The 2-SW distance is also defined more generally between two measures x, E P2(
SW /
JOESH-1 W Po#y)do(0).
(SW)
In addition to its computational accessibility, the SW distance enjoys a dimension-free sample com- plexity (Nadjahi al. (020)). Additional statistical, computational and robustness properties of SW have been explored by Nietert al.] (2022). MMreover, central-limit results have been shown by Xu & Huang 2022) for 1-SW and the 1-max-SW distance (a variant of SW introduced by Deshpande 2019 and related work by Xi & Niles- Weed (2022) shows the convergence of the sliced error process Vn (Wg (Potoooxx Po
al.
Po)y)), where the samples X ~ g and Y ~ y are drawn for each 0. Another salient field of research for SW is its metric properties, and while it has been shown to be weaker than the Wasserstein distance in general by Bonnotte (2013), and metric comparisons with Wasserstein and max-SW have been undergone by Bayraktar Guo (2021) and Paty & Cuturi (2019).
1.3 Related Works
Our subject of interest is the theoretical properties of SW as a loss for implicit generative modelling, which lcads to minimising (T, in the parameters u, where y is the target distribution, and Ta#x the image by the NNH of x, low-dimensional input distribution (often chosen as Gaussian or uniform noise). In order to train NN in this manner, at each iteration one draws samples from x and y (denoted 8x and by as discrete measures with n points), as well projection 0 (or batch of projections) and performs SGD step on the sample loss
[(u) = SW3 (Po Po#or) = - (0 TT 0 (r))))
1 n k=1
(2)
Taking the expectation of this loss over the samples yields the minibatch Sareed--asersttiin discrepancy, a member of the minibatch variants of the OT distances, introduced formally by Fatras et al. Fatras et al.
Similarly the ID case, #a is the push-forward measure of by Tu, f.e. the law of Tu(x) when
2

---

Published in Transactions on Machine Learning Research (11/2023)
2021p. b. The framework (2) fits several Machine Learning applications, for instance, Deshpande et al. (2018)
trains GANs and auto-encoders with this method, and Wu et al. (2019) consider related dual formulation Other examples within this formalism include the synthesis of images by minimising the SW distance between
features of the optimised image and a target image, as done by Heitz et al. (2021) for textures with ncural leatures, and by Tartavel al. 2016) with wavelet. features (amongst other methods).
The general study of convergence of SGD in the context of non-smooth, non-convex functions (as is the case of L from [2) is active field of research: Majewski et al. (2018) and Davis et al. (2020) show the convergence of diminishing-stey SGD under regularity constraints, while Bolte & Pauwels (2021) leverage conservative field theory to show convergence results for training with back-propagation Finally, the recent work by Bianchi et al. (2022) shows the convergence of fixed-step SGD schemes on a general function F under weaker regularity assumptions.
More spocifically, the study of convergence for OT-based generative NNs has been tackled by Fatras et al. (2021], who prove strong convergence results for minibatch variants of classical OT distances, namely the Wasserstein distance, the Entropic OT and the Gromov Wasserstein distance (another OT variant introduced by [émolij ()011)) ). A related study on GANs by Huang et al. (2023) derive optimisation properties for one layer and one dimensional Wassorstein-GAN and generalise to higher dimensions by turning to SW-GANs. Another work by BBreheet et al. (2023) focuses on the theoretical propert of linear NNs trained with the Wassersessr loss (introduced by Bures [1699;; see also (Bhatia et al. 2017 for reference on this metric). Finally, the regularity and optimisation properties of the simpler energy SW2 have been studied by Tanguy al. (0023).
In practice, it has been observed that SGD in such scttings always converges (in the loose numerical sense, see (Deshpande et al. 2018), Section 5, or (Heitz et al. 2021), Figure 3), yet this property is not known theoretically. The aim of this work is to bridge the gap between theory and practical observation by proving convergence results for SGD on (minibatch) Sliced Wasserstein generative losses of the form F(u)= = E S...
1.4 Contributions
Convergence of Interpolated SGD Under Practical Assumptions Under practically realistic as- sumptions, we prove in Theorem that piecewise affine interpolations (defined in Equation [10] of constant- step SGD schemes on u ++ F(u) (formalised in Equation (7) converge towards the set of sub-gradient flow solutions (see Equation 9 as the gradient step decreases. This results signifies that with very small learning
rates, SGD trajectories will be close to sub- aabdgradient flows, which themselves converge to critical points of F
(omitting serious technicalities).
The assumptions for this result are practically reasonable: the input measure and the true data measure
are assumed to be compactly supported. As for the network (u,x) + T(u, x),, we assume that for a fixed datum x, T() is piecewise C2-smooth and that it is Lipschitz jointly in both variables. We require additional assumptions on T which are more costly, but are verified as long as T is a NN composed of typical activations and linear unils, with the constraint that the parameters u and data x slay both stay within a fixed bounded domains. We discuss a class of neural networks that satisfy all of the assumptions of the paper in the Appendix (Section D). Furthermore, this result can be extended to other orders p + 2 of SW: we present the tools for this generalisation in Section E
Stronger Convergence Under Stricter Assumptions In order to obtain stronger convergence result, consider variant of SGD where each iteration receives an additive noise (scaled by the learning rate) which allows for better space exploration, and where cach itcration is projected on ball B(0,r) in order to ensure boundedness. This alternative SGD scheme remains within the realm of practical applications, and we show in Theorem 2) that long-run limits of such trajectories converge towards a set of generalised critical points of F, as the gradient step approaches 0. This result is substantially stronger, and can serve as an explanation of the convergence of practical SGD trajectories, specifically towards a set of critical points which amounts to the stationary points of the energy (barring theoretical technicalities).
3

---

Published in Transactions on Machine Learning Research (11/2023)
Unfortunaelly, require additional assumptions in order to obtain this stronger convergence result, the most important of which is that the input data measure and the dataset measure y are discrete. For the latter, this is always the case in practice, however the former assumption is more problematic, since it is common to cnvision generative NNs as taking an argument from a continuous space (the input is often Gaussian of Uniform noise), thus a discrete setting is a substantial theoretical drawback. For practical concerns, one may argue that the discrete x can have an arbitrary fixed amount of points, and leverage strong sample complexity results to ascertain that the discretisation not costly if the number of samples is large enough.
2 Stochastic Gradient Descent with SW as Loss
Training Sliced- Wassersessn generative models consists in training neural network
T: Rdd x Rdd
Rdy
(u,x)
:-
by minimising the SW minibatch loss u m E Ex &-yan Descent (as described in AAgoriihhm 1). The probability distribution x E P2((R)) the law of the input of the generator T(u,)) -)- The distribution y P2((R))) is the data distribution, which T aims to simulate. Finally, a will denote the uniform measure on the unit sphere of 1 denoted by gdv-1, Given a list of points == (z1,-.. ,.n) E R" R= denote the associated discrete uniform measure ox := E, 6a.- By abuse of notation, we write (X) = (T. (x1)... ... , TT (xx)) E Raxd, The reader may find a summary of this paper's notations in Table
or)] through Stochastic Gradient
Algorithm 1: Training a NN on the SW loss with Stochastic Gradient Descent
Data: Learning rate a > 0, probability distributions x E (2 and y E P2(4))
Initialisation: Draw (0 Rdd; for t E [00, Tmax 1] do
Draw (((-1) ~ 0, (((+1) ~ En y(+1) you, SGD update:
m((+1) (0))
W3 W2 (P #Ta #Ttoxui,, Pau Porwnn)
Du
end
In the following, we will apply results from (Binnchi et al. 2022), and we pave the way to the application of these results by presenting their theoretical framework. Consider a sample loss function f Rdd x 2 - R that is locally Lipschitz in the first variable, and z a probability measure on C Rd which is the law of the samples drawn at each SGD iteration. Consider sp Rdd 2 - Rdd an ehmost-ereryyhhrr gradiend of f, which is to say that for almost every (u, 2) Rd. x Z, p(u,2) = 0.f(u, 2) (since each f(.,2) is locally Lipschitz, it is diffcrentiable almost-cverywhere by Rademacher's theorem). The complete loss function is the expectation of the sample loss, F
Sz su, =)de(z). An SGD trajectory of step > 0 for F is a
sequence (1) E (Rd.)) y of the form:
n (1+1) = ( ()) , (((+1)),
(=())N)))) un @ SN
where up is the distribution of the initial position ()) Within this framework, we define an SGD scheme described by Algorithm 1 with z = w0n @ By0r & a and the minibatch SW sample loss
Rdd Baxd Raxdy -
Rd.
(4)
(Y,
W Po
With this definition for f, we have
F(u) -(xxx,)) (m, X. V. = ((xx))=aass
or)]
(5)
4

---

Published in Transactions on Machine Learning Research (11/2023)
thus the population loss compares the 'true' data with the model's generation Tu Tu# using (minibatch) SW. We now wish to define an almost-everywherr gradient of f. To this end, notice that one may write (,))) (x,, where for X, E guxd, and 0 E 844-1, wa(XXY) = We(Po Pr) The differentiability propertics of wo(- Y) are already known (Tanguy et al. 2023 particular one has the following almost-everywhheet gradient. v))
Bonneel et al. 2015), in
Owa 0x
2 200 (IN MoMn-M0g -oxxx(1)) -16p)r()))
(X,Y)
E RRnd,
kE[1,n]
X.Y E G, is TΦ (T with 7f Gn being a sorting permutation of the list where the permutation 0g"
(0 yı...., 0 yn). The sorting permutations are chosen arbitrarily when there is ambiguity. To define an almost-everywhere gradient, we must differentiate f(.. = u + wn(T(u, X), for which we need regularity assumptions on T: this is the goal of Assumption 1 In the following, A denotes the topological closure of a set A, 0A its boundary, and lgd,, denotes the Lebesgue mcasure of Rd..
Assumption 1. For every x E IRd,, there exists a family of disjoint connected open sets (U)(z)) such that Vj E J(x), T(,x) E (U) U Ug(x) Rdd and hygas jEJ(a)
(U 04;(o)) =0.
JEJ(z)
Note that for measure-theoretie reasons, the sets (zx are assumed countable. One may understand this assumption broadly as the neural networks being piecewise smooth with respect to the parameters u, where the pieces depend on the input data œ. In practice, Assumption iss an assumption on the activation functions of the neural network. For instance, it is of course satisfied in the case of smooth activations, or the common case of piecewise polynomial activations. We detail suitable neural networks in the Appendix Section D).
Assumption implics that given X, Y,Y fixed, f(,, X.Y. is differentiable almost-everywhere, and that one may define the following almost-everywhere gradient (6]
Rdd Ruxd= x Raxd, x 8d
Rdd.
5p:
(,,X,Y,0)
2 OT (u, ark)
(1,
00 O0 (T(u, yt.xe. (N))
(6)
Ou
k=1
AT where for E Rdd,, du (u, x) Rdy denotes the matrix of the diffcrential of u m T(u, x), which is defined for almost-every u. Given u (z) (a point of potential non-diferrnttabilitt)). take instead 0. (Any choice at such points would still define an a.c. gradient, and will make no difference).
Given a step a > O, and an initial position u ~ vo, we may now define formally the following fixed-step SGD scheme for F:
(1+1) = (0 app("!") x(++1), y(++1) g(++1)), Jue (rE)) (0()) un @ xN
An important tochnicality that we must verify in order to apply Bianchi et al. 2022 's results is that u. + [u, X, Y,O and F are locally Lipschitz. Before proving those claims, we reproduce a useful Property from Tanguy et al. 2023). In the following, ||X | denotes max kE[1,n] and By(xr for A' a norm on Rdd,, x E Rdd and r > 0 shall denote the open ball of Rdd. of centre x and radius r for the norm N (if A is omitted, then B is an euclidean ball).
|| given X (#1, ,Ix) E Radd
Proposition 1. The (we(,, (o)(,,))) ()ss)) are uniformly locally Lipschitz (Tanguy et al. 2023) Prop. 2.2.1. Let Kir X, Y) = 4n(r + X , 2 + ||]lo,,)), for X, v,Ye xxd,, and r 0. Then wol, Y) is Kw X, Y)- Lipschilz in the neighbourhood B la (X,r):
AA y" E Bl (X,r), 40 E gd |Wo(Y', "oo""" 5 (r, (,,,,,,)''' Y"
5

---

Published in Transactions on Machine Learning Research (11/2023)
In order to deduce regularity results on f and F from Proposition 1 we will make the assumption that T is globally Lipschitz in (u. In practice, this is the case when both parameters are enforced to stay within a fixed bounded domain, for instance by multiplying a typical NN with the indicator of such set. We present this in detail in the Appendix (Section Dj.
Assumption 2. There exists L > O such that
v(u1, u2, 11, 12) E (R&x)² (R=x)² )2, (T( T(22, x2)ll S L (1lu1 u2l12 + || 22l12).
Proposition 2. Under Assumption 2 for > O, uo E Rd,, X E Ruxds Y E Baxdy and A E Sdh-1, let Ky no, X) = 4Ln(EL 1. + +17 X) lloc, + [|||,,,2) Then f(, XY,0) is Ky(e, uo, X, -IIipschill in B(
Vu,u's E B(u,,e), X.Y,0) [u,, X,Y,0) S K Ky(e, uo x,Y)lu u'112.
Proof. Let e e> O, un Rd., X Raxd. Y Ruxdy and 0 gh.-1. Let u, u' E B(uo, Using Assumption 2. we have T(u, X), T(w,, X) E B b1 (T(uo, x),r), with r r==el.
Denoting L == B(vo B(0nt ()) we apply successively Proposition (first inequality), then Assump-
Lion (second inequality):
If(u, X,, Y,0 f(,,,x,,,0) X,) = (wo(T X),Y) wo('''',
S (r, T(uo, X), (u,, X) - -T'', (()lo,) S 4n(eT: + [T(mo, l0oo,,,2 + 11 u' 1''12-
Proposition 2 shows that f is locally Lipschitz in We now assume some conditions on the measures x and in order to prove that F is also locally Lipschitz. Specifically, we require that the data measures x and y be supported on bounded domains, which imposes little restriction in practice.
Assumption 3. and y are Radon probability measures on Red and Rd, respectively, supported by the compacts X' and y respectively. Denote Rg := sup | and Ry == sup ||vl|2. =ex
yEy'
Proposition 3. Assume Assumption 2 and Assumption For e > 0, no E Rdd, let Ci(uo) = Jxe
xxm [T(uo, X) 2d, 2xx on (X) and C2 02: == J 2d
Let (o)e,uo) ALm(=L + + C2). We have Vu, u' B(uo,e), [F(u) - F(u') & KF(E, uo) u 2'112.
Proof. Let 0, up Rdd and u, u' B(uo, s). We have
[F(u) F(u'l
Xx If(u, -fa--f//ua X,Y, (X)) J Kyle, uo X, Y) u' 22 ( (X)
< | xxxxx" 4Ln(EL + (T X) + || 21l u'
Now by Assumption
[[uuu X 2ooo, is continuous on the compact X", thus upper-boundee by a cer-
tain M (uo) 0. We can define (wp) xM Since X' is compact and x is a Radon probability measure by Assumption 3 x(X²) is well-defined and finite, thus C1(1g) (ux) finite. Likewise, let C2 = Ym
[ X) xol,,,, 2dx 0n (X), which verifies Ci (wo) < M M(uo)xxX))
l on (Y) +00.
Finally, [F(u) F(v()) & 4Tn(eT++ + (uo) + C2)llu u'll2

---

Published in Transactions on Machine Learning Research (11/2023)
Having shown that our losses are locally Lipschitz, we can now turn to convergence results. These conclusions are placed in the context of non-smooth and non-convex optimisation, thus will be tied to the Clarke sub- differential of F, which we denote 0cF. The set of Clarke uub-gradients at a point u is the convex hull of the limits of gradients of F:
F() = conv { | 3(u!)) E (Dp) u())
and VF(u())
i++00 vj.
(8)
where DF is the set of cifferentiability of F. At points u where is differentiable, OcF(u)= (VF(u)), and if F is convex in neighbourhood of then the Clarke differential at is the set of its convex sub- uub-gradients. The interested reader may turn to Section C for further context on non-smooth and non-convex optimisation.
3 Convergence of Interpolated SGD Trajectories F
In general, the idea behind SGD is discretisation of the gradient flow equation ú(s) FF(((s)). In our non-smooth setting, the underlying contimuous-time problem is instead the Clarke differential inclusion i(s) E (u(ss). Our objective is to show that in a certain sense, the SGD trajectories approach the set of solutions of this inclusion problem, as the step size decreases. We consider solutions that are absolutely continuous (we will write u(-) E Caba (R,, R4x)) and start within I. C Rdd, a fixed compact set. We can now define the solution set formally as
S de F(C)= = {u Cahs (R,,, Vs E E,, (iss -QcF(u())) u(0) E K3.
99)
where we write for "almost every". In order compare the discrete SGD trajectories to this set. of continuous-ting trajectories, we interpolate the discrete points in an affine manner: Equation (10) defines the piecewise-affiee interpolated SGD brajectory associated to a discrete SGD trajectory (uh )rev of learning rate a.
+ ( -1 - -w Vs E [ (1 + 1)ab, VIEN.
(x) = a
(10)
In order to compare our interpolated trajectorics with the solutions, consider the metric of uniform convergence all segments
2 min (1, max [u(s)) u u(())l2))
dc(uu()):= =
Σ KEN
(1,
(11)
In order to prove convergence result on the interpolated trajectories, we will leverage the work of Bianchi et al. 2022 which hinges three conditions on the loss F that reproduce and verify successively. Firstly, Condition assumes mild regularity the sample loss function f-
Condition 1.
i) There exists Rdd + R4 measurable such that each k(u,-) is z z-integrable, and:
E O, Vu, u' B(uo, E), Vz Z, If(u,2) f(u', 2) S k<uo, 2)Ilu uu 1|2
i) There exists u e Idd such that f(u,,) is z-integrable.
Our regularity result on f Proposition allows us to verify Condition 1 by letting == and k<uo, =) == K1 X, Condition fii is immediate since for all u E Rd,, (X,Y.0) + ww(T(u,, Y) is con- timuous in each variable separately, thanks to the regularity of T provided by Assumption 2 and to the regularities of w. This continuity implies that. all f(,, are z-integrable, since x x xa e y yon 69 a is a compactly supported probability measure under Assumption 3 Secondly, [ondition 2 concerns the local Lipschitz constant introduced in Condition 1 it is assumed to increase slowly with respect to the network parameters
Condition 2. The function k of Condition verifies:

---

Published in Transactions on Machine Learning Research (11/2023)
i) There crists c 2 0 such that Vu Rde, / x(u,z)dz(2) c(1 + [ull)).
uekt / x(u, d(z,z) +00.
i) For every compact K. C Rdy, sup
Condition 2 ii) is verified by x given its regularity. However, Condition 2 requires that T(u, x) increase slowly as [ increases, which more costly.
Assumption There erists an x-integrable function g : Fd R, such that Yu E R Va e Rd [T 2)122 S g(z)(1+ [|u]]2).
Assumption satisfied in particular as with bounded activation functions), or if is of the form T(u, = r) (u), i.e. limiting the network parameters u be bounded. This second case does not yield substantial restrictions in practice (sce /Section D for class of NNs that satisfy all of the assumptions), yet vastly simplifics thcory. Under Assumption 4 we have for any Rd,, with x(u,2) Ky(,, X,) from Proposition and C2 from
T( is bounded (which is the case for neural network
Proposition 3
K, u, X, V,d) (X)d S 4Im el. + (1 [Iu]l2) f ke[l.n]
max (x))
x xxxxxxx0
$ c(1 + ||ull)).
As a consoquence, Condition holds under our assumptions. We now consider the Markov kernel associatod to the SGD schemes:
Rdd x B(RR..)
[0,1]
Pa
/ 1n(u auplu, z))dz(z)
u, B
Given E Rd,, (u, :) is probability measure Rd, which dictates the law of the positions of the next SGD iteration (t(11), conditionally to u(t) = u. With dga, denoting the Lebesgue mcasure on I let r := {a E J0, + Vt Agd,. uPa < Ape.. }. F is the set of learning rates a for which the kernel Pa maps any absolutely continuous probability measure u to another such measure. We will verify the following condition, which can be interpreted as the SGD trajectories continuing to explore the entire space for a small enough learning rate a:
Condition 3. The closure T contains 0.
In order satisfy Condition 3, we require an additional regularity condition on the neural network T which we formulate in Assumption
Assumption 5. There erists constant M 0, such that (with the notations of Assumpttoo 1 and As- sumption 3) Vx X, Vj E [(),, Vu (z)),, v(d1, i2 is i4) E [1, x d1,d,1
g² Du, Oui
QPT ou, Ouiz
(T(u,)))e, (T,, x)1x.))))) s M, and
(u,
$ M.
The upper bounds in Assumption bear strong consequences on the behaviour of T for ||u||2 > 1, and are only practical for networks of the form T(u,x) ((,, B(o,)) similarly to Assumption 4 We detail the technicalities of verifying this assumption along with the others in the Appendix (Section
Proposition 4. Under Assumption 1 Assumption 3 and Assumption 5 for the SGD trajectories contains j0, aol, where ag := ((dy) + RRyydyM
7
r
We postpone the proof to /Section B. Now that we have verified [ondition 1 Condition 2 and Condition we can apply (Bianchi al.] 2022), Theorem to F, showing a convergence result on interpolated SGD trajectories.
8

---

Published in Transactions on Machine Learning Research (11/2023)
Theorem 1. Consider a neural network T and measures x, y satisfying Assumption 1 Assumplion 3 Assumplion 4 and Assumplion 5 Let < oo (see Proposition 40.
Assumption 2
Let (u²),, a ej0, ail. N a collection of SGD trajectories associated interpolations. For any compact K. C TRo and any 1) > 0, have:
Consider (ua) their associated
lim
40
@N
2 o @N Se > n) = 0.
(12)
+0 af [0,a]
The distance de is defined in As the learning rate decreases, the interpolated trajectories approach the trajectory set F- which is essentially a solution of the gradient flow equation ú(s) = F()((()) (ignoring the set of non-differentiability which is Aga.- To get a tangible idea of the concepts at play, if F was C2 and had a finite amount of critical points, then one would have the convergence of a solution u(s) to a critical point of F, as 8 - +00. These results have implicit consequences on the value of the parameters at the "end" of training for low learning rates, which is why we will consider variant of SGD for which we can say more precise results on the convergence of the parameters.
4 Convergence of Noised Projected SGD Schemes on F
In practice, it is seldom desirable for the parameters of a neural network to reach extremely large values during training. Weight clipping is a common (although contentious) method of enforcing that T(u, -) stay Lipschitz, which desirable for theoretical reasons. For instance the 1-Wasserstein duality Wasserstein GANs (Arjovsky al. (017)) requires Lipschitz networks, and similarly, wreed--assrrttein GANs (Deshpande ett al.) 2018) use weight clipping and enforce their networks be Lipschitz.
0, consider SGD schemes that are restricted to B(0,r) =: B,, by performing projected SGD. At each step E, we also add a noise ag(t+1), where c((+1) is an additive noise of law < lga. which is often taken as standard Gaussian in practice. These additions yield the following SGD scheme:
Given radius
((+11) n, (u()) ap((l(), x(++1), y(+11) g((+l)) + aae(tt))).
(13)
(w))) (X(") Jee (Y(()) (B()) (e()))e)) up 2 NN @ yoN
2 ca
where Tr RR - B, denotes the orthogonal projection on the ball B, = B(0,r). Thanks to Condition Condition 2 and the additional noise, we can verify the assumptions for (Bianchi et al., 2022) Theorem 4, yielding the same result as Theorem 1 for the noised projected scheme [13] In fact, under additional assumptions, we shall prove a stronger mode of convergence for the aforementioned trajectories. The natural context in which to perform gradient descent is on functions that admit a chain rule, which is formalised in the case of elmost-everywweee differentiability by the notion of path differentiability, as studiod thoroughly in (Bolte& & Pauwels 2021]. We also provide brief presentation in the Appendíx (Section C.T..
Condition 4. F is path differentiable, which is to say that for any u E 8 > O, Vu E F(F(ss))), TT i(ss) = (F u))s.
Rdd), for almost. all
Remark 1. There are alternate equivalent formulations for Condition 4 Indeed, as presented in further detail in Section C. is path differentiable if and only if Fc is a conservative field for F if and only if F has a chain rule for dc (the latter is the formulation chosen above in Condition A).
In order to satisfy Condition 4 we need to make the assumption that the NN input measure x and the data measure y are discrete measures, which is the case for in the case of generative neural networks, but is less realistic for x in practice. We define Σn the n-simplex: its clements are the a E I s.t. Vi E [1,n]. as 2 0
and 2= 1.
Assumption 6. One may write x = Σ En und y Σ by with the coefficient vectors E En b E
k=1
1=1
En,, X == for,... Fn,} } C IRdz and y (ur, } C d
9

---

Published in Transactions on Machine Learning Research (11/2023)
There is little practical reason consider non-uniform measures, however the generalisation to any discrete measure makes no theoretical difference. Note that Assumption 3 is clearly implied by Assumption 6
In order to show that F is path differentiable, we require the natural assumption that each 7(: x) be path differentiable. Since T(..z) a vector-valued function, we need to extend the notion of path-differeniaabliy Thankfully, Bolte & Pauwels (2021) define conservalive mappings for vector-valued locally Tipschitz functions (Definition 4), which allows us define naturally path differentiability of vector-valued function as the path-differentiabiity of all of its coordinate functions. See Secction C.2 for detailed presentation.
Assumption 7. For any x E RRdx, T(.xx) is path differentiable.
Assumption holds soon each the neural network has the typical structure of compositions of linear units and typical activations, was proved by Davis et al. 2020), Corollary 5.11 and Bolte & Pauwels (2021), Section 6.2. We provide more specific class of NNs that are path differentiable and satisfy all our other assumptions in Section D
Proposition 5. Under Assumption
Assumption 6 and Assumption 7 F is path differentiable.
Proof. We shall use repeatedly the property that the composition of path differentiable functions remains path differentiable, which is proved in (Bolte & Pauwels 2021 Lemma 6.
Let :: { Buxdy x Ruxdy Y,Y"
{
R+
Sw RY.) By Tanguy et al.
20231. Proposition 2.4.3, each E(, Y)
is semi-concave and thus is path differentiable (by (Taggyy et al. 2023 Proposition 4.3.3).
Thanks to Assumption 6 x and yeu are discrete measures on Raxds and Rxxdy respectively, allowing one to write x0n Ex6x, and yon E, Then F - Cx. Xx), Y) is path differentiable as sum ((oltte & Pauwels, 2021 Lemma 6) of path differentiable functions.
Corollary 4) of compositions ((Bolte & Pauwels 202]
We have now satisfied all the assumptions to apply (Bianchi et al. 2022 Theorem 6, showing that trajectories of (3 converge lowards to a set of generalised critical points C, defined as
{0} if ||u]|e <r {su 2 0} if |u||2=r
C, = fu= | -oF - -,(())))NN(()) (²(u)= =
(14)
Ø if |ull2 >
where (N(uu) refers to the normal cone of the ball B(0,r) at x. The term (G in (14) only makes a difference in the pathological case ||u||2 =r, which never happens in practice since the idea behind projecting to do
so on a very large ball, in order to avoid gradient explosion, limit the Lipschitz constant and satisfy
theoretical assumptions. Omilting the (u) term, and denoting D the points where F is differentiable, (14) simplifies to C, D {u D FF 0}, i.e. the critical points of F for the usual differential. Like
in Theorem
we let ai ag, where ag is defined in [Proposiiion 4 We have met the conditions to apply Bianchi et al. 2022). Theorem 6, showing a long-run convergence results on the SGD trajectories (13)
Theorem 2.
Consider a neural network T and measures x, y satisfying Assumption 1
Assumption 2)
Assumption Assumption 5 Assumption 6 and Assumption Let ( be SGD trajectories defined by 13 for and E/0,a1! One has
Vn 0, lim
Up
x
ON
SN & ca (d(ul) Cr) ((u(9,c))>0
0.
+000
oE)0,a]]
The distance d above is the usual euclidean distance. Theorem 2 shows essentially that as the learning rate approaches 0, the long-run limits of the SGD trajectories approach the set of C, in probability. Omitting the points of non-differeniiabiiit and the pathological case ||u||| r, the general idea is that u.(() - {u F((u) = 0}. which is the convergence that would be achieved by the gradient flow of F, in the simpler case of e smoothness.
a-3
TTypically referred the set of Karush- Kobn- Tucker points the differential inclusion u(s) E -(gF(((()) NV(a(s)).
10

---

Published in Transactions on Machine Learning Research (11/2023)
5 Conclusion and Outlook
Under reasonable assumptions, we have shown that SGD trajectories of parameters of generative NNs with minibatch SW loss converge towards the desired sub-gradient flow solutions, implying in weak sense the convergence of said trajectories. Under stronger assumptions, we have shown that trajectories of mildly modified SGD scheme converge towards set of generalised critical points of the loss, which provides missing convergence result for such optimisation problems.
The core limitation of this theoretical work is the assumption that the input data measure is discrete Assumption 6), which we required in order to prove that the loss is path differentiable. In order to generalise to a non-discrete measure, one would need to apply or show a result on the stability of path differentiability through integration: in our case, we want show that fx (TT ()d d) is path differentiable, knowing that u ~ ((, X), Y) is path differentiable by composition (see the proof of Proposition 5) for the justification). Unfortumately, in general if each g(-,2) is path differentiable, it is not always the case that 9(; x:)dr is path differentiable (at the very least, there is no theorem stating this, even in the simpler case of another sub-class of path differentiable functions, see (Bianchi et al., 2022). Section 6.1). However, there is such a theorem (specifically (Clarke) 1990), Theorem 2.7.2 with Remark 2.3.5) for Clorke regular functions (see [Section C.3 for a presentation of this regularity class), sadly the composition of Clarke regular functions is not always Clarke regular, it is only known to be the case excessively restrictive cases (see (Clarke) 1990), Theorems 2.3.9 and 2.3.10). Similarly to the continuous case, the simpler generalisation in which has countable support adds substantial difficulty, since all of the typical tools (path differentiability itself, Clarke regularity or even definability (see (Bolte & Pauwels, 2021) Section 4.1 for first introduction) do not have rcadily applicable results for infinite operations, to our knowledge. As a result, we leave the generalisation to a non-discrete input measure x for future work.
Our studies focus on the 2-SW distance, but our results from [ecction 3 can be extended to p E [L, +oo| as presented in the appendix (Section Ep. However, also discussed the Appendix, the generalisation of /eccionn is still an open problem, since it has not yet be proven that X SWE(ex, is path differentiable for p #
This paper studies the use of the average SW distance as a loss, and an extension to related distances would be worth considering. The average SW distance aggregates the projected distances through an expectation, while the closely-related mar-Sliced Wasserstein distance introduced by Deshpande et al. (0199) aggregates the projections via a maximisation on the axis 0 E gd-1, The training paradigm presented in DDehpannde) et al. 2019) differs strongly from our formalism since it applies to GANs, however one could consider an extension of our formalism in which the optimal projection 0 becomes learned parameter of the neural
network. A related extension is the Subspace-Robust Wasserstein distance (Paty & Cuturi 2019 which can take the following formulation
(x max 0-00-a (aac([Q)=
Wa(n /1/2
for which one could consider similar extension where the positive semi-definite Ω becomes learned parameter of T.
Another avenue for future study would be to tie the flow approximation result from Theorem to Sliced Wasserstein Flows (Liutkus et al. 2019 Bonet el al. (022). The difficulty in seeing the differential inclusion (9) as a flow F lies in the non-differentiable nature of the functions at play, as well as the presence of the composition between SW and the neural network T, which bodes poorly with Clarke sub-differentials.
Acknowledgements
We thank Julie Delon for proof-reading and general foodback, as well Rémi Flamary and Alain Durmus for fruitful discussions.
11

---

Published in Transactions on Machine Learning Research (11/2023)
References
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Leurning Research, pp. 214-223. PMLR, 06 11 Aug 2017.
URL https://proceediing mlr. pxeess//0/arjovvkkts 1tmll
Erhan Bayraktar and Gaoyue Guo. Strong equivalence between metries of Wasserstein type. 2021.
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the Bures- Wasserstein distance belween positive definite matrices. arXiv, December 2017. doi: 10.48550/arXiv 1712.01504.
Pascal Bianchi, Walid Hachem, and Sholom Schechtman. Convergence of constant step stochastic gradient descent for non-smooth non-convex functions. Set- Valued and Variational Analysis, 30(3):1117-1147, 2022.
Jérôme Bolte and Edouard Pauwels. Conservative set valued fields, automatic differentiation, stochastic gradient methods and deep learning. Mathematioal Programming, 188:19-51, 2021.
Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Efficient gradient flows in sliced- Wasserstein space. Transactions on Machine Learning Research, 2022.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein barycen- ters of measures. Journal of Mathematical Imaging and Vision, 51(1):22 45, 2015.
Nicolas Bonnotte. Unidimensional and evolution methods for optimal transportation. PhD Thesis, Paris 11, 2013.
Pierre Bréchet, Katerina Papagiannouli, Jing An, and Guido Montúfar. Critical points and conver- gence analysis of generative deep linear networks trained with Bures Wasserstein loss. arXiv preprint arXiv:2303.0300 2023.
Donald Bures. An extension of Kakutani's theorem on infinite product measures to the tensor product of semifinite w-aaggebras. Transactions of the American Mathemalical Sociely, 135:199 212, 1969.
Frank H Clarke. Optimization and nonsmooth analysis. SIAM, 1990.
M. Coste. An introduction to o-minimal geometry, inst. rech. RAAG Notes, Institut de Recherche Mathéma- lique de Rennes, 1999.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Damek Davis, Dmitríy Drusvyatskiy, Sham Kakade, and Jason D Lee. Stochastic subgradient method converges tame functions. Foundations of computational mathematics, 20(1):119- 154, 2020.
Ishan Deshpande, Ziyu Zhang, and Alexander G. Schwing. Generative modeling using the sliced Wasserstein distance. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Solt Lake City, UT, USA, June 18-22, 2018, pp. 3483 3491. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 201109/CVPR.201 2018 2018/htm1/Deshpandee https//asspppanddeeeeeraitteeeee
URL http://openacces: thectf creecvf.comm/ootent.. Generative SeweerrttiveeMdeelnnggg Buiigggu__nnnnnnnnonnnenoo Cvwp__cVP_____ _018_papee html
Ishan Deshpande, Yuan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced Wasserstein distance and its use for gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10648-10656, 2019.
Richard Mansficld Dudley. The speed of mean Clivenko-Cantell convergence. The Annals of Mathematical Slatistics, 40(1):40-50) 50, 1969.
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Mini- batch optimal transport distances; analysis and applications. Xiv preprint 2r 2021.
12

---

Published in Transactions on Machine Learning Research (11/2023)
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of Sinkhorn divergences. In The 22nd international conference on arlificial intelligence and slatistics, pp- 1574 1583. PMLR, 2019.
Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alex Smola. kernel method for the two-sample-probler Advances in newral information processing systems, 19, 2006.
Eric Heitz, Kenneth Vanhoey, Thomas Chambon, and Laurent Belcour. A sliced Wasserstein loss for neu- ral texture syuthesis. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 9412-9420, 2021.
Yu-Jui Huang, Shih-Chun Lin, Yu-Chih Huang, Kuan-Hui Lyu, Hsin-Hua Shen, and Wan-Yi Yi Lin. On characterizing optimal Wasserstein GAN solutions for non-Gaussian data. 2023.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and
Robert Stöter. Sliced- Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In Kamalika Chaudhuri and Ruslan Salakhutdinov (cds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceediiggs of Machine Learning Research, pp. 4104 4113. PMLR, 09 15 Jun 2019.
URL https::/proceeddingg mlr. Dreess/v9//iutkung
html
Szymon Majewski, Blažej Miasojedow, and Eric Moulines. Analysis nonsmooth stochastic approximation: the differential inclusion approach. arXiv preprint arXiv:1805.01916, 2018
Facundo Mémoli. Gromov Wasserstein distances and the metric approach to object matching. Foundations of computational mathematics, 11:417-487, 2011.
Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli. Statistical and topological properties of sliced probability divergences. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 20802 20812. Curran Associates, Inc., 2020. URL https://proceedinng www.eee....ppaeeeeeeeeeeeeeeeeee file/eercee..eecca33333202222222 Papee//766--ppr...ee pdf
Sloan Nietert, Ziv Goldfeld, Ritwik Sadhu, and Kengo Kato. Statistical, robustness, and computational guarantees for sliced Wasserstein distances. Advances in Neural Information Processing Systems, 35: 28179-28193, 2022.
François-Pierre Paty and Marco Cuturi. Subspace robust Wasserstein distances. In International conference on machine learning, pp. 5072-508L. 5081. PMLR, 2019.
G. Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 51(1):1 44, 2019. doi: (0.1561/2000000073. URL http:::://arivoor/abb/////////// (aabb///883.005667
Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernot. Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Fin-Gedi, Tsrael, May 29 June 2, 2011, Revised Selected Papers 8, pp. 435-446 Springer, 2012.
Filippo Santambrogio. Optimal transport for applicd mathematicians. Birkäuser, NY, 55(58-633:94, 2015.
Eloi Tanguy, Rémi Flamary, and Julie Delon. Properties of discrete sliced Wasserstein losses. arXiv preprint arXiv:2307.10352, 2023.
Guillaume Tartavel, Gabriel Peyré, and Yann Gousseau. Wasserstein loss for image synthesis and restoration. SIAM Journal on Imaging Sciences, 944):1726-1755, 2016. doi: 10.1137/16M1067494.. URL https:
//doi. org/ 1177116116000
Lou Van Den Dries and Chris Miller. Geometric categories and o-minimal structures. Duke Mathematical Journal, (4(2):497-54) 540, 1996.
13

---

Published in Transactions on Machine Learning Research (11/2023)
Cédric Villani. Optimal transport: old and new / Cédric Villani. Grundlehren der mathematischen Wis- senschaften. Springer, Berlin, 2009. ISBN 978-3-540-71049-3.
Seiichiro Wakabayashi. Remarks on semi-alge alni-algebraic functions, January 2008. Online Notes,
J. Wu, Z. Huang, D. Acharya, W. Li, J. Thoma, D. Paudel, and L. Van Gool. Sliced Wasserstein generative models. In 2019 IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3708- 3717, Los Alamitos, CA, USA, jun 2019. IEEE Computer Society. doi:
URL
wtt.://....g/....11/.//..... 2tt.://C.....2.....0 00383
Jiaqi Xi and Jonathan Niles-Weed Distributional convergence of the sliced Wasserstein process. Advances in Neural Information Processing Systems, 35:13961-13973, 2022.
Xianliang Xu and Zhongyi Huang. Central limit theorem for the sliced 1- Wasserstein distance and the max-sliced 1-Wasserstein distance. Xiv preprint or Xrviv:2205. 14624, 2022.
A Table of Notations
Table 1: List of Notations
Symbol
Explanation
ox
Given X (#1,''' ; ,nn)
Rnnd Dx 6.,
X
(II, Fn) Raxd, an input data sample of law xun
input data probability measure on R supported on X
Y
(Ui, Un) Raaxy target data sample of law yon
y
larget data probability measure Rd,, supported on y
direction in g4,
uniform measure on S
2 O))
sample in X, Y and 0
2: x8" & y 20 a
probability measure for the samples 2, supported on Z X" ) S4,-
neural network parameters in R.
T(u, X)
neural network function defined in
fu, X, 0,
sample loss function defined in
F(u)
population loss function defined
wa (Y, Y!)
discrete and projected 2- Wasserstein distance W (Pooooy Pa#ty.)
(u, X, K,Y,0)
almost-everywwhere gradient of X,, Y,0) delined in (6)
K w. Ky, KF
local Lipschitz constants of w. ʃ, respectively (see Propositions 1, 2, SGD learning rate; noise level
OÇ a
Yudd; P E hud
Lebesgue measure on R4 measure P absolutely continuous w.r.l. Apa
Clarke differential, defined in
(8)
probability measure of SGD initialisation (1)
Up
(!))
additive noise in Id. at SGD step
c B B111(2 (x., R), B (or R)
additive noise probability measure on Rd
open (resp. closed) ball of centre x and radius R for the norm ||
B Postponed Proofs
Proof of Proposition 4
14

---

Published in Transactions on Machine Learning Research (11/2023)
Proof. Let < l and B (( such that &(B) = 0. We have, with a' == 2a/n. (X,Y,0), z x 80 y 80 0 and 2 == x x 3m x 8d,-1
(B) /
1B
OT
u, mkk)
00 (T(u, (ogie.)s du(u)dz(z)
Σ
/ [=2)d2(2))
JRdu x2
du
rEE,
T (u, (u, Ik)
/
OT
where (x() == Jgt
1B (()) du(u)), with 0r,2 == u
00 (T(, - Vr(x)).
du
Vr.aim
Ele 1 1B (r)) de(u), JEJ (x,(x)
Let T E 6, and (X,Y,0) E Z. Using Assumption 1. separate 1z(2)
where the differentiability structure (U jeex(x obtained using the respective differentiability struc- tures: for each & E [I, Assumption 1 depends on Ik, hence the k indices.
yields structure (U,, (TN))))))) of + Tuu, which
To be precise, define for j (i (n)) E J1(11) ... x Jn(En), (,(X)
Nu, Un (au), and J(X) == k=1 [L, T(; is C² on not necessarily defined on the border (X)), which is of Lebesgue
(0i,... "ii...,)nnee E (1)) x.. x (x.)) (x)(X) 00). In particular, for any Ux))). Notice that the derivatives measure O by Assumption 1 thus the values of the derivatives on the border do not change the value of the integrals (the integrals may have the value +oc, depending on the behaviour of Ør,0, but we shall see that they are all finite when a is small enough).
We drop the T index the notation, and focus on the properties of Φ and & as functions of u. Our first objective to determine constant K 0, independent of u, 2, T, such that is K-Lipschitz on Ug (X).
(x(X)
Rdd
First, let X
OT Ou
00 T(u, xx) The function X is of class C² therefore we
u, rk)
determine its Lipschitz constant by upper-bounding the || 1|2-induued operator norm of its differential,
0x u) Du
1 a
denoted by
Notice that x(u)
(T))))))))
20u
02 TT (x.
Now
(0
du max (11,2))
(0 T(T(
using Assumption and |0,| S 1,
Ouz
dus Suiz
82 Duc, Duis
82
(0 (O))(x,))))
S
Σ
0sg Ju,, Jusz
(l) S dy M.
We obtain that x is dy 2 M Lipschitz.
OT Du
Ow (u)
Second, let w E (x(X)
(u, mw)
00 Vr(k), also of class C². We re-write
Ou
8T
VT(1) T
(u, Rk), and conclude similarly by Assumption 5 that w is (2r(0))l2du M-T.i
Du, Ouiz
Finally, vi E(x -
- wk), and therefore : (Hdd 2 +
with Ry from Assumption
1=1
Oy. tu
We have proven that
(u)
s for any u E UX),), and that K does not depend on X,Y,0, or u.
15

---

E using conservative field Published in Transactions on Machine Learning Research (11/2023)
We now suppose that a' < k: which is to say < zh Under this condition, (X( - Fd is injective.
Indeed, if ((u)) 0(u2), then |lw -u2ll2 a' ((ll)(u1) ((2))l2 I a' KKllu - u2212 thus 11 = u2. Furthermore,
Ov 0Φ (u) = ldg ( (w), with for any u E (x)),,
OΦ (u) is invertible
< 1, thus the matrix
du
Ou
Ou
Ou
(using the Neumann series method). By the global inverse function theorem, Φ ; 0= - xX) is a C³ c'-diffeomorphism.
Using the change-of-variables formula, we have (x,( In(o(u))du(u) = / (uiix) 1 (n(()do)uuu)) (B).
we have now shown that is C-diifeomorphism,, thus since u < 2, d#f < h. (a < B denoting that is absolutely continuous with respect to 3). Since x(B) 0, it follows that the integral is 0 then by sum over j. (z) = 0 and finally uPa O by integration over and sum over 7.
C Background on Non-Smooth and Non-Convex Analysis
This work is placed within the context of non-smooth optimisation, field of study in part introduced by Clarke with the so-called Clarke differential, which we introduced in Equation (8) (see (Clarke) 1990) for a general reference on this object). The purpose of this appendix is to present several adjacent objects that can useful to the application of our results, even though we do not need them in order to prove our theorems.
C.1 Conservative Fields
The Clarke differential 0c of a locally Lipschitz function IR - R (defined in Equation (8) is an example of a set-valued map. Such a map is a function D : R" = IR from the subsets of IR to the subsets of RR, for instance in the case of the Clarke differential, we have the signature dcg Rdd = Rd A set-valued map D is graph closed its graph {(u, v) to be conservative field, when it is graph closed, has non-empty compact values and for any absolutely continuous loop 6 ((O, 1], RP) with 7() 7(1), have
IP, E D(u)}
closed set of IR++ set-valued map D is said
max (f(s), v)ds 0. vED(G(S))
Similarly to primitive functions in calculus, one may define function g : Rd D: 1Rd = ddd up to an additive constant through following expression:
h max vED()(s))
= (0)
(+(s), Vy Cabs ((0, 1], R²) such that 7(0) and >(1)
(15)
In this case, say that is a polential function for the field D. This notion allows us to define new regularity class: a function IRd - IR is called path differentiable when there exists conservative field of
which it a potential. A standard result in non-smooth optimisation is the following equivalence between different notions of regularity:
Proposition Bolte & Pawwels 2021) Corollary 2. Let g Rd - R locally Lipschitz. We have the equivalence between the following statements:
g path differentiable
dcg is conservative field
g has a chain rule for the Clarke differential ac:
Vu E Caba ,R, Rd), Vs 0, Vo E dog(u(s)), v Tū((s) (g us).
(16)
This equivalence justifies the terminology used in Condition The reader seeking a complete presentation of conservative field theory may refer to (Bolte & Pauwels 2021
16

---

Published in Transactions on Machine Learning Research (11/2023)
C.2 Conservative Mappings
The nolion of conservalive fields for real-valued locally Lipschitz functions g I2 - I can be generalised to conservative mappings for vector-valued locally Lipschitz functions g I&P -> IR", which one may see a generalised Jacobian matrix (see (Bolte & Pauwels (221)), Section 3.3 for further details). A set-valued map J I." = IRxP is a conservative mapping for such a if
Vu E Cabs (R+, R²), Vs > O, (g ()) = Mü(t), VM E J(u(s))
(17)
In this case, we shall say that g is path differentiable. Note that if each coordinate function gi is the potential of a conservative field D,, then the set-valued map
[(u)
Vi [1. [1,g1 E Di(u)
is conservative mapping for g (although not all conservative mappings for g can be written in this manner). As consequence, one could interpret (simplistically) voctor-valued path differentiability as coordinato- path differentiabilityy
C.3 Clarke Regularity
Another notion of regularity for locally Lipschitz functions is that of Clarke regularity. Let g R - R and u E R2, is said to be Clarke regular at if the two quantities
g (u: limsup glu' + tv) g(u') and g'(u;v) lim 2
g(u + tv) g(u)
t
exist and are equal for all RR. Note that this notion implies path differentiability by (Bolte & Pauwels 2021 Proposition 2. Clarke regularity is the central concept of Clarke's monograph (Clarke 1990
C.4
Semi-Alge Aegebrlge Functions
In non-smooth analysis, one of the simplest regularity cases the class of semi-algebraic functions, which are essentially piecewise polynomial functions defined on polynomial pieces. To be precise, a set A Rd is semé-algebrate if it can be written under the form
A = 4=Un
UN {u E Pis(t)) < 0, Qis(u) = 0},
where the Pt, and Qij are real multivariate polynomials. A function g RP - R is semi-algebraic if its graph g = {(u, g(u)) RP} is semi-algebraic.
A locally Lipschitz real-valued semi-algebraic function is path differentiable (see for instance [olte & Pauwels, 2021), Proposition 2), and in the light of (Bolte & Pauwels, 2021), Lemma 3, this is also the case in the vector-valued case. Another useful property of semi-algebraic functions is that their class is stable by composition and product. The interested reader may consult (Wakabayashi 2008) for additional properties of semi-algebraic objects, CCotte] 1999. Van Den Dries & Miller, 1996). for a presentation of o-minimal structures, generalisation of this concept.
D Suitable Neural Networks
In this section, we detail our claim that typical NN structures satisfy our conditions. To this end, we define class of practical neural networks whose properties sufficient (not all NNs that satisfy assumptions
17

---

Published in Transactions on Machine Learning Research (11/2023)
are within this framework). Consider T the set of NNs T of the form
T: { Rd x 12d=
R Rd
(u,
T(u, x) B(.. (u) ((o
with Rus Re > 0 and > O. The function 1 (), is smoothed version of the usual indicator function 1Bo.m: it any function that has value 1 B(0, R - E), 0 outside B(0, R + E) and is C²-smooth (see Remark 2 for possible construction). Given that one may take arbitrarily large radii, these indicators are added for theoretical purposes and impose no realistic constraints in practice. Additionally, T = hN, the N-th layer of a recursive NN structure defined by
Rdd x Rd
Rd
ho(u,x) =x,
Vn E [1, hn
(u, x)
an Σ (u)(u))((u,z) + Bau
where:
All functions o, IR - R are -"-smooth, or all locally Lipschitz semi-algebraic activation functions (applied entry-wise). The former condition is satisfied by the common sigmoid, hyperbolic tangent or softplus activations. The latter condition applics to the non-difeerentiable ReLU activation, its 'Leaky ReLU' extension, and continuous piecewise polynomial activations. Note that other non- linearities such as softmax can also be considered under the same regularity restrictions, but we limit ourselves to entry-wise non-linearities for notational consistency.
Each dimension d, is positive integer, with obviously dx dy. the output dimension.
Each An,i linear map: Rd= - xd, which maps a parameter vector u to dn x d, matrix. Since the entire parameter vector u is given each layer, this allows the architecture only use certain parameters at each layer (as is more typical in practice). One may this map 3-tensor of shape (dn, di, du), specified in the formulation
Vu E Rd, vh E Rd. An,(u)h=
(n.) Jida 2/
Rd
(18)
j==1
j1E[1,d.]]
The matrix B, E Rd. xd, determines the intercept from the full parameter vector u.
In this model, each layer depends all the previous layers, allowing for residual inputs for instance. Overall, all typical networks fit this description, once bounded using the indicator functions, with only a technicality on the regularities of the activations which need to be all C²-smooth, or all semi-algebraic. One could extend this class of NNs to those with definable activations within the same o o-minimal structure (similarly to Davis et al. 2020 and Bolte & Pauwels (2021)).
Remark 2. We mention that we may construct a 8 -smooth 1s R) in Rdd explicitly as follows:
R
[0, (R 5)² |
e if> else
f(s)
f(s) =
g(s) = f(s) (1-s))
B =
u
4R
Before proving the properties of NNs from the class T, we require a technical result on path differentiable functions.
Proposition 7. Let f Rdd R path differentiable, and g Rd R of class C². Then their product fg is path differentiable.
18

---

Published in Transactions on Machine Learning Research (11/2023)
Proof. Our objective is to apply (Bolte & Pauwels 2021] Corollary 2 (stated in Proposition 6 which is to say that h == Ig admits a chain rule for 0ch). First, we apply the definition of the Clarke differential and compute
VUER" Ocf(u)= + - + g(uje I dcs(u))
Note that used the smoothness of g. We now consider absolutely continuous curve Caha (R+, (R+,IE)) IR). By Bolte & Pauwels (2021) Lemma 2, since f is path differentiable, f fou is differentiable almost everywhere. Let D the associated set of difforentiability, then let E D and u Och(((s)), writing v V(((())) (((()))w with w E (u(s). We compute (h o ()(()) = (f uu) (s)(s)((((s)) + ((
Now since f is path differentiable and w E @cf(u(s)), by Proposition 6 item 3, we have (f o u)(()) = (w, i(s)) On the other hand, (g o ()))) = (Vg(u(s)), ú(s) since C. Finally by definition of v and bilinearity of (:-),
()) (ue, + [(())) (v, ú(s))
We now have all the tools to prove that the class of NNs T satisfics all of the assumptions of our paper. Proposition 8. All networks of the class T verify Assumption Assumption and Assumplion 7
2
Assumption 4
Assumption
Proof. Let T E T,, and ī its associated underlying network. We begin with regularity considerations.
Verifying, Assumptions and in the C² Case In the case where the activations are C²-smooth, then each ((, is also of class C². Furthermore, the smooth indicator my, is -moohh, thus we can conclude that x() is C2-smooth, and thus satisfics Assumption trivially. In particular, T(,,x) is path differentiable for any x E Rd=, thus T also satisfies Assumplion 7
Verifying Assumptions 1 and 7 in the Semi- Aegi-Algebraic Case In the where the activations are locally Lipschitz and semi-algebraic, it follows that each T(,x)) x) is semi-algebraic, which yields naturally a differentiability structure associated to the polynomial pieces, satisfying [Assumption I Furthermore, this regularity yields path differentiability by (Bolte & Pauwels (021)) Proposition 2. By product with the smooth indicator function, is path differentiable by Proposition therefore it satisfies Assumption
Verifying Assumption in the C2 Case In the case where the activations are C²-smooth, is clear that by composition and product (u, x) - T(u, x) is jointly C2-smooth. As consequence, it is Lipschitz jointly in (u, on any compact of Rdd x R&d, and by product with the smooth indicators, so is T. Since T is zero outside B(, R. + e) x B(0, Rs + +e), we conclude that it is globally Lipschitz in (u,
Verifying Assumption in the Semi- Algeb-igee Case the case of locally Lipschitz and semi- algebraic activations, we prove that T' is jointly Lipschitz on any compact K. by strong induction on [[,,]] Let K Ki Ka product compact of Rdd Rd,, and P., '3L,, > hn is LipSiischitz on K". The initialisation Po is trivial, since 2(u,z) = Let E [V] and assume P. to hold true for i [0, 1]. In particular, the hy are jointly continuous in (u, x), allowing the definition of
= max
Σ
(n, + B, 1=0
(u.
19

---

Published in Transactions on Machine Learning Research (11/2023)
Since an is locally Lipschitz, covering argument shows that there exists La, 0 such that an is Lan"
Lipschitz on [- M-. Now let u2) KJ and (2), E K2. We have
|| 21) (w2, x2)]12 Lan
Σ Ana(u (1x + Bu
Σ
An,((2))i((33 - Bn u2
Lan
[Bullop ll U2ll2
Σ An, (4n -n An 22)112
(19)
where Il-llop denotes the - I|2-ind operator norm. Let E [0, n - 11, we separate the norm:
(An,(ui) 01) (22 22)1l2 S (An,
An, An, æ1)ll2 =: A1
An, (u2 (u .2)112 =: D2.
(20)
For Ai, use the tensor form (18) and the inequality ||I||| S vdl|=l|> for RRd, then ||l|o S |uul2
d Σ
A S dn
(n.d) iiJads
x1) - -u )
1==1
jn jiE[1,d]]
S v dn max
(n
max (v.))RiiX2,
[1)((u,))loo | - u21 1011-
Ji.22ds
$ La |lu - U2112.
(21)
For 42, leverage P; and obtain
A S max [|A((u))lop [1,((u,æ1) - h4(u2, x2)lla S max [[A(W))lop L, (llur - u2llz + || - x2ll2).
(22)
uEK]
weKi
Combining 19)
20 21) and (22) shows Ps and concludes the induction, which turn shows that f is
jointly Lipschitz any compact. Like in the smooth case, we conclude that T is globally Lipschitz, and thus that Assumption holds.
Verifying Assumption 4 Under both cases of regularity for the activations,
g==
max 2ll,
is measurable and bounded, Furthermore, observe that for u, x Rdd x Rd, [T(,,)) S g(x). As a consequence, Assumption holds.
Verifying Assumption in the C2 case If all activations -m-smohth, both T and its coordinate- wise products T, T3 are C"-smooth jointly in xw) As result, one may bound these terms on (u, E B(0, R2 + B(0, R₂ + E) by constant M, independent of u,z, and the assumption is verified.
Verifying Assumption in the semi-algebraic case In the semi-algebraic case, there exists a structure (Uf)ees of open sets Rd x Rd whose closures cover the entire space, such that T is polynomial in (u,z) on each Uj, with J finite (this is possible since / is jointly semi-algebraic). The NN can be written T(u,x) = ((,, " (u) (xx), and is therefore C²-smooth on each Uy. Furthermore, its restriction to U, is extendable C2-smoothly to Uy (we shall not introduce different notation to these extensions, for legibility). As result, one may introduce the following bounds on the derivatives of the coordinate functions on the intersection of the compact K. = B(0, R. + x B(O, R. + s) and Uy: there exists an Mj > 0 such that
82 V(u,z) K n Uy Ou,, Ous,
8²T Ou, Our
< M, and
(u,
$ M,
20

---

Published in Transactions on Machine Learning Research (11/2023)
Since J is finite and the (U))Ee cover K, we deduce that this bound holds for (u,x) E K. for common constant M > 0. Moreover, since T is the zero function outside of K, this bounds also holds for any (u, x) E Rd x Rdd. Finally, this shows that Assumption 5 holds.
E Generalisation to Other Sliced Wasserstein Orders
In this section, we shall discuss how some of our results can be extended by replacing the 2-SW term SW2 with Swp for [oo+
Determining Lipschitz Constants The first difficulty lies in showing that the functions w == (X - Wp Po still have a locally Lipschitz regularity similar to Proposition 1 (this propo- sition only shown for p = 2 in (Tanguy et al. 2023 ). We generalise their result the lollowing proposition.
Proposition 9. Let K!P (r, Y) == 2pn(r + 1| + 20 for X, E Raxd, and r > 0. Then wa (p) (, K (r, X, ) in the neighbourhood 2
VY', E 2X, vo E E (wa(Y",Y) ww Y) s (r, XY)Y'" -Y
Proof. Let X,Y Y E Raxd, ,r> 0 and Y Y" E B 1-1 (X,r). ( By Tangny et al. 2023 Lemma 2.2.1, we have (P) (Y')= - wg (p) (Y"") s 2 C"Ile, where || || denotes the Frobenius norm, and C' a n n matrix of entrics C¹21- _ 1/2 H "ll,, with similarly Ch,1 = l v/ 0 yi.. Now consider the function
9m -
{
Rdd
R l y A y1
which satisfies Ch,1 = 9y (vk), and is differentiable almost-everywhere, with V (2) = p/0 -Ty11--1 For almost every E B(Ix, r), we have
VV (2)lla S pllv-uull v = p - S <p p(lv æxl12 + || + 11 12) S p(r+||XX|lee |loc. (1000))))----
As result, Ivr is p(r || + ( in B(xx,r). Now since Y , E 111llboaaX, we have sk 56 E B(I,.T), thus
|[C'|,, - [C"]a,l (9h) 9r, (wh)) S p(r || ||llll.,++++|||00./0011| ||11 1121 Vl
Then ||C" (C'I=" 110"1 < np(r |X 1o,, + ||Y 20 111' Y"1,,,,
Our results regarding the local Lipschitz property of f and F adapt immodiately using the same method with the different constant k(()) (r, X, Y), we obtain the following constant for (with L from Assumption 2
k/P (e, uo, X, 2pnL (e (1o X) + 100,))))
then the following constant for F:
k(P) (e, uo) = 2pnL f Jxnxyn (EL +
(EL+IT((Uo- [(uo, X) + ||oll,0,)) "- dx (X)d(X))))*()),
In order satisfy Condition item i) the case p * 2, one needs to modify Assumption 4 require [[T(u, 2)2 S //=x)//>=>1))))0 + [[=]))) which in realistic cases is not much more expensive than asking for T to be bounded, which is property of the class of NNs that we present in Section [)]
21

---

Published in Transactions on Machine Learning Research (11/2023)
Almost-Ei Mnnost-Everywhere Gradient A second difficulty lies in defining aleest-eeerywhere gradient since in our main text we rely on the formulation of an elmost-eeerywhere gradient of w/ (., which was derived only for p 2 by Bonneel et al. (2015) and Tanguy et al. (2023). In fact, for 0, Y fixed wg? (X,Y) is piecewise smooth, like 2 Y) is piecewise quadratic. As a result, one may show that the following is an aleosteeverywhere gradient of wg (,:Y):
@rlg) ax (X,Y)
(~ssign (8 (T - 0 yxxxx [TTxx Is AT yo (8)
pp-
E Ruxdy
kE[1,n]
The chain rule now yields the following wleosteeeeryyhere gradient for
(ou, X,,,) = Σ P sign (TTT(u,z)) T y (o..,.. 8 T(, - 0 y.
p-1 OT du (u. xx
Adapting Proposition 4 Moving on to adapting Proposition 4 the general case * 2 makes things substantially more technical, but one may still show that the up functions are Lipschitz using restrictions on T its first and second-order derivatives (which can be formulated in more technical version of Assumption 5 In conclusion, Proposition 4 can be adapted to apply to p E too|, and it follows that Theorem also generalises to this case.
Path Differentiability Regarding the results from Section 4 the only substantial difference lies in showing that T(, is path differentiable. The only missing link in the composition chain is the path differentiability of g(r) X + fge-1 w/ (X, Y)do In the case p 2, the difficuuly of the integral
be circumvented by noticing that & is semi-concave (Tanguy et al. 2023 Proposition 2.4.3, which implies path differentia- bility. This argument does not generalise to E 1, 0o naturally, hence our Theorem 2 only generalises to E [1. +oo under the conjecture that g(p) is indeed path differentiable.
22


Reviews:
[{"review_id": "Uo0v6rYDbp", "reviewer": "Reviewer_Lzga", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n(1) The paper is well-motivated (as training generative neural networks on SW loss is relatively popular), and the authors partially address the need for a theoretical converge guarantee of such a method.\n\n(2) Some limitations of current work are properly discussed, such as Assumption 6.\n\n**Weaknesses**:\n\n(1) In contrast to Strength (1), on which the author claims “it has been observed that SGD in such settings always converges”, there should be some explicit citations supporting this claim.\n\n(2) There is the lack of novelty in the techniques: it appears as a straight-forward aggregation of previous works. Two main theorems are direct applications of results in [Bianchi et al, 2022], and thus most of the current paper is spent on deriving conditions aligned with those in that prior work. Besides, properties of SW loss are also employed directly from [Tanguy et al, 2023].\n\n(3) The authors should make clear how the structure of SW loss plays a role in the convergence. A few places where this kind of statement has appeared to my observation are at the regularity of the one-dimensional Wasserstein distance for Condition 1 (which, I believe, is not particularly distinctive), and at Proposition 4 and Proposition 5 (but it is unclear if these properties of SW loss all hold for, for example, other Wasserstein-type losses).\n\n(4) Though the authors do present some examples how each assumption can be satisfied, there is an obvious question left unaddressed: are assumptions satisfied in papers training with SW loss?\n\n(5) The paper is quite difficult to read. While the introduction is fine, latter sections consists of mostly contrived formulas/notations and consecutive assumptions/propositions/theorems with little insight. I suggest the authors add a section for notations, or a table representing them. Assumptions/Conditions may be put together for a better flow. Proof of Proposition 4 should be put into the appendix (it is not clear why it appears/what function it serves in the main text - unlike the proof of Proposition 5 which is short and highlights the technicality of the discreteness assumption).\n\n(6) Could we get the similar mode of convergence in Theorem 2 for the case of SGD in Section 3, using Theorem 3 of [Bianchi et al, 2022]?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper provides a theoretical convergence guarantee of the trajectory of stochastic gradient descent (SGD) on training neural networks with sliced Wasserstein (SW) loss. In particular, the authors specify a set of assumptions that imply some conditions (C2-smoothness, Lipschitzness, compactness, boundedness, ...) on (the function defined by) the neural network and the loss function so that, based on a recent result of [Bianchi et al, 2022], the corresponding linearly interpolated process converges in probability as the step size goes to zero. They also investigate the case of projected SGD (onto some ball), and specify some additional assumptions (discreteness of measures, path differentiability of neural net) that lead to the convergence of the trajectory (to a set of Karush-Kuhn-Tucker points of the differential inclusion).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Minor Typos\n\n“With this definition for f ...“: should add F(u) = E[f(u, ...)]\n\n“To this end, notice that one may write...“: change “w_theta(Y, Z)” to “w_theta(Z, Y)”\n\n“To define an almost-everywhere gradient...“: change “=” to “:”\n\n“In our non-smooth setting, ...“: “is is” to “is”", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oQb5cbQrTl", "reviewer": "Reviewer_FjBK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n1. The considered problem would be of interest at least to a subset of the TMLR community. The setting is well-motivated by generative modeling applications and the paper provides theory to close an existing gap in the literature.\n2. The level of technical writing is decent, although I would like to see some improvements (see below). The derivations are details, explain well, and seem sounds.\n3. All the assumptions are clearly stated, and the authors are honest about the limitations of their work. I appreciate this style of presentation.\n4. The paper is structured well and follows a clear flow. Every time a condition needed to invoke the results from [Bianchi et al. 2022] is introduced, it is followed by the postulated assumptions needed to satisfy it and a formal statement of the relevant result.\n\n**Weaknesses:** \n1. The paper somewhat lacks in novelty in terms of the technical contributions. It strongly relies on the theory developed in [Bianchi et al. 2022] and mostly goes about verifying the conditions therein by imposing assumptions on the model and distributions. That said, I do find the derived results to be interesting. \n2. The paper significantly lacks in its literature review and misses many citations of related topics/works. Please seem Item 1 in `Requested Changes' below.\n3. Some of the assumptions are quite strong and should be rigorously verified for the appropriate class of neural networks. See Item 2 below.\n4. Continuing on the above point: while the notion of convergence in Section 4 is indeed stronger, the discreteness assumption is too restrictive and renders the utility of these results questionable.\n5. There are various natural extensions of the studied problem that should be addressed, as discussed in Item 3 below.\n\n3. The are some stylistic and notational inconsistencies in the text that should be address. Weaknesses of the work are listed in the `Requested Changes' section below. There, I first describe my concern/reservation and then request the appropriate correction.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies convergence of SGD schemes for training neural networks with sliced Wasserstein distance loss. Specifically, the authors consider the problem of minimizing the average-sliced 2-Wasserstein distance between a target probability distribution and a model given by the pushforward of a through reference measure (e.g., Gaussian or uniform) through a neural network $T_u$, with parameters $u$. To establish convergence of SGD to critical points of the objective, the authors adopt the framework of [Bianchi et al. 2022] and verify the conditions required to invoke the theorems therein. To that end, the authors impose various technical conditions on the distributions and the neural network and provide proofs for the regularity conditions needed for the convergence results from [Bianchi et al. 2022].", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The authors focus on SGD optimization of the average-sliced Wasserstein distance, but several related problems are overlooked. Specifically, are similar results available for the standard Wasserstein distance or for its entropically regularized versions. While computation of the standard distance is indeed hard, the differentiability in $u$ can be established and a similar analysis of convergence should apply. Have these results been reported before or are they serious roadblocks towards obtaining them? Perhaps more 4importantly, the entropic optimal transport distance (i.e., where the transport cost is regularized by the negative entropy of the coupling $\\pi$) can be efficiently computed using the Sinkhorn algorithm. As such, it enjoys similar virtues as the sliced version and, again, differentiability of the corresponding objective can be established. Is there literature on this? I would like to see the authors add a `Related Literature' subsection to their introduction, where they would discuss any such related work.\n\nAlso, the authors cite [Nadjahi et al. (2020)] in the context of a statistical theory for the sliced Wasserstein distance, but there are many more relevant works that are overlooked. A partial list includes:\n\n\n* Lin, T., Zheng, Z., Chen, E., Cuturi, M., & Jordan, M. I. (2021, March). On projection robust optimal transport: Sample complexity and model misspecification. In the International Conference on Artificial Intelligence and Statistics (pp. 262-270). PMLR.\n\n* Niles-Weed, J., & Rigollet, P. (2022). Estimation of Wasserstein distances in the spiked transport model. Bernoulli, 28(4), 2663-2688.\n\n* Goldfeld, Z., Kato, K., Rioux, G., & Sadhu, R. (2022). Statistical inference with regularized optimal transport. arXiv preprint arXiv:2205.04283.\n\n* Xi, J., & Niles-Weed, J. (2022). Distributional convergence of the sliced Wasserstein process. Advances in Neural Information Processing Systems, 35, 13961-13973.\n\n* Xu, X., & Huang, Z. (2022). Central limit theorem for the sliced 1-Wasserstein distance and the max-sliced 1-Wasserstein distance. arXiv preprint arXiv:2205.14624.\n\n* Nietert, S., Goldfeld, Z., Sadhu, R., & Kato, K. (2022). Statistical, robustness, and computational guarantees for sliced Wasserstein distances. Advances in Neural Information Processing Systems, 35, 28179-28193.\n\nIn addition, I believe that the notion of sliced Wasserstein distances dates back to \n\n* Rabin, J., Peyré, G., Delon, J., & Bernot, M. (2012). Wasserstein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29–June 2, 2011, Revised Selected Papers 3 (pp. 435-446). Springer Berlin Heidelberg.\n\nOverall, I find the literature review insufficient and would like to see the authors expand it to appropriately cover related work. \n\n2. A key piece that I found missing in this paper is a rigorous verification of the imposed assumptions for a concrete neural network class. Throughout the text, the authors verbally comment on their assumptions being valid for neural nets with bounded parameters/outputs, but I find this to be insufficient. The utility of this theory hinges on its ability to cover neural networks that would be used in practice for training such sliced-Wasserstein losses. I would like the authors to add a formal statement(s), where they specify the considered neural network class (depth/width/nonlinearities/parameter bounds/etc.) and then state and prove that it satisfies all the assumptions needed for Theorem\n\n3. There are various extensions of the work that the authors may consider:\n* Other orders $p\\in[1,\\infty)$ of the sliced Wasserstein distance?\n* Max-sliced Wasserstein distance $\\mathrm{MSW}\\_p(\\mu,\\nu):=\\sup_{\\theta\\in\\mathbb{S}^{d-1}}\\mathrm{W}_p(P\\_{\\theta\\ \\sharp}\\mu,P\\_{\\theta\\ \\sharp}\\nu)$. In this case, the projection can be absorbed into the neural net $T_u$ and they can be jointly optimized.\n* Higher dimensional projection directions a la the so-called projection robust Wasserstein distance, where the unit sphere is replaced with the Stiefel manifold $V\\_k(\\mathbb{R}^d)$ of $d \\times k$ matrices with orthonormal columns. \n* In the more restricted setting of Section 4 (noised projected SGD), the finite support assumption on the distributions is quite restrictive. Is it possible to at least extend the results to countably infinite supports? The extension to continuous distributions seems beyond the authors' reach at the moment, but countably infinite supports could be treatable.\n\nI understand that space may be too limited to address all these items, but it would be appropriate to at least acknowledge them and comment on possible extensions. If certain results extend directly, the authors may add a remark to say so. If other extensions are more challenging, that can be discussed in the `Summary and Outlook' section.\n\n4. Please correct notational inconsistencies that are present throughout the paper. Distributions are sometimes denoted by Greek letters (e.g., $\\nu_1$, \\gamma_X$, etc.) and at other times by fancy lowercase letters. The distinction between when each type of notation is used is not clear to me. I suggest picking a single convention and sticking with it throughout the text. Titles of propositions are also inconsistent, e.g., Proposition 2 has a title but various other ones in the text do not. Again, please adopt a consistent style.  \n\n5. Question: the authors consider neural nets of the form $T_u(x)\\mathbb{1}_B(u)$ and argue that this is not a limitation in practice since the radius of $B$ can be taken arbitrarily large (end of 2nd paragraph on page 3). Is it a concern that $B$ is generally data-dependent?", "broader_impact_concerns": "This is a theory paper, and I don't see any ethical concerns stemming from this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rVJ6vdcNVo", "reviewer": "Reviewer_MSYk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- The authors are clear about their main claims and those claims are well-substantiated. The results they prove are clearly presented and provide satisfying theoretical justification for the claim that SGD can reasonably be viewed as an approximation to (sub-)gradient flow in their setting.\n- The assumptions required for the argument of Section 3 are pretty mild, making the results generally applicable.\n- The authors provide significant intuition throughout as to the correct interpretation of their assumptions and results. This was much appreciated.\n- The authors are clear how their methods relate to the previous work of Bianchi et al. This clarity made it much easier to follow the structure of the paper.\n\n### Weaknesses\n\n- My impression is that the core of the paper essentially boils down to checking that the conditions of a result already proved by Bianchi et al. (2022) are satisfied in the particular context of the SGD procedure induced by considering an SW loss. As a result, much of the work is a straightforward, if somewhat technically intricate, check of those conditions. The results could reasonbly be viewed as just a particular instantiation of those previous results, and it's unclear how much this work contributes to the literature on top of those previous works.\n- The discretization requirement for the result of Section 4 seems like a significant restriction on it's applicability. It seems unlikely that distributions we would want to approximate in practice statisfy that assumption.\n- In Sections 4 and 5, the authors use many concepts that they imply are established in the previous literature, but which they don't explain themselves: “conservative field for $F$”, “$F$ has a chain rule for $\\partial_C$, “conservative mappings”, “Karush-Kahn-Tucker points”, “Clark regular functions”,  etc. At several points it is unclear whether the authors intend particular definitions to be stand-alone definition or if they intend to merely recall a concept from the literature and redirect the reader for further details. While I appreciate that the authors have given references for each of these concepts, it might be convenient for readers who are less familiar with this material to include some background on the relevant concepts (e.g. in an Appendix).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides theoretical results showing that the paths followed by a standard SGD procedure for training generative neural networks with a sliced Wasserstein loss approximates the set of solutions to a sub-gradient flow equation as the step size becomes small. The main idea of the paper is to specialise general results of this style proved by Bianchi et al. (2022) to this case. This requires checking that all the conditions of Bianchi et al.'s theorems hold in the particular case the authors are studying, which the authors do carefully. The outcome is a pair of results, the first showing that SGD approximates sub-gradient flow under a relatively weak set of assumptions, and the second showing under much stronger assumptions that the set of limiting points of the SGD procedure approximates the set of generalized critical points of the loss function.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Most of the changes I have to recommend are minor:\n\n- On page 2, in the third paragraph, there's a conflation between the notation $\\mathbb{y}$ and $\\mu$.\n- On page 4, in Algorithm 1, it is not clear how the noise level $a$ and the convergence threshold $\\beta$ feature.\n- On page 5, it might be nice to have a comment after the final line of Proposition 2 to indicate that you have used Assumption 2 and Proposition 1. (This took me some time to work out.)\n- On pages 5-7, I found distinguishing between $\\kappa_\\varepsilon$ and $\\kappa_r$ quite confusing; they're performing similar but different roles in your proofs, and currently the reader has to be quite careful to keep track from context which one you mean. In addition, on page 7, $\\kappa_1$ now seems to depend on $\\theta$?\n- One page 6, there appears to be a formatting error in the first line.\n- On page 6, in the third line you write $f(u)$. Perhaps I have misunderstood, but should this be $f(u', X, Y, \\theta)$? (Otherwise, I'm not sure how $f(u)$ is defined.)\n- On page 6, you say \"$\\mathbb{x}$ is a Radon probability measure by Assumption 3, thus C1 is finite\"; I'm not sure why this should be true?\n- On page 6, \"In order to prove that the interpolated trajectories...\"; it seems like this clause is incomplete.\n- On page 8, the definition of the notation $\\chi$ is unclear to read.\n- On page 8, in the fourth-from-bottom line and final line, I don't see where the terms $\\theta_{i_3}, \\theta_{i_4}$ have gone.\n- On page 9, I don't follow the final couple of lines of the proof (from \"Rewriting...\"). I think this could do with being explained in more detail or rewritten.\n", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: As explained above, reviewers found the paper interesting and all the concerns have been addressed successfully by the authors.
The authors have already submitted several revisions with minor changes, in order to handle the suggestions by the reviewers. 

Overall, the problem under consideration is likely of interest to some of the TMLR community. It is well-motivated, particularly in the context of generative modeling applications, and the paper addresses a gap in the existing literature with theoretical contributions.
The technical writing is generally good, but some improvements are suggested and nicely handled. The derivations are detailed, well-explained, and appear to be sound.

The authors are transparent about their assumptions and acknowledge the limitations of their work, which is appreciated.
The paper is well-structured with a clear flow. Each time a condition required to apply results from [Bianchi et al. 2022] is introduced, it is accompanied by the necessary assumptions and a formal statement of the relevant result.
Decision: Accept (poster)