Stochastic Wasserstein Gradient Flows using Streaming Data with an Application in Predictive Maintenance

[footnoteinfo]e-mails: ,
. This research is supported by the Swiss National Science Foundation through NCCR Automation (Grant Number 180545).

ETH]Nicolas Lanzetti
ETH]Efe C. Balta
UBC]Dominic Liao-McPherson
ETH]Florian Dörfler

[ETH]Automatic Control Laboratory, ETH Zürich, Physikstrasse 3, 8092 Zürich, Switzerland.
[UBC]University of British Columbia Department of Mechanical Engineering, 2054-6250 Applied Science Lane,
Vancouver, BC V6T 1Z4, Canada.

We study estimation problems in safety-critical applications with streaming data. Since estimation problems can be posed as optimization problems in the probability space, we devise a stochastic projected Wasserstein gradient flow that keeps track of the belief of the estimated quantity and can consume samples from online data.
We show the convergence properties of our algorithm. Our analysis combines recent advances in the Wasserstein space and its differential structure with more classical stochastic gradient descent.
We apply our methodology for predictive maintenance of safety-critical processes: Our approach is shown to lead to superior performance when compared to classical least squares, enabling, among others, improved robustness for decision-making.

Wasserstein gradient flows, streaming data, predictive maintenance

§ INTRODUCTION

Providing performance guarantees for parameter estimation algorithms operating with streaming data is a key challenge when developing methods for safety-critical applications across various domains of engineering and data science. Ideally, one should be able to (i) efficiently handle streaming data in real time, without resorting to computationally expensive one-shot numerical routines, and (ii) rigorously quantify the uncertainty related to the estimated quantity. In this paper, we focus on probabilistic approaches to uncertainty quantification rather than set-based ones [Combettes, 1993].

A prominent approach for parameter estimation with streaming data is acr:rls. In acr:rls, the online solution is obtained by “updating” the previous solution with the latest measurement. acr:rls  avoids the need to store and invert large data matrices and provides probabilistic guarantees on its estimate when the process is linear and all distributions are Gaussian.  acr:rls is a special case of Bayes filter (see e.g., [Särkkä, 2013, Sullivan, 2015]), whose many variants (e.g., particle filters, extended Kalman filters, etc.) are the dominant approaches for inference using non-Gaussian distributions. Bayes filter is powerful but inflexible, it can be challenging to integrate side information and it can be difficult to implement due to the need to compute high-dimensional integrals.

In this work, we propose a different approach based on the theory of optimization in Wasserstein probability spaces [Jordan et al., 1998, Ambrosio et al., 2005, Lanzetti et al., 2022]. We pose the parameter estimation problem as an optimization problem in the probability space and devise a stochastic projected gradient flow to iteratively compute its optimal solution using samples obtained from streaming data. Our approach maintains and iteratively improves an estimate of the probability measure of the estimated quantities and does not require a-priori assumptions on the probability measures (e.g., Gaussianity), but rather works in the space of all probability measures with finite second moment. Our proposed framework is more flexible than the Bayes filter in the sense that its intuitive to add side information e.g., constraints on the support of the final distribution or on the variance (indeed, it can be used to recover maximum likelihood estimator for stochastic least squares problems [Rigollet and Weed, 2018]).

A motivating application of interest is acr:pdm, where the goal is to efficiently maintain a safety-critical process (e.g., with minimal interruption) before an unsafe event occurs (see [Pech et al., 2021] for a recent survey of results).

The acr:pdm problem is challenging from an online algorithmic perspective since, in practice, there is often only historical data on the nominal operation and little or no data on the unsafe operation.

Moreover, the problem calls for careful risk analysis: Too conservative decisions impact performance and efficiency, while unsafe events, if they occur, might lead to catastrophic failures.
The current state-of-the-art consists of rule-based methods and estimation strategies that rely on predetermined distribution models [Hu and Chen, 2020, Kanso et al., 2022].

With our work, we learn the model of the underlying process without a-priori assumptions on its probability measure to improve the overall performance by reducing conservativeness.

Our contributions are twofold. First, we propose a novel stochastic projected gradient flow for optimization in the probability space that operates on streaming data and study its convergence properties. Our analysis combines tools from optimal transport and differential calculus in the probability space with more classical projected stochastic gradient descent. We prove that similar to the Euclidean setting, our scheme yields convergence to a ball around the optimal solution.
Second, we apply our scheme to the predictive maintenance of the damping ratio of a second-order system and demonstrate improved performance relative to a classical least-squares approach.

