DDAE++: Enhancing Diffusion Models Towards Unified Generative and Discriminative Learning

TL;DR

This paper tackles the challenge of learning strong, transferable representations from diffusion models without sacrificing generation quality.It introduces self-conditioning, a lightweight mechanism that aggregates intermediate semantic features and re-injects them into decoding to concentrate high-level semantics across UNet, UViT, and DiT backbones.By pairing this architectural change with non-leaky augmentations and contrastive self-distillation (MoCo v3) via an EMA teacher, the approach achieves concurrent gains in generation metrics (FID/IS) and discriminative metrics (linear accuracy, segmentation), and demonstrates scalable improvements on ImageNet-1k.Analyses via CKA reveal a pronounced semantic bottleneck and a decoupled decoding phase induced by self-conditioning, supporting its role as a bridge between representation learning and generative refinement.

Abstract

While diffusion models excel at image synthesis, useful representations have been shown to emerge from generative pre-training, suggesting a path towards unified generative and discriminative learning. However, suboptimal semantic flow within current architectures can hinder this potential: features encoding the richest high-level semantics are underutilized and diluted when propagating through decoding layers, impeding the formation of an explicit semantic bottleneck layer. To address this, we introduce self-conditioning, a lightweight mechanism that reshapes the model's layer-wise semantic hierarchy without external guidance. By aggregating and rerouting intermediate features to guide subsequent decoding layers, our method concentrates more high-level semantics, concurrently strengthening global generative guidance and forming more discriminative representations. This simple approach yields a dual-improvement trend across pixel-space UNet, UViT and latent-space DiT models with minimal overhead. Crucially, it creates an architectural semantic bridge that propagates discriminative improvements into generation and accommodates further techniques such as contrastive self-distillation. Experiments show that our enhanced models, especially self-conditioned DiT, are powerful dual learners that yield strong and transferable representations on image and dense classification tasks, surpassing various generative self-supervised models in linear probing while also improving or maintaining high generation quality.

Figures (13)

• Figure 1: Trade-offs and gains in generation and representation. We enhance standard diffusion models in both domains, without requiring major framework overhauls or external knowledge.
• Figure 1: Self-conditioning applied to backbones based on adaptive normalization. Originally, time $t$ (and optional condition $c$) are specified via a global conditioning pathway, where their embedding $e$ is injected into all layers. Here, we collect features from a specific layer by average pooling, and add them to the pathway of decoding layers, after being projected and time-adaptively scaled.
• Figure 2: Self-conditioning applied to backbones based on in-context attention. Based on the all-as-tokens design, we leverage an additional token to automatically interact with patch tokens and the time token, eliminating the need for manual feature selection, pooling, modulation or rerouting.
• Figure 3: Component-wise analysis and comparison with other SSL methods. Self-conditioning, when using in conjunction with Aug and/or CL, provides additional improvements to surpass other self-supervised or diffusion-based models. Best results for each backbone are in bold.
• Figure 4: Detailed performance evolution. Our method maintains a consistent lead over the corresponding baselines, with continuously widening accuracy gaps. FID/IS-10k is used for efficiency.