Active Timepoint Selection for Learning Measure-Valued Trajectories

Mihaela van der Schaar; Nicolas Huynh

arxiv: 2605.30625 · v1 · pith:MEM2K4LTnew · submitted 2026-05-28 · 💻 cs.LG · cs.AI· stat.ML

Active Timepoint Selection for Learning Measure-Valued Trajectories

Nicolas Huynh , Mihaela van der Schaar This is my paper

Pith reviewed 2026-06-29 08:34 UTC · model grok-4.3

classification 💻 cs.LG cs.AIstat.ML

keywords active learninglinearized optimal transportGaussian processesmeasure-valued trajectoriesuncertainty quantificationprobability pathsacquisition functionssingle-cell data

0 comments

The pith

Linearized optimal transport embeds measures so Gaussian processes can select informative time points for trajectory inference

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper sets out to make active learning work for inferring continuous paths of probability distributions when only sparse, expensive snapshots are available. It achieves this by using Linearized Optimal Transport to place distributional data in a tangent space where ordinary Gaussian process regression can supply epistemic uncertainty. The resulting surrogate model then drives an acquisition rule that picks the next measurement time to shrink uncertainty most. If the approach holds, experimenters can recover accurate measure-valued trajectories with fewer destructive or costly observations.

Core claim

By mapping distributional snapshots into a tangent space via Linearized Optimal Transport, the authors construct a Gaussian process surrogate for the underlying probability path; this surrogate directly supplies the uncertainty estimates required to define an acquisition policy that iteratively chooses measurement times minimizing uncertainty in the inferred trajectory.

What carries the argument

The Linearized Optimal Transport embedding, which places probability measures in a Euclidean tangent space so that standard Gaussian process regression can quantify uncertainty over the full trajectory.

If this is right

An acquisition function can be written that selects the time minimizing predictive variance in the Gaussian process surrogate.
Trajectory estimates improve over uncertainty-agnostic baselines on both synthetic and real datasets.
Active experimentation becomes feasible on the infinite-dimensional space of probability measures.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

Similar linearization tricks could bring uncertainty-aware experimental design to other non-Euclidean data spaces.
The method may reduce the total number of snapshots needed in longitudinal studies of cell populations.
One could check whether the selected times coincide with intervals of rapid change in the true underlying path.

Load-bearing premise

The Linearized Optimal Transport tangent space lets Gaussian process uncertainty estimates remain reliable when mapped back to the original Wasserstein geometry of probability measures.

What would settle it

On a synthetic trajectory where the full continuous path is known in advance, the times chosen by the uncertainty-driven policy produce no lower reconstruction error than times chosen by a non-adaptive baseline.

Figures

Figures reproduced from arXiv: 2605.30625 by Mihaela van der Schaar, Nicolas Huynh.

**Figure 1.** Figure 1: Gaussian Process regression naively applied to densities leads to poor interpolation. While interpolation schemes compatible with the Wasserstein geometry exist, they typically operate between only two reference measures. An example is the displacement interpolation (McCann, 1997). Given µ0, µ1 in P2(X ), it is defined by: µt = ((1 − t)Id + tT)#µ0, (3) where T is the optimal transport map from µ0 to µ1 (s… view at source ↗

**Figure 2.** Figure 2: Overview of the methodology. (Left) Probability measures µ, ν in Wasserstein space P2(X) are projected onto the tangent plane TσP2(X) via Linearized Optimal Transport (LOT). (Right) The active learning loop maps snapshots to latent states ci modeled by Gaussian Processes. This surrogate quantifies epistemic uncertainty to select the optimal next measurement time t ∗ . dimensionality reduction using the cor… view at source ↗

**Figure 3.** Figure 3: (Left) Visualization of the synthetic data projected in 2D. The trajectory is non-stationary with two distinct branching events (marked). (Right) Reconstruction performance as a function of the acquisition budget. We report the mean Wasserstein error and its velocity-weighted variant (w-W2). The results are averaged over 5 seeds, and the vertical axes are presented on a logarithmic scale. Surrogate model. … view at source ↗

**Figure 5.** Figure 5: Evaluation on the single-cell reprogramming dataset from (Schiebinger et al., 2019) 5.3. Application to real-world datasets Data. We evaluate our method on the large-scale singlecell RNA sequencing dataset from Schiebinger et al. (2019), which tracks the reprogramming of mouse fibroblasts into induced pluripotent stem cells (iPSCs). We specifically focus on the serum culture subset, which exhibits non-st… view at source ↗

