Towards Understanding Gradient Dynamics of the Sliced-Wasserstein Distance via Critical Point Analysis

Thank you for your positive feedback and relevant questions.

• Concerning the limitations of the analysis in Section 5 : it is indeed true that most of the examples we discuss are in 2D, with the exception of our Proposition 5.1(b) which gives some examples of critical points in arbitrary dimension $d > 1$. We will make the introduction clearer about the limitations of this section.

• Regarding the limitations of our assumption of absolute continuity of the target measure $\rho$, we refer to the discussion in our answer to reviewer BmZL. We will make sure to discuss this more extensively in a revised version of our article. In particular, note that one of the advantages of the Sliced-Wasserstein distance is precisely that the Laguerre tesselations are easy to compute in 1D. Indeed, if we assume that $\rho$ is discretized using $M$ points on the real line, then the Laguerre cells are computed by sorting the $M$ points, and the $i$-th Laguerre cell consists of the $i$-th chunk of $\frac{M}{N}$ points in the sorted list.

• Thank you for bringing these additional references to our attention. In particular, the paper by Li and Moosmueller appears to be extremely relevant to our analysis as their assumption (A3) essentially means that the target $\mu$ is the only barycentric Lagrangian critical point in the compact set $K_{\sigma_0,\mu}$ in which the iterates of their scheme remain. This indeed ties our theoretical analysis of Lagrangian critical points to considerations of convergence of practical schemes optimizing $SW$ objectives. On the other hand, as we understand it, the paper by Altekrüger at al considers perturbations by velocity plans, which is a larger class of perturbations than in our Lagrangian framework (as a velocity plan allows to split atoms, while a perturbation by a vector field $\xi \in L^2(\mu,\mathbb{R}^d)$ cannot)

• We agree that it would be fruitful to establish clear links between the different notions of critical points and the different types of tangent spaces investigated by Gigli et al. We will include a discussion on this in a revised version of the article, maybe in an appendix, since for the sake of clarity we have avoided to make use of the theory of gradient flows developed by Ambrosio, Gigli, Savaré in the main body of the article.

