PaperlayerView exact revision on arXiv

Source-linked AI summary

Denoising Diffusion Autoencoders are Unified Self-supervised Learners

Weilai Xiang, Hongyu Yang, Di Huang, Yunhong Wang

arXiv:2303.09769v2cs.CVcs.LG

TL;DR

Whether diffusion models can unify generative and discriminative vision learning remained open. This paper evaluates DDAE representations learned through diffusion pre-training and finds them competitive with established recognition approaches across CIFAR-10 and Tiny-ImageNet.

  • Problem

    Whether diffusion models can unify generative and discriminative vision learning remained open, as image synthesis and recognition relied on approaches not considered unified models.

    Source: §1 · pp. 1–2

  • Method

    The paper uses end-to-end diffusion pre-training and evaluates noise-perturbed intermediate DDAE features for recognition.

    Source: §1 · p. 2

  • Results

    Across CIFAR-10 and Tiny-ImageNet, DDAE is comparable to supervised WideResNet, SimCLR, and MAE.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings position DDAEs as unified vision foundation-model candidates with generative capability and competitive recognition performance.

    Source: §1 · p. 2

    Truncated diffusion UNets with global pooling are not optimal classification backbones and may have limited DDAE performance.

    Source: §4.1.2 · p. 7

Abstract

from arXiv · show

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 introduction asks whether diffusion pre-training can unite image generation and recognition, presenting DDAEs as a way to learn discriminative features without auxiliary encoders. It motivates this approach by the gap between generative and recognition models in vision and reports evidence that diffusion representations support classification.

Introduction

  • In NLP, generative pre-training by GPTs and T5 supports both generation and discriminative tasks, motivating the search for a comparable bridge between image generation and recognition.BERT's masked language modeling is presented as an encoder pretext task that cannot generate full text samples.Source: §1 · p. 1
  • Unlike generative models such as GANs and autoregressive Transformers, which synthesize images without significant recognition benefits, contrastive learning and masked autoencoders target features but are not unified full-generation models.The introduction also notes that masked image modeling is less natural for convolutional networks and does not directly model the complete data distribution.Source: §1 · pp. 1–2
  • Diffusion is proposed as a promising visual learner because full generative pre-training may support understanding, multi-level denoising demands visual concepts, and prior work links generation quality with feature quality.DDPM-seg [4] is cited as evidence that diffusion can capture pixel-wise semantic information.Source: §1 · p. 2
  • The approach evaluates intermediate DDAE activations with linear probes, then truncates the network at its best representation layer and fine-tunes it without additional noising.Linear probes use images perturbed at specific noise scales or timesteps.Source: §1 · p. 2
  • DDAEs learn strongly linearly separable features in intermediate up-sampling layers from noised images, with classification results comparable to WideResNet, SimCLRs, and MAE and transfer verified from ImageNet pre-trained models.The transfer evaluation includes pixel-space UNets and latent-space Vision Transformers such as DiT.Source: §1 · p. 2
  • The paper frames DDAEs as potential unified vision foundation models whose generative and competitive recognition capabilities could inform pre-training and future transfer of discriminative knowledge from large-scale AIGC models.Stable Diffusion [48] is given as a possible source of such transferable knowledge.Source: §1 · p. 2

2. Related work

Diffusion models have become a leading generative paradigm, while prior representation-learning approaches using diffusion have faced recognition limitations. This work studies diffusion models for classification on complex multi-class datasets without modifying the diffusion framework.

Diffusion models

  • Diffusion models offer high-fidelity generation of complex visual concepts [26] [53] [29] [48] [46] [50], and improvements have established state-of-the-art results across unconditional CIFAR-10, class-conditional ImageNet, and MS-COCO text-to-image.The related work also notes diffusion applications including Stable Diffusion [48] and ControlNet [65], and scalable Vision Transformer backbones [16] [2] [41].Source: §2 · pp. 2–3

Representation learning with generative models

  • Earlier generative representation learners include VAEs [32] [54], BigBiGAN [15], and iGPT, but recognition is limited by generation focus, reduced GAN data diversity, or weak image inductive bias.iGPT nevertheless achieves results competitive with contrastive methods, while its downstream applications are described as restricted and inefficient.Source: §2 · p. 3

