用于学习测度值轨迹的主动时间点选择

# 用于学习测度值轨迹的主动时间点选择
来源：https://arxiv.org/html/2605.30625
###### 摘要

从稀疏快照中推断连续概率路径是单细胞生物学等领域的一项基本挑战。在这些领域中，高保真数据采集通常是破坏性的，并受到高昂测序成本的制约。这激发了对主动学习策略的需求，以策略性地选择最优测量时间点。然而，为这一场景设计主动学习策略仍然是一个开放问题：目标对象位于无限维的Wasserstein空间中，标准欧几里得度量在此处失效，并且当前的插值方法缺乏认知不确定性量化。我们引入了一个框架，将主动实验扩展到测度空间。通过利用线性化最优传输（LOT），我们将分布快照映射到一个切空间，该空间适用于高斯过程建模，从而为潜在的概率路径构建一个易于处理的概率代理模型。这产生了一种采集策略，能够迭代选择测量时间点以最小化不确定性。实证结果表明，我们的策略在合成数据集和真实世界数据集上均优于不考虑不确定性的基线方法。

机器学习，ICML

## 1 引言

**测度值轨迹**。推断概率分布的时间演化（即在\[0,1\]区间上的概率路径\{\mu_t\}_{t\in[0,1]}）是科学领域中的一项基本挑战，范围从流体动力学 (Benamou and Brenier,2000) 到宏观经济学 (Achdou et al.,2022)。这个问题在单细胞生物学 (Wagner et al.,2016) 中尤为突出，其中细胞分化被建模为高维基因表达空间 (Trapnell et al.,2014) 中的一个动态过程。在此类情况下，完整的连续轨迹通常无法直接观测到。相反，我们只能有限地获取破坏性快照，即在离散时间点上的经验测度。因此，核心任务是**分布插值**：根据在多个时间点观测到的一组有限边际分布，恢复潜在的连续轨迹 t↦μ_t。

**主动时间点选择**。然而，在实践中，数据采集受到成本的严重限制。例如，在单细胞转录组学中，生成高保真度快照需要**破坏性**采样，并产生高昂的费用，即**通常每个时间点花费数千美元** (Ziegenhain et al.,2017)，从而排除了密集时间采样。在这种预算限制下，观测时间点变得至关重要。这激发了针对测度值过程的**主动学习**框架，旨在根据过去的观测结果迭代选择下一个测量时间点 t^*∈[0,1]，以最好地帮助估计潜在的概率路径¹。

¹ 本工作关注的是**昂贵**快照下的主动获取场景，而非实时选择场景。

**挑战**。在此背景下，主动学习面临着独特的挑战。首先，输出空间的几何结构本质上是非欧几里得的。标准的主动学习方法，例如基于高斯过程 (GP) 的方法 (Schulz et al.,2018; Williams and Rasmussen,1995)，假设输出是向量值并配备欧几里得度量。相比之下，概率测度存在于一个非线性空间中，该空间更适合用Wasserstein几何来描述 (Ambrosio et al.,2005)。其次，这使得对测度进行不确定性量化变得特别困难。用于回归和分类的主动学习通常依赖于需要认知不确定性概念的采集函数，但这种不确定性在当前分布插值方法中并不容易获得 (Lipman et al.; Rohbeck et al.,2025)。最后，测度值动态通常具有很强的非平稳性。例如，细胞发育的速度可能发生剧烈变化：长时间的稳态可能被快速、短暂的支化事件所中断 (Haghverdi et al.,2016)。因此，均匀间隔的采集时间点可能非常次优。

