High Precision Score-based Diffusion Models

Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, Il-chul Moon

Keywords: Diffusion Model, Score-based Model

Abstract: Recent advances in diffusion models bring the state-of-the art performance on image generation tasks. However, the image generation is still an arduous task in high resolution, both theoretically and practically. From the theory side, the difficulty arises in estimating the high precision diffusion because the data score goes to $\infty$ as $t \rightarrow 0$ of the diffusion time. This paper resolves this difficulty by improving the previous diffusion models from three aspects. First, we propose an alternative parameterization for such unbounded data score, which theoretically enables the unbounded score estimation. Second, we provide a practical soft truncation method (ST-trick) to handle the extreme variation of the score scales. Third, we design a reciprocal variance exploding stochastic differential equation (RVESDE) to enable the sampling at the high precision of $t$. These three improvements are applicable to the variations of both NCSN and DDPM, and our improved versions are named as HNCSN and HDDPM, respectively. The experiments show that the improvements result in the state-of-the-art performances in the high resolution image generation, i.e. CelebA-HQ. Also, our ablation study empirically illustrates that all of alternative parameterization, ST-trick, and RVESDE contributes to the performance enhancement.

One-sentence Summary: We improve the diffusion models capable of estimating the unbounded data score, and provide practical training method for the stable diffusion training.