\section{Introduction}
Trajectory inference aims at recovering the dynamics of a population of particles given  samples from its temporal marginals at various time-points. This problem arises notably in the analysis of single-cell RNA-sequencing data~\citep{schiebinger2019optimal,tong2020trajectorynet,farrell2018single}, where one has access---via a destructive measurement process---to the cell state of samples from a large population of cells that evolve in time. In this application, one would like to recover the overall dynamics of the population as well as the trajectories of individual cells so as to improve our understanding of certain biological processes, such as embryonic development or tumor progression.

Trajectory inference can be cast as a regression problem where the unknown is the law of a continuous stochastic process. 
Let $\mathcal{X}$ be the ambient space containing the particles and $\Omega\coloneqq \mathcal{C}([0,1],\mathcal{X})$ the path-space, i.e. the set of all possible continuous trajectories over the time interval $[0,1]$.  
We consider the problem of recovering a probability distribution over paths $P\in \mathcal{P}(\Omega)$, i.e. the law of a continuous stochastic process $(X_t)_{t\in [0,1]}$. 
For the inference to be well-behaved, one should look for a distribution $P$ such that its time marginals are consistent with the observed snapshots and such that the overall dynamics it represents satisfies some notion of regularity.

The problem of trajectory inference has received significant attention over the past several years.
However, while pioneering work has focused on creating new methods, the theoretical treatment has remained limited. 
One class of methods focuses on recovering a potential energy landscape that best fits the observed marginals. For example, Hashimoto \emph{et al.}~\cite{hashimoto2016learning} encode the potential with a neural network, TrajectoryNet~\cite{tong2020trajectorynet} encodes the potential as a neural ODE, and JKOnet~\cite{bunne2021jkonet} encodes the potential with a neural network architecture based on the JKO scheme~\cite{jordan1998variational} and input convex neural networks.
However, these current approaches are all nonconvex, and therefore it can be difficult to establish rigorous guarantees. 

Some guarantees have been established for the equilibrium case by Weinreb \emph{et al.}~\cite{weinreb2018fundamental}, and recently consistency was established for the non-equilibrium case by Lavenant \emph{et al.}~\cite{lavenant2021towards}, who, as a regularity prior, penalize the entropy of $P$ relative to the Wiener measure on $\Omega$ (i.e. the law of the Brownian motion).
Moreover, Bunne \emph{et al.} have recently shown how to produce an improved reference process, beyond ordinary Brownian motion \cite{bunne2022recovering}.


In this work we focus on the estimator introduced by Lavevant \emph{et al.}~\cite{lavenant2021towards}, which reconstructs the process $P$ by minimizing the entropy relative to the Wiener measure over $\Omega$. 
By trading off data fitting with regularization, this approach generalizes the approach of Waddington-OT~\cite{schiebinger2019optimal} to datasets with many time-points yet possibly few samples per time-point, which is required for consistency~\cite{lavenant2021towards}, yields better performance in practice~\cite{schiebinger2021recovering}, and is achievable through single-embryo profiling~\cite{MITTNENZWEIG20212825}.
While the approach of Lavevant \emph{et al.} leads to a consistent estimator, it also leads to a challenging infinite dimensional convex optimization problem. It is tackled in the original paper by discretizing the space and then applying convex optimization methods. The goal of the present paper is to show that the specific structure of this estimator makes it amenable to a newly introduced class of grid-free stochastic methods called \emph{Mean-Field Langevin} (MFL) dynamics~\citep{hu2019mean, mei2018mean}, a non-linear generalization of Langevin dynamics which enjoys quantitative global convergence guarantees~\citep{nitanda2022convex, chizat2022mean}. Instantiated in our context, this method consists in a family of point clouds (one per snapshot) coupled via entropy regularized optimal transport, a.k.a.~Schr\"odinger bridges~\cite{leonard2012schrodinger}, which evolve with noisy gradient descent.
Intuitively, these dynamics can be interpreted as a (non-linear) Langevin diffusion over the path space $\Omega$.
Our approach leads to an inference method with end-to-end theoretical guarantees that solves an interpretable model for trajectory inference. We illustrated on Figure~\ref{fig:cover} the recovered estimator and on Figure~\ref{fig:dynamics} the MFL optimization dynamics.



\paragraph{Organization of the paper} In \S~\ref{sec:estimator}, we present the estimator of~\citet{lavenant2021towards}. The heart of our theoretical contributions is in \S~\ref{sec:algorithm} where we introduce the MFL dynamics and show its convergence towards the min-entropy estimator at an exponential rate. Numerical experiments are shown in \S~\ref{sec:numerics}.




\paragraph{Notation and blanket assumptions} For two probability measures $\mu,\nu$, their relative entropy is $H(\mu|\nu)=\int \log(\d\mu/\d\nu)\d\mu$ if $\mu \ll \nu$ and $+\infty$ otherwise. For $n\in \mathbb{N}$, let $[n]\coloneqq \{1,\dots,n\}$. Throughout, the ambient space $\mathcal{X}$ is a compact convex subset of $\mathbb{R}^d$ and $(B_t)_t$ is a reflected Brownian motion on $\mathcal{X}$ with $B_0$ distributed uniformly over $\mathcal{X}$. The path space is $\mathcal{C}([0,1];\mathcal{X})$ and laws on path space are noted by capital letters, $P$ or $R$. We denote by $P_{t_1,\dots, t_T} = (e_{t_1},\dots,e_{t_T})_\# P$ the marginal of such laws under the (joint) evaluation maps $e_t:\omega \mapsto \omega(t)$. We use boldface letters for family of probability measures $\boldsymbol{\mu}\in \mathcal{P}(\mathcal{X})^T$. 

