Convergence of SGD for Training Neural Networks with Sliced Wasserstein
  Losses
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.
Eloi Tanguy
# Convergence of SGD for Training Neural Networks with Sliced Wasserstein Losses

###### 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.

## Table of Contents

* 1 Introduction
	* 1.1 The Sliced Wasserstein Distance in Machine Learning
	* 1.2 Contributions
* 2 Stochastic Gradient Descent with \(\mathrm{SW}\) as Loss
* 3 Convergence of Interpolated SGD Trajectories on \(F\)
* 4 Convergence of Noised Projected SGD Schemes on \(F\)
* 5 Conclusion and Outlook

## 1 Introduction

### The Sliced Wasserstein Distance in Machine Learning

Optimal Transport (OT) allows the comparison of measures on a metric space by generalising the use of the ground metric. Typical applications use the so-called 2-Wasserstein distance, defined as

\[\forall\nu_{1},\nu_{2}\in\mathcal{P}_{2}(\mathbb{R}^{d}),\;\mathrm{W}_{2}^{2} (\nu_{1},\nu_{2}):=\inf_{\pi\in\Pi(\nu_{1},\nu_{2})}\int_{\mathbb{R}^{d} \times\mathbb{R}^{d}}\|x_{1}-x_{2}\|^{2}\mathrm{d}\pi(x_{1},x_{2}),\] (W2)

where \(\mathcal{P}_{2}(\mathbb{R}^{d})\) is the set of probability measures on \(\mathbb{R}^{d}\) admitting a second-order moment and where \(\Pi(\nu_{1},\nu_{2})\) is the set of measures of \(\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d})\) of first marginal \(\nu_{1}\) and second marginal \(\nu_{2}\). One may find a thorough presentation of its properties in classical monographs such as Peyre and Cuturi (2019); Santambrogio (2015); Villani (2009)
The ability to compare probability measures is useful in probability density fitting problems, which are a sub-genre of generation tasks. In this formalism, one considers a probability measure \(\mu_{u}\), parametrised by \(u\) which is designed to approach a target data distribution \(\mu\) (typically the real-world dataset). In order to determine suitable parameters, one may choose any probability discrepancy (Kullback-Leibler, Ciszar divergences, f-divergences or Maximum Mean Discrepancy), or in our case, the Wasserstein distance. In the case of Generative Adversarial Networks, the optimisation problem which trains the "Wasserstein GAN" (Arjovsky et al., 2017) stems from the Kantorovitch-Rubinstein dual expression of the 1-Wasserstein distance.

