Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model

TL;DR

This paper introduces Diffusion as Self-Distillation (DSD), a novel framework that unifies the encoder, decoder, and diffusion network into a single end-to-end trainable model, overcoming stability issues and achieving high-quality image generation.

Contribution

The paper proposes DSD, a new training framework that stabilizes end-to-end latent diffusion training by drawing an analogy with self-distillation, enabling unified models.

Findings

Stable end-to-end training of a single network for encoding, decoding, and diffusion.

Achieved state-of-the-art conditional image generation on ImageNet 256x256.

Reduced model complexity with fewer parameters and training epochs.

Abstract

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.