Home

Awesome

Accelerating Score-based Generative Models with Preconditioned Diffusion Sampling

Paper

Accelerating Score-based Generative Models with Preconditioned Diffusion Sampling,
Hengyuan Ma, Li Zhang, Xiatian Zhu, and Jianfeng Feng
ECCV 2022

News

Abstract

Score-based generative models (SGMs) have recently emerged as a promising class of generative models. However, a fundamental limitation is that their inference is very slow due to a need for many (e.g., 2000) iterations of sequential computations. An intuitive acceleration method is to reduce the sampling iterations which however causes severe performance degradation. We investigate this problem by viewing the diffusion sampling process as a Metropolis adjusted Langevin algorithm, which helps reveal the underlying cause to be ill-conditioned curvature. Under this insight, we propose a model-agnostic preconditioned diffusion sampling (PDS) method that leverages matrix preconditioning to alleviate the aforementioned problem. Crucially, PDS is proven theoretically to converge to the original target distribution of a SGM, no need for retraining. Extensive experiments on three image datasets with a variety of resolutions and diversity validate that PDS consistently accelerates off-the-shelf SGMs whilst maintaining the synthesis quality. In particular, PDS can accelerate by up to 29x on more challenging high resolution (1024x1024) image generation.

Demo

demo_vanilla FFHQ facial images (1024x1024) generated by vanilla NCSN++ using 2000, 200, 133, 100 and 66 sampling iterations (from left to right).

demo_ours FFHQ facial images (1024x1024) generated by NCSN++ w/ PDS using 2000, 200, 133, 100 and 66 sampling iterations (from left to right).

Preconditioned diffusion sampling

<img src="src/demo.png" height = "200" align=center />

The Langevin dynamics applied by vanilla SGMs for generating samples is

<img src="src/Langevin_dynamics.png" height = "50" align=center />

To accelerate the convergence while keeping the steady state distribution, our PDS use the following preconditioned Langevin dynamics

<img src="src/preconditioned_Langevin_dynamics.png" height = "50" align=center /> where the preconditioning matrix M can be constructed by priori knowledge from the target dataset in frequency domain or space (pixel) domain, S can be any skew-symmetric matrix.

Running the Experiments

Sampling using PDS through main.py.

$ python main.py --config: Model configuration. (default: 'None')
                 --workdir: Working directory. (default: 'None')
                 --eval_folder: Output directory. (default: 'eval')
                 --speed_up: Speedup ratio. 
                 --freq_mask_path: Path to the frequency filter. (default: 'None')
                 --space_mask_path: Path to the space filter. (default: 'None') 
                 --alpha: The setting of the parameter alpha. (default: 5)

Example: sampling FFHQ (1024x1024) dataset with both frequency and space filter, under 30 times acceleration.

$ python main.py --config ./configs/ve/ffhq_ncsnpp_continuous.py --workdir  yourworkdir --speed_up 30 --freq_mask_path ./ffhq_freq.npy --space_mask_path ./ffhq_space.npy

Example: sampling LSUN bedroom (256x256) dataset with only frequency filter, under 12 times acceleration.

$ python main.py --config ./configs/ve/bedroom_ncsnpp_continuous.py --workdir  yourworkdir --speed_up 12 --freq_mask_path ./bedroom_freq.npy

License

MIT

Reference

@inproceedings{ma2022pds,
  title={Accelerating Score-based Generative Models with Preconditioned Diffusion Sampling},
  author={Ma, Hengyuan and Zhang, Li and Zhu, Xiatian and Feng, Jianfeng},
  booktitle={European Conference on Computer Vision},
  year={2022}
}

Acknowledgement

Thanks to previous open-sourced repo: NCSN++