Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation
Authors: Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, Fabian J Theis
ICLR 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | To validate the versatility of our approach, we showcase its applicability on both synthetic and real-world data, utilizing different neural Monge map estimators. We highlight the critical role of incorporating unbalancedness to infer trajectories in developmental single-cell data using ICNN-based estimators (OT-ICNN) (Makkuva et al., 2020) and to predict cellular responses to perturbations with the Monge gap (Uscidda & Cuturi, 2023). ... We demonstrate that unbalancedness is crucial for obtaining meaningful matches when translating digits to letters in the EMNIST dataset (Cohen et al., 2017). Additionally, we benchmark OT-FM on unpaired natural image translation and show that it achieves competitive results compared to established methods. |
| Researcher Affiliation | Academia | 1T ubingen AI Center 2Helmholtz Munich 3TU Munich 4CREST-ENSAE 5Munich Center for Machine Learning (MCML) |
| Pseudocode | Yes | Algorithm 1: Neural Unbalanced Monge maps; Algorithm 2: Unbalanced Optimal Transport Flow Matching (UOT-FM) |
| Open Source Code | Yes | Additionally, the code to reproduce our experiments can be found at https://github.com/Explainable ML/uot-fm. |
| Open Datasets | Yes | We demonstrate that unbalancedness is crucial for obtaining meaningful matches when translating digits to letters in the EMNIST dataset (Cohen et al., 2017). ... Celeb A dataset (Liu et al., 2015)... CIFAR-10 (Krizhevsky et al.). |
| Dataset Splits | No | The paper describes using test sets and training, but does not provide specific percentages, counts, or detailed methodologies for training/validation/test splits across all experiments, often referring to common practices or implicitly defining test sets without full split details. |
| Hardware Specification | No | The paper mentions running experiments and GPU batch size, but does not specify the exact models of GPUs, CPUs, or other hardware used for the experiments. |
| Software Dependencies | No | Our code is based on Jax (Bradbury et al., 2018) while utilizing parts of the Deep Mind Jax ecosystem (Babuschkin et al., 2020). To compute the entropic unbalanced coupling we leverage the ott-jax library (Cuturi et al., 2022)... Additionally, we use the Adam optimizer (Kingma & Ba, 2014)... During inference, we solve for pt at t = 1 using the adaptive step-size solver Tsit5 with atol=rtol=1e-5 implemented in the diffrax library (Kidger, 2021). |
| Experiment Setup | Yes | For each of the Celeb A image translation tasks, we use a very similar hyperparameter setup based upon the UNet architecture used in Dhariwal & Nichol (2021), where all FM models were trained with the exact same setup. We report these in Table 12. For EMNIST we leverage the MLPMixer architecture (Tolstikhin et al., 2021) with hyperparameters detailed in Table 13. ... For each drug, we model the neural map Tθ using an MLP with hidden layer size [max(128, 2d), max(128, 2d), max(64, d), max(64, d)] where d is the dimension of the data. We train it for niter = 10, 000 iterations with a batch size n = 1, 024 and the Adam optimizer (Kingma & Ba, 2014) using a learning rate η = 0.001, along with a polynomial schedule of power p = 1.5 to decrease it to 10 5. |