Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model (2511.14716v1)

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.

• The paper introduces a unified framework that integrates the encoder, decoder, and diffusion backbone into a single ViT model, addressing latent collapse in joint training.
• It employs self-distillation techniques such as stop-gradient operations and loss transformation to reframe velocity prediction as a denoising task, ensuring stable latent space dynamics.
• Empirical results demonstrate that DSD achieves competitive FID scores while enhancing parameter and compute efficiency compared to traditional modular latent diffusion models.

Diffusion as Self-Distillation: A Unified Framework for End-to-End Latent Diffusion

Motivation and Background

Contemporary latent diffusion models (LDMs) have achieved substantial progress in high-fidelity image synthesis but are encumbered by architectural fragmentation: a separately pre-trained encoder (typically a VAE), a decoder, and a large diffusion model, each trained independently. This modular approach incurs several practical and theoretical drawbacks, including suboptimal joint optimization, resource overhead, and incompatibility with the unified single-backbone paradigms that now dominate unsupervised vision representation learning.

The central problem addressed by "Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model" (2511.14716), is whether it is possible to unify the encoder, decoder, and diffusion backbone in a single, end-to-end trainable architecture—realizing latent diffusion with the parameter and sampling efficiency of a monolithic Vision Transformer (ViT) backbone. A naïve attempt at such unification was empirically shown to catastrophically fail due to "latent collapse," motivating a systematic diagnosis and a new theoretical connection to self-distillation (SD) paradigms in unsupervised learning.

Diagnosing Latent Collapse: A Self-Distillation Lens

The authors draw a rigorous analogy between diffusion training objectives and self-distillation-based unsupervised learning. Classical SD methods—DINO, BYOL, I-JEPA—rely on a predictor matching online and exponential moving average (EMA) target encoders, with stability conferred by effective rank ("rank differentiation") mechanisms. Conversely, standard latent diffusion models use L2 or velocity-based loss on a VAE-pretrained latent space, with gradients backpropagated through only the diffusion backbone; any naïve attempt to backpropagate into the encoder triggers latent dimensional collapse.

Figure 1: Analogy between Self-Distillation (SD) based unsupervised learning and the standard diffusion model loss, comparing four architectural paradigms. Collapse mechanisms and the unified DSD are systematically identified.

Dimensional collapse is quantitatively diagnosed via effective rank measures. Joint VAE/diffusion training with a standard L2 loss rapidly suppresses latent space variance—the effective dimensionality drops to nearly unity, and the autoencoding task fails (reconstruction loss does not decrease).

This pathologic behavior is attributed to two mechanisms:

• Latent Variance Suppression: The L2-based diffusion objective introduces gradients that explicitly drive the variance of the latent towards zero, inducing collapse of the representation space.
• Failure of Rank Differentiation: The standard diffusion objective enforces high-rank noise prediction, violating the SD stability condition in which the predictor’s output rank must be no greater than (preferably less than) its targets.

  Figure 2: Trajectory of latent effective rank and reconstruction loss during training. Standard joint training causes swift collapse and failure to decrease reconstruction error; DSD maintains high rank and stable loss reduction.

The DSD Framework: Technical Innovations

Variance Decoupling

The first stabilization is achieved by applying a stop-gradient operation to the target latent in the diffusion loss. This cuts the feedback path penalizing latent variance, thereby protecting expressiveness:

$$\mathcal{L}_{\text{DSD}} = \|v(\mathbf{z}_t, t) - (\text{sg}(\mathbf{z}_2) - \boldsymbol{\epsilon})\|^2$$

The stop-gradient ensures that reconstruction and diffusion remain coupled, but the latent’s expressivity is no longer directly suppressed by the diffusion loss.

Loss Transformation to Denoising Objective

Even with decoupled gradients, standard diffusion loss still violates the rank differentiation mechanism due to full-rank additive noise in the target. By analytically transforming the loss, the task is reframed from velocity prediction to denoising (predicting the clean latent), restoring low-pass filtering and SD-style stability:

$$\mathcal{L}_{\text{z}} = w_t \|\tilde{v}(\mathbf{z}_t, t) - \mathbf{z}\|^2$$

where $w_t = (1-t)^{-2}$. The denoising predictor $\tilde{v}$ approximates the clean latent, functioning as a low-pass filter, satisfying the required rank order for stable SD training.

