Dynamical Persistent Homology via Wasserstein Gradient Flow

Gradient flow in Hilbert space (Ambrosio et al., 2008) is a fundamental concept in the analysis of variational problems and partial differential equations. It describes the evolution of a function in a Hilbert space under the influence of its gradient, leading to a minimization of the associated energy functional.

Let $\mathcal{H}$ be a Hilbert space with inner product $\langle\cdot,\cdot\rangle$ and norm $\|\cdot\|$. Consider an energy functional $E:\mathcal{H}\to\mathbb{R}$, which is typically assumed to be Fréchet differentiable. The gradient flow of $E$ is the solution to the differential equation

where $u(t)\in\mathcal{H}$ represents the state of the system at time $t$, and $\nabla E(u)$ denotes the gradient of $E$ at $u$.

The gradient $\nabla E(u)$ is defined by the property that for all $v\in\mathcal{H}$,

This ensures that the direction of the gradient $\nabla E(u)$ is the direction of the steepest ascent of the functional $E$. Consequently, the negative gradient $-\nabla E(u)$ points in the direction of the steepest descent, which is why it appears in Equation (1).

One of the key properties of gradient flow is that it decreases the energy functional over time. Specifically, if $u(t)$ is a solution to Equation (1), then

This implies that $E(u(t))$ is non-increasing along the curve $u(t)$. Moreover, since $\frac{d}{dt}E(u(t))=0\Leftrightarrow\|\nabla E(u(t))\|=0$, if $E$ has a unique stationary point $u^{\star}$, $u(t)$ converges to it as $t\to\infty$ under appropriate conditions.

From a computational perspective, discretization schemes are essential for approximating gradient flows in Hilbert spaces. The implicit Euler scheme, also known as the minimizing movement scheme or proximal map (Bubeck et al., 2015), is particularly effective for this purpose. This scheme approximates the continuous-time gradient flow by a sequence of discrete minimization problems:

where $\tau>0$ is the time step, $u_{k}$ represents the approximation at the $k$-th time step, $E$ is the energy functional, and $\mathcal{H}$ is the Hilbert space. This scheme not only provides a practical method for numerical computations but also offers insights into the theoretical properties of gradient flows, such as their connection to proximal operators and convex analysis.

Gradient flow in Wasserstein space extends the concept of gradient flow from Hilbert spaces to the space of probability measures. This framework is particularly valuable for studying partial differential equations and variational problems involving mass transport. While this subsection provides a brief introduction, readers are encouraged to consult Ambrosio et al. (2008) for a more comprehensive treatment.

The Wasserstein space is defined as the space of probability measures with finite second moments, equipped with the Wasserstein distance. More formally, let $(X,d)$ be a metric space. The Wasserstein space of order $p$ over $(X,d)$, denoted by $\mathcal{P}_{p}(X)$, is the set of all probability measures $\mu$ on $X$ with finite $p$-th moment:

where $x_{0}$ is a fixed reference point in $X$. This definition ensures that the measures in $\mathcal{P}_{p}(X)$ have well-behaved moments, which is crucial for various analytical properties.

The Wasserstein distance of order $p$ between two probability measures $\mu,\nu\in\mathcal{P}_{p}(X)$ is defined as:

where $\Pi(\mu,\nu)$ represents the set of all joint probability measures on $X\times X$ with marginals $\mu$ and $\nu$. This distance metric quantifies the optimal cost of transporting mass from one probability distribution to another, providing a natural way to compare and analyze probability measures in the context of gradient flows.

With a series of works by Benamou and Brenier (2000); Otto (2001); Ambrosio et al. (2008), the concept of gradient flow in Wasserstein space has been rigorously established.

Consider an energy functional $J:\mathcal{P}_{2}(\mathbb{R}^{d})\to\mathbb{R}$ defined on a Wasserstein space, where $\mathcal{P}_{2}(\mathbb{R}^{d})$ denotes the space of probability measures with finite second moments as in Equation (5). A curve $(\mu_{t})_{t\geq 0}$ of probability measures is called a Wasserstein gradient flow for the functional $J$ if it satisfies the following continuity equation in a weak sense:

Here, $\frac{\delta J}{\delta\mu}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ represents the first variation of $J$ at $\mu$, defined by

