Denoising Diffusion Autoencoders are Unified Self-supervised Learners

Diffusion models [26, 53, 29] define a series of data corruptions which apply Gaussian noise to data $x_{0}$. Given timestep $t=1,...,T$ which indicates noise levels, the corruption is defined as $q(x_{t}|x_{0})=\mathcal{N}(\alpha_{t}x_{0},\sigma_{t}^{2}\mathbf{I})$, where $\alpha_{t}$ and $\sigma_{t}$ are hyper-parameters controlling the signal-to-noise ratio. When $T$ is large enough, data will be approximately corrupted to $x_{T}\sim\mathcal{N}(0,\mathbf{I})$. With the reparameterization trick, a noised version of $x_{0}$ at an arbitrary level of $t$ can be obtained, by sampling $\epsilon\sim\mathcal{N}(0,\mathbf{I})$ and taking:

Diffusion models aim to invert the corruption and reconstruct samples. Specifically, a random $x_{T}$ is drawn from $\mathcal{N}(0,\mathbf{I})$, and the model samples $x_{t-1}\sim q(x_{t-1}|x_{t})$ iteratively until getting $x_{0}$. Diffusion models employ trainable networks to approximate the transitions with $p_{\theta}(x_{t-1}|x_{t})=\mathcal{N}(\mu_{\theta}(x_{t},t),\Sigma_{t}^{2}\mathbf{I})$, whose mean $\mu_{\theta}(x_{t},t)$ is predicted by networks and variance $\Sigma_{t}^{2}$ is a constant. By simplifying the maximize likelihood objective, networks are equivalently trained with the denoising autoencoder objective [29]:

where $D_{\theta}(x_{t},t)$ is a function of $x_{t}$ and $\mu_{\theta}(x_{t},t)$. Therefore, a denoising network trained by diffusion modeling can be seen as a multi-level and level-conditional version of DAEs, specifically Denoising Diffusion Autoencoders (DDAE).

DDPM [26] proposes the “Variance Preserving” parameterization, where $\alpha_{t}=\sqrt{\Pi_{i=1}^{t}{(1-\beta_{i})}}$ and $\alpha_{t}^{2}+\sigma_{t}^{2}=1$. $\beta_{1...T}$ are determined by a linear schedule from $\beta_{min}$ to $\beta_{max}$. The network is trained to minimize the noise prediction error $\|\epsilon_{\theta}(x_{t},t)-\epsilon\|^{2}$, which is a re-weighted version of Eq. 2, since the denoiser can be derived from $D_{\theta}(x_{t},t)=\frac{x_{t}-\sigma_{t}\epsilon_{\theta}(x_{t},t)}{\alpha_{t}}$. DDPM uses a UNet with 35.7M parameters, and achieves competitive performance on CIFAR-10.

EDM [29] proposes to use the “Variance Exploding” parameterization, where $\alpha_{t}=1$ and $\sigma_{t}$ is sampled by an improved schedule. Based on score-based stochastic differential equations [53], which extend the corruptions to infinite timesteps, EDM yields state-of-the-art results on CIFAR-10 using a larger DDPM++ network [53] with 56M parameters.

Apart from the mathematical formulation improvements, some studies explore efficient and scalable architectures for DDAEs. In particular, Latent Diffusion Models (LDM) [48] perform diffusion modeling within the compressed VAE latent space and outperform pixel-space models on high-resolution text-to-image synthesis. DiT [41] explores scalable backbones under the LDM framework, which replaces UNets with Vision Transformers and achieves state-of-the-art results on class-conditional ImageNet generation.

In this paper, we consider (i) the original DDPM UNet, (ii) DDPM++ trained by EDM, and (iii) latent-space DiT as representative DDAE implementations. DDPM(++) uses UNets with timestep embeddings to parameterize $\epsilon_{\theta}(x_{t},t)$ or $D_{\theta}(x_{t},t)$. DiT uses latent-space ViTs with timestep and label embeddings to learn $\epsilon_{\theta}(x_{t},y,t)$ for conditional generation, and jointly learns an unconditional model $\epsilon_{\theta}(x_{t},\O,t)$ to achieve classifier-free guidance [27] sampling.

Extracting meaningful and discriminative representations from DDAEs is not trivial. Although deterministic inference methods like DDIM [51] are able to derive uniquely identifiable encodings, the contained information is not compact enough for classification. There also exists a trend to employ additional encoders to learn representations for attribute manipulation [42, 66] or classification [1, 37]. In contrast, inspired by iGPT [10] and DDPM-seg [4] which evaluate learned features in autoregressive Transformers and diffusion UNets, we propose to directly take the intermediate activations in pre-trained DDAEs. This approach does not require modification to common diffusion frameworks and is compatible with all existing models.

Drawing from the connections with denoising autoencoders, it is possible that DDAEs can produce linear-separable representations at some implicit encoder-decoder interfaces, resembling MAE. Driven by this, we extend previous investigations on GPTs and DDPM [10, 4] to various network backbones (UNets and DiT) under different frameworks (DDPM and EDM). Considering that UNet with skip connections has been the de-facto design, we avoid splitting the encoder-decoder explicitly to prevent diminishing the generation performance. However, the best layer to extract features remains unknown. Additionally, to prevent the gap between pre-training and deploying, images have to be noised by certain scales for linear evaluations. Considering both the aforementioned facts, we investigate the relationship between feature qualities and layer-noise combinations through grid search, following DDPM-seg.

