Two-layers neural networks for Schrödinger eigenvalue problems

Dus Mathias*†, Ehrlacher Virginie*†

August 21, 2024

Abstract

The aim of this article is to analyze numerical schemes using two-layer neural networks with infinite width for the resolution of high-dimensional Schrödinger eigenvalue problems with smooth interaction potentials and Neumann boundary condition on the unit cube in any dimension. More precisely, any eigenfunction associated to the lowest eigenvalue of the Schrödinger operator is a unit L² norm minimizer of the associated energy. Using Barron's representation of the solution with a probability measure defined on the set of parameter values and following the approach initially suggested by Bach and Chizat [1], the energy is minimized thanks to a constrained gradient curve dynamic on the 2-Wasserstein space of the set of parameter values defining the neural network. We prove the existence of solutions to this constrained gradient curve. Furthermore, we prove that, if it converges, the represented function is then an eigenfunction of the considered Schrödinger operator. At least up to our knowledge, this is the first work where this type of analysis is carried out to deal with the minimization of non-convex functionals.

AMS classification: 35P99, 65N75, 49Q22

# 1 Introduction

The main interest of this work is to introduce a rigorous framework for solving eigenvalue problems of a certain type of elliptic operators i.e. Schrödinger operators, using neural networks. A first step towards this objective came from ideas introduced in [1, 2] where authors propose a numerical method to solve a different but close problem:

$$ \begin{cases} -\Delta u^* = f \text{ on } \Omega, \\ \partial_n u^* = 0 \text{ on } \partial \Omega \end{cases} \qquad (1) $$

where f is a source term. Under some conditions on $f \in L^2(\Omega)$, it is possible to approximate numerically the minimizer $u^*$ as a convex combination of what we call, feature functions. In the same article, it is shown that minimizing a certain variational energy $E_\tau$, it is possible to solve Poisson-Neumann problem efficiently when the source term f belongs to the Barron's space $B^0(\Omega)$. To do so, authors propose to analyze the convergence of the gradient curve associated to $E_\tau$. More precisely, they prove that if the gradient curve converges then it is necessarily towards an optimum. Some numerical experiments shows the potential of the method on simple canonical examples.

For the Schrödinger eigenvalue problem, things are a bit different. It consists of finding one/the ground eigencouple of the elliptic operator $-\Delta + W$ where W is a potential, supplemented with boundary condition that we do not specify here. In [3], authors adapted the work [4] to give conditions on the potential W such that the ground eigenstate belongs to the Barron space of order $s > 0$ denoted $B^s$. More specifically, they prove that if the potential W belongs to $B^s$ and if the ground eigenvalue is negative, then the eigenvector $u^*$ belongs to $B^s$ and a two-layers neural approximation is relevant when $s \ge 2$.

The ultimate goal of studying the Schrödinger problem is to get the ground state of the multi-body Schrödinger operator given by:

<footnote>*École des ponts Paritech, Marne-la-Vallée, France</footnote>

<footnote>†INRIA Paris, France, MATHERIALS team-project, emails: mathias.dus@enpc.fr, virginie.ehrlacher@enpc.fr</footnote>
$$H := -\frac{1}{2} \sum_{n=1}^{N} \Delta_{x_n} + \sum_{n=1}^{N} \sum_{m=1}^{M} \frac{-Z_m}{|x_n - X_m|} + \sum_{1 \le j < n \le N} \frac{1}{|x_j - x_n|} \quad (2)$$

where:

* $(x_n)_{1 \le n \le N}$ is the position of $N$ electrons,

* $(X_m)_{1 \le m \le M}$ is the position of $M$, nucleis of charge $(Z_m)_{1 \le m \le M}$

* The domain of $H$ is $H^1(\mathbb{R}^{dN})$.

The difficulties posed by the Coulombic singularities, the unbounded domain, and the antisymmetry condition (due to spins) on the domain of definition of $H$ are not addressed in this paper. However, it is hoped that this work will serve as a theoretical starting point for the solution of these issues. Furthermore, the so-called Jastrow factor ansatz demonstrates that the ground eigenstate exhibits a high degree of regularity in the $W^{2,\infty}$ sense [5]. Additionally, the mixed derivative regularity described in [6] provides further evidence for a neural approximation in the manner described by [7]. It must be acknowledged that the Coulomb potential is not of Barron regularity. However, experimental results presented in references [8, 9] indicate that for at least simple systems, the eigenstate can be represented by complex neural networks for which no mathematical framework is clearly defined.

In this work, we will follow and adapt the strategy presented in [1, 2]. First, we introduce the Schrödinger energy, which we restrict to functions written as an expectation of some simple functions. This energy, rewritten as a function over probability measures, is then minimized using a gradient curve with respect to the 2-Wasserstein metric. A result of convergence of the curve towards a measure representing an eigenfunction is presented, and numerical experiments demonstrate the efficiency of the method, which surprisingly converges towards a ground state.

The article is decomposed as follow. In Section 1, the problem is presented in a rigorous way with a rigorous definition of the eigenvalue problem and energy for which we introduce the penalized gradient flow. In Section 3, the existence of such flow is proved and in Section 4, the convergence towards a minimizer is shown under the hypothesis of convergence of the trajectory. Finally in Section 5, some simulations illustrate the potential of the method.

**Notation:** The notation $C$ is used for universal constants. If there exists a dependence on some parameter $\iota$, the constant is denoted $C(\iota)$.

# 2 Preliminaries and notation

The aim of this section is to present the eigenvalue problem considered in this work and fix some notation.

## 2.1 Problem setting

Let $\Omega := [0, 1]^d$ be the $d$-dimensional hypercube and $\partial\Omega$ denote the boundary of $\Omega$. We first introduce the quadratic functional $\mathcal{E}: H^1(\Omega) \to \mathbb{R}$ defined as follows:

$$\forall u \in H^1(\Omega), \quad \mathcal{E}(u) := \int_{\Omega} |\nabla u|^2 + W u^2 dx$$

where $W \in L^\infty(\Omega; \mathbb{R})$. The functional $\mathcal{E}$ is then the quadratic form associated to the Schrödinger operator $-\Delta + W$ with Neumann boundary condition, which is a self-adjoint operator with compact resolvent on $L^2(\Omega)$ with domain $H^2(\Omega)$. We are interested in the computation of a solution $u^* \in H^1(\Omega)$ of the following minimization problem

$$u^* \in \underset{u \in H^1(\Omega), \|u\|_{L^2(\Omega)}=1}{\operatorname{argmin}} \mathcal{E}(u), \quad (3)$$
solution of the Schrödinger eigenvalue problem with Neumann boundary condition: find $(u^*, \lambda) \in H^1(\Omega) \times \mathbb{R}$ solution to:

$$ \begin{cases} -\Delta u^* + W u^* = \lambda u^* & \text{on } \Omega, \\ \partial_n u^* = 0 & \text{on } \partial\Omega. \end{cases} \qquad (4) $$

Moreover, since $u^*$ is a minimizer of (3), it holds that $\lambda$ is the lowest eigenvalue of the operator $-\Delta + W$. For convenience, we use the following notation for the constraint:

$$ \forall u \in L^2(\Omega), \quad C(u) := \int_{\Omega} |u|^2 dx. $$

The aim of the present paper is to propose and analyze a neural-network based numerical method for the resolution of (3) in the spirit of [2, 1], based on the use of infinite-width two-layer neural networks.

## 2.2 Activation function

We introduce here the particular choice of activation function we consider in this work, which is the same as the one used in [2]. Let $\sigma: \mathbb{R} \to \mathbb{R}$ be the classical Rectified Linear Unit (ReLU) function so that for all $y \in \mathbb{R}$, $\sigma(y) := \max(y, 0)$. Let $\rho: \mathbb{R} \to \mathbb{R}$ be defined by

