Active Learning of Conditional Generative Models via the Transport Neural Tangent Kernel¶

会议: NeurIPS 2026
arXiv: 2610.35652
OpenReview: https://openreview.net/forum?id=WioQm89GJJ
代码: https://github.com/biologicalML/transport_ntk
领域: 优化/理论
关键词: 条件生成模型、主动学习、神经正切核 (NTK)、最优传输、Wasserstein 距离

一句话总结¶

针对高通量科学实验中多条件概率分布输出的数据获取难题，本文提出传输神经正切核（tNTK）刻画生成分布对参数微扰的敏感度，严格证明其为后验 Fréchet 方差的上界，并基于高斯过程贪心最小化后验方差实现分布感知的条件生成模型主动学习。

研究背景与动机¶

高通量生物与物理实验（例如单细胞基因扰动测序、大规模药物篮子筛选、特定细胞类型的染色质开放性分析等）的典型输出往往不是单个标量或固定向量，而是高维连续空间上的条件概率分布。在实际科研场景中，制备或测量新条件（如测试全新抗原、未知小分子或未探究细胞系）的湿实验开销极其高昂，因而必须借助主动学习（Active Learning, AL）与自动化实验设计算法，在有限的采样预算下从海量候选条件池中智能筛选出最具信息量的实验输入，以使代理预测模型在新输入上的泛化性能最优化。

然而，现存的主动学习范式（包括分类不确定性、回归方差最小化、基于经典神经正切核 NTK 的方法等）从根本上假设输出空间为标量或低维向量。当面对分布型输出时，传统策略往往只能被动地将复杂概率分布压缩为人工指定的经验标量泛函（如条件均值或特定阈值事件发生率），这完全忽视了分布的多模态结构、方差异方差性及高阶形状特征。本文在理论上证明，2-Wasserstein 距离下的条件生成泛化风险可严格正交分解为均值误差项与分布形状误差项；仅基于条件均值的主动学习在形状变化剧烈的区域会产生系统性盲区，进而导致实验样本选择的严重次优。

为了打破这一理论与算法瓶颈，亟需一种能够直接在概率测度空间上运作的分布感知主动学习理论。核心 idea：提出传输神经正切核（tNTK），通过在参考测度上积分传输映射关于模型参数的雅可比内积，构建生成模型预测分布对参数微扰的几何敏感度度量；证明其不仅为参数后验下的 Fréchet 方差提供紧致上界，且能作为高斯过程协方差函数直接驱动全池后验方差最小化的主动数据获取。

方法详解¶

整体框架¶

本文方法的输入为无标注候选条件池 \(\mathcal{X}_{\text{pool}}\) 及当前已采集的带有分布样本的实验条件集 \(\mathcal{X}_{\text{train}}\)，输出为用于指导下一轮湿实验测量的最优批量条件子集 \(\mathcal{B}\)，并最终拟合出泛化误差最低的条件生成模型（CGM）。系统涵盖四类典型传输生成架构：仿射高斯传输映射、非交叉单调分位数网络、基于连续常微分方程（ODE）的条件流匹配（CFM）以及非凸条件 Monge 间隙网络。

方法的整体工作流分为四个级联阶段：首先，建立条件生成模型预测分布在 2-Wasserstein 空间上的泛化风险解耦公式，量化均值项与形状项的内在冲突；其次，在参数后验的高斯/拉普拉斯近似下，推导传输映射参数敏感度矩阵 \(G_\theta(x)\)，证明其迹上界控制后验 Fréchet 方差，进而定义双输入间的传输神经正切核（tNTK）；再次，针对仿射变换导出封闭解析解，针对深层网络与流模型构建无需昂贵自动微分的高效倒数第二层特征积分与 Monte Carlo 估计器；最后，将 tNTK 作为等效高斯过程的协方差核，采用顺序贪心算法最小化全池候选条件的预测后验方差，完成主动批次闭环迭代。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["候选实验条件池与生成模型初始化<br/>X_pool 与参数 θ_t"] --> B["Wasserstein 风险解耦与几何敏感度分析<br/>均值误差 R_mean + 形状误差 R_shape"]
    B --> C["传输神经正切核构建与高效估计<br/>定义 tNTK 并通过倒数第二层特征/ODE求解"]
    C --> D["分布感知的主动实验设计与贪心采样<br/>基于 GP 后验方差最小化逐点扩展批次 B"]
    D --> E["高通量实验反馈与条件生成模型更新<br/>扩充 X_train 并重训参数 θ_{t+1}"]

关键设计¶

1. 风险均值-形状正交解耦：揭示均值主动学习的盲区

传统回归主动学习聚焦于最小化均方误差，但在生成建模中，真实目标是缩小预测测度与真实测度之间的 2-Wasserstein 距离。本文针对这一痛点，引入中心化平移算子 \(\tau_v(y)=y-v\)，利用耦合优化技巧证明了如下严格的正交分解定理：

