1 Introduction↩︎

Denoising Diffusion models have achieved notable success in high-dimensional data generation and rapidly applied to fields such as images, videos and 3D. Recently, the introduction of the Transformer architecture has further enhanced the scalability of the diffusion model in these fields [1]–[3], demonstrating its outstanding performance. As the scale of the diffusion model continues to increase, how to efficiently achieve excellent generation results has become a highly focused research topic.

The training objective of the diffusion model is to cover the entire (conditional) data distribution. The model is severely penalized for failing to cover low-probability areas, but it lacks sufficient training and data to learn to generate high-quality images corresponding to these areas. Several studies introduce additional guidance strategies during the sampling stage to encourage the generative process to focus on high-probability areas of the data distribution. Among them, classifier-free guidance (CFG) [4] leads to a better alignment between the prompt and the result by avoiding the unconditional results. Nevertheless, an excessively high CFG coefficient may push the sampling trajectory beyond the expected range of the conditional distribution, resulting in over-simplified or distorted samples [5]. An alternative line of research attempts to guide a well-trained model using a bad version of itself [6] instead of relying on CFG. By avoiding the results of the corresponding degraded version of the generative model during the sampling stage, the generation quality is improved while avoiding overly simplified results. Although promising, these approaches still require carefully designed degradation strategies [7], [8], extra training [6] or additional sampling steps [9], [10], thereby limiting their potential for the large-scale application.

To address this, we propose a simple yet effective sampling guidance strategy, named Internal Guidance (IG), which utilizes the intermediate-layer outputs within the Diffusion Transformer to enhance generation quality. Specifically, we introduce an auxiliary supervisory signal at an intermediate layer during training process, enabling the model to produce a weaker generative output from mid-level representations. Building upon this, during the sampling stage, we leverages the relationship between intermediate and deep outputs during sampling to extrapolate. This internal guidance achieves an Autoguidance-like [6] effect that maintains diversity and improves the generation quality but without any additional sampling steps. Moreover, we show that IG can also combine with the guidance interval [5] or CFG to further enhance generative performance. Besides providing intermediate results for guidance, we observe that the introduction of intermediate supervision alone can alleviate partial vanishing gradients. It’s convergence performance comparable to, or even better than, those achieved with self-supervised regularization applied to intermediate layers. Further leveraging the insight in IG, we discuss how modeling the discrepancy between intermediate and deep outputs within the loss function can inspire the design of new training acceleration strategies.

We conduct comprehensive experiments to evaluate the effect of our proposed internal guidance. Our IG brings significant performance improvements to both DiTs, SiTs and LightningDiT. Ultimately, we compared our generative performance with the current state-of-the-art models on SiT-XL/2 and LightningDiT-XL/1 with IG. For SiT-XL/2 training, we show the model achieves FID = 5.31 and FID = 1.75 on class-conditional ImageNet generation only using 80 and 800 training epochs (without CFG) which already exceeds the FID of the vanilla SiT-XL at 1400 epoch and REPA at 800 epochs that uses pre-trained representation model to align. Moreover, with classifier-free guidance and guidance interval, our scheme shows an improved FID at 1.46 at 800 epochs. And for LightningDiT-XL/1 model which uses self-supervised representation-based tokenizer, it achieves FID = 2.42 using 60 epochs and further achieves FID = 1.34 at 680 epochs, which achieves a large margin between all of these methods. Combined with classifier-free guidance and guidance interval, our scheme achieves state-of-the-art results of FID = 1.19 at 680 epochs.

2 sec:Related32Work↩︎

We discuss with the most relevant literature here and provide a more discussion in Supplementary Material.

Sampling Guidance in Diffusion Models. The sampling guidance after the training has become the key to improving the quality of the visual generation. Classifier-free guidance (CFG) [4] has become the standard method which forms a guided score by extrapolating from an unconditional one. While highly effective, CFG will cause the problem of reduced generation diversity when its guiding coefficient becomes larger. Several studies [5], [11]–[14] have reduced the downsides of CFG by making the guidance weight noise level-dependent. Another research direction separates the class guidance part from the CFG and only uses a weaker version of the current model for guidance. This approach improves image quality while maintaining diversity. While promising, these methods are limited by the complex design of degradation strategies [7], [8], extra training [6] or additional sampling steps [9], [10], which are not available for current large-scale image generators in practice. In this paper, our proposed internal guidance can be seamlessly integrated into the current advanced diffusion Transformer with almost no additional cost, demonstrating its potential for large-scale application.

Intermediate Representation Regularization for Diffusion Transformers. Many recent works have attempted to introduce intermediate representation regularization in diffusion transformer training. MaskDiT [15] and SD-DiT [16] combine MAE’s [17] and IBOT’s [18] training strategy into the DiT’s training process. TREAD [19] designs a token routing strategy to speed up Diffusion Transformer’s training. More recently, self-supervised learning has been introduced into diffusion training to further enhance training efficiency. These methods can be broadly categorized into: (1) aligning intermediate network representations with pre-trained self-supervised models [20]–[22], and (2) incorporating self-supervised representation learning as a regularization mechanism during training [23], [24]. In this paper, we employ the simplest approach of introducing an auxiliary supervision loss to alleviate the gradient vanishing problem in deep networks. This design bears conceptual similarity to DeepFlow [25], but it mainly focuses on introducing second-order ODEs in the training of the diffusion Transformer. Interestingly, we observed that the convergence performance of this simple method is comparable to that achieved by incorporating self-supervised learning as a form of regularization.

3 Methodology↩︎

3.1 Preliminaries↩︎

