Score-based Generative Neural Networks for Large-Scale Optimal Transport. (on ArXiv)
Max Daniels, Tyler Maunu, Paul Hand.

This repository contains code for running synthetic experiments, including those used to generate Table 2 (Gaussian to Gaussian BW-UVP scores) and Figure 4 (Gaussian to Swiss Roll transportation in 2D). A sister repository, used to run all large-scale experiments, can be found here

The required packages can be found in requirements.txt. To create a new conda environment with these packages installed, use

conda create --name <env> --file requirements.txt.

The two main entry points are gaussian_to_gaussian.py and qualitative.py. These files are preconfigured and can be run out-of-the-box, or one can set up custom experiments by changing the configuration inputs, which may be found in

BW-UVP Experiments: run python gaussian_to_gaussian.py to recreate Table 2, which compares our sampling algorithm to the ground truth in BW-UVP distance.

The code uses instances of GaussianConfig (in config.py) to configure experiments. To run customized experiments, such as changing the dimensionality or sampling parameters, pass a customized instance of the config object as in the example above.

Swiss-roll to Gaussian: run python qualitative.py to recreate Figure 4, which simulates our sampling algorithm for a synthetic transportation task between low-dimensional synthetic datasets.

This code uses instances of Config (in config.py) to configure experiments. To run customized experiments, such as changing the source and target datasets or the regularization parameters, pass a customized instance of the configuration object as above.

Note: sampling requires training from scratch a SCONES model for the target distribution. The SCONES model has multiple components: the score-based generative model, the compatibility function, and (optionally) a barycentric projector. Pretrained models for transport to the Swiss Roll dataset, which were used to generate Figure 4, can be found here (merge with the existing pretrained/ directory).