Ex2MCMC: Local-Global MCMC kernels: the best of both worlds (NeurIPS 2022) [Paper]

Implementation of Ex2MCMC, FlEx2MCMC and experiments

[ArXiv]

Authors: Sergey Samsonov, Evgeny Lagutin, Marylou Gabrié, Alain Durmus, Alexey Naumov, Eric Moulines.

Abstract: In the present paper we study an Explore-Exploit Markov chain Monte Carlo strategy (Ex2MCMC) that combines local and global samplers showing that it enjoys the advantages of both approaches. We prove V-uniform geometric ergodicity of Ex2MCMC without requiring a uniform adaptation of the global sampler to the target distribution. We also compute explicit bounds on the mixing rate of the Explore-Exploit strategy under realistic conditions. Moreover, we also analyze an adaptive version of the strategy (FlEx2MCMC) where a normalizing flow is trained while sampling to serve as a proposal for global moves. We illustrate the efficiency of Ex2MCMC and its adaptive version on classical sampling benchmarks as well as in sampling high-dimensional distributions defined by Generative Adversarial Networks seen as Energy Based Models.

• Ex2MCMC: Local-Global MCMC kernels: the best of both worlds (NeurIPS 2022) [Paper]
  • Implementation of Ex2MCMC, FlEx2MCMC and experiments
    • Sampling from GAN as Energy-Based Model with MCMC
    • Algorithms
    • Quick start
    • Requirements
    • Installation
    • Checkpoints
    • Usage
      • Experiments with synthetic distributions:
      • Experiments with GANs on MNIST dataset
      • Experiments with GANs on CIFAR10 dataset
      • Sampling and FID computation
    • Results
      • FID and Inception Score (CIFAR10)
      • Energy landscape approximation (MNIST)
      • Sampling trajectories (CIFAR10)
    • Citation

Sampling from GAN as Energy-Based Model with MCMC

Metrics:

Samples from SNGAN with FlEx2MCMC:

Algorithms

Ex²MCMC:

FlEx²MCMC:

Quick start

pip install ex2mcmc
pip install git+https://github.com/kwotsin/mimicry.git

Requirements

• Python >= 3.8
• PyTorch >= 1.8.0
• torchvision
• pyro-ppl
• Jax
• POT

Installation

Create environment:

conda create -n ex2mcmc python=3.8
conda activate ex2mcmc

Install poetry (if absent):

curl -sSL https://install.python-poetry.org | python3 -
poetry config virtualenvs.create false

Install the project:

poetry install --with dev
poetry add git+https://github.com/kwotsin/mimicry.git@a7fda06c4aff1e6af8dc4c4a35ed6636e434c766

Checkpoints

CIFAR10 checkpoints:

GAN	Steps	Path, G	Path, D
DCGAN NS	100k	netG_100000_steps.pth	netD_100000_steps.pth
SNGAN, Hinge	100k	netG.pth	netD.pth

MNIST checkpoints:

GAN	Path
Vanilla	vanilla_gan.pth
WGAN CP	wgan.pth

You also can run script to download checkpoints:

chmod +x get_ckpts.sh
./get_ckpts.sh

Download statistics for FID cimputation for CIFAR10 dataset:

mkdir -p stats & gdown 1jjgB_iuvmoVAXPRvVTI_hBfuIz7mQgOg -O stats/fid_stats_cifar10.npz

Usage

Demonstration on SNGAN:

Experiments with synthetic distributions:

FlEx²MCMC vs NUTS:

Experiment	Path	Colab
Toyish Gaussian	experiments/exp_synthetic/toyish_gaussian.ipynb	TBD
Ex2MCMC for Mixture of Gaussians	experiments/exp_synthetic/ex2mcmc_mog.ipynb
FlEx2MCMC for Mixture of Gaussians	experiments/exp_synthetic/flex_mog.ipynb
FlEx2MCMC for banana-shaped distribution	experiments/exp_synthetic/flex_banana.ipynb
FlEx2MCMC for Neal's funnel distribution	experiments/exp_synthetic/flex_funnel.ipynb

To reproduce the experimets on banana-shaped and funnel distributions:

python experiments/exp_synthetic/banana_funnel_metrics.py --distribution {banana,funnel} --device cuda:0

Experiments with GANs on MNIST dataset

experiments/exp_mnist/JSGAN_samples.ipynb TBD

experiments/exp_mnist/WGAN_samples.ipynb

Experiments with GANs on CIFAR10 dataset

experiments/exp_cifar10_demo/DCGAN_samples.ipynb

experiments/exp_cifar10_demo/SNGAN_samples.ipynb

Sampling and FID computation

python experiments/exp_cifar10_fid/run.py configs/mcmc_configs/{ula,mala,isir,ex2mcmc,flex2mcmc}.yml configs/mmc_dcgan.yml

To run a full experiment:

chmod +x experiments/exp_cifar10_fid/run.sh & ./experiments/exp_cifar10_fid/run.sh

Results

FID and Inception Score (CIFAR10)

GAN	MCMC	Steps	Inception Score	FID
DCGAN	none	0	6.3	28.4
DCGAN	i-SIR	1k	6.96	22.7
DCGAN	MALA	1k	6.95	23.4
DCGAN	Ex2MCMC (our)	1k	7.56	19.0
DCGAN	FlEx2MCMC (our)	1k	7.92	19.2
DCGAN	FlEx2MCMC (our)	180	7.62	17.1

Energy landscape approximation (MNIST)

Projection of GAN samples onto the energy landsape when trained on MNIST dataset:

Sampling trajectories (CIFAR10)

Generation trajectories for DCGAN.

• ULA:

• MALA:

• i-SIR:

• Ex²MCMC:

• FlEx²MCMC:

Citation

@article{samsonov2022local,
  title={Local-Global MCMC kernels: the best of both worlds},
  author={Samsonov, Sergey and Lagutin, Evgeny and Gabri{\'e}, Marylou and Durmus, Alain and Naumov, Alexey and Moulines, Eric},
  journal={Advances in Neural Information Processing Systems},
  volume={35},
  pages={5178--5193},
  year={2022}
}