Entropic Neural Optimal Transport via Diffusion Processes

Authors: Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry P. Vetrov, Evgeny Burnaev

NeurIPS 2023 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental Empirically, we show the performance of the method on several large-scale EOT tasks. The code for the ENOT solver can be found at https://github.com/ngushchin/Entropic Neural Optimal Transport. In this section, we qualitatively and quantitatively illustrate the performance of our algorithm in several entropic OT tasks.
Researcher Affiliation Collaboration Nikita Gushchin Skoltech Moscow, Russia n.gushchin@skoltech.ru Alexander Kolesov Skoltech Moscow, Russia a.kolesov@skoltech.ru Alexander Korotin Skoltech , AIRI Moscow, Russia a.korotin@skoltech.ru Dmitry Vetrov HSE University , AIRI Moscow, Russia vetrovd@yandex.ru Evgeny Burnaev Skoltech , AIRI Moscow, Russia e.burnaev@skoltech.ru
Pseudocode Yes Algorithm 1: Entropic Neural OT (ENOT) Input: samples from distributions P0, P1; Wiener prior noise variance ϵ 0; drift network fθ : RD [0, 1] RD; beta network βϕ : RD R; number of steps N for Eul-Mar (App C); number of inner iterations Kf. Output: drift f θ of Tf θ solving DSB (11). Algorithm 2: Euler-Maruyama algorithm
Open Source Code Yes The code for the ENOT solver can be found at https://github.com/ngushchin/Entropic Neural Optimal Transport. The code is written in Py Torch and is publicly available at https://github.com/ngushchin/Entropic Neural Optimal Transport
Open Datasets Yes We consider the entropic OT problem between colorized MNIST digits of classes '2' (P0) and '3' (P1). We consider the problem of unpaired image super-resolution for the 64 64 aligned faces of Celeb A dataset [35].
Dataset Splits No We do the unpaired train-test split as follows: we split the dataset into 3 parts: 90k (train A1), 90k (train B1), 20k (test C1) samples. For training in the unpaired setup, we use parts A0 (degraded faces, P0) and B1 (clean faces, P1). The paper specifies train and test splits, but no explicit validation split.
Hardware Specification Yes In the most challenging task (M5.4), ENOT converges in one week on 2 A100 GPUs.
Software Dependencies No The code is written in Py Torch. The paper mentions PyTorch but does not provide a specific version number for it or any other software dependency.
Experiment Setup Yes In all the cases, we use N = 10 discretization steps for solving SDE by Euler-Maruyama Algorithm 2, Adam with lr = 10 4, batch size 512. We train the model for 20000 total iterations of βϕ, and on each of them, we do Kf = 10 updates for the SDE drift function fθ. We use Adam with lr = 10 4, batch size 64 and 10:1 update ratio for fθ/βϕ.