Our work is based on Flow-matching [26] and Diffusion [27], [2] and [28] have proposed a unified perspective to understand these two types of generative modeling methods. We first introduce the relevant preliminaries. The standard Flow-matching and Diffusion-based models both start from Gaussian noise \(\epsilon \sim \mathcal{N}(0, I)\) and gradually transform it into data samples \(\mathbf{x}_* \sim p(\mathbf{x})\) through stochastic processes. This process can be expressed as \[\mathbf{x}_t = \alpha_t \mathbf{x}_* + \sigma_t \epsilon,\] where \(\alpha_t\) and \(\sigma_t\) are decreasing and increasing functions of time \(t \in [0, T]\). The difference between the generative models based on Flow-matching and Diffusion lies in that the former typically interpolates noise and data within a finite time interval \(t \in [0, 1]\), while the latter defines a forward stochastic differential equation (SDE), which converges to a Gaussian distribution \(\mathcal{N}(0, I)\) as \(t \to \infty\). The denoising network \(D_\theta\) estimates the original data by minimizing the denoising diffusion objective: \[\mathbb{E}_{\mathbf{x}_* \sim p_{\mathbf{x}}} \mathbb{E}_{\epsilon \sim \mathcal{N}(0, \mathbf{I})} \| D_\theta(\mathbf{x}_t, t) - \mathbf{x}_0 \|^2 .\]

As the network size of generative models continues to increase, how to efficiently obtain high-quality generation results has become the focus of attention. Next, we will introduce our framework from two aspects: the loss function and the corresponding sampling guidance after training completion. Furthermore, we will discuss some related properties of the proposed sampling guidance, as well as the design of new training acceleration methods based on this.

3.2 Internal Guidance: Guiding a Diffusion Model with Its Own Internal Dynamics↩︎

Training a Diffusion Model with Intermediate Supervision. The difficulty of training deep networks has been explored in many aspects in various fields [29]–[32]. Recently, in the field of generation, people have introduced self-supervised representation learning as a regularization method for training diffusion models [23], [24]. In this paper, we attempt the simplest solution, which is to apply an auxiliary supervision in the intermediate layers of the network. Specifically, we additionally define an output layer after the intermediate part of the network. Accordingly, the intermediate and final losses \(\mathcal{L}_{\textit{inter}}\) and \(\mathcal{L}_{\textit{final}}\) are defined as \[\begin{cases} \mathcal{L}_{\textit{inter}} = \| D_i(\mathbf{x}_t, t) - \mathbf{x}_0 \|^2 ,\\ \mathcal{L}_{\textit{final}} = \| D_f(\mathbf{x}_t, t) - \mathbf{x}_0 \|^2 , \end{cases}\] where \(D_i(\cdot)\) and \(D_f(\cdot)\) denote the outputs of the intermediate and final layers, respectively. The overall framework is illustrated in Figure 2. For instance, for the training of a deep Diffusion Transformer, its final training loss becomes: \[\mathcal{L} = \mathcal{L}_{\textit{final}} + \lambda\mathcal{L}_{\textit{inter}},\] where \(\lambda > 0\) is a hyperparameter that controls the trade-off between the loss of the intermediate output and that of the final output. As shown in Table 1, We verified the effect of this loss on the SiT-B/2 model [2]. Although this method is very simple, it can surprisingly achieve results that are similar to or even better than those obtained by using the regularization method based on representation learning for complex designs [23], [24].

Sampling a Diffusion Model with Its Internal Outputs. The sampling stage is crucial for obtaining high-quality generated results. The standard classifier-free guidance (CFG) [6] sampling strategy can effectively improve the generation quality, but will cause an undesirable reduction in diversity.

An alternative line of research attempts to maintain diversity by guiding a well-trained model using a bad version of itself [6]. While promising, these approaches are still limited by carefully designed degradation strategies [7], [8], extra training [6] or additional sampling steps [9], [10]. By applying an auxiliary supervision to the intermediate layer of the denoising model in the training stage, we can obtain the output of the final layer and that of its intermediate layer during the sampling stage. This naturally amounts to creating a bad version corresponding to the current final layer’s output. In this way, we propose Internal Guidance (IG), the intermediate output \(D_i\) can be used to guide the final output \(D_f\) to shift away from its own poor distribution, thereby achieving better generation results. Similar to classifier free guidance (CFG), its guiding effect is achieved by extrapolating the two denoised results using a factor \(w\) \[\label{eq:ig} D_w(\mathbf{x};\mathbf{sec:c}) = D_i(\mathbf{x};\mathbf{sec:c}) + w (D_f(\mathbf{x};\mathbf{sec:c})-D_i(\mathbf{x};\mathbf{sec:c})).\tag{1}\]

Unlike the original Autoguidance [6] strategy, our proposed IG is a plug and play approach that no longer requires specific training on particular bad versions for specific experimental setups or additional sampling consumption. And this sampling method can also enhance the quality of image generation while maintaining diversity. Building upon this, we conduct an extensive discussion on the properties of IG.

4 Discussion↩︎

In this section, we conduct a detailed analysis of the specific impacts of IG from two perspectives: (1) the compatibility between IG and CFG and (2) Guidance Interval in IG.

4.1 The compatibility between IG and CFG↩︎

