Diffusion Models Need Visual Priors for Image Generation

## Abstract

Abstract Conventional class-guided diffusion models generally succeed in generating images with correct semantic content, but often struggle with texture details.
This limitation stems from the usage of class priors, which only provide coarse and limited conditional information.
To address this issue, we propose Diffusion on Diffusion (DoD), an innovative multi-stage generation framework that first extracts visual priors from previously generated samples, then provides rich guidance for the diffusion model leveraging visual priors from the early stages of diffusion sampling.
Specifically, we introduce a latent embedding module that employs a compression-reconstruction approach to discard redundant detail information from the conditional samples in each stage, retaining only the semantic information for guidance. We evaluate DoD on the popular ImageNet- 256 × 256 256 256 256\times 256 256 × 256 dataset, reducing 7 × \times × training cost compared to SiT and DiT with even better performance in terms of the FID-50K score.
Our largest model DoD-XL achieves an FID-50K score of 1.83 with only 1 million training steps, which surpasses other state-of-the-art methods without bells and whistles during inference.

## 1 Introduction

Diffusion models have emerged as a paradigm-shifting approach in visual content generation employing an innovative process of iterative noise-to-data transformation. Trained to reverse a gradual noising process, these models leverage deep neural networks to generate high-quality new samples that faithfully represent the training data distribution. Diffusion models have surpassed previous state-of-the-art generative frameworks, such as GANs (Sauer et al., ; Goodfellow et al., ; ) and VAEs (Kingma & Welling, ), offering superior sample quality, improved training stability, and enhanced scalability. This superiority of diffusion models has led to their widespread adoption across a diverse range of conditional generation tasks, including class-guided image generation (Lu et al., ; Wang et al., ; Ma et al., ; Peebles & Xie, ), text-to-image generation (Esser et al., ; Podell et al., ), and image editing (Meng et al., ).

Conventionally, class-guided diffusion models generate images conditioned on learned class embeddings. While this embedding-based approach is widely adopted, such class prior can only provide coarse-grained conditional information for models, which is only able to distinguish different categories. The challenge of constructing detailed images from such limited priors has led to exceptionally long training cycles for current class-guided diffusion models.
For example, DiT (Peebles & Xie, ) and SiT (Ma et al., ) require up to 7 million training steps to achieve convergence.

In contrast, visual priors contain more geometric visual information.
Intuitively, visual priors should be closer to the target distribution of image generation.The efficacy of visual priors in enhancing image generation quality has been demonstrated in various domains, including super-resolution models (Yang et al., ; Ren et al., ) and SD-Edit (Meng et al., ). Inspired by this well-established methodology, our study seeks to explore the integration of visual priors into class-guided image generation models. To this end, we propose an innovative multi-stage diffusion sampling framework. The initial stage adheres to the conventional approach, utilizing fixed class embeddings as priors. However, the subsequent refinement stages employ the image generated in the previous stage as a visual prior to guide further image synthesis.
Our framework’s distinctive feature lies in the reuse of the same diffusion model across multiple stages, leading us to term this method “Diffusion on Diffusion” (DoD).

Figure: Figure 2: ImageNet generation with Diffusion on Diffusion (DoD). Left: DoD is parameter-efficient. The diameter of each circle indicates the model size. Our DoD-S and DoD-B models, despite being much smaller, are comparable to the XL variants of other diffusion models. Right: DoD is sampling-efficient. “-G” indicates the application of classifier-free guidance (CFG), and “-S” denotes the number of sampling steps for each stage of DoD. The starting and ending points of each blue line represent the two stages of DoD, where we only apply CFG in the second stage. With the same sampling steps, DoD achieves lower FID-50K score.
Refer to caption: extracted/5918559/figures/teaser_sota.png

The proposed Diffusion on Diffusion (DoD) framework introduces a recurrent approach, where each stage comprises a complete diffusion sampling procedure.
From the second stage, DoD extracts semantic information from the previous output as additional visual priors.
This mechanism provides rich semantic visual guidance during the early stages of diffusion sampling, facilitating the generation of higher-quality images.
By repeatedly leveraging the generation capabilities of the diffusion model, DoD not only enhances texture details but also refines the object-level geometric.
Although longer sampling steps in diffusion models typically provide more accurate approximations and potentially boost performance, the final performance is still constrained by the model capacity.
In contrast, our method, by incorporating visual priors in extended sampling steps, effectively improves both the sampling efficiency and generation quality.
As shown in Figure 2 (Right), the proposed paradigm yields more efficient sampling compared to simply increasing the sample steps in diffusion models.
Moreover, as illustrated in Figure 2 (Left), DoD is also parameter-efficient.
Unlike models such as SDXL (Podell et al., ) which employs a separate refiner model, DoD integrates both image generation and refinement using shared parameters, significantly reducing the model size and lowering the barrier to real-world applications.

Directly using the generated latents as conditions results in the model collapsing into a complete reconstruction of the input samples, thereby losing the desired refinement effect.
To overcome this issue, we introduce a Latent Embedding Module (LEM), a vision transformer, that compresses the conditional sample into a few low-dimensional vectors, thereby discarding redundant texture details while retaining the essential information.
Our investigation shows that the kept information primarily consists of semantics, which are aligned between the training dataset and the generated samples for a well-trained generation model.
The semantic alignment eliminates the need for collecting specialized refinement data or implementing complex training strategies, allowing DoD to be trained end-to-end on generation datasets.

