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

## Abstract

Abstract While diffusion models have gained prominence in image synthesis, their generative pre-training has been shown to yield discriminative representations, paving the way towards unified visual generation and understanding. However, two key questions remain:
1) Can these representations be leveraged to improve the training of diffusion models themselves, rather than solely benefiting downstream tasks? 2) Can the feature quality be enhanced to rival or even surpass modern self-supervised learners, without compromising generative capability?
This work addresses these questions by introducing self-conditioning , a straightforward yet effective mechanism that internally leverages the rich semantics inherent in denoising network to guide its own decoding layers, forming a tighter bottleneck that condenses high-level semantics to improve generation. Results are compelling: our method boosts both generation FID and recognition accuracy with 1% computational overhead and generalizes across diverse diffusion architectures.
Crucially, self-conditioning facilitates an effective integration of discriminative techniques, such as contrastive self-distillation, directly into diffusion models without sacrificing generation quality.
Extensive experiments on pixel-space and latent-space datasets show that in linear evaluations, our enhanced diffusion models, particularly UViT and DiT, serve as strong representation learners, surpassing various self-supervised models.

## 1 Introduction

Diffusion models [26, 59, 33, 42, 41] have recently emerged as one of the most powerful and popular techniques in generative AI, renowned for their capability to synthesize photorealistic visual data. These models have demonstrated remarkable versatility and high performance across a spectrum of tasks, such as class-conditional generation [18, 34, 49], text-to-image synthesis [51, 55, 20], and text-guided image editing [45, 23, 9], along with flexible control and customization options [75, 53].

Meanwhile, interest has grown in repurposing pre-trained diffusion models for discriminative tasks. Studies have shown that their intermediate representations [6, 68] are effective for various downstream tasks, particularly dense prediction tasks such as image segmentation [46, 39, 70], keypoint detection [69, 71], and depth estimation [77]. These findings highlight the potential of diffusion models to serve as unified generative-and-discriminative foundation models, which gain a deep understanding of data through generative pre-training, analogous to language models [50, 35].

Despite these advances, diffusion models still face challenges in representation learning: achieving optimal performance often necessitates specialized methods for feature extraction [46, 44], distillation [39], or design of dedicated decoders [77]. Moreover, while adept at capturing local semantics crucial for dense prediction, their global feature quality often underperforms modern self-supervised learners, especially in image-level tasks like linear classification [30, 15]. This discrepancy underscores a limitation in their ability to condense high-level semantics into compact, discriminative features.

Figure: Table 1: Trade-offs and gains in generation and representation. We aim to enhance standard diffusion models for both performances, without requiring complex framework modifications or external knowledge.

