On Inductive Biases That Enable Generalization in Diffusion Transformers

Thank you very much for your time, thoughtful feedback, and for highlighting the strengths of our work, specifically our focus on timely advances in diffusion models, the solid methodology, the well-supported experimental results, and the interesting findings. We also appreciate your recognition of the potential practical value of our method in low-data regimes. We address the remaining questions below.

• Q1: Suggestion to extend the analysis to more complex commercial diffusion architectures (e.g., SDXL, SD3, Flux) which include hybrid UNet and ViTs and conditional attention blocks.

  Thanks for the valuable suggestion. As stated in line 40-42 of the appendix, the UNet architecture used in our experiments contains one attention layer at the last encoder stage and one attention layer at the middle block, which is a hybrid model combining convolution and two attention layers. In general, we agree that studying the generalization behaviour of hybrid models and conditional transformers (e.g., MMDiT modules) are important and interesting due to the increasing popularity of these architectures.

• Q2: What could happen if the attention window is somehow enlarged in DiTs trained with sufficient data? Could this modification make them generalize worse?

  For attention layers of a normal DiT, each token has full access to all other tokens, meaning that the attention window size is already at the maximum and cannot be enlarged. On the other hand, a DiT trained with sufficient data already forms a locality pattern in its attention layers as shown in Fig. 4 ($N=10^5$ row). Therefore, we think further enlarging the attention window would have very limited impact on a DiT’s generalization and performance since newly added non-local tokens in attention don’t contribute much.

• Q3: I recommend considering the special ViT-UNet-mixed architecture of Stable Diffusion XL as a possible future topic.

  Thank you for a valuable suggestion. We agree and will note these future directions in the revised paper.

Thanks a lot for your time, feedback, and highlighting that our research question is specific and clearly stated while our analysis in Sec. 2.3 and verification in Sec. 3 are reasonable. We answer the remaining questions below.

• Q1: It seems that the result is not very promising. As shown in Table 3, using the proposed local attention can be harmful to the performance in the $10^5$ setting. Although it can improve the performance in the $10^4$ setting, the FID is still high. This makes the applicability of this work a bit questionable.

  This paper aims to analyze the inductive bias of a DiT that contributes to its generalization, as stated in line 81-86. Based on experiments reported in Sec. 2, we hypothesize that attention locality is an important indicator of generalization. Our experiments hence modify attention patterns by enforcing attention locality while assessing the resulting impact on the generalization of a DiT (FID). This should not be interpreted as a new method, but rather as studies to verify that the hypothesized attention locality bias is indeed at play. Regarding the FID results shown in Tab. 3, when trained with $10^5$ images, as shown in Fig. 4 (N=$10^5$ row), attention locality patterns emerge naturally when training. Hence, it is expected that additionally encouraging attention locality doesn’t lead to any changes in a DiT’s generalization as validated via the PSNR gap in Tab. 2. For FID, we observe a minimal $4.33$% performance fluctuation averaged over seven datasets. Considering that using local attention offers the added benefit of reducing FLOPs, we think it is inaccurate to say local attention is harmful when training a DiT with $10^5$ samples. When using $10^4$ training samples, our experiments consistently show that introducing local attention in early DiT layers improves generalization (PSNR gap in Tab. 2 and Cosine Similarity in Tab. 3 of appendix) and sample quality (FID in Tab. 3, in appendix: IS in Tab. 4, FD-Dinov2 in Tab. 6). In Tab. 3, we observe a consistent $17.59$% FID improvement averaged over seven datasets, which we think is significant compared with the $4.33$% fluctuation when $N=10^5$.

• Q2: The proposed method may not be easy to be adapted to other Transformer-based diffusion models, given sensitivity to window size and layer selection.

  Our experimental results in Tab. 2 and 3 consider two DiTs of different sizes and seven datasets. Using local attention at beginning layers of a DiT with gradually increasing attention window size consistently improves the model generalization. In Tab. 8 of the appendix, we apply the same local attention setting for a latent-space DiT, which also demonstrates reasonable generalization and performance improvement when trained with $N=10^4$ samples, as measured by PSNR gap and FID, respectively. Given these consistent improvements, we think our default setting would work reasonably well for most DiTs. Despite this, a careful analysis of the generalization behaviour is still beneficial for a specific DiT architecture, which we consider one of the key takeaways of this paper.

