[论文解读] Supervised Training of Conditional Monge Maps

Charlotte Bunne, Andreas Krause+1 more|arXiv (Cornell University)|Jun 28, 2022

Adversarial Robustness in Machine Learning被引用 14

一句话总结

该论文提出了CondOT，一种监督学习框架，通过部分输入凸神经网络（PICNNs）训练条件Monge映射，以建模概率测度之间的上下文相关最优传输。通过利用高斯近似进行上下文感知初始化，并借助深度集合组合方法推广至未见的扰动，CondOT能够准确预测单细胞对新型药物组合或基因扰动的响应。

ABSTRACT

Optimal transport (OT) theory describes general principles to define and select, among many possible choices, the most efficient way to map a probability measure onto another. That theory has been mostly used to estimate, given a pair of source and target probability measures $(μ, ν)$, a parameterized map $T_θ$ that can efficiently map $μ$ onto $ν$. In many applications, such as predicting cell responses to treatments, pairs of input/output data measures $(μ, ν)$ that define optimal transport problems do not arise in isolation but are associated with a context $c$, as for instance a treatment when comparing populations of untreated and treated cells. To account for that context in OT estimation, we introduce CondOT, a multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures $\left(μ_i, ν_i ight)$ tagged with a context label $c_i$. CondOT learns a global map $\mathcal{T}_θ$ conditioned on context that is not only expected to fit all labeled pairs in the dataset $\left\{\left(c_i,\left(μ_i, ν_i ight) ight) ight\}$, i.e., $\mathcal{T}_θ\left(c_i ight) \sharp μ_i \approx ν_i$, but should also generalize to produce meaningful maps $\mathcal{T}_θ\left(c_{ ext {new }} ight)$ when conditioned on unseen contexts $c_{ ext {new }}$. Our approach harnesses and provides a novel usage for partially input convex neural networks, for which we introduce a robust and efficient initialization strategy inspired by Gaussian approximations. We demonstrate the ability of CondOT to infer the effect of an arbitrary combination of genetic or therapeutic perturbations on single cells, using only observations of the effects of said perturbations separately.

研究动机与目标

• 解决仅观测到单个扰动效应时，预测新型药物或基因组合对单细胞影响的挑战。
• 通过将传输映射基于治疗方式、剂量或基因扰动等上下文变量进行条件化，将最优传输推广至非固定源-目标对。
• 通过学习一个全局、参数化的条件传输映射，实现对未见上下文的泛化。
• 通过基于高斯近似的闭式解初始化策略，克服输入凸神经网络在训练中的不稳定性。

提出的方法

• 将条件传输映射 $\mathcal{T}_{\theta}(c)$ 参数化为部分输入凸神经网络（PICNN）的梯度，确保在输入空间 $\mu$ 上的凸性。
• 利用上下文向量 $c$ 调制PICNN的权重和偏置，从而构建一个生成上下文特定传输映射的超网络。
• 使用 $\mu_i$ 和 $\nu_i$ 的高斯近似之间最优传输的闭式解初始化PICNN参数，确保训练初始映射的稳定性和合理性。
• 在训练过程中，通过权重裁剪（将负权重设为零）强制PICNN保持凸性，并通过可微罚项对偶网络中的负权重进行松弛，以放宽凸性约束。
• 采用深度集合架构将多个扰动嵌入组合为联合上下文表征，实现对未见扰动组合的泛化。
• 使用主成分分析（PCA）或自编码器将高维单细胞RNA-seq数据投影至低维潜在空间，以提升训练效率与泛化能力。

实验结果

研究问题

• RQ1 在仅观测到单个扰动效应的前提下，能否学习到一个单一、参数化的传输映射，使其能泛化至未见的扰动组合？
• RQ2 如何通过高斯近似的闭式解，稳定输入凸神经网络在最优传输中的训练过程？
• RQ3 与独热编码相比，基于深度集合的组合模块在未见扰动组合上的泛化能力有多强？
• RQ4 所提出的条件最优传输框架在预测单细胞对新型药物或基因扰动响应方面，是否优于现有基线方法？

主要发现

• CondOT通过深度集合模块联合学习单个扰动的嵌入，成功实现对未见扰动组合的泛化，其零样本泛化性能优于独热编码。
• 基于高斯近似的初始化策略显著提升了PICNN基传输映射的训练稳定性与收敛速度。
• 在MOFA+和MOFA++数据集上的验证表明，该框架在预测单细胞对新型药物组合响应方面达到了最先进性能。
• 使用基于自编码器的投影方法比PCA更好地保留了高维结构，从而在下游预测任务中取得了更优性能。
• 该模型在未见扰动上表现出有效的泛化能力，证明了其在零样本设置下的鲁棒性。
• 在对偶网络中引入对负权重的可微罚项，可在保持理论保证的前提下有效放松凸性约束。