Diffusion models have emerged as the \emph{de-facto} technique for image generation, yet they entail significant computational overhead, hindering the technique's broader application in the research community. We propose a prior-based denoising training framework, the first to incorporate the pre-train and fine-tune paradigm into the diffusion model training process, which substantially improves training efficiency and shows potential in facilitating various downstream tasks. Our approach centers on masking a high proportion (e.g., up to 90\%) of the input image and employing masked denoising score matching to denoise the visible areas, thereby guiding the diffusion model to learn more salient features from training data as prior knowledge. By utilizing masked learning in a pre-training stage, we efficiently train the ViT-based diffusion model on CelebA-HQ $256 \times 256$ in the pixel space, achieving a 4x acceleration and enhancing the quality of generated images compared to denoising diffusion probabilistic model (DDPM). Moreover, our masked pre-training technique can be universally applied to various diffusion models that directly generate images in the pixel space, aiding in the learning of pre-trained models with superior generalizability. For instance, a diffusion model pre-trained on VGGFace2 attains a 46\% quality improvement through fine-tuning with merely 10\% data from a different distribution. Moreover, our method shows the potential to serve as a training paradigm for enhancing the privacy protection capabilities of diffusion models. Our code is available at \url{https://github.com/jiachenlei/maskdm}.
翻译:扩散模型已成为图像生成的主流技术,但其巨大的计算开销阻碍了该技术在研究社区的广泛应用。我们提出了一种基于先验知识的去噪训练框架,首次将预训练与微调范式融入扩散模型训练过程,大幅提升了训练效率并展现出促进各类下游任务的潜力。该方法的核心是对输入图像的高比例区域(例如高达90%)进行掩码,并采用掩码去噪分数匹配技术对可见区域进行去噪,从而引导扩散模型从训练数据中学习更显著的先验特征。通过在预训练阶段利用掩码学习,我们在像素空间中对CelebA-HQ $256 \times 256$数据集上基于ViT的扩散模型进行高效训练,相比去噪扩散概率模型(DDPM)实现了4倍加速,并提升了生成图像的质量。此外,我们的掩码预训练技术可普适应用于直接在像素空间生成图像的各种扩散模型,有助于学习具有优越泛化能力的预训练模型。例如,在VGGFace2上预训练的扩散模型仅通过利用不同分布的10%数据进行微调,即可获得46%的质量提升。同时,该方法展现出作为增强扩散模型隐私保护能力的训练范式的潜力。我们的代码已开源至 \url{https://github.com/jiachenlei/maskdm}。