\[\mathcal{R}(x,\mathcal{D})=\mathbb{E}_{\theta\mid\mathcal{D}}[W_2^2(\mu_\theta(x),\mu(x))]=\underbrace{\mathbb{E}_{\theta\mid\mathcal{D}}[\|m_\theta(x)-m(x)\|^2]}_{\mathcal{R}_{\text{mean}}(x,\mathcal{D})}+\underbrace{\mathbb{E}_{\theta\mid\mathcal{D}}[W_2^2(\mu^c_\theta(x),\mu^c(x))]}_{\mathcal{R}_{\text{shape}}(x,\mathcal{D})}\]

其中 \(m(x)\) 为条件均值，\(\mu^c\) 为零均值中心化测度。在一维单调分位数函数表示下，该形状项可严格化为中心化分位数方差在各分位数水平 \(\tau \in [0, 1]\) 上的积分。该结论从理论上证实，当数据生成过程的方差或多模态形状随条件 \(x\) 发生显著变化而均值变化平缓时，传统均值主动学习会错误地将样本均匀分配，导致形状误差长期停滞，因此主动学习策略必须具备对形状变化的感知能力。

2. 传输神经正切核（tNTK）与后验 Fréchet 方差上界

在未标注测试点处，真实分布未知且生成模型后验测度之间的 Wasserstein 距离不可解析计算。为了定量捕捉模型预测的不确定性，本文以测度空间上的后验 Fréchet 方差 \(V_F(x) := \inf_{\rho\in\mathcal{P}_2(\mathcal{Y})} \mathbb{E}_{\theta\mid\mathcal{D}}[W_2^2(\mu_\theta(x),\rho)]\) 为测度不确定性代理。在参数后验满足均值为 \(\hat{\theta}\)、协方差为 \(\Sigma\) 的拉普拉斯近似下，本文通过同步耦合与二阶泰勒展开，证明了后验 Fréchet 方差满足可计算线性上界：

\[V_F(x) \leq \operatorname{Tr}(G_{\hat{\theta}}(x)\Sigma) + C(x)\|\Sigma\|_{\text{op}}^{3/2}, \quad \text{其中 } G_{\hat{\theta}}(x)=\int D_\theta T_{\hat{\theta}}(x)(y)^\top D_\theta T_{\hat{\theta}}(x)(y)\,d\nu(y)\]

