[論文レビュー] 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）を理論的根拠に基づいて組み込むための方法を提案する。これは、出発測度と到達測度を再スケーリングすることで実現され、非対応ドメイン変換の性能を顕著に向上させる。本手法は、単細胞データおよび画像変換の両分野において特徴の保存性と軌道推定を向上させ、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.

研究の動機と目的

• バランス型最適輸送が非対応ドメイン変換において示す限界、特に外れ値や分布シフトに対しての感受性の高さを是正すること。
• 敵対的訓練を必要とせず、任意のニューラル・モンジュ写像推定器に不均衡性を統合可能な汎用的フレームワークの開発。
• 質量のずれをモデル化することで、単細胞データにおける軌道推定と、摂動に対する細胞応答の予測を改善すること。
• OTフローマッチング（OT-FM）フレームワークを不均衡設定に拡張（UOT-FM）し、画像変換タスクにおける性能を向上させること。
• 現実世界の応用において、不均衡性がより生物学的に妥当で正確なマッピングをもたらすことを実証すること。

提案手法

• 質量が異なる測度間のモンジュ写像計算を、再スケーリングされた測度間のバランス型写像に再定式化するための再スケーリング方式を提案。
• 理論的に、測度の再スケーリングを用いることで、任意のモンジュ写像推定器に不均衡性を組み込むことが可能であり、質量のずれに対するペナルティが保証されることを証明。
• OT-ICNN、モンジュギャップ、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 敵対的訓練を一切用いずに、不均衡最適輸送をニューラル・モンジュ写像推定器に体系的に統合できるか？
• RQ2 不均衡性を組み込むことで、単細胞データにおける細胞軌道推定の正確性と生物学的妥当性が向上するか？
• RQ3 不均衡性は、非対応画像変換タスクにおける特徴の保存性と輸送コストにどのように影響を与えるか？
• RQ4 UOT-FMフレームワークは、計算効率を維持しつつ、標準的なOT-FMを上回る性能を示せるか？
• RQ5 不均衡性は、ドメイン変換における分布シフトや外れ値の影響をどの程度軽減できるか？

主な発見

• UOT-FMは、非対応画像変換においてOT-FMを上回る優れた性能を達成し、特に入力特徴の保存性と輸送コストの低減が顕著である。
• EMNISTデータセットでは、不均衡性により8→B、1→Iといった正しいマッピングが正しく実現され、バランス型OTで見られる8→{O,B}のような誤ったマッピングを回避する。
• 単細胞データでは、UOT-ICNNがNgn3高発現細胞が内分泌系へ移行するようなより生物学的に妥当な細胞タイプの遷移を回復し、遷移確率の精度が向上する。
• 画像変換において、UOT-FMはバランス型ベースライン比で輸送コストを最大25％まで低減し、質量輸送の効率性が向上している。
• 自然画像変換ベンチマークでも競争力ある性能を維持しつつ、分布シフト下でもより優れた一般化性能を示す。
• 理論的分析により、出発測度と到達測度の再スケーリングが、不均衡状態下でも正当なモンジュ写像計算を可能にすることを確認し、本手法に理論的根拠を与える。