Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation (2311.15100v2)

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.

• The paper introduces unbalanced optimal transport into neural Monge maps to overcome mass conservation limits in unpaired domain translation.
• It demonstrates significant improvements by accurately aligning distributions in tasks like single-cell trajectory inference, drug perturbation modeling, and image translation.
• Empirical results reveal higher transition correctness and competitive FID scores, highlighting the method's potential for diverse applications.

Analysis of "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation"

The paper under examination, "Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation," explores the integration of unbalanced optimal transport (UOT) into neural Monge map estimators to enhance the performance in unpaired domain translation tasks. The authors argue that traditional optimal transport (OT) assumptions, which require mass conservation, limit the applicability of neural Monge maps in real-world scenarios due to their sensitivity to outliers and distribution shifts. By incorporating unbalancedness, they propose a method that significantly reduces these limitations.

The paper begins with an overview of OT, emphasizing the Monge map as an efficient means of transporting a source distribution to a target distribution. While neural Monge map estimators have been applied successfully in tasks such as computer vision and single-cell biology, the authors point out that traditional OT frameworks enforce mass conservation, complicating applications that involve outliers or imbalances. Their solution is to introduce unbalanced OT, which removes this constraint by allowing for penalization of mass deviations.

With this new methodology, the authors propose a theoretically justified scheme that can incorporate unbalancedness into any existing Monge map estimator. Their framework involves re-scaling the source and target distributions, allowing for a more adaptable and potentially more accurate translation of data between distributions. The authors validate their approach in several key domains.

Key Numerical Results and Empirical Evaluations

1. Single-Cell Data Processing: The application of UOT to single-cell trajectory inference demonstrates significant improvements over traditional OT. The authors provide quantitative comparisons showing that unbalanced Monge maps yield a higher percentage of correct cell type transitions. Enhanced performance was particularly notable in predicting transitions for the Ngn3 EP cell population, which traditionally suffered from distribution shifts.
2. Cellular Responses to Drug Perturbations: In the context of modeling drug perturbations, UOT again shows superiority. The empirical results indicate a robust enhancement in predictive performance across multiple drugs, as measured by the Sinkhorn divergence and L2 drug signature.
3. Unpaired Image Translation: When applied to image translation tasks, especially with the use of the EMNIST dataset, UOT shows a marked improvement in aligning digits with corresponding letters, despite the inherent distribution shifts. For complex tasks on the CelebA dataset, UOT-FM leverages the unbalanced OT framework to achieve competitive FID scores and supports the preservation of relevant input features.

Implications and Future Directions

The integration of unbalanced OT in neural Monge maps opens up new avenues in domain translation applications where traditional methodological limitations exist due to distribution shifts or outliers. This paper's approach could have a broad impact on various fields where OT-based methods are employed. Furthermore, the use of UOT in mini-batch training scenarios could lead to more stable and quicker convergence of neural models, a promising insight for deep learning applications.

Future research could explore the efficacy of different cost functions tailored specifically to various data domains and geometries, as the paper primarily employs the squared Euclidean distance due to its tractability. Moreover, while the empirical results are promising, further work could focus on refining hyperparameter selection, especially concerning unbalancedness parameters, which currently require grid search optimization.

In conclusion, the paper's contribution to the field of optimal transport and neural estimation is significant, providing both theoretical insights and practical improvements across multiple challenging tasks. The proposed framework not only enhances existing methods but also provides a strong foundation for developing even more advanced and adaptable domain translation models in the future.