/Soft-Truncation

A repo for "Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation"

Primary LanguagePythonApache License 2.0Apache-2.0

Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation

This repo contains an official PyTorch implementation for the paper Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation.

We propose a genearlly applicable training method for a general weighted diffusion loss.

schematic

Running Commands

CIFAR-10

  • DDPM++ (VP, NLL) + ST
python main.py --config configs/vp/CIFAR10/ddpmpp_nll_st.py --workdir YOUR_SAVING_DIRECTORY --mode train
  • DDPM++ (VP, FID) + ST
python main.py --config configs/vp/CIFAR10/ddpmpp_fid_st_deepest.py --workdir YOUR_SAVING_DIRECTORY --mode train
  • UNCSN++ (RVE) + ST
python main.py --config configs/ve/CIFAR10/uncsnpp_st.py --workdir YOUR_SAVING_DIRECTORY --mode train

CelebA

  • DDPM++ (VP, NLL) + ST
python main.py --config configs/vp/CELEBA/ddpmpp_nll_st.py --workdir YOUR_SAVING_DIRECTORY --mode train
  • DDPM++ (VP, FID) + ST
python main.py --config configs/vp/CELEBA/ddpmpp_fid_st.py --workdir YOUR_SAVING_DIRECTORY --mode train
  • UNCSN++ (RVE) + ST
python main.py --config configs/ve/CELEBA/uncsnpp_st.py --workdir YOUR_SAVING_DIRECTORY --mode train

ImageNet32

  • DDPM++ (VP, NLL) + ST
python main.py --config configs/vp/IMAGENET32/ddpmpp_st.py --workdir YOUR_SAVING_DIRECTORY --mode train

CelebA-HQ

  • UNCSN++ (RVE) + ST
python main.py --config configs/ve/celebahq/uncsnpp_st.py --workdir YOUR_SAVING_DIRECTORY --mode train

Pretrained checkpoints

We release our checkpoints here.

References

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

@article{kim2021soft,
  title={Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation},
  author={Kim, Dongjun and Shin, Seungjae and Song, Kyungwoo and Kang, Wanmo and Moon, Il-Chul},
  journal={arXiv e-prints},
  pages={arXiv--2106},
  year={2021}
}

This work is heavily built upon the code from

  • Song, Yang, et al. "Score-Based Generative Modeling through Stochastic Differential Equations." International Conference on Learning Representations (ICLR). 2021.
  • Song, Yang, et al. "Maximum likelihood training of score-based diffusion models." Advances in Neural Information Processing Systems (NeurIPS). 2021.
  • Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in Neural Information Processing Systems (NeurIPS). 2020.