Sampling via Föllmer Flow

This repo contains the implementation for the paper Sampling via Föllmer Flow

Brief

We introduce a novel unit-time ODE flow called the Föllmer flow, which efficiently transforms a Gaussian measure into a desired target measure at time 1. We apply Euler's method, where the velocity field is calculated either analytically or through Monte Carlo approximation using Gaussian samples. Time evolution of Föllmer flow on 8-mode gaussian mixtures is illustrated below:

Through numerical experiments on mixture distributions in 1D, 2D, and high-dimensional spaces, we demonstrate that the samples generated by our proposed flow exhibit higher quality compared to those obtained by several existing methods, especially for multimodal problems.

We also propose leveraging the Föllmer flow as a warmstart strategy for existing Markov Chain Monte Carlo (MCMC) methods, aiming to mitigate mode collapses and enhance their performance.

We can also leverage deep neural networks to fit the trajectory of sample evaluations. This allows us to obtain a generator for one-step sampling as a result.

Simple Tries

with Google Colab

demo.ipynb

with Docker

Install Docker first.

docker pull dingz489/fflow:demo-v0.1
docker run --rm -it -p 8888:8888 dingz489/fflow:demo-v0.1

Then enter localhost:8888 in your browser, open the demo.ipynb and run it to reproduce the demo picture!

Dependencies

pip install -r requirements

Assets

Our assets (Föllmer flow/MCMC samples(.npz), neural Föllmer flow checkpoints(.pth), tables(.csv) and figures(.pdf,.png)) have been uploaded to Google drive.

Usage

.
├── Makefile
├── README.md
├── assets
├── distribution.py # abstract base class for general distribution, we implemented n-D Gaussian mixtures here
├── example.py # correspond to examples in table 3 of our paper
├── fflow.py # prototype of Föllmer flow
├── main.py # entrance for different routines (generate samples, train/eval networks)
├── mcmc.py # MCMC samplers prototype, modified from https://github.com/Tom271/LangevinMC/blob/master/langevin_traj.py
├── metric.py # compute MMD and W2 distance
├── misc.py # miscellaneous helper utilities
├── network.py # ResNet prototype
├── reproduce.py # entrance file to reproduce tables and figures from paper
├── requirements.txt
└── run.sh # bash scripts, calling main.py

Setup the name of your python environment in run.sh.

We provide bash scripts for generating samples in run.sh. If you wish to know how to use the entrance file main.py, please refer to it for example.

./run.sh 1d # sampling from example 1 - 3, with Föllmer flow & MCMC
./run.sh 2d # sampling from example 4 - 10, with Föllmer flow & MCMC
./run.sh 2dmc # sampling from example 4 - 10, with MC Föllmer flow
./run.sh precondition # sampling from example 5, with different preconditioners
./run.sh hybrid # sampling from example 7, with predictor-corrector scheme
./run.sh nd # sampling from example 11 - 19, dimension ranges from 1 to 10
./run.sh train # train neural sampler for example 4 - 10
./run.sh eval # evaluate(sampling) neural sampler for example 4 - 10

Once you have generated samples or downloaded our provided assets, you can reproduce tables and figures:

python reproduce.py table1 # compute metrics for example 1 - 3
python reproduce.py table2 # compute metrics for example 4 - 10
python reproduce.py fig1 # evolution
python reproduce.py fig2 # plot for example 1 - 3
python reproduce.py fig3 # plot for example 4 - 10
python reproduce.py fig4 # plot influences of preconditioner
python reproduce.py fig5 # compare original and hybrid MCMC methods
python reproduce.py fig6 # plot n-D moments

How to extend to code

• New Gaussian Mixtures: append to the dictionary _EXAMPLES
• New Flows/Other Distributions: inherit the BaseDistribution, and implement the log of Radon-Nikodym derivative ln_scaled_ratio method to run the MC Föllmer flow, and implement the velocity_closed method to run the closed Föllmer flow. Call fflow_sampler and pass the velocity_fn to it.
• New MCMC samplers: append to the mcmc_sampler.methods

References

This repository is built upon some previous repositories

Schrödinger-Föllmer Sampler https://github.com/Liao-Xu/SFS_py

Langevin Monte Carlo Algorithms https://github.com/Tom271/LangevinMC/blob/master/langevin_traj.py

If you find the code useful for your research, please consider citing

@misc{ding2023sampling,
      title={Sampling via F\"ollmer Flow}, 
      author={Zhao Ding and Yuling Jiao and Xiliang Lu and Zhijian Yang and Cheng Yuan},
      year={2023},
      eprint={2311.03660},
      archivePrefix={arXiv},
      primaryClass={stat.ME}
}