[论文解读] Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

Luca Eyring, Dominik Klein+5 more|arXiv (Cornell University)|Nov 25, 2023

Cancer-related molecular mechanisms research被引用 4

一句话总结

该论文提出了一种理论基础扎实的方法，通过重新缩放源分布和目标分布，将非平衡最优传输（UOT）整合到任意神经Monge映射估计器中，显著提升了无配对域迁移的性能。该方法在单细胞数据和图像迁移中均增强了特征保留和轨迹推断能力，UOT-FM相较于OT-FM显著降低了传输成本，并更好地处理了分布偏移问题。

ABSTRACT

In optimal transport (OT), a Monge map is known as a mapping that transports a source distribution to a target distribution in the most cost-efficient way. Recently, multiple neural estimators for Monge maps have been developed and applied in diverse unpaired domain translation tasks, e.g. in single-cell biology and computer vision. However, the classic OT framework enforces mass conservation, which makes it prone to outliers and limits its applicability in real-world scenarios. The latter can be particularly harmful in OT domain translation tasks, where the relative position of a sample within a distribution is explicitly taken into account. While unbalanced OT tackles this challenge in the discrete setting, its integration into neural Monge map estimators has received limited attention. We propose a theoretically grounded method to incorporate unbalancedness into any Monge map estimator. We improve existing estimators to model cell trajectories over time and to predict cellular responses to perturbations. Moreover, our approach seamlessly integrates with the OT flow matching (OT-FM) framework. While we show that OT-FM performs competitively in image translation, we further improve performance by incorporating unbalancedness (UOT-FM), which better preserves relevant features. We hence establish UOT-FM as a principled method for unpaired image translation.

研究动机与目标

• 为解决平衡最优传输在无配对域迁移中的局限性，特别是其对异常值和分布偏移的敏感性。
• 开发一种可泛化的框架，将非平衡性整合到任意神经Monge映射估计器中，而无需依赖对抗性训练。
• 通过建模质量偏差，改进单细胞数据中的轨迹推断和对扰动的细胞反应预测。
• 将OT流匹配（OT-FM）框架扩展至非平衡设置（UOT-FM），提升图像迁移任务中的性能。
• 证明非平衡性在真实应用场景中可产生更具生物学合理性和准确性的映射。

提出的方法

• 提出一种重缩放方案，将质量不等的测度之间的Monge映射计算重新表述为重缩放后测度之间的平衡映射。
• 理论上证明，可通过测度重缩放将非平衡性整合到任意Monge映射估计器中，确保质量偏差的惩罚。
• 将该框架应用于多种估计器：OT-ICNN、Monge gap和OT-FM，实现用于连续归一化流的UOT-FM。
• 使用小批量最优传输耦合训练UOT-FM，无需模拟或基于流的生成建模。
• 采用带熵正则化和非平衡惩罚的传输成本最小化目标，以稳定训练过程。
• 在合成数据、单细胞轨迹和图像迁移（EMNIST）上验证该方法，与基线OT方法及成熟方法进行比较。

Figure 1: Comparison of balanced and unbalanced Monge map computed on the EMNIST dataset translating digits $\rightarrow$ letters . Source and target distribution are rescaled leveraging the unbalanced OT coupling. The computed balanced mapping includes $8\rightarrow\{O,B\}$ , and $1\rightarrow\{O,I

实验结果

研究问题

• RQ1 非平衡最优传输能否在不依赖对抗性训练的前提下，系统性地整合到神经Monge映射估计器中？
• RQ2 引入非平衡性是否能提升单细胞数据中细胞轨迹推断的准确性和生物学合理性？
• RQ3 在无配对图像迁移任务中，非平衡性如何影响特征保留和传输成本？
• RQ4 UOT-FM框架是否能在图像迁移中超越标准OT-FM，同时保持计算效率？
• RQ5 非平衡性在多大程度上能缓解域迁移中分布偏移和异常值的影响？

主要发现

• UOT-FM在无配对图像迁移中表现优于OT-FM，尤其在保留输入特征和降低传输成本方面。
• 在EMNIST数据集上，非平衡性正确地将数字映射为字母（如8→B，1→I），避免了平衡OT中出现的错误映射（如8→{O,B}）。
• 在单细胞数据中，UOT-ICNN恢复了更多具有生物学合理性的细胞类型转变，如Ngn3高表达细胞向内分泌谱系的进展，且过渡概率准确性更高。
• 在图像迁移中，与平衡基线相比，该方法将传输成本降低了最多25%，表明质量传输更高效。
• UOT-FM在自然图像迁移基准上保持了具有竞争力的性能，同时在分布偏移下展现出更好的泛化能力。
• 理论分析证实，通过重缩放源分布和目标分布，可在非平衡条件下实现有效的Monge映射计算，为该方法提供了坚实的理论基础。