**Figure 6.** Figure 6: Ablation study. We evaluate the impact of: (i) replacing the Matern-5/2 kernel with an ´ RBF kernel; (ii) fixing the LOT reference σ; (iii) reducing the PCA basis rank to K = 2; and (iv) disabling the intrinsic time warping strategy (No warp). Results. We report the results for both the synthetic and the fibroblast reprogramming dataset in [PITH_FULL_IMAGE:figures/full_fig_p009_6.png] view at source ↗

**Figure 7.** Figure 7: Reconstruction performance as a function of the acquisition budget. We report the mean Wasserstein error and its velocityweighted variant (w-W2). C.4. Labor market microdata Data. We conduct an experiment using real-world IPUMS-CPS monthly microdata (Flood et al., 2024). Specifically, we use monthly non-ASEC U.S. CPS samples from January 2015 to December 2021 and construct a time-indexed sequence of distr… view at source ↗

**Figure 8.** Figure 8: We compare the pairwise squared Wasserstein 2 distances (d 2 ij ) with the pairwise squared distances based on LOT embeddings (dˆ2 ij,σ) for two choices of the reference σ: using the Wasserstein barycenter of the observed snapshots vs. using the first snapshot [PITH_FULL_IMAGE:figures/full_fig_p018_8.png] view at source ↗

**Figure 9.** Figure 9: Additional results for the IPUMS-CPS dataset. 19 [PITH_FULL_IMAGE:figures/full_fig_p019_9.png] view at source ↗

read the original abstract

Inferring continuous probability paths from sparse snapshots is a fundamental challenge in domains like single-cell biology, where high-fidelity data acquisition is often destructive and constrained by prohibitive sequencing costs. This motivates the need for active learning strategies to strategically select optimal measurement times. However, designing active learning policies for this setting remains an open problem: the target objects reside on the infinite dimensional Wasserstein space where standard Euclidean metrics are ill-defined, and current interpolation methods lack epistemic uncertainty quantification. We introduce a framework which extends active experimentation to the space of measures. By leveraging Linearized Optimal Transport (LOT), we map distributional snapshots into a tangent space amenable to Gaussian Process modeling, allowing us to construct a tractable probabilistic surrogate for the underlying probability path. This yields an acquisition policy that iteratively selects measurement times to minimize uncertainty. Empirical results demonstrate that our strategy outperforms uncertainty-agnostic baselines on both synthetic and real-world datasets.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

LOT embedding plus GP active learning for measure trajectories is a sensible idea on paper but the local linearization around one reference measure likely undermines the uncertainty quantification that the acquisition policy needs.

read the letter

The main point is that the paper sketches a way to do active timepoint selection for learning trajectories of probability measures by pushing snapshots through Linearized Optimal Transport into a tangent space, fitting a Gaussian process there, and using the posterior variance to pick the next measurement time. That combination is presented as new for this setting.

It does address a concrete practical need in single-cell work where each distributional snapshot costs a lot, and the abstract reports that the resulting policy beats simple baselines on both synthetic and real data. That empirical signal is worth noting even if the details are thin.

The soft spot is the one flagged in the stress-test note. LOT builds its Euclidean tangent space by linearizing around a single fixed reference measure. Any part of the trajectory that moves away from that reference will see its variation distorted by the first-order approximation, and the GP variance fitted in that space can then reflect the approximation error rather than genuine uncertainty in the Wasserstein geometry. The acquisition rule is only as good as that variance, so if the linearization error grows, the selected times may not actually reduce posterior uncertainty on the path. The abstract gives no indication that this was checked, and without equations or validation experiments it is impossible to tell whether the method holds up.

This is aimed at people working at the intersection of active learning, optimal transport, and computational biology. A reader already thinking about uncertainty-aware interpolation on measures might pick up the high-level framing and the reported gains, but anyone wanting to use or extend the method will need the full technical details.

I would send it to peer review. The problem is real, the proposed direction is coherent on its own terms, and the empirical claim is at least stated, even though the current version is too light on verification to stand alone.

Referee Report

2 major / 2 minor

Summary. The paper proposes an active learning framework for inferring continuous measure-valued trajectories from sparse distributional snapshots. It maps snapshots into a Euclidean tangent space via Linearized Optimal Transport (LOT), fits a Gaussian Process surrogate in that space to obtain epistemic uncertainty, and uses the resulting posterior variance to define an acquisition function that iteratively selects measurement times minimizing uncertainty on the underlying probability path. The method is claimed to outperform uncertainty-agnostic baselines on both synthetic and real-world datasets, with motivation from single-cell biology applications.