Although Internal Guidance (IG) can significantly improve the generation results without additional consumption during the sampling stage, its combination with Classifier-Free Guidance (CFG) can further improve the generation quality. To further elucidate the compatibility between our proposed IG and CFG, we conducted a series of visualization experiments to analyze at which levels these guidance strategies improve generative quality. We study a 2D toy example like [6] where a small-scale denoising network is trained to perform conditional diffusion in a synthetic dataset (Figure 3). The dataset is designed to exhibit low local dimensionality (i.e., highly anisotropic and narrow support) and a hierarchical emergence of local details upon noise removal. Both properties that can be expected from the actual manifold of realistic images [33], [34]. Details of the experimental setup are provided in the Supplementary Material.

As shown in Figure 3 (a), sampling from a diffusion model without guidance produces many outliers outside the data distribution. When using CFG, this guidance pulls samples toward higher values of \(\log\left[ p_{con}(\mathbf{x}\mid \mathbf{sec:c})/p_{uncon}(\mathbf{x}\mid \mathbf{sec:c}) \right]\) [6], where \(p_{con}\) and \(p_{uncon}\) are the conditional and unconditional distributions learned by the denoising diffusion model. As shown in Figure 3 (b), increasing the CFG coefficient shifts samples away from the opposite class and reduces outliers, but also suppresses branches near other classes, which enhances inter-class separability while reducing diversity. In contrast, Autoguidance [6] provides class-independent gradients that point inward towards the data manifold.The trajectories of these samples generated by this well trained model generally follow the local orientation and branching of the data manifold, pushing the samples deeper into the “good side” concentrates them at the manifold.

Furthermore, our proposed IG exhibits a performance similar to that of Autoguidance (Figure 3 (c)). When the IG coefficient increases, these samples produce a narrower distribution than the ground truth, but in practice this does not appear to have an adverse effect on the images. However, this reduced intra-class sample distance may introduce new outliers when a large Autoguidance or IG coefficient is applied before the model is fully trained. In this case, the stronger output fails to accurately point toward high-probability regions, rendering the guidance less effective. Because CFG provides class-conditional directionality, it complements IG by further suppressing such outliers. Combining the two achieves superior generative quality across training stages (Figure 3 (d)). In Section 5.2, we present a detailed quantitative validation to demonstrate this effect.

4.2 Guidance Interval in Internal Guidance↩︎

Meanwhile, the guidance interval [5] can also affect the performance of internal guidance. The original guidance interval encourages using large CFG coefficient for sampling within a continuous interval of noise levels, which is helpful for achieving better generation performance. For our proposed internal guidance, we also redefined the guidance function of Equation 1 by replacing \(\omega\) with a piecewise constant function: \[\begin{align} & D_w(\mathbf{x};\mathbf{sec:c}) = D_i(\mathbf{x};\mathbf{sec:c}) + w(\sigma) (D_f(\mathbf{x};\mathbf{sec:c})-D_i(\mathbf{x};\mathbf{sec:c})), \\ & \text{where } w(\sigma) = \begin{cases} w & \text{if } \sigma \in (\sigma_{low}, \sigma_{high}], \\ 1 & \text{otherwise,} \end{cases} \end{align}\] where, \(\sigma\) is the coefficient that controls the intensity of the noise. In the original guidance interval, applying CFG during periods of high noise range drastically reduces the variation in the results, basically leading them towards a handful of “template images” per class. Therefore, the guidance interval used in the CFG indicates that this guidance should only be used in the middle and low-noise ranges. For our proposed IG with the guidance interval, we observed the opposite phenomenon: the internal guidance does not need to be applied in the low-noise range, but applied in the high-noise and middle-noise ranges. In Section 5.2, we have provided a detailed presentation of the generated results after applying the guidance interval to the IG sampling.

5 Experiments↩︎

5.1 Setup↩︎

Implementation details. For experiments on DiT [1] and SiT [2], we strictly follow their original setup. For experiments on LightningDiT [35], we find using the recipe in LightningDiT leads to instability or abnormal training dynamics at early epochs and slow EMA model convergence at early epochs. To solve these issues, we instead use the Muon optimizer [36] instead of AdamW and change the EMA weight from 0.9999 to 0.9995. Other optimization hyperparameters are the same as vanilla LightningDiT. We use ImageNet-1K [37] for diffusion training, where each image is preprocessed to the resolution of 256×256 and follow ADM for other data preprocessing protocols. Each image is then encoded into a compressed vector \(\mathbf{z} \in \mathbb{R}^{32\times32\times4}\) using the Stable Diffusion VAE or \(\mathbf{z} \in \mathbb{R}^{16\times16\times32}\) using the VA-VAE. Detailed experimental details and hyperparameter settings, are provided in the Supplementary Material.

Evaluation protocol. We report Fréchet inception distance (FID) [38], sFID [39], inception score [40], precision(Pre.) and recall(Rec.) [41] using 50000 samples. Specifically, we uniformly randomly sample the 1,000 class labels 50,000 times and generate images accordingly SD, SiT, REPA for a fair comparison. Sampling follows original DiT, SiT and LightningDiT that we use the SDE Euler-Maruyama sampler with 250 steps for the experiments on DiT and SiT, and use the ODE Heun sampler with 125 steps for the experiments on LightningDiT. Full evaluation protocol details are also provided in the Supplementary Material.

5.2 Ablation Studies↩︎

The position of intermediate supervision. We begin by examining the effect of attaching the auxiliary supervision loss to different layers, which are shown in Table ¿tbl:tab:interloss?. We find that regularizing only the first few layers (e.g., 4 in SiT-B/2 [2] at 80 epochs) in training is effective. However, placing the auxiliary supervision loss in the latter half of the network or adding multiple auxiliary supervision losses in multiple intermediate layers cannot improve the convergence speed. We hypothesize that it interferes with the training of the deep layer’s output. Furthermore, we are surprised to find that the convergence effect with the auxiliary supervision loss is comparable to that of methods incorporate self-supervised representation learning regularization (Table 1). In large-scale experiments, we apply the auxiliary supervision loss and internal guidance to the 8-th layer of the large-scale diffusion transformer model.

