用自蒸馏思想统一扩散模型的编码解码与扩散过程,实现端到端训练。
Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model
- 将扩散过程视为自蒸馏,设计新训练目标稳定潜在空间。
- 仅用4200万参数,50轮训练即达FID=13.44,无需分类器引导。
- 适合追求轻量高效、一体化生成架构的研究者。
标准潜空间扩散模型采用编码器、解码器和扩散网络三部分分离的架构,需分阶段训练,计算效率低且性能受限,难以与视觉基础模型的单网络结构融合。本文旨在将这三个组件统一为一个可端到端训练的单一网络。我们发现,直接联合训练会因“潜在空间坍塌”而彻底失败,扩散目标干扰了潜在表示的学习。通过将扩散与自蒸馏方法类比,揭示其根本原因,并提出扩散作为自蒸馏(DSD)的新框架,通过改进训练目标稳定潜在空间。该方法首次实现单一网络在端到端训练中同时完成编码、解码与扩散。在ImageNet $256 imes 256$ 条件生成任务上,仅用4200万/11800万/20500万参数和50轮训练,达到FID=13.44/6.38/4.25,且未使用分类器自由引导。
原文摘要 · Abstract (English)
Standard Latent Diffusion Models rely on a complex, three-part architecture consisting of a separate encoder, decoder, and diffusion network, which are trained in multiple stages. This modular design is computationally inefficient, leads to suboptimal performance, and prevents the unification of diffusion with the single-network architectures common in vision foundation models. Our goal is to unify these three components into a single, end-to-end trainable network. We first demonstrate that a naive joint training approach fails catastrophically due to ``latent collapse'', where the diffusion training objective interferes with the network's ability to learn a good latent representation. We identify the root causes of this instability by drawing a novel analogy between diffusion and self-distillation based unsupervised learning method. Based on this insight, we propose Diffusion as Self-Distillation (DSD), a new framework with key modifications to the training objective that stabilize the latent space. This approach enables, for the first time, the stable end-to-end training of a single network that simultaneously learns to encode, decode, and perform diffusion. DSD achieves outstanding performance on the ImageNet $256\times 256$ conditional generation task: FID=13.44/6.38/4.25 with only 42M/118M/205M parameters and 50 training epochs on ImageNet, without using classifier-free-guidance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。