Figure 3: Geometric interpretation of the loss transformation. Predicting the velocity is mathematically equivalent to fitting the clean end-point latent, enabling denoising supervision.

EMA Target and Augmentation

Generalizing the SD paradigm, the framework further introduces:

• EMA Target Encoder ($E_2$): Provides temporal smoothing, yielding a non-collapsing, high-rank target latent.
• Aggressive Augmentation for Online Encoder ($E_1$): Color jitter, masking, and blurring expand the input’s support, encouraging broader invariance and stronger denoising properties.

  Figure 4: Unified DSD architecture: ViT backbone with task-specific heads, EMA-updated blocks, and three key objective heads (distillation, detached velocity, reconstruction). Gradient propagation and stop-gradient pathways are denoted.

Unified ViT Backbone and Head Design

The DSD architecture consists of a single ViT backbone with modular heads for:

• Self-distillation denoising
• Detached-velocity estimation (for sampling stability)
• Reconstruction (autoencoding)
• Auxiliary alignment and classification

All objectives are implemented in parallel, sharing representations to maximize feature reuse and parameter efficiency. Auxiliary objectives include inter-layer alignment (e.g., early ViT layers with DINOv3), representation-level SD, and standard classification loss (for supervised regimes).

Empirical Results and Analysis

The DSD framework is evaluated on conditional image generation (ImageNet $256\times 256$), with models at three scales (42M, 118M, 205M parameters). The primary metric is FID, with additional structural FID (sFID) and Inception Score (IS) computed on 50K samples. All training protocols strictly follow contemporary Vision Transformer standards (optimizer, batch size, data augmentation, EMA smoothing) for comparability.

DSD achieves:

• DSD-B (205M params, 50 epochs, no classifier-free guidance): FID = 4.25
• DSD-M (118M params): FID = 6.38
• DSD-S (42M params): FID = 13.44

Notably, DSD-B matches or outperforms previous LDM baselines that require up to $3\times$ as many parameters (e.g., SiT+REPA, LightningDiT with pretrained VAE, MDT). Unlike prior methods that separately pretrain tokenizers (often on larger datasets), DSD is fully end-to-end trained on ImageNet alone.

Ablation analysis confirms:

• Stability and non-collapsing training dynamics result only when both decoupling and loss transformation are employed.
• EMA targets and aggressive augmentation are synergistic, preventing subtle forms of collapse and maximizing effective rank.

  Figure 5: Qualitative DSD-B generative results on ImageNet $256\times 256$. Highly diverse, semantically consistent images at high resolution are achieved.

Theoretical and Practical Implications

The DSD framework corrects a fundamental mismatch between high-level generative modeling and unsupervised representation learning, establishing diffusion as a subclass of self-distillation. The architectural unification into a single ViT backbone bridges generative and discriminative pre-training regimes, enabling future "vision foundation models" to natively support both types of tasks without explicit modularization.

From a practical modeling perspective, DSD delivers:

• Maximally parameter- and compute-efficient generative pipelines
• Simplified deployment and maintenance in production settings
• Theoretical guarantees of latent space stability for scaling to larger or multimodal backbones

On the theoretical axis, the explicit diagnosis and resolution of latent collapse in generative modeling stimulates further research into the relationship between denoising-based and self-supervised objectives. The ability to manipulate latent rank via architectural and loss design invites the application of these techniques in text-to-image, video synthesis, and general multimodal scenarios.

Future Directions

DSD provides a powerful foundation for scaling unified latent generative models. Key directions for continued investigation include:

• Extending to larger ViT scales and more extensive long-horizon sampling tasks
• Adapting the DSD objective for text-to-image, video, or audio-visual diffusion models
• Theoretical refinement of the rank differentiation mechanism and empirical exploration of new auxiliary objectives
• Positioning DSD as a standard pretraining objective for foundation models, enabling direct transfer between generative and discriminative tasks

Conclusion

The Diffusion as Self-Distillation (DSD) framework introduces a unified, end-to-end trainable solution to latent diffusion, analytically bridging diffusion modeling and self-distillation via rank differentiation and loss transformation. By resolving the collapse phenomena and enabling parameter-efficient generative modeling within a single ViT backbone, DSD establishes a new regime for scalable, general-purpose generative vision modeling, with strong empirical performance and clear transfer potential for emerging multimodal and foundation model paradigms.