This first variation can be interpreted as the “functional derivative” of $J$ with respect to $\mu$, extending the concept of derivatives to measure spaces. Consequently, $\nabla\frac{\delta J}{\delta\mu}(\mu_{t}):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ for $t\geq 0$ represents a time-dependent family of vector fields. Chewi et al. (2020) refer to this as the Wasserstein gradient, denoted as $\nabla_{W_{2}}(\mu_{t}):=\nabla\frac{\delta J}{\delta\mu}(\mu_{t})$. This term draws an analogy with the standard gradient in Euclidean spaces due to its similar role in describing the flow of probability measures. Additionally, some literature defines the Wasserstein gradient as $\nabla_{W_{2}}(\mu_{t}):=-\nabla\cdot\left(\mu_{t}\nabla\frac{\delta J}{\delta\mu}(\mu_{t})\right)$ to maintain consistency with Equation (1).

The JKO scheme, named after Jordan, Kinderlehrer, and Otto, is a time-discretization method for approximating the Wasserstein gradient flow (Jordan et al., 1998), It involves solving a sequence of minimization problems:

where $\tau$ is the time step. The transition from the continuous Wasserstein gradient flow to the discrete JKO scheme can be understood through the concept of minimizing movement which approximates the continuous flow by discrete steps. The convergence of the JKO scheme to the continuous Wasserstein gradient flow as $\tau\to 0$ is established through $\Gamma$-convergence (Ambrosio et al., 2008). The functional $J(\mu)$ often represents the internal energy or potential energy of the system, and its gradient with respect to the Wasserstein metric drives the evolution of the distribution (Villani et al., 2009; Villani, 2021). The JKO scheme is particularly useful in numerical simulations of diffusion processes and other phenomena described by partial differential equations (PDEs) in the Wasserstein space (Santambrogio, 2015). Applications of the JKO scheme include image processing, machine learning, and fluid dynamics, where the evolution of distributions is of interest (Peyré et al., 2019).

Nigmetov and Morozov (2024) introduces a novel method that directly uses the cycles and chains involved in the persistence computation to define gradients for larger subsets of the data. This contrasts with prior methods that focused solely on individual PD points. The authors also demonstrate that this approach can significantly reduce the number of steps required for optimization. Our work builds upon this idea, allowing it to adapt data based on variations in PDs along Wasserstein gradient flows.

In this subsection, we assume that the target persistence diagram is known a priori. Our objective is to develop a learnable process for generating persistence diagrams guided by a dynamical system. To approach this problem, we draw upon two fundamental theorems from optimal transport theory: the Benamou-Brenier Theorem and the McCann Interpolation Theorem. The Benamou-Brenier Theorem provides a dynamic formulation of optimal transport, allowing us to view the transport problem as a fluid flow minimizing a kinetic energy functional. Meanwhile, the McCann Interpolation Theorem offers a way to construct geodesics in the space of probability measures, which is crucial for understanding the geometry of persistence diagrams. By leveraging these theorems, we can formulate a continuous evolution of persistence diagrams that converges to the target diagram while respecting the underlying topological structure of the data.

The Benamou-Brenier theorem shows that the optimal transport cost can be expressed as the infimum of the action integral over all possible time-dependent probability measures $\mu_{t}$ and velocity fields $v_{t}$ that transport $\mu_{0}$ to $\mu_{1}$.

On the other hand, McCann’s interpolation theorem describes the displacement interpolation between two probability measures in the context of optimal transport, showing that the interpolation is a geodesic in the Wasserstein space.

Therefore, we have the following remark.

The McCann interpolation can be viewed as a specific instance of a dynamical process. Based on this insight, we have developed an algorithm that leverages McCann interpolation to guide the optimization within the context of persistence diagrams.

In the Algorithm 1, the $\operatorname*{OptimalTransportPlan}$ can be efficiently computed using the Sinkhorn algorithm (Cuturi, 2013), known for its rapid convergence and computational efficiency. Alternatively, conventional linear programming methods can also be employed. The optimal transportation plan $\pi$ is represented as an $n\times m$ matrix, where $n$ and $m$ denote the cardinalities of the source persistence diagram $X^{(k)}$ and the target persistence diagram $Z$, respectively. Each entry $\pi_{ij}$ in the matrix indicates the amount of mass to be transported from $x_{i}\in X^{(k)}$ to $z_{j}\in Z$. The mass associated with each point $x_{i}\in X^{(k)}$ and $z_{j}\in Z$ is $1/n$ and $1/m$, respectively.

