Analysis of mean-field models arising from self-attention dynamics in transformer architectures with layer normalization

Abstract

The aim of this article is to provide a mathematical analysis of transformer architectures using a self-attention mechanism with layer normalization. In particular, observed patterns in such architectures resembling either clusters or uniform distributions pose a number of challenging mathematical questions. We focus on a special case that admits a gradient flow formulation in the spaces of probability measures on the unit sphere under a special metric, which allows us to give at least partial answers in a rigorous way. The arising mathematical problems resemble those recently studied in aggregation equations but with additional challenges emerging from restricting the dynamics to the sphere and the particular form of the interaction energy. We provide a rigorous framework for studying the gradient flow, which also suggests a possible metric geometry to study the general case (i.e. one that is not described by a gradient flow). We further analyse the stationary points of the induced self-attention dynamics. The latter are related to stationary points of the interaction energy in the Wasserstein geometry, and we further discuss energy minimizers and maximizers in different parameter settings.

This article is part of the theme issue ‘Partial differential equations in data science’.

1. Introduction

Transformer architectures and the associated (self-)attention dynamics gained strong interest recently due to the success of artificial intelligence relying on them in several applications. Examples include large language models such as GPT-4 [1], multimodal large language models such as vision language transformers [2,3], text-to-image generation like Stable Diffusion [4] and protein folding with AlphaFold [5,6], which won the Nobel Prize in Chemistry in 2024.

The practical success of transformers and (self-)attention dynamics calls for developing detailed mathematical understanding which started recently in [7–19].

An interesting viewpoint on such dynamics is to interpret it as an interacting particle system [8,20,21], which allows for natural continuous-time and mean-field limits. The latter approach already provided valuable insights into feed-forward neural networks and their training dynamics (cf. [22,23]). In the context of transformers, this viewpoint also provides interesting (so far formal [9]) connections to gradient flows and the minimization of interaction energy for the particle measures. The latter is a topic of great recent interest due to various applications in biology and social interactions. Indeed, the self-attention dynamics in transformers share certain mathematical similarities with models used in opinion formation, which also exhibit similar emergence of clusters in certain cases [24–26]. In this work, we focus on cluster formation in the infinite time horizon. However, we note that the formation of metastable states is of special interest. For the case of isotropic interaction, metastability was studied in [27,28].

In this article, we proceed with the work in [9] on analysing transformer dynamics with layer normalization, focusing in particular on the case when the underlying dynamics has a gradient flow structure. Indeed, the continuum limit of the self-attention dynamics leads to a Wasserstein-type gradient flow for probability measures on the unit sphere $\mathcal{S}$ of the form

$${\displaystyle {\mathrm{\partial }}_t{\mu }_t={\mathrm{\nabla }}_{\mathcal{S}}\cdot ({\mu }_t{m}_{{\mu }_t}{\mathrm{\nabla }}_{\mathcal{S}}{\mathcal{E}}^{\prime }({\mu }_t)),}$$ (1.1)

where ${\nabla }_{\mathcal{S}}$ and ${\nabla }_{\mathcal{S}}\cdot$ are the tangential gradient and divergence, respectively, and $m_{\mu }=\frac{1}{{\mathcal{E}}^{\prime }(\mu )}$ is a non-local mobility. The underlying energy in this case is of the form

$${\displaystyle \mathcal{E}(\mu )={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (x)\mathrm{d}\mu (y),}$$ (1.2)

with $D$ being a symmetric matrix and ${\mathcal{E}}^{\prime }$ denoting its first variation. Since $D$ is symmetric and hence diagonalizable, we can equivalently assume that $D$ is a diagonal matrix, since we can use an orthogonal diagonalization and a corresponding transfer of variables to the eigenvectors, which leaves the unit ball unchanged. This will be used in several instances to simplify notation. It also permits a more detailed study of stationary patterns, in particular minimizers and maximizers of the energy.

Compared to the existing literature on such gradient flows, there are three distinct features that motivate our study, namely:

• —restriction of the dynamics to the unit sphere (a consequence of the layer normalization);
• —non-local mobility (a consequence of the self-attention mechanism), which is related to but still distinctly different from other variations of Wasserstein gradient flows studied recently (cf. [29–32]);
• —multiplicative coupling of states in the interaction energy, as opposed to commonly used interaction potentials depending only on the difference of the states (cf., e.g. [33–38]).

We make the gradient flow, formally introduced in [9], rigorous, showing that the transport distance with non-local mobilities is well defined, studying energy dissipation properties of the associated gradient flow and describing the large-time behaviour of the dynamics, specifically the convergence to stationary solutions, at least along subsequences. We further carry out a detailed study of energy minimizers and maximizers of $\mathcal{E}$ (extending the previously studied case of $D$ being a multiple of the identity) as well as stationary points of the energy in a Wasserstein setting, which we prove to be equivalent to stationary solutions of the dynamics. For the energy minimizers, we obtain an interesting picture depending on the structure of $D$ :

• —If there is a positive eigenvalue that is the eigenvalue of maximal absolute value, then a Dirac delta concentrated in the direction of a corresponding eigenvalue is a maximizer.
• —If the smallest eigenvalue is negative, then only a Dirac delta concentrated in the direction of a corresponding eigenvalue is a minimizer.
• —If the smallest eigenvalue is zero, then any measure concentrated on the null space of $D$ is a minimizer.
• —Dirac deltas concentrated in directions of arbitrary eigenvectors are stationary points. We also find some convex combinations of Dirac deltas being stationary points.
• —If the smallest eigenvalue is positive, we conjecture that the minimizer of the energy has full support on the unit sphere. To obtain some insight, we carry out a second-order asymptotic analysis of the minimizers for $D$ being a small perturbation of the identity.

We support our theoretical findings with several computational experiments and investigate the cases when the energy minimizers or maximizers cannot be characterized explicitly.

The rest of this work is organized as follows. In the remainder of the introduction, we recapitulate the simplified softmax transformer model introduced in[8], with additional layer normalization as considered in [9]. In §2, we provide a rigorous derivation of the gradient flow induced by the considered model. Sections 3 and 4 are dedicated to characterizing optimizers or stationary points of the studied energy, respectively. We support our findings by numerical experiments in §5 and summarize our results in §6.

(a). Self-attention

Transformer architectures [39] were developed in the field of natural language processing. Here, the input is usually a sentence, which is decomposed into a sequence of tokens (e.g. words or syllables). Each token (possibly along with its position in the sentence) is represented as a vector in a high-dimensional vector space. Apart from a conventional feed-forward component, the main feature of a transformer layer is the so-called attention mechanism. This mechanism implements interactions between tokens and was first introduced in [40] in the context of neural machine translation as an alternative to encoder–decoder approaches, the performance of which often deteriorates for large input lengths due to the use of latent representations of fixed dimensions.

Like [9], we shall focus on a simple yet widely used form of attention, the so-called self-attention. It can be formalized as follows: consider an input sequence $X=[X_i{]}_{i=1}^N\in {ℝ}^{N\times n}$ , where each $X_i\in {ℝ}^n$ represents an $n$ -dimensional token and $N$ denotes the number of tokens. The self-attention matrix $A\in {ℝ}^{N\times N}$ is given by

$${\displaystyle {\displaystyle A_{ij}=\frac{\exp \left(X_i\cdot D{X}_j\right)}{\sum _{k=1}^N\exp \left(X_i\cdot D{X}_k\right)},}}$$ (1.3)

where we assume $D\in {ℝ}^{n\times n}$ to be symmetric. The latter property does not necessarily hold for learned parameters in transformer architectures, but we expect the symmetric part to determine the asymptotic behaviour of the self-attention dynamics. Since the symmetry of $D$ allows one to interpret the dynamics as a gradient flow corresponding to a certain interaction energy, as observed in [9], it will allow us to analyse the asymptotic behaviour for this subclass; the study of the general case is left for future research. An important example of non-symmetric interaction is given by masked attention, which can be used to model causality. We refer to [41–43] for a mean-field interpretation of such dynamics.

By definition, the matrix $A$ is stochastic, i.e. each of its rows is a probability vector. Roughly speaking, the attention matrix determines how strongly a token is influenced by each other token. To determine how tokens influence each other, another matrix $V\in {ℝ}^{n\times n}$ , called the value matrix, is used. The influence of $X_j$ on $X_i$ can then be written as $A_{ij}V{X}_j$ and the self-attention layer $\mathcal{A}:{ℝ}^{N\times n}\to {ℝ}^{N\times n}$ is given by

$${\displaystyle {\displaystyle \mathcal{A}(X)={\left[X_i+\sum _{j=1}^N{A}_{ij}V{X}_j\right]}_{i=1}^N.}}$$ (1.4)

For our purposes, we assume $V=D$ or $V=-D$ since, in this case, one can show that the particles move along a gradient flow. The general case is the subject of future work.

(b). Normalization method

The normalization of intermediate values is a common practice in machine learning models. In the context of neural networks, so-called batch normalization [44] is a popular method to prevent gradients from blowing up and thus to stabilize (and to improve) the training. Since this form of normalization uses information from the entire training batch, [45] proposes layer normalization (LayerNorm), which translates the mean of an intermediate vector to zero and divides it by its standard deviation, and therefore does not depend on any other vector in the batch. While the original implementation of the transformer [39] uses LayerNorm, some of the more recent publications (e.g. Llama, [46]) use a simplified version called Root Mean Square Layer Normalization (RMSNorm) proposed in [47]. Up to a multiplication with learned weights $[g_i{]}_{i=1}^n$ , called gain parameters, RMSNorm performs a projection on to the unit sphere ${\mathcal{S}}^{n-1}$ (where in the following, we shall suppress the superscript and simply write $\mathcal{S}$ ). More precisely, for $x\in {ℝ}^n$ we write

$$\begin{array}{r}\mathrm{RMSNorm}(x{)}_i=g_i\frac{x_i}{\Vert x{\Vert }_2},\end{array}$$

where, in practice, a division by zero is circumvented by adding a small value $ϵ>0$ into $\Vert x{\Vert }_2$ . In our setting, we can assume the norm to be strictly positive as we consider the dynamics in continuous time. Following the setting of [9], we focus on RMSNorm with fixed gain parameters $g_i=1$ for all $i=1,\dots n$ and denote the projection on to the unit sphere for $x\in {ℝ}^n\setminus \{0\}$ by

$$\begin{array}{r}\mathit{\Pi }(x)=\frac{x}{\Vert x{\Vert }_2}.\end{array}$$

(c). Simplified transformer layer and time-continuous dynamics

Combining the attention layer with a normalization layer, we arrive at the following update step:

$$\begin{array}{r}X\leftarrow \mathit{\Pi }(\mathcal{A}(X)),\end{array}$$

where the projection is applied vector-wise to each row of $\mathcal{A}(X)$ . For the sake of our analysis, we shall deviate from typical practical implementations of transformers and consider the architecture to be a composition of such layers which all share the same matrices $D$ and $V$ in equations (1.3) and (1.4). In [9], it was proposed to study the continuum limit of these updates. This approach has become a popular tool for analyzing residual neural networks [48]: as discussed from various perspectives, e.g. in [49–52], the skip connections (i.e. the residual components) of the residual neural network architecture make it possible to interpret it as a forward Euler discretization of an ordinary differential equation. Introducing a time variable $t>0$ and a small time increment $\mathit{\Delta }t>0$ , we get

$${\displaystyle {\displaystyle X_i(t+\mathrm{\Delta }t)=\mathrm{\Pi }\left(X_i(t)+\mathrm{\Delta }t\sum _{j=1}^N{A}_{ij}(t)V{X}_j(t)\right),i=1,\dots ,N.}}$$ (1.5)

At this point, the residual component is hidden in the attention layer and cannot easily be extracted since the projection is nonlinear. In the continuous time limit $\mathit{\Delta }t\to 0$ , remembering that $\mathit{\Pi }(x)=x$ for any $x\in \mathcal{S}$ , we arrive at the following system of differential equations:

$$\begin{array}{}\text{(1.6)}& {\displaystyle {\displaystyle {\dot{X}}_i(t)=⟨{\mathrm{\nabla }}_x\mathrm{\Pi }\left(X_i(t)\right),\sum _{j=1}^N{A}_{ij}(t)V{X}_j(t)⟩,i=1,\dots ,N,}}\end{array}$$

where the spatial derivatives are understood as derivatives in ${ℝ}^n$ . With a simple computation, one can further show that for any $x\in \mathcal{S}$ and $z\in {ℝ}^n$ it holds that

$$\begin{array}{r}⟨{\nabla }_x\mathit{\Pi }(x),z⟩=P_x^{⟂}(z),\end{array}$$

where, following [9], we define $P_x^{⟂}(z)=z-x\cdot zx$ . Substituting this into equation (1.6), we arrive at the following dynamics:

$${\displaystyle {\displaystyle \{\begin{array}{l}{\dot{X}}_i(t)=P_{X_i(t)}^{\perp }\left(\sum _{j=1}^N{A}_{ij}(t)V{X}_j(t)\right),\text{(1.7a)}\\ X_i(0)=X_{0,i}\in \mathcal{S},\text{(1.7b)}\end{array}}}$$

which serve as a starting point of [9].

(d). Interpretation as an evolution of measures

Instead of studying the dynamics of distinct particles, [9] propose to view equation (1.7) as an evolution of an empirical measure

$${\displaystyle {\displaystyle {\mu }_t=\frac{1}{N}\sum _{i=1}^N{\delta }_{X_i(t)}.}}$$

The right-hand side of equation (1.7a) can be understood as an integral with respect to ${\mu }_t$ ; for a generic probability measure $\mu$ , this can be written as a measure-dependent velocity field:

$$\begin{array}{lrr}& \mathsf{V}[\mu ](x)=\frac{P_x^{⟂}({\int }_{\mathcal{S}}{e}^{x\cdot Dy}Vy\mathrm{d}\mu (y))}{{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (y)},& \end{array}$$ (1.8)

and equation (1.7a) turns into $\dot{X_i}(t)=\mathsf{V}[{\mu }_t](X_i(t))$ . With this notion, we recover the weak continuity equation formulated in [9]: for any test function $\phi \in C^1(\mathcal{S}\times [0,T])$ , one has

$${\displaystyle {\displaystyle \begin{array}{rl}\frac{d}{dt}{\int }_{\mathcal{S}}\varphi (t,x)\mathrm{d}{\mu }_t(x)& =\frac{d}{dt}\frac{1}{N}\sum _{i=1}^N\varphi (t,X_i(t))\\ & =\frac{1}{N}\sum _{i=1}^N{\mathrm{\partial }}_t\varphi (t,X_i(t))+⟨{\mathrm{\nabla }}_x\varphi (t,X_i(t)),\mathsf{V}[{\mu }_t](X_i(t))⟩\\ \text{(1.9)}& & ={\int }_{\mathcal{S}}{\mathrm{\partial }}_t\varphi (t,x)+⟨{\mathrm{\nabla }}_x\varphi (t,x),\mathsf{V}[{\mu }_t](x)⟩\mathrm{d}{\mu }_t(x),\end{array}}}$$

where, in this case, the spatial derivatives of $\phi$ have to be understood as derivatives on $\mathcal{S}$ .

Similarly, Geshkovski et al. [9] propose the interaction energy in equation (1.2), which for an empirical measure ${\mu }_t$ reduces to

$${\displaystyle {\displaystyle \mathcal{E}\left({\mu }_t\right)=\sum _{i,j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)}.}}$$

In this discrete case, a straightforward application of the chain rule and a reordering of the terms yields

$${\displaystyle {\displaystyle \frac{d}{dt}\mathcal{E}\left({\mu }_t\right)=2\sum _{i=1}^N\left(\sum _{j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)}D{X}_j(t)\right)\cdot {\dot{X}}_i(t).}}$$

Under our assumption that the value matrix is given by $V=\pm D$ , we see that, up to an application of $P_{X_i(t)}^{⟂}$ and a division by ${\displaystyle {\sum }_{j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)}}$ , the term in the brackets is given by $\dot{X_i}(t)$ . Since $P_x^{⟂}(z)\cdot z=P_x^{⟂}(z)\cdot P_x^{⟂}(z)$ for any $x\in \mathcal{S}$ , $z\in {ℝ}^n$ , we have that

$${\displaystyle {\displaystyle \frac{d}{dt}\mathcal{E}\left({\mu }_t\right)=\pm 2\sum _{i=1}^N{\Vert {\dot{X}}_i(t)\Vert }^2\sum _{j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)}\underset{\le }{\ge }0,}}$$

and hence the energy $\mathcal{E}$ increases ( $V=D$ ) or decreases ( $V=-D$ ) monotonously along the trajectory of ${\mu }_t$ . A formal derivation of the above formulae for general probability measures on smooth manifolds is provided in §2.

Let us mention that problems with similar energies as $\mathcal{E}$ have been studied in the past. The most prominent is an interaction energy with respect to $D$ with a non-local interaction kernel depending on $x-y$ . Choosing the kernel as Gaussian with covariance matrix $D^{-1}$ (which makes sense only if $D$ is positive definite) results in