\begin{figure}
    \centering
    \begin{subfigure}{0.62\linewidth}
        \includegraphics[width = \linewidth]{fig1_langevin_new.pdf}
    \end{subfigure}
    \begin{subfigure}{0.22\linewidth}
        \includegraphics[width = \linewidth]{fig1_distances.pdf}
  
    \end{subfigure}
    \caption{(left) Sampled time-series data and reconstructions by the proposed Mean-Field Langevin dynamics (MFL) and Global Waddington-OT~\cite{lavenant2021towards} (gWOT). Observed and reconstructed marginals are colored by the measurement time $t_i$, and illustrative sample paths are overlaid in grey. We compare a scenario with limited data at intermediate time-points ($N = 1$) to a uniformly sampled scenario ($N = 64$). (right) RMS Energy Distance to the ground truth over marginals for $N = 2^0, \ldots, 2^6$, illustrating the robustness of MFL to the low-sample regime. See \S\ref{sec:fig1_example} for details.}
    \label{fig:cover}
\end{figure}

\section{Min-Entropy Estimator in Path Space}\label{sec:estimator}

\subsection{Trajectory Inference as Stochastic Process Inference}
\paragraph{Model of population dynamics} The first step of any inference task is to determine a prior on the ground truth. For population dynamics with mass conservation, a natural prior is given by drift-diffusion processes. We thus model the population dynamics as
\begin{equation}\label{eq:SDE}
    \d X_t = -\nabla \Psi(t,X_t) \d t +\sqrt{\tau} \d B_t,\qquad \mathrm{Law}(X_0) = \mu_0,\qquad t\in [0,1],
\end{equation}
where $\tau>0$ is the temperature/diffusivity, assumed known, and $\Psi\in \mathcal{C}^2([0,1]\times \mathcal{X})$ is unknown. As discussed in prior works~\citep{hashimoto2016learning, lavenant2021towards}, it is hopeless to recover the divergence-free component of the drift: it is thus natural to assume that the drift is given by the gradient of a function $\Psi$, called the \emph{Waddington potential} or \emph{epigenetic landscape} in the context of cell development. The noise level $\tau>0$ models the inherent randomness of the ground truth process. It turns out that this noise also has a favorable effect on the algorithm we develop here. This model can be extended to allow particles to branch and die~\citep{baradat2021regularized}, this extension is discussed in Section~\ref{sec:mass-variation}.


\paragraph{Model of measurements} Let $0 \leq t_1<\dots <t_T\leq 1$ be a family of measurement times.
In contrast to the classical field of inference for stochastic processes~\citep{rao2000stochastic} where each realization $(X_t)_{t\in [0,1]}$ is observed $T$ times (with $T$ large), in trajectory inference each realization is only observed once, i.e.~we observe $(X_{t_i,j})_{i\in [T], j\in [n_i]}$ with independence for each couple $(i,j)$. The measurements are summarized by the family of snapshots $\hat \mu_{t_i} = \frac{1}{n_i}\sum_{j=1}^{n_i} \delta_{X_{t_i,j}} \in \mathcal{P}(\mathcal{X})$ for $i\in [T]$, where $n_i\geq 1$ is the number of particles observed at time $t_i$.

\paragraph{Inference problem} The general goal is to recover the law $P=\mathrm{Law}(X)\in \mathcal{P}(\Omega)$ of the stochastic process given the snapshots $(\hat \mu_{t_1},\dots, \hat \mu_{t_T})$. This object $P$ contains all the information about $X$: in particular it contains the marginals $P_t$ of the process and the transition probability between each family of marginals $P_{t_1,\dots,t_T}$. Note that, as discussed in the introduction, some works~\citep{tong2020trajectorynet,hashimoto2016learning,bunne2021jkonet} focus on directly on recovering the Waddington potential $\Psi$ within a parameterized class of functions. This is a different problem, with encouraging empirical results but weaker known theoretical guarantees.

\subsection{Min-Entropy Estimator}
Let $W^\tau \in \mathcal{P}(\Omega)$ be the law of the Brownian motion with temperature $\tau$. The min-entropy estimator\footnote{We propose this name because $R^*$ is the probability measure in path space with the smallest entropy relative to the Wiener measure among all those that fit at least equally well the observations.} introduced in~\cite{lavenant2021towards}, is defined as the unique minimizer $R^*$ of the functional $\mathcal{F}:\mathcal{P}(\Omega)\to \mathbb{R}$ defined as
\begin{align}\label{eq:estimator}
\mathcal{F}(R) \coloneqq \frac{1}{\lambda} \sum_{i=1}^T \Delta t_{i}\mathrm{Fit}_\sigma(R_{t_i}|\hat \mu_{t_i}) + \tau H(R|W^{\tau})
\end{align}
where $\lambda>0$ is the regularization strength, $\Delta t_i \coloneqq (t_{i+1}-t_{i-1})/2$ (with the convention $t_{-1}=0$ and $t_{T+1}=1$) and $\mathrm{Fit}_\sigma$ is a divergence functional that quantifies how much $R_{t_i}$ and $\hat \mu_{t_i}$ differ, parameterized by a bandwidth $\sigma>0$. For a particular choice of fitting functional (see below), the authors prove the following result:
\begin{theorem}[Consistency,~\cite{lavenant2021towards}]\label{thm:consistency}
If $(t_i)_{i\in [T]}$ becomes dense in $[0,1]$ as $T$ grows, then
\[
\lim_{\lambda, \sigma \to 0}\lim_{T\to \infty} R^* = P \qquad \text{weakly, almost surely.}
\]
\end{theorem}
In this paper, we consider the following data fitting term
\begin{align}\label{eq:fit}
\mathrm{Fit}_\sigma(R_t,\hat \mu_t) 
&\coloneqq \int -\log \left[ \int \exp\Big( -\frac{\Vert x-y\Vert^2}{2\sigma^2} \Big) \d R_t(x) )\right]\d\hat \mu_t(y)\\
& = H(\hat \mu_t|R_t \ast g_{\sigma}) + H(\hat \mu_t) + C\label{eq:not-always-defined}
\end{align}
where $g_{\sigma}(x) = (2\pi\sigma^2)^{-\frac{d}{2}}e^{-\Vert x\Vert^2/(2\sigma^2)}$ is the Gaussian density and $C>0$ is a constant. Eq.~\eqref{eq:fit} defines $\mathrm{Fit}_\sigma$ and is always well-defined and finite, but Eq.~\eqref{eq:not-always-defined} only makes sense when $H(\hat \mu_t)<\infty$.
This is the negative log-likelihood under the noisy observation model $\hat X_{t_i,j}=X_{t_i,j}+\sigma Z_{i,j}$ where $\hat X_{t_i,j}$ is the observation and $Z_{i,j}\overset{iid}{\sim} \mathcal{N}(0,I)$ the noise. Intuitively, when minimizing this \emph{soft-min} objective over measures $R_t$ which are mixtures of Dirac masses, each of the observed particles from $\hat \mu_{t}$ tends to attract the particles from $R_t$ that are the closest---even if they are far away in an absolute sense---but barely influences those that are the farthest.