$$ \left\{ \begin{array}{ll} Z \exp\left(-\frac{\tan(\frac{\pi}{2}y)^2}{2}\right) & \text{if } |y| \le 1 \\ 0 & \text{otherwise,} \end{array} \right. \qquad (5) $$

where the constant $Z \in \mathbb{R}$ is defined such that the integral of $\rho$ is equal to one. For all $\tau > 0$, we then define $\rho_{\tau} := \tau \rho(\tau)$ and $\sigma_{\tau}: \mathbb{R} \to \mathbb{R}$ the regularized ReLU function defined by

$$ \forall y \in \mathbb{R}, \quad \sigma_{\tau}(y) := (\rho_{\tau} \star \sigma)(y). \qquad (6) $$

In this work, we will rather use a hat version of the regularized ReLU activation function. More precisely, we define:

$$ \forall y \in \mathbb{R}, \quad \sigma_{H,\tau}(y) := \sigma_{\tau}(y+1) - \sigma_{\tau}(2y) + \sigma_{\tau}(y-1), \qquad (7) $$

which is called hereafter the regularized HReLU (Hat ReLU) activation. When $\tau = +\infty$, the following notation is proposed:

$$ \forall y \in \mathbb{R}, \quad \sigma_H(y) := \sigma(y+1) - \sigma(2y) + \sigma(y-1). \qquad (8) $$

Note that a direct corollary of [2, Lemma 2] is that there exists a constant $C > 0$ such that for all $\tau > 0$,

$$ \|\sigma_H - \sigma_{H,\tau}\|_{H^1(\mathbb{R})} \le \frac{C}{\sqrt{\tau}} \qquad (9) $$

We will also use the fact that there exists a constant $C > 0$ such that for all $\tau > 0$,

$$ \|\sigma_{H,\tau}\|_{L^{\infty}(\mathbb{R})} \le C, \quad \|\sigma'_{H,\tau}\|_{L^{\infty}(\mathbb{R})} \le C, \quad \|\sigma''_{H,\tau}\|_{L^{\infty}(\mathbb{R})} \le C\tau \quad \text{and} \quad \|\sigma'''_{H,\tau}\|_{L^{\infty}(\mathbb{R})} \le C\tau^2. \qquad (10) $$

## 2.3 Infinite width two-layer neural networks

Let us introduce the set of parameter values of the neural network

$$ \Theta := \mathbb{R} \times S^{d-1} \times \mathbb{R}, $$

where $S^{d-1}$ denotes the unit sphere embedded in $\mathbb{R}^d$. Moreover, let $\mathcal{P}_2(\Theta)$ be the set of probability measures on $\Theta$ with finite second-order moments. The space $\mathcal{P}_2(\Theta)$ is equipped with the 2-Wasserstein distance:

$$ \forall \mu, \nu \in \mathcal{P}_2(\Theta), \quad W_2^2(\mu, \nu) := \inf_{\gamma \in \Gamma(\mu, \nu)} \int_{\Theta^2} d(\theta, \tilde{\theta})^2 d\gamma(\theta, \tilde{\theta}), $$
where $\Gamma(\mu, \nu)$ is the set of probability measures on $\Theta^2$ with marginals given respectively by $\mu$ and $\nu$ and where $d$ is the geodesic distance in $\Theta$. The geodesic distance between $\theta := (a, w, b)$, $\tilde{\theta} := (\tilde{a}, \tilde{w}, \tilde{b}) \in \Theta$ is defined as:

$$d(\theta, \tilde{\theta}) = \sqrt{(a - \tilde{a})^2 + d_{S^{d-1}}(w, \tilde{w})^2 + (b - \tilde{b})^2},$$

with $d_{S^{d-1}}$ the geodesic distance on the unit sphere $S^{d-1}$. We also denote by $\Gamma_o(\mu, \nu)$ the subset of probability measures $\gamma \in \Gamma(\mu, \nu)$ such that $W_2(\mu, \nu) = \sqrt{\int_{\Theta^2} d(\theta, \tilde{\theta})^2 d\gamma(\theta, \tilde{\theta})}$. If $\gamma \notin \Gamma_o(\mu, \nu)$, we use the notation $W_{2,\gamma}(\mu, \nu) := \sqrt{\int_{\Theta^2} d(\theta, \tilde{\theta})^2 d\gamma(\theta, \tilde{\theta})}$.

For any $\tau > 0$, we introduce the function $\Phi_{\tau} : \Theta \times \Omega \to \mathbb{R}$ defined such that

$$\forall \theta := (a, w, b) \in \Theta, \forall x \in \Omega, \Phi_{\tau}(\theta; x) := a\sigma_{H,\tau}(w \cdot x + b), \quad (11)$$

and the operator $P_{\tau} : \mathcal{P}_2(\Theta) \to H^1(\Omega)$ defined by:

$$\forall \mu \in \mathcal{P}_2(\Theta), P_{\tau}\mu := \int_{\Theta} \Phi_{\tau}(\theta; \cdot) d\mu(\theta).$$

The aim of this paper is to analyze a numerical method which consists in approximating a solution $u^* \in H^1(\Omega)$ to (3) by a function of the form $P_{\tau}\mu^*$ where $\mu^*$ is a solution of a constrained minimization problem of the form

$$\inf_{\substack{\mu \in \mathcal{P}_2(\Theta) \\ \mathcal{C}_{\tau}(\mu)=0}} \mathcal{E}_{\tau}(\mu). \quad (12)$$

where:

* the energy writes:

$$\mathcal{E}_{\tau} : \left\{ \begin{array}{l} \mathcal{P}_2(\Theta) \mapsto \mathbb{R} \\ \mu \mapsto \mathcal{E}(P_{\tau}\mu). \end{array} \right. \quad (13)$$

* the constraint cost is:

$$\mathcal{C}_{\tau} : \left\{ \begin{array}{l} \mathcal{P}_2(\Theta) \mapsto \mathbb{R} \\ \mu \mapsto \|P_{\tau}\mu\|_{L^2(\Omega)} - 1. \end{array} \right.$$

The associated set of constraint is denoted by $X_{\tau} := \{\mu \in \mathcal{P}_2(\Theta) \mid \mathcal{C}_{\tau}(\mu) = 0\}$.

More precisely, we prove the existence of a constrained gradient curve $(\mu_t)_{t \ge 0}$ associated to (a regularized version of) the constrained minimization problem (12). Then, under the assumption that such a curve converges as $t$ goes to infinity to some measure $\mu_{\infty} \in \mathcal{P}_2(\Theta)$, we prove that $P_{\tau}(\mu_{\infty})$ is an eigenfunction of the Schrödinger operator $-\Delta + W$. We would like to stress on the fact that this eigenfunction may however be associated to an eigenvalue $\lambda_{\infty}$ which may not be the smallest eigenvalue of the operator $-\Delta + W$ in general.

## 2.4 Barron spaces and regularity of eigenfunctions of Schrödinger operators

In this section, we introduce the Barron space to justify our choice of neural network approximation of ground eigenstates. To do so, let us introduce the orthonormal basis in $L^2(\Omega)$ composed of the eigenfunctions $\{\phi_k\}_{k \in \mathbb{N}^d}$ of the Laplacian operator with Neumann boundary condition, where

$$\forall k = (k_1, \dots, k_d) \in \mathbb{N}^d, \forall x := (x_1, \dots, x_d) \in \Omega, \quad \phi_k(x_1, \dots, x_d) := \prod_{i=1}^d \cos(\pi k_i x_i). \quad (14)$$

Notice that $\{\phi_k\}_{k \in \mathbb{N}^d}$ is also an orthogonal basis of $H^1(\Omega)$. Using this basis, we have the Fourier representation formula for any function $u \in L^2(\Omega)$:

$$u = \sum_{k \in \mathbb{N}^d} \hat{u}(k)\phi_k,$$
where for all $k \in \mathbb{N}^d$, $\hat{u}(k) := \langle \phi_k, u \rangle_{L^2(\Omega)}$. This allows to define the (spectral) Barron space [4] as follows:

**Definition 1.** For all $s > 0$, the Barron space $\mathcal{B}^s(\Omega)$ is defined as:

$$ \mathcal{B}^s(\Omega) := \left\{ u \in L^1(\Omega) : \sum_{k \in \mathbb{N}^d} (1 + \pi^s |k|_1^s) |\hat{u}(k)| < +\infty \right\} \quad (15) $$

and the space $\mathcal{B}^2(\Omega)$ is denoted $\mathcal{B}(\Omega)$. Moreover, the space $\mathcal{B}^s(\Omega)$ is embedded with the norm:

$$ \|u\|_{\mathcal{B}^s(\Omega)} := \sum_{k \in \mathbb{N}^d} (1 + \pi^s |k|_1^s) |\hat{u}(k)|. \quad (16) $$

One of the reason we consider problem (12) is that under some regularity hypothesis on the potential, the ground eigenvector is unique and belongs to some Barron space. More precisely, using Krein-Rutman theory, the authors of [3] proved the following result.

**Theorem 1 ([3]).** Let $s > 0$. If $W \in L^\infty(\Omega) \cup \mathcal{B}^s(\Omega)$ is such that $(-\Delta + W, H^1(\Omega))$ has a spectral gap, then there exists a unique ground eigenvector $u^*$ which belongs to $\mathcal{B}^s(\Omega)$.

By the theorem of approximation of Barron function, it is then possible to approximate the ground state $u^*$ by a two-layer neural network. To present such result, we introduce the concept of feature space $\mathcal{F}_{\chi,m}(B)$ defined as:

$$ \mathcal{F}_{\chi,m}(B) := \left\{ \sum_{i=1}^{m} a_i \chi(w_i \cdot x + b_i) : a_i, b_i \in \mathbb{R}, w_i \in \mathbb{R}^d, |w_i| = 1, |b_i| \le 1, \sum_{i=1}^{m} |a_i| \le 4B \right\} \quad (17) $$

where $\chi: \mathbb{R} \to \mathbb{R}$ be measurable, $m \in \mathbb{N}^*$ and $B > 0$.

**Theorem 2 ([4]).** For any $u \in \mathcal{B}(\Omega)$, $m \in \mathbb{N}^*$:

(i) there exists $u_m \in \mathcal{F}_{\sigma_{H,m}}(\|u\|_{\mathcal{B}(\Omega)})$ such that:

$$ \|u - u_m\|_{H^1(\Omega)} \le \frac{C \|u\|_{\mathcal{B}(\Omega)}}{\sqrt{m}}, $$

(ii) there exists $\tilde{u}_m \in \mathcal{F}_{\sigma_{H,m},m}(\|u\|_{\mathcal{B}(\Omega)})$ such that:

$$ \|u - \tilde{u}_m\|_{H^1(\Omega)} \le \frac{C \|u\|_{\mathcal{B}(\Omega)}}{\sqrt{m}}. \quad (18) $$

where for both items, $C$ is a universal constant which does not depend on $d$ neither on $u$.

The problem with Theorem 1 is that we do not have a clear estimate of the Barron norm of $u^*$ wrt the Barron norm of $W$. Consequently, we are not sure that the quantity $\|u^*\|_{\mathcal{B}(\Omega)}$ is not exploding wrt dimension. This constitutes an important question that we do not treat in this paper but remains fundamental to legitimate our choice of numerical approximation.

# 3 Existence of a gradient curve

In the following, we denote by $T\Theta$ the tangent bundle of $\Theta$, i.e.

$$ T\Theta := \bigcup_{\theta \in \Theta} \{\theta\} \times T_{\theta}\Theta, $$

where $T_{\theta}\Theta$ is the tangent space to $\Theta$ at the element $\theta \in \Theta$. It is easy to check that for all $\theta := (a, w, b) \in \Theta$, it holds that $T_{\theta}\Theta = \mathbb{R} \times \text{Span}\{w\}^{\perp} \times \mathbb{R}$, where $\text{Span}\{w\}^{\perp}$ is the subspace of $\mathbb{R}^d$ containing all $d$-dimensional vectors orthogonal to $w$.

For any $\theta := (a, w, b) \in \Theta$ and any function $f: \Theta \to \mathbb{R}$ differentiable at $\theta$, we denote by $\nabla f(\theta) \in \mathbb{R}^{d+2}$ the (euclidean) gradient of $f$. We also denote by $\nabla_{\Theta}f(\theta)$ the orthogonal projection of $\nabla f(\theta)$
onto $T_\theta\Theta$ and by $\nabla_{\Theta^\perp} f(\theta)$ the orthogonal projection of $\nabla f(\theta)$ onto $(T_\theta\Theta)^\perp = \{0\} \times \text{Span}\{w\} \times \{0\} = \text{Span}\{(0, w, 0)\}$. Therefore, we have the unique orthogonal decomposition of the gradient:

$$ \nabla f(\theta) = \nabla_{\Theta} f(\theta) + \nabla_{\Theta^{\perp}} f(\theta) \qquad (19) $$

For $r > 0$, let us introduce $K_r := [-r, r] \times S^{d-1} \times [-r, r]$. Restraining the domain of $C_\tau$, it is relatively easy to get continuity and even differentiability of this functional.

**Lemma 1.** The application $C_\tau : \mathcal{P}_2(K_r) \to \mathbb{R}$ is continuous (with respect to the $W_2$ metric). Moreover, there exists a constant $C(\tau, r) > 0$ such that for all $(\mu, \nu) \in \mathcal{P}_2(K_r)^2$ such that $P_\tau\mu \neq 0, W_2(\mu, \nu) \le C(\tau, \mu)$ and $\gamma \in \Gamma_o(\mu, \nu)$,

$$ C_{\tau}(\nu) = C_{\tau}(\mu) + \int_{\Theta^2} \nabla_{\Theta} C_{\tau, \mu}(\theta) \cdot (\tilde{\theta} - \theta) d\gamma(\theta, \tilde{\theta}) + o(W_2(\mu, \nu)) $$

where for all $\theta \in \Theta$, $C_{\tau,\mu}(\theta) := \frac{\langle P_\tau\mu, \Phi_\tau(\theta) \rangle_{L^2(\Omega)}}{\|P_\tau\mu\|_{L^2(\Omega)}}$ is the smooth constraint potential.

*Proof.* Before developing the proof, we will need the following notation:

$$ M_{\tau}(\theta, \tilde{\theta}; x) := \int_{0}^{1} (\tilde{\theta} - \theta)^{T} H_{\theta} \Phi_{\tau}(\theta + t(\tilde{\theta} - \theta); x) (\tilde{\theta} - \theta) (1 - t) dt. $$

Now, for $\mu, \nu \in \mathcal{P}_2(K_r)$ and $\gamma \in \Gamma(\mu, \nu)$ supported in $K_r^2$, it holds that:

$$ \begin{aligned} C_{\tau}(\nu) &= C \left( \int_{K_r} \Phi_{\tau}(\tilde{\theta}; \cdot) d\nu(\tilde{\theta}) \right) \\ &= C \left( \int_{K_r^2} \Phi_{\tau}(\tilde{\theta}; \cdot) d\gamma(\theta, \tilde{\theta}) \right) \\ &= C \left( \int_{K_r^2} \left[ \Phi_{\tau}(\theta; \cdot) + \nabla_{\theta} \Phi_{\tau}(\theta; \cdot) \cdot (\tilde{\theta} - \theta) + M_{\tau}(\theta, \tilde{\theta}; \cdot) \right] d\gamma(\theta, \tilde{\theta}) \right) \\ &= C_{\tau}(\mu) + dC|_{P_{\tau}(\mu)} \left( \int_{K_r^2} \left[ \nabla_{\theta} \Phi_{\tau}(\theta; \cdot) \cdot (\tilde{\theta} - \theta) + M_{\tau}(\theta, \tilde{\theta}; \cdot) \right] d\gamma(\theta, \tilde{\theta}) \right) \\ &\quad + \int_0^1 (1-t)^2 dC \int_{K_r^2} \left[ \Phi_{\tau}(\theta; \cdot) + t(\nabla_{\theta} \Phi_{\tau}(\theta; \cdot) \cdot (\tilde{\theta} - \theta) + M_{\tau}(\theta, \tilde{\theta}; \cdot)) \right] d\gamma(\theta, \tilde{\theta}) (du, du) dt \end{aligned} $$

where the notation $du := \int_{K_r^2} [\nabla_{\theta}\Phi_{\tau}(\theta; \cdot) \cdot (\tilde{\theta} - \theta) + M_{\tau}(\theta, \tilde{\theta}; \cdot)] d\gamma(\theta, \tilde{\theta})$ is introduced to ease the reading. We have:

$$ \bullet \ dC_v(dv) = \frac{\langle v, dv \rangle_{L^2(\Omega)}}{\|v\|_{L^2(\Omega)}} \quad \text{and consequently} $$

$$ \bullet \ d^2 C_v(dv, dv) = \frac{1}{2} \frac{\|dv\|_{L^2(\Omega)}^2}{\|v\|_{L^2(\Omega)}} - \frac{1}{4} \frac{\langle v, dv \rangle_{L^2(\Omega)}^2}{\|v\|_{L^2(\Omega)}} \quad \text{and consequently} $$

Using estimates taken from [2, Proof of corollary 1] i.e

$$ \|du\|_{L^2(\Omega)}^2 \le C(1 + r^2 + \tau^2 + r^2 \tau^2)W_2^2(\mu, \nu), \qquad (22) $$

one gets:
$$ \|P_{\tau}\mu + tdu\|_{L^2(\Omega)} \geq \|P_{\tau}\mu\|_{L^2(\Omega)} - C\sqrt{1 + r^2 + \tau^2 + r^2\tau^2}W_2(\mu, \nu) $$

so there exists $C(\tau, r)$ such that if $W_2(\mu, \nu) \leq C(\tau, r)$ then $\|P_{\tau}\mu + tdu\|_{L^2(\Omega)} > 0$ for all $0 \leq t \leq 1$. Using this and (20)-(22) one gets:

$$ \mathcal{C}_{\tau}(\nu) = \mathcal{C}_{\tau}(\mu) + \int_{\Theta^2} d\mathcal{C}|_{P_{\tau}(\mu)}(\nabla\Phi_{\tau}(\theta) \cdot (\tilde{\theta} - \theta))d\gamma(\theta, \tilde{\theta}) + o(W_2(\mu, \nu)) $$

which is the desired result.

□

We have an equivalent result for the energy $\mathcal{E}_{\tau}$.

**Lemma 2.** The application $\mathcal{E}_{\tau}: \mathcal{P}_2(K_r) \to \mathbb{R}$ is continuous (with respect to the $W_2$ metric). Moreover, for all $(\mu, \nu) \in \mathcal{P}_2(K_r)^2$ such that for all $\gamma \in \Gamma_o(\mu, \nu)$,

$$ \mathcal{E}_{\tau}(\nu) = \mathcal{E}_{\tau}(\mu) + \int_{\Theta^2} \nabla\Theta V_{\tau,\mu}(\theta) \cdot (\tilde{\theta} - \theta) d\gamma(\theta, \tilde{\theta}) + o(W_2(\mu, \nu)) $$

where for all $\theta \in \Theta$, $V_{\tau,\mu}(\theta) := \langle \nabla P_{\tau}\mu, \nabla\Phi_{\tau}(\theta; \cdot) \rangle_{L^2(\Omega)} + \langle W P_{\tau}\mu, \Phi_{\tau}(\theta; \cdot) \rangle_{L^2(\Omega)}$ is the smooth constraint potential.

*Proof.* It suffices to adapt the proof [2, Proposition 3].

□

The penalized energy $\mathcal{E}_{\tau,r}^C : \mathcal{P}_2(\Theta) \mapsto \mathbb{R} \cup \{+\infty\}$ which is defined by:

$$ \forall \mu \in \mathcal{P}_2(\Theta), \quad \mathcal{E}_{\tau,r}^C(\mu) := \begin{cases} \mathcal{E}_{\tau}(\mu) & \text{if } \mathcal{C}_{\tau}(\mu) = 0 \text{ and } \text{Supp}(\mu) \subset K_r, \\ +\infty & \text{otherwise.} \end{cases} $$

where $\text{Supp}(\mu)$ denotes the support of the probability measure $\mu$. Let us also introduce the notation:

$$ X_{\tau} := \{\mu \in \mathcal{P}_2(\Theta) \mid \mathcal{C}_{\tau}(\mu) = 0\} \quad \text{and} \quad X_{\tau,r} := \{\mu \in \mathcal{P}_2(\Theta) \mid \mathcal{C}_{\tau}(\mu) = 0 \text{ and } \text{Supp}(\mu) \subset K_r\}. $$

This allows to prove smoothness on the constrained energy.

**Proposition 1.** The constrained energy $\mathcal{E}_{\tau,r}^C$ is lower semicontinuous and coercive.

*Proof.* Coercivity is obvious. For the lower semicontinuity, let $(\mu_n)_{n \geq 0}$ be a converging sequence in $(\mathcal{P}_2(\Theta), W_2)$ and let $\mu$ be its limit. There are two possible situations. Either there exists an integer $N \in \mathbb{N}$ such that for all $n \geq N$, $\mu_n \notin X_{\tau,r}$. In this case, it then holds that:

$$ \mathcal{E}_{\tau,r}^C(\mu) \leq \liminf_{n \to \infty} \mathcal{E}_{\tau,r}^C(\mu_n) = +\infty. $$

Otherwise, for any $N \in \mathbb{N}$, there exists $n \geq N$ such that $\mu_n \in X_{\tau,r}$. As a consequence, there exists an extracted subsequence of $(\mu_n)_{n \geq 0}$, still denoted by $(\mu_n)_{n \geq 0}$ for the sake of convenience, the elements of which all belong to $X_{\tau,r}$ and which still converges to $\mu$. As $\mathcal{C}_{\tau}$ is continuous from Lemma 1, $X_{\tau} = \mathcal{C}_{\tau}^{-1}(\{0\})$ is closed and $\mu$ belongs to $X_{\tau}$. Moreover, by Portmanteau theorem, $\text{Supp}(\mu) \subset K_r$ which yields that $\mu \in X_{\tau,r}$. In addition, as the energy $\mathcal{E}_{\tau}$ is lower semi-continuous from Lemma 2,

$$ \mathcal{E}_{\tau,r}^C(\mu) = \mathcal{E}_{\tau}(\mu) \leq \liminf_{n \to \infty} \mathcal{E}_{\tau}(\mu_n) = \liminf_{n \to \infty} \mathcal{E}_{\tau,r}^C(\mu_n). $$

This proves the lower semicontinuity of the constrained energy.

□

We need the following preliminary Lemma of non degeneracy of the constraint to prove the existence of a Lagrange multiplier. It involves the definition of the local slope:

**Definition 2.** Let $E: \mathcal{P}_2(\Theta) \to \mathbb{R}$ be an energy functional. The local slope at a measure $\mu \in \mathcal{P}_2(\Theta)$ is the negative quantity (that can be infinite):

$$ |\partial E|(\mu) := \liminf_{\nu \in \mathcal{P}_2(\Theta), W_2(\mu, \nu) \to 0} \frac{|(E(\nu) - E(\mu))^-|}{W_2(\nu, \mu)} $$

where the superscript is the negative part.
**Lemma 3.** For all $\mu \in X_\tau$, $|\partial C_\tau|^2(\mu) = \|\nabla_\Theta C_{\tau,\mu}\|_{L^2(\Theta;\mu)}^2 \ge \frac{1}{\int_\Theta a^2 d\mu} > 0$.

*Proof.* Let $\mu \in X_\tau$. For $t \in [0, 1]$, the operator $T_t : (a, w, b) \ni \Theta \mapsto ((1+t)a, w, b)$ allows to define the curve $\mu_t := T_t\#\mu$ for which it is simple to prove that:

* $C_\tau(\mu_t) = 1 + t$,

* $W_2(\mu, \mu_t) \le t\sqrt{\int_\Theta a^2 d\mu}$ where $\int_\Theta a^2 d\mu > 0$ since $\mu \in X_\tau$.

Consequently, it holds:

$$|\partial C_\tau|(\mu) \ge \frac{1}{\sqrt{\int_\Theta a^2 d\mu}}$$

and the conclusion immediately follows.

<figure><img src="image_1.png" /></figure>

The following proposition is the main result of the paper where we identify the local slope of the constrained energy.

**Proposition 2.** If $\text{Supp}(\mu) \subset K_{r-\delta}$ for some $\delta > 0$ and $\mu \in X_\tau$, then the local slope of the constrained energy is given by:

$$|\partial \mathcal{E}_{\tau,r}^C|(\mu) = \|\nabla_\Theta V_{\tau,\mu} - \sigma_\mu \nabla_\Theta C_{\tau,\mu}\|_{L^2(\Theta;\mu)}$$

where:

$$\sigma_\mu := \frac{\langle \nabla_\Theta V_{\tau,\mu}, \nabla_\Theta C_{\tau,\mu} \rangle_{L^2(\Theta;\mu)}}{\|\nabla_\Theta C_{\tau,\mu}\|_{L^2(\Theta;\mu)}^2}.$$

is well-defined by Lemma 3.

Before going into the proof of Proposition 2, we need the following notations and proposition. Let $\mathfrak{P}$ be the set of geodesics of $\Theta$, i.e. the set of absolutely continuous curves $\pi : [0, 1] \to \Theta$ such that for all $t_1, t_2 \in [0, 1]$, $d(\pi(t_1), \pi(t_2)) = d(\pi(0), \pi(1))|t_1 - t_2|$. It then holds that for all $0 \le t \le 1$, we have $|\dot{\pi}(t)| = \lim_{\epsilon \to 0} \frac{d(\pi(t), \pi(t+\epsilon))}{|\epsilon|} = d(\pi(0), \pi(1))$. For all $s \in [0, 1]$, we define the application map $e_s : \mathfrak{P} \to \Theta$ such that $e_s(\pi) := \pi(s)$. Owing to this, McCann interpolation gives the fundamental characterization of constant speed geodesics in $\mathcal{P}_2(\Theta)$:

**Proposition 3.** [10, Proposition 2.10] For all $\mu, \nu \in \mathcal{P}_2(\Theta)$ and any geodesic $\kappa : [0, 1] \to \mathcal{P}_2(\Theta)$ between them (i.e. such that $\kappa(0) = \mu$ and $\kappa(1) = \nu$) with respect to the $W_2$ metric, there exists $\Pi \in \mathcal{P}_2(\mathfrak{P})$ such that:

$$\forall t \in [0, 1], \kappa(t) = e_t \# \Pi.$$

Moreover, we introduce the operator $G$:

$$G := \begin{cases} \mathfrak{P} & \to T\Theta \\ \pi & \mapsto (\pi(0), \dot{\pi}(0)) \end{cases}$$

which will be useful later.

*Proof.* Let $\mu, \nu \in X_\tau$ and $\gamma \in \Gamma_o(\mu, \nu)$. Then by Lemma 1:

$$\int_{\Theta^2} \nabla C_{\tau,\mu}(\theta) \cdot (\tilde{\theta} - \theta) d\gamma(\theta, \tilde{\theta}) = o(W_2(\mu, \nu)).$$

Consequently,

$$\mathcal{E}_{\tau,r}^C(\mu) - \mathcal{E}_{\tau,r}^C(\nu) = \int_{\Theta^2} \nabla(V_{\tau,\mu} - \sigma_\mu C_{\tau,\mu}) \cdot (\tilde{\theta} - \theta) d\gamma(\mu, \nu) + o(W_2(\mu, \nu)). \quad (23)$$

The first term in the right hand side of (23) can be decomposed as follows:
$$ \int_{\Theta^2} \nabla(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}) \cdot (\tilde{\theta} - \theta)d\gamma(\mu, \nu) = \int_{\Theta^2} \nabla_{\Theta}(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}) \cdot (\tilde{\theta} - \theta)d\gamma(\theta, \tilde{\theta}) \\ + \int_{\Theta^2} \nabla_{\Theta^{\perp}}(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}) \cdot (\tilde{\theta} - \theta)d\gamma(\theta, \tilde{\theta}). \quad (24) $$