Compatibility between IG and CFG. We also demonstrate that the proposed IG combined with CFG can further enhance the generation quality in the quantitative experiment. We apply IG on the fourth layer of SiT-B/2 to conduct the subsequent experiments. We conducted a rough search for the optimal combination rules of IG and CFG. As shown in Figure 4, after SiT-B/2 combines with our IG and then applies CFG, the best generated metric FID can be reduced from the original optimal 7.30 with only CFG to 6.50. Meanwhile, we find that when combining CFG, using a lower IG coefficient can yield better generation results, which is different from the optimal coefficient when only adding IG alone (Table 2 and Figure 7). And this process basically does not involve significant computational costs, which indicates that the IG can serve as an easily integrable component for modern generative models to enhance the generation quality.

Guidance interval in internal guidance. Next, we examine the role of guidance interval in internal guidance. We also apply IG on the fourth layer of SiT-B/2 to conduct experiments. In the Table 2, we can see that when no guidance interval is used, the optimal FID is approximately around the IG coefficient \(\omega\) of 1.9. After increasing the IG coefficient, we find it actually has a positive effect on the generation quality when the low-noise range does not apply IG. This is somewhat contrary to the conclusion of guidance interval used in CFG [5]. In our rough search, the best generative results are achieved when the IG coefficient is 2.3 and the guidance interval range is \([0.3, 1)\).

Scalability. We investigate the scalability of the auxiliary supervision loss with internal guidance by varying the model sizes of the diffusion transformers. In general, the effect of achieving efficient generation results of IG becomes more significant as the diffusion transformer model increases in size. We demonstrate this by plotting FID-50K of different SiT and DiT models with IG and without CFG in Figure: IG achieves the same FID level more quickly with larger models. To facilitate comparison, we set the IG coefficient to 2.3 and the guidance interval range is \([0.3, 1)\) in all the DiTs and SiTs models. And the IG is applied in the fourth layer of the B/2 scale model and in the eighth layer of other scales for the experiments.

Effect of \(\lambda\). Finally, we examine the effect of the regularization coefficient \(\lambda\) by training SiT-B/2 models for 400K with different coefficients 0.25 to 1.0 and comparing the performance. As shown in Table 3, the performance becomes stable after \(\lambda<= 0.5\).

5.3 Comparsion with State-of-the-art Methods↩︎

8pt

We use several recent diffusion-based generation methods as baselines, each employing different inputs and network architectures. Specifically, we consider the following three types of approaches: (a) Pixel Diffusion: ADM [42], RIN [43], PixelFlow [44], PixNerd [45] and SiD2 [46], (2) Vanilla Latent Diffusion: DiT [1], MaskDiT [15], SiT [2], TREAD [19], MDTv2 [47], (c) Latent Diffusion with Self-supervised Representation Model: REPA [20], REPA-E [21] and REG [22], LightningDiT [35] and DiT\(^{\text{DH}}\) [48]. We initialized our models on SiT and LightningDiT respectively to separately verify the effectiveness of the IG in the latent diffusion model with original VAE and aligned VAE that represents the latent space using the pretrained self-supervised representation model (i.e., DINOv2-B [49]).

We compare the FID values between SiT-XL/2 and LightningDiT-XL/1 with IG with other state-of-the-art methods without CFG. As shown in Table ¿tbl:tab:comparison95perf?, IG shows consistent and significant improvement across all model variants under random sampling. We provide balanced sampling results in the Supplementary Material. In particular, on SiT-XL/2, using IG (1.8) leads to FID=5.31 at 80 epochs, which already exceeds the FID of the vanilla SiT-XL at 1400 epoch and REPA at 800 epochs which uses pre-trained representation model to align. Note that the performance continues to improve with longer training; for instance, with SiT-XL/2, FID becomes 1.75 at 800 epochs. On LightningDiT-XL/1, using IG (1.7) leads to FID=2.42 at 60 epochs, which also already exceeds the FID of the vanilla LightningDiT-XL/1 at 64 epochs. For longer training, the FID becomes 1.34 at 680 epochs with LightningDiT-XL/1, which achieve a large margin between all of these methods.

Finally, we provide a quantitative comparison between SiT-XL/2 or LightningDiT-XL/1 with IG and other recent diffusion model methods using CFG. At 800 epochs, SiT-XL/2 + IG (1.8) achieves FID of 1.46 with a CFG scale of \(w=1.35\) and an extra guidance interval. Meanwhile, LightningDiT/1 + IG (1.4) achieves state-of-the-art FID of 1.19 with a CFG scale of \(w=1.45\) and an extra guidance interval at 680 epochs. We provide visual results of LightningDiT-XL/2 + IG in Figure 1. Moreover, we also provide more uncurated examples and experimental results on ImageNet 512×512 in the Supplementary Material, we show that IG also provides significant improvements in such setup.

6 Further Discussion on Training Acceleration↩︎