DoD is built upon the state-of-the-art diffusion transformer, FiTv2 (Wang et al., ), and follows the latent diffusion model (LDM) (Rombach et al., ) training paradigm.
We conduct comprehensive experiments and strictly evaluate our proposed method on ImageNet-$256\times 256$ benchmark. Compared with the DiT (Peebles & Xie, ) and SiT (Ma et al., ) models, our DoD achieves even better performance in terms of FID, with fewer model parameters and computational complexity, consuming 7$\times$ less training cost. Meanwhile, our method outperforms the previous state-of-the-art methods when no bells and whistles were applied during inference, which achieves an FID-50K score of 1.83 with only 1 million training steps.

## 2 Related Work

### 2.1 Diffusions and Flows

Denoising Diffusion Probabilistic Models (DDPMs) (Ho et al., ; Saharia et al., ; Radford et al., ; Croitoru et al., ; Bond-Taylor et al., ) and score-based models (Hyvärinen & Dayan, ; Song et al., ) have demonstrated significant advancements in image generation tasks (Lu et al., ; Ling et al., ; Rombach et al., ; Saharia et al., ; Meng et al., ; Ramesh et al., ; Ruiz et al., ; Poole et al., ). The Denoising Diffusion Implicit Model (DDIM) (Song et al., ) introduced an accelerated sampling method, while Latent Diffusion Models (LDMs) (Rombach et al., ) set a new standard by applying deep generative models to reverse the noise process in the latent space using Variational Autoencoders (VAEs) (Kingma & Welling, ).
Flow models  (Liu et al., ; Albergo & Vanden-Eijnden, ; Lipman et al., ; Albergo et al., ) present an alternative approach by learning a neural ordinary differential equation (ODE) that transports between two distributions.
The rectified flow model (Liu et al., ) solves a nonlinear least squares optimization problem to learn mappings along straight paths, which represent the shortest distance between two points, leading to improved computational efficiency. In this work, we adopt the rectified flow as the noise scheduler to train our DoD models.

### 2.2 Diffusion Transformer

The Transformer model (Vaswani et al., ) has successfully replaced domain-specific architectures across various fields, including language (Brown et al., ; Chowdhery et al., ), vision (Dosovitskiy et al., ; Han et al., ), and multi-modal learning (Team et al., ). In the realm of visual perception research, numerous studies (Touvron et al., ; ; Liu et al., ; ) have focused on accelerating pretraining by utilizing fixed, low-resolution images. Transformers have also been applied in denoising diffusion probabilistic models (Ho et al., ) for image synthesis. DiT (Peebles & Xie, ), a pioneering work in this space, employs a vision transformer as the backbone for latent diffusion models (LDMs), serving as a strong baseline for subsequent research. MDT (Gao et al., ) introduces a masked latent modeling approach, requiring two forward passes during training and inference. U-ViT (Bao et al., ) tokenizes all inputs and integrates U-Net architectures into the ViT backbone of LDMs.
SiT (Ma et al., ), utilizing the same architecture as DiT, explores various rectified flow configurations. Large-DiT and Flag-DiT (Gao et al., ) scale up diffusion transformers to achieve improved performance. SD3 (Esser et al., ) introduces novel noise samplers for rectified flow models and scales these models to billions of parameters, yielding state-of-the-art text-to-image generation results. FiTv2 (Wang et al., ), based on FiT (Lu et al., ), achieves advanced class-conditional image generation performance by leveraging a flexible diffusion transformer architecture within a rectified flow framework. In this work, built upon the FiTv2 architecture, we propose Diffusion on Diffusion (DoD), an innovative framework which effectively incorporates visual priors into class-guided image generation.

## 3 Method

### 3.1 Preliminaries

Diffusion and Flow Models. Before introducing our Diffusion on Diffusion (DoD) framework, we provide a brief review of diffusion and flow models. Given the noise distribution $\epsilon\sim\mathcal{N}(0,\mathbf{I})$ and data distribution $x_{0}\sim p(x)$, the models use the time-dependent forward process:

$$ $x_{t}=\alpha_{t}x_{0}+\beta_{t}\epsilon,$ (1) $$

where $\alpha_{t}$ is a decreasing function of $t$ and $\beta_{t}$ is an increasing function of $t$.
In this unified perspective, diffusion models (Ho et al., ; Song et al., ; Song & Ermon, ; ) set $\alpha_{t}$ and $\beta_{t}$ based on stochastic differential equation (SDE) formulations, where DDPM (Ho et al., ) is equivalent to variance preserving SDE (VP-SDE) and SMLD (Song & Ermon, ; ) corresponds to variance exploding SDE (VE-SDE). DDIM (Song et al., ) sets $\alpha_{t}$ and $\beta_{t}$ through ordinary differential equations (ODE), which leads to fewer sampling steps but sacrifices the generation quality. Flow models (Liu et al., ; Lipman et al., ; Albergo & Vanden-Eijnden, ; Albergo et al., ) restrict the process 1 on $t\in[0,1]$, and set $\alpha_{0}=\beta_{1}=1,\alpha_{1}=\beta_{0}=0$, interpolating between the two distributions through ODE formulation.