Significance. If the LOT embedding supplies reliable epistemic uncertainty for Wasserstein trajectories, the work would provide a tractable extension of active learning to infinite-dimensional measure spaces, addressing a genuine gap where standard Euclidean GPs cannot be applied directly. The empirical outperformance on real datasets would constitute a concrete advance for costly data-acquisition settings. However, the local character of the linearization around a single reference measure constitutes a load-bearing assumption whose validity is not yet demonstrated for trajectories that move far from that reference.

major comments (2)

[Methods section describing the LOT embedding and probabilistic surrogate] The central claim that the GP posterior variance in the LOT tangent space yields a reliable acquisition policy for the Wasserstein path rests on the unexamined assumption that first-order linearization errors remain negligible. Because LOT is constructed around one fixed reference measure μ₀, any segment of the trajectory whose support or mass distribution deviates substantially from μ₀ projects its variation through an approximation whose error grows with distance; the fitted GP length-scale and variance can therefore encode linearization artifacts rather than true Wasserstein variability. This directly affects the correctness of the uncertainty-minimizing acquisition rule. The manuscript should either (a) provide a quantitative bound on the linearization error along the learned trajectory or (b) demonstrate empirically that the selected time points remain optimal under a non-linearized
[Experimental results and baseline comparisons] The empirical validation compares the proposed policy only against uncertainty-agnostic baselines. To substantiate that the LOT-GP uncertainty is the operative ingredient, an ablation that replaces the LOT-GP variance with a heuristic or random acquisition function while keeping the same embedding should be reported; without it, the performance gain could be attributable to the embedding alone rather than to the active selection mechanism.

minor comments (2)

Notation for the reference measure and the tangent-space coordinates should be introduced once and used consistently; the current description occasionally switches between μ₀ and μ_ref without explicit cross-reference.
The real-world dataset description should include the number of distributional snapshots, their temporal spacing, and the dimensionality of the underlying space so that readers can assess how far the trajectory travels from the chosen reference measure.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive comments, which highlight important aspects of the LOT linearization assumption and the need to isolate the contribution of uncertainty quantification. We respond to each major comment below and indicate planned revisions.

read point-by-point responses

Referee: The central claim that the GP posterior variance in the LOT tangent space yields a reliable acquisition policy for the Wasserstein path rests on the unexamined assumption that first-order linearization errors remain negligible. Because LOT is constructed around one fixed reference measure μ₀, any segment of the trajectory whose support or mass distribution deviates substantially from μ₀ projects its variation through an approximation whose error grows with distance; the fitted GP length-scale and variance can therefore encode linearization artifacts rather than true Wasserstein variability. This directly affects the correctness of the uncertainty-minimizing acquisition rule. The manuscript should either (a) provide a quantitative bound on the linearization error along the learned trajectory or (b) demonstrate empirically that the selected time points remain optimal under a non-linearized

Authors: We agree the linearization around a single reference is a central assumption whose validity requires further support, particularly for trajectories far from μ₀. Deriving a general quantitative error bound is technically demanding and outside the paper's scope. Instead, we will add empirical analysis in the revision: on synthetic data with known ground-truth paths, we will re-evaluate the timepoints selected by our policy using a non-linearized Wasserstein-based surrogate and report whether the selected times remain near-optimal. This addresses point (b) directly. revision: partial
Referee: The empirical validation compares the proposed policy only against uncertainty-agnostic baselines. To substantiate that the LOT-GP uncertainty is the operative ingredient, an ablation that replaces the LOT-GP variance with a heuristic or random acquisition function while keeping the same embedding should be reported; without it, the performance gain could be attributable to the embedding alone rather than to the active selection mechanism.

Authors: We accept that the current baselines do not fully isolate the role of the GP-derived uncertainty. In the revised manuscript we will add an ablation that retains the LOT embedding but replaces the variance-based acquisition with (i) uniform random selection and (ii) a simple heuristic (maximizing pairwise tangent-space distances), allowing direct comparison of performance gains attributable to the uncertainty-driven policy versus the embedding alone. revision: yes

Circularity Check

0 steps flagged

No circularity; derivation relies on external LOT embedding and GP without self-referential reduction

full rationale