Unlike paradigms imposing explicit information bottleneck, such as view alignment in contrastive methods [66] or asymmetric encoder-decoder in masked autoencoders [22], diffusion models distribute semantic information across layers, where each offers representations of varying granularity [6], jointly contributing to denoising. Such a distributed representation scheme largely accounts for the challenge in obtaining one compact feature (elaborated in Fig. [5](https://arxiv.org/html/2505.10999v1#S5.F5)).

For effective generative-and-discriminative dual learning, we argue that these representations, currently distributed diversely within diffusion networks, should not be treated uniformly — the feature encoding the richest high-level semantics is currently underutilized in guiding the decoding process, thereby limiting both generation quality and the ability to form a more condensed bottleneck [30]. Motivated by this, we introduce self-conditioning, a simple mechanism to enhance both capabilities of diffusion models. It aggregates and condenses rich features from within the network, and re-injects them to condition its own decoder. This encourages the model to leverage more high-level semantic cues to guide generation, while also fostering a tighter bottleneck, as later supported by Fig. [5](https://arxiv.org/html/2505.10999v1#S5.F5).

To ensure broad applicability and generalizability, we have implemented our mechanism on diverse architectures, including UNet [52], UViT [3], and DiT [49].
For UNet and DiT, models that both employ adaptive normalization to integrate global conditions, we enhance this existing conditioning pathways. This is achieved by collecting features from a designated layer via global average pooling, and then injecting the condensed features back into it (Fig. [1](https://arxiv.org/html/2505.10999v1#S4.F1)).
In the case of UViT, where conditioning is realized by treating all inputs as tokens that interact via self-attention, we introduce an additional token that automatically aggregates features through Transformer layers, while also interacting with the image tokens to condition them, resulting in a simple and elegant method (Fig. [2](https://arxiv.org/html/2505.10999v1#S4.F2)).

Self-conditioning adds minimal computational cost yet simultaneously enhances both sample and feature quality. Furthermore, it generalizes well across diverse diffusion formulations and network architectures. As shown in Tab. [1](https://arxiv.org/html/2505.10999v1#S1.T1), our architectural tweaks are compatible with standard diffusion models in the unconditional setup, obviating the need for framework reconstructions [15, 30], or the incorporation of external knowledge from foundation models [28, 40, 73]. This makes our study, namely DDAE++, as broadly applicable as the original DDAE [68] for regular diffusion models.

Moreover, by bridging the gap between representation learning and diffusion modeling, we pave the way for integrating discriminative techniques into diffusion models, such as contrastive self-distillation [14] based on non-leaky augmentation [33]. To the best of our knowledge, this represents a pioneering effort in directly embedding contrastive principles within diffusion modeling.

Experiments demonstrate that combining these methods yields significant and consistent improvements in recognition accuracy, while generation FID either improves or maintains state-of-the-art levels, a contrast to works that may sacrifice generation abilities [15, 30]. Particularly noteworthy are UViT and DiT, which, when augmented with our enhancements, exhibit exceptional potential for representation learning. They achieve superior linear probing results across a range of datasets and surpass various self-supervised models, even in a class-conditional setting. Crucially, the most substantial gains from discriminative techniques are observed when paired with self-conditioning, highlighting the synergy facilitated by our approach between generative and discriminative paradigms.

## 2 Related work

Self-supervised representation learning.
The two dominant paradigms in self-supervised learning (SSL), contrastive learning (CL) and masked image modeling (MIM), focus on instance discrimination [67] and denoising autoencoding [65], respectively.
CL aggregates image-level representations into a compact feature by pulling augmented views of the same image closer. This inherently discriminative design makes it particularly well-suited for linear evaluations [21, 14].
Conversely, MIM captures recoverable semantics by reconstructing the input. Although this generative process preserves rich information, the resulting features are often less compact and not optimal for direct linear classification without subsequent fine-tuning [5, 22].
Our work situates diffusion models within the denoising autoencoder framework (in line with [68, 15]), but we adapt them to explicitly enhance the learning of discriminative features, inspired by hybrid approaches that combine CL with MIM [29, 78].

Semantic-enhanced diffusion models.
Unconditional models often lag behind those conditioned on semantic cues in terms of generation quality [2]. To bridge this gap, some studies use pre-trained foundation models to generate pseudo-labels, which are derived either by extracting global features from powerful SSL encoders [40] or by clustering [28, 27]. However, such pseudo-conditional models, by their reliance on external labels, may sacrifice the broader flexibility in unconditional applications (e.g., image translation [60] and domain adaptation [42]) or class-conditional controllability.
Another category of methods like REPA [73] utilizes SSL features for output regularization and as a distillation target, rather than as an input condition.
In contrast, our self-contained approach enhances standard diffusion models in a plug-and-play fashion, fully preserving the inherent flexibility and focusing on improving internal diffusion architectures, potentially orthogonal to external methods like REPA.

## 3 Background

### 3.1 Preliminaries

Diffusion models.
Diffusion [26, 33] and flow-based [41, 42] models are state-of-the-art generative models with strong theoretical connections [20]. These models, hereinafter collectively referred to as diffusion models, construct paths that progressively transport data to noise by a time-forward process:

$$ $x_{t}=\alpha_{t}x_{0}+\sigma_{t}\epsilon\quad\text{where}~{}\epsilon\sim \mathcal{N}(0,I),$ (1) $$

with schedules for $\alpha_{t},\sigma_{t}$ over $t\in[0,T]$, such that $t=0$ corresponds to $p_{data}$ and $t=T$ approximates $\mathcal{N}(0,I)$. To reverse from noise to data, an ordinary differential equation (ODE) is adopted [59, 41]:

$$ $\mathrm{d}x_{t}=v_{\theta}(x_{t},t)\mathrm{d}t,$ (2) $$

where its distribution matches marginal $p_{t}(x_{t})$ derived from Eq. ([1](https://arxiv.org/html/2505.10999v1#S3.E1)). Velocity $\mathrm{d}x_{t}/\mathrm{d}t$ is approximated by a time-conditioned network $v_{\theta}$, which, once trained, enables ODE solvers to sample data reversely.

The estimation of $v_{\theta}(x_{t},t)$ is closely related to denoising autoencoding and score matching [59], and optimizing such $v_{\theta}$ can be equivalently reparameterized to a denoising $\epsilon$-prediction objective $\lVert\epsilon_{\theta}(x_{t},t)-\epsilon\rVert_{2}^{2}$, or a $x_{0}$-prediction objective $\lVert D_{\theta}(x_{t},t)-x_{0}\rVert_{2}^{2}$.
Different formulations of diffusion models may vary in the forward process, training objective, and ODE solver. In this paper, we examine three representative formulations, DDPM, EDM, and RF, and give a brief overview of each.

DDPM [26] operates over a large number of discrete timesteps ($T=1000$), based on a variance preserving schedule, i.e., $x_{t}=\alpha_{t}x_{0}+\sqrt{1-\alpha_{t}^{2}}\epsilon$. The network learns to predict the noise $\epsilon_{\theta}$, and is trained using the $\epsilon$-prediction loss. Euler’s method is used for sampling, also known as DDIM [58].

EDM [33] is a continuous-time model based on the variance exploding schedule, i.e., $x_{t}=x_{0}+\sigma_{t}\epsilon$, where $\sigma_{t}$ spans a continuous range of $[0.002,80]$. The network predicts the denoised image $D_{\theta}$, and is trained with the $x_{0}$-prediction loss. 2^nd order Heun solver can be utilized for efficient sampling.

RF (Rectified Flow) [42] defines the forward process as a linear interpolation between data and noise, i.e., $x_{t}=(1-t)x_{0}+t\epsilon$, with $t$ sampled from $[0,1]$. The network directly estimates the velocity and is trained using a $v$-prediction loss $\lVert v_{\theta}(x_{t},t)-(\epsilon-x_{0})\rVert_{2}^{2}$. Euler’s ODE solver is used for sampling.

**Table 2: A comprehensive comparison of generative and discriminative performances of self-supervised diffusion models. Baseline: We extend to more models and backbones, clarifying the impact of each design factor. + Self-conditioning: We introduce a simple method that leverages discriminative features within denoising network to guide generation by itself. As a generalizable enhancement, it improves both metrics in most cases, especially on the more diverse CIFAR-100. Colored values indicate gains (or degradation). Best results for each backbone are in bold.**
| Unconditional CIFAR-10 Generation & Linear Probing |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Model | Backbone | Baseline | + Self-conditioning |  |  |
| FID$\downarrow$ | Acc.$\uparrow$ | FID$\downarrow$ | Acc.$\uparrow$ |  |  |
| DDPM | ddpm | 3.60 | 90.34 | 3.67 (+0.07) | 90.94 (+0.60) |
| EDM | ddpm | 3.39 | 91.41 | 3.30 (-0.09) | 91.07 (-0.34) |
| RF | ddpm | 3.89 | 90.67 | 3.67 (-0.22) | 90.91 (+0.24) |
| DDPM | ddpmpp | 2.98 | 94.02 | 2.75 (-0.23) | 94.34 (+0.32) |
| EDM | ddpmpp | 2.23 | 94.83 | 2.18 (-0.05) | 94.85 (+0.02) |
| RF | ddpmpp | 2.54 | 93.97 | 2.42 (-0.12) | 93.72 (-0.25) |
| DDPM | UViT-S | 4.48 | 93.67 | 4.13 (-0.35) | 94.30 (+0.63) |

Backbones for diffusion models.
DDPM originally introduces a modified UNet [52] backbone, incorporating features like self-attention and group normalization [26], extending the original convolutional structure. Time $t$ is specified through a global conditioning pathway, by adding a sinusoidal embedding [64] of $t$ to each block after transformation. We refer to this basic backbone as ddpm(^1^11Lowercase, monospaced text denotes backbone names (ddpm), distinguishing from model names (DDPM).).
DDPM++ [59] further builds upon this, enhancing the model capacity by doubling the depth and replacing up-/down-sampling layers with BigGAN-style [8] residual blocks, while preserving the core design of ddpm. We refer to this scaled-up backbone as ddpmpp.

Vision Transformer-based (ViT) backbones have gained traction due to their scalability and flexibility [4, 20]. UViT [3], for instance, treats all inputs (image patches, time, and other conditions) as tokens and uses long skip connections to ease training, excelling in both latent-space [51] and pixel-space generation. Meanwhile, DiT [49] leverages the AdaLN-zero modulation to inject global conditions, analogous to UNet’s AdaGN [18]. It is proved to be particularly effective on latent-space generation when using the RF formulation [43, 72].
Note that the high-level design of UViT and DiT remains close to the original ViT [19], which is known for its success in visual recognition [22, 11]. This suggests that the architectural gap between visual encoding and generation may be narrower than previously assumed [68], and may offer additional benefits for scalable representation learning.

### 3.2 Extending DDAE more comprehensively

Most related studies focus on extracting features from one specific model [6, 71, 46, 44], with few efforts to identify which models learn better representations and investigate factors influencing feature quality through model comparisons [68, 77]. A pioneering study DDAE [68] compares DDPM and EDM, finding that EDM outperforms DDPM in both generation and recognition, and suggests a link between them. However, limitations in its methodology suggest that further exploration is warranted:

(i) The comparison setup lacks proper control over design factors such as model formulation, backbone type and size, and augmentation. For instance, it compares DDPM-ddpm with EDM-ddpmpp trained with data augmentation (Fig. 5 in that paper). This makes it difficult to isolate the individual contributions of each factor and give insights on how to make design choices for optimal performance.

(ii) The scope of models compared is limited, omitting recent advances in diffusion models. For example, the efficacy of flow-based models in representation learning, compared to diffusion-based ones, remains unclear. Additionally, whether denoising can effectively train ViTs (previously regarded as data-hungry) for strong representations, especially on small data like CIFAR, remains unexplored.

To address these limitations and show our method’s generalizability, we establish a more comprehensive study, considering a wider range of factors controllably, including models, backbones, and augmentations. We first conduct experiments on CIFAR [36], then gradually scale to complex and even latent-space datasets like ImageNet1k-256x256 [54].

## 4 Approach

We enhance diffusion models by forging a stronger link between learned internal representations and the denoising process, focusing on two aspects:
First, we exploit the rich-semantic representations already within the network. The self-conditioning mechanism is straightforward to implement by leveraging existing pathways, such as adaptive norm [18, 49] or in-context attention [3].
Second, we seek to further enhance feature quality and explore whether the generation benefits accordingly.
To this end, we adapt two common techniques from discriminative representation learning, i.e., data augmentation [32] and contrastive self-distillation [14], and apply them to diffusion models.

### 4.1 Self-conditioning

Figure: (a) UNet with self-conditioning: an illustration.
Refer to caption: x1.png

We identify, aggregate and condense high-level semantics into a compact global feature, and then condition the decoder on it. Flexible implementations are tailored to different backbone architectures:

For UNet and DiT, while other implementations may also be plausible, we opt for an easiest way by leveraging the existing conditioning pathway, as illustrated in Fig. [1(a)](https://arxiv.org/html/2505.10999v1#S4.F1.sf1). Specifically, we transform the collected features and add them to the embeddings in global conditioning. This enriched conditional signal is subsequently propagated to the decoder layers via adaptive norm.
Note that this approach requires specifying which features to collect, i.e., where encoder-decoder splits. Following DDAE, we select the layer that yields optimal accuracy in baseline as the designated feature extraction point.

Direct addition may be suboptimal for generation, as different timesteps may require varying amounts of (noisy) semantic guidance. To fix this, we learn a time-dependent scaling vector, and add modulated features to the pathway, as depicted by Adaptive add and detailed in the pseudocode (Fig. [1(b)](https://arxiv.org/html/2505.10999v1#S4.F1.sf2)).

For vanilla ViT-style UViT, where all conditional information interacts with patch tokens via self-attention, we find an elegant way to automate the entire self-conditioning process. This is achieved by introducing an additional learnable token — initialized randomly and conceptually empty — to dynamically collect and utilize global semantics as it propagates through the Transformer (Fig. [2(a)](https://arxiv.org/html/2505.10999v1#S4.F2.sf1)). We refer to this as a [CLS] token due to its functional similarity to the class token in ViT [19].
This shares similarities with studies on register tokens in ViTs [17] and shallow visual prompts [31].

Note that despite the nomenclature, we do not apply any regularization to this token (purely learned by ViT itself by now), in contrast to the [CLS] used in supervised training. In the following section, however, we do apply contrastive loss to it, thereby explicitly encouraging it to capture features.

### 4.2 Adapting discriminative learning techniques

Non-leaky augmentation, as proposed in EDM [33] by adapting the pipeline from StyleGAN [32], involves a series of geometric transformations, each parameterized. A vector representing these parameters is used to condition the model, preventing augmentation artifacts from leaking into generation. The technique has been shown by EDM to benefit FID on datasets like CIFAR-10.

We adopt this pipeline to avoid potential degradation in FID often observed with the more aggressive ones typically used in contrastive methods [21, 12]. It also serves as the way to construct views for contrastive self-distillation. However, these EDM-style geometric transformations are weaker than discriminative-oriented ones, and may not always generate sufficiently diverse positive pairs, as views can be similar and share excessive mutual information [63].
We also empirically find that disabling the interaction between [CLS] and the non-leaky vector is beneficial in UViT, as shown in Fig. [2(b)](https://arxiv.org/html/2505.10999v1#S4.F2.sf2).

Contrastive self-distillation. Many contrastive methods can be seen as a self-distillation progress, where features are aligned with a dynamically updated teacher, constructed using exponential moving averaging (EMA), which is also relied on by diffusion models [34]. Intriguingly, our preliminary study reveals that the EMA rate for lower FID often also yields higher feature quality, suggesting a naturally evolving, ready-to-use teacher is already present in diffusion modeling.

Figure: (a) UViT with [CLS] token, that aggregates global semantics, conditions the network, and guides denoising automatically.
Refer to caption: x3.png

We adapt the MoCo v3 [14] framework to further enhance the features our self-conditioning relies on. Online features are transformed by a time-dependent MLP projection head, and are trained to align with target features from EMA teacher. Target features are obtained from the same layer, but are computed using noised images from an augmented view, and correspond to a timestep that is optimal for linear probing. Diffusion loss is combined with a MoCo v3 loss, weighted by a hyperparameter $\gamma$.

Summary. We arrive at an approach that combines denoising autoencoding and instance discrimination. Studies with related concepts, like SODA [30] and SD-DiT [79], are discussed in Appx. [B](https://arxiv.org/html/2505.10999v1#A2).

## 5 Experiments

We present a comprehensive investigation into the efficacy of self-conditioning across diverse diffusion formulations, architectures, and scales. We further assess the synergistic performance gains achieved by integrating discriminative techniques. In particular, we address the following research questions:

- •
How effective is our method in concurrently improving dual metrics on diverse baselines? (Tab. [2](https://arxiv.org/html/2505.10999v1#S3.T2))
- •
How does each component interact and contribute to the overall improvements? (Tab. [3](https://arxiv.org/html/2505.10999v1#S5.T3), Fig. [3](https://arxiv.org/html/2505.10999v1#S5.F3))
- •
Does self-conditioning generalize when scaling to more challenging datasets, and how do different architectures (UViT vs. UNet) respond to it? (Tab. [4](https://arxiv.org/html/2505.10999v1#S5.T4))
- •
Does it remain effective in latent-space generation with modern DiT architectures, and even in class-conditional settings? (Tab. [5](https://arxiv.org/html/2505.10999v1#S5.T5), [6](https://arxiv.org/html/2505.10999v1#S5.T6)) How does it work? (Fig. [5](https://arxiv.org/html/2505.10999v1#S5.F5), [6](https://arxiv.org/html/2505.10999v1#S5.F6))

Implementation details. Our pixel-space experiments build upon DDAE [68], after stabilizing some hyper-parameters in feature evaluation. Consequently, the baseline results presented in Tab. [2](https://arxiv.org/html/2505.10999v1#S3.T2) reflect our controlled and overhauled setup, so we do not report its original results. For latent-space models, we adopt the state-of-the-art training recipe, Lightning-DiT based on VA-VAE [72], for fast convergence. FID [24] and IS [56] are used to measure sample quality [18]. Details on backbones, training, evaluation, layer/timestep choices are provided in Appx. [A](https://arxiv.org/html/2505.10999v1#A1), and qualitative results in [C](https://arxiv.org/html/2505.10999v1#A3).

Figure: Table 3: Component-wise analysis and comparison with other SSL methods. Self-conditioning, when using in conjunction with Aug and/or CL, brings significant improvements and can surpass other self-supervised or diffusion-based models. Best results for each backbone are in bold.
Refer to caption: x5.png

### 5.1 Main properties

Observations from the baseline. Tab. [2](https://arxiv.org/html/2505.10999v1#S3.T2) shows the performance of our re-implemented, unconditional baselines, encompassing combinations of the formulations and backbones described in Sec. [3.2](https://arxiv.org/html/2505.10999v1#S3.SS2). We find UViT to be unstable with EDM and RF, so these two are excluded.
Regarding sample quality, EDM generally achieves best FID, consistently outperforming RF, another continuous-time model, suggesting that design factors such as noise scheduling and preconditioning [33] are more impactful. Meanwhile, UViT still lags behind UNet.
For classification, while larger models tend to yield higher results, the model formulation is also crucial. EDM consistently delivers highest accuracy across model sizes and datasets, even when it does not achieve lowest FID (e.g., EDM-ddpm on CIFAR-100).

Self-conditioning improves both metrics. Our mechanism not only improves sample FID but also boosts feature quality, particularly on the more diverse CIFAR-100, where accuracy increases by up to 1.99%, a non-negligible gain within the SSL context. Remarkably, these improvements are achieved with almost no overhead: a mere 0.5M parameters for UNet and a single token for UViT.

Aug alone is not helpful enough. We now extend our analysis to models incorporating augmentation (Aug) and contrastive learning (CL), with results presented in Tab. [3](https://arxiv.org/html/2505.10999v1#S5.T3).
While Aug improves FID to baselines on CIFAR-10, it may cause degradation on CIFAR-100, suggesting this geometric pipeline tuned by EDM for CIFAR-10 may not generalize well. Indeed, in contrast to discriminative SSL, where Aug strategies have been well-explored [12, 21], the development of effective ones for diffusion remains challenging, and may be a key bottleneck for diffusion-based representation learning.

Fortunately, self-conditioning mitigates the issue when applied jointly. It leads to more FID gains in cases where Aug alone only shows marginal ones (e.g., with ddpm), or counteracts FID degradation (e.g., on CIFAR-100).
In baselines, the fact that FID can worsen even when feature improves with Aug also indicates that these features have not been effectively translated in generation. Ours, however, ensure that the enhanced features are directly utilized for denoising, and achieve even higher accuracy.

CL and diffusion are complementary.
We now investigate whether these benefits can be amplified by incorporating an explicit discriminative regularization via CL. While some configurations exhibit a slight FID degradation, it suggests the information needed for (and generated by) diffusion may differ from that in discrimination, which can actually be beneficial to developing more comprehensive representations: results show that combining self-conditioning with CL leads to further gains.

Putting all together. Fig. [3](https://arxiv.org/html/2505.10999v1#S5.F3) visualizes the evolving generation-discrimination relationship. Originally, the correlation (discussed in [68, 73]) is relatively weak, indicating that developing better generative models won’t naturally lead to better features (in line with [15]). In contrast, our methods provide an orthogonal option to concurrently improve both aspects, as a clearer correlation emerges.
They are also proved to be particularly effective for UViT, surpassing ddpmpp with less compute and achieving highest results in Tab. [3](https://arxiv.org/html/2505.10999v1#S5.T3). This suggests that, despite sub-optimal in pixel generation compared to UNet, UViT possesses unique potential to learn powerful representations, which we successfully unlock, even in this low-data regime where ViT may struggle due to the data-hungry nature [19].

### 5.2 Scaling to more challenging datasets

Figure: Table 4: Comparison on two more challenging pixel-space datasets. Best results in respective column are in bold.
Refer to caption: x7.png

Figure: Table 5: Effectiveness in latent-space DiTs.

Tab. [4](https://arxiv.org/html/2505.10999v1#S5.T4) presents results on Tiny-ImageNet [38] and a low-res version of ImageNet100 [62]. These are particularly challenging due to: 1) a higher number of classes but with fewer examples per class and low image quality, and 2) the inherent difficulty of modeling in pixel space without an aid of latent compressor [51].
However, UViT shows significantly improved efficacy on these datasets, despite previously being inferior to UNets. With our [CLS] design, UViT outperforms DDAE-ddpmpp and contrastive methods on Tiny-ImageNet, without any data augmentation other than horizontal flipping. These findings align with the broader trend of pure attention-based Transformers increasingly outperforming UNets. On ImageNet100, self-conditioning also consistently benefits both architectures, and the feature quality from UViT-S (44M, enc+dec) even exceeds that of MAE-based ViT-L (307M encoder) trained on higher resolution.

Figure: Table 6: Detailed evaluation with respect to training epochs. Self-c. denotes self-conditioning.

Finally, we evaluate on latent-space DiT-B models in Tab. [5](https://arxiv.org/html/2505.10999v1#S5.T5). With 1.1% more parameters, our self-conditioned DiT improves performance in both unconditional and class-conditional settings (e.g., 63.25% to 64.42%).
While our aug-free models do not yet match leading contrastive methods like DINO [11], they are highly competitive against L-DAE and MIM-based ViT-B/L with longer training. We hypothesize that the bottleneck limiting our recognition accuracy is the absence of augmentation regularization within latent-space training (at least crop), which is challenging to implement, and we defer to future work.
Tab. [6](https://arxiv.org/html/2505.10999v1#S5.T6) reveals a more comprehensive picture: our models consistently outperform their respective baselines in FID, IS, and linear probing throughout the progress.
Note that modern state-of-the-art DiTs powered by VA-VAE and Lightning-DiT [72] already operate at a high-performance frontier with extremely fast convergence. Against the backdrop of such already highly optimized systems, the ability of our minimal, plug-and-play tweaks to yield further discernible dual improvements is particularly compelling and valuable.

Insights: How does self-conditioning work? Fig. [5](https://arxiv.org/html/2505.10999v1#S5.F5) and Fig. [6](https://arxiv.org/html/2505.10999v1#S5.F6) illuminate how it operates. Specially, we focus on how it reshapes the layer-wise discriminability: in baseline, while there’s a shift of more discriminative features towards middle layers as training progresses (the model itself gradually learns how to adjust the landscape), the shift is slow and the variation is modest, in line with the distributed representation theory in Sec. [1](https://arxiv.org/html/2505.10999v1#S1). In contrast, we not only accelerate the convergence to middle layers, but also create a more pronounced performance hierarchy. This sharper focus suggests a more effective bottleneck. Intriguingly, while our designated layer is the 10th, peak accuracy emerges at an even earlier layer. We leave the deeper investigation to future work. Additionally, the observed reduction in denoising loss suggests that effective semantic guidance is indeed taking place.

Figure: Figure 5: Self-conditioning helps reshape feature distribution. While discriminative features gradually shift towards middle layers, they remain dispersed with modest differences. We accelerate this shift, leading to a pronounced concentration, indicative of a condensed bottleneck.
Refer to caption: x8.png

### 5.3 Ablation studies

Below we show more ablation experiments on design choices in self-conditioning and self-distillation.
Tab. LABEL:ablation:unet_sc: The adaptive addition of semantic features into the pathway is crucial, which consistently yields greater improvements, whereas a direct addition does not show much benefits.
Tab. LABEL:ablation:mlp_head: Using the original MoCo v3 projection head results in the same FID, but linear probing accuracy decreases by 0.5% compared to our time-dependent one.
Tab. LABEL:ablation:target_timestep: Using target features extracted with the optimal timestep that performs in linear probing, leads in accuracy by 0.7% when compared to SD-DiT’s [79] minimal-noise design.
Tab. LABEL:ablation:contrastive_weight: Either excessive and insufficient contrastive weight leads to degradation in both metrics, suggesting that the contrastive method can contribute positively to dual aspects when appropriately tuned.
Tab. LABEL:ablation:uvit_target: We investigate two potential methods to extract features from UViT as distillation target: by global average pooling (in line with linear probing) or by taking the [CLS] token. The [CLS] approach benefits both metrics, indicating that self-conditioning based on this token works well because it aggregates global discriminative semantics naturally.

**Table 7: Ablation studies on CIFAR-10/100 for 1200 epochs. Our design choices are in gray.**
| Models, approaches<br>FID$\downarrow$<br>Acc.$\uparrow$<br>DDPM, baseline<br>3.02<br>94.02<br>DDPM, addition<br>2.80<br>94.18<br>DDPM, adaptive<br>2.76<br>94.34<br>EDM, baseline<br>2.25<br>94.83<br>EDM, addition<br>2.43<br>94.85<br>EDM, adaptive<br>2.21<br>94.85<br>RF, baseline<br>2.57<br>93.97<br>RF, addition<br>2.69<br>93.60<br>RF, adaptive<br>2.42<br>93.72<br>(a) | MLP head<br>FID$\downarrow$<br>Acc.$\uparrow$<br>Original MoCo v3<br>5.87<br>69.01<br>Time-dependent<br>5.88<br>69.50<br>(b) | Target timestep<br>FID$\downarrow$<br>Acc.$\uparrow$<br>Minimal noise [79]<br>5.81<br>68.76<br>Optimal for linear<br>5.88<br>69.50<br>(c) |
| --- | --- | --- |
| $\gamma$<br>FID$\downarrow$<br>Acc.$\uparrow$<br>0.1<br>6.03<br>69.08<br>0.01<br>5.88<br>69.50<br>0.001<br>5.96<br>68.42<br>(d) | Feature to contrast<br>FID$\downarrow$<br>Acc.$\uparrow$<br>Pooling image tokens<br>8.51<br>73.15<br>[CLS]<br>7.37<br>74.48<br>(e) |  |

## 6 Conclusion

We show that established diffusion architectures can be enhanced, by conditioning the decoding process on features learned by themselves. The idea can be surprisingly simple to implement, yet concurrently improves both generation and representation quality at almost no cost.
It also facilitates the integration of discriminative techniques for further recognition gains.
We hope this straightforward principle inspires continued progress towards unified diffusion-based generation and understanding.

## References

- Bandara et al. [2023]
Wele Gedara Chaminda Bandara, Celso M De Melo, and Vishal M Patel.
Guarding barlow twins against overfitting with mixed samples.
*arXiv:2312.02151*, 2023.
- Bao et al. [2022]
Fan Bao, Chongxuan Li, Jiacheng Sun, and Jun Zhu.
Why are conditional generative models better than unconditional ones?
*arXiv:2212.00362*, 2022.
- Bao et al. [2023a]
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu.
All are worth words: A vit backbone for diffusion models.
In *CVPR*, 2023a.
- Bao et al. [2023b]
Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang, Gang Yue, Yue Cao, Hang Su, and Jun Zhu.
One transformer fits all distributions in multi-modal diffusion at scale.
In *ICML*, 2023b.
- Bao et al. [2021]
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei.
Beit: Bert pre-training of image transformers.
*arXiv:2106.08254*, 2021.
- Baranchuk et al. [2022]
Dmitry Baranchuk, Ivan Rubachev, Andrey Voynov, Valentin Khrulkov, and Artem Babenko.
Label-efficient semantic segmentation with diffusion models.
In *ICLR*, 2022.
- Bardes et al. [2022]
Adrien Bardes, Jean Ponce, and Yann LeCun.
Vicreg: Variance-invariance-covariance regularization for self-supervised learning.
In *ICLR*, 2022.
- Brock et al. [2018]
Andrew Brock, Jeff Donahue, and Karen Simonyan.
Large scale gan training for high fidelity natural image synthesis.
In *ICLR*, 2018.
- Brooks et al. [2023]
Tim Brooks, Aleksander Holynski, and Alexei A Efros.
Instructpix2pix: Learning to follow image editing instructions.
In *CVPR*, 2023.
- Caron et al. [2020]
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.
Unsupervised learning of visual features by contrasting cluster assignments.
In *NeurIPS*, 2020.
- Caron et al. [2021]
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.
Emerging properties in self-supervised vision transformers.
In *ICCV*, 2021.
- Chen et al. [2020]
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.
A simple framework for contrastive learning of visual representations.
In *ICML*, 2020.
- Chen and He [2021]
Xinlei Chen and Kaiming He.
Exploring simple siamese representation learning.
In *CVPR*, 2021.
- Chen et al. [2021]
Xinlei Chen, Saining Xie, and Kaiming He.
An empirical study of training self-supervised vision transformers.
In *ICCV*, 2021.
- Chen et al. [2024]
Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He.
Deconstructing denoising diffusion models for self-supervised learning.
*arXiv:2401.14404*, 2024.
- Da Costa et al. [2022]
Victor Guilherme Turrisi Da Costa, Enrico Fini, Moin Nabi, Nicu Sebe, and Elisa Ricci.
solo-learn: A library of self-supervised methods for visual representation learning.
*JMLR*, 23(56):1–6, 2022.
- Darcet et al. [2024]
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski.
Vision transformers need registers.
In *ICLR*, 2024.
- Dhariwal and Nichol [2021]
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat gans on image synthesis.
In *NeurIPS*, 2021.
- Dosovitskiy et al. [2021]
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.
An image is worth 16x16 words: Transformers for image recognition at scale.
In *ICLR*, 2021.
- Esser et al. [2024]
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al.
Scaling rectified flow transformers for high-resolution image synthesis.
*arXiv:2403.03206*, 2024.
- Grill et al. [2020]
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al.
Bootstrap your own latent-a new approach to self-supervised learning.
In *NeurIPS*, 2020.
- He et al. [2022]
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.
Masked autoencoders are scalable vision learners.
In *CVPR*, 2022.
- Hertz et al. [2022]
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or.
Prompt-to-prompt image editing with cross attention control.
*arXiv:2208.01626*, 2022.
- Heusel et al. [2017]
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
Gans trained by a two time-scale update rule converge to a local nash equilibrium.
In *NeurIPS*, 2017.
- Ho and Salimans [2022]
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
*arXiv:2207.12598*, 2022.
- Ho et al. [2020]
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
In *NeurIPS*, 2020.
- Hu et al. [2023a]
Vincent Tao Hu, Yunlu Chen, Mathilde Caron, Yuki M Asano, Cees GM Snoek, and Bjorn Ommer.
Guided diffusion from self-supervised diffusion features.
*arXiv:2312.08825*, 2023a.
- Hu et al. [2023b]
Vincent Tao Hu, David W Zhang, Yuki M Asano, Gertjan J Burghouts, and Cees GM Snoek.
Self-guided diffusion models.
In *CVPR*, 2023b.
- Huang et al. [2023]
Zhicheng Huang, Xiaojie Jin, Chengze Lu, Qibin Hou, Ming-Ming Cheng, Dongmei Fu, Xiaohui Shen, and Jiashi Feng.
Contrastive masked autoencoders are stronger vision learners.
*TPAMI*, 2023.
- Hudson et al. [2023]
Drew A Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James L McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner.
Soda: Bottleneck diffusion models for representation learning.
*arXiv:2311.17901*, 2023.
- Jia et al. [2022]
Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim.
Visual prompt tuning.
In *ECCV*, 2022.
- Karras et al. [2020]
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila.
Training generative adversarial networks with limited data.
In *NeurIPS*, 2020.
- Karras et al. [2022]
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.
Elucidating the design space of diffusion-based generative models.
In *NeurIPS*, 2022.
- Karras et al. [2024]
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine.
Analyzing and improving the training dynamics of diffusion models.
In *CVPR*, 2024.
- Kenton and Toutanova [2019]
Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova.
Bert: Pre-training of deep bidirectional transformers for language understanding.
In *Proceedings of NAACL-HLT*, 2019.
- Krizhevsky et al. [2009]
Alex Krizhevsky, Geoffrey Hinton, et al.
Learning multiple layers of features from tiny images.
2009.
- Kynkäänniemi et al. [2024]
Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen.
Applying guidance in a limited interval improves sample and distribution quality in diffusion models.
In *NeurIPS*, 2024.
- Le and Yang [2015]
Ya Le and Xuan Yang.
Tiny imagenet visual recognition challenge.
2015.
- Li et al. [2023a]
Daiqing Li, Huan Ling, Amlan Kar, David Acuna, Seung Wook Kim, Karsten Kreis, Antonio Torralba, and Sanja Fidler.
Dreamteacher: Pretraining image backbones with deep generative models.
In *ICCV*, 2023a.
- Li et al. [2023b]
Tianhong Li, Dina Katabi, and Kaiming He.
Self-conditioned image generation via generating representations.
*arXiv:2312.03701*, 2023b.
- Lipman et al. [2023]
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
In *ICLR*, 2023.
- Liu et al. [2023]
Xingchao Liu, Chengyue Gong, and Qiang Liu.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
In *ICLR*, 2023.
- Ma et al. [2024]
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie.
Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers.
In *ECCV*, 2024.
- Meng et al. [2024]
Benyuan Meng, Qianqian Xu, Zitai Wang, Xiaochun Cao, and Qingming Huang.
Not all diffusion model activations have been evaluated as discriminative features.
In *NeurIPS*, 2024.
- Mokady et al. [2023]
Ron Mokady, Amir Hertz, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or.
Null-text inversion for editing real images using guided diffusion models.
In *CVPR*, 2023.
- Mukhopadhyay et al. [2023]
Soumik Mukhopadhyay, Matthew Gwilliam, Yosuke Yamaguchi, Vatsal Agarwal, Namitha Padmanabhan, Archana Swaminathan, Tianyi Zhou, and Abhinav Shrivastava.
Do text-free diffusion models learn discriminative visual representations?
*arXiv:2311.17921*, 2023.
- Ozsoy et al. [2022]
Serdar Ozsoy, Shadi Hamdan, Sercan Arik, Deniz Yuret, and Alper Erdogan.
Self-supervised learning with an information maximization criterion.
In *NeurIPS*, 2022.
- Pan et al. [2023]
Zixuan Pan, Jianxu Chen, and Yiyu Shi.
Masked diffusion as self-supervised representation learner.
*arXiv:2308.05695*, 2023.
- Peebles and Xie [2023]
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In *ICCV*, 2023.
- Radford et al. [2019]
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.
Language models are unsupervised multitask learners.
*OpenAI blog*, 2019.
- Rombach et al. [2022]
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
High-resolution image synthesis with latent diffusion models.
In *CVPR*, 2022.
- Ronneberger et al. [2015]
Olaf Ronneberger, Philipp Fischer, and Thomas Brox.
U-net: Convolutional networks for biomedical image segmentation.
In *MICCAI*, 2015.
- Ruiz et al. [2023]
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman.
Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation.
In *CVPR*, 2023.
- Russakovsky et al. [2015]
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.
Imagenet large scale visual recognition challenge.
*IJCV*, 2015.
- Saharia et al. [2022]
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al.
Photorealistic text-to-image diffusion models with deep language understanding.
In *NeurIPS*, 2022.
- Salimans et al. [2016]
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.
Improved techniques for training gans.
In *NeurIPS*, 2016.
- Shin et al. [2024]
Jeongwoo Shin, Inseo Lee, Junho Lee, and Joonseok Lee.
Self-guided masked autoencoder.
In *NeurIPS*, 2024.
- Song et al. [2021a]
Jiaming Song, Chenlin Meng, and Stefano Ermon.
Denoising diffusion implicit models.
In *ICLR*, 2021a.
- Song et al. [2021b]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
In *ICLR*, 2021b.
- Su et al. [2023]
Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon.
Dual diffusion implicit bridges for image-to-image translation.
In *ICLR*, 2023.
- Szegedy et al. [2016]
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna.
Rethinking the inception architecture for computer vision.
In *CVPR*, 2016.
- Tian et al. [2020a]
Yonglong Tian, Dilip Krishnan, and Phillip Isola.
Contrastive multiview coding.
In *ECCV*, 2020a.
- Tian et al. [2020b]
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola.
What makes for good views for contrastive learning?
In *NeurIPS*, 2020b.
- Vaswani et al. [2017]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
In *NeurIPS*, 2017.
- Vincent et al. [2008]
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol.
Extracting and composing robust features with denoising autoencoders.
In *ICML*, 2008.
- Wang and Isola [2020]
Tongzhou Wang and Phillip Isola.
Understanding contrastive representation learning through alignment and uniformity on the hypersphere.
In *ICML*, 2020.
- Wu et al. [2018]
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin.
Unsupervised feature learning via non-parametric instance discrimination.
In *CVPR*, 2018.
- Xiang et al. [2023]
Weilai Xiang, Hongyu Yang, Di Huang, and Yunhong Wang.
Denoising diffusion autoencoders are unified self-supervised learners.
In *ICCV*, 2023.
- Xu et al. [2024]
Guangkai Xu, Yongtao Ge, Mingyu Liu, Chengxiang Fan, Kangyang Xie, Zhiyue Zhao, Hao Chen, and Chunhua Shen.
Diffusion models trained with large data are transferable visual models.
*arXiv:2403.06090*, 2024.
- Xu et al. [2023]
Jiarui Xu, Sifei Liu, Arash Vahdat, Wonmin Byeon, Xiaolong Wang, and Shalini De Mello.
Open-vocabulary panoptic segmentation with text-to-image diffusion models.
In *CVPR*, 2023.
- Yang and Wang [2023]
Xingyi Yang and Xinchao Wang.
Diffusion model as representation learner.
In *ICCV*, 2023.
- Yao et al. [2025]
Jingfeng Yao, Bin Yang, and Xinggang Wang.
Reconstruction vs. generation: Taming optimization dilemma in latent diffusion models.
In *CVPR*, 2025.
- Yu et al. [2025]
Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie.
Representation alignment for generation: Training diffusion transformers is easier than you think.
In *ICLR*, 2025.
- Zbontar et al. [2021]
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.
Barlow twins: Self-supervised learning via redundancy reduction.
In *ICML*, 2021.
- Zhang and Agrawala [2023]
Lvmin Zhang and Maneesh Agrawala.
Adding conditional control to text-to-image diffusion models.
In *ICCV*, 2023.
- Zhang et al. [2022]
Qi Zhang, Yifei Wang, and Yisen Wang.
How mask matters: Towards theoretical understandings of masked autoencoders.
In *NeurIPS*, 2022.
- Zhao et al. [2023]
Wenliang Zhao, Yongming Rao, Zuyan Liu, Benlin Liu, Jie Zhou, and Jiwen Lu.
Unleashing text-to-image diffusion models for visual perception.
In *ICCV*, 2023.
- Zhou et al. [2022]
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong.
ibot: Image bert pre-training with online tokenizer.
In *ICLR*, 2022.
- Zhu et al. [2024]
Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao, Zhenglong Sun, Tao Mei, and Chang Wen Chen.
Sd-dit: Unleashing the power of self-supervised discrimination in diffusion transformer.
*arXiv:2403.17004*, 2024.

## Appendix A Implementation details

Diffusion models. Our pixel-space implementations are primarily based on the official codebases of DDAE(^2^22[https://github.com/FutureXiang/ddae](https://github.com/FutureXiang/ddae)) and UViT(^3^33[https://github.com/baofff/U-ViT](https://github.com/baofff/U-ViT)), including model and backbone definitions for DDPM, EDM, ddpm, ddpmpp, and UViT. The uncertainty loss weighting based on multi-task learning proposed in EDM2 [34] is also incorporated when non-leaky augmentation is applied to EDM, as it mitigates the slower convergence caused by augmentations. We also implement RF following core principles in Flow Matching [41] and Rectified Flow [42]. Our implementations of ddpm and UViT are equivalent to official versions, while ddpmpp is a simplified variant, adhering to main designs outlined in [59] but omitting certain details like skip connection rescaling, which don’t find helpful. Our latent-space implementations are based on the official codebase of Lightning-DiT(^4^44[https://github.com/hustvl/LightningDiT](https://github.com/hustvl/LightningDiT)). Pre-trained VA-VAE checkpoint is used, also downloaded from its official repository. We fix the “3-channel CFG” bug originated in DiT [49] to enhance sample quality. More hyper-parameters are provided in Tab. [8](https://arxiv.org/html/2505.10999v1#A1.T8) and Tab. [9](https://arxiv.org/html/2505.10999v1#A1.T9).

**Table 8: Details of diffusion model formulations.**
| Pixel-space DDPM [26, 58] |  |
| --- | --- |
| $T$ | 1000 |
| training noise schedule | linear beta schedule $[10^{-4},0.02]$ |
| training loss weighting | none |
| dropout rate | 0.1 |
| EMA decay rate | 0.9999 |
| ODE sampler | Euler (DDIM) |
| sampling NFE | 100 |
| Pixel-space RF [42, 41] |  |
| training noise schedule | $t\in[0,1]$ |
| dropout rate | 0.1 |
| EMA decay rate | 0.9999 |
| ODE sampler | RK45 |
| sampling NFE | 140-160 adaptive |

Figure: Table 9: Details of backbone architectures.

Pre-training. On all pixel-space datasets, we train all models for 2000 epochs, saving checkpoints every 200 epochs for FID measurement. We observe that the model typically achieves lowest FID after 1400-2000 epochs. The training setups for DDPM and RF are identical, while EDM differs slightly in warmup, dropout, and EMA rate. These hyper-parameters are inherited and have not been heavily tuned. For UViT, the official implementation applies weight decay to all parameters. However, we find that excluding bias and positional embedding yields better results. We also apply weight decay to the [CLS] token, contrasting with common practices. All these experiments are conducted on 4 NVIDIA 3080Ti or 4090 GPUs with automatic mixed precision enabled.
On latent-space 256x256 datasets, we train DiT models for either 1400 epochs (unconditional, IN100), 1000 epochs (class-conditional, IN100), or 100 epochs (both, IN1k). These training durations are sufficient for conditional models, but unconditional ones may still benefit from longer training. Horizontal flip is applied as the only data augmentation method, unless specified otherwise (i.e., previously using non-leaky on CIFAR). We cache two flipping versions of latents on IN100 or IN1k datasets during data pre-processing. All these experiments are conducted on 6 NVIDIA 4090 GPUs with automatic mixed precision.

Feature extraction. We determine the optimal timestep (noise scale) and layer index for feature extraction by grid searching over reasonable ranges, evaluated using linear probing on the validation set, following the practice in DDAE [68]. The selected values are summarized in Tab. [10](https://arxiv.org/html/2505.10999v1#A1.T10), and may differ from those of DDAE due to slight changes in linear probing settings. Notably, the optimal layer index is also utilized during the training of UNet and DiT models with self-conditioning. Similarly, the optimal noise scale is adopted for extracting the target feature in contrastive self-distillation.

**Table 10: Time and layer for feature extraction. Out denotes output layers in UNet and UViT.**
| Model | Backbone | On CIFAR-10 | On CIFAR-100 |  |  |
| --- | --- | --- | --- | --- | --- |
| Time | Layer | Time | Layer |  |  |
| DDPM | ddpm | $t=11$ | out_6/12 | $t=11$ | out_6/12 |
| EDM | ddpm | $\sigma_{t}=0.06$* | out_7/12 | $\sigma_{t}=0.06$* | out_6/12 |
| RF | ddpm | $t=0.06$ | out_7/12 | $t=0.06$ | out_7/12 |
| DDPM | ddpmpp | $t=11$ | out_7/15 | $t=11$ | out_8/15 |
| EDM | ddpmpp | $\sigma_{t}=0.06$* | out_9/15 | $\sigma_{t}=0.06$* | out_8/15 |
| RF | ddpmpp | $t=0.06$ | out_8/15 | $t=0.06$ | out_8/15 |
| DDPM | UViT-S | $t=11$ | out_2/6 | $t=11$ | out_2/6 |
| *Corresponds to $t=4$ in 18 sampling steps. |  |  |  |  |  |

Linear probing. We simplify the settings in DDAE by using identical training epochs and learning rates for all models. Additionally, we find that random cropping is unnecessary for linear probing, so we use only horizontal flipping as the augmentation method. The training and evaluation configurations are shown in Tab. [11](https://arxiv.org/html/2505.10999v1#A1.T11). On CIFAR datasets, linear probing accuracy is reported as the highest among the checkpoints at 800, 1000, and 1200 epochs. On other datasets, we find that the linear probing accuracy does not saturated till the end of the training.

**Table 11: Details of pre-training and linear evaluations. All our models (as well as other baselines in comparison) are trained and evaluated within same datasets, without transfer learning.**
| Pixel-space experiments: CIFAR, Tiny-IN, IN100-64x64 |  |  |
| --- | --- | --- |
|  | Pre-training | Linear probing |
| GPUs | 4 * 3080Ti (CIFAR) or 4 * 4090 (others) |  |
| batch size | 512 (ddpm) or 256 (ddpmpp, uvit) |  |
| optimizer | Adam (unet)<br>or AdamW (uvit) | Adam |
| warmup epochs | 13 (DDPM, RF)<br>200 (EDM) | — |
| epochs | 2000 | 15 (CIFAR)<br>30 (others) |
| learning rate | 4e-4 | 4e-3 |
| lr schedule | constant | cosine |
| augmentations | flip (optional non-leaky) | flip |

## Appendix B Methodology comparison

Our full approach with contrastive self-distillation culminates in a self-supervised model that learns semantically meaningful representations through cross-view alignment (discriminative, inherited from CL) and intra-view reconstruction (generative, inherent in diffusion denoising, analogous to MIM) [29, 78].
As for generative modeling, our design leads to a self-supported semantic-conditional framework that uses intermediate features as semantic cues to guide generation, with the feature encoder absorbed into the diffusion network as its first few layers.

Relation to SODA.
SODA [30] learns representations through cross-view reconstruction. Similar to our self-conditioning, it also employs a feature modulation mechanism to impose a tighter bottleneck between the encoder and decoder, thereby learning compact, linearly-separable features. However, SODA focuses on image-conditional tasks like novel view synthesis, so it uses a disentangled encoder separate from diffusion decoder. Additionally, SODA’s features are learned through pure generative pre-training, without investigating the influence of contrastive methods.

Please note that a direct comparison with SODA is not so appropriate. First, since SODA has not released its official code, we made our best effort to re-implement one, and we found that it does not perform well on CIFAR and Tiny-ImageNet datasets. Second, SODA is optimized for representations and cannot function as a regular diffusion model in standard unconditional (or class-conditional) settings.
Third, SODA may not achieve superior FID and Acc simultaneously with a same model, as it employs different augmentations for classification (stronger) and reconstruction (weaker).

Relation to “Guided Diffusion from Self-Supervised Diffusion Features”.
This study [27] also improves unconditional generation without relying on external encoders, by utilizing discriminative features within diffusion models. However, it only uses these features to generate pseudo-labels through Sinkhorn-Knopp cluster assignment, without enhancing them by comparing cluster assignments as in SwAV [10]. In contrast, our work presents a more fundamental combination of diffusion pre-training, self-condition guidance, and contrastive feature enhancement.

Relation to SD-DiT.
SD-DiT [79] is a recent work that aims to accelerate the training convergence of DiT [49] through self-distillation. It aligns features extracted from the visible patches of an image with those extracted from the entire image by an EMA teacher. While this joint optimization of generative and discriminative objectives is similar to our approach, it does not focus on enhancing representation quality. Furthermore, its key design, setting the distillation target to the minimal noise scale, differs from ours that using the linear probing timestep, proven more effective in ablation.

Relation to REPA.
REPA [73] is a concurrent work that accelerates DiT [49] training through representation enhancement. It argues that aligning the intermediate features in diffusion models with powerful representations can significantly ease training. This “representation-for-better-generation” idea is similar to ours, but it uses an external large-scale pre-trained encoder as the teacher, rather than leveraging the diffusion models themselves. Additionally, REPA employs MLP heads and similarity functions, such as cross-entropy, to align features, similar to our MoCo v3-based method. However, it does not introduce two views for contrastive learning, as it is fundamentally a knowledge distillation process rather than a self-supervised one.

## Appendix C Qualitative results

To visualize the generation quality, we present randomly generated samples in Fig. [7](https://arxiv.org/html/2505.10999v1#A3.F7) (CIFAR-100) and Fig. [8](https://arxiv.org/html/2505.10999v1#A3.F8) (ImageNet1k-256x256) using the same set of initial noise inputs.

Figure: Figure 7: Generate samples on CIFAR-100 using UViT-S. Our proposed methods gradually improve the overall structure, semantics, and details (e.g., see the motorbike, chair and bugs).
Refer to caption: x10.png

Figure: Figure 8: Generate samples on ImageNet1k-256x256 using DiT-B. Our method also improves the overall structure and details in latent-space unconditional (e.g., see the human face and the dog) and class-conditional generation (e.g., see the lawn mower and the tripod).
Refer to caption: x11.png

## Appendix D Additional information

### D.1 Limitations and future research directions

Although we conduct extensive experiments to prove our claim that it is possible to enhance both generative and discriminative performance simultaneously through our approaches, our evaluations are limited to datasets up to the ImageNet1k-256x256 scale and models up to the DiT-Base scale. We did not investigate larger datasets with higher resolutions or larger models, such as DiT-XL, due to high computational costs inherently for generative models.

Moreover, some results indicate that our approaches may rely on careful choice of data augmentation strategies, which might require tuning when dataset changes. A key remaining challenge for integrating generative and discriminative learning, we believe, is the development of effective strategies to organize and identify multiple views of the same instance, meriting future research.

Finally, though our work provides initial insights into the formation, distribution, and enhancement of internal representations within diffusion models, the precise dynamics governing how these representations evolve throughout the training remain largely unexplored. Moreover, the consequences of potential misalignment between the layers designated for self-conditioning and those where optimal features naturally arise warrant further investigation.

### D.2 Societal impacts

Although we focus on exploring the frontiers in unconditional (or class-conditional) image generation and representation learning, on datasets that contain natural images (instead of human portraits or faces), our method still brings an improvement in the quality of generative models, which could be used to generate fake images for disinformation.
However, the images generated by our models, restricted by the unconditional (or class-conditional) nature, are not as high-fidelity and photorealistic as those in the text-to-image synthesis.
Therefore, the potential risk of negative impact is very low.

### D.3 Licenses

Datasets:

- •
CIFAR-10/100 [36]: MIT license.
- •
ImageNet1k [54]: Custom terms (for non-commercial research and educational use only).
- •
Tiny-ImageNet [38]: Unclear. Probably inherits the original ImageNet terms.
- •
ImageNet100 [62]: Unclear. Probably inherits the original ImageNet terms.

Pre-trained models:

- •
VA-VAE [72] (latent-space compression): MIT license.
- •
Inception-v3 [61] (FID and IS measurement): Apache-2.0 license.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 2
presentation: 3
questions: - Mathematical Formalization: Can the authors provide a precise, equation-level definition of the “time-adaptive scaling” used in self-conditioning? How is this parametrized and learned? Please specify if there is a generalizable mathematical principle or derived gradient supporting its use.
- Rigorous Bottleneck Selection: Is there theoretical—or at least empirical—evidence that the training-loss-based search for the bottleneck layer always selects optimal features for discriminative purposes? Could the authors provide a more principled justification for this heuristic?
- Limits of Dual-Improvement: In Table 2 / Table 3, there are cases where FID is marginally degraded, or improvements are not observed concurrently. Can the authors break down these exceptions and clarify the precise conditions under which self-conditioning may not help (or may hurt)?
- Feature Concentration Analysis: Beyond qualitative shifts shown in Figure 5, have the authors performed any quantitative analysis (e.g., feature entropy, linear separability, mutual information with ground truth) of the “condensed bottleneck” feature? This would greatly strengthen the core claim.
- Augmentation Sensitivity: Given the importance and variability of data augmentations, can the authors provide more systematic experiments or concrete recommendations on how to choose or tune these for new datasets? What are the consequences if the geometric pipeline is poorly matched to the data?
- Layer Dynamics: The authors mention that the peak accuracy emerges at a different layer than the one chosen for self-conditioning (Figure 5, Section 5.3). Could the authors elaborate on the implications of this misalignment for both generation and discrimination?
- Scalability Limits: Are there any architectural or optimization bottlenecks (beyond compute cost) that emerge at higher dataset resolutions or with substantially larger backbones/latent spaces?
soundness: 2
strengths: - The work identifies and addresses a genuine architectural weakness in standard diffusion models: the absence of a discriminative bottleneck due to distributed semantic flow. The “self-conditioning” mechanism is both simple and effective—clearly illustrated in Figure 1 and Figure 2—and does not require external supervision.

- The empirical evaluation is extensive and well-controlled, covering a variety of diffusion model backbones (UNet-based, UViT, DiT) and datasets (CIFAR-10/100, Tiny-ImageNet, ImageNet at various scales). The methodology section and ablations demonstrate care in validating architectural choices, as well as strong generalization of the performance boosts. The results are thoroughly tabulated in Table 2, Table 3, and Table 4, and consistently show that self-conditioning brings gains to both generative (FID/IS) and discriminative (linear accuracy) metrics.

- The method is lightweight (as detailed in the ablations and parameter increase figures/tables in Sections 5.1 and 5.3) and can be retrofitted to diverse diffusion architectures.
summary: This paper addresses a key limitation in diffusion models—namely, the sub-optimal semantic information flow affecting the quality of learned representations—by introducing a lightweight architectural mechanism called self-conditioning. The method aggregates high-level semantic features from intermediate layers and injects them back into the decoding pathway, forming a more effective bottleneck for discriminative features without relying on external cues. Additional integration of contrastive self-distillation and principled data augmentations is demonstrated to further unify and amplify both generative and discriminative capacities of diffusion models. Empirical evidence on several backbones (UNet, UViT, DiT) and datasets (CIFAR-10/100, ImageNet) shows dual improvements in generation (FID/IS) and representation (linear probing accuracy compared to strong self-supervised baselines).
weaknesses: - Although the experiments are broad, there is a bias towards popular image datasets, especially at lower resolutions. In Section E.1 (Limitations and Future Research Directions), the authors admit that their results stop at ImageNet 256x256 and DiT-base scale due to compute constraints, and do not demonstrate scalability for larger models/datasets relevant to modern AI.

- The ablation studies in Table 5 focus on demonstrating parameter sensitivity and certain hyperparameter effects (e.g., MLP head design, distillation timestep, etc.). However, there is little diagnostic quantification of why self-conditioning enhances the bottleneck: for instance, layer-wise representation transferability, sparsity/concentration metrics, or feature alignment before and after rerouting are not rigorously analyzed.

- The comparison in Table 4 is unconvincing because the FID scores on ImageNet for both the baseline and the self-conditioning variant indicate that the models have not reached an optimal level of convergence. While we acknowledge that computational constraints may limit the number of training epochs, conclusions drawn from comparisons between undertrained models are not reliable.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: The main text contains excessive whitespace, and the formatting in the appendix (around Line 712) is messy and requires careful reorganization.
soundness: 3
strengths: The paper presents a clear and intuitive idea, supported by solid experimental analysis. In particular, the layer-wise feature analysis via linear probing in Figure 5 provides valuable insights into why the proposed DDAE method is effective — this diagnostic perspective is worth learning from. Overall, both Table 3 and Figure 3 convincingly demonstrate that DDAE is an effective and practical approach.
summary: This paper proposes a simple yet effective modification called **self-conditioning** to improve both the generative and discriminative capabilities of diffusion models. The method reroutes semantically rich intermediate features to guide decoding layers, creating a stronger representation bottleneck without external supervision. This architectural enhancement allows diffusion models to better integrate discriminative objectives such as contrastive self-distillation, achieving consistent dual improvements across various backbones (UNet, UViT, DiT) and datasets. Experiments show that DDAE++ boosts both generation quality (lower FID) and feature quality (higher linear probing accuracy) with minimal computational overhead, positioning diffusion models as unified learners for generation and recognition.
weaknesses: The claim in Figure 6 that "Self-conditioning facilitates the optimization and narrows the loss gap between un- and class-conditioning" does not hold. From the curves shown, we can only observe that both the conditional and unconditional loss curves converge faster, but there is no evidence indicating that the gap between them actually decreases.

---

confidence: 3
contribution: 2
initial_rating: 2
presentation: 2
questions: See weaknesses
soundness: 2
strengths: * Analyzing and improving the representation of generative models is an important topic.
* The method is applicable to various network backbones.
* The ablation study part is performed in detail.
summary: This manuscript explores leveraging the generative representation in the diffusion model for downstream discriminative purposes. Specifically, they propose self-conditioning, routing important semantic features obtained by probing into the decoder layer of the diffusion networks, along with the timestep condition. On top of that, the authors follow EDM and MoCo v3 for data augmentation and contrastive losses. Experimental results demonstrate that the proposed technique enables improved discriminative accuracy while maintaining the FID over the self-reported baselines on standard image benchmarks.
weaknesses: **Experimental Evaluation**: The reported generative metric, i.e., FID, in this paper, is worse than expected. For example, the official EDM [1] on CIFAR-10-uncond is 1.97, which is much better than the reported 2.23 in the baseline in Table 2. While I understand that re-training the model from scratch could be computationally intensive, the mismatched generative performance makes it hard to judge the efficacy of the proposed method on a well-trained diffusion model. I strongly suggest considering adding fine-tuning experiments on top of a pre-trained diffusion model with the proposed architectural change, and then examining whether there is a performance gain in the discriminative capacity while maintaining the original generative ability.

**Clarity**: There are multiple inappropriate, vague statements in the paper, which should be revised:
   * While determining the intermediate layers is a crucial procedure, the description on lines 203 -- 208 is not informative. For example, it is unclear what "short, parallel training runs" are. There should be at least a reference to the details of this experiment.
   * On lines 247--249, "a vector of transformation parameters is used..." is not informative. The augmentation used should be explicitly specified, either using a reference (if you are following [1] in the exact same way) or in the paper. 
   * On lines 294 -- 295, "after stabilizing some hyperparameters in feature evaluation" is also not informative.

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
questions: 1. In table 1, the authors demonstrate DDAE++ have sample quality++ and feature quality++ compared to DDAE. Despite performance enhancement, what are the fundamental  improvements of DDAE++? 

2. Did the augmented SSL loss introduce any noticeable challenges in training stability or convergence speed compared to the original DDAE objective? Do the authors need special techniques (e.g., warmer learning rates, different schedules) to stabilize the combined generative and discriminative training?

3. Can the authors visualize or quantitatively measure the semantic content or feature diversity of the latent space (e.g., through Singular Value Decomposition or dimensionality) at the self-conditioning point before and after applying DDAE++?

4. What is the sensitivity to the choice of the internal layer used for extracting the semantic token (CLS/AUG)? Is there a theoretical or empirical justification for the chosen layer depth?

5. Why can SSL improve the generation capabilities of diffusion models. Have the authors explored scenarios where the improved discriminative performance (SSL) comes at a noticeable detriment to generative quality (FID)?

6. Please provide a quantitative comparison of DDAE++'s generative performance (FID) and representation learning performance (Linear Probing and Transfer Learning) against REPA.

7. Could the authors comment on the scalability and expected performance of DDAE++ when applied to a modern, large-scale Text-to-Image (T2I) diffusion model (e.g., Stable Diffusion)? Do the authors anticipate needing further modifications to the self-conditioning mechanism to handle the complexity of the cross-attention layers in T2I models?
soundness: 2
strengths: 1. Unified and Principled Approach: The paper successfully integrates generative and discriminative objectives within a single, coherent diffusion framework. This addresses a critical gap by turning the denoising process into a powerful representation learner, maximizing the utility of the pre-training stage.

2. Architectural Efficiency: The core mechanism is lightweight self-conditioning, which does not require significant changes to the base U-Net structure or a substantial increase in computational complexity, making it practical for adoption.

3. Strong Representation Learning Results: The performance of the learned representations, particularly demonstrated through improved linear probing and transfer learning results (e.g., on ImageNet, CIFAR-100), clearly establishes DDAE++ as a state-of-the-art method among diffusion-based representation learners.
summary: The paper introduces DDAE++, an enhanced architecture and training framework for Diffusion Models (DMs) aimed at achieving unified generative and discriminative learning. The core problem addressed is the sub-optimal utilization of rich semantic features within the standard U-Net architecture of DMs, which hinders the quality of the learned representations.

DDAE++ proposes two main architectural and training modifications:

1. Self-Conditioning Mechanism: It introduces a lightweight self-conditioning mechanism that explicitly routes rich, high-level semantic features (extracted from intermediate layers of the U-Net) directly into the denoising process.

2. Augmented Self-Supervised Learning (SSL) Objective: It trains the model with an auxiliary SSL objective using contrastive learning, guided by different views (data augmentations) of the same instance.

Through these modifications, DDAE++ claims to establish a stronger representation bottleneck within the U-Net, leading to superior performance in both generative tasks and discriminative tasks.
weaknesses: 1. Lack of Component Analysis and Insight: While the effectiveness of the proposed components (e.g., the CLS token and the AUG token) is demonstrated via ablation, the paper lacks a deep analytical insight into why these mechanisms are effective from a feature flow perspective. The claim of establishing a "stronger representation bottleneck" remains largely descriptive rather than mechanistically proven. 

2. Marginal Generative and Discriminative Improvements: The quantitative improvements in generative quality (e.g., FID and Acc.) shown in Table 2 appear to be marginal over baselines, especially considering the added complexity of the SSL objective. This raises questions about the practical utility of DDAE++ specifically for users primarily focused on high-fidelity generation.

3. Incomplete Foundational Comparison (SSL): The paper claims to unify generation and self-supervised learning, yet it omits a crucial empirical comparison with highly relevant semantic-enhanced diffusion methods like REPA or other state-of-the-art masked modeling/denoising approaches like MaskDiT. 

4. Missing Evaluation on Complex Tasks (T2I): The evaluation is primarily conducted on unconditional and class-conditional image synthesis (CIFAR-10, ImageNet). A critical test for a modern diffusion enhancement would be its performance on the more complex, widely-used Text-to-Image (T2I) generation task. The current evaluation leaves uncertainty regarding DDAE++'s utility and scalability in the T2I domain, where high-level semantic flow is paramount.


Decision: Reject