Rectified flow (Liu et al., ) introduces an ODE model and transports between the data distribution and the noise distribution via a straight line path, which is the theoretically shortest route between two points.
Given empirical observations $X_{0}\sim p(x),X_{1}\sim\mathcal{N}(0,\mathbf{I})$, the forward process 1 is defined as: $X_{t}:=tX_{1}+(1-t)X_{0}$, which is the linear interpolation of $X_{0}$ and $X_{1}$.

The ODE model learns the drift at time $t\in[0,1]$:

$$ $\text{d}X_{t}=v(X_{t},t)\text{d}t,$ (2) $$

which converts data $X_{0}$ from $p(x)$ to noise $X_{1}$ from $\pi_{1}$, and the drift follows the direction of $(X_{1}-X_{0})$. In practice, a network $v_{\theta}$ is utilized to predict this drift, and the optimization target is:

$$ $\min_{v}\int_{0}^{1}\mathbb{E}\Big{[}||(X_{1}-X_{0})-v(X_{t},t)||^{2}\Big{]} \text{d}t$ (3) $$

Previous studies (Ma et al., ; Esser et al., ; Wang et al., ; Gao et al., ) have demonstrated the efficiency and stability of rectified flow models. In this work, we fully adhere to the form of rectified flow as our noise scheduler and enhance the diffusion model by adding extra conditions beyond the class label. We adopt the implementation of rectified flow following SiT (Ma et al., ) and use ODE sampler for image synthesis.

FiTv2.
DoD utilizes the state-of-the-art diffusion transformer, FiTv2 (Wang et al., ), as the backbone.
FiTv2 is an advanced diffusion transformer on class-guided image generation, evolving from SiT (Ma et al., ) and FiT (Lu et al., ). The key modules of FiTv2 include $2$-D Rotary Positional Embedding ($2$-D RoPE) (Su et al., ), Swish-Gated Linear Unit (SwiGLU) (Shazeer, ), Query-Key Vector Normalization (QK-Norm), and Adaptive Layer Normalization with Low-Rank Adaptation (AdaLN-LoRA) (Hu et al., ).

Despite these advanced modules, FiTv2 adopts the Logit-Normal sampling (Esser et al., ) strategy to accelerate the model convergence. This sampling strategy puts more attention on the middle part of the sampling process, as recent studies (Karras et al., ; Chen, ) have disclosed that the intermediate part is the most challenging part of the diffusion process.

Figure: Figure 3: Illustration of (a) Multi-Stage Sampling and (b) Latent Embedding Module. We only present the first two stages of DoD, while subsequent stages are derived from the second one.
Refer to caption: x2.png

### 3.2 Diffusion on Diffusion with Visual Priors

We introduce the Diffusion on Diffusion (DoD) framework, which enhances diffusion models by recurrently incorporating previously generated samples as visual priors to guide the subsequent sampling process.
There are two core components in DoD: the Multi-Stage Sampling strategy and the Latent Embedding Module (LEM).
The multi-stage sampling enables the use of visual priors, while the LEM extracts semantic information from the samples generated in the previous stage.

Multi-Stage Sampling.
The sampling process of DoD consists of multiple stages, with each stage being a complete diffusion sampling process conditioned on different information.
As depicted in Figure 3 (a), in the $i$-th stage, the backbone diffusion model, denoted as $F(\cdot)$, takes Gaussian noise $z_{1}^{i}\in\mathbb{R}^{H\times W\times d_{z}}$ as input and progressively denoises it to obtain the sample $z_{0}^{i}\in\mathbb{R}^{H\times W\times d_{z}}$, where $H$, $W$, and $d_{z}$ denote the height, width, and channels of the noised data, respectively. Let $c$ represent the conditional feature that guides the generation direction via the AdaLN-LoRA block, we have the sampling formulation as:

$$ $z_{\tau_{t-1}}^{i}=F(z_{\tau_{t}}^{i},c),t=T,T-1,\dots,1,$ (4) $$

where $T$ represents the total sampling steps, $\tau_{t}$ is a monotone increasing function of $t$ that satisfies $\tau_{t}\in[0,1],\tau_{0}=0,\tau_{T}=1$. The accurate form of $\{\tau_{t}\}_{1:T}$ is determined by the ODE sampler. The formulation of each stage closely follows that of FiTv2 (Wang et al., ).

Typically, for class-conditional generation, the conditional feature $c$ consists only of class and time information.
In DoD, the sample generated in the previous iteration is used as additional conditioning information, as follows:

$$ $\displaystyle c=\begin{cases}e_{c}+e_{t}+\mathbf{S}&\text{ if }i=1\\ e_{c}+e_{t}+e_{s}^{i-1}&\text{ if }i>1\end{cases},$ (5) $\displaystyle e_{s}^{i-1}=\text{LEM}(z_{0}^{i-1}),$ $$

where $e_{c}\in\mathbb{R}^{1\times d}$ and $e_{t}\in\mathbb{R}^{1\times d}$ are the embeddings for the label and time, and $e_{s}^{i-1}\in\mathbb{R}^{(\frac{H}{p}\times\frac{W}{p})\times d}$ represents the embedding extracted by the latent embedding module $\text{LEM}(\cdot)$, with $p$ and $d$ denoting the patch size and hidden size of the backbone, respectively.

To maintain consistency in behavior, in the first stage, we replace $e_{s}$ with a trainable sample token $\mathbf{S}\in\mathbb{R}^{1\times d}$.
Note that when $i>1$, broadcast is operated on $e_{c}$ and $e_{t}$ to align their dimensions with the sample embedding $e_{s}$. We modify the modulation function in the AdaLN-LoRA block to accept conditional features with different dimensions accordingly.

