Source-linked AI summary
Denoising Diffusion Autoencoders are Unified Self-supervised Learners
Weilai Xiang, Hongyu Yang, Di Huang, Yunhong Wang
TL;DR
The paper asks whether diffusion models can bridge image generation and recognition, since existing generative self-supervised learners have not done so effectively. It evaluates denoising diffusion autoencoders and finds they learn competitive discriminative representations alongside generative ability, supporting unified vision foundation models.
Problem
Existing generative self-supervised vision models have not effectively bridged image generation and recognition, motivating whether diffusion models can become unified learners.
Method
The paper uses end-to-end diffusion pre-training and extracts intermediate denoising features for linear evaluation and fine-tuning without auxiliary encoders.
Results
DDAEs achieve competitive recognition performance alongside generation, including 66.3% Tiny-ImageNet linear accuracy after ImageNet transfer with a DiT-XL/2 backbone.
Takeaways & Limitations
Diffusion pre-training can serve as a unified approach for acquiring generative ability and visual understanding, potentially supporting unified vision foundation models.
Takeaways & Limitations
Latent-space diffusion models perform worse than pixel-space models for recognition, and latent compression may hinder downstream tasks such as object detection.
Abstract
from arXiv · showhide
Inspired by recent advances in diffusion models, which are reminiscent of denoising autoencoders, we investigate whether they can acquire discriminative representations for classification via generative pre-training. This paper shows that the networks in diffusion models, namely denoising diffusion autoencoders (DDAE), are unified self-supervised learners: by pre-training on unconditional image generation, DDAE has already learned strongly linear-separable representations within its intermediate layers without auxiliary encoders, thus making diffusion pre-training emerge as a general approach for generative-and-discriminative dual learning. To validate this, we conduct linear probe and fine-tuning evaluations. Our diffusion-based approach achieves 95.9% and 50.0% linear evaluation accuracies on CIFAR-10 and Tiny-ImageNet, respectively, and is comparable to contrastive learning and masked autoencoders for the first time. Transfer learning from ImageNet also confirms the suitability of DDAE for Vision Transformers, suggesting the potential to scale DDAEs as unified foundation models. Code is available at github.com/FutureXiang/ddae.
1. Introduction
The paper asks whether diffusion models can bridge generative and discriminative learning in vision. It proposes evaluating diffusion-pretrained networks as feature extractors and reports strong intermediate representations and competitive recognition performance.
- Limited human supervision motivates self-supervised pre-training, but vision has not matched language models in unifying generation and recognition.
- Diffusion models are proposed as unified learners because generative pre-training can support understanding, while denoising requires recovering images across noise levels.
- DDAE evaluates intermediate diffusion activations using linear probing on noised images and fine-tuning after truncating the network into an image encoder.
- DDAEs learn strongly linear-separable features during end-to-end diffusion pre-training, with representations located in intermediate up-sampling layers.
- DDAEs are comparable to supervised WideResNet, contrastive SimCLR, and masked autoencoder methods on CIFAR-10 and Tiny-ImageNet.
- The study presents diffusion models as unified vision foundation models combining generative capability with competitive recognition representations.
2. Related work
Prior generative representation learners include VAEs, GANs, autoregressive Transformers, and diffusion-based methods, but each has limitations for recognition or transfer. Diffusion offers a promising route because it models complete image distributions while supporting representation learning.
- Diffusion models became a leading generative paradigm through high-fidelity synthesis of complex visual concepts without adversarial instability or mode collapse.
- VAEs, BigBiGAN, and iGPT demonstrate generative representation learning, but face limitations involving recognition utility, data diversity, or image-specific inductive bias.
- Diffusion representation methods have used auxiliary encoders or modified frameworks, while classification-objective approaches have not matched pure recognition methods.
3. Approach
DDAEs treat diffusion as multi-level denoising autoencoding and expose intermediate activations as classification features. The approach evaluates layer and noise choices, then uses alignment and uniformity metrics to monitor feature quality without labels.
- 3.1. Background: DDAEs as generative models: Diffusion models corrupt images with Gaussian noise across timesteps and train networks to invert the corruption and reconstruct clean samples.The corruption uses timestep-dependent signal and noise scales, while sampling iteratively reverses the process.
- 3.1. Background: DDAEs as generative models: Diffusion networks are equivalent to multi-level, level-conditional denoising autoencoders, termed Denoising Diffusion Autoencoders.The denoiser is trained with a denoising autoencoder objective derived from the diffusion likelihood objective.
- 3.1. Background: DDAEs as generative models: The study evaluates DDPM UNets, DDPM++ trained by EDM, and latent-space DiT as representative DDAE implementations.DDPM and DDPM++ use UNets, whereas DiT uses latent-space Vision Transformers with timestep and label embeddings.
- 3.2. Evaluating DDAEs as discriminative learners: Intermediate DDAE activations are evaluated with linear probes and fine-tuning, using selected layers and noising scales as feature-extraction choices.For fine-tuning, clean images are passed through encoders truncated at selected layers, while linear probing examines pooled activations from noised images.
- 3.2. Evaluating DDAEs as discriminative learners: 95%+ linear probe accuracy is achieved by DDPM++ trained by EDM when relatively small noises are applied and features are taken from the middle of up-sampling.Layer depth and noising scale jointly affect feature quality, and the best setting is found empirically.
- 3.3. Label-free monitoring for layer selection: Alignment and uniformity metrics agree with linear-probe quality and improve consistently during training for well-performing layers.The metrics are computed on normalized features with independently sampled noises and may support label-free layer selection.
4. Experiments
Experiments test whether diffusion pre-training yields discriminative features alongside generation, using ablations, classification evaluations, and comparisons across datasets and methods. Results show a positive generative–discriminative relationship, strong DDAE recognition performance, and practical caveats for truncated or latent-space backbones.
- Experimental setup: Diffusion pre-training was evaluated through ablations and linear probing, fine-tuning, and ImageNet-transfer comparisons on CIFAR-10 and Tiny-ImageNet.The experiments also compare generative and discriminative performance across diffusion configurations and prior methods.
- Denoising autoencoding: Better generative models generally learn better representations before recognition overfits, confirming a positive generative–discriminative correlation.Recognition can overfit earlier than generation, while generation continues improving.
- Limitations: Truncated diffusion UNets trained from scratch failed to benefit from scaling, while latent-space diffusion models showed lower recognition accuracy than pixel-space models in preliminary tests.The paper identifies general-purpose vision backbones and unified pixel-space backbones as future directions.
- Noise-conditional classifier: DDAE-based noise-conditional classifiers surpassed supervised models over almost all noise scales while producing denoising outputs and classification logits in one forward pass.A two-layer MLP head was trained on frozen DDAE features.
- Generative pre-training: 95.9% CIFAR-10 linear-probe accuracy and 50.0% Tiny-ImageNet linear-probe accuracy were achieved by EDM-based DDAE.Fine-tuning reached 97.2% on CIFAR-10 and 69.4% on Tiny-ImageNet.
- Comparison with previous methods: EDM-based DDAE was comparable to contrastive methods and, after fine-tuning, surpassed comparable SimCLR models while remaining below the scaled 375M SimCLR model by 0.5% in one setting.On Tiny-ImageNet, DDAE outperformed SimCLR ResNet-18 but was less parameter-efficient than a slightly larger ResNet-50.
5. Discussion and conclusion
The paper proposes diffusion pre-training as a unified approach for generation and visual understanding, while identifying backbone, efficiency, and scalability limitations.
- Diffusion pre-training is proposed to acquire both generation ability and deep visual understanding for potential unified vision foundation models.
- Limitations: The approach relies on truncating DDAEs and probing layers to find suitable representations, which may fail on more complex datasets.
- 50.0% linear and 69.4% fine-tuning accuracy are reported for DDAE (EDM) on Tiny-ImageNet.
- 66.3% linear and 77.8% fine-tuning accuracy are reported for DDAE (DiT-XL/2) with ImageNet transfer on Tiny-ImageNet.
- Experimental scope: The comparison of DDAE (DiT-XL/2) includes a class-conditional model evaluated in an unconditional manner, and uses an extra VAE encoder.
- Limitations: DDAEs require larger model sizes and longer training durations than some pure recognition models, making them costly to scale.
A. Implementation details
The implementation follows official diffusion-model implementations and uses simple probing, fine-tuning, and supervised-from-scratch settings.
- DDPM and EDM use WideResNet-based UNets, while DiT-XL/2 is an ImageNet-pretrained Transformer retrieved from its official implementation.
- The experiments use simple linear-probing and fine-tuning settings without tuning Adam β1/β2, weight decay, or adding normalization layers.
- Latent-space DiT recognition training stores VAE-encoded augmented data and randomly samples one encoded version per epoch.
- Supervised truncated-UNet models are trained from scratch for 200 epochs to reach maximum performance for comparison.
B. Layer-noise combinations in grid search
The study searches layer–noise combinations because representation quality depends heavily on both, using separate schedules for different diffusion models.
- Layer–noise combinations are grid-searched for each model and dataset because they heavily affect representation quality.
- 18-step and 50-step EDM models train linear classifiers for 10 epochs at each tested layer and timestep.
- For 1000-step DDPM and DiT models, timestep increments of 5 or 10 make the search more efficient.
- Table 9 reports the adopted layer–noise combinations, with numbers after “@” denoting image or feature-map resolutions.
- Tables 5–8 document diffusion network specifications and linear-probing, fine-tuning, and supervised-training settings.