Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, Fabian Theis

TL;DR

This work proposes a theoretically grounded method to incorporate unbalancedness into any Monge map estimator and establishes UOT-FM as a principled method for unpaired image translation.

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.

Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

TL;DR

This work proposes a theoretically grounded method to incorporate unbalancedness into any Monge map estimator and establishes UOT-FM as a principled method for unpaired image translation.

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.
Paper Structure (46 sections, 2 theorems, 21 equations, 12 figures, 13 tables, 2 algorithms)

This paper contains 46 sections, 2 theorems, 21 equations, 12 figures, 13 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Unbalanced Neural Monge Maps
  4. Estimation of Unbalanced Monge Maps
  5. Related Work
  6. Experiments
  7. Single-Cell Trajectory Inference
  8. Modeling Perturbation Responses
  9. Unpaired Image Translation
  10. Conclusion
  11. Reproducibility
  12. Acknowledgements
  13. Proofs
  14. Coupling computation
  15. Entropic Regularization
  16. ...and 31 more sections

Key Result

Proposition 3.0

Let $\pi_\mathrm{UOT}$ be the solution of problem eq:unbalanced-kantorovich-problem between $\mu, \nu \in \mathcal{M}^+(\Omega)$, for $\tau_1,\tau_2 >0$. Then, the following holds:

Figures (12)

  • 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\}$ because of the distribution shift between digits and letters. With unbalancedness $8 \rightarrow B$, and $1 \rightarrow I$ are recovered.
  • Figure 2: Different maps on data drawn from a mixture of uniform distribution, where the density in the bottom left and the top right ($\frac{3}{5}$) is higher than in the top left and bottom right ($\frac{2}{5}$) (Appendix \ref{['app:simulated_data']}). Besides the data in Figure \ref{['subfig:unbalanced_data']}, the first row shows results of discrete balanced OT (\ref{['subfig:standard_ot']}), and discrete unbalanced OT with two different degrees of unbalancedness $\tau$ (\ref{['subfig:standard_ot_tau0.99']}, \ref{['subfig:standard_ot_tau0.9']}). The second row shows the maps obtained by FM with independent coupling (\ref{['subfig:uniform_fm']}), OT-FM (\ref{['subfig:uniform_ot-fm']}), and UOT-FM (\ref{['subfig:uniform_uot-fm_099']}, \ref{['subfig:uniform_uot-fm_09']}).
  • Figure 3: Velocity stream embedding plots (Appendix \ref{['app:velstream']}). The orange box highlights the direction of Ngn3 EP cells. With OT-ICNN these move to the "right", which contradicts biological ground truth. This is due to the distribution shift shown in Appendix \ref{['app:pancreas_data']} and demonstrates the need to incorporate unbalancedness.
  • Figure 4: Fitting of a transport map $\hat{T}$ to predict the responses of cell populations to cancer treatments on 4i (upper plot), using balanced (OT-MG) and unbalanced Monge maps (UOT-MG) fitted with the Monge gap. A point below the diagonal indicates that unbalancedeness improves performance.
  • Figure 5: CelebA 256x256 translated test samples with FM, OT-FM, and UOT-FM.
  • ...and 7 more figures

Theorems & Definitions (5)

  • Proposition 3.0: Re-balancing the UOT problem
  • Remark 3.1
  • Definition 3.2: Unbalanced Monge maps
  • Proposition A.0: Re-balancing the UOT problem
  • proof