• Q3: It might be a bit narrow if the paper only focuses on geometry-adaptive harmonic bases as in UNet-based denoisers. Is it possible to study the generalization of DiT through a more general exploration?

  This paper studies the generalization of DiTs from different perspectives. In Sec 2.2, we use a Jacobian analysis to study the representation behaviour of DiTs. Beyond that, in Sec. 2.3, we demonstrate the attention locality bias that is associated with a DiT’s generalization, which is verified by an improved PSNR gap when $N=10^4$ after injecting local attention into a DiT (Sec. 3). Further, in Appendix B, we provide a theoretical analysis connecting the discovered attention locality bias with the simplicity bias (Appendix B.1), and subsequently, low sensitivity bias (Appendix B.2), which indicates good generalization of a Transformer (Appendix B.3). Despite these, we recognize a more general exploration of a DiT’s generalization, an important and interesting future work.

Thanks a lot for your time, feedback, and highlighting that this paper reveals the locality of attention maps as an inductive bias of a DiT, which is verified by our experiments incorporating local attention windows into early layers of a DiT. We answer the remaining questions below.

• Q1: More comprehensive comparisons with state-of-the-art methods and in-depth ablation studies to fully validate both generalization and generation quality.

  This paper aims to analyze the inductive bias of a DiT that contributes to its generalization, as stated in line 81-86. Based on experiments reported in Sec. 2, we hypothesize that attention locality is an important indicator of generalization. Our experiments hence modify attention patterns by enforcing attention locality while assessing the resulting impact on the generalization of a DiT. This should not be interpreted as a new method, but rather as studies to verify that the hypothesized attention locality bias is indeed at play. To our best knowledge, there are no state-of-the-art methods to compare to which aim to control the generalization of a DiT. Nonetheless, in Appendix I, we compare the effect of encouraging attention locality with a known way of improving generalization – reducing trainable parameters. As shown in Tab. 11 of Appendix I, two parameter reduction approaches, i.e., using parameter sharing and composing attention maps with pre-defined PCAs, are not as effective in explaining a DiT’s generalization as the local attention study. Further, in Tab. 4, we show that encouraging attention locality at the first few layers of a DiT can improve its generalization. Encouraging attention locality at the last few layers or interleaving global and local attention do not. In Tab. 5, we show that use of a smaller attention window size can lead to a larger generalization improvement while a larger window size leads to a smaller improvement. All these ablation studies verify that encouraging attention locality impacts a DiT’s generalization, corroborating our hypothesis.

• Q2: For DiT, there is no generated image to demonstrate and compare.

  We present qualitative results in Fig. 2 of Appendix D. We find that DiTs trained with $10^4$ and $10^5$ images while using local attention produce images that are more like each other than DiTs trained without using local attention. Further, in Tab. 9 of Appendix D, we quantitatively verify that the average pixel intensity between generated images of models trained with $10^4$ and $10^5$ images is smaller if local attention is encouraged. This phenomenon shows that applying local attention to a DiT trained with $10^4$ images can make it generate images closer to a well-generalized model trained with $10^5$ images, indicating generalization improvement.

• Q3: Only using PSNR gap and FID to verify the effectiveness of the viewpoints is not enough.

  We consider five evaluation metrics: PSNR gap (Tab. 2), FID (Tab. 3), IS (Tab. 4 in Appendix C.3), FD-Dinov2 (Tab. 5 in Appendix C.4), and Cosine Similarity to the nearest training image (Tab. 2 in Appendix C.2). Prior works studying diffusion model generalization [1,2] only use PSNR gap and Cosine Similarity as evaluation metrics. We hence believe the metrics used in our evaluation extend prior work and comprehensively support our claim.

