Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/CUAI/Equivariant-Manifold-Flows
[NeurIPS 2021] Equivariant Manifold Flows (https://arxiv.org/abs/2107.08596)
https://github.com/CUAI/Equivariant-Manifold-Flows
Last synced: 3 months ago
JSON representation
[NeurIPS 2021] Equivariant Manifold Flows (https://arxiv.org/abs/2107.08596)
- Host: GitHub
- URL: https://github.com/CUAI/Equivariant-Manifold-Flows
- Owner: CUAI
- Created: 2021-10-18T03:10:23.000Z (about 3 years ago)
- Default Branch: main
- Last Pushed: 2021-12-08T22:06:03.000Z (almost 3 years ago)
- Last Synced: 2024-07-03T23:45:23.857Z (4 months ago)
- Language: Python
- Homepage:
- Size: 24.4 KB
- Stars: 9
- Watchers: 2
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# Equivariant Manifold Flows
We provide the code for [Equivariant Manifold Flows](https://arxiv.org/abs/2107.08596) in this repository.
Summary: We lay the theoretical foundations for learning symmetry-invariant distributions on arbitrary manifolds via our introduced equivariant manifold flows. We demonstrate the utility of our approach by using it to learn gauge invariant densities over SU(n), useful for lattice quantum field theory.
Example learned densities, together with baselines, are given below.
SU(2) | SU(3)
:-------------------------:|:-------------------------:
| 
Sphere |
:-------------------------:
|
Below we have visualized how our equivariant manifold flows learn the `boyda1` and `boyda2` densities on SU(3) (first and second rows in the SU(3) figure above).
SU(3) Boyda1 | SU(3) Boyda2
:-------------------------:|:-------------------------:
| 
## Software Requirements
This codebase requires Python 3 and PyTorch 1.8+.
## Usage
### Demo
The following command learns a equivariant manifold flow model for the `boyda2` density on SU(3). The density learned will be symmetric with respect to conjugation by SU(3), which we prove is an isometry in Appendix A.5 of our paper.
```
python main.py --M SU3 --batch_size 1024 --dev cpu --save_viz --contsave --save_freq 1 --lr 4e-3 --epochs 200 --example_density boyda2
```
Note that 1024 samples are used per epoch. 200 epochs were used for the results in the paper. The learned density is given below, in comparison with the ground truth density:
Ground Truth | Ours (50 epochs) | Ours (200 epochs)
:-------------------------:|:-------------------------:|:---------------------:
|  | 
Observe that even after 50 epochs (only approximately 50,000 samples), our model approaches the ground truth.
### Full Usage
All options are given below:
```
usage: main.py [-h] [--example_density EXAMPLE_DENSITY] [--num_drops NUM_DROPS] [--save_viz] [--save_model] [--dev DEV]
[--M {SU3}] [--contsave] [--save_freq SAVE_FREQ] [--epochs EPOCHS] [--batch_size BATCH_SIZE] [--lr LR]
[--weight_decay WEIGHT_DECAY]
optional arguments:
-h, --help show this help message and exit
--example_density EXAMPLE_DENSITY
example density to learn. Options are: boyda1 | boyda2 | boyda3 (different densities defined on
SU(3))
--num_drops NUM_DROPS
number of times to drop the learning rate
--save_viz Save a visualization of the learned density once training is completed
--save_model Save a visualization of the learned sampler
--dev DEV Default device; cpu preferred due to poor support of complex tensor differentiation on cuda
--M {SU3} Manifold over which to learn
--contsave Continuously save intermediate flow visualization in contsave/
--save_freq SAVE_FREQ
frequency of continuous saving of intermediate flows
--epochs EPOCHS
--batch_size BATCH_SIZE
--lr LR
--weight_decay WEIGHT_DECAY
```
## Attribution
If you use this code or our results in your research, please cite:
```
@article{Katsman2021EquivariantMF,
title={Equivariant Manifold Flows},
author={Isay Katsman and Aaron Lou and Derek Lim and Qingxuan Jiang and Ser-Nam Lim and Christopher De Sa},
journal={ArXiv},
year={2021},
volume={abs/2107.08596}
}
```