A less cost-intensive alternative to \(\mathrm{W}_{2}^{2}\) is the Sliced Wasserstein (SW) Distance introduced by Bonneel et al. (2015), which consists in computing the 1D Wasserstein distances between projections of input measures, and averaging over the projections. The aforementioned projection of a measure \(\nu\) on \(\mathbb{R}^{d}\) is done by the _push-forward_ operation by the map \(P_{\theta}:x\longmapsto\theta\cdot x\). Formally, \(P_{\theta}\#\nu\) is the measure on \(\mathbb{R}\) such that for any Borel set \(B\subset\mathbb{R}\), \(P_{\theta}\#\nu(B)=\nu(P_{\theta}^{-1}(B))\). Once the measures are projected onto a line \(\mathbb{R}\theta\), the computation of the Wasserstein distance becomes substantially simpler numerically. We shall illustrate this fact in the discrete case, which arises in practical optimisation settings. Let two discrete measures on \(\mathbb{R}^{d}\): \(\gamma_{X}:=\frac{1}{n}\sum_{k}\delta_{x_{k}},\ \gamma_{Y}:=\frac{1}{n}\sum_{k} \delta_{y_{k}}\) with \(x_{1},\cdots,x_{n},y_{1},\cdots,y_{n}\in\mathbb{R}^{d}\). Their push-borwards by \(P_{\theta}\) are simply computed by the formula \(P_{\theta}\#\gamma_{X}=\sum_{k}\delta_{P_{\theta}(x_{k})}\), and the 2-Wasserstein distance between their projections can be computed by sorting their supports: let \(\sigma\) a permutation sorting \((\theta^{T}x_{1},\cdots,\theta^{T}x_{n})\), and \(\tau\) a permutation sorting \((\theta^{T}y_{1},\cdots,\theta^{T}y_{n})\), one has the simple expression

\[\mathrm{W}_{2}^{2}(P_{\theta}\#\gamma_{X},P_{\theta}\#\gamma_{Y})=\frac{1}{n} \sum_{k=1}^{n}(\theta^{T}x_{\sigma(k)}-\theta^{T}y_{\tau(k)})^{2}. \tag{1}\]

The SW distance is the expectation of this quantity with respect to \(\theta\sim\mathfrak{v}\), i.e. uniform on the sphere: \(\mathrm{SW}_{2}^{2}(\gamma_{X},\gamma_{Y})=\mathbb{E}_{\theta\sim\mathfrak{v }}\left[\mathrm{W}_{2}^{2}(P_{\theta}\#\gamma_{X},P_{\theta}\#\gamma_{Y})\right]\). The 2-SW distance is also defined more generally between two measures \(\mu,\nu\in\mathcal{P}_{2}(\mathbb{R}^{d})\):

\[\mathrm{SW}_{2}^{2}(\mu,\nu):=\int_{\theta\in\mathbb{S}^{d-1}}\mathrm{W}_{2}^ {2}(P_{\theta}\#\mu,P_{\theta}\#\nu)\mathrm{d}\sigma(\theta).\] (SW)

The implicit generative modelling framework is a formalisation of the training step of generative Neural Networks (NNs), where a network \(T\) of parameters \(u\) is learned such as to minimise the discrepancy between \(T_{u}\#x\)1 and \(\mathsf{y}\), where \(\mathsf{x}\) is a low-dimensional input distribution (often chosen as Gaussian or uniform noise), and where \(\mu\) is the target distribution. Our case of interest is when the discrepancy is measured with the SW distance, which leads to minimising \(\mathrm{SW}_{2}^{2}(T_{u}\#\mathsf{x},\mathsf{y})\) in \(u\). In order to train a NN in this manner, at each iteration one draws \(n\) samples from \(\mathsf{x}\) and \(\mathsf{y}\) (denoted \(\gamma_{X}\) and \(\gamma_{Y}\) as discrete measures with \(n\) points), as well as a projection \(\theta\) (or a batch of \(p\) projections) and performs an SGD step on the sample loss

Footnote 1: \(T_{u}\#\mathsf{x}\) is the push-forward measure of \(\mathsf{x}\) by \(T_{u}\), i.e. the law of \(T_{u}(x)\) when \(x\sim\mathsf{x}\).

\[\mathcal{L}(u)=\mathrm{SW}_{2}^{2}(P_{\theta}\#T_{u}\#\gamma_{X},P_{\theta} \#\gamma_{Y})=\frac{1}{n}\sum_{k=1}^{n}(\theta^{T}T_{u}(x_{\sigma(k)})-\theta^ {T}y_{\tau(k)})^{2}, \tag{2}\]

with respect to the parameters \(u\) (see Algorithm 1 for a precise formalisation). In order to compute this numerically, the main complexity comes from determining the permutations \(\sigma\) and \(\tau\) by sorting the numbers \((\theta^{T}T_{u}(x_{k}))_{k}\) and \((y_{k})_{k}\), and summing the results, while the Wasserstein alternative \(\mathrm{W}_{2}^{2}(T_{u}\#\gamma_{X},\gamma_{Y})\) is done by solving a Linear Program, which is substantially costlier.

In this paper, we shall study this training method theoretically and prove convergence results. Theoretical guarantees for this optimisation problem are welcome, since this question has not yet been tackled (to our knowledge), even though its use is relatively widespread: for instance, Deshpande et al. (2018) and Wu et al. (2019) train GANs and auto-encoders with this method. 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 neural features, and by Tartavel et al. (2016) with wavelet features (amongst other methods).

In practice, it has been observed that SGD in such settings always converges (in the loose numerical sense), yet this property is not known theoretically, since the loss function defined in (2) is not differentiable nor convex in general, because \(X\longmapsto\mathrm{SW}^{2}_{2}(\gamma_{X},\gamma_{Y})\) and the neural network do not have such regularities. Several efforts have been made to prove the convergence of SGD trajectories within this theoretically difficult setting: Bianchi et al. (2022) show the convergence of fixed-step SGD schemes on a function \(F\) with some technical regularity assumptions, Majewski et al. (2018) show the convergence of diminishing-step SGD schemes assuming stronger regularity results on \(F\). Another notable theoretical work is by Bolte and Pauwels (2021), which leverages conservative field theory to prove convergence for back-propagated SGD on deep NNs with definable activations and loss functions. In the case of Optimal Transport losses, the only work (that we are aware of) that has tackled this problem is by Fatras et al. (2021), proving strong convergence results for minibatch variants of classical OT distances, namely the Wasserstein, Entropic Wasserstein and Gromov Wasserstein distances. The aim of this work is to bridge the gap between theory and practical observation by proving convergence results for SGD on Sliced Wasserstein generative losses of the form \(F(u)=\mathrm{SW}^{2}_{2}(T_{u}\#\mathsf{x},\mathsf{y})\).

### Contributions

Convergence of Interpolated SGD Under Practical AssumptionsUnder practically realistic assumptions, we prove in Theorem1 that piece-wise affine interpolations (defined in Equation (6)) of constant-step SGD schemes on \(u\longmapsto F(u)\) (formalised in Equation (4)) converge towards the set of sub-gradient flow solutions (see Equation (5)) as the gradient step decreases. This results signifies that with very small learning rates, SGD trajectories will be close to sub-gradient flows, which themselves converge to critical points of \(F\) (omitting serious technicalities).

The assumptions needed for this result are practically reasonable: the input measure \(\mathsf{x}\) and the true data measure \(\mathsf{y}\) are assumed to be compactly supported. As for the network \((u,x)\longmapsto T(u,x)\), we assume that for a fixed datum \(x\), \(T(\cdot,x)\) is piecewise \(\mathcal{C}^{2}\)-smooth and that it is Lipschitz jointly in both variables on any compact. We require additional assumptions on \(T\) which are more costly, but are verified as long as \(T\) is of the form \(T(u,x)=\widetilde{T}(u,x)\mathbb{1}_{B}(u)\), where \(\widetilde{T}\) is any typical NN composed of compositions of definable activations (as is the case for all typical activations, see (Bolte and Pauwels, 2021), SS6.2), and of linear units; and where \(\mathbb{1}_{B}(u)\) is the indicator that the parameter \(u\) be in a fixed ball \(B\). This form for \(T\) is a strong theoretical assumption, but in practice makes little difference, as one may take a fixed ball \(B\) to be arbitrarily large.

Stronger Convergence Under Stricter AssumptionsIn order to obtain a stronger convergence result, we consider a variant of SGD where each iteration receives an additive noise (scaled by the learning rate) which allows for better space exploration, and where each iteration is projected on a ball \(B(0,r)\) in order to ensure boundedness. This alternative SGD scheme remains within the realm of practical applications, and we show in Theorem2 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).

Unfortunately, we require additional assumptions in order to obtain this stronger convergence result, the most important of which is that the input data measure \(\mathsf{x}\) and the dataset measure \(\mathsf{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 envision 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 \(\mathsf{x}\) can have an arbitrary fixed amount of points, and leverage strong sample complexity results such as those of Nadjahi et al. (2020) to ascertain that the discretisation is not costly if the number of samples is large enough.
## 2 Stochastic Gradient Descent with \(\mathrm{SW}\) as Loss

Training Sliced-Wasserstein generative models consists in training a neural network

\[T:\left\{\begin{array}{rcl}\mathbb{R}^{d_{u}}\times\mathbb{R}^{d_{x}}& \longrightarrow&\mathbb{R}^{d_{y}}\\ (u,x)&\longmapsto&T_{u}(x):=T(u,x)\end{array}\right.\]

by minimising \(u\longmapsto\mathrm{SW}^{2}_{2}(T_{u}\#\mathsf{x},\mathsf{y})\) through Stochastic Gradient Descent (as described in Algorithm 1). The probability distribution \(\mathsf{x}\in\mathcal{P}_{2}(\mathbb{R}^{d_{x}})\) is the law of the input of the generator \(T(u,\cdot)\). The distribution \(\mathsf{y}\in\mathcal{P}_{2}(\mathbb{R}^{d_{y}})\) is the data distribution, which \(T\) aims to simulate. Finally, \(\mathsf{o}\) will denote the uniform measure on the unit sphere of \(\mathbb{R}^{d_{y}}\), denoted by \(\mathbb{S}^{d_{y}-1}\). Given a list of points \(X=(x_{1},\cdots,x_{n})\in\mathbb{R}^{n\times d_{x}}\), denote the associated discrete uniform measure \(\gamma_{X}:=\frac{1}{n}\sum_{i}\delta_{x_{i}}\). By abuse of notation, we write \(T_{u}(X):=(T_{u}(x_{1}),\cdots,T_{u}(x_{n}))\in\mathbb{R}^{n\times d_{y}}\).

```
Data: Learning rate \(\alpha>0\), noise level \(a\geq 0\), convergence threshold \(\beta>0\), probability distributions \(\mathsf{x}\in\mathcal{P}_{2}(\mathbb{R}^{d_{x}})\) and \(\mathsf{y}\in\mathcal{P}_{2}(\mathbb{R}^{d_{y}})\). Initialisation: Draw \(u^{(0)}\in\mathbb{R}^{d_{u}}\); for\(t\in[\![0,T_{\max}-1]\!]\)do  Draw \(\theta^{(t+1)}\sim\mathsf{o},\ X^{(t+1)}\sim\mathsf{x}^{\otimes n}\ Y^{(t+1)} \sim\mathsf{y}^{\otimes n}\). SGD update: \(u^{(t+1)}=u^{(t)}-\alpha\left[\frac{\partial}{\partial u} \mathrm{W}^{2}_{2}(P_{\theta^{(t+1)}}\#T_{u}\#\gamma_{X^{(t+1)}},P_{\theta^{(t +1)}}\#\gamma_{Y^{(t+1)}})\right]_{u=u^{(t)}}\)  end for
```

**Algorithm 1**Training a NN on the SW loss with Stochastic Gradient Descent

In the following, we will apply results from (Bianchi 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:\mathbb{R}^{d_{u}}\times\Xi\longrightarrow\mathbb{R}\) that is locally Lipschitz in the first variable, and \(\zeta\) a probability measure on \(\Xi\subset\mathbb{R}^{d}\) which is the law of the samples drawn at each SGD iteration. Consider \(\varphi:\mathbb{R}^{d_{u}}\times\Xi\longrightarrow\mathbb{R}^{d_{u}}\) an _almost-everywhere gradient_ of \(f\), which is to say that for almost every \((u,S)\in\mathbb{R}^{d_{u}}\times\Xi,\ \varphi(u,S)=\partial_{u}f(u,S)\) (since each \(f(\cdot,S)\) is locally Lipschitz, it is differentiable almost-everywhere by Rademacher's theorem). The complete loss function is \(F:=u\longrightarrow\int_{\Xi}f(u,S)\mathrm{d}\zeta(S)\). An SGD trajectory of step \(\alpha>0\) for \(F\) is a sequence \((u^{(t)})\in(\mathbb{R}^{d_{u}})^{\mathbb{N}}\) of the form:

\[u^{(t+1)}=u^{(t)}-\alpha\varphi(u^{(t)},S^{(t+1)}),\quad\left(u^{(0)},(S^{(t) })_{t\in\mathbb{N}}\right)\sim\nu\otimes\zeta^{\otimes\mathbb{N}},\]

where \(\nu\) is the distribution of the initial position \(u^{(0)}\). Within this framework, we define an SGD scheme described by Algorithm 1, with \(\zeta:=\mathsf{x}^{\otimes n}\otimes\gamma^{\otimes n}\otimes\mathsf{o}\) and

\[f:=\left\{\begin{array}{rcl}\mathbb{R}^{d_{u}}\times\mathbb{R}^{n\times d_{x} }\times\mathbb{R}^{n\times d_{y}}\times\mathbb{S}^{d_{y}-1}&\longrightarrow& \mathbb{R}^{d_{y}}\\ (u,X,Y,\theta)&\longmapsto&\mathrm{W}^{2}_{2}(P_{\theta}\#T_{u}\#\gamma_{X},P _{\theta}\#\gamma_{Y})\end{array}\right..\]

With this definition for \(f\), we have \(F(u)=\mathbb{E}_{(X,Y,\theta)\sim\zeta}\left[\mathrm{W}^{2}_{2}(P_{\theta}\#T_ {u}\#\gamma_{X},P_{\theta}\#\gamma_{Y})\right]=\mathrm{SW}^{2}_{2}(T_{u}\# \mathsf{x},\mathsf{y})\): the complete loss compares the data \(\mathsf{y}\) with the model's generation \(T_{u}\#\mathsf{x}\) using SW. We now wish to define an almost-everywhere gradient of \(f\). To this end, notice that one may write \(f(u,X,Y,\theta)=w_{\theta}(T(u,X),Y)\), where for \(Z,Y\in\mathbb{R}^{n\times d_{y}}\) and \(\theta\in\mathbb{S}^{d_{y}-1},\ w_{\theta}(Y,Z):=\mathrm{W}^{2}_{2}(P_{\theta} \#\gamma_{Z},P_{\theta}\#\gamma_{Y})\). The differentiability properties of \(w_{\theta}(\cdot,Y)\) are already known (Tanguy et al., 2023; Bonneel et al., 2015), in particular one has the following almost-everywhere gradient of \(w_{\theta}(\cdot,Y)\) :

\[\frac{\partial w_{\theta}}{\partial Z}(Z,Y)=\left(\frac{2}{n}\theta\theta^{T}(z _{k}-y_{\sigma_{\theta}^{2,Y}(k)})\right)_{k\in[\![1,n]\!]}\in\mathbb{R}^{n \times d_{y}},\]

where the permutation \(\sigma_{\theta}^{Z,Y}\in\mathfrak{S}_{n}\) is \(\tau_{Y}^{\theta}\circ(\tau_{Z}^{\theta})^{-1}\), with \(\tau_{Y}^{\theta}\in\mathfrak{S}_{n}\) being a sorting permutation of the list \((\theta\cdot y_{1},\cdots,\theta\cdot y_{n})\). The sorting permutations are chosen arbitrarily when there is ambiguity. To
define an almost-everywhere gradient, we must differentiate \(f(\cdot,X,Y,\theta)=u\longmapsto w_{\theta}(T(u,X),Y)\) for which we need regularity assumptions on \(T\): this is the goal of Assumption 1. In the following, \(\overline{A}\) denotes the topological closure of a set \(A\), \(\partial A\) its boundary, and \(\lambda_{\mathbb{R}^{d_{u}}}\) denotes the Lebesgue measure of \(\mathbb{R}^{d_{u}}\).

**Assumption 1**.: _For every \(x\in\mathbb{R}^{d_{x}},\;\) there exists a family of disjoint connected open sets \((\mathcal{U}_{j}(x))_{j\in J(x)}\) such that \(\forall j\in J(x),\;T(\cdot,x)\in\mathcal{C}^{2}(\mathcal{U}_{j}(x),\mathbb{R }^{d_{y}})\), \(\bigcup\limits_{j\in J(x)}\overline{\mathcal{U}_{j}(x)}=\mathbb{R}^{d_{u}}\) and \(\lambda_{\mathbb{R}^{d_{u}}}\big{(}\bigcup\limits_{j\in J(x)}\partial \mathcal{U}_{j}(x)\big{)}=0\)._

Note that for measure-theoretic reasons, the sets \(J(x)\) are assumed countable.

Assumption 1 implies that given \(X,Y,\theta\) fixed, \(f(\cdot,X,Y,\theta)\) is differentiable almost-everywhere, and that one may define the following almost-everywhere gradient (3).

\[\varphi:=\left\{\begin{array}{rcl}\mathbb{R}^{d_{u}}\times\mathbb{R}^{n \times d_{x}}\times\mathbb{R}^{n\times d_{y}}\times\mathbb{S}^{d_{y}-1}& \longrightarrow&\mathbb{R}^{d_{u}}\\ (u,X,Y,\theta)&\longmapsto&\sum\limits_{k=1}^{n}\frac{2}{n}\bigg{(} \frac{\partial T}{\partial u}(u,x_{k})\bigg{)}^{T}\,\theta\theta^{T}(T(u,x_{ k})-y_{\sigma_{\theta}^{T(u,X),Y}(k)})\end{array}\right., \tag{3}\]

where for \(x\in\mathbb{R}^{d_{x}},\;\frac{\partial T}{\partial u}(u,x)\in\mathbb{R}^{d_{ y}\times d_{u}}\) denotes the matrix of the differential of \(u\longmapsto T(u,x)\), which is defined for almost-every \(u\). Given \(u\in\partial\mathcal{U}_{j}(x)\) (a point of potential non-differentiability), take instead \(0\). (Any choice at such points would still define an a.e. gradient, and will make no difference).

Given a step \(\alpha>0\), and an initial position \(u^{(0)}\sim\nu\), we may now define formally the following fixed-step SGD scheme for \(F\):

\[\begin{split} u^{(t+1)}=u^{(t)}-\alpha\varphi(u^{(t)},X^{(t+1)},Y ^{(t+1)},\theta^{(t+1)}),\\ \Big{(}u^{(0)},(X^{(t)})_{t\in\mathbb{N}}\;(Y^{(t)})_{t\in \mathbb{N}}\;(\theta^{(t)})_{t\in\mathbb{N}}\Big{)}\sim\nu\otimes\mathbf{x}^ {\otimes\mathbb{N}}\otimes\mathbf{\gamma}^{\otimes\mathbb{N}}\otimes\mathbf{ \sigma}^{\otimes\mathbb{N}}.\end{split} \tag{4}\]

An important technicality that we must verify in order to apply Bianchi et al. (2022)'s results is that \(u\longmapsto f(u,X,Y,\theta)\) and \(F\) are locally Lipschitz. Before proving those claims, we reproduce a useful Property from (Tanguy et al., 2023). In the following, \(\|X\|_{\infty,2}\) denotes \(\max\limits_{k\in[\![1,n]\!]}\;\|x_{k}\|_{2}\) given \(X=(x_{1},\cdots,x_{n})\in\mathbb{R}^{n\times d_{x}}\), and \(B_{\mathcal{N}}(x,r)\) for \(\mathcal{N}\) a norm on \(\mathbb{R}^{d_{x}}\), \(x\in\mathbb{R}^{d_{x}}\) and \(r>0\) shall denote the open ball of \(\mathbb{R}^{d_{x}}\) of centre \(x\) and radius \(r\) for the norm \(\mathcal{N}\) (if \(\mathcal{N}\) is omitted, then \(B\) is an euclidean ball).

**Proposition 1**.: _The \((w_{\theta}(\cdot,Y))_{\theta\in\mathbb{S}^{d_{y}-1}}\) are uniformly locally Lipschitz (Tanguy et al., 2023)._

_Let \(\kappa_{r}(Z,Y):=2n(r+\|Z\|_{\infty,2}+\|Y\|_{\infty,2})\), for \(Z,Y\in\mathbb{R}^{n\times d_{y}}\) and \(r>0\). Then \(w_{\theta}(\cdot,Y)\) is \(\kappa_{r}(Z,Y)\)-Lipschitz in the neighbourhood \(B_{\|\cdot\|_{\infty},(Z,r)}\):_

\[\forall Y^{\prime},Y^{\prime\prime}\in B_{\|\cdot\|_{\infty,2}}(Z,r),\;\forall \theta\in\mathbb{S}^{d_{y}-1},\;|w_{\theta}(Y^{\prime},Y)-w_{\theta}(Y^{ \prime\prime},Y)|\leq\kappa_{r}(Z,Y)\|Y^{\prime}-Y^{\prime\prime}\|_{\infty, 2}.\]

In order to deduce regularity results on \(f\) and \(F\) from Proposition 1, we will make the following assumption, which under Assumption 1 only requires additional regularity with respect to the data argument.

**Assumption 2**.: _For any compacts \(\mathcal{K}_{1}\subset\mathbb{R}^{d_{u}}\) and \(\mathcal{K}_{2}\subset\mathbb{R}^{d_{x}},\;\) there exists \(L_{\mathcal{K}_{1},\mathcal{K}_{2}}>0\) such that \(\forall(u_{1},u_{2},x_{1},x_{2})\in\mathcal{K}_{1}^{2}\times\mathcal{K}_{2}^{2},\;\|T(u_{1},x_{1})-T(u_{2},x_{2})\|\leq L_{\mathcal{K}_{1},\mathcal{K}_{2}} \left(\|u_{1}-u_{2}\|+\|x_{1}-x_{2}\|\right)\)._

**Proposition 2** (Regularity of \(u\longmapsto f(u,X,Y,\theta)\)).: _Under Assumption 2, for \(\varepsilon>0,\;u_{0}\in\mathbb{R}^{d_{u}},\;X\in\mathbb{R}^{n\times d_{x}},\; Y\in\mathbb{R}^{n\times d_{y}}\) and \(\theta\in\mathbb{S}^{d_{y}-1}\), let \(\kappa_{\varepsilon}(u_{0},X,Y):=2Ln(\varepsilon L+\|T(u_{0},X)\|_{\infty,2}+\| Y\|_{\infty,2})\), with \(L:=L_{\overline{B}(u_{0},\varepsilon),\overline{B}(\mathbb{0}_{d_{x}},\|X\|_{ \infty,2})}\). Then \(f(\cdot,X,Y,\theta)\) is \(\kappa_{\varepsilon}(X,Y)\)-Lipschitz in \(B(u_{0},\varepsilon)\):_

\[\forall u,u^{\prime}\in B(u_{0},\varepsilon),\;|f(u,X,Y,\theta)-f(u^{\prime},X,Y, \theta)|\leq\kappa_{\varepsilon}(X,Y)\|u-u^{\prime}\|_{2}.\]
Proof.: Let \(\varepsilon>0,\ u_{0}\in\mathbb{R}^{d_{u}},\ X\in\mathbb{R}^{n\times d_{x}},\ Y\in \mathbb{R}^{n\times d_{y}}\) and \(\theta\in\mathbb{S}^{d_{y}-1}\). Let \(u,u^{\prime}\in B(u_{0},\varepsilon)\). Using Assumption 2, we have \(T(u,X),T(u^{\prime},X)\in B_{\|\cdot\|_{\infty,2}}(T(u_{0},X),r)\), with \(r:=\varepsilon L_{\overline{B}(u_{0},\varepsilon),\overline{B}(0_{\text{d}x},\|X\|_{\infty,2})}\).

By Proposition 1, we have, with \(L:=L_{\overline{B}(u_{0},\varepsilon),\overline{B}(0_{\text{d}x},\|X\|_{ \infty,2})}\)

\[|f(u,X,Y,\theta)-f(u^{\prime},X,Y,\theta)| =|w_{\theta}(T(u,X),Y)-w_{\theta}(T(u^{\prime},X),Y)|\] \[\leq\kappa_{r}(T(u_{0},X),Y)\|T(u,X)-T(u^{\prime},X)\|_{\infty,2}\] \[\leq 2n(\varepsilon L+\|T(u_{0},X)\|_{\infty,2}+\|Y\|_{\infty,2}) L\|u-u^{\prime}\|_{2}.\]

Proposition 2 shows that \(f\) is locally Lipschitz in \(u\). We now assume some conditions on the measures \(\mathtt{x}\) and \(\mathtt{y}\) in order to prove that \(F\) is also locally Lipschitz.

**Assumption 3**.: \(\mathtt{x}\) _and \(\mathtt{y}\) are Radon probability measures on \(\mathbb{R}^{d_{x}}\) and \(\mathbb{R}^{d_{y}}\) respectively, supported by the compacts \(\mathcal{X}\) and \(\mathcal{Y}\) respectively. Denote \(R_{x}:=\sup\limits_{x\in\mathcal{X}}\|x\|_{2}\) and \(R_{y}:=\sup\limits_{y\in\mathcal{Y}}\|y\|_{2}\)._

**Proposition 3**.: _Assume Assumption 2 and Assumption 3. For \(\varepsilon>0,\ u_{0}\in\mathbb{R}^{d_{u}},\) let \(C_{1}(u_{0}):=\int_{\mathcal{X}^{n}}\|T(u_{0},X)\|_{\infty,2}\mathrm{d}\psi^{ \otimes n}(X),\ C_{2}:=\int_{\mathcal{Y}^{n}}\|Y\|_{\infty,2}\mathrm{d}\psi^{ \otimes n}(Y)\) and \(L:=L_{\overline{B}(u_{0},\varepsilon),\overline{B}(0,R_{x})}\). Let \(\kappa_{\varepsilon}(u_{0}):=2Ln(\varepsilon L+C_{1}(u_{0})+C_{2})\). We have \(\forall u,u^{\prime}\in B(u_{0},\varepsilon),\ |F(u)-F(u^{\prime})|\leq\kappa_{ \varepsilon}(u_{0})\|u-u^{\prime}\|_{2}\)._

Proof.: Let \(\varepsilon>0,\ u_{0}\in\mathbb{R}^{d_{u}}u,u^{\prime}\in B(u_{0},\varepsilon)\). First, notice that for any \(X\in\mathcal{X}^{n},\ \|X\|_{\infty,2}\leq R_{x}\), thus \(L_{\overline{B}(u_{0},\varepsilon),\overline{B}(0_{\text{d}x},\|X\|_{\infty,2 })}\leq L_{\overline{B}(u_{0},\varepsilon),\overline{B}(0,R_{x})}=:L\). We have

\[|F(u)-F(u^{\prime})| \leq\int_{\mathcal{X}^{n}\times\mathcal{Y}^{n}\times\mathcal{S} ^{d_{y}-1}}|f(u,X,Y,\theta)-f(u)|\mathrm{d}x^{\otimes n}(X)\mathrm{d}\psi^{ \otimes n}(Y)\mathrm{d}\sigma(\theta)\] \[\leq\int_{\mathcal{X}^{n}\times\mathcal{Y}^{n}}\kappa_{\varepsilon }(u_{0},X,Y)\|u-u^{\prime}\|_{2}\mathrm{d}x^{\otimes n}(X)\mathrm{d}\psi^{ \otimes n}(Y)\] \[\leq\int_{\mathcal{X}^{n}\times\mathcal{Y}^{n}}2Ln(\varepsilon L +\|T(u_{0},X)\|_{\infty,2}+\|Y\|_{\infty,2})\|u-u^{\prime}\|_{2}\mathrm{d}x^{ \otimes n}(X)\mathrm{d}y^{\otimes n}(Y).\]

Now by Assumption 2, \(X\longmapsto\|T(u_{0},X)\|_{\infty,2}\) is continuous on the compact \(\mathcal{X}^{n}\), allowing the definition of the constant \(C_{1}(u_{0}):=\int_{\mathcal{X}^{n}}\|T(u_{0},X)\|_{\infty,2}\mathrm{d}\psi^{ \otimes n}(X)\) (\(\mathtt{x}\) is a Radon probability measure by Assumption 3, thus \(C_{1}\) is finite.) Likewise, let \(C_{2}:=\int_{\mathcal{Y}^{n}}\|Y\|_{\infty,2}\mathrm{d}\mathcal{Y}^{\otimes n}( Y)<+\infty\).

Finally, \(|F(u)-F(u^{\prime})|\leq 2Ln(\varepsilon L+C_{1}(u_{0})+C_{2})\|u-u^{\prime}\|_{2}\). 

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 \(\partial_{C}F\). The set of Clarke sub-gradients at a points \(u\) is the convex hull of the limits of gradients of \(F\):

\[\partial_{C}F(u):=\mathrm{conv}\left\{v\in\mathbb{R}^{d_{u}}:\ \exists(u_{i})\in( \mathcal{D}_{F})^{\mathbb{N}}:u_{i}\xrightarrow[i\longrightarrow+\infty]{}u\ \text{ and }\nabla F(u_{i})\xrightarrow[i\longrightarrow+\infty]{}v\right\},\]

where \(\mathcal{D}_{F}\) is the set of differentiability of \(F\). At points \(u\) where \(F\) is differentiable, \(\partial_{C}F(u)=\{\nabla F(u)\}\), and if \(F\) is convex in a neighbourhood of \(u\), then the Clarke differential at \(u\) is the set of its convex sub-gradients.

## 3 Convergence of Interpolated SGD Trajectories on \(F\)

In general, the idea behind SGD is a discretisation of the gradient flow equation \(\dot{u}(s)=-\nabla F(u(s))\). In our non-smooth setting, the underlying continuous-time problem is is instead the Clarke differential
inclusion \(\dot{u}(s)\in-\partial_{C}F(u(s))\). 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(\cdot)\in\mathcal{C}_{\mathrm{abs}}(\mathbb{R}_{+},\mathbb{R}^{d_{u}})\)) and start within \(\mathcal{K}\subset\mathbb{R}^{d_{u}}\), a fixed compact set. We can now define the solution set formally as

\[S_{-\partial_{C}F}(\mathcal{K}):=\left\{u\in\mathcal{C}_{\mathrm{abs}}( \mathbb{R}_{+},\mathbb{R}^{d_{u}})\ |\ \underline{\forall}s\in\mathbb{R}_{+},\ \dot{u}(s)\in-\partial_{C}F(u(s));\ u(0)\in \mathcal{K}\right\}, \tag{5}\]

where we write \(\underline{\forall}\) for \({}^{*}\)almost every\({}^{*}\). In order to compare the discrete SGD trajectories to this set continuous-time trajectories, we interpolate the discrete points in an affine manner: Equation (6) defines the _piecewise-affine interpolated SGD trajectory_ associated to an SGD trajectory \((u^{(t)}_{\alpha})_{t\in\mathbb{N}}\) of learning rate \(\alpha\).

\[u_{\alpha}(s)=u^{(t)}_{\alpha}+\left(\frac{s}{\alpha}-t\right)(u^{(t+1)}_{ \alpha}-u^{(t)}_{\alpha}),\quad\forall s\in[t\alpha,(t+1)\alpha[,\quad\forall t \in\mathbb{N}. \tag{6}\]

In order to compare our interpolated trajectories with the solutions, we consider the metric of uniform convergence on all segments

\[d_{c}(u,u^{\prime}):=\sum_{k\in\mathbb{N}^{*}}\frac{1}{2^{k}}\mathrm{min} \left(1,\max_{s\in[0,k]}\lVert u(s)-u^{\prime}(s)\rVert_{\infty,2}\right). \tag{7}\]

In order to prove that the interpolated trajectories, we will leverage the results of Bianchi et al. (2022) which hinge on three conditions on the loss \(F\) that we reproduce and verify successively.

**Condition 1**.:
1. _There exists_ \(\kappa:\mathbb{R}^{d_{u}}\times\Xi\longrightarrow\mathbb{R}_{+}\) _measurable such that each_ \(\kappa(u,\cdot)\) _is_ \(\zeta\)_-integrable, and:_ \[\exists\varepsilon>0,\ \forall u,u^{\prime}\in B(u_{0},\varepsilon),\ \forall S\in\Xi,\ |f(u,S)-f(u^{\prime},S)|\leq \kappa(u_{0},S)\lVert u-u^{\prime}\rVert_{2}.\]
2. _There exists_ \(u\in\mathbb{R}^{d_{u}}\) _such that_ \(f(u,\cdot)\) _is_ \(\zeta\)_-integrable._

Our regularity result on \(f\) Proposition2 allows us to verify Condition1, by letting \(\varepsilon:=1\) and \(\kappa(u_{0},S):=\kappa_{1}(u_{0},X,Y,\theta)\). Condition1 ii) is immediate since for _all_\(u\in\mathbb{R}^{d_{u}}\), \((X,Y,\theta)\longmapsto w_{\theta}(T(u,X),Y)\) is continuous in each variable separately, thanks to the regularity of \(T\) provided by Assumption2, and to the regularities of \(w\) (as implied by (Tanguy et al., 2023), Lemma2.2.2, for instance). This continuity implies that all \(f(u,\cdot)\) are \(\zeta\)-integrable, since \(\zeta=\mathfrak{x}^{\otimes n}\otimes\mathfrak{y}^{\otimes n}\otimes\sigma\) is a compactly supported Radon measure under Assumption3.

**Condition 2**.: _The function \(\kappa\) of Condition1 verifies:_

1. _There exists_ \(c\geq 0\) _such that_ \(\forall u\in\mathbb{R}^{d_{u}},\ \int_{\Xi}\kappa(u,S)\mathrm{d}\zeta(S)\leq c(1+ \lVert u\rVert_{2})\)_._
2. _For every compact_ \(\mathcal{K}\subset\mathbb{R}^{d_{u}}\)_,_ \(\sup_{u\in\mathcal{K}}\ \int_{\Xi}\kappa(u,S)^{2}\mathrm{d}\zeta(S)<+\infty\)_._

Condition2.ii) is verified by \(\kappa\) given its regularity. However, Condition2.i) requires that \(T(u,x)\) increase slowly as \(\lVert u\rVert\) increases, which is more costly.

**Assumption 4**.: _There exists an \(\kappa\)-integrable function \(g:\mathbb{R}^{d_{x}}\longrightarrow\mathbb{R}_{+}\) such that \(\forall u\in\mathbb{R}^{d_{u}},\ \forall x\in\mathbb{R}^{d_{x}},\ \lVert T(u,x) \rVert\leq g(x)(1+\lVert u\rVert_{2})\)._

Assumption4 is satisfied in particular as soon as \(T(\cdot,x)\) is bounded (which is the case for a neural network with bounded activation functions), or if \(T\) is of the form \(T(u,x)=\tilde{T}(u,x)\mathbb{1}_{B(0,R)}(u)\), i.e. limiting the network parameters \(u\) to be bounded. This second case does not yield substantial restrictions in practice, yet vastly simplifies theory.
Under Assumption 4, we have for any \(u\in\mathbb{R}^{d_{u}}\), with \(\kappa\) from Proposition 2 and \(C_{2}\) from Proposition 3,

\[\int_{\mathcal{X}^{n}\times\mathcal{Y}^{n}\times\mathbb{S}^{d_{y}-1 }}\kappa_{1}(u,X,Y,\theta)\mathrm{d}\varkappa^{\otimes n}(X)\mathrm{d}\gamma^{ \otimes n}(Y)\mathrm{d}\sigma(\theta) \leq 2Ln\left(\varepsilon L+(1+\|u\|_{2})\int_{\mathcal{X}^{n} \,\mathbb{R}\in[1,n]}\,g(x_{k})\mathrm{d}\varkappa^{\otimes n}(X)+C_{2}\right)\] \[\leq c(1+\|u\|_{2}).\]

As a consequence, Condition 2 holds under our assumptions. We now consider the Markov kernel associated to the SGD schemes:

\[P_{\alpha}:\left\{\begin{array}{rcl}\mathbb{R}^{d_{u}}\times\mathcal{B}( \mathbb{R}^{d_{u}})&\longrightarrow&[0,1]\\ u,B&\longmapsto&\int_{\Xi}\mathbb{1}_{B}(u-\alpha\varphi(u,S))\mathrm{d}\zeta(S )\end{array}\right..\]

With \(\lambda_{\mathbb{R}^{d_{u}}}\) denoting the Lebesgue measure on \(\mathbb{R}^{d_{u}}\), let \(\Gamma:=\{\alpha\in\,]0,+\infty[\,\mid\,\forall\rho\ll\lambda_{\mathbb{R}^{d_ {u}}},\ \rho P_{\alpha}\ll\lambda_{\mathbb{R}^{d_{u}}}\}\). We will verify the following condition:

**Condition 3**.: _The closure of \(\Gamma\) contains 0._

In order to satisfy Condition 3, we require an additional regularity condition on the neural network \(T\) which we formulate in Assumption 5.

**Assumption 5**.: _There exists a constant \(M>0\), such that (with the notations of Assumption 1 and Assumption 3) \(\forall x\in\mathcal{X},\ \forall j\in J(x),\ \forall u\in\mathcal{U}_{j}(x),\ \forall(i_{1},i_{2},i_{3},i_{4})\in \llbracket 1,d_{u}\rrbracket^{2}\times\llbracket 1,d_{y}\rrbracket^{2},\)_

\[\left|\frac{\partial^{2}}{\partial u_{i_{1}}\partial u_{i_{2}}}\Big{(}[T(u,x)] _{i_{3}}[T(u,x)]_{i_{4}}\Big{)}\right|\leq M,\ \mathrm{and}\ \left\|\frac{\partial^{2}T}{ \partial u_{i_{1}}\partial u_{i_{2}}}(u,x)\right\|_{2}\leq M.\]

The upper bounds in assumption bear strong consequences on the behaviour of \(T\) for \(\|u\|_{2}\gg 1\), and are only practical for networks of the form \(T(u,x)=\bar{T}(u,x)\mathbb{1}_{B(0,R)}(u)\), similarly to Assumption 4.

**Proposition 4**.: _Under Assumption 1, Assumption 3 and Assumption 5, for the SGD trajectories (4), \(\Gamma\ \supset\,]0,\alpha_{0}[\), where \(\alpha_{0}:=\frac{1}{(d_{y}^{\ 2}+2R_{y})d_{u}M}.\)_

Proof.: Let \(\rho\ll\lambda\) and \(B\in\mathcal{B}(\mathbb{R}^{d_{y}})\) such that \(\lambda(B)=0\). We have, with \(\alpha^{\prime}:=2\alpha/n,\ S:=(X,Y,\theta),\ \zeta:=x^{\otimes n}\otimes \gamma^{\otimes n}\otimes\sigma\) and \(\Xi:=\mathcal{X}^{n}\times\mathcal{Y}^{n}\times\mathbb{S}^{d_{y}-1}\),

\[\rho P_{\alpha}(B)=\int_{\mathbb{R}^{d_{u}}\times\Xi}\mathbb{1}_{B}\left[u- \alpha^{\prime}\sum_{k=1}^{n}\left(\frac{\partial T}{\partial u}(u,x_{k}) \right)^{T}\theta\theta^{T}(T(u,x_{k})-y_{\sigma_{\theta}^{T(u,X),Y}(k)}) \right]\mathrm{d}\rho(u)\mathrm{d}\zeta(S)\ \ \leq\sum_{\tau\in\mathfrak{S}_{n}}\int_{\Xi}I_{\tau}(S) \mathrm{d}\zeta(S)\]

where \(I_{\tau}(S):=\int_{\mathbb{R}^{d_{u}}}\mathbb{1}_{B}\left(\phi_{\tau,S}(u) \right)\mathrm{d}\rho(u)\), with \(\phi_{\tau,S}:=u-\alpha^{\prime}\underbrace{\sum_{k=1}^{n}\left(\frac{ \partial T}{\partial u}(u,x_{k})\right)^{T}\theta\theta^{T}(T(u,x_{k})-y_{\tau (k)})}_{\psi_{\tau,S}:=}\).

Let \(\tau\in\mathfrak{S}_{n}\) and \((X,Y,\theta)\in\Xi\). Using Assumption 1, separate \(I_{\tau}(S)=\sum_{j\in J}\int_{\mathcal{U}_{j}(X)}\mathbb{1}_{B}\left(u-\psi_{ \tau,S}(u)\right)\mathrm{d}\rho(u)\), where the differentiability structure \((\mathcal{U}_{j}(X))_{j\in J(X)}\) is obtained using the respective differentiability structures: for each \(k\in\llbracket 1,n\rrbracket\), Assumption 1 yields a structure \((\mathcal{U}_{j_{k}}(x_{k}))_{j_{k}\in J_{k}(x_{k})}\) of \(u\longmapsto T(u,x_{k})\), which depends on \(x_{k}\), hence the \(k\) indices.

To be precise, define for \(j=(j_{1},\cdots,j_{n})\in J_{1}(x_{1})\times\cdots\times J_{n}(x_{n}),\ \mathcal{U}_{j}(X):=\bigcap_{k=1}^{n} \mathcal{U}_{j_{k}}(x_{k})\), and \(J(X):=\{(j_{1},\cdots,j_{n})\in J_{1}(x_{1})\times\cdots\times J_{n}(x_{n}) \mid\mathcal{U}_{j}(X)\neq\varnothing\}\). In particular, for any \(k\in\llbracket 1,n\rrbracket,\ T(\cdot,x_{k})\) is \(\mathcal{C}^{2}\) on \(\mathcal{U}_{j}(X)\). Notice that the derivatives are not necessarily defined on the border \(\partial\mathcal{U}_{j}(X)\), which is of Lebesgue measure 0 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 \(+\infty\), depending on the behaviour of \(\phi_{\tau,s}\), but we shall see that they are all finite when \(\alpha\) is small enough).

We drop the \(S,\tau\) index in the notation, and focus on the properties of \(\phi\) and \(\psi\) as functions of \(u\). Our first objective is to determine a constant \(K>0\), independent of \(u,S,\tau\), such that \(\psi\) is \(K\)-Lipschitz on \(\mathcal{U}_{j}(X)\).

First, let \(\chi:=u\in\mathcal{U}_{j}(X)\longmapsto\left(\dfrac{\partial T}{ \partial u}(u,x_{k})\right)^{T}\theta\theta^{T}T(u,x_{k})\). \(\chi\) is of class \(\mathcal{C}^{1}\), therefore we determine its Lipschitz constant by upper-bounding the \(\|\cdot\|_{2}\)-induced operator norm of its differential, denoted by \(\left|\dfrac{\partial\chi}{\partial u}(u)\right|_{2}\). Notice that \(\chi(u)=\dfrac{1}{2}\dfrac{\partial}{\partial u}(\theta\cdot T(u,x_{k}))^{2}\).

\[\text{Now }\left|\dfrac{\partial^{2}}{\partial u^{2}}(\theta \cdot T(u,x_{k}))^{2}\right|\hskip-1.0pt\Bigg{|}_{2}\leq d_{u}\max_{(i_{1},i_ {2})\in[\![1,d_{u}]^{2}]}\;\left|\dfrac{\partial^{2}}{\partial u_{i_{1}} \partial u_{i_{2}}}(\theta\cdot T(u,x_{k}))^{2}\right|\hskip-1.0pt\Bigg{|}, \text{with by Assumption \ref{eq:SGD},}\] \[\left|\dfrac{\partial^{2}}{\partial u_{i_{1}}\partial u_{i_{2}}}( \theta\cdot T(u,x_{k}))^{2}\right|\leq\sum_{(i_{3},i_{4})\in[\![1,d_{u}]^{2}]} \left|\theta_{i_{3}}\theta_{i_{4}}\dfrac{\partial^{2}}{\partial u_{i_{1}} \partial u_{i_{2}}}\!\!\left(\![T(u,x_{k})]_{i_{3}}[T(u,x_{k})]_{i_{4}}\right) \right|\leq d_{y}{}^{2}M.\]

We obtain that \(\chi\) is \(\frac{1}{2}d_{u}d_{y}{}^{2}M\)-Lipschitz.

Second, let \(\omega:u\in\mathcal{U}_{j}(X)\longmapsto\left(\dfrac{\partial T}{ \partial u}(u,x_{k})\right)^{T}\theta\theta^{T}y_{\tau(k)}\), also of class \(\mathcal{C}^{1}\). We re-write \(\left[\dfrac{\partial\omega}{\partial u}(u)\right]_{i_{1},i_{2}}=y_{\tau(k)}^ {T}\theta\theta^{T}\dfrac{\partial^{2}T}{\partial u_{i_{1}}\partial u_{i_{2}} }(u,x_{k})\), and conclude similarly by Assumption 5 that \(\omega\) is \(\|y_{\tau(k)}\|_{2}d_{u}M\)-Lipschitz.

Finally, \(\psi=\sum\limits_{k=1}^{n}(\chi_{k}-\omega_{k})\), and is therefore \(K:=(\frac{1}{2}d_{y}{}^{2}+R_{y})d_{u}nM\)-Lipschitz, with \(R_{y}\) from Assumption 3. We have proven that \(\left|\dfrac{\partial\psi}{\partial u}(u)\right|\hskip-1.0pt\Bigg{|}_{2}\leq K\) for any \(u\in\mathcal{U}_{j}(X)\), and that \(K\) does not depend on \(X,Y,\theta,j\) or \(u\).

We now suppose that \(\alpha^{\prime}<\frac{1}{K}\), which is to say \(\alpha<\frac{n}{2K}\). Under this condition, \(\phi:\mathcal{U}_{j}(X)\longrightarrow\mathbb{R}^{d_{u}}\) is injective. Indeed, if \(\phi(u_{1})=\phi(u_{2})\), then \(\|u_{1}-u_{2}\|_{2}=\alpha^{\prime}\|\psi(u_{1})-\psi(u_{2})\|_{2}\leq\alpha^{ \prime}K\|u_{1}-u_{2}\|_{2}\), thus \(u_{1}=u_{2}\). Furthermore, for any \(u\in\mathcal{U}_{j}(X)\), \(\dfrac{\partial\phi}{\partial u}(u)=\mathrm{Id}_{\mathbb{R}^{d_{u}}}-\alpha^ {\prime}\dfrac{\partial\psi}{\partial u}(u)\), with \(\left|\hskip-1.0pt\Bigg{|}\alpha^{\prime}\dfrac{\partial\psi}{\partial u}(u) \right|\hskip-1.0pt\Bigg{|}_{2}<1\), thus the matrix \(\dfrac{\partial\phi}{\partial u}(u)\) is invertible (using the Neumann series method). By the global inverse function theorem, \(\phi:\mathcal{U}_{j}(X)\longrightarrow\phi(\mathcal{U}_{j}(X))\) is a \(\mathcal{C}^{1}\)-diffeomorphism.

Re-writing \(\int_{\mathcal{U}_{j}(X)}\hskip-1.0pt\mathbb{1}_{B}(\phi(u))\mathrm{d}\rho(u)= \phi\#\rho(B)\), we have now shown that \(\phi\) is a \(\mathcal{C}^{1}\)-diffeomorphism, thus since \(\rho\ll\lambda\), \(\phi\#\rho\ll\lambda\). It then follows that the integral is \(0\), then \(I_{\tau}(S)=0\) and finally \(\rho P_{\alpha}(B)=0\). 

Now that we have verified Condition 1, Condition 2 and Condition 3, we can apply (Bianchi et al., 2022), Theorem 2 to \(F\). Let \(\alpha_{1}<\alpha_{0}\) (see Proposition 4).

**Theorem 1** (Convergence of the interpolated SGD trajectories).: _Consider a neural network \(T\) and measures \(\star\), \(\forall\) satisfying Assumption 1, Assumption 2, Assumption 3, Assumption 4 and Assumption 5. Let \((u_{\alpha}^{(l)}),\alpha\in]0,\alpha_{1}],t\in\mathbb{N}\) a collection of SGD trajectories associated to (4). Consider \((u_{\alpha})\) their associated interpolations. For any compact \(\mathcal{K}\subset\mathbb{R}^{d_{u}}\) and any \(\varepsilon>0\), we have:_

\[\lim_{\begin{subarray}{c}\alpha\longrightarrow 0\\ \alpha\in[0,\alpha_{1}]\end{subarray}}\nu\otimes\pi^{\otimes\mathbb{N}}\otimes \gamma^{\otimes\mathbb{N}}\otimes\mathfrak{o}^{\otimes\mathbb{N}}(d_{c}(u_{ \alpha},S_{-\partial_{C}F}(\mathcal{K}))>\varepsilon)=0. \tag{8}\]
The distance \(d_{c}\) is defined in (7). As the learning rate decreases, the interpolated trajectories approach the trajectory set \(S_{-\partial_{C}F}\), which is essentially a solution of the _gradient flow equation_\(\dot{u}(s)=-\nabla F(u(s))\) (ignoring the set of non-differentiability, which is \(\lambda_{\mathbb{R}^{d_{u}}}\)-null). To get a tangible idea of the concepts at play, if \(F\) was \(\mathcal{C}^{2}\) 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 \(s\longrightarrow+\infty\). 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 a 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,\cdot)\) stay Lipschitz, which is desirable for theoretical reasons. For instance the 1-Wasserstein duality in Wasserstein GANs (Arjovsky et al., 2017) requires Lipschitz networks, and similarly, Sliced-Wasserstein GANs (Deshpande et al., 2018) use weight clipping and enforce their networks to be Lipschitz.

Given a radius \(R_{u}>0\), we consider SGD schemes that are restricted to \(u\in\overline{B}(0,r)=:B_{r}\), by performing _projected_ SGD. At each step \(t\), we also add a noise \(a\varepsilon^{(t+1)}\), where \(\varepsilon^{(t+1)}\) is an additive noise of law \(\eta\ll\lambda_{\mathbb{R}^{u}}\), which is often taken as standard Gaussian in practice. These additions yield the following SGD scheme:

\[\begin{split} u^{(t+1)}=\pi_{r}\left(u^{(t)}-\alpha\varphi(u^{(t) },X^{(t+1)},Y^{(t+1)},\theta^{(t+1)})+\alpha a\varepsilon^{(t+1)}\right),\\ \left(u^{(0)},(X^{(t)})_{t\in\mathbb{N}}\ (Y^{(t)})_{t\in \mathbb{N}},\ (\theta^{(t)})_{t\in\mathbb{N}},\ (\varepsilon^{(t)})_{t\in \mathbb{N}}\right)\sim\nu\otimes\mathbb{x}^{\otimes\mathbb{N}}\otimes\gamma^{ \otimes\mathbb{N}}\otimes\sigma^{\otimes\mathbb{N}}\otimes\eta^{\otimes\mathbb{ N}},\end{split} \tag{9}\]

where \(\pi_{r}:\mathbb{R}^{u}\longrightarrow B_{r}\) denotes the orthogonal projection on the ball \(B_{r}:=\overline{B}(0,r)\). Thanks to Condition 1, 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 (9). 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 almost-everywhere differentiability by the notion of _path differentiability_, as studied thoroughly in (Bolte and Pauwels, 2021). We formulate this condition from (Bianchi et al., 2022) before presenting sufficient conditions on \(T\) under which path differentiability shall hold.

**Condition 4**.: \(F\) _is path differentiable, which is to say that for any \(u\in\mathcal{C}_{\mathrm{abs}}(\mathbb{R}_{+},\mathbb{R}^{d_{u}})\), for almost all \(t>0,\ \forall v\in\partial_{C}F(u(s)),\ v\cdot\dot{u}(s)=(F\circ u)^{\prime}(s)\)._

Note that by (Bolte and Pauwels, 2021) Corollary 2, \(F\) is path differentiable if and only if \(\partial_{C}F\) is a conservative field for \(F\) (in the sense of (Bolte and Pauwels, 2021), Definition 1) if and only if \(F\) has a chain rule for \(\partial_{C}\) (which is the formulation chosen in Condition 4 by (Bianchi et al., 2022)).

In order to satisfy Condition 4, we need to make the assumption that the NN input measure \(\mathpzc{x}\) and the data measure \(\mathpzc{y}\) are discrete measures, which is the case for \(\mathpzc{y}\) in the case of generative neural networks, but is less realistic for \(\mathpzc{x}\) in practice. We define \(\Sigma_{n}\) the \(n\)-simplex: its elements are the \(a\in\mathbb{R}^{n}\) s.t. \(\forall i\in\llbracket 1,n\rrbracket,\ a_{i}\geq 0\) and \(\sum_{i}a_{i}=1\).

**Assumption 6**.: _One may write \(\mathpzc{x}=\sum_{k=1}^{n_{x}}a_{k}\delta_{x_{k}}\) and \(\mathpzc{y}=\sum_{k=1}^{n_{y}}b_{k}\delta_{y_{k}}\), with the coefficient vectors \(a\in\Sigma_{n_{x}},b\in\Sigma_{n_{y}},\ \mathcal{X}=\{x_{1},\cdots,x_{n_{x}}\}\subset\mathbb{R}^{d_{x}}\) and \(\mathcal{Y}=\{y_{1},\cdots,y_{n_{y}}\}\subset\mathbb{R}^{d_{y}}\)._

There is little practical reason to 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 \(T(\cdot,x)\) is path differentiable. Since \(T(\cdot,x)\) is a vector-valued function, we need to extend the notion of path-differentiability. Thankfully, Bolte and Pauwels (2021) define _conservative mappings_ for vector-valued locally Lipschitz functions (Definition 4), which allows us to define naturally path differentiability of a vector-valued function as the path-differentiability of all of its coordinate functions.

**Assumption 7**.: _For any \(x\in\mathbb{R}^{d_{x}}\), \(T(\cdot,x)\) is path differentiable._

Assumption 7 holds as soon as each \(T(\cdot,x)\) is semi-algebraic (i.e. piecewise polynomial, where the pieces are in finite number and can be written through polynomial equations) or more generally definable (see (Davis et al., 2020), Definition 5.10), as proven by (Davis et al., 2020), Theorem 5.8. This is the case for iterated compositions of linear maps and definable activation functions (such as the widespread sigmoid and ReLU), see (Davis et al., 2020), Corollary 5.11, as well as (Bolte and Pauwels, 2021), SS6.2 for further explanations on suitable NNs.

**Proposition 5**.: _Under Assumption 2, 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 and Pauwels, 2021), Lemma 6.

Let \(\mathcal{E}:\left\{\begin{array}{ccc}\mathbb{R}^{n\times d_{y}}\times \mathbb{R}^{n\times d_{y}}&\longrightarrow&\mathbb{R}_{+}\\ Y,Y^{\prime}&\longmapsto&\mathrm{SW}^{2}_{2}(\gamma_{Y},\gamma_{Y^{\prime}}) \end{array}\right..\) By (Tanguy et al., 2023), Proposition 2.4.3, each \(\mathcal{E}(\cdot,Y)\) is semi-concave and thus is path differentiable (by (Tanguy et al., 2023), Proposition 4.3.3).

Thanks to Assumption 6, \(\mathsf{x}^{\otimes n}\) and \(\mathsf{y}^{\otimes n}\) are discrete measures on \(\mathbb{R}^{n\times d_{x}}\) and \(\mathbb{R}^{n\times d_{y}}\) respectively, allowing one to write \(\mathsf{x}^{\otimes n}=\sum_{k}a_{k}\delta\chi_{k_{k}}\) and \(\mathsf{y}^{\otimes n}=\sum_{l}b_{l}\delta_{Y_{l}}\). Then \(F=u\longmapsto\sum_{k,l}a_{k}b_{l}\mathcal{E}(T(u,X_{k}),Y_{l})\) is path differentiable as a sum ((Bolte and Pauwels, 2021), Corollary 4) of compositions ((Bolte and Pauwels, 2021), Lemma 6) of path differentiable functions. 

We have now satisfied all the assumptions to apply (Bianchi et al., 2022), Theorem 6, showing that trajectories of (9) converge towards \(\mathcal{Z}_{r}\), the set of _Karush-Kahn-Tucker_ points related to the differential inclusion tied to the discrete scheme (9):

\[\mathcal{Z}_{r}:=\left\{u\in\mathbb{R}^{d_{u}}\ |\ 0\in-\partial_{C}F(u)- \mathcal{N}_{r}(u)\right\},\quad\mathcal{N}_{r}(u)=\left\{\begin{array}{c} \{0\}\text{ if }\|u\|_{2}<r\\ \{\lambda u\ |\ \lambda\geq 0\}\text{ if }\|u\|_{2}=r\\ \varnothing\text{ if }\|u\|_{2}>r\end{array}\right., \tag{10}\]

where \(\mathcal{N}_{r}(u)\) refers to the _normal cone_ of the ball \(B(0,r)\) at \(x\). The term \(\mathcal{N}_{r}(u)\) in (10) only makes a difference in the pathological case \(\|u\|_{2}=r\), which never happens in practice since the idea behind projecting is to do so on a very large ball, in order to avoid gradient explosion, to limit the Lipschitz constant and to satisfy theoretical assumptions. Omitting the \(\mathcal{N}_{r}(u)\) term, and denoting \(\mathcal{D}\) the points where \(F\) is differentiable, (10) simplifies to \(\mathcal{Z}_{r}\cap\mathcal{D}=\{u\in\mathcal{D}\ |\ \nabla F(u)=0\}\), i.e. the critical points of \(F\) for the usual differential. Like in Theorem 1, we let \(\alpha_{1}<\alpha_{0}\), where \(\alpha_{0}\) is defined in Proposition 4.

**Theorem 2** (Bianchi et al. (2022), Theorem 6 applied to (9)).: _Consider a neural network \(T\) and measures \(\mathsf{x}\), \(\mathsf{y}\) satisfying Assumption 1, Assumption 2, Assumption 4, Assumption 5, Assumption 6 and Assumption 7. Let \((u^{(t)}_{\alpha})_{t\in\mathbb{N}}\) be SGD trajectories defined by (9) for \(r>0\) and \(\alpha\in]0,\alpha_{1}]\). One has_

\[\forall\varepsilon>0,\ \underset{t\longrightarrow+\infty}{\overline{\lim}}\ \nu\otimes\mathsf{x}^{\otimes\mathbb{N}}\otimes\mathsf{y}^{\otimes\mathbb{N}} \otimes\mathsf{e}^{\mathbb{N}}\otimes\eta^{\otimes\mathbb{N}}\left(d(u^{(t)}_{ \alpha},\mathcal{Z}_{r})>\varepsilon\right)\xrightarrow[\alpha\alpha\to 0]{ \alpha\in]0,\alpha_{1}]}0.\]

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 \(\mathcal{Z}_{r}\) in probability. Omitting the points of non-differentiability and the pathological case \(\|u\|_{2}=r\), the general idea is that \(u^{(\infty)}_{\alpha}\xrightarrow[\alpha\longrightarrow 0]{}\{u\ :\ \nabla F(u)=0\}\), which is the convergence that would be achieved by the gradient flow of \(F\), in the simpler case of \(\mathcal{C}^{2}\) smoothness.
## 5 Conclusion and Outlook

Under reasonable assumptions, we have shown that SGD trajectories of parameters of generative NNs with a SW loss converge towards the desired sub-gradient flow solutions, implying in a weak sense the convergence of said trajectories. Under stronger assumptions, we have shown that trajectories of a mildly modified SGD scheme converge towards a set of generalised critical points of the loss, which provides a missing convergence result for such optimisation problems.

The core limitation of this theoretical work is the assumption that the input data measure \(\varkappa\) is discrete (Assumption 6), which we required in order to prove that the loss \(F\) 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 to show that \(\int_{\mathcal{X}^{n}}\mathcal{E}(T(u,X),Y)\mathrm{d}\kappa^{\otimes n}(X)\) is path differentiable, knowing that \(u\longmapsto\mathcal{E}(T(u,X),Y)\) is path differentiable by composition (see the proof of Proposition 5 for the justification). Unfortunately, in general if each \(g(\cdot,x)\) is path differentiable, it is not always the case that \(\int g(\cdot,x)\mathrm{d}x\) is path differentiable (at the very least, there is no theorem stating this, even in the simpler case of tame functions, see (Bianchi et al., 2022), Section 6.1). However, there is such a theorem for _Clarke regular_ functions (specifically (Clarke, 1990), Theorem 2.7.2 with Remark 2.3.5), sadly the composition of Clarke regular functions is not always Clarke regular, it is only known to be the case in excessively restrictive cases (see (Clarke, 1990), Theorems 2.3.9 and 2.3.10). As a result, we leave the generalisation to a non-discrete input measure \(\varkappa\) for future work.

Another avenue for future study would be to tie the flow approximation result from Theorem 1 to Sliced Wasserstein Flows (Liutkus et al., 2019; Bonet et al., 2022). The difficulty in seeing the differential inclusion (5) as a flow of \(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 feedback, as well as Remi Flamary and Alain Durmus for fruitful discussions.