**方法**。我们提出了一种针对测度值轨迹的主动时间点选择策略，以应对上述挑战。我们的核心思想是通过将每个观测到的快照 μ_t 提升到一个切空间来**线性化**Wasserstein空间。具体来说，我们使用线性化最优传输（LOT）(Wang et al.,2013,2025) 将 μ_t 映射到固定参考测度处的切向量。然后，我们将这些切向量压缩为低维表示，并在其时间系数上放置一个**扭曲的**高斯过程（GP）先验。GP后验诱导出一个实用的认知不确定性代理，而扭曲机制通过允许对时间进行重新参数化，从而处理非平稳动态。最后，我们利用GP量化的认知不确定性来确定下一个测量时间点 t^*。

**贡献**。从**概念**上讲，我们提出了测度值轨迹的主动学习问题，将主动实验扩展到测度空间。从**技术**上讲，我们通过将线性化最优传输与多输出高斯过程相结合，在Wasserstein空间中构建了一个易于处理的概率代理模型。从**实证**上讲，我们证明我们的采集策略在合成数据集和真实世界数据集上均优于均匀采样和随机采样等基线方法。

## 2 相关工作

**主动学习**。主动学习在欧几里得空间的标量或分类目标上已相当成熟。早期且广泛使用的启发式方法包括**不确定性采样**，它查询具有最大预测模糊性的点 (Lewis,1995)，以及**委员会查询** (Seung et al.,1992)。一种互补的贝叶斯视角将采集视为最优实验设计，选择能最大化预期信息增益的输入 (Houlsby et al.,2011)。其他工作使用近似贝叶斯推断（例如，MC dropout (Gal et al.,2017)）将基于不确定性和信息的采集扩展到深度模型。更相关的是 (Singh et al.,2005)，它主动选择时间点以更好地拟合欧几里得值的基因表达曲线。然而，这些工作均未涉及分布空间中的主动学习，而这正是我们的重点。

**分布回归**。一个相关（但方向相反）的研究方向是**分布回归**，其中输入是概率测度，输出位于 R^d（或希尔伯特空间）中。经典方法通过核均值嵌入将输入分布嵌入到再生核希尔伯特空间（RKHS）中，然后执行（核）回归 (Póczos et al.,2013; Szabó et al.,2016; Muandet et al.,2017; Law et al.,2018)。相比之下，我们的设置将**时间**视为协变量，将**输出**视为分布。这更直接地关联到**分布插值**和**轨迹推断**方法，这些方法学习与观测到的边际分布一致的概率流或随机过程，包括流匹配和分数匹配公式、多边际扩展以及基于薛定谔桥的方法 (Lipman et al.; Tong et al.,2024; Lee et al.,2025)。然而，这些方法主要是**重建**方法：给定固定的快照，它们通常返回一个学习到的概率流或拟合的随机动态，因此仅产生一条诱导的边际路径。相比之下，主动获取需要关于可能概率路径的**认知不确定性**，以便决定下一步测量哪个时间点。我们的框架对这一不确定性进行建模，使其适用于主动学习。

**线性化最优传输（LOT）**。LOT将测度嵌入到参考分布的切空间中 (Wang et al.,2013; Kolouri et al.,2016)。该技术通过在切平面中使用线性分类器和回归，被证明对测度上的模式识别任务（如分类和重心估计）非常有效 (Moosmüller and Cloninger,2023)。然而，这些工作均未关注测度值轨迹的主动学习问题，而这需要我们引入不确定性概念。

## 3 问题设定

### 3.1 分布空间中的回归

**符号**。令 X ⊆ R^d 为特征空间。我们考虑一个概率路径，即概率测度空间中的一个时变函数 μ: [0,1] → P_2(X)。为方便符号表示，下文我们用 μ_t 表示时间 t 处的测度，而非 μ(t)，并注意到最大时间可在缩放后设为任意 t_max。P_2(X) 表示 X 上具有有限二阶矩的概率测度空间，定义为：

P_2(X) := { ρ ∈ P(X) : ∫_X ‖x‖^2 dρ(x) < ∞ }   (1)

该空间上的一个自然度量是 2-Wasserstein 度量 (Villani,2021)。对于任意两个测度 μ, ν ∈ P_2(X)，该度量定义为：

