[ICML 2025/ArXiv] DCT：让分布输运具备“举一反三”的泛化力

TL;DR

在科学决策和生物医学领域，我们经常需要预测一个群体（分布）在特定干预后的状态。传统的输运模型往往只能处理训练集里见过的分布。本文提出的 Distribution-Conditioned Transport (DCT) 框架，通过引入高度抽象的分布编码器，实现了“任意分布到任意分布”的自动化输运，即便是在完全陌生的观测条件下，也能给出高精度的预测。

1. 痛点：被配对数据“困住”的输运模型

在处理单细胞 RNA 测序或免疫组库数据时，研究者常面临两个棘手问题：

1. OOD 泛化难题：模型能否预测一个从未见过的供体对某种新药物的反应？
2. 数据稀疏性（Orphan Marginals）：许多实验样本只有初始状态或只有结果状态，缺乏配对的动态。

传统的模型（如简单的 GAN 或初级的 Flow Matching）通常将分布视为离散的标签（K-to-K）。一旦遇到第 K+1 个标签，模型就抓瞎了。

2. 核心直觉：将分布映射到连续潜空间

DCT 的核心思想是：与其学习特定分布间的映射，不如学习“分布特征”到“输运向量场”的通用映射。

2.1 分布编码器 (The Encoder)

作者采用了具有 置换不变性 (Permutation Invariant) 和 比例不变性 (Proportional Invariant) 的深度集合 (Deep Sets) 架构。这意味着无论你输入多少个细胞、以什么顺序输入，编码器提取出的 $z$ 都指向同一个概率分布。

2.2 三种学习模式

• 监督模式：已知 (Source, Target) 对，学习固定的转移。
• 任何对任何 (Any-to-Any) 模式：在训练集中自由组合不同的分布对，迫使模型理解 latent space 的几何结构。
• 半监督模式：利用那些只有单点观测的 “孤儿” 分布碎片，补全整体的输运景观。

Figure 1: DCT 框架示意图。源分布与目标分布分别被编码为嵌入向量，输运模型 T 以此为条件生成路径。

3. 实验见证：从高斯混合到 T 细胞演化

3.1 合成实验：K-to-K 的终结

在二维高斯混合模型 (GMM) 实验中，作者对比了传统的“标签式”模型和 DCT。结果发现，当目标分布位于训练支撑集之外时（OOD），DCT 能够顺滑地进行插值和外推，而 K-to-K 模型则完全崩溃。

在 OOD 目标上，DCT (Any-to-any) 的 MMD 和 SWD 误差远低于基线。

3.2 生物学应用：跨越未知的鸿沟

• 药物扰动预测：在处理不同病人的类器官对化疗药物的反应时，DCT 展现了极强的跨病人（Cross-patient）泛化能力。
• TCR 序列预测：针对 COVID-19 患者的免疫受体序列，由于大多数病人只有单次采样，DCT 通过半监督学习捕捉到了跨病人的共同演化规律，预测精度提升了 2 倍以上。

4. 深度洞察：为什么有效？

DCT 的成功在于它从数学上解决了 中心极限定理 (CLT) 在分布空间的应用。作者在附录中证明了：只要编码器满足特定的一致性条件，在小批量 (Minibatch) 上训练得到的损失函数是渐进无偏的。这意味着我们可以用有限的样本，去逼近无限测度空间里的真实流场。

5. 局限与未来

尽管 DCT 在 OOD 任务上表现卓越，但在某些 IID（分布内）任务上由于模型过于通用，偶尔会出现欠拟合（Underfitting）。此外，如何在大规模蛋白质-细胞交互数据上进一步扩展 Encoder 的表达能力，仍是未来的热点。

总结：DCT 为我们提供了一种全新的视角——分布不再是死的标签，而是可以被数字化、被内插、被操纵的连续向量。这对于个性化医疗和精准干预预测具有深远的里程碑意义。