[1] Generalization in diffusion models arises from geometry-adaptive harmonic representation. In ICLR, 2024.

[2] The emergence of reproducibility and consistency in diffusion models. In ICML, 2024.

Thanks a lot for your time and feedback. We are encouraged that our findings about DiT’s locality bias shown in Fig. 4 are recognized: A DiT can form a locality pattern similar to adaptive convolution when it achieves a good generalization, despite that access to global tokens is provided. We also appreciate the suggested related work by Raghu et al., which echoes our finding but in ViTs. We will discuss this new related work in the revised paper. We answer the remaining questions below.

• Q1: DiT lacks local linearity due to the use of GeLU, Softmax, and group norm layers. Hence, the Jacobian analysis in Figs. 1 and 3 is not a valid interpretation tool for DiT’s mappings.

  We agree, a DiT is not locally linear. Note, we use the Jacobian to analyze the model’s sensitivity and representation behavior instead of linearly decomposing the full mapping. Therefore, we do not draw strong conclusions from Fig. 1 and Fig. 3 alone. Rather, we view them as suggestive empirical evidence that first-order Jacobian analysis may not suffice to study the inductive bias of a DiT. This motivates direct visualization of the attention maps of a DiT in Fig. 4, where we observe that attention patterns in DiT converge to localized structures. We’ll clarify in a revised version by adding this discussion and noting that this doesn’t rule out GAHBs.

• Q2: The Jacobian analysis in Fig. 1 is not as informative due to the non-linearity in both DiT and UNet used by Nichol & Dhariwal 2021.

  Thanks a lot for a great argument. To confirm, we use one attention layer at the last encoder stage and in the middle block in the UNet by Nichol & Dhariwal 2021, which is stated in line 40-42 of the appendix. We also use group norm and SiLU layers in the UNet. We think the amount of non-locality introduced by network operators is a reasonable and interesting explanation for the Jacobian eigenvector differences between the UNets and DiT in Fig. 1. We will add this discussion to Sec. 2.2 of the revised paper to provide more insight. Furthermore, for a DiT with the same FLOPs as a UNet, whether the amount of non-linearity is a deciding factor for their generalization difference is another interesting future direction.

• Q3: Minor: Labeling the UNet in Kadkhodaie et al. (2024) with “simplified” is misleading. The UNet used by Nichol & Dhariwal 2021 should be labeled as augmented UNet.

  Thanks for the valuable suggestion. Initially, we considered the UNet by Nichol & Dhariwal 2021 as the "standard" UNet as it has been widely used. For clarity, in the revised paper, we will rename the UNet used by Kadkhodaie et al. (2024) as the locally-linear UNet and the UNet by Nichol & Dhariwal 2021 as a hybrid UNet, recognizing that it contains attention layers.

• Q4: To compute the average PSNR in Eq 4, the sum should be outside of log (the two operations do not commute). i.e compute individual PSNR for each image and then take the average.

  Thank you for pointing this out. Upon reviewing our code, we confirmed that the reported PSNR values are computed by first calculating the PSNR for each image individually and then averaging the results. We will revise Eq. (4) to reflect this accordingly.

• Q5: Suggestion: replacing the deviation from identity in eq 6 with deviation from convolutional kernel with a certain kernel size.

  Thanks for a valuable suggestion. Below we present the deviation score w.r.t. a convolution kernel of $5\times5$ for a DiT-XS/1 trained with CelebA data using $10^3$, $10^4$, and $10^5$ images, respectively.

  Train Data Num	$10^3$	$10^4$	$10^5$
  Layer 1	0.274	0.016	0.010
  Layer 5	0.075	0.055	0.033
  Layer 9	0.070	0.054	0.040

  The results above also show that the deviation score w.r.t. a convolution kernel decreases when a DiT has better generalization, i.e., when it is trained with more images. We will add the new deviation score results into the revised appendix.

• Q6: Do you observe strong generalization (sampling the same image by two networks trained on disjoint datasets) in DiT?

  Yes, we did. We will add images generated by two DiTs trained with disjoint data into the revised appendix.