Latent Embedding Module (LEM).
Since we use previously generated latents as visual priors to guide the generation process, retaining all the information from these samples would cause the diffusion model to collapse into an identity mapping, resulting in a complete reconstruction of the conditional samples.
To avoid this, we propose the Latent Embedding Module (LEM) that filters the conditional information using a compression-reconstruction approach to discard redundant details.

The architecture of LEM is illustrated in Figure 3 (b), comprising an encoder and a decoder, both of which are standard vision transformers (Dosovitskiy et al., ).
The encoder first patchifies the input latent $z_{0}\in\mathbb{R}^{H\times W\times d_{z}}$ through a patch embedding layer into $e_{p}\in\mathbb{R}^{(\frac{H}{p}\times\frac{W}{p})\times d_{l}}$, where $d_{l}$ is the hidden size of the LEM, and the patch size $p$ is consistent with that of the backbone network.
Next, $N$ learnable latent tokens $\mathbf{L}\in\mathbb{R}^{N\times d_{l}}$ are concatenated with $e_{p}$ and fed into the vision transformer.
A linear layer then reduces the channel dimension to $d_{e}$. Only the latent tokens are used as the output of the encoder, formally:

$$ $\hat{\mathbf{L}}=Enc((e_{p}+\textbf{PE})\oplus\textbf{L}),$ (6) $$

where $\oplus$ denotes concatenation, PE is the frozen sine-cosine positional embedding, and $Enc(\cdot)$ indicates the encoder. $\hat{\mathbf{L}}\in\mathbb{R}^{N\times d_{e}}$ represents the output latent tokens, providing a compressed representation of the input latent $z_{0}$.

The decoder seeks to map the compressed latent tokens $\hat{\mathbf{L}}$ back to the shape of the patch embedding $e_{p}$ to enable element-wise conditioning.
Specifically, a sequence of mask tokens $\mathbf{M}\in\mathbb{R}^{(\frac{H}{p}\times\frac{W}{p})\times d_{e}}$, obtained by repeating a shared mask token, is concatenated with $\hat{\mathbf{L}}$.
The combined representation is then projected back to $d_{l}$ and fed into the vision transformer:

$$ $e_{s}=Dec((\textbf{M}+\textbf{PE})\oplus\hat{\mathbf{L}}),$ (7) $$

where $Dec(\cdot)$ denotes the decoder, which has the same number of heads and hidden size as the encoder. The output $e_{s}$ is the extracted latent condition.

Equipped with the Latent Embedding Module (LEM), DoD realizes image refinement effectively through latent space reconstruction.
The encoder of LEM projects the input latents into a more compact feature space, which is then restored by the decoder and the diffusion model.
We empirically find that the compressed latent tokens $\hat{\mathbf{L}}$ primarily retain semantic information, enabling the diffusion model to focus on texture details, thereby enhancing image fidelity, as in section 4.2.

### 3.3 Configurations

Table: Table 1: Details of DoD models. The layers of the latent embedding module are presented as encoder depth $+$ decoder depth. We list the number of parameters for the backbone and the latent embedding module in different colors.

Training.
DoD is trained within the latent space of a pre-trained Variational Autoencoder (VAE) from Stable Diffusion (Rombach et al., ).
The VAE encoder maps RGB images of shape $256\times 256\times 3$ to latents of shape $32\times 32\times 4$.
The new latents sampled by DoD are then decoded back to pixel space by the VAE decoder.

As mentioned above, the Latent Embedding Module (LEM) extracts semantic information from samples through compression and reconstruction to serve as visual priors.
We reasonably assume that the high-level semantic information extracted from generated images is similar to that obtained from real images.
This assumption allows us to use the latents of ground truth images as inputs to LDM during training, simplifying the training strategy. Such simplification allows end-to-end training of DoD on image latents and joint optimization of the backbone model and LEM.
Similar to training with classifier-free guidance, we randomly replace the output of the latent embedding module with a trainable sample token $\mathbf{S}$ at a probability of $p_{s}$, which is set to $0.5$ by default.

Sampling.
Classifier-free guidance (CFG) (Ho & Salimans, ) is well-known for enhancing generation quality and improving the alignment between conditions and generated images.
For visual priors, we employ a large CFG scale to ensure the consistency of the generated samples, while we do not use CFG in the first stage of DoD.
Additionally, DoD employs the adaptive-step ODE sampler (*i.e.*., dopri5) same as SiT (Ma et al., ) for sample synthesis in each stage.

Models.
We use three different model configurations with varying sizes, DoD-{S, B, XL}, as detailed in Table 1.
For the backbone network, *i.e.*, FiTv2, we closely follow the configurations outlined in their paper but reduce the model depth to match those of DiT (Peebles & Xie, ) and SiT (Ma et al., ).
For the latent embedding module, we use an 8-layer encoder and a 4-layer decoder by default.
Our largest model, DoD-XL, uses the same size latent embedding module as DoD-B to reduce the number of parameters.
The patch size for both the backbone network and the latent embedding module is set to $2$, and the channel dimension $d_{e}$ of the latent tokens in the latent embedding module is set to $16$.

## 4 Experiments

### 4.1 Experimental Setup