• Whether the Wasserstein gradient flow starting at a discrete measure with an absolutely continuous target measure remains discrete is an excellent question. For $W_2$ objectives, it is known that the gradient flow diffuses (this is necessary as the flow converges to $\rho$ exponentially in the $W_2$ metric, see for instance Proposition 3.1 in [1]. For the $SW$ objective, although we have not formally shown that the flow of the $SW$ "diffuses" discrete measures, we expect it to do so. This is in fact why we chose to work with the "Lagrangian" framework and Lagrangian critical points. Indeed, what we study is the behavior of the system $dX_t^i/dt = - \nabla F(X_t^i)$ where $F$ is defined in Section 3, which corresponds to the continuous time limit of the gradient descent algorithm implemented in practice. Therefore, considering perturbations by vector fields $\xi$, which can't "split" atoms, allows for a theory better suited to analyze how algorithms relying on particle dynamics work.

[1] Huang, Y. J., & Malik, Z. (2024). Generative Modeling by Minimizing the Wasserstein-2 Loss. arXiv preprint arXiv:2406.13619.

• We have not investigated specifically what happens when the gradient descent is performed using a very small number of directions, such as $L = 1$. However, in our numerical experiments, we did observe that it is extremely easy to escape critical points such as those described in Proposition 5.1 or Proposition 5.2, as long as some stochasticity is introduced in the choice of directions.

Thank you for your relevant remarks and questions.

• Regarding the possible generalizations of our results to other types of optimized or target measures (notably discrete ones), we refer to our answer to reviewer BmZL which addresses these points extensively. Note in particular that even though we assumed throughout the article that the target measure $\rho$ is a density, this is mostly because we started our investigation from the study of the semi-discrete setting, and most of our results can actually be proven when $\rho$ is simply assumed to be without atoms. This covers many types of measures often encountered in machine learning, such as measures supported on lower dimensional manifolds of $\mathbb{R}^d$.

• We did make some attempts at a second order analysis of the $SW$ distance, but it turned out to be a significantly difficult problem. Most of what we know concerns cases such as the ones in Proposition 5.2, where we have explicit expressions of the quantile functions and we can compute explicit Taylor expansions of the $SW$. Even the discrete case is difficult : while formally differentiating under the integral sign does indeed give the expression of $\nabla F$ (Proposition 3.1), we cannot obtain the Hessian of $F$ by this method, as formally differentiating again $\nabla F$ under the integral sign gives $\frac{1}{Nd} I_d$, which cannot be the actual Hessian of $F$ as $F$ is not convex (it is semiconcave by Proposition A.2). Furthermore, attempts to numerically approximate the Hessian were inconclusive, as the computed Hessian would converge to $\frac{1}{Nd} I_d$ when we increased precision or added more directions (which is to be expected from the analysis of Appendix B, which shows that the approximation of $SW$ using a fixed number $L$ of directions $\theta_1,\ldots,\theta_L$ has Hessian $\frac{1}{NL} \sum_l \theta_l \theta_l^T$ everywhere where it is defined, and this converges to $\frac{1}{Nd} I_d$ when $L \to \infty$ and the $\theta_l$ are choosed randomly).

Thank you for your positive remarks and comments. Regarding the questions you raised :

• We meant by "descent lemma" a result that guarantees that, provided some conditions on the initial point and the step size are satisfied, one step of the gradient descent will decrease the loss. In particular, it gives the maximal step-size for which descent is maximal and still guaranteed - this size is related to the inverse of the smoothness constant of the functional. Above this limit, descent is not guaranteed for an explicit time-discretization, as corroborated by our experiments (Figure 2) where step-size below $2Nd$ yields descent while step-size above $2Nd$ diverges. Hence, while this is a theoretical result, it gives practical hints for optimization of the $SW$ distance in practice since below $2Nd$ (which is a known quantity) the descent is guaranteed. The term "descent lemma" is a standard term in optimization, cf for instance [1]. We will revise the article to clarify this point.

[1] Bauschke, H. H., Bolte, J., & Teboulle, M. (2017). A descent lemma beyond Lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research.

• Regarding the meaning of "a suitable alternating vector field $\xi$" at lines 369-371 (second column), we meant was that (using the notations of Proposition 5.2) by approximating the perturbation $\mu^t$ by $(Id+t\xi)\mu$ where $\xi$ rapidly alternates between $\vec{n}$ and $-\vec{n}$ on the segment $S$, we may hope that $SW^2_2((Id+t\xi)\mu,\rho)$ will also have a maximum at $t = 0$. For example, for $S = [-1,1] \times \{0\}$, we may consider $\xi$ such that $\xi(x,0) = \vec{e_2}$ for $x \in [i/n,(i+1)/n)$ with $i$ even, and $\xi(x,0) = -\vec{e_2}$ for $x \in [i/n, (i+1)/n)$ with $i$ odd, with $n$ large. In the experiments shown in Figure 1, we use such alternating perturbations. We will make our formulation clearer in a revised version of the article.

We thank you for the additional references, which we will accordingly cite in a revised version of our article.

Thank you for your positive remarks and relevant questions. Even though we restricted our theoretical analysis of the $SW$ distance to absolutely continuous target measures (and in Section 3 to discrete approximating candidates), this was mostly for the sake of simplicity, and it is indeed possible to extend many of our results to larger classes of measures, often (but not always) without requiring significant adaptation of the proofs. For example :

• Direct extensions :

  • In section 4, the notion of barycentric Lagrangian critical point for $SW^2_2(\cdot,\rho)$ can be defined for arbitrary measures $\mu, \rho \in P_2(R^d)$, as the 1D optimal transport plans $\gamma_\theta$ are always uniquely defined.

  • If we replace the assumption that $\rho$ is absolutely continuous by the assumption that it has no atoms, then Propositions 3.1, 3.2, 4.3, 4.6, Theorem 4.4 and Corollary 4.7 remain true.

• Minor extensions :

  • When $\rho$ is an uniform point cloud $\rho = \frac 1N \sum_{i=1}^N \delta_{Y_i}$ (with the same $N$ as in $\mu$), it should be possible to prove analogues of Propositions 3.1 and 3.2, replacing the barycenters $b_{\theta,i}$ by the reordered projections of the points $Y_i$ in the statements of the propositions (in fact, Theorem 1 in Bonneel et al. 2015 is the analogue of our Proposition 3.1 with $p=2$).

  • Proposition 3.3 remains true under the weaker assumption (A1) on $\rho$ that there exists $\beta > 0$ and $\Theta \subset \mathbb{S}^{d-1}$ with $\mathcal{H}^{d-1}(\Theta) > 0$ such that $\rho_\theta$ is a density bounded from above by $\beta$ for every $\theta \in \Theta$

  • The proof of Proposition 5.2 should be adaptable to the case where we only assume of $\rho$ that it has no atoms, and $\exists b > 0$ and a neighborhood $U \subset \mathbb{S}^1$ of $\vec{n}$ such that $\forall \theta \in U$, $\rho_\theta$ is a density bounded from above by $b$ - let's call this assumption (A2).

• Significant extensions :

  • The proofs of Propositions 3.1 and 3.2 should be adaptable to arbitrary $\rho$. The main difficulty would be that the Power cells $V_{\theta,i}$ are not well-defined. We would instead have to work with a decomposition $\rho_\theta = \sum_i \rho_{\theta,i}$ where $\rho_{\theta,i}$ is coupled to $\langle X_{\sigma_\theta(i)}, \theta \rangle$ by the optimal 1D transport plan (when $\rho_\theta$ has no atoms, we have $\rho_{\theta,i} = \rho_{\theta|V_{\theta,i}}$)

  • Theorem 4.4 also holds under the assumption that $\rho$ has no atoms and that the set of atoms of $\mu$ is closed, but proving this requires more sophisticated methods. The sketch of the proof is roughly the following : the beginning of the proof is the same as in the article. We cannot use Proposition 4.6(c) as $\mu$ may have atoms, but we can prove that Equation (152) still holds is $\xi$ is assumed to cancel at the atoms of $\mu$, and this allows to prove $v_\mu = 0$ $\mu$-ae on the complementary of the set of atoms of $\mu$. Then, by considering perturbations of each individual atom of $\mu$, we show that $v_\mu$ also cancels at the atoms. Since this assumption on $\mu$ seems unnatural and its proof is more complex, we chose to state the theorem in the article under the stronger but more natural assumption that $\mu$ is without atoms.

Moreover, regarding the limitations of our analysis :

• While our original assumption of absolute continuity of $\rho$ does limit the applicability of our results, the extensions we discussed to weaker assumptions on $\rho$ such as atomlessness, (A1) or (A2), allow us to cover a much wider range of target measures, including many types of singular measures which arise in machine learning, such as densities supported on a lower dimensional manifold (for example, $\rho = \frac 12 \mathcal{H}^1_{[-1,1]\times\{0\}} \in \mathcal{P}_2(\mathbb{R}^2)$ has no atoms, and satisfies (A1) and (A2) for $\vec{n} = \vec{e}_2$)

• The assumption of compact support of $\mu, \rho$, which we some of our results require, seem harder to relax. Indeed, it is needed in Proposition 4.6(c) to obtain regularity properties of the Kantorovich potentials, which we use to prove the differentiability, while the proof of Theorem 4.4 makes use of the fact that in a compact space, $SW_2$ and $W_2$ are equivalent distances metrizing the topology of weak convergence.

• Finally, the fact that our numerical experiments, in which the target measures were discretized, exhibit the behaviors of convergence and instability that our theoretical analysis highlighted, suggests that our results should still be relevant in the cases where the target measure $\rho$ is approximated by a discrete measure.

We will add in a revised version of our article a discussion of the generalizability of our results and of their limitations.