Optimal Transport (Spring 2022)

Project (Latest Instructions Update: Apr 12 2022)

The deadline for submission is changed to May 4.

You can either send your

  • pdf

  • colab link with experiments/code

to marcocuturicameto+assignment@gmail.com (if you use a different email alias, your assignment may risk ending “lost” in my inbox).

or a pull-request to the OTT-JAX toolbox if you choose the first project.

Choose one topic from those presented below. If you wish to explore a different direction, send me a proposal by email.

Contribute to the OTT toolbox with a tutorial or a feature

The OTT-JAX toolbox provides implementations of optimal transport tools, coded in the JAX (Python) framework. The toolbox is open source and anyone can contribute to it. A crucial aspect of such toolboxes is to provide tutorials that help users play with these tools. After looking at the code and at available tutorials (it's important that you look at them to get a feel of what the expectation is), I propose that you run the experiments of one of the family of papers below and package this in a nice notebook that can then be used as tutorial material for the toolbox (your contribution will stay there and you will be credited!). All of the experiments below rely on the Sinkhorn algorithm as a workhorse, I propose that you also explore in your notebooks with the option to use low-rank OT approaches (this might speed up things, maybe with slightly worse performance).

  • Supervised Word Mover's Distance: this will involve differentiating the Sinkhorn divergence between point clouds. You should not reimplement the gradient approach in the paper but experiment directly with the OTT-JAX's gradient functionalities.

  • Experiment with datasets provided in the SCOT package, which concerns applications to registrating single cell multi-omics data.

  • Explore the Domain Adaptation for Optimal Transport literature and implement, using OTT, an application of your liking described in that literature.

  • Any other paper of your liking that could illustrate a numerical approach for optimal transport.

Please consider that a tutorial is not a place to add new features to OTT, this is a different goal, contact me if you would prefer your contribution to be the addition of a new functionality (which I expect to be substantially harder because it would require understanding the inner workings of the toolbox and JAX knowledge).

Please consider that I will only be able to accept a few of the PRs to remain in the toolbox, but I am happy to create a special section of miscellaneous examples in the toolbox where I will put your work if it does not pass the bar. Submitting a PR will also likely involve reviews and hence modifications. You are encouraged to raise issues if you encounter unexpected problems when writing the tutorial that you would like to share with the community.

Implement OT-GAN

In this assignment you will implement GAN variants that built on optimal transport, using a different approach than the Wasserstein-GAN paper, as proposed in the OT-Gan paper (you can also start using this variant of the same idea).

Please train your model on simple datasets (e.g. MNIST digits or simpler images). Code everything from scratch using a differentiale programming framework (TF, Jax, Autograd or PyTorch)

Does Wasserstein-GAN approximate Wasserstein distances?

The Wasserstein-GAN paper proposes a proxy for the 1-Wasserstein distance that uses neural networks. While that proxy seems to work for the task of training GANs, it is not well understood whether that approach can approximate, numerically, the Wasserstein distance. In this assignment, you will implement the W-GAN approach to solve OT and benchmark it against other approaches (e.g. Sinkhorn divergence) to study its ability to compute a quantity that is truly similar to “true” optimal transport. You should restrict yourself to low-dimensional settings (e.g. 1/2D) or to settings for which the ground truth OT distance is known (i.e. Gaussians or elliptically contoured distributions).