To apply noise, we randomly sample $\epsilon$ and use Eq. 1 to obtain $x_{t}$, since no obvious differences can be observed between random and deterministic noising. Linear probe accuracies on the features after global average pooling are examined, as illustrated in Figure 1(b). Figure 2 shows that layer depths and noising scales affect feature quality jointly as a concave function, whose global maximum point can be found empirically. For the resolution of $32^{2}$ pixels, the best features lie in the middle of up-sampling, rather than at the lowest resolution as in common practices. Furthermore, we find that perturbing images with relatively small noises improves the linear probe performance, especially on DDPM++ trained by EDM, which achieves $95\%+$ linear probe accuracy and surpasses classical AE or VAEs [1].

These properties have been verified across different datasets and models, but the optimal layer-noise combination may vary under different settings. In the remainder of this paper, linear probe accuracies are reported as the highest found in grid search. For fine-tuning, clean images are passed to DDAE encoders, which are truncated at the optimal layers. The timestep inputs are also fixed to the optimal values. Note that this may not perform best for fine-tuning, since we find the fine-tuning accuracy can be further improved if more layers are used, and it is less sensitive to timestep inputs. However, we keep the layer-noise setup consistent with linear probing to reduce notation overhead.

Although DDAE learns high quality representations for discriminative tasks, it may rely on probing with annotated data to search for proper configurations. Therefore, it would be valuable if label-free metrics can indicate the best performing layers at training. Inspired by feature distribution analysis [58] in contrastive learning, we normalize the features from each layer to be on the unit hypersphere and investigate their alignment and uniformity [58] properties.

In the context of contrastive learning, the alignment loss directly evaluates the distance between positive pairs, while the uniformity loss aims to measure the distribution uniformity on the unit hypersphere:

To apply them to DDAEs, we need to define the positive pairs $(x,y)\sim p_{\rm pos}$ properly. Considering that the feature extraction method proposed in Section 3.2 does not rely on deterministic noising, we assume that the representation should be independent to the noise sample $\epsilon$ used in Eq. 1. To this end, we adapt these two metrics as:

where $f_{t}(\cdot)$ is the DDAE encoder, $x_{t}^{\epsilon}$ denotes the noised image using $\epsilon$, and $\epsilon_{1},\epsilon_{2}$ are independently sampled noise.

We present the metrics with respect to selected layers in Figure 3, to validate whether they align with the feature quality. We evaluate three DDAE implementations on the CIFAR-10 training set: DDPM, DDPM++ trained by DDPM, and DDPM++ trained by EDM. For each model, $t$ is fixed to the respective optimal value, and we monitor the trajectories of metrics during training. Figure 3 shows that overall, $\mathcal{L}_{\rm align}$ and $\mathcal{L}_{\rm uniform}$ agree with the linear probe results (near lower left corners). Moreover, features from the well-performing layers show consistent improvements to both metrics during training. These results indicate that diffusion training may share similarities with contrastive representation learning, and selecting layers by metrics may mitigate the layer searching issue.

Previous research has demonstrated that in Transformers, better generative performance links to better representations, as measured by log-likelihood and linear probe accuracy [10]. Accordingly, we propose to investigate the correlation between generative and discriminative capabilities of DDAEs by plotting evaluation accuracy as a function of image quality calculated on 50k samples. In particular, the widely applied Fréchet Inception Distance (FID) [25] is used as the metric for generation quality. We conduct ablation studies from two perspectives: (i) denoising autoencoding, and (ii) generative pre-training, which correspond to the two sides of DDAEs as discussed in Section 1.

We compare EDM and DiT models with other diffusion-based representation learning methods and self-supervised discriminative methods. Since the generative performance is our first priority, we select checkpoints with the lowest FID for DDAEs, despite the recognition rates may overfit.

Comparison with diffusion-based methods. Table 2 shows that EDM-based DDAE outperforms all previous supervised or unsupervised diffusion-based methods on both generation and recognition. Moreover, our DDAE can be seen as the state-of-the-art hybrid model [61] on CIFAR-10, which can generate and classify (through linear classifier) with a single model. On Tiny-ImageNet, our self-supervised EDM yields significantly better generation FID than the supervised HybViT, despite a lower linear probe accuracy. After fine-tuning, DDAE catches up with supervised WideResNet and surpasses HybViT by large margins.

Comparison with contrastive learning methods. Table 3 presents the evaluation results on CIFAR-10. For linear probing, EDM-based DDAE is comparable with SimCLRs considering model sizes. After fine-tuning, EDM achieves 97.2% (w/o transfer) and 98.1% (w/ transfer) accuracies, outperforming SimCLRs with comparable parameters, despite underperforming the scaled 375M SimCLR model by 0.5%. Table 4 presents results on Tiny-ImageNet. Our EDM-based model significantly outperforms SimCLR pre-trained ResNet-18 under both linear probing and fine-tuning settings. However, DDAE is not as efficient as SimCLR on this dataset, that a slightly larger ResNet-50 can surpass our linear probe result with fewer parameters.

Transfer learning with Vision Transformers. To verify the transfer ability on scalable ViTs, we transfer the DiT model, which is pre-trained on $256^{2}$ ImageNet, to CIFAR-10 and Tiny-ImageNet. Since the DiT codebase only provides the largest DiT-XL/2 checkpoint for class-conditional generation, we use it in an unconditional manner by dropping label to null [27]. Although this may not be strictly fair due to the scale and supervision difference, we mainly aim to confirm the scalability of ViT-based DDAEs.

Table 3 and Table 4 show that the scaled DiT-XL/2 outperforms the smaller MAE ViT-B/16 under all settings by large margins except for linear probing on CIFAR-10. It also catches up with the 375M SimCLR and achieves 98.4% accuracy on CIFAR-10 after fine-tuning. These results indicate that similar to pixel-space ViTs, latent-space DiTs can also benefit from scaling and pre-training on larger datasets. However, diffusion pre-trained DiTs may not be as efficient as MAE pre-trained ViTs on recognition tasks, since the former is specifically designed for advanced image generation without optimizing its representation learning ability.