After conducting the above analysis, we naturally came up with a new question: Can the effect of this IG be replicated during the training process, so as to achieve the goal of training acceleration? We also demonstrated the denoising diffusion model training process through a 2D toy example. As shown in Figure 6 (a) and (b), comprehensive training can significantly reduce many outliers near the branches compared to incomplete training. With our proposed IG, since its guiding direction points towards the overall direction of the stronger model, outliers near the branches can also be effectively reduced (Figure 6c). Based on the above observations, we incorporate the guiding effect of IG into the training process of the denoising diffusion model. Specifically, we integrate the gradient vectors obtained by calculating the differences between the final and intermediate outputs during the sampling process into the loss function, enabling the model to learn the guidance signal during training. Similar to [50], we rewrite the training objectives for the intermediate and final layers during the training process \[\begin{cases} \text{Intermediate objective: } \mathbf{x}_0, \\ \text{Final objective: } \mathbf{x}_0 + \textcolor{purple}{\omega \cdot \text{sg}(D_f(\mathbf{x}_t) - D_i(\mathbf{x}_t))}. \end{cases}\] Where \(\omega > 0\) is a hyperparameter that controls the strength of the guiding signal for the final layer’s training objective. In practice, we utilize the EMA model that updates along with the model to obtain the outputs at both the intermediate and final layers during the training process and set \(\omega = 0.5\) following [50]. As shown in Figure 6 (d), the outliers of all branches have been significantly reduced after modifying the training objective. In the subsequent SiT-B/2 experiment, our designed training loss significantly outperformed the REPA [20] using pre-trained self-supervised representation models after 400K iterations, demonstrating performance comparable to that of IG. Due to the flexibility and tunability of IG in generating images during the reasoning process, we suggest continuing to use IG in actual applications.

7 Conclusion↩︎

In this paper, we have presented Internal Guidance (IG), a simple but effective guidance strategy for enhancing both generation quality and training efficiency. We introduce an auxiliary supervision loss in the deep Diffusion Transformer and utilize the output of its intermediate layers to guide the output results at the deep level during the sampling stage. This plug-and-play method has significantly improved the quality of image generation. We further analyze the properties of each part of our proposed IG, including its compatibility with CFG, its own guidance interval and discuss the design of new training acceleration methods based on its insight. Our IG can significantly improve the generation performance of various scale diffusion transformers with less computational resources. Ultimately, we get the state-of-the-art FID = 1.19 result for image generation when using LightningDiT-XL/1+IG with additional CFG. We hope our work would facilitate the research on more applicable sampling methods for visual generation. Acknowledgement. This work was supported by National Natural Science Foundation of China (No. 62476051).

In this file, we provide more implementation and experimental details which are not included in the main text. In Section 8, we provide the details of all the evaluation metrics and further compare the generation performance under uniform sampling. In Section 9, we provide a detailed comparison of computational resource consumption during the sampling stage between the proposed IG and REPA. In Section 10, we provide a comparison of the generate results under different sampling guidance methods. In Section 11, we provide the experimental results of the proposed IG on the 512 \(\times\) 512 ImageNet dataset. In Section 12, we provide more hyperparameter settings and implementation details for Diffusion Transformer experiments of various scales. In Section 13, We provide hyperparameter settings and implementation details for the 2D toy example. In Section 14, We provide a more detailed discussion on the related work of our approach. In Section 15, we provide more uncurated generation results on ImageNet 256×256 from the SiT-XL+IG with CFG.

8pt

8 Evaluation Metric↩︎

We strictly follow the setup and use the same reference batches of ADM [42] for evaluation, following their official implementation. We use NVIDIA A6000 pro GPUs or 4090Ti GPUs for evaluation and enable tf32 precision for faster generation. In what follows, we explain the main concept of metrics that we used for the evaluation.

• FID [38] measures the feature distance between the distributions of real and generated images. It uses the Inception-v3 network [55] and computes distance based on an assumption that both feature distributions are multivariate gaussian distributions.

• sFID [39] proposes to compute FID with intermediate spatial features of the Inception-v3 network to capture the generated images’ spatial distribution.

• IS [40] also uses the Inception-v3 network but use logit for evaluation of the metric. Specifically, it measures a KL-divergence between the original label distribution and the distribution of logits after the softmax normalization.

• Precision and recall [41] are based on their classic definitions: the fraction of realistic images and the fraction of training data manifold covered by generated data.

As stated in [48], using class-balanced sampling can lead to better results compared to random sampling because it more closely approximates the true label distribution of the training set. As the FID value approaches a lower range, these subtle differences begin to have a greater impact. The vast majority of VQ-based generation methods [51]–[53] have adopt the balanced sampling. To conduct a comprehensive and fair comparison with more baselines, we also adopt balanced sampling to conduct a comprehensive comparison of the current state-of-the-art VQ and Diffusion-based methods, as shown in Table ¿tbl:tab:bc?.

9 Computational Cost Comparison↩︎

We compare the computational efficiency of SiT-XL/2+REPA [20] and SiT-XL+IG under the same model scale in Table ¿tbl:tab:flops?. IG introduces only a marginal increase in parameter count and FLOPs relative to SiT-XL, while maintaining neral identical latency. Despite the minimal computational overhead, IG yields substantial improvements in generation quality, achieving a relative reduction in FID, alongside an increase in IS. These results demonstrate that IG simultaneously improves generation quality and computational efficiency, highlighting its effectiveness as a general-purpose enhancement for generative models.

2pt

10 Comparison of Guidance Methods↩︎

r0.38

To further illustrate the advantages of our proposed IG, we compared it with the original CFG [4] and Autoguidance [6] methods. For CFG method, we use the vanilla LightningDiT-XL/1 model which is trained for 800 epochs and employed its official CFG coefficients. For the Autoguidance method, since it is difficult to determine the optimal bad version model, we initialized a 12-layer LightningDiT model and trained it for 10 epochs. We roughly searched for the optimal coefficients of Autoguidance and determined it to be 1.15. We also compare our IG with original Autoguidance on the vanilla EDM2-S [56] model on ImageNet 64 dataset with the same training setting, which is shown in Table 4. As shown in Table 5, our IG method outperforms both adding only CFG or only adding Autoguidance. Moreover, further combining CFG can achieve the current state-of-the-art in image generation, and has a significant improvement compared to other guidance methods.