This data-fitting functional $\mathrm{Fit}_\sigma$ slightly differs from $H(g_\sigma \ast \hat \mu_t|R_t)$, which is the one for which Thm.~\ref{thm:consistency} has been established in~\cite{lavenant2021towards}, but it preserves its key properties, namely joint convexity in $(R_t,\hat \mu_t)$ and linearity in $\hat \mu_t$. Our choice for $\mathrm{Fit}_\sigma$ is motivated by a more favorable behavior when $R_t$ is a discrete measure and by its natural statistical interpretation. In this paper, we focus on the optimization aspects and leave the statistical analysis of the estimator $R^*$ with this specific choice of $\mathrm{Fit}_\sigma$ for future work. Let us also mention that Thm.~\ref{thm:consistency} has been established for $\mathcal{X}$ a compact manifold without boundary, while under our assumptions $\mathcal{X}$ has boundaries. 

We now turn to the presentation of the algorithm. The data model of Eq.~\eqref{eq:SDE} was introduced only to motivate the min-entropy estimator via Thm.~\ref{thm:consistency}, and plays no role in the rest of the paper.


\section{Optimization Method: Mean Field Langevin Dynamics}\label{sec:algorithm}
In order to compute the min-entropy estimator $R^*\in \mathcal{P}(\Omega)$, one needs to minimize an entropy regularized functional over the space of paths $\mathcal{F}:\mathcal{P}(\Omega)\to \mathbb{R}$ which is of the form
 \begin{align}\label{eq:path-objective}
 \mathcal{F}(R) = \mathrm{Fit}(R_{t_1},\dots,R_{t_T}) +  \tau H(R|W^{\tau})
 \end{align}
where we have posed $\mathrm{Fit}(R_{t_1},\dots,R_{t_T}):=\frac{1}{\lambda} \sum_{i=1}^T \Delta t_{i}\mathrm{Fit}_\sigma(R_{t_i}|\hat \mu_{t_i})$. This is the sum of a convex functional $\mathrm{Fit}$ and the relative entropy with respect to the scaled Wiener measure. If, instead of $\mathcal{P}(\Omega)$, the optimization space was $\mathcal{P}(\mathcal{X})$, a problem with this structure could be solved by Mean-Field Langevin (MFL) dynamics. More explicitly, MFL dynamics are designed to minimize problems of the form $F(\mu) = G(\mu)+\tau H(\mu)$ where $G:\mathcal{P}(\mathcal{X})\to \mathbb{R}$ is convex and ``smooth'' and $H$ is minus the differential entropy. Considering for some $m\in \mathbb{N}^*$ the noisy gradient descent on the function $(x_1,\dots,x_m) \mapsto G(\frac1m \sum_{i=1}^m \delta_{x_i})$,  MFL dynamics are obtained in the mean-field $m\to \infty$ and vanishing step-size limit, see Section~\ref{sec:MFL} for details.

Inspired by this parallel, we now design a drift-diffusion dynamics in path space that converges to the minimizer of $\mathcal{F}$. The main idea is a reformulation of the problem as a ``reduced'' problem over $\mathcal{P}(\mathcal{X})^T$ with a structure amenable to MFL dynamics.

\subsection{Reduced Formulation}
 We first introduce the reduced functional $F$, and then state its connection to $\mathcal{F}$ in Thm.~\ref{thm:representer}.

For $\mu,\nu \in \mathcal{P}(\mathcal{X})$, let $\Pi(\mu,\nu)$ be the set of transport plans between $\mu$ and $\nu$, that is, probability measures on $\mathcal{X}\times \mathcal{X}$ with respective marginals $\mu$ and $\nu$. The entropy regularized optimal transport cost between $\mu$ and $\nu$ is defined, for some $\tau_i>0$, as
\begin{align}\label{eq:ROT}
    T_{\tau_i}(\mu,\nu) &\coloneqq \min_{\gamma \in \Pi(\mu,\nu)} \int c_{\tau_i}(x,y) \d\gamma(x,y) + \tau_i H(\gamma|\mu\otimes \nu) = \min_{\gamma \in \Pi(\mu,\nu)} \tau_i H(\gamma|p_{\tau_i}\mu\otimes \nu).
