Diffusion models have emerged as the \emph{de-facto} generative model for image synthesis, yet they entail significant training overhead, hindering the technique's broader adoption in the research community. We observe that these models are commonly trained to learn all fine-grained visual information from scratch, thus motivating our investigation on its necessity. In this work, we show that it suffices to set up pre-training stage to initialize a diffusion model by encouraging it to learn some primer distribution of the unknown real image distribution. Then the pre-trained model can be fine-tuned for specific generation tasks efficiently. To approximate the primer distribution, our approach centers on masking a high proportion (e.g., up to 90\%) of an input image and employing masked denoising score matching to denoise visible areas. Utilizing the learned primer distribution in subsequent fine-tuning, we efficiently train a ViT-based diffusion model on CelebA-HQ $256 \times 256$ in the raw pixel space, achieving superior training acceleration compared to denoising diffusion probabilistic model (DDPM) counterpart and a new FID score record of 6.73 for ViT-based diffusion models. 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 on only 10\% data from a different dataset. Our code is available at \url{https://github.com/jiachenlei/maskdm}.
翻译:扩散模型已成为图像合成的\textit{事实标准}生成模型,但其训练开销巨大,阻碍了该技术在研究界的广泛采用。我们观察到,这些模型通常被训练来从头学习所有细粒度视觉信息,这促使我们研究其必要性。在本工作中,我们证明只需设立一个预训练阶段,通过鼓励扩散模型学习未知真实图像分布的一些先验分布来初始化模型,即可实现高效微调以完成特定生成任务。为逼近该先验分布,我们的方法核心是对输入图像进行高比例(例如高达90%)掩蔽,并采用掩蔽去噪分数匹配来去噪可见区域。在后续微调中利用所学先验分布,我们在原始像素空间中高效训练了基于ViT的扩散模型(CelebA-HQ $256 \times 256$),相比去噪扩散概率模型(DDPM)实现了显著的训练加速,并以6.73的FID分数创下了ViT扩散模型的新纪录。此外,我们的掩蔽预训练技术可普遍应用于各种在像素空间中直接生成图像的扩散模型,有助于学习具有出色泛化能力的预训练模型。例如,在VGGFace2上预训练的扩散模型仅通过10%的不同数据集数据进行微调,即可获得46%的质量提升。我们的代码已开源在\url{https://github.com/jiachenlei/maskdm}。