Implementation Details.
We conduct experiments on ImageNet (Deng et al., ), using a resolution of $256\times 256$.
All models share the same training strategy.
We use the AdamW (Kingma, ; Loshchilov & Hutter, ) optimizer with a constant learning rate of $1\times 10^{-4}$ and without weight decay.
Models are trained with a batch size of $256$.
Following SiT and FiTv2, we also apply an exponential moving average (EMA) with a decay factor of $0.999$ to the model weights during training, and we report the performance of the EMA checkpoints.

Evaluation Metrics.
We use Fréchet Inception Distance (FID) (Heusel et al., ) as the primary evaluation metric. Inception Score (IS) (Salimans et al., ), sFID (Nash et al., ), improved Precision and Recall (Kynkäänniemi et al., ) are also included to holistically evaluate the generation quality.
Moreover, the inference of DoD involves two processes – sampling with class priors and sampling with visual priors.
Since the latter process is formally similar to image reconstruction, we also adopt reconstruction-FID (rFID), PSNR, and SSIM to assess the reconstruction quality, providing a more comprehensive evaluation of our method.

### 4.2 Preliminary Experiments

Figure: Figure 4: Left: rFID score with different CFG scales. A large CFG scale is required in DoD to effectively leverage visual priors. Right: FID score comparison across different stages. Visual priors are not available in Stage 1, while both stages 2 & 3 utilize the samples from the previous stage as visual priors.
Refer to caption: extracted/5918559/figures/cfg_rfid.png

Visual Prior Drastically Improves Performance.
We first verify how effectively visual priors can guide image generation.

We use the latents of ground truth images as inputs for the latent embedding module and report the reconstruction-FID (rFID) of the generated images.
This experiment is conducted with a DoD-B model trained with $400K$ steps.
As shown in Figure 4 (Left), classifier-free guidance (CFG) is critical for DoD, with a high CFG scale leading to significant improvements. The higher the CFG scale is, the more the model relies on the input condition, which contains the visual prior.
When the CFG scale is set to $5.5$, the model performs best, achieving an rFID score of $3.6$, which is much better than the score without CFG ($10.8$ rFID).
This encouraging result demonstrates the potential of conditioning on visual priors.
In this work, unless otherwise stated, we default to using a CFG scale of $5.5$ when employing image priors as conditions for DoD.

We also present the generation results of DoD models trained for $400K$ steps in Figure 4 (Right).
The visual prior is significantly effective in improving the FID-50K score, since the performance of Stage 2 (w/ visual prior) is obviously better than the performance of Stage 1 (w/o visual prior).

Table: Table 2: Effects of the number of latent tokens $N$ in LEM. We provide a comprehensive evaluation that includes: generation in Stage 1 (S1) and Stage 2 (S2), image reconstruction, and linear probing. The results are based on DoD-B, trained for $400K$ steps.

Table: Table 3: Effects of the joint training strategy in DoD. We report the results based on DoD-B trained for $400K$ steps.

Extracting Visual Priors Through Compression.
In Table 2, we evaluate the effectiveness of the compression-reconstruction approach used by DoD to extract visual prior information from three perspectives: generation, reconstruction, and linear probing.
DoD controls the compression rate through the number of latent tokens $N$ in the latent embedding module, where using 256 tokens means no compression is applied.
Thanks to the generative capabilities of the diffusion model, we can reconstruct images using only $16$ latent tokens.
We draw the following conclusions:

- •
A lower compression rate (*i.e.*, using more tokens) results in better reconstruction performance. This is intuitive, as a lower compression rate means that more information is retained for image restoration.
- •
More semantically rich representations can be learned with fewer tokens. Drawing from self-supervised learning, we measure the quality of the representations using linear probing accuracy. The best performance is achieved with $16$ tokens, demonstrating that the latent embedding module in DoD effectively extracts global semantic information as visual priors through compression.

Interestingly, DoD with $32$ tokens performs best in generation during both the first and second stages, suggesting a trade-off between reconstruction capability and semantic richness.
By default, we utilize 32 latent tokens across all DoD variants.

Table: Table 4: Comparison between DoD and FiTv2. Both of the models are trained with 1.5M steps. We report the performance of DoD in Stage 1 (S1) and Stage 2 (S2) separately. With the same sampling steps and less computation, our DoD demonstrates better FID performance.

Table: Table 5: Benchmarking class-conditional image generation on ImageNet $256\times 256$. “-G” indicates results with classifier-free guidance, “S2” denotes results from the second stage of DoD.

Joint Training in DoD.
In DoD, we implement a joint training strategy that optimizes both the backbone diffusion model and the latent embedding module concurrently.
The training process is controlled by a hyperparameter $p_{s}$, which dictates the probability that the ground truth image priors are not utilized.
The results of generation and reconstruction are detailed in Table 3.
A higher value of $p_{s}$ leads to improved generation performance, while a lower value is more beneficial for reconstruction.
As this work aims to validate the importance of visual priors rather than pursuing state-of-the-art performance, we simply set $p_{s}=0.5$ as the default value.

Additionally, the joint training of DoD is based on the assumption that images of varying fidelity can provide similar visual priors through compression.
We validate this assumption in Figure 4 (Right), where the output from stage 1 is used as a visual prior in stage 2, yielding more realistic images.
However, stage 3, which relies on stage 2, does not show further improvements, indicating that the poor outputs from stage 1 and the good outputs from stage 2 provide similar prior information.
Therefore, in our experiments, we report the performance of DoD from stage 2.

