Diffusion models have emerged as the 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 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 this masked learning process in a pre-training stage, we efficiently train the ViT-based diffusion model on CelebA-HQ 256x256 in the pixel space, achieving a 4x acceleration and enhancing the quality of generated images compared to DDPM. Moreover, our masked pre-training technique is universally applicable to various diffusion models that directly generate images in the pixel space and facilitates learning pre-trained models with excellent generalizability: a diffusion model pre-trained on VGGFace2 attains a 46% quality improvement through fine-tuning with merely 10% local data. Our code is available at https://github.com/jiachenlei/maskdm.
翻译:扩散模型已成为图像生成领域的事实标准技术,但其显著的计算开销阻碍了该技术在研究社区的更广泛应用。我们提出了一种基于先验的去噪训练框架,首次将预训练与微调范式融入扩散模型训练过程,大幅提升了训练效率,并展现出促进各类下游任务的潜力。我们的方法核心在于对输入图像高比例掩膜(例如高达90%),并采用掩膜得分匹配对可见区域进行去噪,从而引导扩散模型从训练数据中学习更显著特征作为先验知识。通过在预训练阶段利用这一掩膜学习过程,我们在像素空间中对基于ViT的扩散模型在CelebA-HQ 256×256数据集上进行高效训练,与DDPM相比实现了4倍加速并提升了生成图像质量。此外,我们的掩膜预训练技术普遍适用于各类在像素空间直接生成图像的扩散模型,并能训练出具有优异泛化能力的预训练模型:在VGGFace2上预训练的扩散模型仅通过10%本地数据微调就实现了46%的质量提升。我们的代码已开源在https://github.com/jiachenlei/maskdm。