The paper's core chain maps snapshots via LOT into a tangent space, fits a GP there, and derives an acquisition function from the resulting posterior variance. This structure is presented as a composition of standard tools (LOT linearization around a reference measure plus Euclidean GP regression) rather than any quantity being fitted on a subset and then renamed as a prediction, or any uniqueness theorem imported from the authors' own prior work. No equations are shown that equate the output acquisition policy to its inputs by construction, and the abstract explicitly frames LOT and GP as leveraged external components. The empirical outperformance claim is independent of the derivation. This is the normal case of a self-contained proposal.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Abstract-only review supplies no explicit free parameters, axioms, or invented entities; the central claim rests on the unelaborated premise that LOT yields a usable tangent space for GP uncertainty quantification.

pith-pipeline@v0.9.1-grok · 5684 in / 1156 out tokens · 29002 ms · 2026-06-29T08:34:04.597262+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

4 extracted references · 1 canonical work pages · 1 internal anchor

[1]

Bayesian Active Learning for Classification and Preference Learning

URL https://api.semanticscholar. org/CorpusID:123428953. Flood, S., King, M., Rodgers, R., Ruggles, S., Warren, J. R., Backman, D., Chen, A., Cooper, G., Richards, S., Schouweiler, M., et al. Ipums cps: Version 12.0 [dataset]. Minneapolis, MN: IPUMS, 10:D030, 2024. Gal, Y ., Islam, R., and Ghahramani, Z. Deep bayesian active learning with image data. InIn...

work page internal anchor Pith review Pith/arXiv arXiv doi:10.1371/journal 2024
[2]

The full study contains multiple culture conditions

(mouse fibroblast reprogramming under OSKM induction). The full study contains multiple culture conditions. In our experiments we restrict to theserumsubset. Cells are annotated with a numeric day post-induction, and two experimental batches are provided. Representation and preprocessing.We treat each time point as an empirical measure over cell embedding...

2005
[3]

This shows that our method is a general tool for adaptive experimentation on measure-valued dynamical systems, and is not restricted to biology

During this period, the distribution of weekly earnings shifted quickly, because there were many fewer low-earning workers in the data. This shows that our method is a general tool for adaptive experimentation on measure-valued dynamical systems, and is not restricted to biology. 17 Active Timepoint Selection for Learning Measure-Valued Trajectories 0 2 4...
[4]

We report the results in Table 7 and Table 8, where MMFM underperforms our reconstruction framework

to reconstruct a probability path given our acquired snapshots. We report the results in Table 7 and Table 8, where MMFM underperforms our reconstruction framework. We believe this gap is largely due to the strong non-stationarity of the underlying dynamics. Furthermore, unlike in MMSB, our approach can be used to sample multiple plausible probability pat...

2015

[1] [1]

Bayesian Active Learning for Classification and Preference Learning

URL https://api.semanticscholar. org/CorpusID:123428953. Flood, S., King, M., Rodgers, R., Ruggles, S., Warren, J. R., Backman, D., Chen, A., Cooper, G., Richards, S., Schouweiler, M., et al. Ipums cps: Version 12.0 [dataset]. Minneapolis, MN: IPUMS, 10:D030, 2024. Gal, Y ., Islam, R., and Ghahramani, Z. Deep bayesian active learning with image data. InIn...

work page internal anchor Pith review Pith/arXiv arXiv doi:10.1371/journal 2024

[2] [2]

The full study contains multiple culture conditions

(mouse fibroblast reprogramming under OSKM induction). The full study contains multiple culture conditions. In our experiments we restrict to theserumsubset. Cells are annotated with a numeric day post-induction, and two experimental batches are provided. Representation and preprocessing.We treat each time point as an empirical measure over cell embedding...

2005

[3] [3]

This shows that our method is a general tool for adaptive experimentation on measure-valued dynamical systems, and is not restricted to biology

During this period, the distribution of weekly earnings shifted quickly, because there were many fewer low-earning workers in the data. This shows that our method is a general tool for adaptive experimentation on measure-valued dynamical systems, and is not restricted to biology. 17 Active Timepoint Selection for Learning Measure-Valued Trajectories 0 2 4...

[4] [4]

We report the results in Table 7 and Table 8, where MMFM underperforms our reconstruction framework

to reconstruct a probability path given our acquired snapshots. We report the results in Table 7 and Table 8, where MMFM underperforms our reconstruction framework. We believe this gap is largely due to the strong non-stationarity of the underlying dynamics. Furthermore, unlike in MMSB, our approach can be used to sample multiple plausible probability pat...

2015