Visual Prior Leads to Efficient Sampling in DoD.
One potential drawback of DoD is that it involves a multi-stage sampling process, which may result in a substantial computational overhead.
Therefore, in Table 4, we demonstrate that DoD outperforms the baseline model (i.e., FiTv2) even with fewer inference GFLOPs.
We employ the Euler sampler and fixed the number of sampling steps.
The total GFLOPs is calculated as GFLOPs multiplied by the number of sampling steps.
For methods utilizing classifier-free guidance (CFG), the total GFLOPs are doubled.
For DoD, we separately present the performance of two sampling stages, where the total GFLOPs for stage 2 are added to those of stage 1.
Our observations are as follows:

- •
Increasing the sampling steps from 60 to 240 reduces the FID score across all cases.
- •
Due to the use of a shallower backbone, DoD at Stage 1 performs worse than FiTv2. However, in Stage 2, with the same number of sampling steps and fewer GFLOPs, DoD fully surpasses the baseline FiTv2. This further demonstrates the efficacy of the visual prior.

This experiment highlights the efficiency of DoD and the effectiveness of leveraging visual priors.

### 4.3 Main Experiments

The comparison against state-of-the-art class-conditional generation methods is shown in Table 5.
To ensure fair comparisons, we use the total number of training images (denoted as “Images” in the table) as a measure of the training cost.
The total number of training images is calculated as $\text{training steps}\times\text{batch size}$.

DoD achieves better FID scores with lower training costs and much fewer parameters.
Specifically, using only $191M$ parameters, DoD-B achieves an FID score of $3.15$, surpassing both ADM ($3.94$ FID with $673M$ parameters) and LDM-4 ($3.60$ FID with $395M$ parameters).
When scaling up the model, DoD-XL further improves performance to a FID of $1.83$, using only $613M$ parameters and $1M$ training steps, outperforming all previous diffusion models. Qualitative results are shown in Figure 1 and Figure 5.

Figure: Figure 5: Qualitative Results. We present images generated by DoD-XL with $1M$ training steps. Across all stages, semantic information is preserved, and image quality improves progressively.
Refer to caption: extracted/5918559/figures/three_stage/1/1.png

## 5 Conclusion

In this work, we propose using visual priors to further enhance diffusion models. To achieve this, we introduce Diffusion on Diffusion (DoD), a novel multi-stage generation framework that provides rich guidance for the diffusion model by leveraging visual priors from previously generated samples.
The latent embedding module in DoD effectively extracts semantic information from generated samples through a compression-reconstruction approach.
DoD demonstrates remarkable training and sampling efficiency while being easy to reproduce.
We conduct extensive experiments to validate the potential of injecting visual priors and explore the design space of the model.
We hope our work will inspire the community to further investigate the use of visual priors in image generation.

## Appendix A Appendix

### A.1 Technologies adopted by FiTv2

FiTv2 Wang et al. () is an advanced diffusion transformer on class-guided image generation, evolving from SiT (Ma et al., ) and FiT (Lu et al., ). The key modules of FiTv2 include $2$-D Rotary Positional Embedding ($2$-D RoPE) (Su et al., ), Swish-Gated Linear Unit (SwiGLU) (Shazeer, ), Query-Key Vector Normalization (QK-Norm), and Adaptive Layer Normalization with Low-Rank Adaptation (AdaLN-LoRA) (Hu et al., ).

$2$-D RoPE. FiTv2 follows FiT and adopts $2$-D RoPE as its positional embedding. RoPE applys a rotary transformation to the embedding, incorporating both absolute and relative positional information into the query and key vectors. Benefiting from such property, RoPE and its high-dimensional variants have been widely adopted in current vision diffusion transformer models (Lu et al., ; Gao et al., ; Esser et al., ).

SwiGLU. SwiGLU is widely used in advanced language models like LLaMA (Touvron et al., ; ; Dubey et al., ). FiTv2 utilizes SwiGLU module as its Feed-forward Neural Network (FFN), rather than normal Multi-layer Perception (MLP). The SwiGLU is defined as:

$$ $\begin{split}\text{SwiGLU}(x,W,V)=\text{SiLU}(xW)\otimes(xV)\\ \text{FFN}(x)=\text{SwiGLU}(x,W_{1},W_{2})W_{3}\end{split}$ (8) $$

QK-Norm. FiTv2 applies LayerNorm (LN) to the Query (Q) and Key (K) vectors before the attention calculation. This technique effectively stabilizes the training process, particularly during the mixed-precision setting, as well as slightly improves the performance.
Formally, the attention is calculated as:

$$ $\text{Softmax}(\frac{1}{d_{k}}\text{LN}(Q_{i})LN(K_{i})^{T}).$ (9) $$

AdaLN-LoRA. FiTv2 utilizes AdaLN-LoRa to reduce the too many parameters occupied by the original AdaLN module in each transformer block. Besides, a global AdaLN module is employed to extract overlapping condition information and reduce the information redundancy in each block. Let $S^{i}=[\beta_{1}^{i},\beta_{2}^{i},\gamma_{1}^{i},\gamma_{2}^{i},\alpha_{1}^{i
},\alpha_{2}^{i}]\in\mathbb{R}^{6\times d}$ denote the tuple of all output scale and shift parameters, $\mathbf{c}\in\mathbb{R}^{d}$ and $\mathbf{t}\in\mathbb{R}^{d}$ represent the embedding for class and time step respectively. In FiTv2 blocks, $S^{i}$ is calculated as:

$$ $\displaystyle S^{i}$ $\displaystyle=\text{AdaLN}_{\text{global}}(\mathbf{c}+\mathbf{t})+\text{AdaLN} _{\text{LoRA}}(\mathbf{c}+\mathbf{t})$ (10) $\displaystyle=W^{\text{g}}(\mathbf{c}+\mathbf{t})+W_{2}^{i}W_{1}^{i}(\mathbf{c }+\mathbf{t}),$ $$

where $W^{\text{g}}\in\mathbb{R}^{(6\times d)\times d},W_{2}^{i}\in\mathbb{R}^{(6
\times d)\times r},W_{1}^{i}\in\mathbb{R}^{r\times d}$, and the bias parameters are omitted for simplicity.

Logit-Normal Sampling. Despite these advanced modules, FiTv2 adopts the Logit-Normal sampling (Esser et al., ) strategy to accelerate the model convergence. Normally, rectified flow models sample times uniformly from the $[0,1]$ interval, which means each part of the noise scheduler is trained equally. FiTv2 samples timesepts from a logit-normal distribution (Atchison & Shen, ), which is formulated as:

$$ $u\sim\mathcal{N}(\mathbf{0},\mathbf{1}),\quad t=\log(\frac{u}{1-u})$ (11) $$

where $\mathcal{N}(\mathbf{0},\mathbf{1})$ denotes the standard normal distribution. This sampling strategy puts more attention on the middle part of the sampling process, as recent studies (Karras et al., ; Chen, ) have disclosed that the intermediate part is the most challenging part in diffusion process.

### A.2 More Qualitative Results

We present additional qualitative results generated by DoD-B and DoD-XL, displayed in Figure 6 and Figure 7, respectively.

Figure: Figure 6: Images generated by DoD-B.
Refer to caption: extracted/5918559/figures/viz/B/1/1.png

Figure: Figure 7: Images generated by DoD-XL.
Refer to caption: extracted/5918559/figures/viz/XL/15/1.png

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: How stable are second and subsequent generation stage results given conditioning input?  It might be useful to visualize the variance of second- or third-stage generated samples given the same conditioning information and different noise inputs.

What is being learned by the LEM?  High-level semantics, patch/texture characteristics?  Given how close subsequent stage outputs appear to the first stage, the latter seems more likely.  If that is the case, why not just use the previously generated image itself as conditioning?

In fact, an alternative baseline system might be one that uses image-based conditioning for refinement via diffusion.  Instead of taking a latent embedding as conditioning, such systems simply take an entire image as conditioning (alongside a noisy input).  Examples of prior work using such designs include diffusion-based super-resolution approaches (e.g., cascaded diffusion); the paper's task would be a special case where the goal is not to increase resolution, but simply denoise again at the same resolution.  How essential is the LEM as compared to direct image-based conditioning?
soundness: 3
strengths: Experiments demonstrate that running the trained network in a multi-stage fashion yields visual quality improvements in the generated images.  The image output by each stage appears to refine the details of the image produced by the previous stage, indicating that the latent embedding acts as a useful guide to the generation process.

Quality improvements (as measured by FID) are achieved even when comparing to baselines at comparable inference cost (in FLOPs), as shown in Table 4.
summary: This paper proposes a multi-stage image generation process using a diffusion model which can optionally accept a latent embedding as conditioning information.  To synthesize an image from noise, this diffusion model is first run for an entire sampling chain without any conditional input.  The resulting synthesized image is then fed to a latent embedding module (LEM), implemented via a vision transformer, which outputs conditioning tokens.  The diffusion synthesis process is then restarted from pure noise, but with this latent embedding available as a conditioning input at each denoising step.  The embedding provides prior information about the desired target output, allowing the second synthesis process to produce a refined image that is consistent with the latent representation of the previously generated image.
weaknesses: Figure 4 suggests that beyond a two-stage system (one unconditional and one conditional generation pass), there is minimal benefit to subsequent conditional generation stages; stage 2 and stage 3 FID scores seem nearly identical.

Linear probing scores (Table 2) suggest that the LEM is learning only a limited semantic representation.  For example, a pre-trained contrastive encoder would score much better on ImageNet linear probing.

I am concerned about the novelty of the contribution.  The overall idea appears quite similar to cascaded diffusion [Ho et al., Cascaded Diffusion Models for High Fidelity Image Generation, JMLR 2022], which is not cited.  The difference seems to hinge on whether conditioning information is an image or a lower-dimensional embedding of an image.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: In Table 4, is there a reason why only FiTv2 is included in the sampling compute comparison? In Table 5, other models appear to have relatively better scores in terms of the evaluation metric. I would be interested to know if your method has been compared with them in terms of the computational cost. Perhaps Table 4 & 5 could be combined?

Additionally, the calculation of sampling GFLOPs is somewhat confusing. In Line 437, the calculation appears to be based on FiTv2-B-G with 60 sampling steps * 2 (due to DFG) * 27.3 GFLOPs per forward pass. However, in Line 439, you compare this to DoD-B-G with only 30 sampling steps instead of 60. Is this a typo? Also, if you calculate the total steps in CFG (e.g., 60 = 30+30) for DoD-B-G, it would be clearer to use the total steps (i.e., 120 = 60 + 60) for FiTv2-B-G to avoid potential reader confusion.