W_2(μ, ν) := ( inf_{π ∈ Π(μ, ν)} ∫_{X×X} ‖x - y‖^2 dπ(x, y) )^{1/2},   (2)

其中 ‖·‖ 表示欧几里得范数，Π(μ, ν) 表示 X×X 上所有边际分布为 μ 和 ν 的联合概率测度的集合。

**目标**。我们假设拥有一个时间快照数据集 D = { (t_i, μ̂_{t_i}) }_{i=1}^N，其中每个 t_i ∈ [0,1] 是一个测量时间点，μ̂_{t_i} 是在该时间观测到的经验测度（由从边际分布 μ_{t_i} 中抽取的样本构成）。我们的目标是利用 D 中的快照来估计潜在的连续概率路径 {μ_t}_{t∈[0,1]}。

该目标适用于多个领域，尤其是计算生物学。在此背景下，特征空间 X 表示基因表达空间（或由其导出的潜在空间）。每个经验测度 μ̂_{t_i} 对应于在时间 t_i 观测到的 n_i 个不同细胞的表达谱，目标是估计细胞在表达空间中的动态。

### 3.2 主动学习问题

在这项工作中，我们关注以下问题：给定固定的测量预算 B，我们如何选择测量时间点 {t_i}_{i=1}^B，以便最好地估计 {μ_t}_{t∈[0,1]}？

这意味着我们寻求一个采集策略 π，该策略根据当前历史数据 D，选择下一个测量时间点 t^* ∈ [0,1]。我们假设测量时间点不受特定顺序约束，且不必是单调递增的。

**单细胞主动测序**。例如，可以以密集的时间间隔收集并冷冻保存生物样本，形成一个潜在数据的“库”。由于对这些样本进行测序是主要的成本瓶颈，处理整个库通常是不可行的，因为每个时间点的成本通常高达数千美元。相反，主动学习策略必须顺序查询这个库，选择最优的时间点进行解冻和测序，以便在固定预算下最小化轨迹的不确定性。

### 3.3 关键挑战

要理解这个主动学习问题的难点，首先回顾一下主动学习在标准回归中如何运作是有帮助的。传统方法依赖于两个要素：一个能够根据输入预测输出的**回归器**，以及**认知不确定性**的量化。在欧几里得空间中，不确定性可以从概率模型（例如高斯过程）的后验方差或集成预测器的经验方差中推导出来。然后，主动学习策略利用这种不确定性来选择下一个查询点 t^*。然而，这种方法在我们的设定中失效，原因在于两个基本挑战。

**➀ 非欧几里得几何**。标准回归模型（例如高斯过程）依赖欧几里得运算在观测值之间进行插值。然而，P_2(X) 不是一个向量空间，即测度的线性组合通常不会产生有效的测度。我们在图1中说明了朴素欧几里得插值的局限性。具体来说，我们考虑一个具有时变均值的一维高斯轨迹。每个分布由其在一固定网格上评估的密度表示，并直接对这些密度向量拟合一个标准高斯过程以进行插值。如图所示，这种欧几里得插值会分裂质量，而不是连贯地传输质量。

参见图注

图 1：将高斯过程回归简单应用于密度会导致较差的插值效果。

虽然存在与Wasserstein几何兼容的插值方案，但它们通常只操作于**两个**参考测度之间。一个例子是位移插值 (McCann,1997)。给定 P_2(X) 中的 μ_0, μ_1，其定义为：

μ_t = ((1-t)Id + tT)_# μ_0,   (3)

其中 T 是从 μ_0 到 μ_1 的最优传输映射（见方程4），# 表示推前操作。将该思想扩展到 N>2 个快照的设置并非易事。虽然近期的研究为此设定提出了插值方案 (Rohbeck et al.,2025; Lee et al.,2025)，但它们本质上仍然是确定性的，并未提供主动学习所需的不确定性估计，我们接下来将讨论这一点。

**➁ 缺乏**