Kernel Neural Optimal Transport (NOT)

This is the official Python implementation of the ICLR 2023 paper Kernel Neural Optimal Transport (paper on openreview) by Alexander Korotin, Daniil Selikhanovych and Evgeny Burnaev.

The repository contains reproducible PyTorch source code for computing optimal transport (OT) plans for the weak kernel costs in high dimensions with neural networks. Examples are provided for toy problems (1D, 2D) and for the unpaired image-to-image translation task for various pairs of datasets. The implementation of the latter contains distance-based, laplacian, exponential and bilinear kernels.

Presentations

• Short Talk by Daniil Selikhanovych at ICLR 2023 (May 2023, EN)

Seminars and Materials

• Seminar and solutions on NOT with weak costs (TO DO);
• Vector *.svg sources of the figures in the paper (use inkscape to edit);

Related repositories

• Repository for Neural Optimal Transport paper (ICLR 2023).
• Repository for Kantorovich Strikes Back! Wasserstein GANs are not Optimal Transport? paper (NeurIPS 2022).
• Repository for Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark paper (NeurIPS 2021).

Citation

@inproceedings{
    korotin2023kernel,
    title={Kernel Neural Optimal Transport},
    author={Korotin, Alexander and Selikhanovych, Daniil and Burnaev, Evgeny},
    booktitle={International Conference on Learning Representations},
    year={2023},
    url={https://openreview.net/forum?id=Zuc_MHtUma4}
}

Application to Unpaired Image-to-Image Translation Task

The unpaired domain translation task can be posed as an OT problem. Our NOT algorithm with kernel costs is applicable here. It searches for a transport map with the minimal transport cost (we use kernel $\ell^{2}$). It naturally aims to preserve certain image attributes during the translation.

Compared to the popular image-to-image translation models based on GANs or diffusion models, our method provides the following key advantages

• controlable amount of diversity in generated samples (without any duct tape or heuristics);
• better interpretability of the learned map.

Qualitative examples are shown below for various pairs of datasets (at resolutions $128\times 128$).

One-to-many translation (weak kernel cost)

We show unpaired translition with NOT with the $\gamma$-weak kernel cost on celeba (female) → anime, outdoor → church, handbags → shoes, texture → shoes, texture → handbags, handbags → shoes, shoes → handbags datasets.

Controlling the amount of diversity

Our method offers a single parameter $\gamma\in[0,+\infty)$ in the weak kernel cost to control the amount of diversity.

Repository structure

The implementation is GPU-based with the multi-GPU support. Tested with torch== 1.9.0 and 1-4 Tesla V100.

All the experiments are issued in the form of pretty self-explanatory jupyter notebooks (notebooks/). For convenience, the majority of the evaluation output is preserved. Auxilary source code is moved to .py modules (src/).

• notebooks/KNOT_toy_1D.ipynb - toy experiments in 1D (weak kernel costs);
• notebooks/KNOT_toy_2D.ipynb - toy experiments in 2D (weak kernel costs);
• notebooks/KNOT_training_weak.ipynb - unpaired image-to-image translation (one-to-many, weak kernel costs);
• notebooks/KNOT_plots.ipynb - plotting the translation results (pre-trained models are needed);
• stats/compute_stats.ipynb - pre-compute InceptionV3 statistics to speed up test FID computation;

Datasets

• Aligned anime faces (105GB) should be pre-processed with datasets/preprocess.ipynb;
• CelebA faces requires datasets/list_attr_celeba.ipynb;
• Handbags, shoes, churches, outdoor datasets;
• Describable Textures Dataset (DTD).

The dataloaders can be created by load_dataset function from src/tools.py. The latter four datasets get loaded directly to RAM.

Credits

• Weights & Biases developer tools for machine learning;
• pytorch-fid repo to compute FID score;
• UNet architecture for transporter network;
• ResNet architectures for generator and discriminator;
• Inkscape for the awesome editor for vector graphics;
• AugCycleGAN repository for CondInstanceNorm layer.