The first term in (24) can be easily bounded:

$$ \left| \int_{\Theta^2} \nabla_{\Theta}(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}) \cdot (\tilde{\theta} - \theta) d\gamma(\theta, \tilde{\theta}) \right| \le \|\nabla_{\Theta}V_{\tau,\mu} - \sigma_{\mu}\nabla_{\Theta}C_{\tau,\mu}\|_{L^2(\Theta;\mu)} \sqrt{\int_{\Theta^2} |\tilde{\theta} - \theta|^2 d\gamma(\theta, \tilde{\theta})} \\ \le \|\nabla_{\Theta}V_{\tau,\mu} - \sigma_{\mu}\nabla_{\Theta}C_{\tau,\mu}\|_{L^2(\Theta;\mu)} W_2(\mu, \nu) \quad (25) $$

where we used the fact that the euclidean distance is bounded by the geodesic distance on the submanifold $\Theta$ to get the last inequality. The second term in (24) is a bit more difficult to bound and to do so, we use McCann's representation of optimal plans. There exists $\Pi \in \mathcal{P}_2(\mathfrak{P})$ such that $\gamma = (e_0, e_1)\#\Pi$ and:

$$ \left| \int_{\Theta^2} \nabla_{\Theta^{\perp}}(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}) \cdot (\tilde{\theta} - \theta) d\gamma(\theta, \tilde{\theta}) \right| = \left| \int_{T\Theta} \nabla_{\Theta^{\perp}}(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}) \cdot (\exp_{\theta}(v) - \theta) dG_{\#}(\theta, v) \right| \\ \le \int_{T\Theta} |\nabla_{\Theta^{\perp}}(V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu})| |v|^2 dG_{\#}(\theta, v) \\ \le C_r \int_{T\Theta} |v|^2 dG_{\#}(\theta, v) \\ = C_r W_2^2(\mu, \nu), \quad (26) $$

where we used the fact that:

* On the sphere $\exp_z(v) = \cos(|v|)z + \sin(|v|)\frac{v}{|v|}$, so that

$$ |\mathbf{P}_{\theta}^{\perp}(\exp_{\theta}(v) - \theta)| = |\cos(|v|) - 1| \le |v|^2 $$

to get the second inequality.

* The support of $\mu$ is included in $K_r$ so $\|V_{\tau,\mu} - \sigma_{\mu}C_{\tau,\mu}\|_{W^{1,\infty}(\Theta)} \le C_r$ for some constant $C_r \ge 0$ to get the last one.

Injecting (25)-(26) in (24) gives:

$$ | \mathcal{E}_{\tau,r}^{C}(\mu) - \mathcal{E}_{\tau,r}^{C}(\nu) | \le \| \nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu} \|_{L^2(\Theta;\mu)} W_2(\mu, \nu) + o(W_2(\mu, \nu)) $$

and finally:

$$ |\partial \mathcal{E}_{\tau}^{C}|(\mu) \le \|\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}\|_{L^2(\Theta;\mu)}. \quad (27) $$

For the reverse inequality, one considers the following curves:

$$ \left\{ \begin{array}{l} \mu_s = \exp(-s\nabla_{\Theta}V_{\tau,\mu})\#\mu \\ \mu_{s,t} = \exp(t\nabla_{\Theta}C_{\tau,\mu_s})\#\mu_s. \end{array} \right. $$

and defines the function $f(s,t) := C_{\tau}(\mu_{s,t})$. We will use the implicit function theorem on $f$ and to do so, it is necessary to prove that it is $C^1$.

**Lemma 4.** The function $f(s,t) := C_{\tau}(\mu_{s,t})$ is $C^1$ on a neighborhood of $(0,0)$.

*Proof.* Let $dt, ds \in \mathbb{R}$, then we have for $\gamma := (i \times \exp(dt\nabla_{\Theta}C_{\tau,\mu_s}))\#\mu_{s,t}$:

$$ \begin{aligned} C_{\tau}(\mu_{s,t+dt}) &= C_{\tau}(\mu_{s,t}) + \int_{\Theta^2} \nabla_{\Theta} C_{\tau,\mu}(\theta) \cdot (\tilde{\theta} - \theta) d\gamma(\theta, \tilde{\theta}) + o(W_{2,\gamma}(\mu_{s,t}, \mu_{s,t+dt})) \\ &= C_{\tau}(\mu_{s,t}) + \int_{\Theta} \nabla_{\Theta} C_{\tau,\mu}(\theta) \cdot (\exp_{\theta}(dt\nabla_{\Theta}C_{\tau,\mu_s})) - \theta d\mu_{s,t}(\theta) + o(W_{2,\gamma}(\mu_{s,t}, \mu_{s,t+dt})). \end{aligned} $$
Hence as $W_{2,\gamma}(\mu_{s,t}, \mu_{s,t+dt}) = dt \|\nabla\Theta C_{\tau,\mu_s}\|_{L^2(\Theta; d\mu_{s,t})}$,

$$ \frac{C_{\tau}(\mu_{s,t+dt}) - C_{\tau}(\mu_{s,t})}{dt} = \int_{\Theta} \nabla_{\Theta} C_{\tau,\mu}(\theta) \cdot \frac{\exp_{\theta}(dt\nabla_{\Theta}C_{\tau,\mu_s}) - \theta}{dt} d\mu_{s,t}(\theta) + o(dt) $$

As $\nabla_{\Theta}C_{\tau,\mu}(\theta)$ is uniformly bounded, we can use the dominated convergence theorem and

$$ \lim_{dt \to 0} \frac{C_{\tau}(\mu_{s,t+dt}) - C_{\tau}(\mu_{s,t})}{dt} = \|\nabla_{\Theta}C_{\tau,\mu_{s,t}}\|_{\Theta;\mu_{s,t}}^{2}. \quad (28) $$

For the $s$-derivative, things get more complex since Jacobi fields are involved. The analysis is based on [11] where authors define the following Jacobi map:

$$ J_v(v_0, v_1)(\theta) := \begin{cases} \text{the value at } \zeta = 1 \text{ of the Jacobi field } \zeta \mapsto j_{\zeta} \in T_{\exp_{\theta}(\zeta v(\theta))}\Theta \\ \text{along the geodesic } \zeta \mapsto \exp_{\theta}(\zeta v(\theta)) \\ \text{having initial conditions } j_0(\theta) = v_0(\theta), j'_0(\theta) = v_1(\theta) \end{cases} \quad (29) $$

and define:

$$ j_v(v_0, v_1) \circ \exp(v(\theta)) := J_v(v_0, v_1)(\theta) \quad (30) $$

meaning that:

$$ \forall \theta \in \Theta, \tilde{\theta} \in \exp(v)^{-1}(\{\theta\}), \quad j_v(v_0, v_1)(\theta) = J_v(v_0, v_1)(\tilde{\theta}). $$

By the same reference [11], the curve $s \mapsto \mu_{s,t}$ is absolutely continuous and its velocity field $(w_s)_s$ writes:

$$ w_s = j_t \nabla_{\Theta} C_{\tau, \mu_s} \left( - \nabla_{\Theta} V_{\tau, \mu}, t \frac{d}{ds} \nabla_{\Theta} C_{\tau, \mu_s} \right) \quad (31) $$

where the $s$-derivative has to be understood in the following sense: $\frac{d}{ds} \nabla_{\Theta} C_{\tau, \mu_s} = \frac{\partial}{\partial s} (\nabla_{\Theta} C_{\tau, \mu_s} \circ \exp(-sV_{\tau, \mu}))$. Consequently, $s \mapsto f(s,t)$ is differentiable a.e and:

$$ \frac{\partial f}{\partial s} = \int_{\Theta} \nabla_{\Theta} C_{\tau, \mu_{s,t}} \cdot w_s d\mu_{s,t}. $$

To conclude, we need to prove the continuity of the rhs of last equation. To do so and by standard calculations, one needs compute the $s$-derivative of $s \mapsto P_{\tau}\mu_s$:

$$ \frac{\partial P_{\tau}\mu_s}{\partial s} = \int_{\Theta} \nabla_{\Theta} C_{\tau, \mu}(\theta) \cdot \nabla_{\Theta} \Phi_{\tau}(\theta, \cdot) d\mu_s(\theta) $$

and consequently $s \mapsto \nabla_{\Theta} C_{\tau, \mu_s}(\theta)$ is differentiable:

$$ \frac{d}{ds} \nabla_{\Theta} C_{\tau, \mu_s} = \nabla_{\Theta} \left( \left( \frac{\langle \frac{\partial P_{\tau} \mu_s}{\partial s}, \Phi_{\tau}(\theta) \rangle_{L^2(\Omega)}}{\|P_{\tau} \mu_s\|_{L^2(\Omega)}} - \langle P_{\tau} \mu_s, \Phi_{\tau}(\theta) \rangle_{L^2(\Omega)} \right) \times \frac{\langle P_{\tau} \mu_s, \frac{\partial P_{\tau} \mu_s}{\partial s} \rangle_{L^2(\Omega)}}{\|P_{\tau} \mu_s\|_{L^2(\Omega)}^3} \right) $$

which is continuous in the $s$ variable. As the manifold $\Theta$ is smooth, the application $J.(\cdot, \cdot)$ is smooth and the application $s \to w_s(\theta)$ is $C^1$. Then finally,

$$ \frac{\partial f}{\partial s}(s + ds, t) - \frac{\partial f}{\partial s}(s, t) = \int_{\Theta} \nabla_{\Theta} C_{\tau, \mu_{s,t}} \cdot w_s d(\mu_{s+ds,t} - \mu_{s,t}) + \int_{\Theta} (\nabla_{\Theta} C_{\tau, \mu_{s+ds,t}} \cdot w_{s+ds} - \nabla_{\Theta} C_{\tau, \mu_{s,t}} \cdot w_s) d\mu_{s,t} \quad (32) $$

and:

* the first term in (32) goes to zero as $\mu_{s,t}$, $\mu_{s+ds,t}$ are compactly supported (since $\mu$ is compactly supported) and function $\nabla_{\Theta}C_{\tau,\mu_{s,t}} \cdot w_s$ is continuous in the theta variable.

* The second term in (32) goes to zero because of the smoothness of the integrand wrt the $s$ variable and dominated convergence theorem which is easily applicable since the integrand is uniformly bounded on $\Theta$ wrt $s$ on the compact support of $\mu_{s,t}$.
This finishes the proof of the lemma.

□

By the boundedness of $\nabla\Theta V_{\tau,\mu}$, $\sigma_\mu\nabla\Theta C_{\tau,\mu}$ in $L^\infty$ and as Supp($\mu$) $\subset K_{r-\delta}$ for some $\delta > 0$, there exists an $\varepsilon > 0$ such that for $(s, t) \in [0, \varepsilon]^2$, Supp($\mu_s$) $\subset K_r$ and Supp($\mu_{s,t}$) $\subset K_r$. As $\mu \in X_\tau$, $f(0,0) = C_\tau(\mu) = 0$. Moreover:

$$ \frac{\partial f}{\partial t} |_{s,t=0} = \frac{d}{dt} C_\tau(\exp(t\nabla\Theta C_\mu)\#\mu) = \|\nabla\Theta C_{\tau,\mu}\|_{L^2(\Theta;\mu)}^2 \neq 0. $$

For the other derivative,

$$ \frac{\partial f}{\partial s} |_{s,t=0} = \frac{d}{ds} C_\tau(\mu_s) = -\langle \nabla\Theta V_{\tau,\mu}, \nabla\Theta C_{\tau,\mu} \rangle_{L^2(\Theta;\mu)}. $$

By the implicit function theorem, there exists an open interval $[0, \tilde{\varepsilon}]$ and a $C^1$ function $t: [0, \varepsilon] \to \mathbb{R}^+$ with $t'(0) = \frac{\frac{\partial f}{\partial s}|_{s,t=0}}{\frac{\partial f}{\partial t}|_{s,t=0}} = -\sigma_\mu$ such that for all $0 \le s < \tilde{\varepsilon}$, $f(s, t(s)) = 0$. To conclude, we need the two following estimates that will be derived using the notation $\psi_{s,t} = \exp(t\nabla\Theta C_{\mu_s}) \circ \exp(-s\nabla\Theta V_{\tau,\mu})$:

• We first estimate $W_2(\mu_{s,t(s)}, \mu)$:

$$ \frac{W_2^2(\mu_{s,t(s)}, \mu)}{s^2} \le \int_{\Theta} \frac{d(\psi_{s,t(s)}(\theta), \theta)^2}{s^2} d\mu(\theta) \quad (33) $$

passing to the limit:

$$ \limsup_{s \to 0} \frac{W_2^2(\mu_{s,t(s)}, \mu)}{s^2} \le \|\nabla\Theta V_{\tau,\mu} - \sigma_\mu \nabla\Theta C_{\tau,\mu}\|_{L^2(\Theta;\mu)}^2. \quad (34) $$

To justify this, we need a bit of work involving the dominated convergence theorem.

– The derivatives of $\psi$ are given below at $t, s = 0$:

$$ \left\{ \begin{array}{ll} \frac{\partial \psi}{\partial t} |_{s,t=0} (\theta) = & \nabla \Theta C_{\tau, \mu} (\theta) \\ \\ \frac{\partial \psi}{\partial s} |_{s,t=0} (\theta) = & - \nabla \Theta V_{\tau, \mu} (\theta). \end{array} \right. $$

Consequently, the pointwise limit of the rhs of (33) is $\left| \frac{\partial \psi}{\partial s} |_{s,t=0} (\theta) + \frac{\partial \psi}{\partial t} |_{s,t=0} (\theta) t'(0) \right| = |\nabla\Theta V_{\tau,\mu}(\theta) - \sigma_\mu \nabla\Theta C_{\tau,\mu}(\theta)|$.

– Now we will calculate the velocity field of the smooth curve $s \mapsto \psi_{s,t(s)}$. To do so consider the variation of geodesic:

$$ \iota(s, u) := \exp(ut(s)\nabla\Theta C_{\mu_s}) \circ \exp(-s\nabla\Theta V_{\tau,\mu}). $$

We have:

$$ \left\{ \begin{array}{ll} \frac{\partial \iota}{\partial s} |_{s,u=0} = & -\nabla\Theta V_{\tau,\mu} \\ \\ \frac{\partial^2 \iota}{\partial s \partial u} |_{s,u=0} = & \frac{\partial}{\partial s} (t(s)\nabla\Theta C_{\tau,\mu_s} \circ \exp(-s\nabla\Theta V_{\tau,\mu}))|_{s=0} = t'(0)\nabla\Theta C_{\tau,\mu} \end{array} \right. $$

where the last equality comes from the fact that $t(0) = 0$. Hence, the velocity field of the smooth curve $s \mapsto \psi_{s,t(s)}(\theta)$ is the Jacobi field $J(\theta)$ evaluated at $u = 1$ with initial condition given by:

$$ \left\{ \begin{array}{l} J_0(\theta) = -\nabla\Theta V_{\tau,\mu}(\theta) \\ J'_0(\theta) = -\sigma_\mu \nabla\Theta C_{\tau,\mu}(\theta) \end{array} \right. $$

By the boundedness of $J_0, J'_0$ on the support of $\mu$ (which is compact), we deduce that $J(\theta)$ is uniformly bounded on the support of $\mu$ and:
$$ \frac{d(\psi_{s,t(s)}(\theta), \theta)}{s} \leq \|J\|_{L^{\infty}(\text{supp}(\mu))} $$

so we can apply the dominated convergence theorem.

• Then for the energy:

$$ \mathcal{E}_{\tau,r}^{C}(\mu_{s,t(s)}) - \mathcal{E}_{\tau,r}^{C}(\mu) = \int_{\Theta} \nabla_{\Theta} V_{\tau,\mu}(\theta) \cdot (\psi_{s,t(s)}(\theta) - \theta) d\mu(\theta) + o(W_{2}(\mu, \mu_{s,t(s)})) $$

and using the same techniques as above, one can pass to the limit:

$$ \begin{aligned} \lim_{s \to 0} \frac{\mathcal{E}_{\tau,r}(\mu_{s,t(s)}) - \mathcal{E}_{\tau,r}(\mu)}{s} &= -\int_{\Theta} \nabla_{\Theta} V_{\tau,\mu}(\theta) \cdot (\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}) d\mu(\theta) \\ &= -\int_{\Theta} (V_{\tau,\mu}(\theta) - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}) \cdot (\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}) d\mu(\theta) \quad (35) \\ &= -\|\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}\|_{L^2(\Theta; \mu)}^2 \end{aligned} $$

where the first equality comes from the fact that $\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}$ is orthogonal to $\nabla_{\Theta} C_{\tau,\mu}$ by the definition of $\sigma_{\mu}$.

To conclude with (34)-(35):

$$ \begin{aligned} \liminf_{s \to 0} \frac{\mathcal{E}_{\tau,r}^{C}(\mu_{s,t(s)}) - \mathcal{E}_{\tau,r}^{C}(\mu)}{W_{2}(\mu, \mu_{s,t(s)})} &= \lim_{s \to 0} \frac{\mathcal{E}_{\tau,r}^{C}(\mu_{s,t(s)}) - \mathcal{E}_{\tau,r}^{C}(\mu)}{s} \lim_{s \to 0} \frac{s}{W_{2}(\mu, \mu_{s,t(s)})} \\ &\geq -\|\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}\| \end{aligned} $$