In cases where $n$ and $m$ are not equal, the $\operatorname*{TargetDgm}$ function computes an appropriate target persistence diagram. This is achieved by using the barycenter to determine the target locations, calculated as follows:

The Algorithm 1 demonstrates the use of McCann interpolation to guide optimization on persistence diagrams. The core idea is to obtain the next persistence diagram through McCann interpolation and employ a learnable scheme, specifically the critical set method (Nigmetov and Morozov, 2024), to determine the subsequent filtration. This approach allows for the construction of the entire dynamical process of the persistence diagram.

The results $X^{(k)},Y^{(k)},f^{(k)}$ represent the persistence diagram, the next target persistence diagram, and the filtration at step $k$, respectively. We explicitly store $Y^{(k)}$ to maintain the next step’s persistence diagram because the learnable scheme cannot guarantee that the persistence diagram is always achievable (Nigmetov and Morozov, 2024). This is also why we do not compute the full trajectories before learning the filtration. Instead, we compute the McCann interpolation at each step to ensure that the persistence diagram progresses towards the target persistence diagram. In other words, this algorithm optimizes the filtration towards the final target rather than intermediate states.

The proposed algorithm is particularly well-suited for applications where a target persistence diagram is explicitly defined. Additionally, in the context of McCann interpolation, our algorithm guarantees that the measure $\mu_{t}$ evolves along the geodesic in Wasserstein space, provided that the target persistence diagram is achievable at each step. This property ensures that the topological features of the data, as captured by persistence diagrams, change towards the target diagram in the best possible manner.

The previous subsection demonstrated the dynamical persistence diagrams along McCann interpolation trajectories, given a target persistence diagram. In this subsection, we extend our methodology to scenarios without a predefined target diagram. As introduced in Subsection 3.3, a dynamical system describes the temporal evolution of a state according to a set of rules or equations. In the context of Wasserstein gradient flow, the probability measure $\mu_{t}$ evolves following the gradient flow of an energy functional $J$ in the Wasserstein space. This evolution follows a path that locally minimizes $J$ most efficiently with respect to the Wasserstein distance. The choice of functional $J$ determines the system’s ultimate configuration, allowing for diverse outcomes depending on the specific energy landscape.

To perform the dynamical persistence diagram through Wasserstein gradient flows, we propose the following Algorithm 2.

Algorithm 2, akin to Algorithm 1, outputs $X^{(k)}$, $Y^{(k)}$, and $f^{(k)}$, which represent the current persistence diagram, the target persistence diagram, and the filtration at step $k$, respectively. The key distinction lies in Algorithm 2’s utilization of the JKO scheme to compute the target persistence diagram at each iteration. Importantly, we initialize the target persistence diagram with the current persistence diagram, ensuring equal cardinalities. This approach enables us to leverage the energy functional to guide the optimization process without prior knowledge of the final target persistence diagram.

Our first illustration concerns the denoising of a circle. The input data consists of a set of points in a 2D plane, sampled from a circle with added Gaussian noise. This noisy data forms a perturbed circle, as depicted in Figure 1. The objective is to remove the noise and highlight the primary 1-dimensional topological feature, namely the circle itself.

Given a set of random 2D points on a plane, we observe that the first persistence diagram is not always empty. In this case, the objective is to enhance the circles represented by the first persistence diagram. In the experiment, we independently and identically sample 250 points from the uniform distribution on the square $[-1,1]\times[-1,1]$. Similar to the previous case, we apply the McCann algorithm with a target of $(0,0.08)$ on the 0th persistence diagrams. Additionally, we apply Algorithm 2 with the functional

to encourage the first persistence pairs above a certain threshold to evolve away from the diagonal line and drift to the left. The results of this process are illustrated in Figure 8. We also present 0th and 1st persistence diagrams in Figure 9 and Figure 10, respectively. The results demonstrate how the persistence diagrams evolve over time, highlighting the emergence of the circle from the random 2D points.