Representation learning with diffusion models

  • Prior diffusion methods use auxiliary encoders for attribute manipulation, modified frameworks that trail contrastive baselines, classification objectives that hurt generation, or likelihood estimates below ResNet on CIFAR-10.The paper contrasts these approaches with its framework-preserving method, which it describes as comparable to typical self-supervised and supervised models.Source: §2 · p. 3
  • DDPM-seg studies DDPM [26] on single-super-class datasets for segmentation; this paper says it is the first to explore various diffusion models for classification on complex multi-class datasets.The proposed approach is described as comparable to typical self-supervised and supervised models without modifying diffusion frameworks.Source: §2 · p. 3

3. Approach

The approach treats diffusion denoisers as multi-level, level-conditional denoising autoencoders and evaluates their intermediate features for discrimination. It also explores alignment and uniformity metrics for monitoring layer quality without labels.

3.1. Background: DDAEs as generative models

  • Diffusion models corrupt inputs with timestep-dependent Gaussian noise and iteratively reverse the corruption; their trained denoisers act as multi-level, level-conditional denoising autoencoders.The reverse process starts from Gaussian noise and uses learned transitions to reconstruct samples.Source: §3.1 · p. 3

3.2. Evaluating DDAEs as discriminative learners

  • The method evaluates pooled intermediate activations directly, searching layer-depth and noise-scale combinations because the best feature layer is unknown and deterministic encodings are not compact enough for classification.For fine-tuning, clean images are passed to encoders truncated at the selected layers, with timestep inputs fixed to the selected values.Source: §3.2 · pp. 3–4
  • 95%+ linear probe accuracy is achieved by DDPM++ trained by EDM, with best features in the middle of up-sampling under relatively small noise.Feature quality depends jointly on layer depth and noising scale, and their optimal combination can vary across settings.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3. Label-free monitoring for layer selection

  • Alignment and uniformity of normalized layer features on the unit hypersphere agree with linear-probe quality, and both metrics improve consistently for well-performing layers during training.The metrics are adapted under the assumption that representations are independent of the sampled noise used to form noised images.Source: §3.3 · pp. 4–5, Figure 3 · p. 5

4. Experiments

Experiments ablate denoising and generative pre-training, compare DDAE with prior methods, and test noise-conditional classification and ImageNet-pretrained DiT transfer. The studies examine how diffusion generation quality relates to learned representations and recognition.

4.1.1 Denoising autoencoding

  • The full DDPM noise-level count and broadest noise-scale range give the best generation and recognition, while reducing either weakens both performances.Despite large FID declines, the highest linear probe accuracy falls by less than 3%; T = 64 yields stronger features than other handicapped models despite failing to generate meaningful images.Source: §4.1.1 · p. 6, Figure 4 · p. 6
  • Recognition tends to overfit earlier than generation: linear probe accuracy tends to overfit after 1000 DDPM epochs while generation performance continues improving.The paper proposes that noisy versions of the same image act as aligned positive views, while denser and wider noise increases view diversity but can require longer training: about 1000 epochs for half configurations versus 2000 for the full model.Source: §4.1.1 · p. 6

4.1.2 Generative pre-training

  • On CIFAR-10, EDM improves linear probe accuracy over DDPM from 88.5% to 95.9% and reaches 97.2% after fine-tuning, outperforming supervised WideResNet-28-10.On Tiny-ImageNet, EDM raises recognition rates over DDPM by 3.2% for linear probing and 5.3% for fine-tuning; the study notes its models and hyperparameters there are not optimized.Source: Figure 5 · p. 7, §4.1.2 · p. 7

4.1.3 Noise-conditional classifier

  • A two-layer classifier on frozen DDAE features outperforms supervised WideResNets across almost all noise scales and predicts classification logits alongside diffusion outputs in one forward pass.The method extracts noise-perturbed features and conditions the classifier head on timestep embeddings.Source: §4.1.3 · pp. 7–8, Figure 6 · p. 8

