Jiachen Lei Keli Liu Julius Berner Haiming Yu Hongkai Zheng Jiahong Wu Xiangxiang Chu

摘要

像素空间生成模型通常比其隐空间对应模型更难训练，且整体性能也较差，导致性能与效率之间长期存在差距。本文提出一种新颖的两阶段训练框架，成功弥合了像素空间扩散模型与一致性模型在这一方面的差距。在第一阶段，我们预训练编码器，使其从干净图像中捕捉有意义的语义信息，同时将这些编码器与同一确定性采样轨迹上的点对齐——该轨迹从先验分布逐步演化至数据分布。在第二阶段，我们将预训练的编码器与随机初始化的解码器相结合，并对完整的模型进行端到端微调，以同时适用于扩散模型和一致性模型。该训练框架在ImageNet数据集上展现出强劲的实证性能。具体而言，我们的扩散模型在ImageNet-256上达到2.04的FID分数，在ImageNet-512上达到2.35的FID分数（使用75次函数评估次数，NFE），在生成质量和效率两方面均显著超越此前所有像素空间方法，且在相近训练成本下媲美当前领先的基于VAE的模型。此外，在ImageNet-256上，我们的一致性模型仅需单步采样即实现8.82的FID分数，显著优于其隐空间对应模型。据我们所知，这是首个在不依赖预训练VAE或扩散模型的前提下，成功直接在高分辨率图像上训练一致性模型的实例。