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

[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 bost of both worlds (NeurIPS 2022) [Paper]
  • Single chain mixing
  • Sampling from GAN as Energy-Based Models with MCMC
  • Algorithms
  • Installation
    • Developement installation
    • Download checkpoints and stats
  • Usage
    • Demonstration on SNGAN
    • 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)
    • Sampling trajectories (CIFAR10)
    • Energy landscape approximation (MNIST)
  • Citation

Single chain mixing

Sampling from GAN as Energy-Based Models with MCMC

Algorithms

Ex²MCMC:

FlEx²MCMC:

Installation

Create environment:

conda create -n ex2mcmc python=3.8
conda activate ex2mcmc

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

Developement installation

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

Download checkpoints and stats

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

Try with colab:

Experiments with synthetic distributions:

FlEx²MCMC vs NUTS:

Experiment	Path	Colab
Toyish Gaussian	experiments/exp_synthetic/toyish_gaussian.ipynb	TBD
Gaussian mixture	experiments/exp_synthetic/gaussian_mixture.ipynb
FlEx for Mixture of Gaussians	experiments/exp_synthetic/flex_mog.ipynb
FlEx for banana-shaped distribution	experiments/exp_synthetic/flex_banana.ipynb
FlEx 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

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

Sampling trajectories (CIFAR10)

Generation trajectories for DCGAN.

• ULA:

• MALA:

• i-SIR:

• Ex²MCMC:

• FlEx²MCMC:

Energy landscape approximation (MNIST)

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

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}
}