分布条件传输：跨分布泛化的生成传输框架 Distribution-Conditioned Transport
将传输模型条件化在源/目标分布的嵌入上，实现跨分布泛化与半监督传输
前置知识
流匹配（Flow Matching）
流匹配是一类连续时间生成建模方法，通过学习一个时变速度场 $v_\theta(x,t,z)$ 来描述从源分布到目标分布的传输轨迹。训练时采样时刻 $t \sim \text{Uniform}[0,1]$、源样本 $x_0$、目标样本 $x_1$，构造线性插值 $x_t=(1-t)x_0+tx_1+\sigma\epsilon$，并用 $\mathcal{L}_{FM}=\mathbb{E}\|v_\theta(x_t,t,z)-(x_1-x_0)\|^2$ 训练；推理时用 dopri5 等 ODE 求解器从 $x(0)=x_0$ 积分到 $x(1)$ 得到目标样本。它是 DCT 框架中使用的核心传输机制之一。
本文三种生成器家族（SWD、Energy、Flow Matching）中，FM 因其样本级别配对损失能学到最强点对点耦合，是理解 DCT 条件化机制的关键背景。
分布编码器（Distribution Encoder）
分布编码器 $E: \mathcal{S} \to \mathbb{R}^d$ 把一组样本 $S=\{x_j\}_{j=1}^m$ 映射成一个固定维度的向量 $z$，该向量只反映底层分布 $P$ 而非采样噪声。它必须满足置换不变性（重排样本不变）和比例不变性（均匀复制样本不变），从而保证 $E(S)$ 仅依赖于经验测度 $\hat P = \frac{1}{m}\sum_j\delta_{x_j}$。在 DCT 中采用 Fishman et al. (2025) 的均值池化 DeepSets/GNN 架构。
分布编码器是 DCT 框架的基石，本文所有传输模型都以源/目标分布嵌入作为条件，没有它就无法在分布空间做条件化。
中心极限定理（CLT）与 plug-in 损失
当编码器 $E$ 是置换+比例不变时，存在可测函数 $\phi$ 使 $E(S_m)=\phi(\hat P_m)$。在 Hadamard 可微条件下，$\sqrt{m}(\phi(\hat P_m)-\phi(P))\overset{d}{\Rightarrow}\mathcal{N}(0,\Sigma)$，即编码器有 CLT。基于此，对任何依赖编码器状态的下游目标 $\ell(x,z)$，plug-in 损失 $\ell(x,E(\hat S))$ 是均值相合的，偏差为 $O(m^{-1})$。
Prop. 2.1 是本文 DCT 训练的理论基石——它证明对 minibatch 求梯度不会偏离总体目标 $O(m^{-1/2})$，因而可以直接用小批量联合训练分布编码器与传输模型。
Wasserstein 与能量距离
Wasserstein 距离 $W_2(P,Q)=(\inf_{\pi\in\Pi(P,Q)}\int\|x-y\|^2d\pi)^{1/2}$ 度量两个分布间的传输成本。Sliced Wasserstein 距离 (SWD) 通过在 $L$ 个随机方向上投影后做 1D Wasserstein 求平均。能量距离 $D_{\text{energy}}(P,Q)=2\mathbb{E}\|X-Y\|-\mathbb{E}\|X-X'\|-\mathbb{E}\|Y-Y'\|$ 等价于用能量核 $-\|x-y\|$ 计算的 MMD，无需带宽选择。两者都是分布间的严格度量，值越小表示分布越接近。
SWD 与 Energy 都是 DCT 中使用的传输机制与评估指标；理解它们有助于看懂 Table 1、Table 12 中的具体数值。
scRNA-seq 批次效应与谱系示踪
scRNA-seq（单细胞 RNA 测序）实验在不同生物样本、测序批次、处理条件下产生的细胞基因表达谱会带有技术差异（批次效应），需要校正。谱系示踪（lineage tracing）通过给每个细胞打上克隆条形码追踪其分化命运，例如 Weinreb et al. (2020) 在 $t\in\{2,4,6\}$ 天测量同一克隆在多个时间点的细胞状态。
本文四个生物应用（batch effect transfer、organoid perturbation、clonal dynamics、TCR evolution）都建立在这些概念之上，是把抽象 DCT 框架落到真实科学的桥梁。
研究动机
在现代生物医学数据中，研究者越来越频繁地遇到"分层"数据集：每个 donor、batch、时间点或克隆对应一个分布 $P_i$，而这些分布本身又从一个元分布 $Q$ 中采样。这种多尺度结构在 scRNA-seq 中尤其常见——Yazar et al. (2022) 的 eQTL 数据集包含数百个 donor，Weinreb et al. (2020) 的谱系示踪数据约 $6\cdot10^3$ 个克隆中只有约 $2\cdot10^3$ 个在多时间点被观测到（即"孤儿"边际）。现有传输工具（GAN、normalizing flows、flow matching、stochastic interpolants）都只针对一对源-目标分布，不能泛化到训练时未见过的 $(P_i,P_{i'})$。多边际随机插值 MMSI（Albergo et al., 2023a）只能处理训练时固定的 K 个分布；Meta Flow Matching (Atanackovic et al., 2024) 只能 source-conditioning 而无法利用孤儿边际——任何只在单一时间点观测到的细胞群都会被浪费。此外，单细胞药物扰动预测里需要在新患者（OOD）上做迁移，但现有 scGen/CellOT 等方法没有 patient-level 表示，IID 与 OOD 性能差距巨大。
本文的目标是本文提出一个统一的框架——分布条件传输 (Distribution-Conditioned Transport, DCT)，将分布编码器 $E$ 与任意传输模型 $T$ 耦合，使 $T(x\mid z_{src},z_{tgt})$ 能：(1) 在监督 (one-to-one)、无监督 (any-to-any)、半监督三种设定下学习传输；(2) 泛化到训练时未见的源/目标分布；(3) 利用"孤儿边际"——即只在单一条件下观测到的分布。框架对底层传输机制（flow matching、SWD、Energy/MMD）保持 agnostic，目标是把这些机制提升为可泛化的条件生成器。
与已有工作不同的是，本文的独特切入角度是 plug-in 损失理论 + 半监督接口。理论方面，本文证明存在一个 CLT（Prop. 2.1/Thm. A.5）：只要分布编码器满足 Hadamard 可微与不变性，对任何下游损失 $\ell(x,z)$，plug-in 估计 $\ell(x,E(\hat S))$ 都均值相合且偏差为 $O(m^{-1})$，从而保证用 minibatch 训练不偏离总体目标。半监督方面，本文给出两阶段方法：先用无监督 any-to-any 训练分布编码器，再用轻量级（线性）回归 $z_{tgt}\approx Wz_{src}+b$ 在可用配对上学习潜在预测器，把无监督编码与监督任务解耦。这种组合既能让 OOD 泛化，又不需要修改原有传输机制的损失——这是对 MFM 等方法的实质性推广。同时本文还诊断了 product-coupling 下 transport 可能"忽略源样本"的失败模式（Prop. A.7），并给出对齐诊断（Tab. 5）。
核心方法
DCT 的核心思想是：把传输模型 $T$ 的条件从单一分布身份（如 K 个 one-hot 标签）替换为从样本集合学到的分布嵌入。具体地，先训练一个置换+比例不变的分布编码器 $E:\mathcal{S}\to\mathbb{R}^d$，把任意样本集 $S=\{x_j\}_{j=1}^m\sim P$ 编码为 $z=E(S)\approx\phi(P)$；然后用一个底层传输模型（FM/SWD/Energy）学条件生成 $T(x\mid z_{src})$（监督）或 $T(x\mid z_{src},z_{tgt})$（无监督/半监督）。直觉上：分布嵌入是分布的"指纹"，相邻分布的指纹相近，因而 $T$ 能在嵌入空间内插外推；plug-in CLT 理论保证对 minibatch 求梯度不会偏离总体目标。技术路线有三条：(a) 监督 source-conditioned（条件仅 $z_{src}$），对应一对一的配对传输；(b) 无监督 source-target-conditioned（条件 $z_{src},z_{tgt}$），对应任意两分布间的传输；(c) 半监督——先用 (b) 学到好的 $E$，再用轻量级回归把监督任务投影到潜在空间。
核心创新在于三点：1) 把已有传输机制提升为"分布条件"接口——任何一个能从样本集合算 $\hat y$ 的生成器都可以通过把 $(z_{src},z_{tgt})$ 拼到输入或时间条件里变成条件生成器，无需重新设计损失；2) CLT 化的 plug-in 训练理论——证明只要分布编码器是 GDE (Fishman et al. 2025) 那一类，minibatch 训练是无偏且高效的，每个梯度步成本与基础传输模型同阶（只多一个 $E$ 的前向传播）；3) 半监督解耦——把分布编码器的训练（无监督，依赖大量未配对边际）与下游监督任务（用线性回归 $z_{src}\to z_{tgt}$）分开，使 OOD 任务能借助未配对数据。与已有方法的本质区别：DCT 不是新传输模型，而是把 GDE 嵌入与任意传输机制"插接"起来的统一框架；与 MFM 的区别是 DCT 支持 source-target conditioning 和利用孤儿边际；与 MMSI 的区别是 DCT 用连续嵌入而非 K 维 simplex 顶点，因而能泛化到 $(K+1)$-th 分布。
方法步骤详情
完整方法包含以下步骤：(1) 分布编码器 $E$ 的构造——对每个样本 $x_j$ 用 MLP 得到初始特征 $h_j^{(0)}=\text{MLP}_{in}(x_j)$，经 $L$ 层均值池化块更新 $h_j^{(\ell)}=\text{MLP}_{pool}^{(\ell)}([h_j^{(\ell-1)};\bar h^{(\ell-1)}])$，最后做平均并线性投影 $z=\sigma(W\cdot\bar h^{(L)}+b)$ 到 $d_z$ 维，$\sigma=\text{SELU}$，可选择 $\ell_2$ 归一化；该编码器满足置换+比例不变性。 (2) 监督传输——给定配对 $(S^{src}_i,S^{tgt}_i)$，编码源侧 $z^{src}_i=E(S^{src}_i)$，训练 $T(x\mid z^{src})$ 使 $T(S^{src}_i\mid z^{src}_i)\overset{d}{\Rightarrow}P^{tgt}_i$；损失 $\mathcal{L}_{sc}=\ell(S^{tgt}_i,T(S^{src}_i\mid z^{src}_i))$ 可用 $\ell_{FM}$、SWD regression 或 Energy regression 中任一种。 (3) 无监督传输——任意 $(S_i,S_{i'})$ 配对（甚至来自独立边际），编码双侧 $(z_i,z_{i'})$，训练 $T(x\mid z_i,z_{i'})$；损失 $\mathcal{L}_{stc}=\ell(S_{i'},T(S_i\mid z_i,z_{i'}))$；对每个梯度步独立采样 $(u,v)\sim Q^{joint}$，$\hat S_u\subset S_u,\hat S_v\subset S_v$，只增加 $E$ 一次前向。 (4) 半监督——先用 (3) 在所有可用边际上训 $E$ 与 $T$；再用可用配对 $(S^{src},S^{tgt})$ 拟合线性预测器 $\hat z^{tgt}=Wz^{src}+b$（带交叉验证选正则）；测试时把 $\hat z^{tgt}$ 喂给 $T$。 (5) 训练协议——Adam，lr$=2\times 10^{-4}$，batch=256 分布对，MVN 实验训 200 epoch，Trellis 训 20000 epoch；FM 用 dopri5 求解器，绝对/相对容差 $10^{-4}$，噪声 $\sigma=0.5$。
技术新颖性
技术新颖性体现在三方面：(a) 理论方面，Prop. 2.1/Thm. A.5 给出 plug-in 损失 CLT 的 DCT 特化版本，明确处理 paired 与 unpaired minibatch 两种情形，并证明在 minibatch 内复用有限个点不影响渐近分布；Cor. A.6 证明 plug-in 目标均值相合且偏差 $O(m^{-1})$。(b) 接口方面，把 GDE 嵌入与三类生成器（SWD、Energy、FM）"插接"成一个统一框架——此前 Meta Flow Matching 只针对 FM，He et al. (2025) 只针对 MMD；本文第一次把它们放进同一框架做横向对比。(c) 应用方面，展示了 DCT 能把"任意源-任意目标"传输与"半监督利用孤儿边际"这两个能力同时用于 4 个真实生物问题，且在 PDO 与 lineage tracing 两个任务上 STC 显著优于 SC（patient holdout OOD MMD：PDO 上 Energy 从 0.0298 降到 0.0191；Weinreb 上从 14.10 降到 9.67）。此外还贡献了对 product-coupling 下"忽略源样本"失败模式的诊断（Prop. A.7/Tab. 5）：FM 因为样本级配对损失而保持强对齐（$d_{pair}/d_{rand}=0.017,\rho=1.0$），而 deterministic SWD/Energy（$d_{pair}/d_{rand}\approx 0.89$）与 stochastic energy sampler（$d_{pair}/d_{rand}\approx 1.0$）的对齐较弱。
实验结果
DCT 在 6 个数据集上一致优于 K-to-K 或 source-conditioned 基线。(1) MVN/GMM 合成实验（Tab. 7-10）：固定源 $P_i$，对 OOD 目标 $P_{i'}$ 测 W₂/Energy 距离。当 $K=10$ 时 K-to-K 的 OOD Energy 距离高达 0.91（MVN），而 any-to-any 仅 0.15；$K=10^4$ 时 any-to-any 在 IID 与 OOD 都接近 0。(2) scRNA-seq 批次效应迁移（Tab. 1/Tab. 12）：用 56 只小鼠胰腺数据（Hrovatin et al. 2023），留出"老年"组作分布漂移。DCT 在 3 个 held-out donor 间 MMD：Energy 生成器从 K-to-K 的 0.80 降至 0.13，FM 从 0.49 降至 0.15；Harmony MMD 0.18、scVI MMD 1.81。(3) 半监督 PDO 药物扰动（Tab. 2）：Trellis 数据集 10 个 colorectal 患者、11 种药物、927 对。监督 SC 在 patient holdout（OOD）MMD=2.27-2.98（×10⁻²），半监督 STC 降至 1.91-2.08；Oracle（用真实 $z_{tgt}$）进一步降至 0.22-0.53；IID（replicate holdout）上 Oracle 反超 SC：MMD 0.10 vs 2.19（scGen）、1.98（CellOT）。(4) 谱系示踪（Tab. 3）：Weinreb 数据集约 $6\cdot10^3$ 克隆但仅约 $2\cdot10^3$ 在多时间点观测到。STC 在 628 held-out 克隆上的 MMD：SWD 8.67（vs SC 9.87）、Energy 8.54（vs 9.42）、FM 9.67（vs 14.10）；Oracle 进一步降至 2.88-5.08。(5) TCR 库预测（Tab. 4）：37 COVID-19 患者中仅 10 个多时间点，DFM 半监督把 MMD 从 0.058 降到 0.022（降幅 > 50%），而 ProGen2 因 encoder collapse（cosine similarity > 0.998）几乎无差别。(6) 半监督泛化曲线（Fig. 5）：在 $\|\mu\|_\infty\leq 2.5$ 内监督与半监督相当，超出后监督 SWD/Energy 误差从 0 跳到 0.017-0.037，半监督维持在 0.002-0.003，逼近 Oracle。整体规律：K-to-K 在 IID 小 K 下有竞争力，但 OOD 完全失败；any-to-any 在 IID/OOD 都稳定，且 K 越大优势越显著。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| scRNA-seq batch effect transfer (3 held-out donors) | MMD distance (↓) | 0.0733±0.0119 (FM, any-to-any) | 0.2429±0.0385 (FM, K-to-K); 0.9074 (scVI); 0.0903 (Harmony) | DCT FM 较 K-to-K FM 降低约 70%；接近 Harmony 水平且显著优于 scVI |
| scRNA-seq batch effect transfer (Energy generator) | MMD distance (↓) | 0.0639±0.0046 (any-to-any) | 0.4000±0.0453 (K-to-K) | 降幅约 84% |
| Drug perturbation (PDO) — held-out patients OOD | MMD × 10² (↓) | 1.97±0.22 (SWD STC) | 2.59±0.08 (SWD SC) | 约 24% 改进；Energy STC 1.91±0.49 vs SC 2.98±0.19；FM STC 2.08±0.18 vs SC 2.27±0.06 |
| Drug perturbation (PDO) — held-out replicates IID | MMD × 10² (↓) | 0.94±0.18 (SWD STC) | 0.64±0.22 (SWD SC) | IID 下 STC 略差于 SC（Oracle 仅 0.10，说明编码器仍有较大提升空间） |
| Clonal dynamics forecasting (628 held-out clones) | MMD (↓) | 8.67±0.22 (SWD STC); 8.54±0.21 (Energy STC); 9.67±0.25 (FM STC) | 9.87±0.27 (SWD SC); 9.42±0.26 (Energy SC); 14.10±0.28 (FM SC); 2.19 (scGen); 1.98 (CellOT) | DCT 各生成器均优于 SC；FM STC 较 SC 改进 31%；Oracle 显示若预测器更好可达 2.88-5.08 |
| TCR repertoire forecasting (held-out patients) | MMD (↓) | 0.0215±0.0043 (DFM STC) | 0.0579±0.010 (DFM SC); 0.0588±0.009 (ProGen2 SC) | DFM STC 较 DFM SC 改进 63%；ProGen2 因 encoder collapse 几乎无效 |
| MVN transport, K=10000, OOD targets | Energy distance (↓) | -0.000 (Energy any-to-any); 0.011 (FM any-to-any) | 0.019 (Energy K-to-K); 0.025 (FM K-to-K) | DCT 在 OOD 目标上保持近零误差，K-to-K 远逊 |
局限与改进
作者明确指出的局限：(1) STC 在 IID 上有时弱于 SC，例如 PDO replicate holdout 上 MMD 从 SC 的 0.64 (SWD) 变为 STC 的 0.94，作者认为可能因 equal compute budget 下 STC 容量更大但需更多训练（scaling 不同）。(2) Product-coupling 下随机生成器可能退化为"忽略源"的退化模式（Prop. A.7/Tab. 5），FM 因样本级配对损失能学到强对齐（$d_{pair}/d_{rand}=0.017$），但 deterministic SWD/Energy 对齐较弱（$\rho=0.1$），stochastic energy sampler 完全退化（$d_{pair}/d_{rand}=1.0$）。(3) TCR 实验中 ProGen2 出现 encoder collapse（cosine similarity 0.998），原因是 prefix-conditioning 信息通道太窄。(4) Weinreb 实验的 SWD/Energy/FM Oracle MMD 仍在 2.88-5.08，与潜在空间预测器的简单 ridge 回归（8.54-9.67）差距大，说明 latent predictor 仍是主要瓶颈；尝试 co-train 与 stratified sampling（Tab. 16）以及用 MLP/RF 替换 ridge（Tab. 17）都只带来小幅提升。 我的观察：(a) 框架对 K-to-K 基线在低 K 时其实不占优（K=10 时 Tab. 7 MVN IID 下 K-to-K Energy 仅 0.001 vs any-to-any 0.077），因 any-to-any 需要把 K 个分布"塞"进连续嵌入会欠拟合。(b) IID/OOD tradeoff 在 PDO 实验最明显：SC 在 replicate holdout 上 MMD 仅 0.64-1.01 但在 patient holdout 上飙到 2.27-2.98，暗示 SC 把 patient-specific 模式过拟合到了监督配对上。(c) 半监督预测器的样本量仅 $10$-$10^3$（生物数据集），小样本下线性模型的方差是另一个隐患。(d) FM 在合成 Gaussian 上有反常的 OOD 泛化能力（Tab. 9 Supervised FM OOD Energy=-0.000），但真实数据不复现，作者承认这可能是低维线性变换的 artifact。
独立分析的弱点
独立分析三个可改进的弱点：(1) 半监督预测器瓶颈——STC 的核心假设是用一个轻量级映射 $z^{src}\to z^{tgt}$ 把无监督编码转成监督任务，但 Oracle 与 STC 的差距（PDO 上 Energy Oracle 0.22 vs STC 1.91，约 9 倍）说明这条捷径还不够。改进方向：可探索非线性预测器（如 Tab. 17 的 MLP）或直接在 STC 训练中混入监督对（cotrain + stratified sampling，Tab. 16 在 PDO 上仅把 OOD MMD 从 0.0208 微降到 0.0202，需要更激进的方法如 meta-learning 或在 latent 空间做 prototypical network）。(2) x-ignoring 失败模式——deterministic SWD/Energy 在 product coupling 下学到 $T(x\mid z^{src},z^{tgt})\approx G(z^{tgt})$ 而忽略源样本（Tab. 5 中 $d_{pair}/d_{rand}\approx 0.89$），这在生物学上是危险的：药物扰动预测需要给出 patient-specific 反例，而不是任意匹配目标分布的细胞。改进方向：可加样本级对齐正则（如 OT pairing loss、Tong et al. 2023b 的 mini-batch OT）或在 inference 时去噪约束 $T$；具体包括在 latent 空间做 contrastive pairing、或用 He et al. (2025) 那种显式 MMD 加 paired objective。(3) 框架对 K 小时的 IID 性能——当 $K\leq 10$ 时 any-to-any 比 K-to-K 差（MVN/GMM IID 下 Tab. 7/8），因为连续嵌入空间把信息"压扁"了。改进方向：混合架构——小 K 时退化为 one-hot embedding 或加 hierarchical bias；或对编码器加 kernel-based inductive bias 让它在小样本下不欠拟合。
未来方向
作者在 Limitations and extensions 一节明确说：(a) STC 的 scaling properties 可能与 SC 不同，需要更多算力预算；(b) 该框架可向更多领域（如 image style transfer、discrete sequence、optimization）推广。 基于本文成果可延伸的方向：(1) 与生成式蛋白设计结合——既然 DFM 在 TCR 上把 MMD 砍掉 63%，可以推广到抗体亲和力成熟、酶设计等真正需要"从一个序列分布到另一个"的场景；(2) 与最优传输理论结合——本文 Fig. 6 显示 latent interpolant 几乎与 OT displacement interpolation 重合，可以发展 latent-space Wasserstein Wormhole 类的工作，把 GDE 嵌入与 OT 几何耦合起来；(3) 把分布编码器扩展到图结构（蛋白质相互作用网络、单细胞邻域图）或多模态（RNA+ATAC+Protein），目前 ESM2-based encoder 是 frozen backbone，未来可用 end-to-end 微调；(4) 因果推断扩展——DCT 的 source-target 条件化天然适合反事实推断（如何用 source 分布 + treatment embedding 推断反事实结果），可与 do-calculus 或 instrumental variable 结合；(5) 大规模预训练——GDE 与 FM 都可作为预训练目标，先在百万级分布上预训练分布编码器，再下游微调；(6) 与 LLM agent 结合——把 DCT 当作工具，让 LLM 在科学推理时调用 DCT 做"如果改变实验条件，细胞分布会怎样"的预测。
复现评估
复现评估较高：作者开源了 codebase（abstract 末尾 "The codebase is available here" 注明 GitHub 链接），附录提供了完整超参（Tab. 6 Gaussian、Tab. 20 TCR）与模型架构（App. C.2 标准分布编码器、App. C.3/C.4 SWD/Energy/FM 详细定义、App. D 三种评估指标、App. E Gaussian 实验、App. F MNIST-Colors、App. G scRNA-seq、App. H PDO、App. I lineage tracing、App. J TCR）。数据来源全部公开（cellxgene、Weinreb et al. 2020 论文网站、Schultheiß et al. 2020）。算力要求中等：连续数据实验用 4 层 MLP + 2 层 DeepSets encoder，单 GPU（如 A100）即可；TCR 实验用 8M 参数 ESM2 + ProGen2-medium，batch=1, lr=$2\times10^{-5}$，可在单卡上跑。复现难度等级：合成实验（MVN/GMM/MNIST-Colors）⭐⭐ 简单；scRNA-seq/PDO/lineage tracing ⭐⭐⭐ 中等（数据预处理较繁琐）；TCR ⭐⭐⭐⭐ 较难（依赖 ESM2 与 ProGen2 预训练权重，且 ProGen2 出现 encoder collapse 需要调试）。主要风险点：(a) FM 在合成 Gaussian 上的"异常"OOD 泛化可能让复现者困惑，需要参考作者的解释（可能是低维线性变换的 artifact）；(b) PDO 与 TCR 的小样本预测器方差较大，建议多次 seed 平均；(c) 不同 PyTorch 版本下 SWD 投影的随机性需要固定种子。
论文图表