* Modeling and data analysis for online applications (e.g., control, analysis, verification, etc.) are relevant and ubiquitous in various domains of engineering and science. (why we're looking into this problem at all).

* However, a lot of the existing methods either rely on a restricted class of functions to model the behavior of the underlying process or utilize a black-box model that lacks interpretability and requires a lot of training data. (why modeling in a distribution space, people usually assume a model on the distribution, etc.).

* Even when we learn a model in our framework, the model may change over time, have dynamics that can only be observed in the run-time (e.g., the predictive maintenance example we have), or we may not have enough data to learn a good enough model. Thus, we need methods that can consume online data (motivating the online aspect).

* One application of interest in this context is predictive maintenance (PdM). In PdM, we have a safety-critical process that we wish to maintain before an unsafe event occurs. This problem is tricky, because we may not have prior data on what happens in an unsafe situation (i.e., an imbalanced dataset). The current state of the art consists of rule-based methods and estimation strategies that rely on predetermined distribution models and heuristics. Nevertheless, we wish to learn the model of the underlying process without assumptions on its distributional characteristics, etc., to improve the overall performance by reducing conservativeness.

§ BACKGROUND

In this section, we briefly review our notation, basics of measure theory and optimal transport, geodesic convexity, and Wasserstein gradients. For more details, we refer the reader to [Villani, 2009, Ambrosio et al., 2005, Santambrogio, 2015, Lanzetti et al., 2022].
finalsubmissionAll proofs of our theoretic results are relegated to the online extended version of this paper <cit.>.

§.§.§ Notation

We consider the Euclidean space $\reals^d$, with the usual Euclidean norm $\norm{\cdot}$.
For a matrix $A\in\reals^{m\times n}$, we denote by $\sigma_\mathrm{min}(A)$ and by $\sigma_\mathrm{max}(A)$ its minimum and maximum singular value, respectively. If $m=n$, we use the notation $\lambda_\mathrm{min}(A)$ and by $\lambda_\mathrm{max}(A)$ for the minimum and maximum eigenvalue of $A$ and $\trace(A)$ for its trace.

§.§.§ Basics in Measure Theory

We denote by $\Pp{}{\reals^d}$ the space of (Borel) probability measures over $\reals^d$ and by $\Pp{2}{\reals^d}\coloneqq\{\mu\in\Pp{}{\reals^d}:\int_{\reals^d}\norm{x}^2\d\mu(x)<+\infty\}$ the space of probability measures with finite second moment.
We denote the Dirac measures at $x\in\reals^d$ by $\diracMeasure{x}$, defined by $\diracMeasure{x}(A)=1$ if and only if $x\in A$.
We denote the support of a probability measure $\mu\in\Pp{2}{\reals^d}$ by $\support{\mu}\subset\reals^d$.
The pushforward of a measure $\mu\in\Pp{}{\reals^d}$ via a (Borel) map $T:\reals^d\to\reals^d$ is denoted by $\pushforward{T}\mu$ and defined by $(\pushforward{T}\mu)(B)=\mu(T^{-1}(B))$ for all $B\subset\reals^d$ Borel. For any $f:\reals^d\to\reals$, $\pushforward{T}\mu$-integrable it holds

\begin{equation*}
\int_{\reals^d}f(x)\d(\pushforward{T}\mu)(x)
\int_{\reals^d}f(T(x))\d\mu(x).
\end{equation*}

A sequence of probability measures $(\mu_n)_{n\in\naturals}\subset\Pp{}{\reals^d}$ converges narrowly to $\mu\in\Pp{}{\reals^d}$ if $\int_{\reals^d}\phi(x)\d\mu(x)\to\int_{\reals^d}\phi(x)\d\mu(x)$ for all bounded continuous $\phi:\reals^d\to\reals$.

§.§.§ Wasserstein distance

The (type 2) Wasserstein distance between two probability measures $\mu, \nu\in\Pp{}{\reals^d}$ is

\begin{equation*}
\wassersteinDistance{2}{\mu}{\nu}
\coloneqq\left(\min_{\gamma\in\setPlans{\mu}{\nu}}\int_{\reals^d\times\reals^d}\norm{x-y}^2\d\gamma(x,y)\right)^{\frac{1}{2}},
\end{equation*}

where $\setPlans{\mu}{\nu}$ is the set of transport plans, that is, of probability measures on $\reals^d\times\reals^d$ whose first marginal is $\mu$ and second marginal is $\nu$; i.e.,  $\setPlans{\mu}{\nu}=\{\gamma\in\Pp{}{\reals^d\times\reals^d}:\pushforward{(\proj_1)}\gamma=\mu,\pushforward{(\proj_2)}\gamma=\nu\}$ where $\proj_1$ and $\proj_2$ are projection operators (e.g., $\proj_1(x,y)=x$). We denote by $\setOptimalPlans{\mu}{\nu}$ the (non-empty) set of optimal couplings between $\mu$ and $\nu$. It is well-known that the Wasserstein distance is a distance on $\Pp{2}{\reals^d}$.

§.§.§ Geodesic convexity

Given $\mu_0\in\Pp{}{\reals^d}$ and $\mu_1\in\Pp{}{\reals^d}$, we define the geodesic between them by $\mu_t=\pushforward{((1-t)\proj_1+t\proj_2)}\gamma$, where $\gamma\in\setOptimalPlans{\mu_0}{\mu_1}$ is an optimal transport plan between $\mu_0$ and $\mu_1$. Since optimal transport plans are generally not unique, there might exist multiple geodesics between $\mu_0$ and $\mu_1$.
Accordingly, a functional $J:\Pp{2}{\reals^d}\to\reals$ is $\alpha$-geodesically convex if for all $\mu_0,\mu_1\in\Pp{2}{\reals^d}$ there exists a geodesic $\mu_t$ so that $J(\mu_t)\leq (1-t)J(\mu_0)+tJ(\mu_1)-\frac{\alpha}{2}t(1-t)\wassersteinDistance{2}{\mu_0}{\mu_1}^2$. For instance, $\mu\mapsto \expectedValue{\mu}{V}$ is $\alpha$-geodesically convex if and only if $V:\reals^d\to\reals$ is $\alpha$-convex (i.e., convex with convexity parameter $\alpha$) and $\mu\mapsto\variance{\mu}{x_i}$ is geodesically convex (with $\alpha=0$), where $\variance{\mu}{x_i}$ denotes the variance of $x_i$.

§.§.§ Wasserstein gradient

A function $\gradient{\mu}J(\mu)\in\Lp{2}{\reals^d,\reals^d;\mu}$ is a Wasserstein gradient of a real-valued functional over the probability space $J:\Pp{2}{\reals^d}\to\reals$ if it approximates $J$ “linearly”; i.e., for all $\gamma\in\setOptimalPlans{\mu}{\nu}$

\begin{equation*}
\int_{\reals^d}\transpose{\gradient{\mu}J(\mu)(x)}(y-x)\d\gamma(x,y)+\onotation{\wassersteinDistance{2}{\mu}{\nu}},
\end{equation*}

where $\onotation{\wassersteinDistance{2}{\mu}{\nu}}$ denotes high-order term.
Wasserstein gradients are well-defined for many functionals of practical interest. In particular, we have $\gradient{\mu}\expectedValue{\mu}{V}=\gradient{}V$ for any smooth $V:\reals^d\to\reals$ with at most quadratic growth (i.e., the Wasserstein gradient of an expected value is simply the gradient of the function in the expected value) and $\gradient{}\variance{\mu}{x_i}=2(x_i-\expectedValue{\mu}{x_i})$. For the Wasserstein gradients of more functionals, we refer to [Lanzetti et al., 2022].

§ CONCLUSIONS

In this work, we present a novel stochastic Wasserstein gradient flow method to efficiently perform estimation in probability spaces with streaming data.

Our formal results provide a convergence analysis of our online stochastic optimization method, which provides convergence to a ball around the optimal solution, similar to the standard Euclidean setting.
We illustrate the utility of the proposed method in an application of predictive maintenance to show the benefit over classical approaches such as simple least-squares with Gaussianity assumptions.
Overall, our method provides a flexible online estimation tool to estimate a rich set of processes without any assumptions on the model of the underlying distribution.
Future work will consider providing further results under relaxed settings such as non-strong convexity and applications for predictive maintenance of real-world physical examples.

[Ambrosio et al., 2005]
Ambrosio, L., Gigli, N., and Savaré, G. (2005).
Gradient flows: in metric spaces and in the space of
probability measures.
Springer Science & Business Media.

[Aolaritei et al., 2022]
Aolaritei, L., Lanzetti, N., Chen, H., and Dörfler, F. (2022).
Uncertainty propagation via optimal transport ambiguity sets.
arXiv preprint arXiv:2205.00343.

[Blanchet and Murthy, 2019]
Blanchet, J. and Murthy, K. (2019).
Quantifying distributional model risk via optimal transport.
Mathematics of Operations Research, 44(2), 565–600.

[Bottou et al., 2018]
Bottou, L., Curtis, F.E., and Nocedal, J. (2018).
Optimization methods for large-scale machine learning.
Siam Review, 60(2), 223–311.

[Combettes, 1993]
Combettes, P.L. (1993).
The foundations of set theoretic estimation.
Proceedings of the IEEE, 81(2), 182–208.

[Gao and Kleywegt, 2022]
Gao, R. and Kleywegt, A. (2022).
Distributionally robust stochastic optimization with Wasserstein
Mathematics of Operations Research.

[Gelbrich, 1990]
Gelbrich, M. (1990).
On a formula for the L2 Wasserstein metric between measures on
Euclidean and Hilbert spaces.
Mathematische Nachrichten, 147(1), 185–203.

[Hu and Chen, 2020]
Hu, J. and Chen, P. (2020).
Predictive maintenance of systems subject to hard failure based on
proportional hazards model.
Reliability Engineering & System Safety, 196, 106707.

[Jordan et al., 1998]
Jordan, R., Kinderlehrer, D., and Otto, F. (1998).
The variational formulation of the Fokker–Planck equation.
SIAM journal on mathematical analysis, 29(1), 1–17.

[Kanso et al., 2022]
Kanso, S., Jha, M.S., Galeotta, M., and Theilliol, D. (2022).
Remaining useful life prediction with uncertainty quantification of
liquid propulsion rocket engine combustion chamber.
IFAC-PapersOnLine, 55(6), 96–101.

[Lanzetti et al., 2022]
Lanzetti, N., Bolognani, S., and Dörfler, F. (2022).
First-order conditions for optimization in the Wasserstein space.
arXiv preprint arXiv:2209.12197.

[Mohajerin Esfahani and Kuhn, 2018]
Mohajerin Esfahani, P. and Kuhn, D. (2018).
Data-driven distributionally robust optimization using the
Wasserstein metric: Performance guarantees and tractable reformulations.
Mathematical Programming, 171(1), 115–166.

[Pech et al., 2021]
Pech, M., Vrchota, J., and Bednář, J. (2021).
Predictive maintenance and intelligent sensors in smart factory.
Sensors, 21(4), 1470.

[Rigollet and Weed, 2018]
Rigollet, P. and Weed, J. (2018).
Entropic optimal transport is maximum-likelihood deconvolution.
Comptes Rendus Mathematique, 356(11-12), 1228–1235.

[Santambrogio, 2015]
Santambrogio, F. (2015).
Optimal transport for applied mathematicians.
Birkäuser, NY, 55(58-63), 94.

[Särkkä, 2013]
Särkkä, S. (2013).
Bayesian filtering and smoothing.
3. Cambridge university press.

[Sullivan, 2015]
Sullivan, T.J. (2015).
Introduction to uncertainty quantification, volume 63.

[Villani, 2009]
Villani, C. (2009).
Optimal transport: old and new, volume 338.

§ PROOFS

§.§ Proof of Proposition <ref>

We split the proof in three parts.

Existence: Since the Wasserstein is lower semi-continuous (w.r.t. narrow convergence), and has compact (w.r.t. narrow convergence) level sets, it suffices to prove that the set $\{\nu\in\Pp{}{\reals}:\support{\nu}\subset\Theta\}\subset\Pp{}{\reals^d}$ is closed (w.r.t. narrow convergence). Let $(\nu_n)_{n\in\naturals}\subset\Pp{}{\reals^d}$ so that $\support{\nu_n}\subset\Theta$ for all $n\in\naturals$, and assume that $(\nu_n)_{n\in\naturals}$ converges narrowly to $\nu\in\Pp{}{\reals^d}$. We seek to prove that $\support{\nu}\subset\Theta$.
For $m\in\naturals$, let $f_m(x)$ be continuous and bounded so that (i) $f_m(x)=1$ if $x\in\Theta$ and (ii) converges pointwise to $1_\Theta(x)$ as $m\to\infty$. As $\Theta$ is closed, such $f_m$ can always be constructed.
Since $\support{\nu_n}\subset\Theta$ for every $m\in\naturals$ and every $n\in\naturals$

\begin{equation*}
\int_{\reals^d}f_m(x)\d\nu_n(x)=\int_{\Theta}f_m(x)\d\nu_n(x)=\nu_n(\Theta)=1.
\end{equation*}

Moreover, as $f_m$ is continuous and bounded, for any fixed $m\in\naturals$ the definition of narrow convergence gives

\begin{equation*}
\end{equation*}

We can now deploy dominated convergence ($|f_m|$ is uniformly dominated by an $\nu$-integrable function) to conclude

\begin{equation*}
\begin{aligned}
\lim_{m\to\infty}\int_{\reals^d}f_m(x)\d\nu(x)
\int_{\reals^d}\lim_{m\to\infty}f_m(x)\d\nu(x)
\\
\int_{\reals^d}1_\mathrm{\Theta}(x)\d\nu(x) =\nu(\Theta).
% \\
% &=\nu(\Theta)
\end{aligned}
\end{equation*}

Thus, $\support{\nu}\subset\Theta$, and the set $\{\nu\in\Pp{}{\reals^d}:\support{\nu}\subset\Theta\}$ is closed (w.r.t. narrow convergence), as desired.

Uniqueness: Assume two non-equal projections $\bar\mu_1,\bar\mu_2\in\Pp{2}{\reals^d}$ exist. Let $\gamma_1\in\setOptimalPlans{\bar\mu}{\bar\mu_1}$ and $\gamma_2\in\setOptimalPlans{\bar\mu}{\bar\mu_2}$, and let $\mu_{1/2}$ be the generalized geodesics with $\mu_{1/2}=\pushforward{\left(\frac{1}{2}\proj_2+\frac{1}{2}\proj_3\right)}\gamma$,
where $\gamma\in\setPlans{\bar\mu,\bar\mu_1}{\bar\mu_2}\subset\Pp{}{\reals^d\times\reals^d\times\reals^\d}$ results from gluing $\gamma_1$ and $\gamma_2$ (via the Gluing Lemma, e.g., Lemma 5.3.2 in Ambrosio et al., 2005; more generally, see Chapter 9.2 in Ambrosio et al., 2005 for an introduction to generalized geodesics.
Since $\gamma_1$ and $\gamma_2$ are both concentrated on $\Theta\times\Theta$, $\gamma$ must also have support in $\Theta\times\Theta\times\Theta$. Thus,

\begin{equation*}
\begin{aligned}
\support{\mu_{1/2}}
\\
\frac{1}{2}\Theta\oplus\frac{1}{2}\Theta = \Theta,
% \\
% &=\Theta,
\end{aligned}
\end{equation*}

where the last equality follows from the convexity of $\Theta$ and the definition of Minkovsky sum of sets. This shows that $\mu_{1/2}$ is feasible.
The squared Wasserstein distance from $\mu$ is known to be 2-convex along this geodesic. Thus,

\begin{equation*}
\begin{aligned}
\wassersteinDistance{2}{\mu_{1/2}}{\bar\mu}^2
\frac{1}{2}\wassersteinDistance{2}{\bar\mu_{0}}{\bar\mu}^2\!+\!\frac{1}{2}\wassersteinDistance{2}{\mu_{1}}{\bar\mu}\!-\!\frac{1}{4}\wassersteinDistance{2}{\bar\mu_0}{\bar\mu_1}^2
\\
\end{aligned}
\end{equation*}

Since $\bar\mu_0\neq\bar\mu_1$, $\wassersteinDistance{2}{\bar\mu_0}{\bar\mu_1}>0$.
However, this implies $\wassersteinDistance{2}{\mu_{1/2}}{\bar\mu}<\wassersteinDistance{2}{\bar\mu_{0}}{\bar\mu}$, which contradicts optimality of $\bar\mu_{0}$ and $\bar\mu_1$.

Equation (<ref>): We will use Kantorovich duality. Let $\nu\in\Pp{2}{\reals^d}$ with $\support{\nu}\subset\Theta$ and let $1_\Theta(x)$ be 0 in $\Omega$ and $+\infty$ outside. Clearly, $1_\Theta$ is zero $\nu$-a.e., and so $\int_{\reals^d}1_{\Theta}(x)\d\nu(x)=0$ for any $\nu\in\mathcal{A}$. Thus, Kantorovich duality (e.g., Chapter 5 in [Villani, 2009]), with $f^{c}(y)=\sup_{x\in\reals^d}f(x)-\norm{x-y}^2$, gives

\begin{equation}\label{eq:proof projection lb}
\begin{aligned}
\wassersteinDistance{2}{\nu}{\mu}^2
\int_{\reals^d}1_{\Theta}(x)\d\nu(x)-\int_{\reals^d}(1_\Theta)^{c}(y)\d\mu(y)
\\
\\
\int_{\reals^d}\inf_{x\in\Theta}\norm{x-y}^2\d\mu(x).
\end{aligned}
\end{equation}

Moreover, $\proj_\Theta[\cdot]$ is trivially a (possibly suboptimal) transport map from $\mu$ to $\pushforward{(\proj_\Theta[\cdot])}\mu$. Thus,

\begin{equation}\label{eq:proof projection ub}
\begin{aligned}
\wassersteinDistance{2}{\pushforward{(\proj_\Theta[\cdot])}\mu}{\mu}^2
\int_{\reals^d}\norm{x-\proj_{\Theta}[x]}^2\d\mu(x)
\\
\int_{\reals^d}\inf_{x\in\Theta}\norm{x-y}^2\d\mu(x),
\end{aligned}
\end{equation}

where the last equality follows from the definition of projection. We can now combine (<ref>) and (<ref>) to conclude that for all $\nu\in\Pp{2}{\reals^d}$ with $\support{\nu}\subset\Theta$

\begin{equation*}
\wassersteinDistance{2}{\nu}{\mu}
\geq
\wassersteinDistance{2}{\pushforward{(\proj_\Theta[\cdot])}\mu}{\mu}.
\end{equation*}

Since $\support{{\pushforward{(\proj_\Theta[\cdot])}\mu}}\subset\Theta$ and projections are unique, we establish (<ref>). This concludes the proof.

§.§ Proof of Theorem <ref>

We start with the proof of (<ref>). The other statements then follow.
Let $\gamma\in\setOptimalPlans{\mu_k}{\mu^\ast}$, where $\mu^\ast$ is well-defined; indeed, if $J$ is $\alpha$-geodesically convex with $\alpha>0$ and lower semi-continuous w.r.t. the convergence induced by the Wasserstein distance, a unique minimizer exists (e.g., see Section 11.2 in [Ambrosio et al., 2005]).

\begin{align*}
\\
\overset{\clubsuit}&{=}
\expectedValue{\xi_k}{\wassersteinDistance{2}{\pushforward{(\proj_{\Theta})\circ(\Id-\tau\xi_k)}\mu_k}{\mu^\ast}^2}
\\
\overset{\heartsuit}&{\leq}
\begin{aligned}[t]
\mathbb{E}^{\xi_k}&\left[\int_{\reals^d}\norm{\theta_k-\theta^\ast}^2\right. \\
\end{aligned}
\\
\overset{\square}&{=}
\expectedValue{\xi_k}{\int_{\reals^d}\norm{\proj_{\Theta}[\theta_k-\tau\xi_k(\theta_k)]-\theta^\ast}^2\d\gamma(\theta_k,\theta^\ast)}
\\
\overset{\triangle}&{\leq}
\expectedValue{\xi_k}{\int_{\reals^d}\norm{\theta_k-\tau\xi_k(\theta_k)-\theta^\ast}^2\d\gamma(\theta_k,\theta^\ast)}
\\
\begin{aligned}[t]
&+\tau^2\expectedValue{\xi_k}{\int_{\reals^d}\norm{\xi_k(\theta_k)}^2\d\gamma(\theta_k,\theta^\ast)} \\
\end{aligned}
\\
\overset{\spadesuit}&{\leq}
\begin{aligned}[t]
+\tau^2\expectedValue{\xi_k}{\int_{\reals^d}\norm{\xi_k(\theta_k)}^2\d\mu_k(\theta_k)} \\
\end{aligned}
\\
\begin{aligned}[t]
+\tau^2\expectedValue{\xi_k}{\int_{\reals^d}\norm{\xi_k(\theta_k)}^2\d\mu_k(\theta_k)} \\
\end{aligned}
\\
\begin{aligned}[t]
+\tau^2\expectedValue{\xi_k}{\int_{\reals^d}\norm{\xi_k(\theta_k)}^2\d\mu_k(\theta_k)} \\
&-2\tau\int_{\reals^d}\innerProduct{\theta_k-\theta^\ast}{\gradient{\mu} J(\mu)(\theta_k)}\d\gamma(\theta_k,\theta^\ast)
\end{aligned}
\\
\overset{\diamondsuit}&{\leq}
\begin{aligned}[t]
&\wassersteinDistance{2}{\mu_k}{\mu^\ast}^2+\tau^2(\sigma^2+C(J(\mu_k)-J(\mu^\ast))) \\&+2\tau\left(J(\mu^\ast)-J(\mu_k)-\frac{\alpha}{2}\wassersteinDistance{2}{\mu_k}{\mu^\ast}^2\right)
\end{aligned}
\\
\begin{aligned}[t]
&(1-\alpha\tau)\wassersteinDistance{2}{\mu_k}{\mu^\ast}^2 + \tau^2\sigma^2 \\ &+(J(\mu_k)-J(\mu^\ast))(C\tau^2-2\tau),
\end{aligned}
\end{align*}

where in $\clubsuit$ we used the definition of $\mu_{k+1}$; in $\heartsuit$ we used that $\pushforward{(\proj_\Theta\circ(\Id-\tau\xi_k)\times\Id)}\gamma$ is a (possibly sub-optimal) transport plan between $\pushforward{(\proj_{\Theta}\circ(\Id-\tau\xi_k))}\mu_k$ and $\mu^\ast$ (by Lemma 3.3 Aolaritei et al., 2022), i.e.,

\begin{equation*}
\pushforward{(\proj_\Theta\circ(\Id\!-\tau\xi_k)\times\Id)}\gamma
\in
\setPlans{\pushforward{(\proj_{\Theta}\!\circ(\Id\!-\tau\xi_k))}\mu_k}{\mu^\ast}
\end{equation*}

is candidate (but generally suboptimal) plan for the Wasserstein distance between $\pushforward{(\proj_{\Theta}\circ(\Id-\tau\xi_k))}\mu_k$ and $\mu^\ast$; in $\square$ we used $\int g\d\pushforward{f}\mu=\int g\circ f\d\mu$; in $\triangle$ we used non-expansivness of the projection operator (together with $\proj_{\Theta}[\theta^\ast]=\theta^\ast$ for all $\theta^\ast\in\support{\mu^\ast}\subset\Theta$); $\spadesuit$ follows from the definition of $\gamma$; and in $\diamondsuit$ we used properties of Wasserstein gradients of $\alpha$-convex functionals (e.g., see Proposition 2.8 in [Lanzetti et al., 2022]), together with Assumption <ref>.
Since by assumption $\tau\leq 2/C$, $C\tau^2-2\tau\leq 0$, and

\begin{equation*}\label{eq:proof onestep}
\expectedValue{\xi_k}{\wassersteinDistance{2}{\mu_{k+1}}{\mu^\ast}^2}
\leq
(1-\alpha\tau)\wassersteinDistance{2}{\mu_k}{\mu^\ast}^2 + \tau^2\sigma^2.%\tag{$\ast$}
\end{equation*}

We can now proceed iteratively to obtain

\begin{equation*}
\begin{aligned}
\\
\end{aligned}
\end{equation*}

This establishes (<ref>).

We now prove (<ref>). By assumption $0<\tau<1/\alpha$, $1-\alpha\tau\in(0,1)$, and so the limit $k\to\infty$, (<ref>) gives

\begin{equation*}
\limsup_{k\to\infty}\expectedValue{\{\xi_j\}_{j=0}^k}{\wassersteinDistance{2}{\mu_{k+1}}{\mu^\ast}^2}
\leq
\frac{\tau\sigma^2}{\alpha}.
\end{equation*}

By Jensen inequality, together with continuity and monotonicity of $x\mapsto x^2$, we have

\begin{equation*}
\begin{aligned}
\\
\limsup_{k\to\infty}\left(\expectedValue{\{\xi_j\}_{j=0}^k}{\wassersteinDistance{2}{\mu_{k+1}}{\mu^\ast}}\right)^2
\\
\limsup_{k\to\infty}\expectedValue{\{\xi_j\}_{j=0}^k}{\wassersteinDistance{2}{\mu_{k+1}}{\mu^\ast}^2} \leq
\frac{\tau\sigma^2}{\alpha}.
% \\
% &\leq
% \frac{\tau\sigma^2}{\alpha}.
\end{aligned}
\end{equation*}

Monotonicity $x\mapsto\sqrt{x}$ establishes (<ref>).

We now prove (<ref>), it suffices to observe that, in virtue of Gelbrich's bound ([Gelbrich, 1990]), we have

\begin{equation*}
\wassersteinDistance{2}{\mu}{\nu}
\geq
\sqrt{\norm{m_{\mu}-m_{\nu}}^2+d(S_{\mu},S_{\nu})^2}
\end{equation*}

with $m_\mu$ and and $S_{\mu}$ ($m_\nu$ and $S_{\nu}$) being the mean and covariance matrices of $\mu$ ($\nu$), and $d$ being the Bures distance on symmetric positive semidefinite matrices,  defined in (<ref>).
Thus, (<ref>) follows from (<ref>).

Finally, (<ref>) follows from <cit.>, observing that $\reals^d$ is locally compact and replacing $\varphi$ via $|\varphi|$.
This concludes the proof.

§.§ Proof of Corollary <ref>

The proof follows directly from Theorem <ref>, together with the well-known fact convergence in the Wasserstein distance is equivalent to weak convergence in $\Pp{2}{\reals^d}$; see Chapter 6 in [Villani, 2009].

§.§ Proof of Lemma <ref>

We prove the statements separately.

The proof of (<ref>) follows from Section 2 in [Lanzetti et al., 2022].
To prove (<ref>) observe

\begin{equation*}
\begin{aligned}
\expectedValue{}{\xi_k}
\expectedValue{}{W^\top \left(W\theta -\hat y_i\right)+\rho\left(\theta-\expectedValue{\mu}{\theta}\right)}
\\
W^\top \left(W\theta-\expectedValue{}{\hat y_i}\right)
\rho\left(\theta-\expectedValue{\mu}{\theta}\right)
\\
\left(W\theta-W\theta^\ast\right)
\rho\left(\theta-\expectedValue{\mu}{\theta}\right)
\gradient{\mu}J(\mu).
\end{aligned}
\end{equation*}

For the proof of (<ref>), observe that

\begin{equation*}
\begin{aligned}
\\
\expectedValue{}{\LpNorm{W^\top(W\theta-W\theta^\ast-w) + \rho\left(\theta-\expectedValue{\mu}{\theta}\right)}{2}{\reals^d,\reals^d;\mu}^2}
\\
\begin{aligned}[t]
&2\expectedValue{}{\int_{\Theta}\norm{W^\top \left(W\theta-W\theta^\ast-w\right)}^2\d\mu(\theta)}
\\
\end{aligned}
\\
\begin{aligned}[t]
&4\sigma_\mathrm{max}(W)^2\frac{1}{2}\expectedValue{}{\int_{\Theta}\norm{ W\theta-W\theta^\ast-w}^2\d\mu(\theta)}
\\
\end{aligned}
\\
4\max\{\sigma_\mathrm{max}(W)^2,\rho\} J(\mu),
\end{aligned}
\end{equation*}

where we used the definition of $J$ for the last inequality.

\begin{equation*}
\begin{aligned}
\\
4\max\{\sigma_\mathrm{max}(W)^2,\rho\} \left(J(\mu)-J(\mu^\ast) + J(\mu^\ast)\right).
\\
4\max\{\sigma_\mathrm{max}(W)^2,\rho\} \left(J(\mu)-J(\mu^\ast) + \int_{\mathcal W}\norm{w}^2\d\nu(w)\right)
\\
\begin{aligned}[t]
&4\max\{\sigma_\mathrm{max}(W)^2,\rho\} \left(J(\mu)-J(\mu^\ast)\right) \\
&+ 4\max\{\sigma_\mathrm{max}(W)^2,\rho\}\sigma^2.
\end{aligned}
\end{aligned}
\end{equation*}

This concludes the proof.

§.§ Proof of Corollary <ref>

We just need to evaluate the convexity parameters of $J$ and $\sigma^2, C$, as defined in Assumption <ref>. Since $\mu\mapsto\expectedValue{\mu}{f}$ is $\alpha$-geodesically convex if and only if $f$ is $\alpha$-convex, we have that $J$ is geodesically convex with $\alpha=\lambda_\mathrm{min}(\transpose{W}W)=\sigma_\mathrm{min}(W)^2$. Moreover, by Lemma <ref>, we have

\begin{equation*}
\begin{aligned}
\sigma^2 &=4\max\{\sigma_\mathrm{max}(W)^2,\rho\}\sigma_w^2
\\
C &=4\max\{\sigma_\mathrm{max}(W)^2,\rho\}.
\end{aligned}
\end{equation*}

Then, the result follows from Theorem <ref>, as

\begin{equation*}
\frac{\sigma^2\tau}{\alpha}=\frac{4\max\{\sigma_\mathrm{max}(W)^2,\rho\}\sigma_w^2\tau}{\sigma_\mathrm{min}(W)^2}=\eta\sigma_w^2\tau.
\end{equation*}

Also, $\alpha\leq 2C$. Thus, the condition on $\tau$ simplifies to $\tau\in (0,1/(2\max\{\sigma_\mathrm{max}(W)^2,\rho\}))$. This concludes the proof.