\end{align}
where $p_t(x,y)$ is the transition probability density of the Brownian motion on $\mathcal{X}$ over the time interval $[0,t]$ and $c_{\tau_i}(x,y)\coloneqq -\tau_i \log(p_{\tau_i}(x,y))$. In numerical experiments we use the approximation suggested by Varadhan's formula when $\tau_i$ is small: $c_{\tau_i}(x,y)\approx \frac12 \Vert y-x\Vert^2$ (up to an additive constant which is irrelevant when one is interested in minimizers only)~\cite{norris1997heat}.
This optimization problem is also known as the \emph{Schr\"odinger bridge problem}: its solution gives the most likely evolution of a cloud of particles following a Brownian motion, conditioned on being distributed as $\mu$ at $t=0$ and $\nu$ at $t=1$~\citep{leonard2012schrodinger}. Some background on this problem is given in Appendix~\ref{app:EOT}, including the definition of the Schr\"odinger potentials $(\varphi,\psi)$ used hereafter. Consider the function $G:\mathcal{P}(\mathcal{X})^T\to \mathbb{R}$ defined for $\boldsymbol{\mu} = (\boldsymbol{\mu}^{(1)},\dots,\boldsymbol{\mu}^{(T)})$ that represents a family of $T$ reconstructed temporal marginals, by
\begin{align}\label{eq:G-functional}
    G(\boldsymbol{\mu}) \coloneqq \mathrm{Fit}(\boldsymbol{\mu}) +  \sum_{i=1}^{T-1} \frac{1}{t_{i+1}-t_i} T_{\tau_i}(\boldsymbol{\mu}^{(i)},\boldsymbol{\mu}^{(i+1)})
\end{align}
where $\tau_i \coloneqq (t_{i+1}-t_i)\tau$.
 We now introduce the \emph{reduced objective} $F:\mathcal{P}(\mathcal{X})^T\to \mathbb{R}$, defined as 
\begin{align}\label{eq:reduced-objective}
F(\boldsymbol{\mu}) \coloneqq G(\boldsymbol{\mu}) + \tau H(\boldsymbol{\mu})
\end{align}
where $H(\boldsymbol{\mu})=\sum_{i=1}^T \int \log(\boldsymbol{\mu}^{(i)}(x))\d\boldsymbol{\mu}^{(i)}(x)$ is minus the differential entropy of the family of measures $\boldsymbol{\mu}$. 

The next result makes the link between minimizing $\mathcal{F}$, the objective in path space~\eqref{eq:path-objective}, and $F$, the reduced objective~\eqref{eq:reduced-objective}. It can be interpreted as a Representer Theorem for the min-entropy estimator.\footnote{We stretch a bit the term ``representer theorem'' which is usually reserved to finite-dimensional reductions.} This theorem is straightforward to deduce from~\cite[Cor.~3.5]{benamou2019entropy}, but the specific form of $F$ that we use here is central for our subsequent algorithmic developments.\footnote{The key difference with~\cite[Prop.~B.1]{lavenant2021towards} is that here $T_\tau$ involves the entropy of $\gamma$ relative to the product measure (instead of the volume measure), which makes $H(\boldsymbol{\mu})$ appear with a \emph{positive} sign in $F$ (instead of negative).}

\begin{theorem}[Representer Theorem]\label{thm:representer}
Let $\mathrm{Fit}:\mathcal{P}(\mathcal{X})^T\to \mathbb{R}$ be \emph{any} function.\begin{enumerate}[label=(\roman*)]
    \item If $\mathcal{F}$ admits a minimizer $R^*$ then $(R^*_{t_1},\dots,R^*_{t_T})$ is a minimizer for $F$.
    \item Conversely, if $F$ admits a minimizer $\boldsymbol{\mu}^* \in \mathcal{P}(\mathcal{X})^T$ then a minimizer $R^*$ for $\mathcal{F}$ is built as
\[
R^*(\cdot) = \int_{\mathcal{X}^T} W^\tau(\cdot| x_1,\dots,x_T )\d R_{t_1,\dots,t_T}(x_1,\dots,x_T)
\]
 where $W^\tau(\cdot|x_1,\dots,x_T)$ is the law of $W^\tau$ conditioned on passing through $x_1,\dots,x_T$ at times $t_1,\dots,t_T$ respectively and $R_{t_1,\dots,t_T}$ is the composition of the transport plans $\gamma_{i,i+1}$ which are optimal in the definition of $T_{\tau_i}(\boldsymbol{\mu}^{*(i)},\boldsymbol{\mu}^{*(i+1)})$, for $i=1,\dots,T-1$. 
\end{enumerate} 
\end{theorem}

The composition of the transport plans is obtained as
\[
R_{t_1,\dots,t_T}(\d x_1,\dots, \d x_T) = \gamma_{1,2}(\d x_1,\d x_2) \gamma_{2,3}(\d x_3|x_2) \dots \gamma_{T-1,T}(\d x_{T}|x_{T-1})
\]
where we have introduced the conditional probability (a.k.a.~``disintegrations'') characterized by $\gamma_{i,i+1}(\d x_{i},\d x_{i+1}) = \gamma_{i,i+1}(\d x_{i+1}|x_i)\mu_i(\d x_i)$. 
The proof of Theorem~\ref{thm:representer} is given in Appendix~\ref{app:proof-representer}. 

The importance of the reduced problem comes from the following facts:
\begin{itemize}
    \item The optimization space has been ``reduced'' from $\mathcal{P}(\Omega)$ to $\mathcal{P}(\mathcal{X})^T$. This reduction is enabled by the Markovian property of $W^\tau$. Moreover, Theorem~\ref{thm:representer} gives an explicit method to construct a minimizer for $\mathcal{F}$ from a minimizer for $F$ and the associated optimal transport plans $(\gamma_{i,i+1})_{i\in [T-1]}$. This fact was already exploited in~\citet{lavenant2021towards}. 
    \item The objective function $F$ is the sum of two terms: a ``smooth'' convex function $G$ and a  differential entropy term $\tau H$. This is precisely the structure tackled by MFL dynamics. Observe how the entropy in path space $H(P|W^\tau)$ is split into two parts: the Schr\"odinger bridges $T_\tau$, included in the ``smooth'' term $G$, and minus the differential entropy $H$.
\end{itemize}

Let us now describe some useful properties of $G$ and $F$. Hereafter, the \emph{first-variation} of $G:\mathcal{P}(\mathcal{X})^T\to \mathbb{R}$ at $\boldsymbol{\mu}$ is the unique function $V[\boldsymbol{\mu}]\in \mathcal{C}(\mathcal{X})^T$ such that for all $\boldsymbol{\nu} \in \mathcal{P}(\mathcal{X})^T$,
\begin{align}\label{eq:first-variation}
\lim_{\epsilon \downarrow 0} \frac1\epsilon \big(G((1-\epsilon)\boldsymbol{\mu}+\epsilon \boldsymbol{\nu}) -G(\boldsymbol{\mu})\big) = \sum_{i=1}^T \int V^{(i)}[\boldsymbol{\mu}](x) \d(\boldsymbol{\nu}-\boldsymbol{\mu})^{(i)}(x).
\end{align}

\begin{proposition}\label{prop:first-properties}
The function $G$ is convex, weakly continuous and its \emph{first-variation} is given for $\boldsymbol{\mu}\in \mathcal{P}(\mathcal{X})^T$ and $i\in [T]$ by 
\begin{align*}
    V^{(i)}[\boldsymbol{\mu}] = \frac{\delta \mathrm{Fit}}{\delta \boldsymbol{\mu}^{(i)}}[\boldsymbol{\mu}]+\frac{\varphi_{i,i+1}}{t_{i+1}-t_i} + \frac{\psi_{i,i-1}}{t_{i}-t_{i-1}},\quad \frac{\delta \mathrm{Fit}}{\delta \boldsymbol{\mu}^{(i)}}[\boldsymbol{\mu}]:x \mapsto  - \frac{\Delta t_i}{\lambda} \int \frac{g_\sigma(x-y)}{(g_{\sigma}\ast \boldsymbol{\mu}^{(i)})(y)}\d\hat \mu_{t_i}(y)
\end{align*}
where $(\varphi_{i,j},\psi_{i,j})\in \mathcal{C}^\infty(\mathcal{X})$ are the Schr\"odinger potentials for $T_{\tau_i}(\boldsymbol{\mu}^{(i)},\boldsymbol{\mu}^{(j)})$, with the convention that the corresponding term vanishes when it involves $\psi_{1,0}$ or $\varphi_{T,T+1}$.
The function $F_\tau$ admits a unique minimizer $\boldsymbol{\mu}^*$, which has an absolutely continuous density (again denoted by $\boldsymbol{\mu}^*$) characterized by
\[
(\boldsymbol{\mu}^*)^{(i)} \propto e^{-V^{(i)}[\boldsymbol{\mu}^*]/\tau}, \quad \text{for $i\in [T]$}.
\]
\end{proposition}

From now on, we thus focus on minimizing the reduced functional $F$ in Eq.~\eqref{eq:reduced-objective}, keeping in mind that this is sufficient to minimize $\mathcal{F}$, which is our main goal.



\subsection{Mean-Field Langevin Dynamics}\label{sec:MFL}
As previously mentioned, the MFL dynamics is natural when it comes to minimize functionals of the form $F=G+\tau H$. Using the first-variation $V[\boldsymbol{\mu}]$ of $G$ given in Prop.~\ref{prop:first-properties}, the MFL dynamics is defined as the solution of the following non-linear SDE of McKean-Vlasov type, for $s\geq 0$:
\begin{align}\label{eq:MFL-SDE}
\left\{
\begin{aligned}
    dX_{s}^{(i)} &= - \nabla V^{(i)}[\boldsymbol{\mu}_s](X_{s}^{(i)})\d s + \sqrt{2\tau} dB_{s}^{(i)},\quad \mathrm{Law}(X_{0}^{(i)})=\boldsymbol{\mu}_0^{(i)}\\
    \boldsymbol{\mu}_{s}^{(i)} &= \mathrm{Law}(X_{s}^{(i)}), \quad i\in [T]
    \end{aligned}
    \right.
\end{align}
with reflected boundary conditions in $\mathcal{X}$. Beware that the pseudo-time $s$ of the optimization dynamics in Eq.~\eqref{eq:MFL-SDE} should not be mistaken with the pseudo-time $t$ of the process in Eq.~\eqref{eq:SDE}, which is now represented by the discrete exponents $i\in [T]$. 

The family of laws $(\boldsymbol{\mu}_{s})_{s\geq 0}$ of this stochastic process are characterized by the following system of PDEs (understood in the sense of distributions, with no-flux boundary conditions):
\begin{align}\label{eq:MFL}
    \partial_s \boldsymbol{\mu}_{s}^{(i)} = \nabla \cdot (\nabla V^{(i)}[\boldsymbol{\mu}_s]\boldsymbol{\mu}_{s}^{(i)}) + \tau \Delta \boldsymbol{\mu}_{s}^{(i)},\quad i \in [T]
\end{align}
which are coupled with the quantity $\nabla V^{(i)}[\boldsymbol{\mu}_s]$. The link between~\eqref{eq:MFL-SDE} and~\eqref{eq:MFL} follows from Ito's lemma, see e.g.~\cite{hu2019mean}. This is a multi-species PDE where each of the species $\boldsymbol{\mu}^{(i)}$ attempts to minimize $\frac{\Delta t_i}{\lambda}\mathrm{Fit}_\sigma(\cdot|\hat \mu_{t_i}) +\tau H$ via a drift-diffusion dynamics, and is connected to $\boldsymbol{\mu}^{(i-1)}$ and $\boldsymbol{\mu}^{(i+1)}$ via Schr\"odinger bridges.


\subsection{Exponential Convergence}\label{sec:theory}
Let us now state the main convergence result proved in Appendix~\ref{app:convergence}. We note that extending this result to the non-compact case raises interesting challenges and is left as an open question.

\begin{theorem}[Convergence]\label{thm:convergence}
Let $\boldsymbol{\mu}_0\in \mathcal{P}(\mathcal{X})^T$ be such that $F(\boldsymbol{\mu}_0)<\infty$. Then there exists a unique solution $(\boldsymbol{\mu}_s)_{s\geq 0}$ to the MFL dynamics~\eqref{eq:MFL} and it holds
\[
F(\boldsymbol{\mu}_s) - F(\boldsymbol{\mu}^*) \leq e^{-Cs}\big(F(\boldsymbol{\mu}_0) - F(\boldsymbol{\mu}^*)\big).
\]
where $C:= 2\tau D^{-2} e^{-(2\mathrm{osc}(c)/\Delta t + e^{D^2/(2\sigma^2)}\Delta t/\lambda)/\tau}$ with $\Delta t \coloneqq \max_{i} \{ t_{i+1}-t_i\}$ and $D$ the diameter of $\mathcal{X}$, and $\mathrm{osc}(c) \coloneqq \max_i (\sup c_{\tau_i} - \inf c_{\tau_i})$ which, converges to $D^{2}/2$ when $\max_i \tau_i\to 0$. 
\end{theorem}

We can make the following comments:
\begin{itemize}
    \item Under these assumptions, the convergence of $\boldsymbol{\mu}_s$ to $\boldsymbol{\mu}^*$ in relative entropy and in Wasserstein distance, also holds, with the same rate~\cite{nitanda2022convex, chizat2022mean}.
    \item By Thm.~\ref{thm:representer}, one can map $(\boldsymbol{\mu}_s)_{s\geq 0}$ to a dynamics in $\mathcal{P}(\Omega)$ which converges to $R^*$. It is this dynamics which we refer to as ``Mean-Field Langevin in Path Space'' in the title.
    \item This result gives a convergence rate for an optimization dynamics in a non-convex landscape, so it is not even obvious to have global convergence. The convergence rate depends exponentially on $1/\tau$, a standard drawback of Langevin-like dynamics in absence of log-concavity. The very poor dependency in $D^2/\sigma^2$ could be mitigated by replacing the Gaussian smoothing kernel $g_\sigma$ by a kernel with a slower decay.
\end{itemize}
 
 
\paragraph{Summary of theoretical guarantees:}
Overall, the MFL dynamics of Eq.~\eqref{eq:MFL} converges to the unique minimizer $\boldsymbol{\mu}^*$ of the reduced problem in Eq.~\eqref{eq:reduced-objective}, from which we can build  by Thm~\ref{thm:representer} the min-entropy estimator $R^*\in \mathcal{P}(\Omega)$ which is optimal for~\eqref{eq:path-objective}. In this sense, our method leads to an inference method with end-to-end theoretical guarantees for trajectory inference. 
However, this statement is about an idealized dynamics which in practice has to be discretized in time and with particles. 
In the next section (\S \ref{sec:discretization}), we will see that the resulting discretization error can be controlled, leveraging the long history of work on mean-field dynamics. 

\subsection{Discretization}\label{sec:discretization}
In practice, an approximation of the MFL dynamics is obtained by running noisy gradient descent on the function $G_m:(\mathcal{X}^m)^T\to \mathbb{R}$ defined as
\begin{align}\label{eq:Gm}
    G_m(\hat X) \coloneqq G(\hat \boldsymbol{\mu}_{\hat X}) &&\text{where}&& \hat \boldsymbol{\mu}_{\hat X}^{(i)} = \frac1m \sum_{j=1}^m \delta_{\hat X_j^{(i)}}
\end{align}
where $m\in \mathbb{N}$ is be number of particles used to discretized each of the time marginals $\boldsymbol{\mu}^{(i)}$. Since it can be shown that $m \nabla_{X^{(i)}_j} G_m(\hat X) = \nabla V^{(i)}[\hat \boldsymbol{\mu}_{\hat X}](\hat X^{(i)}_{j})$ (see e.g.~\cite[Prop.~2.4]{chizat2022mean}), this leads to the following update equations, for $k\geq 0$, $i\in [T]$ and $j\in [m]$:
\begin{align}\label{eq:NPGD}
\left\{
\begin{aligned}
    \hat X_{j}^{(i)}[k+1] &= \hat X_{j}^{(i)}[k] - \eta \nabla V^{(i)}[\hat \boldsymbol{\mu}[k]](\hat X^{(i)}_{j}[k]) + \sqrt{2\eta \tau} Z^{(i)}_{j,k},\quad \hat X_{j}^{(i)}[0] \overset{iid}{\sim} \boldsymbol{\mu}_0^{(i)} \\
    \hat \boldsymbol{\mu}^{(i)}[k] &= \frac{1}{m} \sum_{j=1}^m \delta_{\hat X^{(i)}_j[k]}, \quad i\in [T]
    \end{aligned}
    \right.
\end{align}
where $\eta>0$ is a step-size. The convergence of such a numerical scheme towards the MFL dynamics is standard in the mean-field literature~\cite{sznitman1991topics}, and holds here thanks to the stability of $\nabla V$ which follows from~\cite{anonymous2022stability}. 

\begin{figure}
    \centering
    \begin{subfigure}{0.22\linewidth}
     \centering
        \includegraphics[width = \linewidth]{fig2_objective.pdf}
      
    \end{subfigure}
    \begin{subfigure}{0.77\linewidth}
        \includegraphics[width = \linewidth]{fig2_trainingdynamics.pdf}
     
    \end{subfigure}
    \caption{Optimization dynamics in the setting of Figure~\ref{fig:cover} with $N=64$. (left) Evolution of the objective function $F(\boldsymbol{\mu}_s)$ (Eq.~\eqref{eq:reduced-objective}) (right) Evolution of the reconstructed marginals $\hat \boldsymbol{\mu}^{(i)}[k]$ (Eq.~\eqref{eq:NPGD}) with the iteration number $k$, starting from isotropic Gaussians at $k=1$ and colored by measurement time $t_i$.}
    \label{fig:dynamics}
\end{figure}

In the related context of wide neural networks, quantitative bounds on the discretization error were studied in~\cite[Thm.~2]{mei2019mean}. It is shown that for $s=k\eta$, with high probability, $\vert G(\hat \boldsymbol{\mu}_k) - G(\boldsymbol{\mu}_{s})\vert = \tilde O(e^{Cs}(m^{-1/2}+\eta^{1/2}))$ for some $C>0$ independent of $k,\eta$ and $m$. Inspecting their proof, it can be seen that it goes through as long as $G(\mu)$ and $\nabla V[\mu]$ can be estimated at the ``fast rate'' from independent samples of $\mu$, which is our case here~\cite[Thm.~4.5]{delbarrio2022improved}. However, the exponential dependence in the pseudo-time $s$ of these bounds make them not precise enough to obtain interesting long-time guarantees. We leave for future work a direct convergence analysis of the fully discretized MFL dynamics (note that ~\cite[Cor.~2]{nitanda2022convex} gives a guarantee for the \emph{time}-discretized MFL dynamics).

Each iteration of Eq.~\eqref{eq:NPGD} requires to solve an entropic optimal transport problem in order to obtain $\nabla V$. As discussed in Appendix~\ref{app:sinkhorn}, this can be solved to precision $\epsilon$ in time $O(m^2/\epsilon)$ using Sinkhorn's algorithm.

\section{Numerical Experiments}\label{sec:numerics}

\subsection{Simulated data}\label{sec:fig1_example}
In our first experiment, we compare the behaviour of the MFL dynamics to that of Global Waddington-OT (gWOT)~\cite{lavenant2021towards} in a setting with few samples per time-point. Although both methods minimize a functional of the form \eqref{eq:estimator}, gWOT works on a fixed, discretized support: the union of all observed sample points. Thus, gWOT should perform poorly when the support set has missing regions or ``gaps''. To simulate this, we simulated a dataset with $N_i$ samples at time $t_i$, with $N_1 = N_{10} = 64$ and $N_i = N, 2 \leq i \leq 9$ for $N \in \{ 2^0, \ldots, 2^6 \}$. The samples are drawn from the marginals of a bifurcating stochastic differential equation (see Appendix~\ref{sec:simdetails} for details). 

In Figure \ref{fig:cover}(left) we illustrate two extreme cases: $N = 1$ (very few samples at intermediate timepoints) and $N = 64$ (uniform sampling over time) respectively, with $\lambda = 0.05, \lambda_{\mathrm{gWOT}} = 0.0025$. 
Visually, it seems that the output of MFL is robust to the few-sample regime, with relatively little qualitative difference between the reconstructed trajectories for $N = 1, 64$. On the other hand, the performance of gWOT degrades visibly once the set of support points is a poor reflection of the underlying process. 

To examine this quantitatively, we applied both MFL dynamics and gWOT for various values of $N$ and computed the root-mean-square (RMS) Energy Distance \cite{szekely2013energy} over time to an approximate ground truth (see Appendix~\ref{sec:simdetails}).
As shown in Figure \ref{fig:cover}(b), for large values of $N$ we observe that both MFL and gWOT perform similarly, but with MFL outperforming gWOT for small $N$ as expected.

Finally, Figure \ref{fig:dynamics} illustrates the evolution of the reduced objective $F$ \eqref{eq:reduced-objective} and the reconstructed marginals $\hat{\boldsymbol{\mu}}^{(i)}[k]$ over the course of MFL dynamics for $N = 64, \lambda = 0.05$. To estimate the entropy term $H(\boldsymbol{\mu})$, we used the nearest neighbour estimator of \cite[Section 6]{singh2016analysis}.

\subsection{Dealing with Branching}\label{sec:mass-variation}
In the context of single cell RNA-sequencing data it is crucial to be able to deal with the birth and death of cells, a.k.a.~branching, which does not occur homogeneously in the domain $\mathcal{X}$. Assume that we dispose of a function $g\in \mathcal{C}^1([0,1]\times \mathcal{X})$ such that $g(t,x)$ is the prior knowledge about the instantaneous growth rate of the distribution of particles in $x$ at time $t$ in the model of Eq.~\eqref{eq:SDE}. That is, in time $\d t$ the probability of a particle $X_t$ dividing is $g(t, X_t) \d t$~\cite{baradat2021regularized}. 
We would like to incorporate this knowledge, and also allow for additional mass variations to account for the inaccuracy of our prior $g$. 

To this end, we proceed heuristically by simply replacing $T_{\tau_i}(\mu_{t_i},\mu_{t_{i+1}})$ in Eq.~\eqref{eq:ROT} with the following ``unbalanced'' Schr\"odinger bridge problem 
\begin{align}\label{eq:unbalanced}
   
    \min_{\gamma \in \mathcal{M}_+(\mathcal{X}^2)} \int c_{\tau_i}(x,y)\d\gamma(x,y) + \rho H(\gamma_1|\tilde \mu_{t_i}) + \rho H(\gamma_2|\tilde \mu_{t_{i+1}}) + \tau_i H(\gamma | \tilde \mu_{t_i}\otimes \tilde \mu_{t_{i+1}})
\end{align}
where $\tilde \mu_{t_i} \propto \exp(-g(x)(t_{i+1}-t_i)/2)\mu_{t_i}$ and $\tilde \mu_{t_{i+1}}\propto \exp(g(x)(t_{i+1}-t_i)/2)\mu_{t_{i+1}}$ are probability measures, $\rho > 0$ is a parameter, $\mathcal{M}_+(\mathcal{X}^2)$ is the set of nonnegative measures on $\mathcal{X}$ and $\gamma_1,\gamma_2$ are the marginals of $\gamma$. This problem can be solved with a variant of Sinkhorn's algorithm~\cite{chizat2018scaling,frogner2015learning}.

The rationale behind this formula is as follows: (i) the modifications $\tilde \mu_{t_{i}}$ of $\mu_{t_i}$ are intended to approximate the growth process $\partial_t \hat \mu_t = g(t,\cdot) \hat \mu_t$ over the time interval $[t_i,t_{i+1}]$ and (ii) we relax the marginal constraints, with a parameter $\rho>0$, to account for the potential inaccuracy of our prior $g$. We illustrate on Figure~\ref{fig:growth} the practical advantage of taking into account branching in the algorithm. 

\begin{figure}
   
    \begin{subfigure}{0.49\linewidth}
        \includegraphics[width = \linewidth]{fig3_growth_a.pdf}
    \end{subfigure}
    \begin{subfigure}{0.49\linewidth}
        \includegraphics[width = \linewidth]{fig3_growth_b.pdf}
    \end{subfigure}
    \caption{Accounting for branching rates allows for the separation of dynamics from growth. (left) Ground truth paths and sample points. (right) Reconstruction produced by MFL dynamics with and without accounting for branching. Fraction of paths terminating in the upper (resp. lower) branch is annotated in the plot.}
    \label{fig:growth}
\end{figure}

\paragraph{Simulated data experiment} Here we consider a modified version of the bistable process from \cite[Figure 18]{lavenant2021towards} (see Appendix~\ref{sec:simdetails} for details).  
Figure \ref{fig:growth} shows paths sampled from the ground truth, sample points, and the inferred trajectories. Since the lower potential well 
is closer to the initial condition, we observe in the ground truth that $\sim 72\%$ of paths terminate in the lower branch. This effect is opposed by the high rate of branching in the upper branch, which ultimately results in more particles being sampled in the upper branch. By accounting for branching in the MFL dynamics we are able to reconstruct dynamics that isolate the effect of the potential, as evidenced by the comparable proportion (77\%) of inferred paths ending near the lower branch. On the other hand, neglecting the presence of branching results in growth effects being confounded with the potential, and the proportion of paths ending near each branch are roughly equal.

\subsection{Reprogramming dataset}

\begin{figure}
    \centering
    \begin{subfigure}{0.745\linewidth}
        \includegraphics[width = \linewidth]{reprogramming_snapshots.pdf}
    \end{subfigure}
    \begin{subfigure}{0.245\linewidth}
        \includegraphics[width = \linewidth]{reprogramming_distances.pdf}
    \end{subfigure}
    \caption{(left) Subsampled timepoints (Subsample) from the iPSC reprogramming dataset of \cite{schiebinger2019optimal}, with 100 cells at the first and last timepoints and 10 cells at each intermediate timepoint, shown beside inferred marginals obtained using MFL dynamics and the full dataset for those timepoints. (right) Best performance as measured by Energy Distance of inferred marginals to the full dataset for MFL dynamics (Langevin) and Global Waddington-OT~\cite{lavenant2021towards} (gWOT), with the subsampled timepoints (Sample) for reference.}
    \label{fig:reprogramming}
\end{figure}


We now consider the stem cell reprogramming dataset of \cite{schiebinger2019optimal}, in which a population of differentiating cells were profiled over a time course using single-cell RNA sequencing. For the purpose of this example, we restrict our attention to days 2.5-6 inclusive making a total of 8 timepoints. As previously, we consider a regime where few samples are observed as snapshots of the time-series and we apply MFL dynamics to infer reconstructed marginals. We reason that if the underlying process is well described by a Schr\"odinger bridge (as was empirically validated in \cite{schiebinger2019optimal}), then MFL dynamics should be able to ``improve'' on the sample marginals. As a proxy for ground truth, we used the full dataset (consisting of 59154 cells over the 8 timepoints). From this, we subsampled timepoints consisting of 100 cells at days 2.5 and 6, and 10 cells at the remaining intermediate timepoints. After carrying out a series of preprocessing steps (see Appendix~\ref{sec:preprocessing} for details), we applied MFL dynamics to the subsampled data to produce reconstructed marginals. Figure \ref{fig:reprogramming}(left) shows an example of the sampled points, MFL reconstruction, and the full dataset. 

A reconstruction ``error'' was then computed as the Energy Distance \eqref{eq:energydistance} of each reconstructed marginal to the corresponding snapshot in the full dataset. Figure \ref{fig:reprogramming}(right) shows the error for each timepoint for each of MFL and gWOT, for the best performing parameters found in a parameter sweep. We observe that both MFL and gWOT improve upon the subsampled snapshots, but MFL by a larger margin. 

\section{Conclusion}
We introduced a grid-free numerical method for trajectory inference that computes the min-entropy estimator introduced in~\cite{lavenant2021towards}, with global and quantitative convergence guarantees in the mean-field limit. This method arises naturally when decomposing the optimization problem in a suitable way and, in practice, outperforms the fixed-grid method of~\cite{lavenant2021towards}. 

Concerning limitations, our method shares those of the min-entropy estimator : since it does not incorporate fine structural prior on the structure of the Waddington potential (i.e. it is fully nonparametric), it may suffer from a limited statistical efficiency in high dimension. An interesting research direction is to quantify this statistical performance and to adapt our algorithm to learn the Waddington potential $\Psi$ in a structured parameterized class of functions, jointly with the law on paths $P$.

On a more abstract level, our main insight is that min-entropy problems in Wiener space can be solved via a multi-species diffusion dynamics coupled with Schr\"odinger bridges. This point of view raises interesting theoretical questions, e.g.~can we rigorously interpret our dynamics as a diffusion in path space?   

\section*{Acknowledgements}

GS was supported by a Career Award at the Scientific Interface from the Burroughs Wellcome Fund, a NFRF Exploration Grant, a NSERC Discovery Grant, and a CIHR Project Grant. 

{
\small

