Noise-based Regularized Training for Diffusion Models
Yuzhang Shang, Yubin Lu, Jinchao Feng, Ming Jun Zhong, Yan Yan · 2025
Denoising diffusion (score-based) generative models have recently achieved significant accomplishments in generating realistic and diverse high-dimensional data. These approaches define a forward diffusion process for transforming data into noise and a backward denoising process for sampling from noise via an approximated score function (i.e., score network). Analyzing the score function through the lens of conditional expectation, we notice that near the end of the denoising process, the score function has a blow-up (i.e., singularity), preventing the score network from approximating it properly. To address this issue, we introduce a regularity theorem for the score function and correspondingly propose our Training-Stabilized (TRAST) Diffusion Models. A key feature of our approach is noise manipulation, i.e. the addition of an unnoticeably small perturbation to the training images, particularly designed to prevent the score network from approximating the score function near its singularity. Subsequently, we incorporate a noise-filtering module at the end of the denoising process to refine the generated images. Training diffusion models can be regularized by early stopping and noise correction. Note that the designs of TRAST are theoretically supported by two theorems (i.e., singularity and regularity of the score function). Therefore, our method can stabilize diffusion models’ training by avoiding the singularity, and improve their performance in terms of fidelity and diversity.