Active Timepoint Selection for Learning Measure-Valued Trajectories

Technical Summary: Active Timepoint Selection for Learning Measure-Valued Trajectories

Problem Statement
The paper addresses the challenge of inferring continuous probability paths $\{\mu_t\}_{t \in [0,1]}$ from sparse, destructive snapshots. This problem is critical in domains like single-cell biology, where high-fidelity data acquisition (e.g., sequencing) is costly and destructive, preventing dense temporal sampling. The goal is to perform distributional interpolation: recovering the underlying continuous trajectory given a finite set of observed empirical measures at discrete time points.

The authors identify three primary challenges in applying active learning to this setting:

1. Non-Euclidean Geometry: Probability measures reside in the Wasserstein space $\mathcal{P}_2(\mathcal{X})$, which is non-linear. Standard active learning methods relying on Euclidean metrics (e.g., standard Gaussian Processes) fail because linear combinations of measures do not yield valid measures, and naive interpolation splits mass rather than transporting it.
2. Lack of Epistemic Uncertainty: Current distribution interpolation methods (e.g., flow matching, Schrödinger bridges) are typically deterministic. They return a single fitted path and lack a mechanism to quantify epistemic uncertainty, which is essential for guiding active learning policies.
3. Non-Stationarity: Measure-valued dynamics often exhibit varying speeds (e.g., periods of homeostasis punctuated by rapid branching events). Uniformly spaced acquisition times are suboptimal for capturing these transient, high-velocity regions.

Methodology
The authors propose a framework that extends active experimentation to the space of measures by constructing a tractable probabilistic surrogate. The method consists of four key stages:

1. Linearization via Linearized Optimal Transport (LOT):
   To handle the non-Euclidean geometry, the framework maps observed empirical measures $\hat{\mu}_t$ into a common tangent space $T_\sigma \mathcal{P}_2(\mathcal{X})$ defined at a fixed reference measure $\sigma$ (chosen as the Wasserstein barycenter of observed snapshots). Using the LOT logarithmic map, each measure is represented as a displacement vector field $v_\mu(x) = T_{\sigma \to \mu}(x) - x$, where $T$ is the optimal transport map. This linearizes the Wasserstein geometry, allowing measures to be treated as vectors in a Hilbert space.

2. Low-Dimensional Representation:
   The high-dimensional displacement fields are compressed into a low-dimensional latent space. The authors flatten the reweighted displacement fields and apply Principal Component Analysis (PCA) to obtain a $K$-dimensional representation $c_t$ for each time point $t$.

3. Probabilistic Modeling with Gaussian Processes (GP):
   A Multi-Output Gaussian Process (MOGP) is placed over the temporal evolution of the latent coefficients $c_t$. The model assumes a zero-mean prior with a matrix-valued kernel. The GP posterior provides a predictive mean and, crucially, a covariance matrix that quantifies epistemic uncertainty over the latent trajectory.

   • Reconstruction: To recover the predicted measure $\hat{\mu}_t$ at any time $t$, the latent mean is mapped back through the inverse PCA projection and the inverse LOT map (adding the displacement to the reference landmarks).

4. Handling Non-Stationarity via Time Warping:
   To address the non-stationary nature of biological processes, the authors introduce an intrinsic time parametrization. A warping function $\Phi(t)$ maps physical time to an intrinsic time $\tau$, representing the cumulative arc-length of the trajectory in Wasserstein space. A stationary kernel is applied in this warped domain, effectively inducing a non-stationary kernel in physical time that adapts to local speeds of distributional change.

5. Acquisition Policy:
   The framework employs an acquisition function to select the next measurement time $t^*$. The policy iteratively selects $t^*$ to maximize information gain, specifically targeting regions of high epistemic uncertainty in the intrinsic time domain. Two acquisition strategies are proposed: point-wise uncertainty (maximizing trace of posterior covariance) and expected integrated risk reduction.

Key Contributions

• Conceptual: The paper formulates the problem of active learning for measure-valued trajectories, extending active experimentation from Euclidean spaces to the space of probability measures.
• Technical: It constructs a tractable probabilistic surrogate in Wasserstein space by combining Linearized Optimal Transport with Multi-Output Gaussian Processes. This combination enables the quantification of epistemic uncertainty over probability paths, a capability missing in prior distribution interpolation methods.
• Empirical: The authors demonstrate that their acquisition strategy outperforms uncertainty-agnostic baselines (Uniform and Random sampling) on both synthetic datasets mimicking branching events and real-world single-cell transcriptomics data.

Results

• Synthetic Data: On a synthetic dataset with non-stationary dynamics and branching events, the active strategy consistently outperformed Uniform and Random baselines, particularly at low-to-moderate budgets. The method successfully concentrated measurements in high-velocity regions (branching events), reducing the Mean Wasserstein Error and velocity-weighted error. Sensitivity analysis showed the method is most effective when branching events are highly localized.
• Real-World Data: Applied to a single-cell reprogramming dataset (mouse fibroblasts to iPSCs), the active strategy achieved the lowest reconstruction error in the low-to-moderate budget regime ($B \le 12$). The advantage diminished as the budget increased and the sampling became dense, confirming the method's utility in resource-constrained scenarios.
• Ablation Studies: Removing intrinsic time warping degraded performance, validating the need for non-stationary modeling. Reducing the latent dimension ($K=2$) caused significant performance drops, highlighting the need for sufficient dimensionality to capture complex variations. Fixing the reference measure (rather than using the barycenter) increased error, emphasizing the importance of minimizing linearization error.

Significance and Claims
The paper claims that its primary contribution is enabling active learning in the space of measures by solving the dual problems of non-Euclidean geometry and the lack of uncertainty quantification. By leveraging LOT and GPs, the framework provides a practical mechanism to select optimal measurement times that minimize uncertainty in the reconstructed probability path.

The authors note limitations: the approach relies on a tangent space approximation, which may introduce errors under large distributional shifts, and the computational cost is dominated by Optimal Transport subroutines, making it practical for the intended regime (up to $\sim 10^5$ samples per snapshot) but potentially challenging for million-point datasets without further scalable approximations. Future work is suggested to explore Riemannian geometry tools (multiple charts) and multidimensional inputs.