and

$$ |\partial \mathcal{E}_{\tau,r}^{C}|(\mu) \geq \|\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}\|_{L^2(\Theta; \mu)} \quad (36) $$

With (27)-(36), we have:

$$ |\partial \mathcal{E}_{\tau,r}^{C}|(\mu) = \|\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}\|_{L^2(\Theta; \mu)} $$

□

**Remark 1.** If $\text{Supp}(\mu) \not\subset K_{r-\delta}$ for all $\delta > 0$ and $\mu \in X_{\tau}$, one can easily reproduce the proof above to show that:

$$ |\partial \mathcal{E}_{\tau,r}^{C}|(\mu) = \|\Pi_{r}(\nabla_{\Theta} V_{\tau,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu})\|_{L^{2}(\Theta; \mu)} $$

where $\Pi_r$ is the projection on the tangent space of the manifold with boundary $K_r$.

The form of the local slope given by Proposition 2 and Remark 1 allows to prove its regularity.

**Proposition 4.** The local slope $|\partial \mathcal{E}_{\tau,r}^C|$ is continuous.

*Proof.* The map $P_{\tau} : \mathcal{P}_{2}(\Omega) \to H^{1}(\Omega)$ is continuous by [2, Corollary 1]. This easily implies that maps $\mu \ni \mathcal{P}_{2}(\Omega) \mapsto \nabla_{\Theta} V_{\tau,\mu} \ni L^{2}(\Theta; \mu)$ and $\mu \ni \mathcal{P}_{2}(\Omega) \mapsto \nabla_{\Theta} C_{\tau,\mu} \ni L^{2}(\Theta; \mu)$ are continuous. Also by definition of $\sigma_{\mu}$ and Lemma 3, the map $\mu \ni \mathcal{P}_{2}(\Omega) \mapsto \sigma_{\mu} \ni \mathbb{R}$ is also continuous. Finally, the projection map $\Pi_r : L^2(\Theta; \mu) \to L^2(\Theta; \mu)$ is continuous as it is 1-Lipschitz. □

Finally, we can prove the existence of the gradient curve in Proposition 5.

**Proposition 5.** For all $\mu_0 \in X_{\tau,r}$ for some $r > 0$, there exists a locally absolutely continuous curve $(\mu_t^r)_{t \ge 0}$ of maximal slope for $\mathcal{E}_{\tau,r}^C$ with respect to $|\partial \mathcal{E}_{\tau,r}^C|$. Moreover for almost all $t \ge 0$, there exists a vector field $v_t^r \in L^2(\Theta; \mu_t)^{d+2}$ such that:

$$ \int_{\Theta} \|v_t^r\|^2 d\mu_t = \|v_t^r\|_{L^2(\Theta; \mu_t)}^2 < +\infty $$
and:

$$ \left\{ \begin{array}{l} \partial_t \mu_t^r + \operatorname{div}(v_t^r \mu_t^r) = 0 \\ \mu(0) = \mu_0. \end{array} \right. \qquad (37) $$

**Remark 2.** Proposition 5 does not give the uniqueness of the curve. To do so, one should restrict the space to measure in $X_\tau$ and define a metric derived from the classic Wasserstein one for which the functional $\mathcal{E}_{\tau,r}$ is $\lambda$-geodesically convex. An example of such procedure is the work [12] where authors restrict the space of measures to those with fixed mean and variance. In addition, they construct a metric derived from $W_2$ using a midpoint iterative approach. Under such new metric, they prove the geodesically convexity of the entropy and a potential to state the existence and the uniqueness of the Fokker-Planck equation constrained by mean and variance. Transferring such an argument to our context is more difficult since the condition $\|P_\tau \mu\|_{L^2(\Omega)}^2 = 1$ has a much less obvious interpretation with respect to the measure $\mu$.

*Proof.* It suffices to apply [13, Theorem 2.3.1] owing the coercivity and the lower semicontinuity of $\mathcal{E}_{\tau,r}^C$ by Proposition 1 and the lower semicontinuity of $|\partial \mathcal{E}_{\tau,r}^C|$. The existence of the vector field $v_t$ for almost all $t \ge 0$ is given by the absolute continuity of the curve $[0, 1] \ni t \mapsto \mu_t$ (because it is a gradient curve) and by [14, Proposition 2.5]. □

**Lemma 5.** The velocity field $v_t^r$ in (39) is orthogonal to $\nabla \Theta C_{\mu_t^r}$ in $L^2(\Theta; \mu_t^r)$.

*Proof.* Using the proof of Proposition 5 of [2], one gets the differentiability of $\mathbb{R}_+ \ni t \to \mathcal{C}_\tau(\mu_t^r)$ almost everywhere and for almost all $t \ge 0$:

$$ \frac{d}{dt} [\mathcal{C}_{\tau}(\mu_t^r)] = \int_{\Theta} \nabla \Theta C_{\tau, \mu_t^r}(\theta) \cdot v_t^r(\theta) d\mu_t^r(\theta). $$

As $\mu_t^r$ belongs to $X_\tau$ for almost all time, then $\frac{d}{dt} [\mathcal{C}_\tau(\mu_t^r)] = 0$ for almost all time which proves the orthogonality. □

From now on we introduce the notation:

$$ V_{\tau, \mu}^{C} := V_{\tau, \mu} - \sigma_{\mu} C_{\tau, \mu}. \qquad (38) $$

and characterizes the velocity field $v_t^r$.

**Proposition 6.** The velocity field $v_t^r$ in (39) is equal to $\nabla \Theta V_{\tau, \mu_t^r}^C = (\nabla \Theta V_{\mu_t^r} - \sigma_{\mu_t^r} \nabla \Theta C_{\mu_t^r}) \mu_t^r$ - almost everywhere as long as there exists $\delta > 0$ such that $\operatorname{Supp}(\mu^r(t)) \subset K_{r-\delta}$.

*Proof.* Using again the proof of Proposition 5 of [2], one gets the differentiability of $\mathbb{R}_+ \ni t \to \mathcal{E}_\tau(\mu_t^r)$ almost everywhere and for almost all $t \ge 0$:

$$ \begin{aligned} \frac{d}{dt} [\mathcal{E}_{\tau,r}^C(\mu_t^r)] &= \int_{\Theta} \nabla \Theta V_{\tau, \mu_t^r}(\theta) \cdot v_t^r(\theta) d\mu_t^r(\theta) \\ &= \int_{\Theta} (\nabla \Theta V_{\tau, \mu_t^r} - \sigma_{\mu_t^r} \nabla \Theta C_{\tau, \mu_t^r})(\theta) \cdot v_t^r(\theta) d\mu_t^r(\theta) \\ &= \int_{\Theta} \nabla \Theta V_{\tau, \mu_t^r}^C(\theta) \cdot v_t^r(\theta) d\mu_t^r(\theta) \end{aligned} $$

where the second equation comes from the fact that $v_t^r$ is orthogonal to $\nabla \Theta C_{\mu_t^r}$ for the scalar product of $L^2(\Theta; \mu_t^r)$ by Lemma 5. As $\mu^r$ is a curve of maximal slope with respect to the upper gradient $|\partial \mathcal{E}_{\tau,r}^C|$ of $\mathcal{E}_{\tau,r}^C$, one has:

$$ \begin{aligned} \frac{d}{dt} [\mathcal{E}_{\tau,r}^C(\mu_t)] &= \int_{\Theta} \nabla \Theta V_{\tau, \mu_t^r}^C \cdot v_t^r(\theta) d\mu^r(t)(\theta) \le -\frac{1}{2} \|v_t^r\|_{L^2(\Theta; d\mu^r(t))}^2 - \frac{1}{2} |\partial \mathcal{E}_{\tau,r}^C|^2(\mu^r(t)) \\ &\le -\frac{1}{2} \|v_t^r\|_{L^2(\Theta; d\mu^r(t))}^2 - \frac{1}{2} \|\nabla \Theta V_{\tau, \mu_t^r}^C\|_{L^2(\Theta; d\mu_t^r)}^2 \end{aligned} $$
where we used Proposition 2. As a consequence,

$$ \int_{\Theta} \left( \frac{1}{2} |\nabla_{\Theta} V_{\tau, \mu_t^r}^C|^2 + \frac{1}{2} |v_t^r(\theta)|^2 - \nabla_{\Theta} V_{\tau, \mu_t^r}^C \cdot v_t^r(\theta) \right) d\mu^r(t)(\theta) \le 0 $$

and

$$ v_t^r = -\nabla_{\Theta} V_{\tau, \mu_t^r}^C \quad \mu_t^r \text{-a.e.} $$

□

In fact the gradient flow writes as the pushforward of the initial measure by the flow generated with the velocity field $-\nabla_{\Theta} V_{\tau, \mu_t^r}^C$.

**Proposition 7.** Let $\chi^r: \mathbb{R}_+ \times \Theta \to \Theta$ be the flow associated to the velocity field $-\nabla_{\Theta} V_{\tau, \mu_t^r}^C$:

$$ \left\{ \begin{array}{l} \partial_t \chi^r(t) = -\nabla_{\Theta} V_{\tau, \mu_t^r}^C \\ \chi^r(0; \theta) = \theta. \end{array} \right. $$

Then $\chi^r$ is uniquely defined, continuous, and for all $t \ge 0$, $\chi^r(t)$ is Lipschitz on $K_r$. Moreover, as long as $\text{Supp}(\mu_t^r) \subset K_{r-\delta}$ for some $\delta > 0$:

$$ \mu_t^r = \chi^r(t) \# \mu_0. $$

*Proof.* [Checked] This is a direct consequence of the fact that the vector field $\nabla_{\Theta} V_{\tau, \mu_t^r}^C$ is $C^\infty$.

□

Before stating the theorem of existence without support limitation, we need the following preliminary lemmas 6-7:

**Lemma 6.** For all $\mu \in X_{\tau,r}$, it holds:

$$ |\sigma_{\mu}| \le C(\tau) \left( \int_{\Theta} a^2 d\mu \right) \|P_{\tau}\mu\|_{H^1(\Omega)} $$

*Proof.* By the definition of $\sigma_{\mu}$ and Cauchy Schwartz inequality,

$$ |\sigma_{\mu}| \le \frac{\|\nabla_{\Theta} V_{\tau, \mu}\|_{L^2(\Theta; \mu)}}{\|\nabla_{\Theta} C_{\tau, \mu}\|_{L^2(\Theta; \mu)}}. $$

Then,

•

$$ \|\nabla_{\Theta} V_{\tau, \mu}\|_{L^2(\Theta; \mu)} \le C_{\tau} \sqrt{\int_{\Theta} a^2 d\mu} \|P_{\tau}\mu\|_{H^1(\Omega)} $$

• By Lemma 3,

$$ \|\nabla_{\Theta} C_{\tau, \mu}\|_{L^2(\Theta; \mu)} \ge 1/\sqrt{\int_{\Theta} a^2 d\mu} $$

Therefore combining both inequalities:

$$ |\sigma_{\mu}| \le C_{\tau} \left( \int_{\Theta} a^2 d\mu \right) \|P_{\tau}\mu\|_{H^1(\Omega)} $$

□

**Lemma 7.** For all $\mu \in X_{\tau,r}$, it holds:

$$ \forall \theta \in \Theta, \quad |\partial_a V_{\tau, \mu}^C| \le C_{\tau} r^2. $$
Proof.

$$
\begin{align*}
|\partial_a V_{\tau, \mu}^C| &= |d \mathcal{E}_{P_\tau \mu}(\sigma_\tau(w \cdot +b)) - \sigma_\mu \langle P_\tau \mu, \sigma_\tau(w \cdot +b) \rangle_{L^2(\Omega)}| \\
&\le C (\|P_\tau \mu\|_{H^1(\Omega)} + |\sigma_\mu|) \\
&\le C_\tau r^2
\end{align*}
$$

where we used Lemma 6 to get last inequality.

□

Finally, Theorem 3 gives the existence of a gradient curve removing the support penalization in the energy.

**Theorem 3.** Let $T > 0$, $r_0 > 0$ and $\mu_0 \in \mathcal{P}_2(\Theta)$ such that $\text{Supp}(\mu_0) \subset K_{r_0}$. There exists a curve $(\mu_t)_{t \ge 0}$ such that :

$$
\left\{
\begin{array}{l}
\partial_t \mu_t + \operatorname{div}((-\nabla_{\Theta} V_{\tau, \mu_t}^C)\mu_t) = 0 \\
\mu(0) = \mu_0.
\end{array}
\right.
\qquad (39)
$$

and for almost all $t \ge 0$ :

$$
\int_{\Theta} |\nabla_{\Theta} V_{\tau, \mu_t}^C|^2 d\mu_t = \|\nabla_{\Theta} V_{\tau, \mu_t}^C\|_{L^2(\Theta; \mu_t)}^2 < +\infty.
$$

Moreover, the solution satisfies :

$$
\forall t \ge 0, \mu_t = \chi(t) \# \mu_0
$$

with $\chi : \mathbb{R}_+ \times \Theta \to \Theta$ solution to

$$
\left\{
\begin{array}{l}
\partial_t \chi(t; \theta) = -\nabla_{\Theta} V_{\tau, \mu_t}^C(\theta) \\
\chi(0; \theta) = \theta.
\end{array}
\right.
$$

Proof. Let $T > 0$ and:

• $r_0 > 0$ be such that $\text{Supp}(\mu_0) \subset K_{r_0}$,

• $\mu^r : \mathbb{R}_+ \ni t \mapsto \mu^r(t)$ be a gradient curve associated to $\mathcal{E}_{\tau,r}$ for $r > r_0$.

The only way that mass escapes from $K_r$ is through the variable $a$, so we need to focus on the velocity
in this variable only. Now defines:

$$
r_t := \sup\{a \mid \theta = (a, w, b) \in \text{Supp}(\mu_t^r)\} \le r.
$$

By Lemma 7, it holds that $|\partial_a V_{\mu_t^r}^C(\theta)| \le C_\tau r_t^2$ for all $t \ge 0$ and by definition of $r_t$ :

$$
|\partial_a V_{\tau, \mu_t^r}^C(\theta)| \le C_\tau r_t r. \tag{40}
$$

This implies that:

$$
r_t \le r_0 e^{C_r t} \tag{41}
$$

Indeed, for all $(a, w, b) \in \text{Supp}(\mu_t^r)$ and by (40), there exists $a_0$ with $|a_0| \le r_0$ such that:

$$
|a| \le |a_0| + C_{\tau} r \int_{0}^{t} r_s ds
$$

and taking the sup:

$$
r_t \le r_0 + C_{\tau} r \int_{0}^{t} r_s ds
$$

and Gronwall lemma gives (41).

Therefore for all time $t \in [0, T_r := \frac{1}{C_\tau r} \ln(\frac{r+r_0}{2r_0})]$, $\text{Supp}(\mu^r(t)) \subset K_{(r+r_0)/2} \subset K_r$. Hence we do not hit the boundary of $K_r$ and the velocity field associated to the curve $\mu_t^r$ is $-\nabla V_{\tau, \mu_t}^C$ during such period of time (cf Proposition 6).
Defining for all $n \in \mathbb{N}^*$

$$r_n := (n + 1)r_0,$$

we can build inductively a gradient curve $(\mu_t)_{t \ge 0}$ on $[0, \frac{1}{Cr_0} \sum_{i=1}^n \frac{1}{(i+1)} \log \left( \frac{i+2}{2(i+1)} \right)]$. As the width of this interval is diverging, it is possible to construct a gradient curve on $\mathbb{R}^+$. By construction, the velocity field associated to this curve is $-\nabla V_{\tau, \mu_t}^C$. The rest of the proof comes from Propositions 5, 6 and 7. $\square$

# 4 Convergence

Our convergence result towards a global optimum is based on the following hypothesis on the initial measure $\mu_0$ for which we need to define the projection operator:

$$\pi_{w,b} := \begin{cases} \Theta \to S^{d-1} \times \mathbb{R} \\ (a, w, b) \mapsto (w, b). \end{cases}$$

**Hypothesis 1.** The support of the measure $\pi_{w,b}\#\mu_0$ verifies:

$$S_{\mathbb{R}^d}(1) \times [-\sqrt{d} - 2, \sqrt{d} + 2] \subset \text{Supp}(\pi_{w,b}\#\mu_0)$$

and $\mu_0 \in X_\tau$

As we will use Hypothesis 1, we need to be sure it is not empty.

**Proposition 8.** There exists a measure $\mu_0 \in \mathcal{P}_2(\Theta)$ satisfying Hypothesis 1.

*Proof.* Define $\mu_{00} := \delta_1 \times \delta_{(w,b)}$ for some arbitrary $(w,b) \in S^{d-1} \times [-\sqrt{d} - 2, \sqrt{d} + 2]$. Finally, set $\mu_0 := \delta_{1/\|P_\tau \mu_{00}\|_{L^2(\Omega)}} \times \delta_{(w,b)}$ such that $\|P_\tau \mu_0\|_{L^2(\Omega)} = 1$ and $\mu_0$ belongs to $X_\tau$. $\square$

Under such hypothesis, one gets a result of convergence in the spirit of a previous work from Bach and Chizat [1]:

**Theorem 4.** If $\mu_0$ satisfies Hypothesis 1 and $\mu(t)$ converges towards $\mu^* \in \mathcal{P}_2(\Theta)$ as $t$ goes to infinity, then $V_{\tau, \mu^*}^C = 0$ everywhere.

*Proof.* The proof is identical to the proof of [2, Theorem 7] and given in Appendix for completeness. $\square$

In fact the equilibrium $\mu^*$ given by Theorem 4 represents an eigenstate:

**Theorem 5.** If $\mu \in X_\tau$ is such that $V_{\tau, \mu}^C = 0$ everywhere, then $P_\tau \mu$ is an eigenvector of the Poisson-Neumann equation with $\sigma_\mu$ as eigenvalue.

*Proof.* $V_{\tau, \mu}^C = 0$ is equivalent to:

$$\forall \theta \in \Theta, \langle \nabla_x P_\tau \mu, \nabla_x \Phi(\theta; \cdot) \rangle_{L^2(\Omega)} + \langle W P_\tau \mu, \Phi(\theta; \cdot) \rangle_{L^2(\Omega)} - \sigma_\mu \langle P_\tau \mu, \Phi(\theta; \cdot) \rangle_{L^2(\Omega)} = 0.$$

As the space of features is dense in $H^1(\Omega)$, the conclusion is straightforward. $\square$

# 5 Numerical simulations

In this section, we present some numerical experiments to show that the method has good performance in practice. For all tests, we use the following parameters:

* We use the tensorflow/keras framework.

* Two-layer neural networks are utilized with a network width of either $m = 100$ or $m = 1000$.

* The dataset is made of $10^5$ points sampled uniformly from the domain $\Omega$.
* Batches are made of $n := 100$ points taken from the dataset.

* The energy $\mathcal{E}$, the constraint $\mathcal{C}$ and its derivatives are computed by Monte-Carlo approximation and automatic differentiation.

* The optimizer is the classical stochastic gradient descent (SGD) with learning rate equal to $\frac{1}{\tau m}$.

* At each time step, we normalize the last linear layer to impose the unit $L^2$ norm condition.

* To evaluate the performance of our method, we use a finite difference algorithm that is capable of computing eigenvectors and eigenvalues in the case where the potential depends only on the first two variables.

* Note that because of Monte-Carlo sampling, the neural network algorithm is stochastic. This is why, we ran our algorithms 8 times to evaluate the mean and the variance of our results.

In the first test, we take a potential such that the solution behaves well with respect to tensor trains method [15] i.e.:

$$W(x) = 100 \cos(2\pi x_1).$$

Figures 1-4 show the energy and $L^2$ error during optimization. The red horizontal line corresponds to the first eigenvalue computed by the finite difference algorithm while the $L^2$ error is computed with respect to the first eigenvector computed by this same algorithm. The thick blue line represents the mean value and the shaded area shows the variance of the quantity.

<figure><img src="image_2.png" /><figcaption>Figure 1: The tensorized case when $d = 2$</figcaption></figure>

Next when one increases the dimension, curves are given in Figure 2. Remark that for $m = 1000$, the time of optimization was taken larger to reach convergence.
<figure><img src="image_3.png" /><figcaption>Figure 2: The tensorized case when $d = 8$</figcaption></figure>

For both cases $d = 2$ and $d = 8$, we see that the approximation works well and that the width has little influence on the result. Next we give a less obvious test where tensor techniques behave well. We call this test the "cos diagonal test".

$$W(x) = -100 \cos(2\pi(x_1 - x_2))$$

Plots are given here for $d = 2$ and $d = 8$.

<figure><img src="image_4.png" /><figcaption>Figure 3: The cos diagonal test</figcaption></figure>

Finally, a test where tensor methods fail as the potential:
$$W(x) = -100e^{-\frac{1}{2}(x_1 - x_2)^2}.$$

is not separable, is studied. For obvious reasons, we call this test the "exponential diagonal" one. Plots are given in figure 4.

<figure><img src="image_5.png" /><figcaption>Figure 4: The exp diagonal test</figcaption></figure>

Our conclusions are given below:

* On each example, the algorithm seems to converge towards the fundamental eigenvalue of the operator involved.

* For the three tests globally speaking, neural networks behave well when dimension increases.

* The dependence with respect to the width is not very clear sometimes giving better results in the tensorized case (Figures 1-2) with less variance. On the contrary in the exponential diagonal case, a wider network deteriorates the error and the variance.

* Neural networks algorithms allow to have good approximation of the exact solution when tensor methods fail.

To compare more precisely the effect of increasing the width or the batch size, we compare the $L^2$ error after convergence for different parameters in Figure 5 for the cos diagonal test. It seems that increasing the width drastically does not improve the error and a threshold phenomenon occurs. In contrast, increasing the batch size appears to be beneficial to get more precise results with less variance.
<figure><img src="image_6.png" /><figcaption>Figure 5: $L^2$ error for different $n$ and $m$</figcaption></figure>

For the sake of completeness, we show a last case with a double well potential. First define $f(z) = (z^2 - 1)^{-2}$ and $g(z) = f(4(z - 0.5))$, the potential $W$ is given by:

$$W(x) = 100 \exp(-g(x_1)).$$

A visual representation of the potential is given in Figure 6.

<figure><img src="image_7.png" /><figcaption>Figure 6: The double well potential</figcaption></figure>

Owing to remarks related to Figure 5, we take a smaller network of size $m = 200$ and a larger batch size with $n = 1000$. The results given in Figure 7, are as precise as the single well cases.
<figure><img src="image_8.png" /><figcaption>Figure 7: The double well case</figcaption></figure>

# 6 Conclusion and perspectives

In this work, an eigencouple of the Schrödinger operator is approximated by a two-layer neural network of infinite width, which can be represented by a probability measure. Subsequently, a Wasserstein gradient flow with respect to a constrained energy is introduced, and an existence result is provided. Assuming the convergence of the gradient flow, it is shown that the gradient flow converges to a suitable measure that represents an eigenfunction. Finally, some numerical experiments demonstrate the efficacy of the method, which surprisingly converges to a minimal energy state. At this time, no proof exists to substantiate this fact theoretically. Furthermore, the convergence result assumes the convergence itself, which represents a technical limitation. Finally, numerical tests were conducted on very smooth potentials for which the result from [3] holds, meaning that the solution can be approximated by two-layer neural networks. In order to address the Schrödinger multibody problem presented in the introduction, it is essential to consider less regular potentials and, consequently, more complex neural networks. Furthermore to be relevant for applications, the results obtained should be subsequently more accurate than the one obtained in the present framework. One potential avenue for improvement is the use of the natural gradient technique, as introduced in [16] and the reference therein. We intend to investigate this path in a future research work.
References

[1] L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In *Advances in Neural Information Processing Systems*, volume 31. Curran Associates, Inc., 2018.

[2] M. Dus and V. Ehrlacher. Numerical solution of Poisson partial differential equations in high dimension using two-layer neural networks, 2023.

[3] J. Lu and Y. Lu. A priori generalization error analysis of two-layer neural networks for solving high dimensional Schrödinger eigenvalue problems. *Comm. Amer. Math. Soc.*, 2:1–21, 2022.

[4] Y. Lu, J. Lu, and M. Wang. A priori generalization analysis of the deep ritz method for solving high dimensional elliptic partial differential equations. In *COLT*, volume 134 of *Proceedings of Machine Learning Research*, pages 3196–3241. PMLR, 2021.

[5] S. Fournais, M. Hoffmann-Ostenhof, T. Hoffmann-Ostenhof, and H. Siedentop. Sharp regularity results for coulombic many-electron wave functions. *Commun. Math. Phys.*, 255(1):183–227, 2005.

[6] H. Yserentant. On the regularity of the electronic Schrödinger equation in Hilbert spaces of mixed derivatives. *Numer. Math.*, 98(4):731–759, 2004.

[7] H. Montanelli and Q. Du. New error bounds for deep relu networks using sparse grids. *SIAM J. Math. Data Sci.*, 1(1):78–92, 2019.

[8] D. Pfau, J. S. Spencer, A. G. D. G. Matthews, and W. M. C. Foulkes. Ab initio solution of the many-electron schrödinger equation with deep neural networks. *Phys. Rev. Res.*, 2(3), 2020.

[9] J. Hermann, Z. Schätzle, and F. Noé. Deep-neural-network solution of the electronic schrödinger equation. *Nat. Chem.*, 12(10):891–897, 2020.

[10] J. Lott and C. Villani. Ricci curvature for metric-measure spaces via optimal transport. *Ann. of Math. (2)*, 169(3):903–991, 2009.

[11] N. Gigli. Second order analysis on $(P_2(M), W_2)$. *Mem. Am. Math. Soc.*, 2009.

[12] E. A. Carlen and W. Gangbo. Constrained steepest descent in the 2-Wasserstein metric. *Ann. of Math. (2)*, 157(3):807–846, 2003.

[13] L. Ambrosio, N. Gigli, and G. Savaré. *Gradient Flows In Metric Spaces And In The Space Of Probability Measures*. Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition, 2008.

[14] M. Erbar. The heat equation on manifolds as a gradient flow in the Wasserstein space. *Ann. Inst. Henri Poincaré Probab. Stat.*, 46(1):1–23, 2010.

[15] M.S Dupuy. Tensor trains for high-dimensional problems, master notes, 2024.

[16] J. Müller and M. Zeinhofer. Achieving high accuracy with PINNs via energy natural gradient descent. In *Proc. 40th ICML*, volume 202, pages 25471–25485. PMLR, 2023.

[17] A. Hatcher. *Algebraic Topology*. Cambridge University Press, Cambridge, 2002.
# A Proof of Theorem 4

## A.1 Convergence of the measure a critical point

In the following, a LaSalle's principle argument is invoked in order to prove Theorem 4. For simplicity, we note $\mathcal{E}_{\tau, \infty}^C := \mathcal{E}_{\tau, \infty}^C$ for $0 < \tau < +\infty$.

### A.1.1 Escape from non suitable critical points

In this section, we use the notation $\omega := (w, b)$ so that:

$$ \theta = (a, w, b) = (a, \omega) $$

to make the difference between "linear" variables and "nonlinear" ones. Moreover in order to ease the reading, we introduce the notation:

$$ v_{\mu} := \nabla_{\Theta} V_{\tau, \mu}^{C}. $$

Let us first prove the following preliminary lemma.

**Lemma 8.** For all $\mu$ in $P_2(\Theta)$, it holds that:

$$ \max(|V_{\tau,\mu}(\theta)|, |\nabla_{\Theta}V_{\tau,\mu}(\theta)|) \le C(\tau) \sqrt{\int_{\Theta} |\theta|^2 d\mu} \times |\theta| $$

and

$$ \max(|C_{\tau,\mu}(\theta)|, |\nabla_{\Theta}C_{\tau,\mu}(\theta)|) \le C(\tau) \sqrt{\int_{\Theta} |\theta|^2 d\mu} \times |\theta|. $$

*Proof.* It is a direct application of [2, Corollary 1]. □

The velocity potential and the derived vector field are smooth in the sense given by Lemma 9.

**Lemma 9.** For all $\mu, \nu$ in $X_{\tau}$, it holds that

$$ \forall \theta \in \Theta, |V_{\tau,\mu}^{C}(\theta) - V_{\tau,\nu}^{C}(\theta)| \le C \frac{\max(1, M_{2,\mu\nu}^{7/2})}{\min(1, \sqrt{\int_{\Theta} a^2 d\mu}, \sqrt{\int_{\Theta} a^2 d\nu})^3} W_2(\mu, \nu) |\theta| $$

$$ \forall \theta \in \Theta, |v_{\mu}(\theta) - v_{\nu}(\theta)| \le C \frac{\max(1, M_{2,\mu\nu}^{7/2})}{\min(1, \sqrt{\int_{\Theta} a^2 d\mu}, \sqrt{\int_{\Theta} a^2 d\nu})^3} W_2(\mu, \nu) |\theta| $$

where:

$$ M_{2,\mu\nu} := \int_{\Theta} |\theta_1|^2 d\mu(\theta_1) + \int_{\Theta} |\theta_2|^2 d\nu(\theta_2) \quad (42) $$

*Proof.* Here we focus on $v_{\mu}$, the proof for $V_{\tau,\mu}^C$ being very similar. One can decompose $v_{\mu}$ as

$$ v_{\mu} := v_{\mu,1} + v_{2,\mu} - \sigma_{\mu} \nabla_{\Theta} C_{\tau,\mu}, \quad (43) $$

with

$$ v_{\mu,1} := \nabla_{\Theta} [\langle \nabla_x P_{\tau}\mu, \nabla_x \Phi_{\tau}(\theta; \cdot) \rangle_{L^2(\Omega)}], $$

$$ v_{\mu,2} := \nabla_{\Theta} [\langle W P_{\tau}\mu, \Phi_{\tau}(\theta; \cdot) \rangle_{L^2(\Omega)}]. $$

Using standard integral derivation and Fubini theorems, it holds that for all $\gamma \in \Gamma_o(\mu, \nu)$,

$$ v_{\mu,1}(\theta) - v_{\nu,1}(\theta) = \int_{\Theta^2} \int_{\Omega} \nabla_{\Theta} \nabla_x \Phi_{\tau}(\theta; x) (\nabla_x \Phi_{\tau}(\theta_1; x) - \nabla_x \Phi_{\tau}(\theta_2; x)) dx d\gamma(\theta_1, \theta_2). $$
Owing to (10), one gets

$$
\begin{aligned}
|v_{\mu,1}(\theta) - v_{\nu,1}(\theta)| &\le C(\tau) \int_{\Theta^2} \max(|\theta_1|, |\theta_2|) |\theta_1 - \theta_1| |\theta| d\gamma(\theta_1, \theta_2) \\
&\le C(\tau) \sqrt{\int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu} W_2(\mu, \nu) |\theta|,
\end{aligned}
$$

where we used the Cauchy-Schwartz inequality. For the second term in the decomposition (43), one
has:

$$
v_{\mu,2} - v_{\nu,2} = \int_{\Theta^2} \int_{\Omega} W(x)(\Phi_{\tau}(\theta_1; x) - \Phi_{\tau}(\theta_2; x)) \nabla_{\Theta} \Phi_{\tau}(\theta; x) dxd\gamma(\theta_1, \theta_2)
$$

Owing to (10), one gets:

$$
\begin{aligned}
|v_{\mu,2}(\theta) - v_{\nu,2}(\theta)| &\le C(\tau) \int_{\Theta^2} \int_{\Omega} \max(|\theta_1, \theta_2|) |\theta_1 - \theta_1| |dx d\gamma(\theta_1, \theta_2)| |\theta| \\
&\le C(\tau) \sqrt{\int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu} W_2(\mu, \nu) |\theta|
\end{aligned}
$$

where we used again the Cauchy-Schwartz inequality. For $\nabla_{\Theta} C_{\tau,\mu}$, it is bounded according to Lemma 8 and doing the same calculation as above with $W = 1$:

$$
|\nabla_{\Theta} C_{\tau,\mu}(\theta) - \nabla_{\Theta} C_{\tau,\mu}(\theta)| \leq C(\tau) \sqrt{\int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu} |\theta| W_2(\mu, \nu).
$$

The remaining term to analyze is $\sigma_\mu$ for which by Lemma 3 and Lemma 8:

$$
|\sigma_{\mu}| \leq C(\tau) \frac{\int |\theta|^2 d\mu}{\sqrt{\int a^2 d\mu}}.
$$

Moreover, it is Lipschitz:

$$
\begin{aligned}
|\sigma_{\nu} - \sigma_{\mu}| &\le \left| \frac{1}{\|\nabla C_{\tau,\mu}\|_{L^2(\Theta;\mu)}^2} - \frac{1}{\|\nabla C_{\tau,\nu}\|_{L^2(\Theta;\nu)}^2} \right| \langle \nabla_{\Theta} V_{\tau,\mu}, \nabla_{\Theta} C_{\tau,\mu} \rangle_{L^2(\Theta;\mu)} \\
&\quad + \frac{1}{\|\nabla C_{\tau,\nu}\|_{L^2(\Theta;\nu)}^2} \left| \langle \nabla_{\Theta} V_{\tau,\mu} - \nabla_{\Theta} V_{\tau,\nu}, \nabla_{\Theta} C_{\tau,\mu} \rangle_{L^2(\Theta;\mu)} \right| \\
&\quad + \frac{1}{\|\nabla C_{\tau,\nu}\|_{L^2(\Theta;\nu)}^2} \left| \langle \nabla_{\Theta} V_{\tau,\nu}, \nabla_{\Theta} C_{\tau,\mu} - \nabla_{\Theta} C_{\tau,\nu} \rangle_{L^2(\Theta;\mu)} \right|
\end{aligned}
\tag{44}
$$

For the first term denoted $A_1$ and introducing an arbitrary $\gamma \in \Gamma_o(\mu, \nu)$, it can be bounded by:

$$
\frac{|\int_{\Theta^2} (\nabla_{\Theta} C_{\tau,\mu}(\theta_1) - \nabla_{\Theta} C_{\tau,\nu}(\theta_2)) \cdot (\nabla_{\Theta} C_{\tau,\mu}(\theta_1) + \nabla_{\Theta} C_{\tau,\nu}(\theta_2)) d\gamma(\theta_1, \theta_2)|}{\|\nabla C_{\mu}\|_{L^2(\Theta;\mu)} \|\nabla C_{\nu}\|_{L^2(\Theta;\nu)}^2} \|\nabla V_{\mu}\|_{L^2(\Theta;\mu)}.
$$

and the term on the numerator can be decomposed as $A_1 = B_1 + B_2$ where:

$$
B_1 := \left| \int_{\Theta^2} (\nabla_{\Theta} C_{\tau,\mu}(\theta_1) - \nabla_{\Theta} C_{\tau,\nu}(\theta_1)) \cdot (\nabla_{\Theta} C_{\tau,\mu}(\theta_1) + \nabla_{\Theta} C_{\tau,\nu}(\theta_2)) d\gamma(\theta_1, \theta_2) \right|
$$

and:

$$
B_2 := \left| \int_{\Theta^2} (\nabla_{\Theta} C_{\tau,\nu}(\theta_1) - \nabla_{\Theta} C_{\tau,\nu}(\theta_2)) \cdot (\nabla_{\Theta} C_{\tau,\mu}(\theta_1) + \nabla_{\Theta} C_{\tau,\mu}(\theta_2)) d\gamma(\theta_1, \theta_2) \right|
$$

Reusing the estimate on $\nabla_{\Theta} C_{\tau,\mu}$ computed above and Lemma 8, one gets:

$$
\begin{aligned}
B_1 &\le C(\tau) \left( \int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu \right) W_2(\mu, \nu) \int_{\Theta^2} |\theta_1| (|\theta_1| + |\theta_2|) d\gamma(\theta_1, \theta_2) \\
&\le C(\tau) \left( \int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu \right)^2 W_2(\mu, \nu)
\end{aligned}
$$
For $A_2$, it suffices to use the uniform boundedness of $H_\Theta\Phi_\tau$ to prove that:

$$
\begin{aligned}
B_2 &\le C(\tau) \left( \int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu \right) \int_{\Theta^2} |\theta_1 - \theta_2|(|\theta_1| + |\theta_2|) d\gamma(\theta_1, \theta_2) \\
&\le C(\tau) \left( \int_{\Theta} |\theta_1|^2 d\mu + \int_{\Theta} |\theta_2|^2 d\nu \right)^{3/2} W_2(\mu, \nu).
\end{aligned}
$$

Using Lemma 3, one has:

$$
\frac{A_1 \|\nabla V_\mu\|_{L^2(\Theta; \mu)}}{\|\nabla C_\mu\|_{L^2(\Theta; \mu)} \|\nabla C_\nu\|_{L^2(\Theta; \nu)}^2} \le C(\tau) \frac{\max(1, M_{2,\mu\nu}^3)}{\sqrt{\int_{\Theta} |a_1|^2 d\mu \int_{\Theta} |a_2|^2 d\nu}} W_2(\mu, \nu).
$$

The bound of other terms in (44) are similar, easier to compute and left for the reader.

To conclude, using estimate derived above and Lemma 8:

$$
\begin{aligned}
|\sigma_\mu \nabla_\Theta C_{\tau,\mu} - \sigma_\nu \nabla_\Theta C_{\tau,\nu}| &\le |\sigma_\mu - \sigma_\nu| |\nabla_\Theta C_{\tau,\nu}| + |\sigma_\mu| |\nabla_\Theta C_{\tau,\mu} - \nabla_\Theta C_{\tau,\nu}| \\
&\le C(\tau) \frac{\max(1, M_{2,\mu\nu}^3)}{\sqrt{\int_{\Theta} |a_1|^2 d\mu \int_{\Theta} |a_2|^2 d\nu}} W_2(\mu, \nu) \sqrt{M_{2,\mu\nu}} |\theta| \\
&\quad + C(\tau) \frac{\int |\theta|^2 d\mu}{\sqrt{\int a^2 d\mu}} \sqrt{M_{2,\mu\nu}} |\theta| W_2(\mu, \nu) \\
&\le C(\tau) \frac{\max(1, M_{2,\mu\nu}^{7/2})}{\min(1, \sqrt{\int_{\Theta} a^2 d\mu}, \sqrt{\int_{\Theta} a^2 d\nu})^3} W_2(\mu, \nu) |\theta|.
\end{aligned}
$$

□

**Proposition 9.** Let $\mu \in \mathcal{P}_2(\Theta)$ such that there exists $\theta \in \Theta$, $\phi_\mu(\theta) \neq 0$. Then there exist a set $A \subset \Theta$ and $\varepsilon > 0$ such that if there exists $t_0 > 0$ with $W_2(\mu_{t_0}, \mu) \le \varepsilon$ and $\mu_{t_0}(A) > 0$, then there exists a time $0 < t_0 < t_1 < +\infty$ such that $W_2(\mu_{t_1}, \mu) > \varepsilon$.

*Proof.* As $\phi_\mu$ is linear in $a$, it can be written under the form

$$
V_{\tau, \mu}^{C}(\theta) := a\psi_{\mu}(\omega).
$$

By hypothesis, the set

$$
A_0 := \{ \theta \in \Theta \mid \phi_{\mu}(\theta) \neq 0 \}
$$

is a non empty (open set). This is equivalent to say that either there exists $\omega$ such that $\psi_\mu(\omega) \neq 0$ or
$r_\mu \neq 0$. Suppose that $\psi_\mu \neq 0$ is non zero somewhere. For all $\alpha \in \mathbb{R}$, we denote by

$$
\begin{cases}
A_{\alpha}^{+} = \psi_{\mu}^{-1}([\alpha], +\infty[), \\
A_{\alpha}^{-} = \psi_{\mu}^{-1}([ -\infty, \alpha[).
\end{cases}
$$

Now we focus on $A_0^-$ and suppose that this set is non empty. The case where $A_0^+$ is non empty is
similar to handle and left to the reader.

By Lemma 10 and the regular value theorem, there exists $\eta > 0$ such that $\partial A_{-\eta}^- = \psi_{\mu}^{-1}(\{-\eta\})$ is a
$(d+1)$-orientable manifold on which $\nabla_{\omega}\psi_{\mu}$ is non zero. With our choice of activation function $\sigma_{H,\tau}$,
it is easy to prove that $A_{-\eta}^-$ is a bounded set. Indeed, if $b$ is large enough, then $\Omega \ni x \mapsto \sigma_{H,\tau}(w \cdot x + b)$
is zero and $\psi_{\mu}(w, b)$ is zero.

On $A_{-\eta}^{-}$, the gradient $\nabla_{\omega}\psi_{\mu}$ is pointing outward $A_{-\eta}^{-}$ and, denoting by $n_{\text{out}}$ the outward unit vector to $A_{-\eta}^{-}$, there exists $\beta > 0$ such that $|\nabla_{\omega}\psi_{\mu} \cdot n_{\text{out}}| > \beta$ for on $\partial A_{-\eta}^{-}$, since this continuous function is nonzero on a compact set. Hence, defining:

$$
A := \{(a, \omega) \in \Theta \mid \omega \in A_{-\eta}^{-}, a \ge 0\}
$$

and owing to the fact that $v_\mu = (v_{\mu,a}, v_{\mu,\omega})$ with $v_{\mu,a} = \psi_\mu(\omega)$ and $v_{\mu,\omega} = a\nabla_\omega\psi_\mu(\omega)$, it holds:

$$
\left\{
\begin{array}{l}
v_{\mu,a} < \eta \text{ on } A \\
v_{\mu,\omega} \cdot n_{\text{out}} > \beta a \text{ on } \mathbb{R}_{+} \times \partial A_{-\eta}^{-}.
\end{array}
\right.
\qquad (45)
$$
By contradiction, suppose that $\mu_{t_0}$ has non zero mass on A and that $W_2(\mu, \mu_t) \le \varepsilon$ (with $\varepsilon$ fixed later) for all time $t \ge t_0$. Then using Lemma 9, one has:

$$|v_{\mu_t}(\theta) - v_{\mu}(\theta)| \le C(\tau, \mu)|\theta|\varepsilon \quad (46)$$

and

$$|V_{\tau, \mu_t}^C(\theta) - V_{\tau, \mu_t}^C(\theta)| \le C(\tau, \mu)|\theta|\varepsilon.$$

One takes $\varepsilon := \frac{\eta}{2C(\tau, \mu)R}$ where $R > 0$ satisfies:

$$R^2 \mu_{t_0}(A) > \int |\theta|^2 d\mu + \frac{\eta^2}{4C^2(\tau, \mu)R} \quad (47)$$

which exists since $\mu_{t_0}(A) > 0$ by hypothesis. On the set $\{\theta \in A \mid |\theta| \le R\}$ and by (46), we have:

$$|v_{\mu_t}(\theta) - v_{\mu}(\theta)| \le \frac{\eta}{2}$$

and so by (45):

$$\begin{cases} -v_{\mu_t, a} > \eta/2 & \text{on } A \\ -v_{\mu_t, \omega} \cdot n_{\text{out}} < -\beta/2 \times a & \text{on } \mathbb{R}^+ \times \partial A^-_{-\eta}. \end{cases}$$

The general picture is given by Figure 8. As a consequence, there exists a time $t_1$ such that the set $\{\theta \in A \mid |\theta| \le R\}$ has no mass and

$$\int |\theta|^2 d\mu_t(\theta) \ge R^2 \mu_t(A) \ge R^2 \mu_{t_0}(A).$$

At the same time, as $W_2(\mu, \mu_t) \le \varepsilon$:

$$\int |\theta|^2 d\mu_t(\theta) \le \int |\theta|^2 d\mu(\theta) + \varepsilon^2 = \int |\theta|^2 d\mu(\theta) + \frac{\eta^2}{4C(\tau, \mu)^2 R^2}$$

and this a contradiction with the condition (47) on R.

□

**Remark 3.** The set A constructed in the proof of previous lemma is of the form:

$$A := \{(a, \omega) \in \Theta \mid \omega \in A^-_{-\eta_1}\} \cup \{(a, \omega) \mid \omega \in A^+_{\eta_2}\} \quad (48)$$

where $\eta_1, \eta_2$ are strictly positive.

**Lemma 10.** For all $\mu \in \mathcal{P}_2(\Theta)$, if $\psi_\mu < 0$ somewhere, there exists a strictly negative regular value $-\eta$ ($\eta > 0$) of $\psi_\mu$.

*Proof.* As $\psi_\mu < 0$ somewhere and by continuity, there exists a non empty open $O \subset] -\infty, 0[$ such that $O \subset \text{range}(\psi_\mu)$. Next, we use the Sard-Morse theorem recalled below:

**Theorem 6 (Sard-Morse).** Let $\mathcal{M}$ be a differentiable manifold and $f: \mathcal{M} \to \mathbb{R}$ of class $C^n$, then the image of the critical points of $f$ (where the gradient is zero) is Lebesgue negligible in $\mathbb{R}$.

This result applies to $V_{\tau, \mu}^C$ and the image of critical points of $\phi_\mu$ is Lebesgue negligible. As a consequence, there exists a point $o \in O$ which is a regular value of $V_{\tau, \mu}^C$. As $o \in O$, it is strictly negative and this finishes the proof of the lemma.

□
<figure><img src="image_9.png" /><figcaption>Figure 8: The escape of mass towards large values of a</figcaption></figure>

### A.1.2 Convergence

This preliminary lemma gives an insight of why Hypothesis 1 is useful:

**Lemma 11.** For all $\mu \in \mathcal{P}_2(\Theta)$, $\theta \notin \mathbb{R} \times S_{\mathbb{R}^d}(1) \times [-\sqrt{d}-2, \sqrt{d}+2[$, $\tau > 1$, the potential writes:

$$V_{\tau, \mu}^{C}(\theta) = 0.$$

In particular, $V_{\tau, \mu}^{C}(\theta)$ does not depend on $a, w, b$.

*Proof.* For all $x \in \Omega$, $|b| > \sqrt{d} + 2, \tau > 1$:

$$|w \cdot x + b| \geq |b| - |x|_{\infty} |w|_{1} > 2$$

and

$$\sigma_{H,\tau}(w \cdot x + b) = 0.$$

This implies that for $|b| \geq \sqrt{d} + 2, \mu \in \mathcal{P}_2(\Theta)$, the potential $V_{\tau, \mu}^{C} = 0$. $\square$

In fact Hypothesis 1 is verified by the gradient curve $(\mu(t))_{t \geq 0}$ for all time. This is proved in the next lemma.

**Lemma 12.** If $\mu_0$ satisfies Hypothesis 1 then for all $t \geq 0$ and all open set $O \subset S_{\mathbb{R}^d}(1) \times [-\sqrt{d}-2, \sqrt{d}+2]$, $\mu_t(\mathbb{R} \times O) > 0$

The arguments of the proof of last lemma are based on fine tools of algebraic topology. One can find a nice introduction to the topic in the reference book [17]. With simple words, we enjoy the homotopy properties on the sphere to prove that the measure $\mu(t)$ keeps a large enough support.

*Proof.* For all $t \geq 0$, as $\mu_t = (\chi(t))\#\mu_0$, we have [1, Lemma C.8]:

$$\mathrm{Supp}(\mu_t) = \overline{\chi(t) (\mathrm{Supp}(\mu_0))}. \qquad (49)$$

Now let $\xi_t(w, b) := (P_{S_{\mathbb{R}^d}(1) \times \mathbb{R}} \circ \chi(t))((0, w, b))$ where $P_{S_{\mathbb{R}^d}(1) \times \mathbb{R}}$ is the projection on $S_{\mathbb{R}^d}(1) \times \mathbb{R}$ ($w, b$ variables). We claim that the choice of the function of activation lets the extremal spheres invariant ie $\xi_t(w, \pm(\sqrt{d}+2)) = (w, \pm(\sqrt{d}+2))$. Indeed, by Lemma 11 for $\theta = (a, w, \pm(\sqrt{d}+2))$, $V_{\tau, \mu}^{C}(\theta) = 0$ giving:

$$\begin{cases} v_{\mu,w}(\theta) = 0, \\ v_{\mu,b}(\theta) = 0 \end{cases}$$

and the claim is proven. Consequently by Lemma 13, the continuous map $\xi_t$ is surjective.

Now let $O \subset S_{\mathbb{R}^d}(1) \times [-\sqrt{d}-2, \sqrt{d}+2]$ be an open set. By what precedes, there exists a point $\omega \in S_{\mathbb{R}^d}(1) \times [-\sqrt{d}-2, \sqrt{d}+2]$ such that $\xi_t(\omega) \in O$ and $\chi(t)((0, \omega)) \in \mathbb{R} \times O$. As $(0, \omega)$ belongs to the support of $\mu_0$ by hypothesis then $\chi(t)((0, \omega))$ belongs to the support of $\mu(t)$ by (49) and:
$$ \mu_t(\mathbb{R} \times O) > 0 $$

which finishes the proof of the lemma.

<figure><img src="image_10.png" /></figure>

Lemma 13 gives conditions for the surjectivity of a continuous map on a cylinder.

**Lemma 13.** Let $f$ be a continuous map $f : S_{\mathbb{R}^d}(1) \times [0, 1] \to S_{\mathbb{R}^d}(1) \times [0, 1] =: C$, homotopic to the identity such that:

$$ \forall w \in S_{\mathbb{R}^d}(1), \quad \begin{cases} f(w, 0) = (w, 0), \\ f(w, 1) = (w, 1). \end{cases} $$

Then $f$ is surjective.

*Proof.* Suppose that $f$ misses a point $p$, then necessarily $p = (w, t)$ with $0 < t < 1$. We can write:

$$ g : C \to C \setminus \{p\} $$

the restriction of $f$ on its image. The induced homomorphism on homology groups writes:

$$ g_{\star} : H_{d-1}(C) \to H_{d-1}(C \setminus \{p\}). $$

Aside that, we have the classic information on homology groups of $C$ and $C \setminus \{p\}$:

$$ \left\{ \begin{array}{l} H_{d-1}(C) = H_{d-1}(S_{\mathbb{R}^d}(1)) \simeq \mathbb{Z}, \\ H_{d-1}(C \setminus \{p\}) = H_{d-1}(S_{\mathbb{R}^d}(1) \cup S_{\mathbb{R}^d}(1)) \simeq \mathbb{Z}^2 \end{array} \right. $$

where $\vee$ designates the wedge sum. Thus, the homomorphism $g_{\star}$ can be written as:

$$ g_{\star} : \mathbb{Z} \to \mathbb{Z}^2. $$

As $g$ lets the two spheres $w \to (w, 0)$, $w \to (w, 1)$ invariant, we have:

$$ g_{\star}(1) = (1, 1). $$

Now we note $i : C \setminus \{p\} \to C$ the canonical inclusion map. For all $(a, b) \in \mathbb{Z}^2$,

$$ i_{\star}(a, b) = a + b. $$

By hypothesis, $f$ is homotopic to the identity so $f_{\star} = I_{\star}$ and $f_{\star}(1) = 1$ but at the same time:

$$ f_{\star}(1) = i_{\star}g_{\star}(1) = i_{\star}((1, 1)) = 2 $$

which gives a contradiction.

<figure><img src="image_11.png" /></figure>

It allows to conclude on the convergence and prove Theorem 4.

*Proof of Theorem 4.* Suppose $V_{\tau, \mu^{\star}}^{C} \neq 0$ somewhere. Reusing the separation of variables (see the proof of Proposition 9):

$$ V_{\tau, \mu^{\star}}^{C}(\theta) = a\psi_{\mu}(w, b). $$

Hence $\psi_{\mu}$ is not identically zero and the set $A$ defined in (48) is not empty and verifies:

$$ A \subset \mathbb{R} \times S_{\mathbb{R}^d}(1) \times [-\sqrt{d}-2, \sqrt{d}+2] \quad (50) $$

by Lemma 11.

By Proposition 9, there exists $\varepsilon > 0$ such that if $W_2(\mu_{t_0}, \mu^*) \leq \varepsilon$ for some $t_0$ and $\mu_{t_0}(A) > 0$ then there exists a further time $t_1$ with $W_2(\mu_{t_0}, \mu^*) > \varepsilon$. As $(\mu_t)_{t \geq 0}$ converges towards $\mu^*$, there exists $t_0$ such that:
$$ \forall t \ge t_0, W_2(\mu_{t_0}, \mu^*) \le \varepsilon. $$

But by Lemma 12 and (50), for all time $\mu_t(A) > 0$ and consequently there exists a time $t_1 > t_0$ with:

$$ W_2(\mu_{t_0}, \mu^*) > \varepsilon $$

which gives the contradiction.

□