Manifold-learning flows
Johann Brehmer and Kyle Cranmer 2019-2020
In the paper Flows for simultaneous manifold learning and density estimation we introduce manifold-learning flows or ℳ-flows, a new class of generative models that simultaneously learn the data manifold as well as a tractable probability density on that manifold. This repository contains our implementation of ℳ-flows, as well as the code for our experiments with them.
Getting started
Please make sure your Python environment satisfies the requirements in the environment.yml. To use the OT training, please also follow the installation instructions for geomloss.
Data sets
| Data set | Data dimension | Manifold dimension | Model parameters | Arguments to train.py, and evaluate.py |
|---|---|---|---|---|
Gaussian on an n-sphere |
d |
n |
- | --dataset spherical_gaussian --truelatentdim n --datadim d --epsilon eps |
Conditional Gaussian on a n-sphere |
d |
n |
2 | --dataset conditional_spherical_gaussian --truelatentdim n --datadim d |
| Mixture model on a polynomial manifold | 3 | 2 | 1 | --dataset power |
| Lorenz system | 3 | 2 | 0 | --dataset lorenz |
| Particle physics | 40 | 14 | 2 | --dataset lhc40d |
| 2-D StyleGAN image manifold | 64 x 64 x 3 | 2 | 0 | --dataset gan2d |
| 64-D StyleGAN image manifold | 64 x 64 x 3 | 64 | 1 | --dataset gan64d |
| CelebA-HQ | 64 x 64 x 3 | ? | 0 | --dataset celeba |
| ImageNet | 64 x 64 x 3 | ? | 0 | --dataset imagenet |
The data from most data sets should automatically download when required. It is not necessary to generate any data yourself anymore. If there is a problem with that, please let us know.
Training
See experiments/train.py -h. The configurations for the models in the paper can be found in experiments/configs.
Note that the algorithms have different internal names from the acronyms in the paper:
| Model (algorithm) | Arguments to train.py |
|---|---|
| Ambient flow (AF) | --algorithm flow |
| Flow on manifold (FOM) | --algorithm mf --specified |
| Pseudo-invertible encoder (PIE) | --algorithm pie |
| ℳ-flow, simultaneous training (not recommended) | --algorithm mf |
| ℳ-flow, alternating M/D training | --algorithm mf --alternate |
| ℳ-flow, sequential M/D training | --algorithm mf --sequential |
| ℳ-flow, Optimal Transport training | --algorithm gamf |
| ℳ-flow, alternating Optimal Transport training | --algorithm gamf --alternate |
| ℳ_e-flow, simultaneous training (not recommended) | --algorithm emf |
| ℳ_e-flow, alternating M/D training | --algorithm emf --alternate |
| ℳ_e-flow, sequential M/D training | --algorithm emf --sequential |
Evaluation
See experiments/evaluate.py -h and the notebooks in experiments/notebooks. Note that the algorithms have different internal names from the acronyms in the paper:
| Model (algorithm) | Arguments to train.py |
|---|---|
| Ambient flow (AF) | --algorithm flow |
| Flow on manifold (FOM) | --algorithm mf --specified |
| Pseudo-invertible encoder (PIE) | --algorithm pie |
| ℳ-flow (except when OT-trained) | --algorithm mf |
| ℳ-flow, OT training | --algorithm gamf |
Acknowledgements
The code is largely based on the excellent Neural Spline Flow code base by C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios, see 1906.04032 for their paper.