3pt

11 512×512 ImageNet↩︎

To further validate IG’s effectiveness, we conduct experiments at 512×512 resolution following REPA’s protocol [20] with Muon optimizer [36]. The RGB images are processed through the VAE [57] to yield 64×64×3 latents. As demonstrated in Table, IG surpasses the performance of REPA trained for 200 epochs and SiT-XL/2 trained for 600 epochs in terms of FID at only 60 epochs, demonstrating its superior effectiveness.

12 Hyperparameter and More Implementation Details↩︎

12.0.0.1 More implementation details.

We implement our models based on the original DiT, SiT and LightningDiT implementation. To speed up training and save GPU memory, we use mixed-precision (fp16) with gradient clipping and FlashAttention [58] operation for attention computation. We also pre-compute compressed latent vectors from raw pixels via stable diffusion VAE and VA-VAE and use these latent vectors. We do not apply any data augmentation as discussed in MaskDiT and REPA. For stable diffusion VAE, we use stabilityai/sd-vae-ft-ema for encoding images to latent vectors and decoding latent vectors to images. For the output head, we use a Linear layer to map the feature to the output size.

12.0.0.2 The use of Muon

In our large-scale experiment with the LightningDiT-XL/1 model, we adopt the Muon optimizer [36]. We find that in the environment of PyTorch, when using the pre-trained self-supervised representation model (such as DINOv2-B [49]) to align the encoder of the VAE, there will be unstable training issues in the early stage of training. After careful investigation and experimentation, we discover that replacing the AdamW optimizer with Muon could effectively alleviate such issue. Compared to AdamW, the Muon optimizer has a certain acceleration convergence effect in the early stage of model training, but we observe that the gap between them gradually narrowed as the training time increased, reaching similar final convergence results. In the later stage of training, there is a small probability of encountering training instability issues again. Resuming the training with full precision can solve this problem.

12.0.0.3 Computing resources.

We use NVIDIA A6000 pro 96 GB GPUs for training largest model (LightningDiT-XL/1, SiT-XL/2); and uses NVIDIA 4090 24GB GPUs for training smaller model (L, B). When sampling, we use either NVIDIA A6000 pro 96 GB GPUs or NVIDIA RTX 4090 24GB GPUs to obtain the samples for evaluation.

13 Details of the 2D toy example↩︎

In this Section, we describe the construction of the 2D toy dataset used in the analysis of Section 4, as well as the associated model architecture, training setup, and sampling parameters.

13.0.0.1 Dataset.

For each of the two classes \(\mathbf{sec:c}\), we model the fractal-like data distribution as a mixture of Gaussians \(\mathcal{M}_\mathbf{sec:c} = \big( \{\phi_i\}, \{\mathbf{\mu}_i\}, \{\mathbf{\Sigma}_i\} \big)\), where \(\phi_i\), \(\mathbf{\mu}_i\), and \(\mathbf{\Sigma}_i\) represent the weight, mean, and \(2{\times}2\) covariance matrix of each component \(i\), respectively. This lets us calculate the ground truth scores and probability densities analytically and, consequently, to visualize them without making any additional assumptions. The probability density for a given class is given by

\[p_{\text{data}}(\mathbf{x}\mid \mathbf{sec:c}) = \sum_{i \in \mathcal{M}_c} \phi_i \, \mathcal{N}(\mathbf{x}; \mu_i, \Sigma_i), \quad\text{where}\]

\[\mathcal{N}(\mathbf{x}; \mu, \Sigma) = \frac{1}{\sqrt{(2\pi)^2 \det(\Sigma)}} \exp\left( -\frac{1}{2}(\mathbf{x}-\mu)^{\top} \Sigma^{-1} (\mathbf{x}-\mu) \right).\]

Applying heat diffusion to \(p_\text{data}(\mathbf{x} | \mathbf{sec:c})\), we obtain a sequence of increasingly smoothed densities \(p(\mathbf{x} | \mathbf{sec:c}; \sigma)\) parameterized by noise level \(\sigma\): \[\label{eq:GT} p(\mathbf{x}\mid \mathbf{sec:c}; \sigma) = \sum_{i \in \mathcal{M}_c} \phi_i \, \mathcal{N}(\mathbf{x}; \mu_i, \Sigma_{i,\sigma}^*), \quad\text{where } \Sigma_{i,\sigma}^* = \Sigma_i + \sigma^2 \mathbf{I}.\tag{2}\] The score function of \(p(\mathbf{x} | \mathbf{sec:c}; \sigma)\) is then given by \[\nabla_{\mathbf{x}} \log p(\mathbf{x}\mid \mathbf{sec:c}; \sigma) = \frac{ \sum_{i \in \mathcal{M}_c} \phi_i \, \mathcal{N}(\mathbf{x}; \mu_i, \Sigma_{i,\sigma}^*) \, (\Sigma_{i,\sigma}^*)^{-1} (\mu_i - \mathbf{x}) }{ \sum_{i \in \mathcal{M}_c} \phi_i \, \mathcal{N}(\mathbf{x}; \mu_i, \Sigma_{i,\sigma}^*) } .\]