4.2. Comparison with previous methods

  • CIFAR-10 DDAE (EDM) achieves 95.9% linear and 97.2% fine-tuned accuracy; Tiny-ImageNet DDAE (EDM) achieves 50.0% linear and 69.4% fine-tuned accuracy.On CIFAR-10, the fine-tuned model surpasses WideResNet-28-10; on Tiny-ImageNet, fine-tuning brings DDAE to WideResNet-level accuracy and above HybViT.Source: §4.2 · p. 8
  • EDM-based DDAE is comparable to SimCLR on CIFAR-10 and outperforms SimCLR ResNet-18 on Tiny-ImageNet, though ResNet-50 can exceed its linear probe with fewer parameters.On CIFAR-10, fine-tuned DDAE outperforms similarly sized SimCLRs but trails the 375M SimCLR by 0.5%.Source: §4.2 · p. 8, Table 3 · p. 9
  • Scaled DiT-XL/2 exceeds MAE ViT-B/16 in all reported settings except CIFAR-10 linear probing and reaches 98.4% CIFAR-10 accuracy after fine-tuning.It catches up with 375M SimCLR; the authors note DiT was class-conditional pretrained but evaluated unconditionally, so the comparison is not strictly fair.Source: §4.2 · pp. 8–9, Table 3 · p. 9

5. Discussion and conclusion

The paper proposes diffusion pre-training as a unified route to image generation and visual understanding, with potential for vision foundation models, while identifying open architectural, efficiency, and learning questions.

Backbone issues

  • Mid-network truncation is non-ideal and requires probing for the best layer; metric-based selection works on CIFAR-10 but may fail on more complex datasets.The authors suggest explicit encoder-decoder disentanglement as an ideal backbone design and say latent-based networks’ ability to rival pixel-space models on more recognition tasks needs further study.Source: §5 · p. 9

Efficiency issues

  • DDAEs can match some pure recognition models’ accuracies, but their larger model sizes and long training durations make them inefficient and costly to scale.Diffusion models require longer training to achieve optimal generative performance.Source: §5 · p. 9

Relation to other self-supervised methods

  • The authors hypothesize that noising-trajectory alignment acts like positive-only contrastive learning and propose adding consistency or self-predictive constraints to DDAE features.The concurrent Consistency Models study [52] applies a related idea to denoising outputs and shares similarities with self-distillation-based contrastive learning [9]; combining MAEs with contrastive methods [28] motivates constraints at a designated encoder-decoder interface, which may also address prior limitations and improve efficiency.Source: §5 · p. 9

A. Implementation details

The section describes the diffusion models and recognition-training procedures used in the experiments. It also notes limitations of offline latent-space training and the duration used for supervised-from-scratch comparisons.

Diffusion pre-training

  • DDPM and EDM use WideResNet-based UNets, with CIFAR-10 checkpoints retrieved and Tiny-ImageNet unconditional models trained using official or equivalent implementations.Their UNets include convolutional down- and up-sampling stages and single-head self-attention in residual blocks at some resolutions.Source: A. Implementation details · p. 12
  • DiT-XL/2 pre-trained on 256^2 ImageNet and a Stable Diffusion VAE with down-sample factor 8 are retrieved off the shelf; smaller DiTs are not trained due to computational cost.Source: A. Implementation details · p. 12

Linear probing and fine-tuning

  • Linear probing and fine-tuning use simple settings without hyperparameter tuning or normalization layers before linear classifiers.The authors report that classifiers work well without the additional normalization layers.Source: A. Implementation details · p. 12
  • Latent-space DiTs are trained offline on cached VAE codes from 10 augmented training-set versions, sampling one version per epoch.The authors note that augmentation may be insufficient and that more versions or an online VAE encoder may improve recognition accuracy.Source: A. Implementation details · p. 12

Supervised training from scratch

  • Supervised-from-scratch truncated UNet encoders are compared with supervised Wide ResNets after 200 training epochs to reach maximum performance for fair comparison.Source: A. Implementation details · p. 12

B. Layer-noise combinations in grid search

The section describes a grid search over layer–noise combinations to find suitable settings for each model and dataset, because these combinations strongly affect representation quality. It reports the search procedure and points to Table 9 for the combinations adopted in Section 4.

B. Layer-noise combinations in grid search

  • For 18- or 50-step EDM models, the search trained a linear classifier for 10 epochs at each layer and timestep; for 1000-step DDPM or DiT, it sampled timesteps in increments of 5 or 10.The larger timestep increments were used to search more efficiently.Source: B. Layer-noise combinations in grid… · p. 12
  • The grid search aimed to find a suitable, though not necessarily best, layer–noise combination for each model and dataset, and Table 9 lists the combinations adopted in Section 4.The paper states that layer-noise combinations substantially affect representation quality; the supplied table caption does not provide the specific combinations.Source: B. Layer-noise combinations in grid… · p. 12, Table 9 · p. 12

Related papers

Loading 2303.09769v2…