In Table 5, is there a reason why all other methods use cfg = 1.5, whereas your method uses cfg = 5.5? cfg = 1.5 seems like a very low value to me.
soundness: 2
strengths: Overall, the image quality in the paper seems fine with good details.
Overall presentation of the paper is clear.
summary: The authors introduce Diffusion on Diffusion (DoD), a method that extracts visual priors from samples generated in the initial sampling stage and uses them as guidance for later stages through an embedding module that stores semantic information.
weaknesses: The method includes multi-stage sampling, which makes the already slow sampling process even more time-consuming and computationally inefficient. However, the experimental results and design that demonstrate your method's computational efficiency are somewhat confusing and unconvincing to me. I have addressed some of these points in the next section.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: (1) The authors only present results for the conditional generation task; how does the model perform in unconditional generation?
(2) What is the performance of S1 and S2 models without classifier-free guidance, and how does this compare to the baseline FiTv2?
(3) In Table 4, why does the S1 model perform significantly worse than FiTv2? This might be due to the reduced steps; would it be possible to evaluate both models with the same number of steps?
(4) Paper [1] also uses generated results as conditional input to the diffusion model. How does this approach compare to the proposed methods?


[1] Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
soundness: 2
strengths: (1) The proposed method enhances image generation quality compared to the baseline.
(2) The authors conduct a comprehensive experimental analysis of model configurations for conditional image generation tasks, including visualizations for qualitative comparisons.
(3) The multi-stage generative model utilizes a shared backbone, efficiently reducing the total parameter count as additional stages are added.
summary: This paper introduces a multi-stage generative model in which the initial stage employs a standard diffusion model, followed by a conditional diffusion model in subsequent stages. The output of the prior stage serves as the conditional input for the following stage, processed as a latent vector to retain abstract semantic information while minimizing low-level detail. In conditional image generation experiments on the ImageNet-256x256 dataset, the two-stage approach achieves a lower FID score than the single-stage model, though performance plateaus after two stages. Additionally, the authors assess model scalability by experimenting with different model sizes.
weaknesses: (1) As shown in Figure 5, the images generated across multiple stages appear nearly identical, making it difficult to visually assess the advantages of the multi-stage system.
(2) Although the multi-stage system incorporates an efficient shared-parameter design, its FLOPs increase with each additional stage. Performance seems to plateau after two stages, raising questions about the scalability of adding more stages.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 4
questions: In relation to my described weaknesses, the top areas that I see the most important for a strong submission and would need to be improved for me to raise my rating:

1) provide experiments against SOTA comparable models, meaning a) other post-hoc diffusion methods for DoD and b) other image feature conditioning methods for LEM. See W1 and W2 for specific suggestions of comparison methods. I am not particularly tied to these exact methods, but just giving them as a starting place to further understand what kinds of methods I mean.

2) Include better analysis within the main results (table 5), including a) more in-depth explanations in writing, b) better justification for why DoD is useful even though it is not the top-performing model on 3 / 5 metrics (without even including SOTA, as described in Q1)

Given convincing experiments, I would be included to greatly raise my rating, as I like the paper from a theoretical level.
soundness: 1
strengths: 1) From a theoretical level, I find the paper very nice. It is a) simple / elegant and b) fast. From a novelty perspective, I have no concerns.

2) The paper includes multiple strong results concerning the practical usefulness of the method (num parameters, steps, GFLOPS)

3) This paper is well-written, and well-presented
summary: Authors propose to use previous-stage generations as conditioning, thus re-generating improved images. They show via FID and linear probing that the generated images are higher quality than those generated with a single step. They also do extensive experiments to show that DoD can achieve lower FID faster and with less parameters.
weaknesses: Points are in order of my perceived importance (most to least), indicating how heavily they weigh in my rating.

1) Table 5 (main results) seem to primarily compare against basic diffusion models, but miss comparisons against other post-hoc diffusion methods (e.g. MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer, Gao 2023 already cited in this paper; ReNO: Enhancing One-step Text-to-Image Models Through Reward-based Noise Optimization, Eyring 2024; ElasticDiffusion: Training-free Arbitrary Size Image Generation through Global-Local Content Separation, Haji-Ali 2023). These should be the true competing methods.

2) LEM is a novel contribution, correct? There are other methods that use image features as conditioning (e.g. Diffusion Feedback Helps CLIP See Better, Wang 2024) and address the issue of making the task too easy if too much information is used--LEM should be justified by comparing to other methods of this type, both theoretically and quantitatively.

3) In Table 5 (main results), DoD is only the best on 3 / 5 metrics. This in itself is not convincing to me--however, I would be more convinced if other results were provided to show that DoD is much more efficient (e.g. the GFLOP or num parameter experiments) than the better performing methods, for relatively little performance drop. At the moment, these efficiency experiments are not done on the out-performing methods (StyleGAN-XL, BigGAN-deep, LDM-4-G, Flag-DiT-G).

4) Top-1 accuracy is included in Table 2--I believe it would be useful in the main results in Table 5 as well, as it is another way to show image generation quality.

5) The main results comparing to other SOTA (table 5) are very hidden, only at the end of the experimental section. They are also very slim on analysis. These need to be brought out more. 

6) I have a hard time interpreting the qualitative results, as all images seem very similar. It would be helpful to provide the reader with some specific guidance as to what you perceive as the qualitative improvements.