We construct \(\mathcal{M}_\mathbf{sec:c}\) to represent a thin tree-like structure by starting with one main “branch” and recursively subdividing it into smaller ones. Each branch is represented by 8 anisotropic Gaussian components and the subdivision is performed 6 times, decaying \(\phi\) after each subdivision and slightly randomizing the lengths and orientations of the two resulting sub-branches. This yields \(127 {\times} 8 = 1016\) components per class and \(1016 {\times} 2 = 2032\) components in total. We define the coordinate system so that the mean and standard deviation of \(p_\text{data}\), marginalized over \(\mathbf{sec:c}\), are equal to \(0\) and \(\sigma_\text{data} = 0.5\) along each axis, respectively, matching the recommendations by Karras et al. Karras2022elucidating?.

13.0.0.2 Models.

We implement the denoiser models \(D\) as simple multi-layer perceptrons, utilizing the magnitude-preserving design principles from EDM2 [56]. Following Autoguidance [6], we design the model interface so that for a given noisy sample, each model outputs a single scalar representing the logarithm of the corresponding unnormalized probability density, as opposed to directly outputting the denoised sample or the score vector. Concretely, let us denote the output of a given model by \(G_\theta(\mathbf{x}; \sigma, \mathbf{sec:c})\). The corresponding normalized probability density is then given by \[p_{\theta}(\mathbf{x}\mid \mathbf{sec:c}; \sigma) = \exp\!\big( G_{\theta}(\mathbf{x}; \sigma, \mathbf{sec:c}) \big) \;\Big/\; \int \exp\!\big( G_{\theta}(\mathbf{x}; \sigma, \mathbf{sec:c}) \big)\, d\mathbf{x}.\] By virtue of defining \(G_\theta\) this way, we can derive the score vector, and by extension, the denoised sample, from \(G_\theta\) through automatic differentiation: \[\label{eq:EBM} \nabla_x \log p_{\theta}(x | c; \sigma) = \nabla_x G_{\theta}(x; \sigma, c)\tag{3}\]

\[D_{\theta}(x; \sigma, c) = x + \sigma^2 \nabla_x G_{\theta}(x; \sigma, c).\] Besides Equation 3 , according to the description of Autoguidance, trying out the alternative formulations where the model outputs the score vector or the denoised sample directly is qualitatively more or less identical.

To connect the above definition of \(G_\theta\) to the raw network layers, we apply preconditioning using the same general principles as in EDM [59]. Denoting the function represented by the raw network layers as \(F_{\theta}\), we define \(G_\theta\) as \[\label{eq:precond} G_{\theta}(\mathbf{x}; \sigma, c) = -\frac{1}{2}\|\mathbf{x}\|_2^2 - \frac{g_{\theta}}{\sigma n} \sum_{i=1}^{n} F_{\theta, i}\left(\mathbf{x}^{*}; \frac{1}{4}\log \sigma, c\right)^2, \quad \text{where} \quad \mathbf{x}^{*} = \frac{\mathbf{x}}{\sqrt{\sigma^2 + \sigma_{\text{data}}^2}}\tag{4}\] and the sum is taken over the \(n\) output features of \(F_\theta\). We scale the output of \(F_\theta\) by a learned scaling factor \(g_\theta\) that we initialize to zero. The goal of Equation 4 is to satisfy the following three requirements:

• The input of \(F_\theta\) should have zero mean and unit magnitude. This is achieved through the division by \(\sqrt{\sigma^2 + \sigma_\text{data}^2}\).

• After initialization, \(G_\theta\) should represent the best possible first-order approximation of the correct solution. This is achieved through the \(-\frac{1}{2}\|\mathbf{x}^{*}\|_2^2\) term, as well as the fact that \(g_\theta = 0\) after initialization.

• After training, \(\sqrt{g_\theta} \cdot F_\theta\) should have approximately unit magnitude. This is achieved through the division by \(\sigma n\).

In practice, we use an MLP with one input layer and four hidden layers, interspersed with SiLU [60] activation functions and implemented using the magnitude-preserving primitives from EDM2 [61]. Additionally, we defined an output layer after the first hidden layer to produce a weaker version of the output. The input is a 4-dimensional vector \(\big[ \mathbf{x}^\ast_x; \mathbf{x}^\ast_y; \tfrac{1}{4} \log \sigma; 1 \big]\) and the output of each hidden layer has \(n\) features, where \(n = 64\).

13.0.0.3 Training.

Given that we have the exact score function of the ground truth distribution readily available (Equation 2 ), we train the models using exact score matching [62] for simplicity and increased robustness. We thus define the loss function and the additional supervision loss as \[\mathcal{L}(\theta_i) = \mathbb{E}_{\sigma \sim p_{\text{train}}, \mathbf{x} \sim p(\mathbf{x}; \sigma)} \sigma^2 \left\|\nabla_{\mathbf{x}} \log p_{\theta_i}(\mathbf{x}; \sigma) - \nabla_{\mathbf{x}} \log p(\mathbf{x}; \sigma)\right\|_2^2,\]