$${\displaystyle \begin{array}{}\text{(1.10)}& {\mathcal{E}}^{\text{inter}}(\mu )={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{-\frac{1}{2}(x-y)\cdot (D(x-y))}\mathrm{d}\mu (x)\mathrm{d}\mu (y).\end{array}}$$

For $D=\pm \mathrm{Id}$ , the minimizers and maximizers of the expressions in equations (1.2) and (1.10) are equivalent as $\mp \frac{1}{2}(x-y)\cdot (x-y)=\mp \frac{1}{2}(x\cdot x+y\cdot y)\pm x\cdot y=\mp 1+\pm x\cdot y$ for all $x,y\in \mathcal{S}$ . The important difference between equations (1.2) and (1.10) is the rotation-(in)variance of the interaction functions $e^{x\cdot (Dy)}$ and $e^{-\frac{1}{2}(x-y)\cdot (D(x-y))}$ . In the general case, this is not true, but we shall use an analogy to the interaction energy to rewrite

$${\displaystyle \mathcal{E}(\mu )=e^{\lambda }{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{-\frac{\lambda }{2}|x{|}^2-\frac{\lambda }{2}|y{|}^2+x\cdot Dy}\mathrm{d}\mu (x)\mathrm{d}\mu (y)=e^{\lambda }{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{-\frac{\lambda }{2}|x-y{|}^2+x\cdot ((D-\lambda \mathrm{Id})y)}\mathrm{d}\mu (x)\mathrm{d}\mu (y).}$$

(e). Understanding $x\cdot Dy$ on the sphere

For our further analysis, it is crucial to understand the implications of restricting the problem to the unit sphere and the behaviour of the bilinear form $x\cdot Dy$ on it. For $D=\mathrm{Id}$ , it is clear that the minimizer of $f_y(x)=x\cdot Dy$ is given by $x=-y$ and the maximizer by $x=y$ . This changes for a general $D$ and as a result, the minimizer of the energy in equation (1.2) is not given by the uniform distribution on $\mathcal{S}$ anymore. For a diagonal matrix $D$ , the maximizer/minimizer of $f_y$ for a fixed $y\in \mathcal{S}$ with $Dy\ne 0$ is given by $x_{\pm }=\pm \frac{Dy}{\Vert Dy\Vert }$ . Therefore, we know that $x\cdot Dy=0$ if and only if $x\cdot x_{\pm }=0$ (same for $>$ and $<$ ). For $Dy=0$ , we already have $f_y(x)=0$ for any $x\in \mathcal{S}$ , i.e. each point is a minimizer, maximizer and orthogonal to $y$ w.r.t. $D$ . A further consequence is that

$$\begin{array}{r}\underset{x,y\in \mathcal{S}}{\max }x\cdot Dy=\underset{y\in \mathcal{S}}{\max }\frac{Dy\cdot Dy}{\Vert Dy\Vert }=\underset{y\in \mathcal{S}}{\max }\Vert Dy\Vert =\left|\lambda \right|,\end{array}$$

where $\lambda$ denotes the eigenvalue of maximum absolute value of $D$ . We further note that all of the following results on minimizers/maximizers as well as stationary points of ${\mathcal{E}}_D$ can be generalized to probability measures concentrated on an ellipsoid instead of a sphere. To see this, we consider the ellipsoid

$$\begin{array}{r}C\mathcal{S}=\{x\in {ℝ}^n:\Vert C^{-1}x\Vert =1\},\end{array}$$

where $C\in {ℝ}^{n\times n}$ is invertible, and the corresponding energy

$$\begin{array}{r}{\mathcal{E}}_D^C(\mu )={\int }_{C\mathcal{S}}{\int }_{C\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (x)\mathrm{d}\mu (y).\end{array}$$

Since $C$ is invertible, any measure $\mu$ is uniquely determined by the pushforward measure $\nu =C_{\#}^{-1}\mu$ , as $\mu =C_{\#}\nu$ . Thus, we can rewrite the energy as

$$\begin{array}{r}{\mathcal{E}}_D^C(\mu )={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{Cx\cdot DCy}\mathrm{d}\nu (x)\mathrm{d}\nu (y)={\mathcal{E}}_{C^{T}DC}(\nu ),\end{array}$$

and equivalently optimize the energy ${\mathcal{E}}_{C^{T}DC}$ on the sphere. A special case that leads to measures concentrated on an ellipsoid corresponds to RMSNorm normalization with non-vanishing gain parameters $g_i\ne 0$ . In this case, the ellipsoid is given by $G\mathcal{S}$ , where $G$ is a diagonal matrix with entries $[g_i{]}_{i=1}^n$ .

2. Gradient flow

As shown above, the particle dynamics can be ‘lifted’ by the use of empirical measures to the space of probability measures $\mathcal{P}(\mathcal{S})$ over the sphere. As mentioned in [9, Remark 3.3], for arbitrary probability measures, the connection between the partial dynamics and a corresponding continuity equation can be made by a mean field limit approach. Hence, instead of the particle dynamics, one can study the continuity equation:

$${\displaystyle \begin{array}{rl}{\mathrm{\partial }}_t\mu +\mathrm{d}\mathrm{i}\mathrm{v}(\mathsf{V}[\mu ]\mu )=0& \text{on }[0,T]\times \mathcal{S},\\ {\mu }_{|t=0}=\mu (0)& \text{on }\mathcal{S},\end{array}}$$ (2.1)

with the velocity field given by equation (1.8), which holds in the sense of distributions. Note that, in this section, we scale the energy by a factor of $1/2$ to be consistent with [9]. It was remarked in [9, ch. 3.3] that for $V=\pm D$ , the energy,

$${\displaystyle \mathcal{E}(\mu )=\pm \frac{1}{2}{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (x)\mathrm{d}\mu (y),}$$

is monotonic along these dynamics, and the partial differential equation (2.1) can be interpreted as a gradient flow for a modified optimal transport distance. However, as the authors of [9] acknowledge, there is a gap in the literature that prevents them from making this observation rigorous.

In this section, we aim to close this gap. We show that $\mathcal{P}(\mathcal{S})$ equipped with this new distance is a geodesic space with properties similar to the classical $2$ -Wasserstein space and prove that solutions of equation (2.1) are curves of maximal slope of $\mathcal{E}$ with respect to this distance and thus satisfy the energy dissipation equality

$$\begin{array}{r}\frac{\mathrm{d}}{\mathrm{d}t}\mathcal{E}({\mu }_t)=-{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}{\mu }_t(y)|\mathsf{V}[{\mu }_t](x){|}^2\mathrm{d}{\mu }_t(x)\text{for a.e. }t.\end{array}$$

Finally, we study the long-time behaviour of the dynamics and show that subsequences of the flow converge to stationary points of the energy $\mathcal{E}$ .

Let us mention that the basic analysis of this section related to the novel transport distance can be generalized in a rather straightforward way to the more general case of $D$ being non-symmetric and can thus provide the basis for future analysis of the non-gradient flow case with $V$ arbitrary and $D$ non-symmetric.

(a). Continuity equation on manifolds

Let $M$ be a compact $n$ -dimensional Riemannian manifold without a boundary, e.g. the sphere $\mathcal{S}\subset {ℝ}^n$ . The tangent bundle $TM={\bigsqcup }_{x\in M}{T}_{x}M$ is given by the disjoint union of all tangent spaces of all $x\in M$ . We denote by $\mathcal{P}(M)$ the space of Borel probability measures on $M$ , equipped with the standard narrow topology (e.g. [53, ch. 5.1]). The symbol $\rightharpoonup$ is used to indicate convergence in this topology. Let $I=(0,T)$ be an open interval, $\mu :t\to {\mu }_t\in \mathcal{P}(M)$ a narrowly continuous curve and $V:(x,t)\in M\times I\mapsto v_t(x)\in TM$ a Borel velocity field such that ${\int }_0^T{\int }_M|v_t(x)|\mathrm{d}{\mu }_t\mathrm{d}t<\infty$ . The continuity equation holds in the sense of distributions if

$$\begin{array}{lrr}& {\int }_{(0,T)}{\int }_M{\partial }_t\phi (x,t)+⟨\mathcal{D}\phi (x,t),v_t(x)⟩\mathrm{d}{\mu }_t\mathrm{d}t=0,\forall \phi \in C_c^1(M\times (0,T)).& \end{array}$$ (2.2)

Here, $\mathcal{D}$ denotes the differential on the manifold $M$ . Sometimes, we shall use ${\mathcal{D}}_x$ to clarify with respect to which variable the differential is taken. We define the set of solutions to the continuity equation as follows:

$${\displaystyle \begin{array}{r}CE(0,T):=\left\{(\mu ,v):\begin{array}{c}\mu :I\mapsto \mathcal{P}(M)\text{ is narrowly continuous},\\ {\int }_0^T{\int }_M|v_t(x)|\mathrm{d}{\mu }_t\mathrm{d}t<\mathrm{\infty },\\ (\mu ,v)\text{ satisfy the continuity equation,}\end{array}\right\}\end{array}}$$

Furthermore, we define $CE(0,T;\nu \to \eta )$ as the subset $(\mu ,v)$ such that ${\mu }_0=\nu$ , ${\mu }_T=\eta$ . For more details, we refer to appendix A(a).

(b). Distance

To interpret equation (2.1) as a gradient flow on $\mathcal{P}(M)$ , we need to modify the well-known dynamic formulation of the $2$ -Wasserstein distance [54] and introduce the following mobility:

$$\begin{array}{r}{m}_{\mu }(x)={\int }_{M}K(x,y)\mathrm{d}\mu (y).\end{array}$$

With this, the modified transport distance between ${\mu }_0,{\mu }_1\in \mathcal{P}(M)$ is defined as follows (see [9, Section 3.4.2]):

$$\begin{array}{lrr}& W_{m,2}^2({\mu }_0,{\mu }_1)=\inf \{{\int }_0^1{\int }_M{m}_{{\mu }_t}(x)|v_t(x){|}^2\mathrm{d}{\mu }_t(x)\mathrm{d}t:(\mu ,v)\in CE(\mathrm{0,1};{\mu }_0\to {\mu }_1)\}.& \end{array}$$ (2.3)

For $K\equiv 1$ , we recover the classical $2$ -Wasserstein distance. The dynamic (2.1) corresponds to the kernel $K(x,y)=e^{x\cdot Dy}$ , but for the sake of generality, we carry out the analysis for a more general class of kernels $K$ .

Assumption 1. The kernel $K(x,y)\in C(M\times M)$ is continuous, and there exists a constant $C>0$ such that $K(x,y)⩾C$ for all $x,y\in M$ .

Remark 2.1. The assumption that $K$ is bounded from below is vital for our analysis and covers the cases of interest in this article. Nonetheless, it would be interesting to see whether this assumption can be relaxed. For example, instead of a compact manifold $M$ , we could consider ${ℝ}^d$ as the underlying space and take $K$ to be a Gaussian or a bounded confidence kernel $K(x,y)={\mathsf{1}}_{|x-y|⩽1}$ as studied in [ 55 ].

As the next theorem shows, the infimum in equation (2.3) is actually attained by some $(\mu ,v)\in CE(\mathrm{0,1};{\mu }_0\to {\mu }_1)$ . The proof can be found in appendix A(b).

Theorem 2.2 (Existence of minimizers). For every pair ${\mu }_0,{\mu }_1\in \mathcal{P}(M)$ with $W_{m,2}({\mu }_0,{\mu }_1)<+\infty$ , there exists a couple $(\mu ,v)\in CE(\mathrm{0,1})$ such that

$$\begin{array}{r}{W}_{m,2}^2({\mu }_0,{\mu }_1)={\int }_0^1{\int }_M{m}_{{\mu }_t}(x)|v_t(x){|}^2\mathrm{d}{\mu }_t(x)\mathrm{d}t.\end{array}$$

Furthermore, such minimizers can be equivalently characterized as those of

$$\begin{array}{}\text{(2.4)}& {\displaystyle W_{m,2}\left({\mu }_0,{\mu }_t\right)=inf\left\{{\int }_0^T{\left({\int }_M{m}_{{\mu }_t}(x){\left|v_t(x)\right|}^2\mathrm{d}{\mu }_t(x)\right)}^{\frac{1}{2}}\mathrm{d}t:(\mu ,v)\in CE\left(0,T;{\mu }_0\to {\mu }_T\right)\right\}.}\end{array}$$

Using the theorem above, it is easy to show that $W_{m,2}$ is a distance on $\mathcal{P}(M)$ .

Theorem 2.3. The space $\mathcal{P}(M)$ equipped with $W_{m,2}$ is a complete metric space and its topology is equivalent to the one induced by the $2$ -Wasserstein distance which, since $M$ is compact, is equivalent to the topology of narrow convergence.

Proof. First, we check that $W_{m,2}$ is a distance. Indeed, (i) symmetry follows from simply rescaling time by ${\displaystyle \tilde{t}:t\in [0,T]\mapsto T-t\in [0,T]}$ ; (ii) definiteness: Since $m_{{\mu }_t}$ is bounded from below, $W_{m,2}(\mu ,\nu )=0$ implies that $v_t=0$ for $\mathit{\mu }$ -a.e. $(p,t)\in M\times (0,T)$ . Thus by equation (A 3) ${\displaystyle \mu =\nu }$ ; (iii) the triangle inequality follows from the characterization in equation (2.4) and the gluing property from proposition A.1. To show the equivalence of the distances, we observe that by assumption 1, $K(x,y)⩾C$ and since $M\times M$ is compact and $K(x,y)$ is continuous, we can also find a $\stackrel{\mathrm{~}}{C}$ such that $K(x,y)⩽\stackrel{\mathrm{~}}{C}$ . This implies that

$${\displaystyle \begin{array}{r}\frac{1}{C}{W}_2(\mu ,\nu )⩽W_{m,2}(\mu ,\nu )⩽\tilde{C}{W}_2(\mu ,\nu )<+\mathrm{\infty }\mathrm{\forall }\mu ,\nu \in \mathcal{P}(M),\end{array}}$$

and the distances are equivalent. Since $(\mathcal{P}(M),W_2)$ is complete, $(\mathcal{P}(M),W_{m,2})$ has to be complete as well.∎

Let us recall that in a general complete metric space $(X,d)$ , a curve $\gamma :[0,T]\to X$ is called absolutely continuous if there exists a function $m\in L^1(0,T)$ such that

$${\displaystyle \begin{array}{r}d({\gamma }_s,{\gamma }_r)⩽{\int }_s^{r}m(t)\mathrm{d}t\mathrm{\forall }s,r\in [0,T]\text{ with }s⩽r.\end{array}}$$ (2.5)

For an absolutely continuous curve $\gamma (t)$ , its metric derivative is defined by

$${\displaystyle \begin{array}{r}|\dot{\gamma }|(t):=\underset{h\to 0}{lim}\frac{d({\gamma }_{t+h},{\gamma }_t)}{h},\end{array}}$$

and it exists for a.e. $t\in (0,T)$ . It can be shown that $\left|\dot{\gamma }\right|$ is minimal in the sense that for all $m(t)$ satisfying equation (2.5), it holds that $\left|\dot{\gamma }\right|(t)⩽m(t)$ for a.e. $t\in (0,T)$ . The next lemma, which is proven in appendix A(c), characterizes absolutely continuous curves in $(\mathcal{P}(M),W_{m,2})$ .

Lemma 2.4. Let ${\mu }_t$ be an absolutely continuous curve w.r.t. $W_{2,m}$ . Then there exists a Borel velocity field $(v_t{)}_{t\in (0,T)}$ such that $(\mu ,v)\in CE(0,T)$ and

$${\displaystyle {\left({\int }_M{m}_{{\mu }_t}(x)|v_t(x){|}^2\mathrm{d}{\mu }_t(x)\right)}^{1/2}=|\dot{\mu }|(t)fora.e.t\in (0,T).}$$

Conversely, if $(\mu ,v)\in CE(0,T)$ and ${\int }_0^T{\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}t<+\infty$ then $t\to {\mu }_t$ is absolutely continuous and

$${\displaystyle |\dot{\mu }|(t)⩽{\left({\int }_M{m}_{{\mu }_t}(x)|v_t(x){|}^2\mathrm{d}{\mu }_t(x)\right)}^{1/2}fora.e.t\in (0,T).}$$

A metric space is called a length space if

$$d(x,y)=\inf {\int }_0^1\left|\dot{\gamma }\right|(t)\mathrm{d}t,$$

where the infimum is taken over all absolutely continuous curves $\gamma :[\mathrm{0,1}]\to X$ with $\gamma (0)=x$ and $\gamma (1)=y$ . If this infimum is obtained by a minimal curve, also called geodesic, we say that $(X,d)$ is a geodesic space. As it turns out, the minimal curves obtained in theorem 2.2 are such geodesics. This can be immediately deduced from equation (A 9) and the definition of the metric velocity,

Corollary 2.5. The space $(\mathcal{P}(M),W_{m,2})$ is a geodesic space.

(c). Gradient flows of the interaction energy

Let $W(x,y)\in C^1(M\times M)$ be a symmetric interaction kernel. The interaction energy is given by

$$\begin{array}{r}\mathcal{E}(\mu )≔\frac{1}{2}{\int }_{M\times M}W(x,y)\mathrm{d}\mu (x)\mathrm{d}\mu (y).\end{array}$$

Let us consider the following inverse duality map:

$${\displaystyle {\mathcal{J}}_2:x\in T{M}_p^{\ast }\mapsto |x{|}_{\ast }\underset{y\in T{M}_p:|y|=1}{\text{arg max}}x(y).}$$

Since all tangent spaces are finite-dimensional, this map is well defined. The application of ${\mathcal{J}}_2$ to a 1-form on $M$ (in particular, a differential of a function) yields a velocity field on $M$ . Below we show that gradient flows of the energy $\mathcal{E}$ with respect to the metric $W_{m,2}$ are given by weak solutions to PDEs of the form

$${\displaystyle \begin{array}{}\text{(2.6)}& {\mathrm{\partial }}_t\mu +\mathrm{d}\mathrm{i}\mathrm{v}\left(\frac{1}{m_{\mu }}{J}_2(\mathcal{D}W[\mu ])\mu \right)=0,\end{array}}$$

where $W[\mu ](x)={\int }_{M}W(x,y)\mathrm{d}\mu (y)$ . For $M=\mathcal{S}$ , $K(x,y)=e^{x\cdot Dy}$ and $W(x,y)=\pm e^{x\cdot Dy}$ equation (2.6) corresponds precisely to equation (2.1) if $V=\pm D$ . The sole difference between equation (2.6) and classical Wasserstein gradient flows is the presence of the factor ${\displaystyle \frac{1}{m_{\mu }}}$ . It arises since the modified transport distance punishes the movement of particles with a high mobility $m_{\mu }(x)$ . When we interpret $K(x,y)$ as an interaction kernel between particles, those particles interacting strongly with others are slowed down, while particles with low interaction are sped up.

Lemma 2.6 (Chain rule). Let $t\to {\mu }_t$ be an absolutely continuous curve in $W_{2,m}$ . Then $t\mapsto \mathcal{E}({\mu }_t)$ is absolutely continuous and

$${\displaystyle \begin{array}{}\text{(2.7)}& \frac{\mathrm{d}}{\mathrm{d}t}\mathcal{E}({\mu }_t)={\int }_M⟨\mathcal{D}W[{\mu }_t](x),v_t(x)⟩\mathrm{d}{\mu }_t(x)fora.et\in (0,T).\end{array}}$$

Proof. Let us consider an absolutely continuous curve $(\mu ,v)\in CE(\mathrm{0,1};\mu \to \nu )$ and the function $\eta :(x,t)\in M\times [0,T]\mapsto \frac{1}{2}{\int }_{M}W(x,y)\mathrm{d}{\mu }_t(y)$ . In the case when $\eta \in C^1(M\times [0,T])$ , we could use it as a test function in equation (A 3) and immediately obtain

$$\begin{array}{rl}\mathcal{E}({\mu }_T)-\mathcal{E}({\mu }_0)& ={\int }_M\eta (x,T)\mathrm{d}{\mu }_T(x)-{\int }_M\eta (x,0)\mathrm{d}{\mu }_0(x)\\ & ={\int }_0^T{\int }_M{\partial }_t\eta (x,t)\mathrm{d}{\mu }_t(x)+{\int }_M⟨\mathcal{D}\eta (t,x),v_t(x)⟩\mathrm{d}{\mu }_t(x)\mathrm{d}t\\ & ={\int }_0^T{\int }_M{\int }_M⟨{\mathcal{D}}_{x}W(x,y),v_t(x)⟩\mathrm{d}{\mu }_t(y)\mathrm{d}{\mu }_t(x)\mathrm{d}t<+\infty .\end{array}$$

The finiteness follows from the fact that we can bound $|{\mathcal{D}}_{x}W(x,y){|}_{\ast }$ uniformly on $M\times M$ . In the general case, we have to use a rather lengthy time mollification argument, see appendix A(d).∎

Equation (2.7) is reminiscent of the classical chain rule $\frac{\mathrm{d}}{\mathrm{d}t}F(x(t))=\nabla F(x(t))\cdot \dot{x}(t)$ for a function $F:{ℝ}^d\to ℝ$ and a curve $x:[0,T]\to {ℝ}^d$ . The velocity field $v_t$ can be viewed as the ‘derivative’ of the curve ${\mu }_t$ , while $\mathcal{D}W[{\mu }_t]$ is the corresponding ‘gradient’ of the interaction energy. Using this chain rule, we can estimate how fast the energy can decrease along a curve ${\mu }_t$ . Therefore, curves reaching this bound dissipate the energy as fast as possible and satisfy the so-called energy dissipation equality.

Lemma 2.7. For any absolutely continuous w.r.t. $W_{2,m}$ curve $({\mu }_t{)}_{t\in (0,T)}$ , we have that

$$\begin{array}{lrr}& \mathcal{E}({\mu }_T)-\mathcal{E}({\mu }_0)+\frac{1}{2}{\int }_0^T{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t\mathrm{d}t+\frac{1}{2}{\int }_0^T{\int }_M\frac{1}{m_{{\mu }_t}}|\mathcal{D}W[{\mu }_t]{|}_{\ast }^2\mathrm{d}{\mu }_t\mathrm{d}t⩾0.& \end{array}$$ (2.8)

Moreover, we have equality if and only if $({\mu }_t{)}_{t\in (0,T)}$ is a weak solution to equation (2.6) .

Proof. We can estimate the right-hand side of equation (2.7) by Hölder’s and Young’s inequalities:

$${\displaystyle \begin{array}{rl}{\int }_M\mathcal{D}W[{\mu }_t](v_t)\mathrm{d}{\mu }_t& ⩾-\sqrt{{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t}\sqrt{{\int }_M\frac{1}{m_{{\mu }_t}}|\mathcal{D}W[{\mu }_t]{|}_{\ast }^2\mathrm{d}{\mu }_t}\\ & ⩾-\frac{1}{2}{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t-\frac{1}{2}{\int }_M\frac{1}{m_{{\mu }_t}}|\mathcal{D}W[{\mu }_t]{|}_{\ast }^2\mathrm{d}{\mu }_t.\end{array}}$$

Integrating both sides of equation (2.7) from 0 to T, we obtain equation (2.8). Moreover, equality holds if and only if for a.e. $t$ and ${\mu }_t$ -a.e. we have $v_t=\frac{1}{m_{{\mu }_t}}{J}_2(\mathcal{D}W[{\mu }_t]$ ). Hence, ${\mu }_t$ is a weak solution to equation (2.6).∎

(d). Metric gradient flows

Let us put the previous calculations into the context of curves of maximal slope [53, ch. 1], which can be viewed as a way to generalize gradient flows to general metric spaces. We assume $(X,d)$ to be a complete metric space. Let $\mathcal{E}:X\to ℝ$ . A function $g:X\to [0,+\infty ]$ is called a strong upper gradient of $\mathcal{E}$ if for any absolutely continuous curve $x:[0,T]\to X$ the concatenation $g\circ x$ is Borel and

$$\begin{array}{r}|\mathcal{E}(x(t))-\mathcal{E}(x(s))|⩽{\int }_s^{t}g(x(r))\left|\dot{x}\right|(r)\mathrm{d}r\forall 0⩽s⩽t⩽T.\end{array}$$

If $\mathcal{E}(x(t))$ is non-increasing in $t$ then the application of Young’s inequality yields

$$\begin{array}{r}\mathcal{E}(x(t))-\mathcal{E}(x(s))+\frac{1}{2}{\int }_s^{t}g(x(r){)}^2+|\dot{x}|(r{)}^2\mathrm{d}r⩾0\forall 0⩽s⩽t⩽T.\end{array}$$

This observation allows us to define curves of maximal slope as those that decrease the energy as fast as possible.

Definition 2.8 (Curve of maximal slope). An absolutely continuous curve $x:[0,T]\to X$ is called a curve of maximal slope of $\mathcal{E}$ with respect to its strong upper gradient $g$ if $t\mapsto E(x(t))$ is non-increasing and

$$\begin{array}{r}\mathcal{E}(x(t))-\mathcal{E}(x(s))+\frac{1}{2}{\int }_s^{t}g(x(r){)}^2+|\dot{x}|(r{)}^2\mathrm{d}r⩽0\forall 0⩽s⩽t⩽T.\end{array}$$

Lemma 2.9. The map

$$\begin{array}{r}g:\mu \mapsto \sqrt{{\int }_M\frac{1}{m_{\mu }}|\mathcal{D}W[{\mu }_t]{|}_{\ast }^2\mathrm{d}\mu }\end{array}$$

is a strong upper gradient of $\mathcal{E}$ and solutions of equation (2.6) coincide with curves of maximal slope of $\mathcal{E}$ with respect to the strong upper gradient $g$ .

Proof. For an absolutely continuous w.r.t. $W_{2,m}$ curve ${\mu }_t$ , we can find, by lemma 2.4, a velocity field $(v_t{)}_{t\in (0,T)}$ such that $(\mu ,v)\in CE(0,T)$ and

$${\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}=\left|\dot{\mu }\right|(t)\text{for a.e. }t\in (0,T).$$

Then, the chain rule, lemma 2.6 yields

$${\displaystyle \begin{array}{r}|\mathcal{E}({\mu }_t)-\mathcal{E}({\mu }_s)|⩽{\int }_s^t|⟨\mathcal{D}W[{\mu }_t],v_r⟩|\mathrm{d}r⩽{\int }_s^{t}g({\mu }_r)|\dot{\mu }|(r)\mathrm{d}r,\end{array}}$$

and $g$ is a strong upper gradient. The coincidence of solutions of equation (2.6) and curves of maximal slope follows from lemma 2.7.∎

(e). Energy dissipation and large-time behaviour

Due to the missing geodesic convexity properties of the energy, we cannot expect convergence of the evolution to a unique minimizer in the large time limit. However, we can obtain some weaker results by further analysing the energy dissipation property:

$$\begin{array}{lcr}& \mathcal{E}({\mu }_t)+\frac{1}{2}{\int }_0^t{\int }_M{m}_{{\mu }_s}(x)|\nabla {\mathcal{E}}^{\prime }({\mu }_s){|}^2\mathrm{d}{\mu }_s(x)\mathrm{d}s⩽\mathcal{E}({\mu }_0).& \end{array}$$ (2.9)

As $s\to \infty$ , we can pick narrowly convergent subsequences of ${\mu }_s$ (i.e. converging weakly star in the Banach space of Radon measures). Moreover, the entropy dissipation inequality above implies

$${\int }_0^{\infty }{\int }_M{m}_{{\mu }_s}(x)|\nabla {\mathcal{E}}^{\prime }({\mu }_s){|}^2\mathrm{d}\mu {}_s(x)\mathrm{d}s<\infty ,$$

hence, along suitable subsequences, the entropy dissipation,

$${\displaystyle D(s)={\int }_M{m}_{{\mu }_s}(x)|\mathrm{\nabla }{\mathcal{E}}^{\prime }({\mu }_s){|}^2\mathrm{d}{\mu }_s(x),}$$

converges to zero since it is non-negative and bounded. To establish the existence of subsequences converging to stationary solutions, we need to identify the limit in suitable spaces. Under appropriate regularity assumptions on the interaction kernel $W$ (satisfied, for example, for the exponential kernel), this is a direct consequence of the Arzelà–Ascoli theorem.

Lemma 2.10. Let $M$ be a compact manifold without a boundary, $W\in C^{1,\alpha }(M\times M)$ for some $\alpha >0$ and symmetric. Moreover, let ${\mu }^n$ be a sequence of probability measures on $M$ . Then the sequences

$$m_{{\mu }^n}={\int }_{M}W(\cdot ,y)\mathrm{d}{\mu }^n(y)\text{and}\nabla {\mathcal{E}}^{\prime }({\mu }^n)={\int }_M{\nabla }_{x}W(\cdot ,y)\mathrm{d}{\mu }^n(y)$$

have uniformly convergent subsequences. If ${\mu }^n$ converges narrowly to ${\mu }^{\ast }$ , then $m_{{\mu }^n}$ converges uniformly to $m_{{\mu }^{\ast }}$ and $\mathrm{\nabla }{\mathcal{E}}^{\prime }({\mu }^n)$ converges uniformly to $\mathrm{\nabla }{\mathcal{E}}^{\prime }({\mu }^{\ast }).$

Lemma 2.10 combined with the entropy dissipation inequality (2.9) yields the following result.

Corollary 2.11. Let $M$ be a compact manifold without a boundary, $W\in C^{1,\alpha }(M\times M)$ for some $\alpha >0$ and symmetric. Then each weak solution ${\mu }_t$ of equation (2.1) with the velocity field given by equation (1.8) has a narrowly convergent subsequence ${\mu }_{t_n}$ as $t_n\to \infty$ , the limit of which is a stationary solution.

The following example connects the general results of this section with the transformer dynamics.

Example 2.12. The transformer dynamics for a finite number of particles described by equation (1.7) with $V=\pm D$ correspond to the choice $M=\mathcal{S}$ , $K(x,y)=e^{x\cdot Dy}$ and $W(x,y)=\pm e^{x\cdot Dy}$ . As discussed in §1d, the corresponding empirical measures ${\mu }_t$ fulfil the continuity equation (1.9). Thus, they solve equation (2.1) in the weak sense with the velocity field given by equation (1.8), and all requirements of corollary 2.11 are fulfilled. Therefore, there exists a subsequence of ${\mu }_t$ that converges narrowly to a stationary solution of the interaction energy ${\mathcal{E}}_D$ defined in equation (1.2).

This section establishes the relation between the particle model in equation (1.7) and gradient flows of interaction energies for the special cases $V=\pm D$ . The energy dissipation property equation (2.8) and convergence property from corollary 2.11 motivate the study of stationary solutions of the energy ${\mathcal{E}}_D$ , which we carry out in §§3 and 4. We shall start with minimizers and maximizers.

3. Explicit energy minimizers and maximizers

In this section, we compute explicit minimizers and maximizers of the energy ${\mathcal{E}}_D$ (from equation (1.2), i.e. without the factor $1/2$ ) in different scenarios, depending on the properties of the interaction matrix $D$ . We make the dependence on the matrix $D$ explicit by employing it as a subscript of the energy. The case $D=\mathrm{Id}$ has already been covered in [9, Proposition 3.4], where it is stated that a measure is a maximizer if and only if it is a Dirac delta placed at any point on the sphere, and a minimizer if and only if it is the uniform distribution. As we show below, for more general matrices, the position of optimal Diracs depends strongly on the eigenvalues of the matrix $D$ . We further derive a symmetry condition for minimizers of energies with a positive definite interaction matrix $D$ . This property yields an alternative, simpler proof that the uniform distribution is the only minimizer for $D=\mathrm{Id}$ .

(a). Maximal eigenvalue and related maximizers or minimizers

Like for $D=\mathrm{Id}$ , there are several cases in which the minimizers or maximizers of the energy ${\mathcal{E}}_D$ are given by Diracs concentrated at a single point. We start with the maximizers when the largest eigenvalue of $D$ is also an eigenvalue of the largest absolute value (or, respectively, minimizers when the smallest eigenvalue of $D$ is also an eigenvalue of the largest absolute value).

Theorem 3.1. Let $\lambda$ be an eigenvalue of maximal absolute value of $D$ and $Z_{\lambda }\subseteq \mathcal{S}$ the set of associated normalized eigenvectors. If ${\displaystyle \lambda >0}$ then ${\mu }^{\ast }={\delta }_z$ with $z\in Z_{\lambda }$ are the only maximizers of the energy ${\mathcal{E}}_D$ . If ${\displaystyle \lambda <0}$ then ${\mu }^{\ast }={\delta }_z$ with $z\in Z_{\lambda }$ are the only minimizers.

Proof. We consider the case ${\displaystyle \lambda >0}$ ; the case ${\displaystyle \lambda <0}$ can be treated similarly. For all $x,y\in \mathcal{S}$ , we have $e^{x\cdot Dy}⩽e^{\lambda }$ with equality if and only if $x=y=\pm z$ . Thus,

$${\mathcal{E}}_D(\mu )={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (x)\mathrm{d}\mu (y)⩽{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{\lambda }\mathrm{d}\mu (x)\mathrm{d}\mu (y)=e^{\lambda }={\mathcal{E}}_D({\mu }^{\ast }),$$

where the inequality is strict if $\mu$ is not concentrated on an eigenvector associated with $\lambda$ .∎

An example of the above setting is maximizing the energy for $D=\mathrm{Id}$ [9, Proposition 3.4], where the authors make a connection between the existence of concentrated maximizers and the so-called mode collapse of transformers often observed in practice. For a positive definite $D\ne \mathrm{Id}$ , theorem 3.1 shows that the set of maximizers is not only restricted to Dirac measures, but that it is actually finite. We summarize this insight in the following example and refer to §5a for an illustrating numerical example.

Example 3.2. If $D=\mathrm{Id}$ then ${\mu }^{\ast }={\delta }_z$ is a maximizer of the energy ${\mathcal{E}}_{\mathrm{Id}}$ for any $z\in \mathcal{S}$ . Similarly, for $D=-\mathrm{Id}$ , ${\mu }^{\ast }={\delta }_z$ is a minimizer for any $z\in \mathcal{S}$ . If $D\ne \mathrm{Id}$ is positive definite then ${\mu }^{\ast }={\delta }_z$ is a maximizer of ${\mathcal{E}}_D$ only if $Dz=\lambda z$ and $\lambda$ is the largest eigenvalue of $D$ . Similarly, for a negative definite $D\ne \mathrm{Id}$ , ${\mu }^{\ast }={\delta }_z$ is a minimizer only if $Dz=\lambda z$ and $\lambda$ is the smallest eigenvalue of $D$ .

In the remainder of this section, we study minimizers for matrices that do not fulfil the conditions of theorem 3.1.

(b). Minimizers for indefinite matrices

We now generalize the statement in theorem 3.1 to minimizers of energies where the matrix $D$ has at least one non-positive eigenvalue. In particular, we do not assume that the smallest eigenvalue is the eigenvalue of maximal absolute value. A key property is the following result that gives a lower bound on the energy in terms of the smallest eigenvalue of $D$ .

Lemma 3.3. Let $\overline{x}$ be the expected value of $x$ under $\mu$ , i.e. $\overline{x}:={\int }_{\mathcal{S}}x\mathrm{d}\mu (x)$ . Then

$$\begin{array}{lcr}& {\mathcal{E}}_D(\mu )⩾e^{\overline{x}\cdot D\overline{x}}.& \end{array}$$ (3.1)

If $D$ is not positive definite and $\lambda {}_{\min }$ is its smallest eigenvalue, it further holds that

$$\begin{array}{}\text{(3.2)}& {\displaystyle {\mathcal{E}}_D(\mu )⩾e^{{\lambda }_{min}}.}\end{array}$$

Proof. We use the convexity of exponential functions of the form $x\mapsto e^{x\cdot a}$ and $y\mapsto e^{b\cdot y}$ for arbitrary $a,b\in {ℝ}^n$ , which, with two applications of Jensen’s inequality, implies

$$\begin{array}{lrlr}& {\mathcal{E}}_D(\mu )& ={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (y)\mathrm{d}\mu (y){\int }_{\mathcal{S}}{e}^{x\cdot D\overline{x}}\mathrm{d}\mu (x)⩾e^{\overline{x}\cdot D\overline{x}}.& \end{array}$$ (3.3)

Since, further, ${\displaystyle \overline{x}\cdot D\overline{x}⩾{\lambda }_{min}\Vert \overline{x}{\Vert }^2}$ and $0⩽\Vert \overline{x}\Vert ⩽1$ , the monotonicity of the exponential function gives us

$${\displaystyle {\mathcal{E}}_D(\mu )⩾e^{min\left\{{\lambda }_{min},0\right\}}.}$$

If $D$ is not positive definite, we know that $\lambda {}_{\min }⩽0$ and the above inequality reduces to inequality (3.2).∎

A direct consequence of lemma 3.3 for indefinite matrices is that a Dirac measure that is concentrated on an eigenvector corresponding to the smallest eigenvalue is a minimizer of the energy. If the smallest eigenvalue is negative, we can even show that all minimizers are of this form. In the case of a vanishing smallest eigenvalue, it is necessary and sufficient that the measure is concentrated on the null space of $D$ .

Theorem 3.4. Consider a matrix $D$ that is not positive definite with the smallest eigenvalue $\lambda {}_{\min }⩽0$ . If $\lambda {}_{\min }<0$ , a measure minimizes the energy if and only if it is a Dirac measure placed at an eigenvector corresponding to $\lambda {}_{\min }$ . If $\lambda {}_{\min }=0$ , a measure minimizes the energy if and only if it is concentrated on the null space of $D$ .

Proof. We first assume $\lambda {}_{\min }<0$ . It follows directly from equation (3.2) that every Dirac measure concentrated on an eigenvector corresponding to $\lambda {}_{\min }$ is a minimizer. We further see that $\overline{x}\cdot D\overline{x}=\lambda {}_{\min }$ if only if $\overline{x}$ is an eigenvector corresponding to $\lambda {}_{\min }$ and $\Vert \overline{x}\Vert =1$ . This can only hold for Dirac measures. Thus, there are no other minimizers.

For $\lambda {}_{\min }=0$ , it also follows directly from equation (3.2) that every measure concentrated on the null space of $D$ minimizes the energy. However, $\overline{x}\cdot D\overline{x}=\lambda {}_{\min }$ holds for all measures that fulfil $\overline{x}=0$ . Still, the estimate in equation (3.3), obtained using Jensen’s inequality, is only an equality if $x\cdot Dx=\overline{x}\cdot D\overline{x}=0$ for $\mu$ -a.e. $x\in \mathcal{S}$ . Therefore, all minimizers are concentrated on the null space of $D$ .∎

Remark 3.5. In general, theorem 3.4 does not transfer to maximizers for matrices $D$ that are not negative definite. To see this, consider $D$ with the largest eigenvalue $\lambda {}_{\max }⩾0$ , the smallest eigenvalue $\lambda {}_{\min }<0$ and corresponding eigenvectors $z_{\min }$ and $z_{\max }$ . If further ${\displaystyle e^{{\lambda }_{max}}<\cosh ({\lambda }_{min})}$ , it holds that

$${\mathcal{E}}_D\left({\delta }_{z_{\max }}\right)=e^{{\lambda }_{\max }}<\cosh \left({\lambda }_{\min }\right)={\mathcal{E}}_D\left(\frac{{\delta }_{z_{\min }}+{\delta }_{-z_{\min }}}{2}\right)$$

and thus, ${\delta }_{z_{\max }}$ is not a maximizer. In the special case $\lambda {}_{\max }=0$ , the above inequality holds for all measures concentrated on the null space of $D$ and all $\lambda {}_{\min }<0$ .

At this point, we further note that the above strategy does not work for analysing minimizers for positive definite interaction matrices $D$ . In this case, lemma 3.3 not only gives us ${\mathcal{E}}_D(\mu )⩾e^0=1$ , but also $x\cdot Dx>0$ for all $x\in \mathcal{S}$ , so the inequality is strict for all measures $\mu \in \mathcal{P}(\mathcal{S})$ .

(c). Symmetry property for positive definite matrices

The remainder of this section gives the first characterization of minimizers of the energy when the interaction matrix is positive definite. More precisely, we can show that, in this case, all minimizers are symmetric, and the symmetry axes are determined by the eigenvectors of $D$ . The first step towards this is to show that the energy ${\mathcal{E}}_D$ is strictly convex if $D$ is positive definite.

Lemma 3.6. If $D$ is positive semi-definite (resp. positive definite) then ${\mathcal{E}}_D$ is convex (resp. strictly convex).

Proof. Since ${\mathcal{E}}_D$ is quadratic, convexity (resp. strict convexity) follows from the non-negativity (resp. positivity) of the quadratic form:

$${\displaystyle \mathcal{F}(\mu )={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\mu (x)\mathrm{d}\mu (y),}$$

for arbitrary signed Radon measures $\mu$ , e.g. [56, Proposition 2.11]. For $D$ positive semi-definite, there exists a unique positive semi-definite matrix square root $D^{1/2}$ and we can use the transformation $T(x)=D^{1/2}x$ . We denote by $T_{\#}\mu$ the pushforward of $\mu$ by $T$ , so that

$${\displaystyle \begin{array}{rl}\mathcal{F}(\mu )& ={\int }_{T(\mathcal{S})}{\int }_{T(\mathcal{S})}{e}^{x\cdot y}\mathrm{d}{T}_{\mathrm{\#}}\mu (x)\mathrm{d}{T}_{\mathrm{\#}}\mu (y)\\ & ={\int }_{T(\mathcal{S})}{\int }_{T(\mathcal{S})}{e}^{-\frac{1}{2}|x-y{|}^2}{e}^{\frac{1}{2}|x{|}^2}\mathrm{d}{T}_{\mathrm{\#}}\mu (x)e^{\frac{1}{2}|y{|}^2}\mathrm{d}{T}_{\mathrm{\#}}\mu (y).\end{array}}$$

Let $\mathrm{d}\eta =e^{\frac{1}{2}|x{|}^2}\mathrm{d}{T}_{\#}\mu (x)$ , then

$${\displaystyle \mathcal{F}(\mu )={\int }_{T(\mathcal{S})}{\int }_{T(\mathcal{S})}{e}^{-\frac{1}{2}|x-y{|}^2}\mathrm{d}\eta (x)\mathrm{d}\eta (y).}$$

The fact that the Gaussian kernel is positive definite (e.g. [57]) yields that $\mathcal{F}(\mu )>0$ unless $\nu$ vanishes. This can only happen if $\mu =0$ or, in the case of a semi-definite matrix $D$ , if $\mu$ is concentrated on the null space $\mathcal{N}(D)$ and $\mu (\mathcal{N}(D))=0$ . This yields the assertion.∎

Remark 3.7. The previous convexity result does not guarantee the convergence of the gradient flow in ( equation (2.6) ) to a global minimizer of $\mathcal{F}$ . For such results, usually a slightly different notion of convexity is required, the so-called geodesic convexity. The following example shows that besides the case of $D$ being a multiple of the identity, we do not have geodesic convexity for the classical $2$ -Wasserstein distance. We do not expect any improvements for our modified optimal transport distance.

Example 3.8. We consider a simple counterexample in ${\mathcal{S}}^1$ (equipped with the spherical distance) to show that $\mathcal{F}$ is not convex along $2$ -Wasserstein geodesics. Choose

$$D=\left(\begin{array}{cc}2& 0\\ 0& 1\end{array}\right)\text{and the curve}\gamma :t\in [\mathrm{0,1}]\mapsto \left(\begin{array}{c}\cos (-\frac{\pi }{4}+t\frac{\pi }{2})\\ \sin (-\frac{\pi }{4}+t\frac{\pi }{2})\end{array}\right).$$

Then ${\mu }_t≔{\delta }_{\gamma (t)}$ is a constant-speed geodesic in the $2$ -Wasserstein space connecting ${\delta }_{\gamma (0)}$ and ${\delta }_{\gamma (1)}$ . Clearly, the map $[\mathrm{0,1}]\ni t\mapsto \mathcal{F}(\gamma (t))$ is not convex, since

$$\mathcal{F}(\gamma (0))=\mathcal{F}(\gamma (1))=e^{1.5}<e^2=\mathcal{F}\left(\gamma \left(\frac{1}{2}\right)\right).$$

Such a counterexample can always be constructed as long as $D$ has two different eigenvectors. Lemma 3.6 does not contradict this counterexample, however, as it only implies the convexity of

$$[\mathrm{0,1}]\ni t\mapsto \mathcal{F}((1-t){\mu }_0+t{\mu }_1).$$

Having established convexity, we can show that reflecting a measure along the eigenvectors of $D$ and then normalizing it does not increase the energy. Moreover, if $D$ is positive definite and $\mu$ is not symmetric with respect to all eigenvectors of $D$ , one can always construct a symmetric measure with a smaller energy.

Lemma 3.9. Let $z$ be an eigenvector related to an eigenvalue $\lambda$ of a positive semi-definite matrix $D$ . For a measure $\mu$ , we define $\stackrel{\mathrm{~}}{\mu }$ as

$$\begin{array}{r}\stackrel{\mathrm{~}}{\mu }:=\frac{1}{2}(\mu +{H_z}_{\#}\mu ),H_z(x)=x-2(x\cdot z)z,\end{array}$$

where $H_z$ denotes a reflection. Then, ${\mathcal{E}}_D(\stackrel{\mathrm{~}}{\mu })⩽{\mathcal{E}}_D(\mu )$ and the inequality is strict if $D$ is positive definite and $\stackrel{\mathrm{~}}{\mu }\ne \mu$ .

Proof. Since $e^{x\cdot Dy}=e^{H_z(x)\cdot D{H}_z(y)}$ , it is straightforward to see that ${\mathcal{E}}_D(\mu )={\mathcal{E}}_D({H_z}_{\#}\mu )$ . The (strict) convexity of the energy yields the assertion.∎

As a direct consequence, we obtain a symmetry property of minimizers for positive definite $D$ .

Corollary 3.10. If $D$ is positive definite then each minimizer is symmetric with respect to its eigenvectors.

If $D$ is a positive multiple of the identity, one can easily show using the above result that the uniform distribution is the unique energy minimizer. This has been shown already in [9, Proposition 3.4] using properties of Gegenbauer polynomials [58, Proposition 2.2]. The symmetry property from corollary 3.10 gives an alternative—and straightforward—proof of this fact.

Proposition 3.11. If $D=\lambda \mathrm{Id}$ for $\lambda >0$ then the uniform distribution is the unique energy minimizer.

Proof. If $\mu$ is not uniform, we can find a unit vector $z$ such that with $H_z$ as in lemma 3.9, we have

$$\stackrel{\mathrm{~}}{\mu }=\frac{1}{2}(\mu +{H_z}_{\#}\mu )\ne \mu .$$

However, for $D=\lambda \mathrm{Id}$ , every unit vector is an eigenvector and lemma 3.9 implies that ${\mathcal{E}}_D(\stackrel{\mathrm{~}}{\mu })<{\mathcal{E}}_D(\mu )$ . Hence, the uniform distribution is the only minimizer of the energy.∎

Remark 3.12. The statement in proposition 3.11 does not transfer to maximizers for negative multiples of the identity. To see this, consider $D=\lambda \mathrm{Id}$ with $\lambda <0$ and let ${\mu }_0$ denote the uniform distribution on $\mathcal{S}$ . The symmetry of ${\mu }_0$ yields

$${\displaystyle {\mathcal{E}}_D\left({\mu }_0\right)=2{\int }_{{\mathcal{S}}^{+}}{\int }_{{\mathcal{S}}^{+}}{e}^{\lambda x\cdot y}+e^{-\lambda x\cdot y}\mathrm{d}{\mu }_0(x)\mathrm{d}{\mu }_0(y)=4{\int }_{{\mathcal{S}}^{+}}{\int }_{{\mathcal{S}}^{+}}\cosh (\lambda x\cdot y)\mathrm{d}{\mu }_0(x)\mathrm{d}{\mu }_0(y),}$$

where ${\mathcal{S}}^{+}:=\{x\in \mathcal{S}:x{}_1>0\}$ . Since $|x\cdot y|<1$ ${\mu }_0\times {\mu }_0$ -almost everywhere on ${\mathcal{S}}^{+}\times {\mathcal{S}}^{+}$ the integrand can be strictly bounded from above by $4\cosh (\lambda )$ . Since ${\mu }_0({\mathcal{S}}^{+})=1/2$ it follows that

$$\begin{array}{r}{\mathcal{E}}_D({\mu }_0)<\cosh (\lambda )={\mathcal{E}}_D(1/2({\delta }_z+{\delta }_{-z})),\end{array}$$

with $z\in \mathcal{S}$ . Therefore, ${\mu }_0$ cannot be a maximizer of ${\mathcal{E}}_D$ .

Remark 3.13. The above argument can be used to show that for arbitrary $D$ , one has

$$\begin{array}{r}{\mathcal{E}}_D(\mu )⩽{\mathcal{E}}_D\left(\frac{{\delta }_z+{\delta }_{-z}}{2}\right)\end{array}$$

for all symmetric measures $\mu$ if and only if $z$ is an eigenvector that corresponds to the eigenvalue of the largest absolute value. In the upcoming section, we use this insight to show that such measures are maximizers of ${\mathcal{E}}_D$ for negative semi-definite $D$ .

If $D$ has non-positive eigenvalues, theorems 3.1 and 3.4 still show that all minimizers are invariant with respect to reflections $H_z$ , where $z$ corresponds to a positive eigenvalue. However, if $D$ has negative eigenvalues, such reflections can increase the energy when they are applied to general, non-minimizing measures. This is illustrated by the following example.

Example 3.14. Consider the two-dimensional case with ${\displaystyle D=\text{diag}(\lambda ,1)}$ and ${\displaystyle \lambda <0}$ . For any $\theta \in [\mathrm{0,2}\pi )$ , denote by ${\delta }_{\theta }$ the Dirac delta placed at $(\cos (\theta ),\sin (\theta ))$ . Fix $\phi \in [\mathrm{0,2}\pi )$ and let

$$\begin{array}{r}\mu =\frac{1}{2}({\delta }_{\phi }+{\delta }_{\pi +\phi }).\end{array}$$

In the two-dimensional setting, the symmetrization is given by

$$\begin{array}{r}\stackrel{\mathrm{~}}{\mu }=\frac{1}{4}({\delta }_{\phi }+{\delta }_{\pi +\phi }+{\delta }_{-\phi }+{\delta }_{\pi -\phi }).\end{array}$$

Denoting, for convenience, $\cos (\phi )=c$ , we have

$${\displaystyle {\mathcal{E}}_D(\mu )-{\mathcal{E}}_D(\tilde{\mu })=\frac{1}{2}(\cosh \left|(\lambda -1)c^2+1\right|-\cosh \left|(-\lambda -1)c^2+1\right|).}$$

Since $t\mapsto \cosh (t)$ is strictly increasing for $t⩾0$ , we get that ${\mathcal{E}}_D(\mu )⩽{\mathcal{E}}_D(\stackrel{\mathrm{~}}{\mu })$ since

$$|(\lambda -1)c^2+1|=|-|\lambda |c^2+1-c^2|⩽\left|\lambda \right|c^2+1-c^2=\left|\right|\lambda |c^2+1-c^2|=|(-\lambda -1)c^2+1|$$

for any $0⩽c⩽1$ and ${\displaystyle \lambda ⩽0}$ , and the inequality is strict if and only if $0<c<1$ and ${\displaystyle \lambda <0}$ .

(d). Maximizers for negative semi-definite matrices

There is no apparent way to use the proof strategy from the previous section for showing that maximizers for negative definite matrices are symmetric, since the kernel $(x,y)\mapsto e^{x\cdot Dy}$ is not negative definite for a negative definite $D$ . However, we can show that the quadratic form $\mathcal{F}$ used to prove lemma 3.6 is non-positive for anti-symmetric measures. This yields a symmetry property of maximizers for negative semi-definite matrices.

Lemma 3.15. Let $D$ be a negative semi-definite matrix and $\mu$ a measure on the sphere. Define $\stackrel{\mathrm{~}}{\mu }$ as

$${\displaystyle \begin{array}{r}\mathrm{d}\tilde{\mu }(x)=\frac{1}{2}(\mathrm{d}\mu (x)+\mathrm{d}\mu (-x)).\end{array}}$$

Then ${\mathcal{E}}_D(\stackrel{\mathrm{~}}{\mu })⩾{\mathcal{E}}_D(\mu )$ and the inequality is strict if $\stackrel{\mathrm{~}}{\mu }\ne \mu$ and either $D$ is negative definite or $\stackrel{\mathrm{~}}{\mu }=\mu$ on the null space $\mathcal{N}(D)$ .

Proof. We denote by $N(x)=-x$ the negation and define

$$\begin{array}{r}{\mu }^{+}:=\mu ,{\mu }^{-}:=N_{\#}\mu ,\zeta :=1/2({\mu }^{-}-{\mu }^{+}).\end{array}$$

This yields that $\mathrm{d}\zeta (-x)=2(\mathrm{d}\mu (-x)-\mathrm{d}\mu (x))=-\mathrm{d}\zeta (x)$ and

$${\displaystyle \begin{array}{rl}{\mathcal{E}}_D(\zeta )=& {\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}\mathrm{d}\zeta (x)\mathrm{d}\zeta (y)={\int }_{{\mathcal{S}}^{+}}{\int }_{{\mathcal{S}}^{+}}{e}^{x\cdot Dy}\mathrm{d}\zeta (x)\mathrm{d}\zeta (y)+{\int }_{{\mathcal{S}}^{+}}{\int }_{{\mathcal{S}}^{+}}{e}^{x\cdot Dy}\mathrm{d}\zeta (-x)\mathrm{d}\zeta (-y)\\ & +2{\int }_{{\mathcal{S}}^{+}}{\int }_{{\mathcal{S}}^{+}}{e}^{-x\cdot Dy}\mathrm{d}\zeta (-x)\mathrm{d}\zeta (y)=2{\int }_{{\mathcal{S}}^{+}}{\int }_{{\mathcal{S}}^{+}}{e}^{x\cdot Dy}-e^{-x\cdot Dy}\mathrm{d}\zeta (x)\mathrm{d}\zeta (y)=-{\mathcal{E}}_{-D}(\zeta ).\end{array}}$$

Since $-D$ is positive semi-definite, the proof of lemma 3.6 shows that ${\mathcal{E}}_{-D}(\zeta )⩾0$ and thus ${\mathcal{E}}_D(\zeta )⩽0$ . The inequality is strict if $\zeta \ne 0$ and either $D$ is negative definite or $\zeta$ is concentrated on $\mathcal{N}(D{)}^{⟂}$ . The symmetry of the kernel yields ${\mathcal{E}}_D({\mu }^{-})={\mathcal{E}}_D({\mu }^{+})$ . Further, by substituting ${\mu }^{+}=\stackrel{\mathrm{~}}{\mu }+\zeta$ and ${\mu }^{-}=\stackrel{\mathrm{~}}{\mu }-\zeta$ , we see that

$${\displaystyle \begin{array}{rl}{\mathcal{E}}_D(\tilde{\mu })& =\frac{1}{4}{\mathcal{E}}_D\left({\mu }^{+}\right)+\frac{1}{4}{\mathcal{E}}_D\left({\mu }^{-}\right)+\frac{1}{2}{\mathcal{E}}_D\left({\mu }^{+},{\mu }^{-}\right)\\ & =\frac{1}{2}{\mathcal{E}}_D(\mu )+\frac{1}{2}{\mathcal{E}}_D(\tilde{\mu }+\zeta ,\tilde{\mu }-\zeta )=\frac{1}{2}{\mathcal{E}}_D(\mu )+\frac{1}{2}{\mathcal{E}}_D(\tilde{\mu })-\frac{1}{2}{\mathcal{E}}_D(\zeta ).\end{array}}$$

Reordering the terms leads to

$${\displaystyle {\mathcal{E}}_D(\tilde{\mu })={\mathcal{E}}_D(\mu )-{\mathcal{E}}_D(\zeta )⩾{\mathcal{E}}_D(\mu ).}$$

From the conditions on $\zeta$ and $D$ that lead to ${\mathcal{E}}_D<0$ , we derive that the above inequality is strict if $\stackrel{\mathrm{~}}{\mu }\ne \mu$ and either $D$ negative definite or $\stackrel{\mathrm{~}}{\mu }=\mu$ on $\mathcal{N}$ .∎

Corollary 3.16. Let ${\mu }^{\ast }$ be a maximizer of ${\mathcal{E}}_D$ for a negative definite $D$ . Then $\mathrm{d}{\mu }^{\ast }(x)=\mathrm{d}{\mu }^{\ast }(-x)$ .

This symmetry property is the missing ingredient for showing that the discrete measures introduced in remarks 3.12 and 3.13 are maximizers for negative semi-definite matrices $D$ .

Theorem 3.17. Let $D$ be negative semi-definite and $\lambda {}_{\min }<0$ its smallest eigenvalue. Then, a measure $\mu$ maximizes ${\mathcal{E}}_D$ if and only if ${\mu }^{\ast }=1/2({\delta }_z+{\delta }_{-z})$ where $z\in \mathcal{S}$ is an eigenvector associated with $\lambda {}_{\min }$ .

Proof. By lemma 3.15, it suffices to consider $\mu$ satisfying $\mathrm{d}\mu (x)=\mathrm{d}\mu (-x)$ . Denoting ${\mathcal{S}}^{+}:=\{x\in \mathcal{S}:x_1>0\}$ and using the symmetry property of $\mu$ , with the arguments from remark 3.12, we have

$${\mathcal{E}}_D(\mu )⩽\cosh {\lambda }_{\min }={\mathcal{E}}_D\left({\mu }^{\ast }\right),$$

where equality is only obtained if $|x\cdot Dy|=\lambda {}_{\min }$ holds $\mu \times \mu$ -almost everywhere on ${\mathcal{S}}^{+}\times {\mathcal{S}}^{+}$ . Since $\mu$ is symmetric, this is equivalent to $\mu ={\mu }^{\ast }$ . For a negative definite $D$ , we already know from corollary 3.16 that there are no other measures that maximize ${\mathcal{E}}_D$ . In the negative semi-definite case, we have that any $\mu$ that fulfils ${\mathcal{E}}_D(\mu )=\cosh {\lambda }_{\text{min}}$ has to be concentrated on $\mathcal{N}(D{)}^{⟂}$ and, therefore, also in this case, there are no other maximizers.∎

4. Energy variation and stationary points

To study stationary points or local maximizers/minimizers, it is useful to consider the first and second variations of the energy on the Wasserstein space of probability measures on the sphere, as studied previously for Vlasov-type interactions, e.g. the mean-field aggregation equation, cf. [36,59,60]. The first variation of ${\mathcal{E}}_D$ is given by

$${\displaystyle \mathrm{d}{\mathcal{E}}_D(\mu ;V)=\frac{\mathrm{d}}{\mathrm{d}t}{\mathcal{E}}_D({\mu }_t){|}_{t=0},}$$ (4.1)

where ${\mu }_t$ satisfies

$$\begin{array}{lcr}& {\partial }_t{\mu }_t+\nabla \cdot ({\mu }_t{P}_x^{⟂}V)=0,{\mu }_0=\mu ,& \end{array}$$ (4.2)

and $P_x^{⟂}=\mathrm{Id}-x{x}^T$ is the projection to the tangent space of the unit ball at $x$ . Here, the velocity field $V$ is an arbitrary Lipschitz function on ${ℝ}^n$ ; by the projection $P_x^{⟂}$ , we restrict it further to admissible velocities that keep the distribution on the unit sphere.

The following weak formulation, where $\phi$ is a continuously differentiable test function, will be useful later:

$$\frac{\mathrm{d}}{\mathrm{d}t}{\int }_{\mathcal{S}}\phi (x)\mathrm{d}\mu {}_t(x)={\int }_{\mathcal{S}}{P}_x^{⟂}\nabla \phi (x)\cdot V(x)\mathrm{d}\mu {}_t(x).$$

Similar to the first variation, the second variation of ${\mathcal{E}}_D$ can be defined as

$$\begin{array}{lcr}& \mathrm{d}{}^2{\mathcal{E}}_D(\mu ;V.W)=\frac{\mathrm{d}}{\mathrm{d}t}\mathrm{d}{\mathcal{E}}_D({\mu }_t,W){|}_{t=0}& \end{array}$$ (4.3)

if the derivative on the right-hand side exists. The computation of the first variation is completely analogous to the case of the aggregation equation (cf. [59]) and thus omitted here.

Lemma 4.1. For any Lipschitz continuous vector field $V$ , the first variation of the energy ${\mathcal{E}}_D$ in the direction $V$ exists and is given by

$$\begin{array}{lcr}& \mathrm{d}{\mathcal{E}}_D(\mu ;V)={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot (Dy)}{P}_x^{⟂}Dy\cdot V(x)\mathrm{d}\mu (x)\mathrm{d}\mu (y).& \end{array}$$ (4.4)

It is straightforward to see that the first variation vanishes at the extremal points of the energy:

Proposition 4.2. Let ${\mu }^{\ast }$ be a minimizer or maximizer of the energy. Then $\mathrm{d}{\mathcal{E}}_D(\mu ;V)=0$ for all Lipschitz vector fields $V$ .

Proof. Let ${\mu }^{\ast }$ be the initial value for the transport equation (4.2). For Lipschitz-continuous vector fields, there is a unique solution ${\mu }_t$ of the transport equation, and for all times $t>0$ , it is an admissible distribution on the sphere. Hence, if ${\displaystyle {\mu }^{\ast }}$ is a minimizer, then

$${\mathcal{E}}_D({\mu }^{\ast })⩽{\mathcal{E}}_D({\mu }_t)$$

for all $t>0$ , which implies that $\mathrm{d}{\mathcal{E}}_D({\mu }^{\ast };V)⩽0$ in the limit $t\downarrow 0$ . Since $V$ is arbitrary and $\mathrm{d}{\mathcal{E}}_D$ is linear in $V$ , we have that $\mathrm{d}{\mathcal{E}}_D(\mu ;V)=0$ . The case of a maximizer is treated in the same way, with an opposite inequality initially.∎

The connection between the transformer dynamics and the energy variations in Wasserstein spaces is readily established in the following.

Lemma 4.3. A probability measure $\mu$ is a stationary solution of equation (2.1) with the velocity field given by equation (1.8) if and only if $\mathrm{d}{\mathcal{E}}_D(\mu ;W)=0$ for all Lipschitz continuous $W$ .

Similarly to lemma 4.1, one can obtain an expression for the second variation.

Lemma 4.4. For $V,W$ being Lipschitz continuous, the second variation of the energy ${\mathcal{E}}_D$ in the directions $V$ , $W$ exists and is given by

$${\displaystyle \mathrm{d}{\mathcal{E}}_D(\mu ;V,W)={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}((P_x^{\perp }Dy\cdot V(x))(P_x^{\perp }Dy\cdot W(x))+(Dy{)}^T\mathrm{\nabla }(P_x^{\perp }V(x)))\mathrm{d}\mu (x)\mathrm{d}\mu (y).}$$

(a). Energy variation at concentrated distributions

From lemma 4.1, we see that any measure $\mu$ that fulfils

$$\begin{array}{lrr}& {\int }_{\mathcal{S}}{e}^{x\cdot Dy}{P}_x^{⟂}Dy\mathrm{d}\mu (y)=0\text{for }\mu \text{-almost all }x\in \mathcal{S},& \end{array}$$ (4.5)

is a stationary point of ${\mathcal{E}}_D$ . Here and in the following, with a slight abuse of notation, we denote the $0$ -vector by $0$ . For concentrated measures, the above condition is also necessary and rather easy to verify, as we see in what follows. We first show that single Dirac measures can only be stationary points if they align with an eigenvector of the matrix $D$ .

Lemma 4.5. A Dirac measure ${\mu }^{\ast }={\delta }_z$ is a stationary point of ${\mathcal{E}}_D$ if and only if $z$ is an eigenvector of $D$ .

Proof. The first variation is given by

$$\mathrm{d}{\mathcal{E}}_D({\mu }^{\ast };V)=-e^{z\cdot Dz}{P}_{z}Dz\cdot V(z).$$

Since $V(z)$ is an arbitrary vector, ${\mu }^{\ast }$ is a stationary point if and only if

$$0=P_z^{⟂}Dz=Dz-(z^{T}Dz)z,$$

which holds if and only if $z$ is an eigenvector of $D$ .∎

Intuitively speaking, $P_z^{⟂}Dz=0$ means that the force emerging from the interaction of a particle located at eigenvector $z$ with itself is orthogonal to the tangent space of $\mathcal{S}$ at point $z$ and is thus cancelled out by the projection. The same effect can be observed for convex combinations of a Dirac measure and its reflection.

Lemma 4.6. For any $t\in [\mathrm{0,1}]$ , we have that $t{\delta }_z+(1-t){\delta }_{-z}$ is a stationary point of ${\mathcal{E}}_D$ if and only if $z$ is an eigenvector of $D$ .

Proof. Using the expression in lemma 4.1, we obtain for any Lipschitz continuous $V$ , using the abbreviation $\iota =z\cdot Dz$ , that

$${\displaystyle \begin{array}{rl}\mathrm{d}{\mathcal{E}}_D(t{\delta }_z+(1-t){\delta }_{-z};V)=& t^{2}e{}^{\iota }{P}_z^{\perp }DzV(z)+(1-t{)}^{2}e{}^{\iota }{P}_{-z}^{\perp }D(-z)V(-z)\\ & +t(1-t)e{}^{-\iota }{P}_{-z}^{\perp }DzV(z)+t(1-t)e{}^{-\iota }{P}_z^{\perp }D(-z)V(-z).\end{array}}$$

We first observe that for any $x,y$ one has that $P_x^{⟂}y=P_{-x}^{⟂}y=-P_x^{⟂}(-y)$ . By comparing the coefficients in the above equation, we obtain that

$${\displaystyle \begin{array}{rl}\mathrm{d}{\mathcal{E}}_D(t){\delta }_z& +(1-t){\delta }_{-z};V)=0\text{ for all }V\text{ Lipschitz }\iff P_z^{\perp }Dz=0\\ & \iff Dz-(z\cdot Dz)z=0\iff z\text{ is an eigenvector. }\end{array}}$$

∎

For the symmetric case $t=1/2$ in the above lemma, we can further show that any convex combination of such stationary points is again a stationary point.

Lemma 4.7. Let $Z_D$ be a finite subset of eigenvectors of $D$ such that $w\cdot z=0$ for all $z\in Z_D\backslash \{w\}$ . Then for any choice of parameters $t:Z_D\to {ℝ}_0^{+}$ such that ${\displaystyle {\sum }_{z\in Z_D}t(z)=1}$ the following measure is a stationary point of ${\mathcal{E}}_D$ :

$${\displaystyle {\displaystyle \mu =\frac{1}{2}{\sum }_{z\in Z_D}t(z)\left({\delta }_z+{\delta }_{-z}\right).}}$$

Proof. We prove the statement by showing that equation (4.5) holds. For any $w\in Z_D$ , it holds that

$$\begin{array}{r}{P}_w^{⟂}Dw=-P_w^{⟂}Dw=0,\end{array}$$

since $Z_D$ only contains eigenvectors of $D$ . On the other hand, since we also require $w\cdot z=0$ for all $z\in Z_D\backslash \{w\}$ it follows that $z\cdot Dw=-z\cdot Dw=0$ and therefore,

$$\begin{array}{r}{e}^{w\cdot Dz}=e^{-w\cdot Dz}\end{array}$$

for all $z\in Z_D\backslash \{w\}$ . In total, this yields

$${\displaystyle {\displaystyle {\int }_{\mathcal{S}}{e}^{w\cdot Dy}{P}_x^{\perp }Dy\mathrm{d}\mu (y)=\sum _{z\in Z_D}t(z)(e^{w\cdot Dz}-e^{-w\cdot Dz})P_w^{\perp }(Dz)=0,}}$$

for all $w\in Z_D$ and thus also for $\mu$ -almost all $w\in \mathcal{S}$ .∎

The above proof strategy works only for Dirac measures aligned with the eigenvectors of $D$ . However, there exist other discrete measures that are stationary points, as the following example shows. For the sake of simplicity, we restrict ourselves to the two-dimensional case with a positive definite matrix $D$ and a symmetric combination of four Dirac measures. We further assume that $D$ is diagonal; the case of a general symmetric $D$ can be treated similarly with a rotation argument.

Lemma 4.8. Let $n=2$ , $\phi \in [\mathrm{0,2}\pi )$ and $D$ be diagonal and positive definite. A discrete measure:

$${\displaystyle {\displaystyle {\mu }_{\phi }=\frac{1}{\left|X_{\phi }\right|}{\sum }_{x\in X_{\phi }}{\delta }_x,where{X}_{\phi }=\{X(\phi ),X(\pi -\phi ),X(\pi +\phi ),X(2\pi -\phi )\},}}$$ (4.6)

is a stationary point of ${\mathcal{E}}_D$ if and only if either $\phi \in \{0,\pi /2,\pi \}$ or

$$\begin{array}{lrr}& \frac{\tanh (\lambda {}_1{\cos }^2\phi )}{\tanh (\lambda {}_2{\sin }^2\phi )}=\frac{\lambda {}_2}{\lambda {}_1},& \end{array}$$ (4.7)

where $\lambda {}_1,\lambda {}_2$ denote the diagonal entries of $D$ . For any choice of $\lambda {}_1,\lambda {}_2>0$ , there exists exactly one $\phi \in (0,\pi /2)$ that fulfils the condition in equation (4.7).

Proof. Without loss of generality, we prove the statement for $\phi \in [0,\pi /2]$ , since otherwise it holds that $(\psi \mathrm{mod}2\pi )\in [0,\pi /2]$ for a $\psi \in \{\pi -\phi ,\pi +\phi ,2\pi -\phi \}$ , and thus ${\mu }_{\phi }={\mu }_{\psi }$ .

It follows directly from lemma 4.6 that ${\mu }_{\phi }$ is a stationary point if $\phi \in \{0,\pi /2\}$ . Therefore, it remains to show that ${\mu }_{\phi }$ is a stationary point if and only if equation (4.7) is fulfilled. This means that we have to see when there exists a Lipschitz continuous $V$ such that $\mathrm{d}{\mathcal{E}}_D({\mu }_{\phi },V)\ne 0$ .

We first fix $x\in \mathcal{S}$ and consider

$${\displaystyle \begin{array}{rl}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}{P}_x^{\perp }Dy\mathrm{d}\mu {}_{\phi }(y)=& \frac{1}{4}((e^{x\cdot DX(\phi )}-e^{-x\cdot DX(\phi )})P_x^{\perp }DX(\phi )\\ & +(e^{x\cdot DX(\pi -\phi )}-e^{-x\cdot DX(\pi -\phi )})P_x^{\perp }DX(\pi -\phi )).\end{array}}$$ (4.8)

Since $n=2$ , we can further write $P_x^{⟂}y=x^{⟂}\cdot yx{}^{⟂}$ , where $x^{⟂}=(-x_2,x_1{)}^T$ . We factor out $x^{⟂}$ to rewrite equation (4.8) as $E(x;{\mu }_{\phi })x{}^{⟂}$ with

$$\begin{array}{r}E(x;{\mu }_{\phi })=(1/2)(\sinh (x\cdot DX(\phi ))x{}^{⟂}\cdot DX(\phi )+\sinh (x\cdot DX(\pi -\phi ))x{}^{⟂}\cdot DX(\pi -\phi )).\end{array}$$

Lemma 4.1 now gives us that

$${\displaystyle {\displaystyle \mathrm{d}{\mathcal{E}}_D\left({\mu }_{\phi },V\right)=\sum _{x\in X_{\phi }}E\left(x;{\mu }_{\phi }\right)x^{\perp }\cdot V(x),}}$$

which can become zero for all admissible $V$ if and only if $E(x;{\mu }_{\phi })=0$ for all $x\in X_{\phi }$ . Due to the symmetry properties of our measures ${\mu }_{\phi }$ , it further holds that $E(x;{\mu }_{\phi })$ is constant on $X_{\phi }$ ; therefore, it suffices to consider $x=X(\phi )$ . Remembering that $X(\phi )=(\cos \phi ,\sin \phi {)}^T$ , we derive

$${\displaystyle \begin{array}{rl}2E(X(\phi );{\mu }_{\phi })=& \sinh (\lambda {}_1\cos {}^2\phi +\lambda {}_2\sin {}^2\phi )(-\lambda {}_1+\lambda {}_2)\sin \phi \cos \phi \\ & +\sinh (-\lambda {}_1\cos {}^2\phi +\lambda {}_2\sin {}^2\phi )(\lambda {}_1+\lambda {}_2)\sin \phi \cos \phi .\end{array}}$$

Since $\phi \in (0,\pi /2)$ , the factor $\sin \phi \cos \phi$ cannot vanish, and the zeros of $E(X(\phi );{\mu }_{\phi })$ coincide with those of

$${\displaystyle \begin{array}{rl}& \sinh (\lambda {}_1\cos {}^2\phi +\lambda {}_2\sin {}^2\phi )(-\lambda {}_1+\lambda {}_2)+\sinh (-\lambda {}_1\cos {}^2\phi +\lambda {}_2\sin {}^2\phi )(\lambda {}_1+\lambda {}_2)\\ \text{(4.9)}& & =\sinh (\lambda {}_1+(-\lambda {}_1+\lambda {}_2)\sin {}^2\phi )(-\lambda {}_1+\lambda {}_2)+\sinh (-\lambda {}_1+(\lambda {}_1+\lambda {}_2)\sin {}^2\phi )(\lambda {}_1+\lambda {}_2).\end{array}}$$

This function obtains its minima at $(\phi \mathrm{mod}2\pi )\in \{0,\pi \}$ and its maxima at $(\phi \mathrm{mod}2\pi )\in \{\pi /2,3\pi /2\}$ and strictly increases or decreases, respectively, in between. Substituting these points into equation (4.9), we see that the minima are strictly negative and the maxima are strictly positive since $\lambda {}_1,\lambda {}_2>0$ . Therefore, there exists exactly one zero in the interval $(0,\pi /2)$ . Using the hyperbolic identity $\sinh (x+y)=\sinh x\cosh y+\cosh x\sinh y$ in equation (4.9), we arrive at the criterion in equation (4.7).∎

Remark 4.9. Importantly, the angle $\phi$ that fulfils equation (4.7) depends not only on the ratio of the eigenvalues of $D$ but also on their magnitude since they appear separately within the hyperbolic tangent.

Although the ratio of the eigenvalues does in general not determine the angle $\phi$ that fulfils equation (4.7), we can still make a qualitative prediction based on the ratio. The left-hand side of equation (4.7) decreases monotonically for $\phi \in [0,\pi /2)$ ; for $\lambda {}_1=\lambda {}_2$ , the condition is fulfilled for $\phi =\pi /4$ . Therefore, the condition is fulfilled by some $\phi \in [0,\pi /4)$ if $\lambda {}_2>\lambda {}_1$ and by some $\phi \in (\pi /4,\pi /2]$ if $\lambda {}_1>\lambda {}_2$ . The numerical experiments in §5b show that the measures characterized by equation (4.7) are not only stationary points but also minimizers among empirical measures consisting of at most four Dirac measures. In the remainder of this section, we aim to characterize minimizers for positive definite matrices $D$ in arbitrary dimensions $n⩾2$ .

(b). Energy variation at the uniform distribution

To characterize minimizers for positive definite $D$ , we start by identifying the cases when the uniform distribution is a stationary state. As we show in the following lemma, this can only be the case if the strength of the interaction does not depend on the direction, i.e. the eigenvalues of $D$ all have the same absolute value.

Lemma 4.10. The uniform distribution $\mu =\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ is a stationary point of ${\mathcal{E}}_D$ if and only if all eigenvalues $(\lambda {}_i{)}_{i=1}^n$ of $D$ have the same absolute value, i.e. $\left|\lambda {}_i\right|=\lambda$ for some $\lambda \in ℝ$ .

Proof. To keep the notation simple, we treat here the case $n=2$ , leaving the general proof for $n>2$ to appendix C(a). Let us fix $x\in \mathcal{S}$ and determine $\phi \in [\mathrm{0,2}\pi )$ such that $Dx/\Vert Dx\Vert =(\cos \phi ,\sin \phi {)}^T$ . Consider the integral

$$\begin{array}{rl}{\int }_{\mathcal{S}}{e}^{x\cdot Dy}{P}_x^{⟂}Dy\mathrm{d}\mathcal{H}{}^2(y)& ={\int }_0^{2\pi }{e}^{\Vert Dx\Vert \cos (\psi -\phi )}{P}_x^{⟂}(D(\cos \psi ,\sin \psi {)}^T)\mathrm{d}\psi =(\ast ),\end{array}$$

which can be rewritten with a change of variables $\theta =\psi -\phi$ as follows (recall that $P_x^{⟂}=\mathrm{Id}-x{x}^T$ ):

$${\displaystyle \begin{array}{rl}\ast & ={\int }_0^{2\pi }{e}^{\Vert Dx\Vert \cos \theta }\left(\cos \theta \left(D^{2}x/\Vert Dx\Vert -\Vert Dx\Vert x\right)+\sin \theta (Dx/\Vert Dx\Vert {)}^{\perp }\right)\mathrm{d}\theta \\ & =\left(D^{2}x/\Vert Dx\Vert -\Vert Dx\Vert x\right)\underset{>0}{\underbrace{{\int }_0^{2\pi }{e}^{\Vert Dx\Vert \cos \theta }\cos \theta \mathrm{d}\theta }}+(Dx/\Vert Dx\Vert {)}^{\perp }\underset{=0}{\underbrace{{\int }_0^{2\pi }{e}^{\Vert Dx\Vert \cos \theta }\sin \theta \mathrm{d}\theta }}.\end{array}}$$

From the above derivations, we see that $(\ast )=0$ if and only if $x$ is an eigenvector of $D^2$ . This holds true for $\mu$ -almost all $x\in \mathcal{S}$ if and only if $\left|\lambda {}_1\right|=\left|\lambda {}_2\right|$ . This automatically yields $\mathrm{d}{\mathcal{E}}_D(\mu ,V)=0$ if $\left|\lambda {}_1\right|=\left|\lambda {}_2\right|$ . It remains to show that this is also a necessary condition.

Without loss of generality, we assume that $\left|\lambda {}_1\right|>\left|\lambda {}_2\right|$ , where $\lambda {}_1$ and $\lambda {}_2$ are the eigenvalues corresponding to the eigenvectors $z_1$ and $z_2$ , respectively. Then, $(D^{2}x/\Vert Dx\Vert -\Vert Dx\Vert x)\cdot z_2$ is strictly negative on the set

$$\begin{array}{r}A=\{x\in \mathcal{S}|(x\cdot z_1)\in (\left|\lambda {}_2/\lambda {}_1\right|,1),(x\cdot z_2)>0\}.\end{array}$$

Since $\mu (A)>0$ we can find a Lipschitz continuous $V$ such that $V\cdot z_1=0$ for $\mu$ -a.e. on $\mathcal{S}$ and

$${\displaystyle \begin{array}{r}V(x)\cdot z_2\{\begin{array}{ll}>0& \text{for a.e. }x\in A,\\ =0& \text{for a.e. }x\in \mathcal{S}\mathrm{\setminus }A.\end{array}\end{array}}$$

For all such $V$ it holds that $\mathrm{d}{\mathcal{E}}_D(\mu ,V)>0$ , which concludes the proof.∎

Since we already know that minimizers for $D$ with at least one negative eigenvalue are Dirac measures, we can conclude that the uniform distribution is only a minimizer for $D=\mathrm{Id}$ .

Corollary 4.11. The uniform distribution $\mu =\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ minimizes ${\mathcal{E}}_D$ if and only if $D=\lambda \mathrm{Id}$ for ${\displaystyle \lambda ⩾0}$ .

Proof. We only need to show that there are no other matrices $D$ such that ${\mathcal{E}}_D$ is minimized by $\mu$ ; the other direction has been treated in proposition 3.11. The measure $\mu$ can only be a minimizer if it is a stationary point. By lemma 4.10, this implies that all eigenvalues of $D$ have to have the same absolute value. If such $D$ has at least one negative eigenvalue, it is also the smallest eigenvalue. Thus, by theorem 3.1, the only minimizers are Dirac deltas placed at eigenvectors corresponding to the negative eigenvalue.∎

(c). Perturbation of the identity

It is not clear whether an explicit computation of stationary points for an arbitrary positive definite matrix $D$ with at least two distinct eigenvalues is possible, but some insight can be gained with asymptotic analysis. We consider the following perturbed energy:

$$\begin{array}{r}\mathcal{E}{}_{\epsilon }(\mu ):={\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot (\mathrm{Id}+\epsilon M)y}\mathrm{d}\mu (x)\mathrm{d}\mu (y),\end{array}$$

where $M$ is a diagonal matrix and $\left|\epsilon \right|\ll 1$ is a small parameter. Using the second-order Taylor expansion of the exponential function, we can write

$$\begin{array}{lcr}& {\mathcal{E}}_{\epsilon }(\mu )\approx {\mathcal{E}}_D(\mu )+\epsilon {\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}x\cdot My\mathrm{d}\mu (x)\mathrm{d}\mu (y)+{\epsilon }^2{\int }_{\mathcal{S}}{e}^{x\cdot y}(x\cdot My{)}^2\mathrm{d}\mu (x)\mathrm{d}\mu (y).& \text{(4.10)}\end{array}$$

For $\epsilon =0$ we know that the unique minimizer ${\mu }_0$ is the uniform distribution on the sphere. Therefore, we use the following second-order asymptotic ansatz:

$$\begin{array}{lcr}& {\mu }_{\epsilon }:={\mu }_0+\epsilon \nu +{\epsilon }^{2}w,{\int }_{\mathcal{S}}\mathrm{d}\nu ={\int }_{\mathcal{S}}\mathrm{d}w=0.& \end{array}$$ (4.11)

We stress that here we consider the energy as a function on the space of signed Radon measures on the sphere $\mathcal{M}(\mathcal{S})$ with the total variation norm and not on the space of probability measures $\mathcal{P}(\mathcal{S})$ with the Wasserstein metric as in §4a. For this reason, the perturbation here is a measure and not a vector field (cf. equation (4.1)).

Substituting equation (4.11) into equation (4.10) and neglecting higher-order terms, we derive

$$\begin{array}{r}\mathcal{E}{}_{\epsilon }({\mu }_{\epsilon })-\mathcal{E}{}_{\epsilon }({\mu }_0)\approx \epsilon {\mathcal{E}}_D({\mu }_0,\nu )+{\epsilon }^2{\mathcal{E}}_D({\mu }_0,w)+{\epsilon }^2{\mathcal{E}}_D(\nu )+2{\epsilon }^2{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}x\cdot My\mathrm{d}\mu {}_0(x)\mathrm{d}\nu (y).\end{array}$$

Since further $y\mapsto {\int }_{\mathcal{S}}{e}^{x\cdot y}d{\mu }_0(x)$ is constant on $\mathcal{S}$ , it follows that

$${\displaystyle \begin{array}{r}{\mathcal{E}}_D({\mu }_0,\nu )=C(n){\int }_{\mathcal{S}}\mathrm{d}\nu =0\text{and}{\mathcal{E}}_D({\mu }_0,w)=C(n){\int }_{\mathcal{S}}\mathrm{d}w=0.\end{array}}$$

In particular, we see that the term ${\epsilon }^2\omega$ from equation (4.11) does not contribute to the second-order expansion of the energy. Therefore, minimizing $\mathcal{E}{}_{\epsilon }$ over all possible ${\mu }_{\epsilon }$ satisfying equation (4.11) is equivalent to minimizing

$$\begin{array}{r}{\stackrel{\mathrm{~}}{\mathcal{E}}}_{\epsilon }(\nu ):={\epsilon }^2({\mathcal{E}}_D(\nu )+2{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}x\cdot My\mathrm{d}\mu {}_0(x)\mathrm{d}\nu (y))\end{array}$$

over all signed measures $\nu$ with $\nu (\mathcal{S})=0$ . The first variation in the direction ${\nu }^{\prime }$ satisfying ${\int }_{\mathcal{S}}\mathrm{d}\nu {}^{\prime }=0$ is given by

$$\begin{array}{lrr}& \mathrm{d}{\stackrel{\mathrm{~}}{\mathcal{E}}}_{\epsilon }(\nu ,{\nu }^{\prime })=2{\epsilon }^2({\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}\mathrm{d}\nu (x)\mathrm{d}\nu {}^{\prime }(y)+{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}x\cdot My\mathrm{d}\mu {}_0(x)\mathrm{d}{\nu }^{\prime }(y)).& \end{array}$$ (4.12)

Our goal is now to find an optimal measure $\nu$ , such that its first variation vanishes in any direction ${\nu }^{\prime }$ such that ${\int }_{\mathcal{S}}\mathrm{d}\nu {}^{\prime }=0$ . To do so, we shall need the following two technical lemmas. To make the definition of the uniform distribution on the sphere rigorous, we denote by ${\mathcal{H}}^n$ the $n$ -dimensional Hausdorff measure and write ${\mathcal{S}}^{n-1}$ instead of $\mathcal{S}$ .

Lemma 4.12. Let $n⩾2$ and ${\mu }_0=\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ . It holds that

$${\displaystyle {\int }_{{\mathcal{S}}^{n-1}}{e}^{x\cdot y}x\mathrm{d}{\mu }_0(x)=C_{1}y}$$ (4.13)

for any $y\in {\mathcal{S}}^{n-1}$ , where the constant $C_1$ is positive and depends only on the dimension $n$ .

Proof. For the sake of simplicity, here we present the (more intuitive) proof for $n=2$ , leaving the general case $n>2$ to appendix C(b). We write $x=(\cos \phi ,\sin \phi {)}^T$ and $y=(\cos \psi ,\sin \psi {)}^T$ and derive that

$${\displaystyle \begin{array}{rl}2\pi {\int }_{\mathcal{S}}{e}^{x\cdot y}x\mathrm{d}\mu {}_0(x)& ={\int }_0^{2\pi }{e}^{\cos (\phi -\psi )}(\cos \phi ,\sin \phi {)}^T\mathrm{d}\phi ={\int }_0^{2\pi }{e}^{\cos \theta }(\cos (\psi +\theta ),\sin (\psi +\theta ){)}^T\mathrm{d}\theta \\ & =(\cos \psi ,\sin \psi {)}^T{\int }_0^{2\pi }{e}^{\cos \theta }\cos \theta \mathrm{d}\theta +(-\sin \psi ,\cos \psi {)}^T{\int }_0^{2\pi }{e}^{\cos \theta }\sin \theta \mathrm{d}\theta ,\end{array}}$$

where we use the coordinate transform $\theta =\phi -\psi$ and two trigonometric identities to separate the summands inside sine and cosine. Since ${\displaystyle {\int }_S{e}^{\cos \theta }\sin \theta \mathrm{d}\theta =0}$ , this yields equation (4.13) with

$${\displaystyle C_1=\frac{1}{2\pi }{\int }_0^{2\pi }{e}^{\cos \theta }\cos \theta \mathrm{d}\theta >0.}$$

∎

Lemma 4.13. Let $n⩾2$ and ${\mu }_0=\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ . It holds that for any $y\in {\mathcal{S}}^{n-1}$

$$\begin{array}{}\text{(4.14)}& {\displaystyle {\int }_{{\mathcal{S}}^{n-1}}{e}^{x\cdot y}{x}_i^2\mathrm{d}{\mu }_0(x)=C_2{y}_i^2+C_3,1⩽i⩽n,}\end{array}$$

where the constants $C_2$ and $C_3$ are positive and depend only on the dimension $n$ .

Proof. For the sake of simplicity, we again present the proof for $n=2$ ; the general case $n>2$ is treated in appendix C(c). Using the same arguments as in the previous proof, we derive

$${\displaystyle \begin{array}{rl}2\pi {\int }_{\mathcal{S}}{e}^{x\cdot y}{x}^2\mathrm{d}\mu {}_0(x)& ={\int }_0^{2\pi }{e}^{\cos \theta }({\cos }^2(\psi +\theta ),{\sin }^2(\psi +\theta ){)}^T\mathrm{d}\theta \\ & =({\cos }^2\psi ,{\sin }^2\psi {)}^T{\int }_0^{2\pi }{e}^{\cos \theta }{\cos }^2\theta \mathrm{d}\theta +({\sin }^2\psi ,{\cos }^2\psi {)}^T{\int }_0^{2\pi }{e}^{\cos \theta }{\sin }^2\theta \mathrm{d}\theta ,\end{array}}$$

where the mixed terms containing $\cos \theta \sin \theta$ vanish due to symmetry. Further, since ${\cos }^2\psi +{\sin }^2\psi =1$ , we can write

$$\begin{array}{r}({\sin }^2\psi ,{\cos }^2\psi {)}^T=(\mathrm{1,1}{)}^T-({\cos }^2\psi ,{\sin }^2\psi {)}^T.\end{array}$$

This yields equation (4.14) with positive constants:

$${\displaystyle \begin{array}{r}{C}_2=\frac{1}{2\pi }{\int }_0^{2\pi }{e}^{\cos \theta }({\cos }^2(\theta )-{\sin }^2(\theta ))\mathrm{d}\theta ,C_3=\frac{1}{2\pi }{\int }_0^{2\pi }{e}^{\cos \theta }{\sin }^2(\theta )\mathrm{d}\theta .\end{array}}$$

∎

Lemma 4.12 allows us to rewrite the second summand in equation (4.12) such that it contains $y\cdot My$ . Using lemma 4.13, we can then deduce that, up to constants, the measure $-(x\cdot Mx)\mu {}_0(x)$ is a stationary point of ${\stackrel{\mathrm{~}}{\mathcal{E}}}_{\epsilon }$ .

Theorem 4.14. The measure

$$\begin{array}{r}\mathrm{d}{\nu }^{\ast }(x)=(\alpha x\cdot Mx+\beta )\mathrm{d}{\mu }_0(x),\text{where \hspace{0.25em}\hspace{0.05em} }\alpha =-C_1/C_2\text{ \hspace{0.25em}\hspace{0.05em} and \hspace{0.25em}\hspace{0.05em} }\beta =-{\int }_{\mathcal{S}}\alpha x\cdot Mx\mathrm{d}{\mu }_0(x),\end{array}$$

fulfils ${\int }_{\mathcal{S}}\mathrm{d}{\nu }^{\ast }=0$ and $\mathrm{d}\mathcal{E}{}_{\epsilon }({\nu }^{\ast },{\nu }^{\prime })=0$ for all ${\nu }^{\prime }$ satisfying ${\int }_{\mathcal{S}}\mathrm{d}{\nu }^{\prime }=0$ .

Proof. From the definition of $\beta$ and ${\int }_{\mathcal{S}}\mathrm{d}{\mu }_0=1$ , it follows that ${\int }_{\mathcal{S}}\mathrm{d}{\nu }^{\ast }=0$ . With lemma 4.12, we write the optimality condition derived from equation (4.12) as

$$\begin{array}{r}{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}\mathrm{d}\nu (x)\mathrm{d}\omega (y)=-C_1{\int }_{\mathcal{S}}y\cdot My\mathrm{d}\omega (y).\end{array}$$

Substituting ${\nu }^{\ast }$ into the left-hand side and using lemma 4.13, we get

$${\displaystyle \begin{array}{rl}{\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}\mathrm{d}{\nu }^{\ast }(x)\mathrm{d}\omega (y)& ={\int }_{\mathcal{S}}\alpha \left(C_{2}y\cdot My+\mathrm{Tr}(M)C_3\right)\mathrm{d}\omega (y)+\beta {\int }_{\mathcal{S}}{\int }_{\mathcal{S}}{e}^{x\cdot y}\mathrm{d}{\mu }_0(x)\mathrm{d}\omega (y)\\ & =\alpha C_2{\int }_{\mathcal{S}}y\cdot My\mathrm{d}\omega (y),\end{array}}$$

where all terms that do not depend on $y$ , including ${\int }_{\mathcal{S}}{e}^{x\cdot y}\mathrm{d}\mu {}_0(x)$ , vanish due to ${\int }_{\mathcal{S}}\mathrm{d}\omega =0$ . Substituting $\alpha =-C_1/C_2$ completes the proof.∎

Theorem 4.14 gives us the following intuitive characterization. The measure ${\mu }_{\epsilon }$ that optimizes the perturbed energy is obtained by taking mass from the uniform distribution where $(x\cdot Mx)$ is large and adding it where $(x\cdot Mx)$ is small. In other words, we expect minimizers of the energy ${\mathcal{E}}_D$ with a positive definite matrix $D$ to have more mass in regions that correspond to small eigenvalues of $D$ than in regions that correspond to large ones. This intuition is in line with the results of the particle approximation in figure 3. Furthermore, in figure 5, we also observe that the density obtained in equation (4.11) with the measure ${\nu }^{\ast }$ from above can indeed be seen as a first-order approximation for small values of $\epsilon$ .

5. Numerical examples

To illustrate the obtained theoretical results, we perform a series of numerical experiments using a particle approximation of the energy from equation (1.2) with an ensemble of $N$ particles $X=(X_1,\dots ,X_N)$ ,

$${\displaystyle {\displaystyle {\mathcal{E}}_D\left({\mu }_N(X)\right),\text{ where }{\mu }_N(X)=\frac{1}{N}\sum _{i=1}^N{\delta }_{X_i}.}}$$

We consider the following particle flow, introduced in [9],

$${\displaystyle {\displaystyle {\dot{X}}_i(t)=P_{X_i(t)}^{\perp }\left(\pm \frac{1}{J_i(X)}\sum _{j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)}D{X}_j(t)\right),}}$$

with normalization factors $J_i(X)$ . If we choose the constant normalization

$$\begin{array}{lrr}& J_i(X)=N,& \end{array}$$ (5.1)

this corresponds merely to a step-size rescaling of a standard gradient descent scheme for ${\mathcal{E}}_D$ , which is called the (USA) flow in [9]. Choosing the normalization as the partition function

$${\displaystyle {\displaystyle \begin{array}{r}{J}_i(X)=\sum _{j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)},\end{array}}}$$ (5.2)

corresponds more closely to the self-attention dynamics and is labelled the SA flow in [9]. In what follows, we mostly use the normalization in equation (5.2), highlighting minor differences between the two formulations as appropriate. We use the explicit Euler discretization from equation (1.5) with step size $\tau >0$ to obtain the following update:

$${\displaystyle {\displaystyle \begin{array}{r}{X}_i(t+\tau )=\mathrm{\Pi }\left(X_i(t)\pm \frac{\tau }{J_i(X)}\sum _{j=1}^N{e}^{X_i(t)\cdot D{X}_j(t)}D{X}_j(t)\right).\end{array}}}$$ (5.3)

Remark 5.1. For $N=1$ and this scheme reduces to the following power iteration in the limit $\tau \to \infty$ :

$$\begin{array}{r}{X}_1(t+\tau )=\mathit{\Pi }(D{X}_1(t)).\end{array}$$

In this regard, the iteration in equation (5.3) can be seen as a method for approximating the largest eigenvalue and the corresponding eigenvector. We leave further analysis of this connection to future work.

The source code for the experiments here is available at https://github.com/TimRoith/TransformerDynamics and uses Python [61], mainly building upon the packages NumPy [62], SciPy [63] and PyTorch [64].

(a). Maximizers for positive definite matrices

To validate our results on maximizers, we first consider a simple set-up of a one-particle system, $N=1$ . We choose $\tau =0.075$ and run the scheme in equation (5.3) for $1500$ iterations. We only report the results for the adaptive normalization from equation (5.2), those for the constant normalization from equation (5.1) being essentially the same. For $D=\mathrm{Id}$ , we know that every single Dirac is a maximizer, which is indeed observed in figure 1a. Here, each random initialization on the sphere leads to a different final state. In fact, in this case, there is no evolution at all, and the particle stays at its initial position. If $D$ is positive definite and has a strictly largest eigenvalue $\lambda {}_{\text{max}}$ , theorem 3.1 shows that only Diracs at eigenvectors $z_{\text{max}}$ corresponding to $\lambda {}_{\text{max}}$ are maximizers. This can be observed in figure 1b where the final state is either at $z_{\text{max}}$ or $-z_{\text{max}}$ .

Figure 1.

Discrete maximizers on the sphere for $N=1$ particles. The colour indicates the value of $x\cdot Dx$ at each point on the sphere. (a) For ${\displaystyle D=\mathrm{I}\mathrm{d}}$ every single Dirac is a maximizer. We show the results for 30 different initializations (b) For ${\displaystyle D=\mathrm{d}\mathrm{i}\mathrm{a}\mathrm{g}(1,3,4)}$ the final state is either (0, 0,1) or (0,0,−1).

For multiple particle systems with $N>1$ , lemma 4.6 suggests also that linear combinations of an eigenvector with its negative are stationary points. These linear combinations are not maximizers, but their basin of attraction depends on the eigenvalues of the matrix. In figure 2 (left), we plot the probability (i.e. the proportion of random initializations) of converging to a single cluster versus two clusters as function of the eigenvalues. We fix $\lambda {}_1=1$ and vary $\lambda {}_2$ between $1$ and $1.5$ . Note that, as discussed in lemma 4.8 and remark 4.9, the actual values of the eigenvalues matter; not just their ratio. For ${\displaystyle \lambda {}_2\approx 1}$ , the probability of converging to a single cluster is high, whereas for larger values $\lambda {}_2\gtrsim 1.4$ , most trajectories converge to two clusters. The results in figure 2 were obtained with the adaptive normalization from equation (5.2); however, we observed the same quantitative behaviour with the constant normalization from equation (5.1).

Figure 2.

We study the trajectories for a symmetric positive definite matrix $D=\text{diag}(1,\lambda {}_2)$ with $\lambda {}_2\in [\mathrm{1.,1.5}]$ and $100$ different initializations using $100$ particles. We evaluate the number of clusters at the final iteration with the $k$ -means implementation of the SciPy package [63]. The centre of each cluster is close to an eigenvector corresponding to an eigenvalue of maximal absolute value. For $\lambda {}_2\approx 1$ , the evolution converges to the optimal state with a single cluster (blue, solid), while for bigger values, it tends to get stuck in the suboptimal stationary state with two clusters (red, hatched) from lemma 4.6.

(b). Minimizers for positive (semi-)definite matrices

We now study discrete minimizers for positive definite matrices. In figure 3, we show how the matrix $D$ influences the particle configuration to which the scheme in equation (5.3) converges. Here, too, we used the adaptive normalization from equation (5.2); the results for the constant one from equation (5.1) are largely the same.

Figure 3.

Final states for the minimization scheme after 10 000 steps with $N=400$ particles. The colour indicates the value of $x\cdot Dx$ at each point on the sphere. In (a), the uniform distribution is the minimizer of the energy. In (b), the particles do not form clusters at single Diracs but rather follow a smooth distribution on the sphere. In (c), any configuration with $(X_i{)}_1=(X_i{)}_3=0$ for all $i$ is a minimizer. In (d), any configuration with $(X_i{)}_3=0$ for all $i$ is a minimizer.

Furthermore, in figure 4, we illustrate the results of lemma 4.8 for matrices $D=\text{diag}(1,\lambda {}_2)$ with varying values $\lambda {}_2\in [\mathrm{0.5,8}]$ . We initialize $N=4$ particles as

Figure 4.

We consider minimizers for the matrix $D=\text{diag}(1,\lambda {}_2)$ . Starting with the initial configuration described in equation (5.4) , we compute the mean of $\tanh ({\cos }^2{\phi }_i)/\tanh (\lambda {}_2{\sin }^2{\phi }_i)$ over all particles. For a small step size, the resulting curve is very close to the identity, as predicted by lemma 4.8. If $\lambda {}_2\tau$ is too big, the dynamics converge to a suboptimal stationary point. We also compare the normalizations given by equations (5.1) and (5.2). We see that with the same step size $\tau =0.2$ , the adaptive normalization in equation (5.1) yields faster convergence than the constant one in equation (5.2).

$${\displaystyle \begin{array}{}\text{(5.4)}& X_i=X({\phi }_i)\text{with}{\phi }_i=(i-1)\cdot \pi +\pi /4\text{for}i=1,\dots ,4,\end{array}}$$

and let the scheme in equation (5.3) run for 10 000 iterations. From the final particle state, we compute the value $\tanh ({\cos }^2{\phi }_i)/\tanh (\lambda {}_2{\sin }^2{\phi }_i)$ for each particle separately; lemma 4.8 tells us that this should be equal to $\lambda {}_2$ for the minimizer. In figure 4, we observe that this holds true for the particle configurations computed with the discrete scheme. However, if the step size is too big compared to the value $\lambda {}_2$ , the system instead converges to the two-cluster stationary point from figure 2. Here, we notice a slight difference between the two normalizations. The adaptive normalization from equation (5.2) allows choosing bigger step sizes compared to the constant normalization from equation (5.1), enabling faster convergence to the large-time limit.

We further investigate the validity of the asymptotic solution from theorem 4.14 in the two-dimensional case. Here, we deviate from the particle approximation and instead discretize the interval $[-\pi ,\pi )$ with $N$ equidistant grid points $\mathit{\Theta }\in [-\pi ,\pi ]{}^N$ and the associated points on the sphere $x_1,\dots ,x_N\in {\mathcal{S}}^1$ . In this setting, we then aim to minimize

$$\begin{array}{}\text{(5.5)}& {\displaystyle {\displaystyle {\tilde{E}}_{\epsilon }(m)=\sum _{i,j=1}^N{e}^{x_i\cdot (\mathrm{Id}+\epsilon M)x_j}{m}_i\cdot m_j,}}\end{array}$$

where $m\in {ℝ}^N$ is a probability vector. Note that already, for $n=3$ , a more sophisticated quadrature rule would be required, e.g. the Lebedev quadrature on the sphere [65]. To deal with the simplex constraint for the vector $m$ , we use exponentiated gradient descent, specifically mirror descent with the negative log-entropy as the distance generating function [66], which yields the update

$$\begin{array}{lcr}& m(\epsilon {)}_i\leftarrow \frac{m_i{e}^{-\tau \nabla {\stackrel{\mathrm{~}}{E}}_{\epsilon }(m(\epsilon ){)}_i}}{{\sum }_{j=1}^{N}m(\epsilon {)}_j{e}^{-\tau \nabla \stackrel{\mathrm{~}}{E}(m(\epsilon ){)}_j}}=\mathrm{SoftMax}(\log (m(\epsilon ))-\tau \nabla {\stackrel{\mathrm{~}}{E}}_{\epsilon }(m(\epsilon ){)}_i.& \text{(5.6)}\end{array}$$

We take the perturbation matrix as $M=\text{diag}(\mathrm{0,1})$ , that is, the perturbed matrix $D$ is given by $D_{\epsilon }=\text{diag}(\mathrm{1,1}+\epsilon ).$ Recall the asymptotic expansion in equation (4.10). As noted in §4c, the contribution of the term ${\epsilon }^2\omega$ vanishes in the second-order expansion of the energy, and we are left with a solution:

$$\begin{array}{lcr}& {\mu }_{\epsilon }^{\ast }={\mu }_0+\epsilon {\nu }^{\ast },& \end{array}$$ (5.7)

where ${\nu }^{\ast }$ is as in theorem 4.14. We note that this measure has a Lebesgue density that can be evaluated at the grid points in $\mathit{\Theta }$ ; we denote the resulting vector by $\mathrm{d}{\mu }_{\epsilon }^{\ast }{|}_{\mathit{\Theta }}$ . In figure 5, we compare this solution to the vector $m(\epsilon )$ obtained by solving equations (5.5)–(5.6). The vector $m(\epsilon )$ for different values of $\epsilon$ is shown in figures 5a and 5b, we plot the ${\ell }^2$ error ${|m(ϵ)-d{\mu }_{\epsilon }^{\ast }{|}_{\mathit{\Theta }}|}_2$ .

Figure 5.

Numerical study of the asymptotic solution from theorem 4.14 in two dimensions. (a) The probability vectors ${\displaystyle m(ϵ)}$ computed using equation (5.5) with 500 steps for ${\displaystyle \tau =\mathrm{0.1.}}$ (b) The ${\displaystyle l^2}$ approximation error for the first-order expansion in equation (5.7) (blue, solid) and the conjectured form in equation (5.8) (green, dotted)

Beyond the first-order expansion in equation (5.7), we conjecture that $m(\epsilon )$ behaves as follows:

$$\begin{array}{lrr}& \mathrm{d}{\mu }_{\epsilon }^{\text{guess}}(\theta )\sim \exp (\mathit{{\rm Y}}(\epsilon )\cos (2\theta )),& \end{array}$$ (5.8)

where $\mathit{{\rm Y}}(\epsilon )$ is a function to be determined. Taking a second-order Taylor expansion $\mathit{{\rm Y}}(\epsilon )$ , we estimate the coefficients via linear regression with the given vectors $m(\epsilon )$ as data points and obtain $\mathit{{\rm Y}}(\epsilon )\approx 1/5\epsilon {}^2+e/2\epsilon$ . The ${\ell }^2$ error of this approximation is shown in figure 5b and is lower than that of the first-order expansion in equation (5.7). We leave the analysis of this ansatz to future work.

(c). Maximizers for negative definite and indefinite matrices

We proceed to numerical examples for §3d, i.e. maximization of the energy corresponding to a negative definite matrix. We take a system of $N=100$ particles and consider the two matrices from figure 1 multiplied by $-1$ . The results are shown in figure 6. We observe that a single final state consists of clusters at $\pm z$ , where $z$ is an eigenvector corresponding to the smallest eigenvalue, in agreement with theorem 3.17. As shown there, the behaviour does not change if one of the eigenvalues is zero, as only the eigenvectors corresponding to the smallest eigenvalue are relevant. For this reason, we do not consider the semi-definite case separately. The results here are not affected by the choice of the normalization; we only show the ones obtained with that in equation (5.2).

Figure 6.

Discrete maximizers on the sphere for negative definite matrices obtained with $N=100$ particles. We visualize the two-cluster final states by connecting the two components of each cluster corresponding to the same run with a line, assigning different colours to the two opposite clusters. The colour of the sphere indicates the value of $x\cdot Dx$ at each point on the sphere. (a) For D = −Id a single final state has clusters at both ${\displaystyle z}$ and ${\displaystyle -z}$ for any ${\displaystyle z\in S}$ . For clarity, we only show results for 6 different initializations. (b) For D = −diag(1,3,4) a single final state has clusters both at (0,0,1) and (0,0,−1). We show the results for 100 different initializations.

Finally, we turn to the case of indefinite matrices. As noted in remark 3.5, for a matrix $D$ that is not negative definite, a Dirac delta placed at the eigenvector corresponding to the largest eigenvalue may not be a maximizer. This can be observed numerically as shown in figure 7 where we plot the energies of one- and two-cluster states for $D=\text{diag}(-1,\lambda {}_2)$ with $\lambda {}_2\in [-\mathrm{1,1}]$ .

Figure 7.

Energies of the states $X^{\text{single}}=((\mathrm{0,1}))$ in blue, $X^{\text{two},1}=((\mathrm{0,1}),(0,-1))$ in red and $X^{\text{two},2}=((\mathrm{1,0}),(-\mathrm{1,0}))$ in green for the matrix $D=\text{diag}(-1,\lambda {}_2)$ with varying values of $\lambda {}_2$ .

6. Conclusion

In this work, we studied a mathematical model of self-attention layers used in the transformer architecture. Building upon [9], we analysed a continuum limit in the space of probability measures on a sphere. To understand the underlying geometry, we studied a new optimal transport distance $W_{m,2}$ with non-local mobility. We proved that the space of probability measures with this distance is a geodesic space and characterized absolutely continuous curves in this space. This allowed us to interpret the continuity equation (2.5) as curves of maximal slope of the interaction energy and to analyse the large-time behaviour using the energy dissipation property, showing that the dynamics converge to a stationary point of the interaction energy.

We analysed these critical points (in particular, minimizers and maximizers) for various types of interactions determined by the matrix $D$ in equation (1.2). These results are listed in table 1. We find that the positions of stationary points are strongly connected to normalized eigenvectors of $D$ , which form a strict subset of $\mathcal{S}$ in the case $D\ne \lambda \mathrm{Id}$ . In other words, the regions where clusters appear do not only depend on the initial configuration, but also on the interaction matrix itself. This could be related to mode collapse often observed in practice. It is an interesting question to understand whether an alternative, rotation-invariant architecture could prevent mode collapse.

Table 1.

Summary of results on minimizers/maximizers of the interaction energy in equation (1.2). We denote by $z_{\min }$ and $z_{\max }$ the eigenvectors that correspond to the smallest, respectively largest, eigenvalue of $D$ .

property of $D$	minimizers	maximizers
top rule positive definite	symmetric w.r.t. all eigenvectors (corollary 3.10 and §5b)	$\mu ={\delta }_{z_{\max }}$ (theorem 3.1 and §5a)
mid-rule positive semi-definite	any $\mu$ concentrated on $\mathcal{N}(D)$ (theorem 3.4 and §5b)	$\mu ={\delta }_{z_{\max }}$ (theorem 3.1 )
negative (semi-)definite	$\mu ={\delta }_{z_{\min }}$ (theorem 3.1)	$\mu =1/2({\delta }_{z_{\min }}+{\delta }_{-z_{\min }})$ (corollary 3.16 and §5c)
indefinite	$\mu ={\delta }_{z_{\min }}$ (theorem 3.4)	$\left|{\lambda }_{\max }\right|$ maximal: $\mu ={\delta }_{z_{\max }}$ (theorem 3.1 and §5c)

Several further questions remain open for future work: as already discussed, it would be interesting to study the optimal transport distance for mobilities $m_{\mu }$ that cannot be bounded from below, which is the case, for example, in problems of opinion dynamics where the Gaussian kernel on the Euclidean space is often used. In this case, the metric $W_{m,2}$ is no longer equivalent to $W_2$ . So far, we have only shown that equation (2.6) represents gradient flows in $(\mathcal{P}(M),W_{m,2})$ using the concept of curves of maximal slope. We do not know if these curves satisfy the slightly stronger energy variational inequality, which would yield an easy stability estimate for solutions of equation (2.6).

From a practical point of view, an even more interesting direction is studying more general flows in $W_{m,2}$ that correspond to non-symmetrical matrices $D$ in equation (1.2), which is common in transformer architectures. As mentioned above, basic properties of the distance carry over to the non-symmetric case, but characterizing the stationary states is non-trivial; one possibility is splitting the effective velocity fields into a dissipative and a (generalized) divergence-free part, similar to non-symmetric Fokker–Planck equations.

Finally, to justify the use of the continuum limit for studying the practical behaviour of transformers, one needs to establish convergence of discrete time-stepping in arbitrary time intervals. Moreover, it is worth studying how the step size influences the behaviour of the system and what effect weight-sharing would have.

Appendix A. Proofs of Section 2

A.1. Continuity equation on manifolds

Let $M$ be a compact, $n$ -dimensional Riemannian manifold and $TM={\bigsqcup }_{x\in M}{T}_{x}M$ its tangent bundle. Although $TM$ is not a vector space, the tangent bundle $TM$ itself can be considered as $2n$ -dimensional Riemannian manifold. For its proper definition and the topology on $TM$ , we refer to [67, ch. 3 (The Tangent Bundle)]. Velocity fields on manifolds are maps $V:M\to TM$ such that $\pi \circ V=I{d}_M$ , where $\pi :TM\to M$ is the projection map sending each vector in $T_{x}M$ to $x$ . We shall regularly commit the mild crime of interpreting $V(x)$ as an element in $T_{x}M$ instead of $TM$ . Let $I=(0,T)$ be an open interval, ${({\mu }_t)}_{t\in I}$ be a Borel family of probability measures on $M$ and $v:(x,t)\in M\times I\mapsto v_t(x)\in T_{x}M$ be a time-dependent Borel velocity field such that

$$\begin{array}{lrr}& \int \int |v_t(x)|\mathrm{d}{\mu }_t\mathrm{d}t<\infty ,& \end{array}$$ (A 1)

where $|\cdot |:T{M}_x\to [0,+\infty )$ denotes the norm induced by the inner product of the Riemannian structure. The continuity equation holds in the sense of distributions if

$$\begin{array}{lrr}& {\int }_{(0,T)}{\int }_M{\partial }_t\phi (x,t)+⟨\mathcal{D}\phi (x,t),v_t(x)⟩\mathrm{d}{\mu }_t\mathrm{d}t=0\forall \phi \in C_c^1(M\times (0,T)).& \end{array}$$ (A 2)

Here, $\mathcal{D}\phi$ denotes the differential of the map $x\in M\mapsto \phi (t,x)$ for a fixed $t\in [0,T]$ .

Proposition A.1 (Properties). Solutions to the continuity equation have the following properties:

• — Continuous representative: Let ${\mu }_t$ be a Borel family of probability measures satisfying equation (A 2) for a Borel vector field $v_t$ satisfying equation (A 1) . Then there exists a narrowly continuous curve $t\in [0,T]\to {\stackrel{\mathrm{~}}{\mu }}_t\in \mathcal{P}(M)$ such that ${\mu }_t={\stackrel{\mathrm{~}}{\mu }}_t$ for a.e. $t\in (0,T)$ . Moreover, if $\phi \in C_c^1(M\times [0,T])$ and $s⩽r\in [0,T]$ we have [53, Lemma 8.1.2]:
  $$\begin{array}{lrr}& {\int }_M\phi (x,r)\mathrm{d}{\stackrel{\mathrm{~}}{\mu }}_r-{\int }_M\phi (x,s)\mathrm{d}{\stackrel{\mathrm{~}}{\mu }}_s={\int }_s^r{\int }_M{\partial }_t\phi +\mathcal{D}\phi (v_t)\mathrm{d}{\mu }_t\mathrm{d}t.& \end{array}$$ (A3)
• — Time rescaling: Let ${\displaystyle t:s\in [0,T^{\prime }]\to t(s)\in [0,T]}$ be a strictly increasing absolutely continuous map with absolutely continuous inverse $s≔t^{-1}$ . Then $({\mu }_t,v_t)$ is a distributional solution of the continuity equation if and only if [68, Lemma 8.1.3]

  ${\displaystyle \begin{array}{r}\hat{\mu }:=\mu \circ t,\hat{v}=t^{\prime }v\circ t\end{array}}$ is a distributional solution of the continuity equation on ${\displaystyle (0,T^{\mathrm{\prime }})}$ .

• — Gluing solutions: Let $\{{\mu }_t{\}}_{t\in [0,T_1]},\{{\nu }_t{\}}_{t\in [0,T_2]}$ be two narrowly continuous curves in $\mathcal{P}(M)$ with ${\mu }_{T_1}={\nu }_0$ . Let further $\{v{\}}_{t\in [0,T_1]},\{w{\}}_{t\in [0,T_2]}$ be the corresponding Borel velocity fields such that equation (A 3) is satisfied. Then $\{{\eta }_t{\}}_{t\in [0,T_1+T_2]}$ and $\{u_t{\}}_{t\in [0,T_1+T_2]}$ defined by
  $${\displaystyle \begin{array}{r}{\eta }_t:=\{\begin{array}{ll}{\mu }_t& \text{if }t\in [0,T_1],\\ {\nu }_{t-T_1}& \text{if }t\in (T_1,T_1+T_2],\end{array}{u}_t:=\{\begin{array}{ll}{v}_t& \text{if }t\in [0,T_1],\\ w_{t-T_1}& \text{if }t\in (T_1,T_1+T_2],\end{array}\end{array}}$$

  satisfy equation (A 3) [ 69 , Lemma 4.4].

A.2. Proof of Theorem 2.2

We follow the proof strategy from [69] for the ‘flat’ Euclidean case, but since $TM$ is not a vector space, modifications are required. We start by establishing a compactness result for solutions of continuity equations with finite energy. For our purposes, we define the ‘lifted’ flux $J_t\in \mathcal{P}(TM\times M)$ in duality with $C_c(TM\times M)$ (see [70, Theorem 7.2]) by

$$\begin{array}{lrr}& {\int }_{TM\times M}\phi (w,y)\mathrm{d}{J}_t(w,y)={\int }_M{\int }_M\phi (v_t(x),y)\mathrm{d}{\mu }_t(x)\mathrm{d}{\mu }_t(y)\forall \phi \in C_c(TM\times M).& \end{array}$$ (A 4)

Notably, $({\mu }_t,J_t)$ solve the continuity equation in the sense that for all $s⩽r\in [0,T]$ :

$$\begin{array}{lrr}& {\int }_M{\phi }_r\mathrm{d}{\mu }_r-{\int }_M{\phi }_s\mathrm{d}{\mu }_s={\int }_s^r{\int }_M{\partial }_t\phi \mathrm{d}{\mu }_t\mathrm{d}t+{\int }_s^r{\int }_{TM\times M}\stackrel{\mathrm{~}}{\mathcal{D}\phi }\mathrm{d}{J}_t\mathrm{d}t\forall \phi \in C^1(M\times [0,T]),& \end{array}$$ (A 5)

where $\stackrel{\mathrm{~}}{\mathcal{D}\phi }:(w,y)\mapsto ⟨\mathcal{D}\phi (\pi (w)),w⟩$ is the extension of $\mathcal{D}\phi$ on to $TM\times M$ that is constant along $y\in M$ . Further, we define $\mathit{J}\in \mathcal{P}(TM\times M\times [0,T])$ in duality with $C_c(TM\times M\times [0,T])$ by

$${\int }_{TM\times M\times (0,T)}\phi \mathrm{d}\mathit{J}={\int }_0^T{\int }_{TM\times M}\phi \mathrm{d}{J}_t\mathrm{d}t\forall \phi \in C_c(TM\times M\times [0,T]).$$

Lemma A.2. Let $({\mu }^n,v^n)$ be a sequence in $CE(0,T)$ with

$$\begin{array}{r}\underset{n}{\sup }\{{\int }_0^1{\int }_M{m}_{\mu }(x)|v_t^n(x){|}^2\mathrm{d}{\mu }_t^n(x)\mathrm{d}t\}<+\infty .\end{array}$$

Then there exists a subsequence and a couple $(\mu ,J)$ satisfying the continuity equation in the sense of equation (A 5) such that

$${\displaystyle \begin{array}{r}{\mu }_t^n\rightharpoonup {\mu }_t\mathrm{\forall }t\in [0,T]\text{and}{\mathit{J}}^n\rightharpoonup \mathit{J},\end{array}}$$

and for the map $g:(v,p)\in TM\times M\mapsto (\pi (v),p)$ one has

$$\begin{array}{lrr}& g_{\#}{J}_t={\mu }_t\otimes {\mu }_t\text{for a.e.}t\in (0,T).& \end{array}$$ (A 6)

Proof. Step 1 (Convergence of $\mathit{J}$ ):

The estimate

$$\begin{array}{r}\underset{n}{\sup }{\int }_0^T{\int }_{TM\times M}|w{|}^2\mathrm{d}{J}_t^n(w,x)\mathrm{d}t⩽\frac{1}{C}\underset{n}{\sup }{\int }_0^T{\int }_{TM}{m}_{{\mu }_t^n}|v_t^n{|}^2\mathrm{d}{\mu }_t^n\mathrm{d}t<\infty \end{array}$$

combined with the fact that $M$ is compact and [53, Remark 5.1.5] implies tightness of $\mathit{J}\in \mathcal{P}(TM\times M\times [0,T])$ . By disintegrating $\mathit{J}$ , we obtain a Borel family $J_t$ such that ${\displaystyle \mathrm{d}\mathit{J}=\mathrm{d}{J}_t\mathrm{d}t}$ . Since $M$ is compact ${\mu }_0^n$ is tight, and we extract a further subsequence such that ${\displaystyle {\mu }_0^n\rightharpoonup {\mu }_0}$ .

Step 2 (Convergence of ${\mu }_t$ ):

Consider a function $\phi \in C^1(M)$ and for $t\in [0,T]$ set $\mathit{\zeta }:(v,y,t)\in TM\times M\times [0,T]\mapsto {\chi }_{[0,t]}⟨\mathcal{D}\phi (\pi (v)),v⟩$ . Since the discontinuity set of $\mathit{\zeta }$ is concentrated on $N=TM\times M\times \{0,t\}$ and $\left|F\right|(N)=0$ , general convergence theorems (see, e.g. [3, Prop. 5.1.10]) imply

$${\displaystyle \begin{array}{rl}\underset{n\to \mathrm{\infty }}{lim}{\int }_0^t{\int }_{TM\times M}\tilde{\mathcal{D}}\phi \mathrm{d}{J}_t^n\mathrm{d}t& =\underset{n\to \mathrm{\infty }}{lim}{\int }_{TM\times M\times [0,T]}\mathit{\zeta }\mathrm{d}{\mathit{J}}^n\\ \text{(A 7)}& & ={\int }_{TM\times M\times [0,T]}\mathit{\zeta }\mathrm{d}\mathit{J}={\int }_0^t{\int }_{TM\times M}\tilde{\mathcal{D}}\phi \mathrm{d}{J}_t\mathrm{d}t.\end{array}}$$

Let us fix a $t\in (0,T]$ . Since $M$ is compact, ${\mu }_t^n$ is tight, and we can extract from any subsequence a further subsequence such that ${\mu }_t^n$ converges narrowly. Then by equations (A 5) and (A 7) and the fact that $C^1$ is dense in $C^0$ , we know that all subsequences have the same limit. Therefore, ${\mu }_t^n\rightharpoonup {\mu }_t\in \mathcal{P}(M)$ for a particular ${\mu }_t$ . By the previous calculations, we also immediately obtain that $(\mu ,J)$ satisfy the continuity equation in the sense of equation (A 5). To show equation (A 6), we observe that since $M$ is compact

$$\begin{array}{r}{g}_{\#}{J}_t^n={\mu }_t^n\otimes {\mu }_t^n\rightharpoonup {\mu }_t\otimes {\mu }_t\forall t\in [0,T].\end{array}$$

∎

Proof of Theorem 2.2. Step 1:

Let $({\mu }^n,v^n)\in CE(\mathrm{0,1})$ be a minimizing sequence of the functional in equation (2.3) for some ${\mu }_0,{\mu }_1$ . Then the conditions of lemma A.2 are met, and we obtain

$$\begin{array}{r}{\mu }_t^n\rightharpoonup {\mu }_t\in \mathcal{P}(M)\forall [0,T]\text{and}{\mathit{J}}^n\rightharpoonup \mathit{J}\in \mathcal{P}(TM\times M\times [0,T]),\end{array}$$

where the limit satisfies the continuity equation in the sense of equation (A 5). Equation (A 6) in particular implies that $\mathit{J}$ can be disintegrated in the following way:

$$\mathrm{d}\mathit{J}=\mathrm{d}{u}_{t,x}(v)\mathrm{d}{\mu }_t(x)\mathrm{d}{\mu }_t(y)\mathrm{d}t,$$

where $u_{t,x}(v)\in \mathcal{P}(T{M}_p={\pi }^{-1}(x))$ . Using [53, Lemma 5.1.7], we now show that for ${\displaystyle {\overline{u}}_{x,t}=\mathbf{M}\mathbf{e}\mathbf{a}\mathbf{n}(u_{t,x})={\int }_{T{M}_x}v\mathrm{d}{u}_{t,x}(v)}$ it holds that

$${\displaystyle \begin{array}{rl}{W}_{m,2}({\mu }_0,{\mu }_1{)}^2=& \underset{n\to \mathrm{\infty }}{lim}{\int }_0^1{\int }_M{m}_{{\mu }_t^n}(x)|v_t^n(x){|}^2\mathrm{d}{\mu }_t^n(x)\mathrm{d}t\\ =& \underset{n\to \mathrm{\infty }}{lim}{\int }_0^1{\int }_{TM\times M}K(\pi (v),y)|v{|}^2\mathrm{d}{J}_t^n(v,y)\mathrm{d}t={\int }_0^1{\int }_{TM\times M}K(\pi (v),y)|v{|}^2\mathrm{d}{J}_t(v,y)\mathrm{d}t\\ =& {\int }_0^1{\int }_M{\int }_{M}K(x,y)\mathrm{d}{\mu }_t(y){\int }_{T{M}_p}|v{|}^2\mathrm{d}{u}_{x,t}(v)\mathrm{d}{\mu }_t(x)\mathrm{d}t\\ ⩾& {\int }_0^1{\int }_M{m}_{{\mu }_t}(x){\int }_{T{M}_x}|v{|}^2\mathrm{d}{\delta }_{{\overline{u}}_{x,t}}(v)\mathrm{d}{\mu }_t(x)\mathrm{d}t={\int }_0^1{\int }_M{m}_{{\mu }_t}(x)|{\overline{u}}_{x,t}{|}^2\mathrm{d}{\mu }_t(x)\mathrm{d}t,\end{array}}$$

where in the last line, we used Jensen’s inequality. Since $\mathcal{D}\phi (x):T{M}_x\to ℝ$ is linear and $(\mu ,\mathit{J})$ satisfy equation (A 5), this implies that $(\mu ,v=({\overline{u}}_{x,t}{)}_{t\in [0,T]})\in CE(\mathrm{0,1})$ and for this couple the infimum in equation (2.3) is obtained.

Step 2:

Proposition A.1 and a linear time rescaling show that

$$\begin{array}{lrr}& W_{m,2}^2({\mu }_0,{\mu }_T)=\inf \left\{T{\int }_0^T{\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t\mathrm{d}t:({\mu }_t,v_t)\in CE(0,T;{\mu }_0\to {\mu }_T)\}.& \end{array}$$ (A 8)

We denote by ${\overline{W}}_{m,2}(\mu ,\nu )$ the infimum in equation (2.4) and show that indeed ${\overline{W}}_{m,2}(\mu ,\nu )=W_{m,2}(\mu ,\nu )$ . By Hölder’s inequality, we immediately obtain that ${\overline{W}}_{m,2}(\mu ,\nu )⩽W_{m,2}(\mu ,\nu )$ . To show the reverse, we follow the arguments of [69, Theorem 5.4] and define for $(\mu ,v)\in CE(0,T;\mu \to \nu )$ :

$$\begin{array}{r}{s}_{ϵ}(t)≔{\int }_0^t{(ϵ+{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}r\text{for}t\in [0,T].\end{array}$$

Then $s_{ϵ}$ is strictly increasing, $s_{ϵ}^{\prime }⩾ϵ$ and $s_{ϵ}(0,T)=(0,S_{ϵ})$ with $S_{ϵ}≔s_{ϵ}(T)$ , so that its inverse map $t_{ϵ}:[0,S_{ϵ}]\to [0,T]$ is well defined and Lipschitz continuous and

$${\displaystyle \begin{array}{r}{t}_{ϵ}^{\prime }\circ s_{ϵ}:={\left(ϵ+{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t\right)}^{-1/2}\text{for a.e. }t\in (0,T).\end{array}}$$

By proposition A.1, we have that for ${\displaystyle {\mu }^{ϵ}:=\mu \circ t_{ϵ}}$ , ${\displaystyle v^{ϵ}:=t_{ϵ}^{\prime }v\circ t_{ϵ}}$ the couple $({\mu }^{ϵ},v^{ϵ})\in CE(0,S_{ϵ};\mu ,\nu )$ and

$$\begin{array}{rl}{W}_{m,2}^2(\mu ,\nu )& ⩽S_{ϵ}{\int }_0^{S_{ϵ}}{\int }_M{m}_{{\mu }_t^{ϵ}}\left|v_t^{ϵ}\right|\mathrm{d}{\mu }_t^{ϵ}\mathrm{d}s\\ & =S_{ϵ}{\int }_0^T\frac{{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t}{ϵ+{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t}{(ϵ+{\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}t,\end{array}$$

with the last term being smaller or equal to $S_{ϵ}^2$ . Sending $ϵ\to 0$ , we obtain

$$\begin{array}{r}{W}_{m,2}(\mu ,\nu )={\int }_0^T{\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}t\text{for all}(\mu ,v)\in CE(0,T;\mu \to \nu )\end{array}$$

and hence, $W_{m,2}(\mu ,\nu )={\overline{W}}_{m,2}(\mu ,\nu )$ . This in particular implies that for every minimizer $(\mu ,v)\in CE(\mathrm{0,1};\mu \to \nu )$ of the functional in equation (2.3), the equality

$$\begin{array}{r}{\left({\int }_0^1{\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t\mathrm{d}t)}^{1/2}={\int }_0^1{\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}t\end{array}$$

holds, which is only the case when ${\int }_M{m}_{{\mu }_t}|v_t{|}^2\mathrm{d}{\mu }_t$ is constant for a.e. $t\in (0,T)$ , implying by a further time rescaling argument∎

$${\displaystyle W_{m,2}({\mu }_s,{\mu }_t)=|s-t|W_{m,2}({\mu }_0,{\mu }_1)\mathrm{\forall }0⩽s⩽t⩽1.}$$ (A 9)

A.3. Proof of Lemma 2.4

Proof of Lemma 2.4. If $(\mu ,v)\in CE(0,T)$ and ${\int }_0^T{\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}t<+\infty$ then by equation (2.4), we have

$$\begin{array}{r}{W}_m({\mu }_s,{\nu }_r)⩽{\int }_s^r{\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\mathrm{d}t\forall 0⩽s⩽r⩽T.\end{array}$$

On the other hand, if ${\mu }_t$ is an absolutely continuous curve, then by a standard reparametrization argument [53, Lemma 1.1.4], we may assume ${\mu }_t$ to be Lipschitz. For $N\in \mathbb{N}$ , we set the step size as $\tau =T{2}^{-N}$ and choose a family of constant-speed geodesics $({\mu }^{k,N},v^{k,N})\in CE((k-1)\tau ,k\tau ;{\mu }_{(k-1)\tau }\to {\mu }_{k\tau })$ , $k\in \{1,...,N\}$ such that for $t\in ((k-1)\tau ,k\tau )$

$${\displaystyle \tau {\int }_M{m}_{{\mu }_t}{\left|v_t\right|}^2\mathrm{d}{\mu }_t\stackrel{(A8)}{=}\frac{1}{\tau }{W}_m^2\left({\mu }_{(k-1)\tau },{\mu }_{k\tau }\right)⩽\frac{1}{\tau }{\left({\int }_{(k-1)\tau }^{k\tau }|\dot{\mu }|(t)\mathrm{d}t\right)}^2\stackrel{\text{ Hölder }}{⩽}{\int }_{(k-1)\tau }^{k\tau }|\dot{\mu }|(t{)}^2\mathrm{d}t.}$$

Gluing all geodesics together by proposition A.1, we obtain a curve $({\mu }^N,v^N)\in CE(\mathrm{0,1})$ . Lemma A.2 gives us a subsequence, still denoted by $N$ , and a couple $(\stackrel{\mathrm{~}}{\mu },\stackrel{\mathrm{~}}{v})\in CE(\mathrm{0,1})$ such that ${\displaystyle {\mu }_t^N\rightharpoonup {\tilde{\mu }}_t}$ and ${\displaystyle \mathit{J}\rightharpoonup \tilde{\mathit{J}}}$ . By construction, ${\stackrel{\mathrm{~}}{\mu }}_t$ and ${\mu }_t$ coincide on the dense (in $[0,T]$ ) set $\{0\}\cup \{\frac{T}{M}{2}^{-N}:M,N\in ℝ,M⩽N\}$ . Since both ${\stackrel{\mathrm{~}}{\mu }}_t$ and ${\mu }_t$ are narrowly continuous ${\stackrel{\mathrm{~}}{\mu }}_t={\mu }_t$ must hold. Again, equation (A 6) implies that $\mathit{J}$ can be disintegrated in the following way:

$${\displaystyle \mathrm{d}\mathit{J}=\mathrm{d}{u}_{t,x}(v)\mathrm{d}{\mu }_t(x)\mathrm{d}{\mu }_t(y)\mathrm{d}t,}$$

where $u_{t,x}(v)\in \mathcal{P}(T{M}_x={\pi }^{-1}(x))$ . Then $(\mu ,\stackrel{\mathrm{~}}{v})\in CE(0,T)$ with ${\stackrel{\mathrm{~}}{v}}_t≔{\int }_{T{M}_x}w\mathrm{d}{u}_{t,x}(w)$ and

$${\displaystyle \begin{array}{r}\begin{array}{rl}{\int }_0^T{\int }_M{m}_{{\mu }_t}|{\tilde{v}}_t{|}^2\mathrm{d}{\mu }_t\mathrm{d}t& \stackrel{\text{Jensen}}{⩽}{\int }_0^T{\int }_M{\int }_{M}K(x,y)\mathrm{d}{\mu }_t(y){\int }_{T{M}_p}|v{|}^2\mathrm{d}{u}_{x,t}(v)\mathrm{d}{\mu }_t(x)\mathrm{d}t\\ & ⩽{\int }_0^T{\int }_{TM\times M}K(\pi (v),y)|v{|}^2\mathrm{d}{J}_t(v,y)\mathrm{d}t\\ & ⩽\underset{n\to \mathrm{\infty }}{lim inf}{\int }_0^T{\int }_{TM\times M}K(\pi (v),y)|v{|}^2\mathrm{d}{J}_t^n(v,y)\mathrm{d}t\\ & =\underset{n\to \mathrm{\infty }}{lim inf}{\int }_0^T{\int }_M{m}_{{\mu }_t^n}|v_t^n{|}^2\mathrm{d}{\mu }_t\mathrm{d}t⩽{\int }_0^T|\dot{\mu }{|}^2(t)\mathrm{d}t.\end{array}\end{array}}$$ (A 10)

Since $(\mu ,\stackrel{\mathrm{~}}{v})\in CE(0,T)$ , we have that

$$\left|\dot{\mu }\right|(t)⩽{\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}\text{for a.e.}t\in (0,T).$$

Finally, for equation (A 10) to hold $\left|\dot{\mu }\right|(t)={\left({\int }_M{m}_{{\mu }_t}\right|v_t{|}^2\mathrm{d}{\mu }_t)}^{1/2}$ must hold for a.e. $t\in (0,T)$ ∎

A.4. Proof of Lemma 2.6

Proof of Lemma 2.6. From theorem 2.3, we know that the distances $W_2$ and $W_{m,2}$ are equivalent. Therefore, we can assume absolute continuity with respect to $W_2$ . Further, by a standard rescaling argument (e.g. [53, Lemma 1.1.4] or [53, Lemma 8.1.3]), it is enough to prove equation (2.7) for $1$ -Lipschitz curves (w.r.t. $W_2$ ), i.e. we only need to consider absolutely continuous curves $({\mu }_t,v_t)\in CE(\mathrm{0,1};\mu \to \nu )$ such that

$$\begin{array}{r}{\int }_M|v_t(x){|}^2\mathrm{d}{\mu }_t(x)=1\text{for a.e.}t\in (0,T).\end{array}$$

For convenience, we shall set ${\mu }_t={\mu }_0$ for $t⩽0$ and ${\mu }_t={\mu }_T$ for $t⩾T$ as well as $v_t=0$ for $t\notin [0,T]$ . We define the function $\eta :(x,t)\in M\times ℝ\mapsto \frac{1}{2}{\int }_{M}W(x,y)\mathrm{d}{\mu }_t(y)$ for which

$${\displaystyle {\mathrm{\partial }}_t\eta (t,x)=\{\begin{array}{ll}0& \text{if }t\notin [0,T],\\ \frac{1}{2}{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_t(y)⟩\mathrm{d}{\mu }_t(y)& \text{else},\end{array}}$$

in the distributional sense. Using the mollifier $g_{ϵ}$ , as described in [71, ch. C.5], one can smooth out $\eta$ in the time direction by setting

$$\begin{array}{r}{\eta }_{ϵ}(t,x)≔{\int }_{ℝ}\eta (\tau ,x)g_{ϵ}(t-\tau )\mathrm{d}\tau .\end{array}$$

By [71, ch. C.5, Theorem 7 (iii)], we have that ${\eta }_{ϵ}\to \eta$ pointwise, and with the use of the dominated convergence theorem with the upper bound $\left|{\eta }_{ϵ}\right|⩽{\sup }_{(x,y)\in M\times M}|W(x,y)|<\infty$ , we calculate

$$\begin{array}{r}\mathcal{E}({\mu }_T)-\mathcal{E}({\mu }_0)={\int }_M\eta \mathrm{d}{\mu }_T-{\int }_M\eta \mathrm{d}{\mu }_0=\underset{ϵ\to 0}{\lim }{\int }_M{\eta }_{ϵ}\mathrm{d}{\mu }_T-{\int }_M{\eta }_{ϵ}\mathrm{d}{\mu }_0.\end{array}$$

We further have that

$$\begin{array}{rl}+\infty & >\frac{1}{2}{\int }_0^T{\int }_M{\int }_M⟨{\mathcal{D}}_{y}W(y,x),v_t(y)⟩\mathrm{d}{\mu }_t(y)\mathrm{d}{\mu }_t(x)\mathrm{d}t\\ & \stackrel{(\ast \ast )}{=}\underset{ϵ\to 0}{\lim }\frac{1}{2}{\int }_0^T{\int }_{ℝ}{\int }_M{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_t⟩\mathrm{d}{\mu }_t(y)g_{ϵ}(t-\tau )(t)\mathrm{d}\tau \mathrm{d}{\mu }_t(x)\mathrm{d}t\\ & \stackrel{(\ast )}{=}-\underset{ϵ\to 0}{\lim }{\int }_0^T{\int }_M{\int }_{ℝ}\eta (\tau ,x){\partial }_{\tau }{g}_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\mu }_t(x)\mathrm{d}t\\ & =\underset{ϵ\to 0}{\lim }{\int }_0^T{\int }_M{\int }_{ℝ}\eta (\tau ,x){\partial }_t{g}_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\mu }_t(x)\mathrm{d}t=\underset{ϵ\to 0}{\lim }{\int }_0^T{\int }_M{\partial }_t{\eta }_{ϵ}(t,x)\mathrm{d}{\mu }_t(x)\mathrm{d}t,\end{array}$$

where for $(\ast )$ we use the definition of the distributional derivative and rearrange the integral using the Fubini–Tonelli theorem. To prove $(\ast \ast )$ , we need to define a piecewise constant approximation of ${\mu }_t$ . We fix a $N\in \mathbb{N}$ $\tau =\frac{T}{N}$ and set for $k\in \{1,N\}$

$$\begin{array}{r}{\overline{\mu }}_t≔{\mu }_{k\tau }\text{for}t\in [k\tau ,(k+1)\tau ),{\overline{\mu }}_T≔{\mu }_T.\end{array}$$

Since ${\mu }_t$ is $1$ -Lipschitz, we have $W_2({\mu }_t,{\overline{\mu }}_t)⩽\tau$ for all $t\in [0,T]$ . Then, we estimate

$${\displaystyle \begin{array}{rl}& |{\int }_0^T{\int }_M{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_t(y)⟩\mathrm{d}{\mu }_t(y)\mathrm{d}{\mu }_t(x)\mathrm{d}t-{\int }_0^T{\int }_M{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_t(y)⟩\mathrm{d}{\mu }_t(y)\mathrm{d}{\overline{\mu }}_t(x)\mathrm{d}t|\\ & ⩽{\int }_0^T{\int }_{M\times M}{\int }_M|⟨{\mathcal{D}}_{y}W(x_1,y),v_t(y)⟩-⟨{\mathcal{D}}_{y}W(x_2,y),v_t(y)⟩|\mathrm{d}{\mu }_t(y)\mathrm{d}{\pi }_t(x_1,x_2)\mathrm{d}t\\ & ⩽{\int }_0^T{\int }_{M\times M}{\int }_M|{\mathcal{D}}_{y}W(x_1,y)-{\mathcal{D}}_{y}W(x_2,y){|}_{\ast }|v_t(y)|\mathrm{d}{\mu }_t(y)\mathrm{d}{\pi }_t(x_1,x_2)\mathrm{d}t\\ & ⩽C{\int }_0^T{\int }_{M\times M}{\int }_M|x_1-x_2||v_t(y)|\mathrm{d}{\mu }_t(y)\mathrm{d}{\pi }_t(x_1,x_2)\mathrm{d}t\\ & =C{\int }_0^T({\int }_{M\times M}|x_1-x_2|\mathrm{d}{\pi }_t(x_1,x_2))\left({\int }_M|v_t(y)|\mathrm{d}{\mu }_t(y)\right)\mathrm{d}t\\ & ⩽C({\int }_0^T{W}_2^2({\mu }_t,{\overline{\mu }}_t)\mathrm{d}t{\int }_0^T{\int }_M|v_t(y){|}^2\mathrm{d}{\mu }_t(y)\mathrm{d}t{)}^{1/2}\\ \text{(A 11)}& & =C(T{\int }_0^T{W}_2^2({\mu }_t,{\overline{\mu }}_t)\mathrm{d}t{)}^{1/2}⩽C(T{\int }_0^T{\tau }^2\mathrm{d}t{)}^{1/2}=C\frac{T^2}{N},\end{array}}$$

where ${\pi }_t\in \mathcal{P}(M\times M)$ is the optimal transport plan between ${\mu }_t$ and ${\overline{\mu }}_t$ and $|\cdot {|}_{\ast }$ denotes the dual norm of $|\cdot |$ . (For more details on the static formulation of Wasserstein distances via optimal transport plans, we refer to [53, ch. 6]). We can argue similarly in the mollified case:

$${\displaystyle \begin{array}{rl}& |{\int }_0^T{\int }_M{\int }_{\mathbb{R}}{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_{\tau }(y)⟩\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\mu }_t(x)\mathrm{d}t-\\ & {\int }_0^T{\int }_M{\int }_{\mathbb{R}}{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_{\tau }(y)⟩\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\overline{\mu }}_t(x)\mathrm{d}t|\\ ⩽& {\int }_0^T{\int }_{M\times M}{\int }_{\mathbb{R}}{\int }_M|⟨{\mathcal{D}}_{y}W(x_1,y),v_{\tau }(y)⟩-\\ & ⟨{\mathcal{D}}_{y}W(x_2,y),v_{\tau }(y)⟩|\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\pi }_t(x_1,x_2)\mathrm{d}t\\ ⩽& {\int }_0^T{\int }_{M\times M}{\int }_{\mathbb{R}}{\int }_M|{\mathcal{D}}_{y}W(x_1,y)-{\mathcal{D}}_{y}W(x_2,y){|}_{\ast }|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\pi }_t(x_1,x_2)\mathrm{d}t\\ ⩽& C{\int }_0^T{\int }_{M\times M}{\int }_{\mathbb{R}}{\int }_M|x_1-x_2||v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\pi }_t(x_1,x_2)\mathrm{d}t\\ =& C{\int }_0^T({\int }_{M\times M}|x_1-x_2|\mathrm{d}{\pi }_t(x_1,x_2))\left({\int }_{\mathbb{R}}{\int }_M|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \right)\mathrm{d}t\\ ⩽& C{\int }_0^T{W}_2({\mu }_t,{\overline{\mu }}_t){\int }_{\mathbb{R}}{\int }_M|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}t\\ ⩽& C\frac{T}{N}{\int }_0^T{\int }_{\mathbb{R}}{\int }_M|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}t=C\frac{T}{N}{\int }_{\mathbb{R}}{\int }_M|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y){\int }_0^T{g}_{ϵ}(t-\tau )\mathrm{d}t\mathrm{d}\tau \\ \le & C\frac{T}{N}{\int }_{\mathbb{R}}{\int }_M|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)\mathrm{d}\tau =C\frac{T}{N}{\int }_0^T{\int }_M|v_{\tau }(y)|\mathrm{d}{\mu }_{\tau }(y)\mathrm{d}\tau \\ \text{(A 12)}& \le & C\frac{T}{N}{\int }_0^T({\int }_M|v_{\tau }(y){|}^2\mathrm{d}{\mu }_{\tau }(y){)}^{1/2}\mathrm{d}\tau ⩽C\frac{T^2}{N}.\end{array}}$$

We denote $\stackrel{\mathrm{~}}{C}={\sup }_{(x,y)\in M\times M}|{\mathcal{D}}_{y}W(x,y){|}_{\ast }<+\infty$ and combine equations (A 11) and (A 12) to estimate

$${\displaystyle {\displaystyle \begin{array}{rl}& |{\int }_0^T{\int }_M{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_t(y)⟩\mathrm{d}{\mu }_t(y)\mathrm{d}{\mu }_t(x)\mathrm{d}t\\ & -{\int }_0^T{\int }_M{\int }_{\mathbb{R}}{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_{\tau }(y)⟩\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}{\mu }_t(x)\mathrm{d}t|\\ & ⩽2C\frac{T^2}{N}+|{\int }_0^T{\int }_M\underset{:=f(x,t)}{\underbrace{{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_t(y)⟩\mathrm{d}{\mu }_t(y)}}\mathrm{d}{\overline{\mu }}_t(x)\mathrm{d}t-\\ & {\int }_0^T{\int }_M\underset{:=f_{ϵ}(x,t)}{\underbrace{{\int }_{\mathbb{R}}{\int }_M⟨{\mathcal{D}}_{y}W(x,y),v_{\tau }(y)⟩\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau }}\mathrm{d}{\overline{\mu }}_t(x)\mathrm{d}t|\\ & ⩽2C\frac{T^2}{N}+\sum _{i=1}^N{\int }_{(i-1)\tau +ϵ}^{i\tau -ϵ}{\int }_M\left|f-f_{ϵ}\right|\mathrm{d}{\overline{\mu }}_t\mathrm{d}t+{\int }_{(i-1)\tau }^{(i-1)\tau +ϵ}{\int }_M\left|f-f_{ϵ}\right|\mathrm{d}{\overline{\mu }}_t\mathrm{d}t+\\ & {\int }_{i\tau -ϵ}^{i\tau }{\int }_M\left|f-f_{ϵ}\right|\mathrm{d}{\overline{\mu }}_t\mathrm{d}t\\ & ⩽2C\frac{T^2}{N}+\sum _{i=1}^N{\int }_{(i-1)\tau +ϵ}^{i\tau -ϵ}{\int }_M\left|f-f_{ϵ}\right|\mathrm{d}{\overline{\mu }}_t\mathrm{d}t+{\int }_{(i-1)\tau }^{(i-1)\tau +ϵ}{\int }_{M}2\tilde{C}\mathrm{d}{\overline{\mu }}_t\mathrm{d}t+{\int }_{i\tau -ϵ}^{i\tau }{\int }_{M}2\tilde{C}\mathrm{d}{\overline{\mu }}_t\mathrm{d}t\\ & ⩽2C\frac{T^2}{N}+\sum _{i=1}^N{\int }_{(i-1)\tau +ϵ}^{i\tau -ϵ}{\int }_M\left|f-f_{ϵ}\right|\mathrm{d}{\overline{\mu }}_t\mathrm{d}t+4Nϵ\tilde{C}⩽\frac{\delta }{3}+\sum _{i=1}^N\frac{\delta }{3N}+\frac{\delta }{3},\end{array}}}$$

where, first, $N$ is chosen such that $N⩾6C\frac{T^2}{\delta }$ and, second, $ϵ$ such that $ϵ⩽\frac{\delta }{12N\stackrel{\mathrm{~}}{C}}$ and for each $i\in \{1,...,N\}$ , it holds ${\displaystyle {\int }_{(i-1)\tau +ϵ}^{i\tau -ϵ}{\int }_M|f-f_{ϵ}|\mathrm{d}{\overline{\mu }}_t\mathrm{d}t⩽\frac{\delta }{3N}}$ (by lemma A.3). Therefore $(\ast \ast )$ is proven.

Finally, by lemma A.4, we obtain $n_{ϵ}\in C^1(M\times [0,T])$ that we can use as a test function in equation (A 3) and send $ϵ\to 0$ to obtain

$\begin{array}{rl}\mathcal{E}\left({\mu }_T\right)-\mathcal{E}\left({\mu }_0\right)& ={\int }_M\eta \mathrm{d}{\mu }_T-{\int }_M\eta \mathrm{d}{\mu }_0={\int }_0^T{\int }_M{\mathrm{\partial }}_t\eta \mathrm{d}{\mu }_t+{\int }_M⟨\mathcal{D}\eta ,v_t⟩\mathrm{d}{\mu }_t\mathrm{d}t\\ & ={\int }_0^T{\int }_{M\times M}⟨{\mathcal{D}}_{x}W(x,y),v_t(x)⟩\mathrm{d}{\mu }_t(x)\mathrm{d}{\mu }_t(y)\mathrm{d}t.\end{array}$

∎

Lemma A.3. Let $f:M\times [0,T]\to ℝ$ be Borel measurable and $\mu \in \mathcal{P}(M)$ with

$${\int }_a^b{\int }_M\left|f\right|\mathrm{d}\mu \mathrm{d}t<\infty \text{for}0⩽a<b⩽T.$$

For

$${\mu }_a^b(A)≔\mu \otimes \mathcal{L}(a,b)$$

it holds

$$\Vert f_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}⩽\Vert f{\Vert }_{L^1({\mu }_a^b)}\text{and}{f}_{ϵ}\to f\text{in}{L}^1({\mu }_{a+ϵ}^{b-ϵ}).$$

Proof. We adapt [71, ch. 5, Theorem 7] to our case and start by showing

$${\displaystyle \begin{array}{rl}\Vert f_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}⩽& {\int }_M{\int }_{a+ϵ}^{b-ϵ}{\int }_a^b|f(x,\tau )|g_{ϵ}(t-\tau )\mathrm{d}\tau \mathrm{d}t\mathrm{d}\mu (x)\\ =& {\int }_M{\int }_a^b|f(x,\tau )|{\int }_{a+ϵ}^{b-ϵ}{g}_{ϵ}(t-\tau )\mathrm{d}t\mathrm{d}\tau \mathrm{d}\mu (x)={\int }_M{\int }_a^b|f(x,\tau )|\mathrm{d}\tau \mathrm{d}\mu (x)=\Vert f{\Vert }_{L^1({\mu }_a^b)}.\end{array}}$$

We approximate $f$ in $L^1({\mu }_a^b)$ by $\gamma \in C_c(M\times [a,b])$ (see [70, Proposition 7.9]) and calculate

$${\displaystyle \begin{array}{rl}\Vert f-f_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}⩽& \Vert f-\gamma {\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}+\Vert \gamma -{\gamma }_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}+\Vert {\gamma }_{ϵ}-f_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}\\ ⩽& 2\Vert f-\gamma {\Vert }_{L^1({\mu }_a^b)}+\Vert \gamma -{\gamma }_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}.\end{array}}$$

From [71, ch. C.5, Theorem 7], we know that ${\gamma }_{ϵ}\to \gamma$ for all $(x,t)\in M\times [a,b]$ because $\gamma$ is continuous. Choosing $\gamma$ such that $\Vert f-\gamma {\Vert }_{L^1({\mu }_a^b)}<\delta$ and using the dominated convergence theorem, we get ${\mathrm{lim\hspace{0.17em}sup}}_{ϵ\to 0}\Vert f-f_{ϵ}{\Vert }_{L^1({\mu }_{a+ϵ}^{b-ϵ})}⩽2\delta$ . As $\delta$ can be chosen arbitrarily small, we obtain convergence.∎

Lemma A.4. We have ${\eta }_{ϵ}\in C^1(M\times [0,T])$ .

Proof. Let $\gamma :V\subset {ℝ}^d\to U(x)$ be a smooth local chart for an open set $U(x)$ containing $x$ . Then, since $W(x,y)\in C^1(M\times M)$ , the function $z\mapsto {\int }_M{\partial }_{z_i}W(\gamma (z),y)\mathrm{d}{\mu }_{\tau }$ is continuous in $z$ and the product

$$(z,t)\mapsto {\int }_M{\partial }_{z_i}W(\gamma (z),y)\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )$$

is continuous on $V\times ℝ$ . Taking any sequence $(z_n,t_n)\to (z,t)$ , we can use the dominated convergence theorem to obtain

$${\displaystyle \begin{array}{rl}\underset{n\to \mathrm{\infty }}{lim}{\mathrm{\partial }}_t{\eta }_{ϵ}(\gamma (z_n),t_n)& =\underset{n\to \mathrm{\infty }}{lim}{\int }_{\mathbb{R}}{\int }_M{\mathrm{\partial }}_{z_i}W(\gamma (z_n),y)\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t_n-\tau )\mathrm{d}\tau \\ & ={\int }_{\mathbb{R}}{\int }_M{\mathrm{\partial }}_{z_i}W(\gamma (z),y)\mathrm{d}{\mu }_{\tau }(y)g_{ϵ}(t-\tau )\mathrm{d}\tau ={\mathrm{\partial }}_t{\eta }_{ϵ}(\gamma (z),t).\end{array}}$$

An upper bound is given by the function ${\sup }_{V\times W}{\partial }_{z_i}W(\gamma (z),y){\chi }_{[{\inf }_n{t}_n-ϵ,{\sup }_n{t}_n+ϵ]}(\tau )$ . Thus, ${\partial }_t{\eta }_{ϵ}(\gamma (z),t)$ is continuous in $V\times [0,T]$ . With the same argument, a similar statement can be shown for

$${\displaystyle \begin{array}{r}{\mathrm{\partial }}_t{\eta }_{ϵ}(x,t)={\int }_{\mathbb{R}}{\int }_{M}W(x,y)\mathrm{d}{\mu }_{\tau }(y){\mathrm{\partial }}_t{g}_{ϵ}(t-\tau )(t)\mathrm{d}\tau .\end{array}}$$

By [72, Theorem 2.8], it follows that ${\eta }_{ϵ}(t,\gamma (z))\in C^1(V\times [0,T])$ and since the local chart was chosen arbitrarily ${\eta }_{ϵ}\in C^1(M\times [0,T])$ .∎

Appendix B. Spherical coordinates

For many computations in §4, we use spherical coordinates. Up to small notational changes, we use the definition provided in [73]. We define the coordinate transform $X_n:\phi \in [0,\pi ]{}^{n-2}\times [\mathrm{0,2}\pi ]\to {\mathcal{S}}^{n-1}$ for $\phi \in [0,\pi ]{}^{n-2}\times [\mathrm{0,2}\pi ]$ as

$${\displaystyle {\displaystyle X_n(\phi )=\cos \left({\phi }_1\right)e_1+\sum _{i=2}^{n-1}\cos \left({\phi }_i\right)\prod _{j=1}^{i-1}\sin \left({\phi }_j\right)e_i+\prod _{j=i}^{n-1}\sin \left({\phi }_i\right)e_n.}}$$

Here and in the following, $e_i\in {ℝ}^n$ denotes the $i$ th standard basis vector.

The Jacobian determinant is given by

$$\begin{array}{r}J{X}_n(\phi )=\prod _{i=1}^{n-2}{\sin }^{n-1-i}({\phi }_i).\end{array}$$

To highlight the recursive character of $X_n$ with respect to $n$ , we further note that

$${\displaystyle \begin{array}{r}{X}_n(\phi {)}_{\hat{1}}=\sin ({\phi }_1)X{}_{n-1}({\phi }_{\hat{1}})\text{and}J{X}_n(\phi )={\sin }^{n-2}({\phi }_1)J{X}_{n-1}({\phi }_{\hat{1}}),\end{array}}$$

where the index $\hat{1}$ denotes that we drop the first element, i.e. for $\phi \in {ℝ}^{n-1}$ , ${\displaystyle {\phi }_{\hat{1}}={\sum }_{i=2}^{n-1}{\phi }_i{e}_{i-1}}$ . A practical consequence of this property is the recursive computation formula for the Hausdorff measure of the $n$ -dimensional sphere.

Lemma B.1. Denote $\left|{\mathcal{S}}^{n-1}\right|:={\mathcal{H}}^n({\mathcal{S}}^{n-1})$ . For $n⩾2$ , it holds that

$$\begin{array}{r}\left|{\mathcal{S}}^{n-1}\right|=\left|{\mathcal{S}}^{n-2}\right|{\int }_0^{\pi }{\sin }^{n-2}\phi \mathrm{d}\phi .\end{array}$$

Proof. For $n=2$ , the proof follows from a simple computation and the fact that $\left|{\mathcal{S}}^0\right|=2$ and $\left|{\mathcal{S}}^1\right|=2\pi$ . For $n>2$ , we have

$${\displaystyle \begin{array}{rl}|{\mathcal{S}}^{\mathcal{n}\mathcal{-}1}\mathcal{|}& ={\int }_{[0,\pi {]}^{n-2}\times [0,2\pi ]}J{X}_n(\phi )\mathrm{d}\phi ={\int }_0^{\pi }{\sin }^{n-2}{\phi }_{1}J{X}_{n-1}({\phi }_{\hat{1}})\mathrm{d}\phi \\ & ={\int }_0^{\pi }{\sin }^{n-2}\phi \mathrm{d}\phi {\int }_{[0,\pi {]}^{n-3}\times [0,2\pi ]}J{X}_{n-1}(\psi )\mathrm{d}\psi =|{\mathrm{\S }}^{n-2}|{\int }_0^{\pi }{\sin }^{n-2}\phi \mathrm{d}\phi ,\end{array}}$$

where we use the recursive property of the Jacobian determinant. ∎

B.1. Definition using Givens rotations

Spherical coordinates can equivalently be defined using Givens rotations (see e.g. [74, ch. 5.1.8]). A Givens rotation for an angle $\phi \in [\mathrm{0,2}\pi )$ and indices $i,j⩽n$ with $i\ne j$ is determined by the rotation matrix $G(i,j,\phi )\in {ℝ}^{n\times n}$ :

$${\displaystyle \begin{array}{r}G(i,j,\phi {)}_{k,l}=\{\begin{array}{ll}\cos (\phi )& \text{if }k=l=i\text{ or }k=l=j,\\ 1& \text{if }k=l\ne i\text{ and }k=l\ne j,\\ \sin (\phi )& \text{if }k=i,l=j,\\ -\sin (\phi )& \text{if }k=j,l=i,\\ 0& \text{otherwise.}\end{array}\end{array}}$$

Applying $G(i,j,\phi {)}^T$ to a vector $x\in R^n$ corresponds to a counterclockwise rotation of $x$ by the angle $\phi$ in the $(i,j)$ -plane. For a given vector of angles $\phi \in [0,\pi ]{}^{n-2}\times [\mathrm{0,2}\pi ]$ , we can thus construct the matrix

$$\begin{array}{lrr}& \begin{array}{rl}R(\phi )=G(n-1,n,{\phi }_{n-1})\circ \dots & \circ G(\mathrm{2,3},{\phi }_2)\\ & \circ G(\mathrm{1,2},{\phi }_1)\circ G(\mathrm{2,3},{\phi }_2{)}^T\circ \cdots \circ G(n-1,n,{\phi }_{n-1}{)}^T.\end{array}& \end{array}$$ (B 1)

The rotation matrix $R(\phi )$ can be written as a two-dimensional rotation of angle ${\phi }_1$ in the $(e_1,X_{n-1}({\phi }_{\hat{1}}))$ -plane, as the following lemma shows.

Lemma B.2. Let $R(\phi )$ be the rotation matrix as described in equation (B 1) . Then, it holds that

$$\begin{array}{r}R(\phi )=UG(\mathrm{1,2},{\phi }_1)U{}^T,\end{array}$$

with $U{U}^T=\text{Id}$ , $U_{1,\cdot }=e_1$ and $U_{2,\cdot }=(0,X{}_{n-1}({\phi }_{\hat{1}}){)}^T$ .

Proof. For $n=2$ , the statement can be verified by inserting $U=\mathrm{Id}$ and the definition of $R(\phi )$ . For $n>2$ , we define

$$\begin{array}{r}U=G(\mathrm{2,3},{\phi }_2)\circ \cdots \circ G(n-1,n,{\phi }_{n-1}).\end{array}$$

With this choice of $U$ , $R(\phi )$ has the claimed form and $U{U}^T=\mathrm{Id}$ due to the orthogonality of Givens matrices. It remains to show that the first two rows of $U$ fulfil $U_{1,\cdot }=e_1$ and $U_{2,\cdot }=(0,X{}_{n-1}({\phi }_{\hat{1}}){)}^T$ . For $n=3$ , $U$ reduces to

$$\begin{array}{r}U=G(\mathrm{2,3},{\phi }_2)=\left(\begin{array}{ccc}1& 0& 0\\ 0& \cos {\phi }_2& \sin {\phi }_2\\ 0& -\sin {\phi }_2& \cos {\phi }_2\end{array}\right),\end{array}$$

and clearly, $U_{1,\cdot }=e_1$ and $U_{2,\cdot }=(0,\cos {\phi }_2,\sin {\phi }_2{)}^T=(0,X_2({\phi }_2){)}^T$ . For $n>3$ , the proof follows from induction over $n$ .∎

Corollary B.3. Let $x=X_n(\phi )$ , $\stackrel{\mathrm{~}}{x}=(0,X_{n-1}({\phi }_{\hat{1}}))$ then

$${\displaystyle \begin{array}{rl}R(\phi {)}^{T}y=y-\left(y\cdot e_1\right)e_1-(y\cdot \tilde{x})\tilde{x}& +\left(\cos \left({\phi }_1\right)\left(y\cdot e_1\right)-\sin \left({\phi }_1\right)(y\cdot \tilde{x})\right)e_1\\ +& \left(\sin \left({\phi }_1\right)\left(y\cdot e_1\right)+\cos \left({\phi }_1\right)(y\cdot \tilde{x})\right)\tilde{x}\end{array}}$$

In particular, if $y\cdot e_1=0$ it holds that

$$\begin{array}{rl}R(\phi {)}^{T}y=y& -(y\cdot \stackrel{\mathrm{~}}{x})\stackrel{\mathrm{~}}{x}+(y\cdot \stackrel{\mathrm{~}}{x})(-\sin ({\phi }_1)e{}_1+\cos ({\phi }_1)\stackrel{\mathrm{~}}{x}).\end{array}$$

With the above results, we obtain

$$\begin{array}{r}{X}_n(\phi )=R(\phi {)}^T{e}_1,\end{array}$$

and since Givens matrices are orthonormal, it also holds that

$$\begin{array}{r}R(\phi )X_n(\phi )=e_1.\end{array}$$

We can also therefore consider rotated spherical coordinates

$$\begin{array}{r}{X}_n^{\theta }(\phi )=R(\theta {)}^T{X}_n(\phi )\end{array}$$

for a reference point $x=X_n(\theta )$ , with the same Jacobian determinant as before, i.e. $J{X}_n^{\theta }(\phi )=J{X}_n(\phi )$ .

Appendix C. Proofs for Section 4

C.1. Proof of Lemma 4.10

Lemma 4.7 (cont.) Let $n>2$ . The uniform distribution $\mu =\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ is a stationary point of $\mathcal{E}$ if and only if all eigenvalues $\{\lambda {}_i{\}}_{i=1}^n$ of $D$ have the same absolute value, i.e. $\left|\lambda {}_i\right|=\lambda$ for some $\lambda \in ℝ$ .

Proof. The proof for $n>2$ uses the same arguments as for $n=2$ ; however, the rotation corresponding to a translation of the angle in two dimensions is technically more complicated. We use the notation and techniques from appendix B (spherical coordinates $X_n$ and rotations $R$ ).

Again, we first fix $x\in {\mathcal{S}}^{n-1}$ and consider the integral

$$\begin{array}{r}{\int }_{{\mathcal{S}}^{n-1}}{e}^{x\cdot Dy}{P}_x^{⟂}Dy\mathrm{d}\mathcal{H}{}^n(y)=(\ast ).\end{array}$$

Similarly to the two-dimensional case, we choose $\phi \in [0,\pi ]{}^{n-2}\times [\mathrm{0,2}\pi ]$ such that

$$\begin{array}{r}{X}_n(\phi )=\frac{Dx}{\Vert Dx\Vert },\end{array}$$

and therefore also

$$\begin{array}{r}R(\phi )Dx=\Vert Dx\Vert R(\phi )R(\phi {)}^T{e}_1=\Vert Dx\Vert e_1,\end{array}$$

where $e_1=(\mathrm{1,0},\dots ,0{)}^T\in {ℝ}^n$ denotes the first standard basis vector. We rewrite the integral using rotated spherical coordinates and substitute it into the above identity to obtain

$${\displaystyle \begin{array}{rl}(\ast )=& {\int }_{[0,\pi {]}^{n-2}\times [0,2\pi ]}{e}^{x\cdot DR(\phi {)}^T{X}_n(\theta )}{P}_x^{\perp }DR(\phi {)}^T{X}_n(\theta )J{X}_n(\theta )\mathrm{d}\theta \\ =& {\int }_{[0,\pi {]}^{n-2}\times [0,2\pi ]}{e}^{\Vert Dx\Vert \cos ({\theta }_1)}(DR(\phi {)}^T{X}_n(\theta )-\Vert Dx\Vert \cos ({\theta }_1)x)J{X}_n(\theta )\mathrm{d}\theta ,\end{array}}$$

where $J{X}_n$ denotes the Jacobian determinant of $X_n$ . To reduce the above integral over the vector $\theta$ to an integral over only the first component ${\theta }_1$ , we write

$$\begin{array}{rl}R(\phi {)}^T{X}_n(\theta )& =\cos ({\theta }_1)R(\phi {)}^T{e}_1+R(\phi {)}^T\left(\begin{array}{c}0\\ X_n(\theta {)}_{\hat{1}}\end{array}\right)\\ & =\cos ({\theta }_1)\frac{Dx}{\Vert Dx\Vert }+\sin ({\theta }_1)R(\phi {)}^T\left(\begin{array}{c}0\\ X_{n-1}({\theta }_{\hat{1}})\end{array}\right),\end{array}$$

where the subscript $\hat{1}$ denotes that we neglect the first component. Inserting this into $(\ast )$ , we get

$${\displaystyle \begin{array}{rl}(\ast )& =\left(D^{2}x/\Vert Dx\Vert -\Vert Dx\Vert x\right){\int }_0^{\pi }{e}^{\Vert Dx\Vert \cos ({\theta }_1)}\cos ({\theta }_1){\sin }^{n-2}({\theta }_1)\mathrm{d}\theta {}_1\\ & +{\int }_0^{\pi }{e}^{\Vert Dx\Vert \cos ({\theta }_1)}{\sin }^{n-1}({\theta }_1)DR(\theta {)}^T\underset{=0}{\underbrace{{\int }_{S^{n-2}}\left(\begin{array}{c}0\\ z\end{array}\right)\mathrm{d}{\mathcal{H}}^{n-1}(z)}}\mathrm{d}\theta {}_1\\ & =C(n,\Vert Dx\Vert )\left(D^{2}x/\Vert Dx\Vert -\Vert Dx\Vert x\right),\end{array}}$$

and due to the symmetry of sine and cosine, we have that $C(n,\Vert Dx\Vert )>0$ for any $n⩾2$ , $\Vert Dx\Vert >0$ . We can thus deduce that $(\ast )=0$ if and only if $x$ is an eigenvector of $D^2$ , exactly as in the case $n=2$ . This holds true for $\mu$ -almost all $x\in {\mathcal{S}}^{n-1}$ if and only if all eigenvalues of $D$ have the same absolute value, which then automatically yields $\mathrm{d}{\mathcal{E}}_D(\mu ,V)=0$ .

Again, it remains to show that this is also necessary. Without loss of generality, we assume $\left|\lambda {}_1\right|>\left|\lambda {}_2\right|$ and $\lambda {}_1$ and $\lambda {}_2$ to be the eigenvalues of largest, respectively second largest, absolute value corresponding to the eigenvectors $z_1$ , respectively $z_2$ .

From here, the strategy is the exact same as in the two-dimensional case, which we restate here for completeness. The factor $(D^{2}x/\Vert Dx\Vert -\Vert Dx\Vert x)\cdot z_2$ is strictly negative on the set

$$\begin{array}{r}A=\{x\in {\mathcal{S}}^{n-1}|(x\cdot z_1)\in (\left|\lambda {}_2/\lambda {}_1\right|,1),(x\cdot z_2)>0\}.\end{array}$$

Since $\mu (A)>0$ , we can find a Lipschitz continuous $V$ such that $V\cdot z_1=0$ for $\mu$ -a.e. on ${\mathcal{S}}^{n-1}$ and

$${\displaystyle \begin{array}{r}V(x)\cdot z_2=\{\begin{array}{ll}>0& \text{for a.e. }x\in A,\\ =0& \text{for a.e. }x\in {\mathcal{S}}^{n-1}\mathrm{\setminus }A.\end{array}\end{array}}$$

For all such $V$ , it holds that $\mathrm{d}{\mathcal{E}}_D(\mu ,V)>0$ , which concludes the proof.∎

C.2. Proof of Lemma 4.12

Lemma 4.8 (cont.) Let $n⩾2$ , and ${\mu }_0=\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ . Then, it holds that

$$\begin{array}{}\text{(C 1)}& {\displaystyle {\int }_{{\mathcal{S}}^{n-1}}{e}^{x\cdot y}x\mathrm{d}{\mu }_0(x)=C_{1}y}\end{array}$$

for any $y\in {\mathcal{S}}^{n-1}$ , where the constant $C_1$ is positive and depends only on the dimension $n$ .

Proof. The proof for $n>2$ goes along the lines of the proof for $n=2$ . However, the rotation corresponding to a translation of the angle in two dimensions is technically more complicated in higher dimensions. For an introduction to rotated spherical coordinates used in this proof, we refer the reader to appendix B.

We first fix $y\in {\mathcal{S}}^{n-1}$ and choose $\theta \in {ℝ}^{n-1}$ such that $y=X_n(\theta )$ . We proceed to write the integral using rotated spherical coordinates $x=X_n^{\theta }(\phi )$ and obtain

$${\displaystyle {\int }_{{\mathcal{S}}^{n-1}}{e}^{x\cdot y}{x}_j\mathrm{d}{\mu }_0(x)=\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_{[0,\pi {]}^{n-2}\times [0,2\pi ]}{e}^{X_n^{\theta }(\phi )\cdot X_n(\theta )}{\left(X_n^{\theta }(\phi )\right)}_{i}J{X}_n(\phi )\mathrm{d}\phi =(\ast ).}$$

Substituting the expressions for $X_n$ and $X_n^{\theta }$ yields

$$\begin{array}{rl}{X}_n^{\theta }(\phi )\cdot X_n(\theta )& =R(\theta {)}^T{X}_n(\phi )\cdot X_n(\theta )=X_n(\phi )\cdot R(\theta )X_n(\theta )=X_n(\phi )\cdot e_1=\cos ({\phi }_1).\end{array}$$

In addition, we note that we can write any $x=x_1{e}_1+(0,x_2,\dots ,x_n{)}^T$ and see that

$${\displaystyle \begin{array}{rl}{X}_n^{\theta }(\phi )=R(\theta {)}^T{X}_n(\phi )=& R(\theta {)}^T\cos ({\phi }_1)e_1+R(\theta {)}^T\left(\begin{array}{c}0\\ X_n(\phi {)}_{\hat{1}}\end{array}\right)\\ =& \cos ({\phi }_1)y+\sin ({\phi }_1)R(\theta {)}^T\left(\begin{array}{c}0\\ X_{n-1}({\phi }_{\hat{1}})\end{array}\right),\end{array}}$$

where $e_1=(\mathrm{1,0},\dots ,0{)}^T\in {ℝ}^n$ denotes the first standard basis vector. Substituting the above equality into the integral, we derive

$${\displaystyle \begin{array}{rl}(\ast )=& \frac{1}{|{\mathcal{S}}^{n-1}|}{\int }_{[0,\pi {]}^{n-2}\times [0,2\pi ]}{e}^{\cos ({\phi }_1)}{\left[\cos ({\phi }_1)y+\sin ({\phi }_1)R(\theta {)}^T\left(\begin{array}{c}0\\ X_{n-1}({\phi }_{\hat{1}})\end{array}\right)\right]}_{j}J{X}_n(\phi )\mathrm{d}\phi \\ =& y_i\frac{|{\mathcal{S}}^{n-2}|}{|{\mathcal{S}}^{n-1}|}{\int }_0^{\pi }{e}^{\cos \phi }\cos \phi {\sin }^{n-2}\phi \mathrm{d}\phi \\ & +\frac{1}{|{\mathcal{S}}^{n-1}|}{\int }_0^{\pi }{e}^{\cos \phi }{\sin }^{n-1}\phi {\left[R(\theta {)}^T\underset{=0}{\underbrace{{\int }_{{\mathcal{S}}^{n-2}}\left(\begin{array}{c}0\\ z\end{array}\right)\mathrm{d}{\mathcal{H}}^{n-1}(z)}}\right]}_j\mathrm{d}\phi .\end{array}}$$

The proof now follows from choosing the constant:

$$\begin{array}{r}{C}_1=\frac{\left|{\mathcal{S}}^{n-2}\right|}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_0^{\pi }{e}^{\cos \phi }\cos \phi {\sin }^{n-2}\phi \mathrm{d}\phi =\frac{\left|{\mathcal{S}}^{n-2}\right|}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_0^{\pi /2}{\sin }^{n-2}\phi \cos \phi \sinh (\cos \phi )\mathrm{d}\phi ,\end{array}$$

which is positive for all $n⩾2$ since the function $t\mapsto t\sinh t$ is positive for $t>0$ and both sine and cosine are positive for $\phi \in (0,\pi /2)$ .∎

C.3. Proof of Lemma 4.13

Lemma 4.9 (cont.) Let $n⩾2$ , and ${\mu }_0=\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\mathcal{H}}^n$ . Then, for all $y\in {\mathcal{S}}^{n-1}$ and $1⩽i⩽n$ , it holds that

$$\begin{array}{}\text{(C 2)}& {\displaystyle {\int }_{{\mathcal{S}}^{n-1}}{e}^{x\cdot y}{x}_i^2\mathrm{d}{\mu }_0(x)=C_2{y}_i^2+C_3,}\end{array}$$

where the constants $C_2$ and $C_3$ are positive and depend only on the dimension $n$ .

Proof. Using the same arguments as in the previous proof, we obtain

$${\displaystyle \begin{array}{rl}{\int }_0^{\pi }{e}^{x\cdot y}{x}_j^2\mathrm{d}{\mu }_0(x)=& y_j^2\frac{\left|{\mathcal{S}}^{n-2}\right|}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_{{\mathcal{S}}^{n-1}}{e}^{\cos \phi }{\cos }^2\phi {\sin }^{n-2}\phi \mathrm{d}\phi \\ \text{(C 3)}& & +\frac{1}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_0^{\pi }{e}^{\cos \phi }{\sin }^n\phi {\int }_{{\mathcal{S}}^{n-2}}{\left[R(\theta {)}^T(\genfrac{}{}{0pt}{}{0}{z})\right]}_j^2\mathrm{d}{\mathcal{H}}^{n-1}(z)\mathrm{d}\phi ,\end{array}}$$

where the mixed term containing $x_i{y}_i$ vanishes due to symmetry. Since the second term still depends on $y$ due to the rotation, we write $\stackrel{\mathrm{~}}{y}=(0,X_{n-1}({\theta }_{\hat{1}}))$ and decompose $\stackrel{\mathrm{~}}{z}=(0,z)$ into its rotation-invariant and rotation-variant part. More precisely, we use corollary B.3 to get

$$\begin{array}{r}R(\theta {)}^T\stackrel{\mathrm{~}}{z}=\stackrel{\mathrm{~}}{z}-(\stackrel{\mathrm{~}}{y}\cdot \stackrel{\mathrm{~}}{z})\stackrel{\mathrm{~}}{y}+(\stackrel{\mathrm{~}}{y}\cdot \stackrel{\mathrm{~}}{z})[-\sin ({\theta }_1)e{}_1+\cos ({\theta }_1)\stackrel{\mathrm{~}}{y}]\end{array}$$

and thus

$$\begin{array}{r}{[R(\theta {)}^T\stackrel{\mathrm{~}}{z}]}^2=(\stackrel{\mathrm{~}}{z}-(\stackrel{\mathrm{~}}{y}\cdot \stackrel{\mathrm{~}}{z})\stackrel{\mathrm{~}}{y}{)}^2+(\stackrel{\mathrm{~}}{y}\cdot \stackrel{\mathrm{~}}{z}{)}^2({\sin }^2({\theta }_1)e{}_1+{\cos }^2({\theta }_1)\stackrel{\mathrm{~}}{y}{}^2)+2\cos ({\theta }_1)(\stackrel{\mathrm{~}}{z}-(\stackrel{\mathrm{~}}{y}\cdot \stackrel{\mathrm{~}}{z})\stackrel{\mathrm{~}}{y})(\stackrel{\mathrm{~}}{y}\cdot \stackrel{\mathrm{~}}{z})\stackrel{\mathrm{~}}{y}.\end{array}$$

Making use of the trigonometric identity ${\cos }^2({\theta }_1)+{\sin }^2({\theta }_1)=1$ , we get

$${\displaystyle \begin{array}{rl}{\left[R(\theta {)}^T\tilde{z}\right]}^2& ={\tilde{z}}^2+(\tilde{y}\cdot \tilde{z}{)}^2{e}_1+2(\tilde{y}\cdot \tilde{z}{)}^2{\tilde{y}}^2-2(\tilde{y}\cdot \tilde{z})\tilde{z}\tilde{y}+2\cos \left({\theta }_1\right)(\tilde{z}-(\tilde{y}\cdot \tilde{z})\tilde{y})(\tilde{y}\cdot \tilde{z})\tilde{y}\\ & -(\tilde{y}\cdot \tilde{z}{)}^2\left({\cos }^2\left({\theta }_1\right)e_1+{\sin }^2\left({\theta }_1\right){\tilde{y}}^2\right),\\ \text{(C 4)}& & ={\tilde{z}}^2+(\tilde{y}\cdot \tilde{z}{)}^2\left(e_1-y^2\right)+2\left(\cos \left({\theta }_1\right)-1\right)(\tilde{z}-(\tilde{y}\cdot \tilde{z})\tilde{y})(\tilde{y}\cdot \tilde{z})\tilde{y},\end{array}}$$

where in the last step, we use the fact that $y^2={\cos }^2{\theta }_1{e}_1+{\sin }^2{\theta }_1{\stackrel{\mathrm{~}}{y}}^2$ . To prove that the integral over the expression in equation (C 4) can be written as claimed, we observe that for all $j=2,\dots ,n$

$$\begin{array}{r}{\int }_{{\mathcal{S}}^{n-2}}{\stackrel{\mathrm{~}}{z}}_j^2\mathrm{d}\mathcal{H}{}^{n-1}(z)=:\stackrel{\mathrm{~}}{C},\end{array}$$

where $\stackrel{\mathrm{~}}{C}$ is positive and depends only on $n$ , and therefore, also ${\int }_{{\mathcal{S}}^{n-2}}(\stackrel{\mathrm{~}}{z}\cdot \stackrel{\mathrm{~}}{y}{)}^2\mathrm{d}\mathcal{H}{}^{n-1}(z)=\stackrel{\mathrm{~}}{C}\Vert \stackrel{\mathrm{~}}{y}{\Vert }^2=\stackrel{\mathrm{~}}{C}$ . With this, we derive that

$$\begin{array}{}\text{(C 5)}& {\displaystyle {\left[{\int }_{{\mathcal{S}}^{n-2}}{\tilde{z}}^2+(\tilde{y}\cdot \tilde{z}{)}^2\left(e_1-y^2\right)\mathrm{d}{\mathcal{H}}^{n-1}(z)\right]}_j=\tilde{C}\left(1-y_j^2\right)}\end{array}$$

for all $j=1,\dots ,n$ and it remains to show that for any $1⩽j⩽n$

$${\displaystyle \begin{array}{}\text{(C 6)}& {\int }_{{\mathcal{S}}^{n-2}}[(\tilde{z}-(\tilde{y}\cdot \tilde{z})\tilde{y})(\tilde{y}\cdot \tilde{z})\tilde{y}{]}_j\mathrm{d}\mathcal{H}{}^{n-1}(z)=0.\end{array}}$$

The case $j=1$ is trivial as ${\stackrel{\mathrm{~}}{y}}_1={\stackrel{\mathrm{~}}{z}}_1=0$ . For $2⩽j⩽n$ , we write out the integrand and obtain

$${\displaystyle {\displaystyle \begin{array}{rl}& [(\tilde{z}-(\tilde{y}\cdot \tilde{z})\tilde{y})(\tilde{y}\cdot \tilde{z})\tilde{y}{]}_j=\left(\sum _{k=1}^n{\tilde{z}}_k{\tilde{y}}_k\right){\tilde{z}}_j{\tilde{y}}_j-\left(\sum _{k,l=1}^n{\tilde{z}}_k{\tilde{z}}_l{\tilde{y}}_k{\tilde{y}}_l\right){\tilde{y}}_j^2\\ & =\left(\sum _{k=1,k\ne j}^n{\tilde{z}}_j{\tilde{z}}_k{\tilde{y}}_j{\tilde{y}}_k\right)-\left(\sum _{k,l=1,k\ne l}^n{\tilde{z}}_k{\tilde{z}}_l{\tilde{y}}_k{\tilde{y}}_l\right){\tilde{y}}_j^2+\left({\tilde{z}}_j^2-(\tilde{z}\cdot \tilde{y}{)}^2\right){\tilde{y}}_j^2,\end{array}}}$$

where we can use the same argument as for equation (C 5) to show that the last summand integrates to zero. Since also ${\int }_{{\mathcal{S}}^{n-2}}{\stackrel{\mathrm{~}}{z}}_j{\stackrel{\mathrm{~}}{z}}_k\mathrm{d}{\mathcal{H}}^{n-1}(z)=0$ for any $j\ne k$ , we derive equation (C 6). Togethepossible to interpret it as a forward Euler discretizationr with equations (C 4) and (C 5), this yields

$$\begin{array}{rl}{\int }_{{\mathcal{S}}^{n-2}}{[R(\theta {)}^T\stackrel{\mathrm{~}}{z}]}_j^2\mathrm{d}\mathcal{H}{}^{n-1}(z)& =\stackrel{\mathrm{~}}{C}(1-y_j^2).\end{array}$$

The statement now follows from substituting the above into equation (C 3), with constants given by

$$\begin{array}{rl}& C_2=\frac{\left|{\mathcal{S}}^{n-2}\right|}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_0^{\pi }{e}^{\cos \phi }{\cos }^2\phi {\sin }^{n-2}\phi \mathrm{d}\phi -C_3,C_3=\frac{\stackrel{\mathrm{~}}{C}}{\left|{\mathcal{S}}^{n-1}\right|}{\int }_0^{\pi }{e}^{\cos \phi }{\sin }^n\phi \mathrm{d}\phi .\end{array}$$

Since $\stackrel{\mathrm{~}}{C}>0$ for all $n⩾2$ , it directly follows that $C_3>0$ . To show that $C_2>0$ for all $n⩾2$ , we first show that $\stackrel{\mathrm{~}}{C}=\left|{\mathcal{S}}^{n-2}\right|/(n-1)$ . For $n=2$ , this follows directly from $\stackrel{\mathrm{~}}{C}=\left|{\mathcal{S}}^0\right|=2$ . For $n>2$ , we have

$${\displaystyle \begin{array}{r}\tilde{C}=|{\mathcal{S}}^{n-3}|{\int }_0^{\pi }{\cos }^2\phi {\sin }^{n-3}\phi \mathrm{d}\phi =|{\mathcal{S}}^{n-3}|{\int }_0^{\pi }{\sin }^{n-3}\phi -{\sin }^{n-1}\phi \mathrm{d}\phi .\end{array}}$$

Using integration by parts, we further derive that

$${\displaystyle \begin{array}{r}{\int }_0^{\pi }{\sin }^{n-1}\phi \mathrm{d}\phi =(n-2)/(n-1){\sin }^{n-3}\phi \mathrm{d}\phi .\end{array}}$$

As shown in lemma B.1, the recursive form of the Jacobian determinant of spherical coordinates yields that

$${\displaystyle \left|{\mathcal{S}}^{n-2}\right|=\left|{\mathcal{S}}^{n-3}\right|{\int }_0^{\pi }{\sin }^{n-3}\phi \mathrm{d}\phi }$$

Combining these equalities, we see that

$$\begin{array}{r}\stackrel{\mathrm{~}}{C}=(1-((n-2)/(n-1)))\left|{\mathcal{S}}^{n-2}\right|=\left|{\mathcal{S}}^{n-2}\right|/(n-1),\end{array}$$

and therefore, with integration by parts, we get

$${\displaystyle \begin{array}{rl}{C}_2& =\frac{|{\mathcal{S}}^{n-2}|}{|{\mathcal{S}}^{n-1}|}{\int }_0^{\pi }{e}^{\cos \phi }\left[{\cos }^2\phi {\sin }^{n-2}\phi -\frac{1}{1-n}{\sin }^n\phi \right]\mathrm{d}\phi \\ & =\frac{|{\mathcal{S}}^{n-2}|}{|{\mathcal{S}}^{n-1}|}{\int }_0^{\pi }{e}^{\cos \phi }{\sin }^{n-2}\phi \cos \phi \left(\cos \phi -1\right)\mathrm{d}\phi .\end{array}}$$

Due to the symmetry of sine and cosine, we get

$${\displaystyle \begin{array}{rl}{C}_2& =\frac{|{\mathcal{S}}^{n-2}|}{|{\mathcal{S}}^{n-1}|}{\int }_0^{\pi /2}{e}^{\cos \phi }{\sin }^{n-2}\phi \cos \phi \left(\cos \phi -1\right)+e^{-\cos \phi }{\sin }^{n-2}\phi \cos \phi \left(\cos \phi +1\right)\mathrm{d}\phi \\ & =2\frac{|{\mathcal{S}}^{n-2}|}{|{\mathcal{S}}^{n-1}|}{\int }_0^{\pi /2}{\sin }^{n-2}\phi \left(\cos (\phi )\cosh (\cos (\phi ))-\sinh (\cos (\phi ))\right)>0,\end{array}}$$

where the positivity follows from the fact that the function $t\mapsto t\cosh (t)-\sinh (t)$ is positive for $t>0$ and both sine and cosine are positive for $\phi \in (0,\pi /2)$ .∎

Data accessibility

This article has no additional data.

Declaration of AI use

We have not used AI-assisted technologies in creating this article.

Authors’ contributions

M.B.: conceptualization, formal analysis, funding acquisition, investigation, methodology, supervision, writing—original draft; S.K.: formal analysis, investigation, methodology, visualization, writing—original draft; Y.K.: conceptualization, formal analysis, funding acquisition, investigation, methodology, writing—original draft; T.R.: funding acquisition, investigation, methodology, software, visualization, writing—original draft; L.W.: formal analysis, investigation, methodology, writing—original draft.

All authors gave final approval for publication and agreed to be held accountable for the work performed therein.

Conflict of interest declaration

We declare we have no competing interests.

Funding

M.B. and T.R. acknowledge funding by the German Ministry of Science and Technology (BMBF) under grant agreement No. 01IS24072A (COMFORT). M.B., S.K., T.R. and L.W. acknowledge support from DESY (Hamburg, Germany), a member of the Helmholtz Association HGF. This research was supported in part through the Maxwell computational resources operated at Deutsches Elektronen-Synchrotron DESY, Hamburg, Germany. M.B. and S.K. acknowledge support from the German Research Foundation, project BU 2327/19-1. M.B. and L.W. acknowledge support from the German Research Foundation, project BU 2327/20-1. Y.K. acknowledges support from the German Research Foundation as visiting fellow within the priority programme Foundations of Deep Learning. Part of this study was carried out while S.K. and T.R. were visiting the California Institute of Technology, supported by the DAAD grant for project 57698811 'Bayesian Computations for Large-scale (Nonlinear) Inverse Problems in Imaging'. Y.K. acknowledges the support of the EPSRC (Fellowship EP/V003615/2 and Programme Grant EP/V026259/1). S.K. and Y.K. are grateful for the hospitality of the University of Bath during the workshop 'Machine Learning in Infinite Dimensions', sponsored by the ICMS, LMS, IMI Bath, ProbAI and Maths4DL, where part of this work was undertaken.