在此基础上，本文将单点敏感度矩阵拓展至任意条件点对 \((x, x')\)，正式定义传输神经正切核（tNTK）：

\[k_\theta(x,x') = \int \operatorname{Tr}\left[D_\theta T_\theta(x)(y)^\top D_\theta T_\theta(x')(y)\right] d\nu(y)\]

当各向同性先验 \(\Sigma=\sigma^2 I\) 时，\(V_F(x) \leq \sigma^2 k_\theta(x,x) + O(\sigma^3)\)。当基础分布退化为狄拉克测度时，tNTK 严格退化为确定性神经网络的标准有限宽度 NTK；而在下三角仿射高斯传输下，\(h^\top G_\theta(x) h\) 恰好精准恢复 Knothe-Rosenblatt 重排度量下的测度导数平方，为主动采样提供了坚实的几何不确定性度量。

3. 轻量化闭式分解与流匹配 ODE 伴随积分

针对深度网络全参数雅可比矩阵在高维空间中计算不可行的工程痛点，本文推导了多种实用计算范式。若传输网络具有线性顶层 \(T_\theta(x)(y)=W^\top h_\theta(x,y)\)，本文证明倒数第二层特征 tNTK 仅需对隐藏表示内积在基础测度上求期望，完全规避了反向传播：

\[k_\theta^{\text{ll}}(x,x') = d \int h_\theta(x,y)^\top h_\theta(x',y)\,d\nu(y)\]

对于连续条件流匹配模型，其传输映射为常微分方程从 \(t=0\) 至 \(t=1\) 的流积分解 \(T_\theta(x)(y)=\phi^{1,x}(y)\)。本文通过对速度场雅可比矩阵沿 ODE 轨迹的时间积分构建核估计器：

\[k_\theta^{\text{Flow}}(x,x') = \int \operatorname{Tr}\left[ \left(\int_0^1 D_\theta v_\theta^{t,x}(\phi^{t,x}(y))\,dt\right)^\top \left(\int_0^1 D_\theta v_\theta^{t,x'}(\phi^{t,x'}(y))\,dt\right) \right] d\nu(y)\]

借助 Tsitouras 5/4 自适应 Runge-Kutta 求解器与 Monte Carlo 采样，算法即使在复杂的动力系统生成模型下仍能实现高保真的不确定性估计。

4. 基于高斯过程后验方差最小化的主动批次选择策略

在获取经验 tNTK 后，本文将其视作条件分布隐变量的高斯过程协方差函数。在已有标注子集 \(\mathcal{X}_S\) 条件下，任意测试点 \(x'\) 的后验方差为 \(\sigma_\theta^2(x'\mid \mathcal{X}_S) = k_\theta(x',x') - k_\theta(\mathcal{X}_S,x')^\top (k_\theta(\mathcal{X}_S,\mathcal{X}_S) + \lambda I)^{-1} k_\theta(\mathcal{X}_S,x')\)。为了实现全池综合不确定性的最快下降，算法采用贪心序列选择准则：

\[x^* = \arg\min_{x\in\mathcal{X}_{\text{pool}}\setminus\mathcal{B}} \sum_{x'\in\mathcal{X}} \sigma_\theta^2(x'\mid \mathcal{X}_{\text{train}} \cup \mathcal{B} \cup \{x\})\]

该策略无需启发式阈值调节，通过封闭矩阵低秩更新高效实施，自动兼顾了候选条件自身的测度敏感度与已有样本间的特征正交性多样性。

实验关键数据¶

主实验¶

论文在四类条件生成任务上进行了广泛评测，涵盖一维与多维高斯混合模拟、真实地面气象站年度温度与露点分布、CLINC 150 意图文本语义潜空间生成、以及两套前沿单细胞扰动测序基准（小分子激酶抑制剂测序 K562-Chemical 与巨噬细胞遗传敲除光学-转录组筛选 BMDM-Genetic）。基线对比了随机采样（Random）、基于条件均值拟合网络的标准 NTK（对应回归主动学习 SOTA 方法 BAIT）、联合均值方差 NTK、分位数 NTK 以及贝叶斯分布分歧方法 BALSA-EMD。

下表展示了非交叉分位数预测任务在第 5 轮迭代时的测试集 Wasserstein 风险改善量（\(\Delta \mathcal{R} \times 10^{-3}\)，数值越低代表改善越显著，即相对初始模型的误差降幅越大）：

数据集	预算配置 (\(B_0+B\))	Random	NTK (Mean)	NTK (Mean-Var)	BALSA-EMD	本文 (tNTK)	相对最优基线提升
高斯混合模拟	5 + 1	\(-296.12 \pm 16.29\)	\(-416.43 \pm 16.15\)	\(-367.99 \pm 17.59\)	\(-332.16 \pm 16.06\)	\(\mathbf{-429.52 \pm 16.39}\)	降低 \(13.09 \times 10^{-3}\)
高斯混合模拟	10 + 1	\(-121.02 \pm 10.93\)	\(-168.28 \pm 11.23\)	\(-162.85 \pm 11.04\)	\(-133.59 \pm 12.12\)	\(\mathbf{-205.01 \pm 12.38}\)	降低 \(36.73 \times 10^{-3}\)
地表温度预测	5 + 5	\(-219.96 \pm 14.43\)	\(-255.65 \pm 15.29\)	\(-258.48 \pm 15.55\)	\(-199.14 \pm 15.87\)	\(\mathbf{-267.93 \pm 15.36}\)	降低 \(9.45 \times 10^{-3}\)
地表温度预测	10 + 1	\(-41.02 \pm 7.44\)	\(-84.57 \pm 9.70\)	\(-86.88 \pm 9.69\)	\(-71.14 \pm 9.57\)	\(\mathbf{-89.84 \pm 10.01}\)	降低 \(2.96 \times 10^{-3}\)
露点分布预测	5 + 5	\(-280.77 \pm 21.75\)	\(-364.19 \pm 21.53\)	\(-372.49 \pm 22.06\)	\(-178.00 \pm 17.91\)	\(\mathbf{-442.25 \pm 22.34}\)	降低 \(69.76 \times 10^{-3}\)
露点分布预测	10 + 1	\(-76.55 \pm 10.15\)	\(-158.94 \pm 10.97\)	\(-159.24 \pm 10.32\)	\(-124.13 \pm 9.39\)	\(\mathbf{-212.74 \pm 11.73}\)	降低 \(53.50 \times 10^{-3}\)

在条件流匹配（CFM）与条件 Monge 间隙单细胞实验中，测试集评价采用 Sinkhorn 散度改善量（\(\Delta \text{Sinkhorn} \times 10^{-3}\)）：在 CLINC 150 意图语言潜空间中，\(5+1\) 配置下 tNTK 取得 \(\mathbf{-672.85 \pm 107.66}\)，显著超越均值 NTK 的 \(-306.01 \pm 111.84\)；在 BMDM 遗传扰动预测中，\(30+5\) 配置下 tNTK 取得 \(\mathbf{-70.59 \pm 4.84}\)，超越均值 NTK 的 \(-54.90 \pm 4.17\) 与随机采样的 \(-48.73 \pm 2.41\)。全论文统计表明，在全部 30 个基准实验配置中，tNTK 在 26 个配置下全面击败均值 NTK。

消融与误差解耦实验¶

论文针对合成二维高斯分布的仿射传输任务，显式消融了解析计算的均值敏感度分量与协方差敏感度分量对测试集 Wasserstein 误差演进的影响：

采样策略	最终总 Wasserstein 误差	最终均值分量误差	最终协方差分量误差	表现定性说明
Random (均匀随机)	\(0.231 \pm 0.014\)	\(0.118 \pm 0.009\)	\(0.113 \pm 0.008\)	探索效率低，收敛缓慢
NTK (仅均值梯度)	\(0.228 \pm 0.012\)	\(\mathbf{0.082 \pm 0.006}\)	\(0.146 \pm 0.009\)	均值误差虽降，但协方差误差恶化（甚至劣于随机）
tNTK (全分布敏感度)	\(\mathbf{0.189 \pm 0.009}\)	\(0.091 \pm 0.006\)	\(\mathbf{0.098 \pm 0.005}\)	均值与协方差协同下降，总 Wasserstein 误差降低 1.2 倍
tNTK (高维可扩展性 \(d=50\))	\(\mathbf{0.642 \pm 0.021}\)	\(\mathbf{0.285 \pm 0.011}\)	\(\mathbf{0.357 \pm 0.014}\)	在 50 维高斯传输下依旧显著优于仅均值基线

关键发现¶

均值与形状的优化权衡悖论：单纯依据均值 NTK 实施主动采样时，模型虽然能以最快速度降低均值估计误差，但会导致高方差或复杂几何变形区域被长期忽视；在协方差变化显著的子流形上，均值 NTK 的形状误差甚至逊于盲目随机采样，直接拖垮全局分布相似度。
理论上界的渐进紧致性验证：在 \(d=1\) 且传输单调时，Fréchet 方差的 tNTK 上界在小方差微扰极限下是严格精确的；当维度提升至 \(d=50\) 时，固定点数值解构表明 tNTK 理论高估程度不足 70%（小于 2 倍），且在相对候选排序上保持单调保序性。
动力学常微分方程积分的可行性：即使在条件流匹配框架下，传输映射依赖于神经速度场的全时间自适应数值积分，tNTK 仍然可以通过伴随灵敏度或前向特征积分提供高度稳健的采样导向。

亮点与洞察¶

从标量主动学习向测度流形的主动范式跨越：将主动学习从传统标量响应面推广到无穷维测度空间，首次严格建立了 2-Wasserstein 几何框架下的 Fréchet 方差上界，统一了解析高斯变换与深层非参数生成模型。
极度轻量化的倒数第二层特征积分技术：证明了线性顶层下 tNTK 仅依赖倒数第二层特征在基准分布下的蒙特卡洛内积，使分布感知的主动学习无需昂贵的全网络二阶反向传播，在大规模生物数据集中极易落地。
生物湿实验高通量闭环设计的实用范本：单细胞基因筛选实验验证了该方法能够精准捕捉药物与敲除诱导的异质性转录组细胞状态转变，为自驱动实验室（Self-Driving Labs）提供核心算法引擎。

局限与展望¶

高维空间的几何松弛缺口：理论分析表明，在 \(d>1\) 时，由于 Wasserstein 测度流形曲率的存在，tNTK 上界并不严格紧致（仅在 Knothe-Rosenblatt 重排度量下严格精确），未来需进一步刻画黎曼曲率修正项。
大批次采样的多样性坍缩隐患：当前采用基于 GP 后验方差的贪心顺序采样，在超大批次（如 \(B > 50\)）实验设计中，可能会倾向于局部聚集，后续可结合行列式点过程（DPP）或几何多样性约束增强批次覆盖。
高维隐式生成模型的扩展：实验主要覆盖了潜变量维度 \(\leq 50\) 的场景，对于高分辨率像素级扩散模型（Pixel Diffusion）的端到端雅可比积分，计算内存开销依然严峻，有待探索自适应随机投影等降维估计技术。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统提出传输神经正切核理论并证明后验 Fréchet 方差上界，开拓了分布感知主动学习新方向。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 种生成架构、合成基准与真实世界单细胞激酶/遗传敲除测序，横跨 30 种实验配置。
写作质量: ⭐⭐⭐⭐⭐ 数学论证严谨，测度几何动机清晰，实验图表与消融分析详实自洽。
价值: ⭐⭐⭐⭐⭐ 能够显著削减高通量分子与单细胞生物学实验的耗材与试剂成本，具有巨大的科学与工业应用价值。