\[\mathcal{L}(\theta_d) = \mathbb{E}_{\sigma \sim p_{\text{train}}, \mathbf{x} \sim p(\mathbf{x}; \sigma)} \sigma^2 \left\|\nabla_{\mathbf{x}} \log p_{\theta_d}(\mathbf{x}; \sigma) - \nabla_{\mathbf{x}} \log p(\mathbf{x}; \sigma)\right\|_2^2,\] where \(\sigma \sim p_\text{train}\) is realized as \(\log(\sigma) \sim \mathcal{N}(P_\text{mean}, P_\text{std})\) Karras2022elucidating? and \(\theta_i\) and \(\theta_d\) represent the intermediate layer and the deep layer network respectively. In particular, for the training acceleration method depicted in Figure 4, our loss function is \[\mathcal{L}(\theta_i) = \mathbb{E}_{\sigma \sim p_{\text{train}}, \mathbf{x} \sim p(\mathbf{x}; \sigma)} \sigma^2 \left\|\nabla_{\mathbf{x}} \log p_{\theta_i}(\mathbf{x}; \sigma) - \nabla_{\mathbf{x}} \log p(\mathbf{x}; \sigma)\right\|_2^2,\] \[\mathcal{L}(\theta_d) = \mathbb{E}_{\sigma \sim p_{\text{train}}, \mathbf{x} \sim p(\mathbf{x}; \sigma)} \sigma^2 \left\|\nabla_{\mathbf{x}} \log p_{\theta_d}(\mathbf{x}; \sigma) - (\nabla_{\mathbf{x}} \log p(\mathbf{x}; \sigma)+0.5*(\nabla_{\mathbf{x}} \log p_{\theta_d}(\mathbf{x}; \sigma)-\nabla_{\mathbf{x}} \log p_{\theta_i}(\mathbf{x}; \sigma)))\right\|_2^2,\] The final training loss is \(\mathcal{L} = \mathcal{L}(\theta_d) + 0.5*\mathcal{L}(\theta_i)\) The more commonly used denoising score matching do not show any significant differences in the model behavior or training dynamics.

We train these models for 4096 iterations using a batch size of 4096 samples. We set \(P_\text{mean} \!= -2.3\) and \(P_\text{std} \!= 1.5\), and use \(\alpha_\text{ref} / \sqrt{\max(t / t_\text{ref}, 1)}\) learning rate decay schedule with \(\alpha_\text{ref} \!= 0.01\) and \(t_\text{ref} \!= 512\) iterations, along with a power function EMA profile [61] with \(\sigma_\text{rel} \!= 0.010\). Overall, the setup is robust with respect to the hyperparameters; the phenomena illustrated in Figures 3 and 4 remain unchanged across a wide range of parameter choices.

13.0.0.4 Sampling.

We use the standard EDM sampler Karras2022elucidating? with \(N = 32\) Heun steps (\(\text{NFE} = 63\)), \(\sigma_\text{min} \!= 0.002\), \(\sigma_\text{max} \!= 5\), and \(\rho = 7\). We chose the values of \(N\) and \(\sigma_\text{max}\) to be much higher than what is actually needed for this dataset in order to avoid potential discretization errors from affecting our conclusions. In Figure 3, we set \(w = 2.5\) for CFG, \(w = 2\) for autoguidance and \(w = 2\) for internal guidance (IG), and \(w_1 = 1\) and \(w_2 = 1.5\) for IG + CFG. In Figure 4, we set \(w = 2\) for internal guidance.

14 More Discussion on Related Work↩︎

14.0.0.1 Denoising transformers.

Many recent works have tried to use transformer backbones for diffusion or flow-based model training. First, several works like U-ViT [63], MDT [47], and DiffiT [64] how transformer-based backbones with skip connections can be an effective backbone for training diffusion models. Intriguingly, DiT and SiT show a pure transformer architecture can be a scalable architecture for training diffusion-based models. Based on these improvements, Stable diffusion 3 [3], FLUX 1 [65] and Nano Banana [66] show pure transformers can be scaled up for challenging text-to-image generation, and Sora, CogvideoX [67] and Wan [68] demonstrate their success in text-to-video generation. Our work analyzes and improves the training of Diffusion Transformer architecture based on a simple auxiliary supervision to the early layers.

14.0.0.2 Sampling guidance in diffusion models.

Achieving better generation results over diffusion models remains challenging yet essential. Early approaches, such as classifier guidance [42], introduce control by incorporating classifier gradients into the sampling process. However, this method requires separately trained classifiers, making it less flexible and computationally demanding. To overcome these limitations, classifier free guidance (CFG) [4] is proposed, enabling guidance without the need for an external classifier. Instead, CFG trains conditional and unconditional models simultaneously and interpolates between their outputs during sampling. Subsequently, several studies [5], [11]–[14], [69], [70] have reduced the downsides of CFG by making the guidance weight noise level-dependent. Another research direction separates the class guidance part from the CFG and only uses a weaker version of the current model for guidance [6]–[10], which also bears conceptual similarity to contrastive decoding [71] used in large language models to reduce the repetitiveness of generations. Our approach does not require any additional sampling steps. It can be applied to any deep diffusion Transformer network in a plug-and-play manner while achieving better generation results.

14.0.0.3 Efficient training of generative models.

The issue of how to accelerate the convergence of generative models during the training process has also received significant attention. Previous studies have explored architectural optimization [1], [19], [72], improved flow-based theories [26], [73], optimized data [74], [75] to accelerate convergence. Meanwhile, there have been several approaches in generative adversarial network [76] that try to accelerate training with better convergence using pretrained visual encoder [77]–[80]. Another line of work tries to exploit the pretrained visual encoders for improving diffusion model training from scratch [81], [82], usually by training two diffusion models where one model generates the pretrained representations and the other model generates the target data conditioned on the generated representation. More recently, representation learning has been introduced into generative model training to further enhance training efficiency [20]–[22], [35], [48], [83]. Our method not only alleviate the problem of gradient vanishing during training, but also can be directly applied in the sampling stage to achieve better generation results on models with less training consumption.

15 More Qualitative Results↩︎

Below we show some uncurated generation results on ImageNet 256×256 from the SiT-XL+IG (1.8). We use classifier-free guidance with w = 2.5.

References↩︎