On Inductive Biases That Enable Generalization in Diffusion Transformers

Authors: Jie An, De Wang, Pengsheng Guo, Jiebo Luo, Alex Schwing
Primary area: generative models
Keywords: Diffusion Model, Diffusion Transformer, Generalization, Inductive Bias

Abstract:
Recent work studying the generalization of diffusion models with UNet-based denoisers reveals inductive biases that can be expressed via geometry-adaptive harmonic bases. However, in practice, more recent denoising networks are often based on transformers, e.g., the diffusion transformer (DiT). This raises the question: do transformer-based denoising networks exhibit inductive biases that can also be expressed via geometry-adaptive harmonic bases? To our surprise, we find that this is not the case. This discrepancy motivates our search for the inductive bias that can lead to good generalization in DiT models. Investigating a DiT’s pivotal attention modules, we find that locality of attention maps are closely associated with generalization. To verify this finding, we modify the generalization of a DiT by restricting its attention windows. We inject local attention windows to a DiT and observe an improvement in generalization. Furthermore, we empirically find that both the placement and the effective attention size of these local attention windows are crucial factors. Experimental results on the CelebA, ImageNet, and LSUN datasets show that strengthening the inductive bias of a DiT can improve both generalization and generation quality when less training data is available. Source code will be released publicly upon paper publication.

Decision: reject
Reviewer scores: 3, 6, 8, 3

Reviews and discussion:
<thread>
creation date:
1730124486182

thread_type:
official_review

<thread_object>
creator: reviewer_mcon

<content>
summary:
This paper investigates the generalization mechanisms of diffusion models, comparing the inductive biases of UNet-based and Transformer-based diffusion models (DiTs). The authors find that while geometry-adaptive harmonic bases drive generalization in UNets, DiTs rely on the locality of attention maps. This paper modifies attention window sizes to enhance DiT generalization, especially in settings with limited training data.

strengths:
This paper considers the locality of attention maps as the key inductive bias that contributes to the generalization of DiTs, which is an interesting perspective. In this way, the authors control generalization of DiTs by local attention window.

weaknesses:
1.	Findings in this paper need to be further verified. Although Fig. 6 shows that using local attentions in a DiT can improve its generalization (measured by PSNR gap) with limited training data conditions, using a pruned DiTs may also improve its generalization ability. It is a well-known generalization phenomenon that an appropriate ratio between model size/complexity and data quantity achieves a better performance. In this way, decreasing the model size/complexity may be the key reason for improving generalization ability with limited training data conditions, where local attentions is just one way for decreasing model size/complexity.

2.	The contribution is limited. The authors introduce local attention windows to improve the generalization of DiTs, which have been widely used in previous studies [cite1-3]. Furthermore, the authors do not provide a theoretical explanation for why local attention windows improve generalization in DiTs. Therefore, the contribution of this paper is limited.

[cite1] Chenglin Yang, Siyuan Qiao, Qihang Yu, Xiaoding Yuan, Yukun Zhu, Alan Yuille, Hartwig Adam, and Liang-Chieh Chen. Moat: Alternating mobile convolution and attention brings strong vision models. In ICLR, 2022  
[cite2] Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M Alvarez, Jan Kautz, and Pavlo Molchanov. Fastervit: Fast vision transformers with hierarchical attention. arXiv preprint arXiv:2306.06189, 2023  
[cite3] Ali Hassani, StevenWalton, Jiachen Li, Shen Li, and Humphrey Shi. Neighborhood attention transformer. In CVPR, 2023

3.	The experimental results are not convincing. The paper relies heavily on FID scores to evaluate the performance of DiTs. Expanding the evaluation to other relevant metrics, such as CLIP Score [cite4-6], could provide a more comprehensive assessment of the model's capabilities.

[cite4] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748–8763. PMLR, 2021.  
[cite 5] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 586–595, 2018.  
[cite6] Alain Hore and Djemel Ziou. Image quality metrics: PSNR vs. SSIM. In 20th International Conference on Pattern Recognition, pp. 2366–2369. IEEE, 2010.

4.	The improvements shown in the paper focus on limited training data scenarios, but the impact of the proposed method on other datasets or high-resolution tasks is not thoroughly explored, e.g., MS COCO and PartiPrompts.
5.	I suggest the authors provide more visual comparisons to demonstrate the generalization improvement of DiTs.

questions:
Please see the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730576471903

thread_type:
official_review

<thread_object>
creator: reviewer_pa74

<content>
summary:
The paper studies generalization of transfomer-based diffusion models (DiT). The authors first make a connection to prior work that discovered that UNet-based diffusion models generalise because of the inductive biases that can be expressed via geometry-adaptive harmonic bases and show that the same analysis is not able to explain the generalisation in DiTs. Authors instead propose that generalisation in DiTs emerges because of the "locality" of the self-attention layers. They first confirm this empirically, by looking into the attention masks at different layers and discover strong local patterns (i.e., pixels mainly attend to their neighbouring pixels). Then they also propose to use the local attention mask (by masking out all the pixels that are outside of the neighbourhood) and show that this can lead to better performance for smaller datasets (both in terms of the PSNR gap and FID).

I like the paper overall and enjoyed reading it. Hence I vote for acceptance. I am not (at all) familiar with the related literature on inductive biases in diffusion models, so my confidence is low (2/3).

strengths:
- The problem studied (i.e., generalisation of diffusion models) is interesting and highly-relevant. I cannot comment on the novelty aspect since I'm not deeply familiar with this field
- I find the main "theoretical" analysis in the paper interesting, both how they show that the differences with the existing theory for UNets and also how they explain the generalisation of DiTs via the local patterns in the attention masks
- I find the experimental results sufficient for confirming their main hypothesis on the importance of local attention masks

weaknesses:
- The paper is not really self-contained. I feel like this could be improved with including the background section where main concepts used throughout the paper are explained. For example, the Jacobian of the denoising network is crucial for the parts of the analysis, however it is never properly introduced

questions:
- What exactly is the Jacobian you are talking about in Section 2.2? Is it the derivative of the denoising NN w.r.t. its parameters? If so, how do you deal with its dimensionality, since it has dimension of (HW) x M where HxW are the dimensions of the output/image and M is the number of model parameters, right? Also, it is not a diagonal matrix, so do you compute eigen or singular values?
- In line 308 you state you remove the autoencoder, but how do you then go from H x W to (HW) x d, where d is the token dimension?
- Out of curiosity, what is the purpose of black boxes over faces in Figure 1? You are not considering an inpainting task, or?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730641714687

thread_type:
official_review

<thread_object>
creator: reviewer_oygr

<content>
summary:
In this work, the authors identified that an inductive bias of DiT lies in the locality of attention which contributes to the generalizability of DiT. The local attention windows are proposed to be incorporated to DiT to improve the generalizability accordingly. Experimental results also demonstrated the enhanced performance from the incorporated local attention windows.

strengths:
1. *Quality*: This work demonstrates good quality on identifying the inductive bias of DiT on generalizability and constructing the local attention map. The preliminary analysis and the model are of coherence and comprehensiveness.

2. *Significance*: Diffusion models have been attracting increasing attention in the recent years. Analysis the generalizability of the prominent diffusion models DiT via inductive biases provides an interesting perspective to improve diffusion models.

weaknesses:
1. *Presentation*: A minor comment is on the grammar and spelling check of the paper, for example, in Section 1, “Their training involves approximates a distribution…” should be “Their training involves approximating a distribution…”.

2. *Clarity*: What does the Jacobian of a UNet or a DiT mean? Could the authors please articulate how it is computed as in [1] within the main content, for example, Section 2? Hence we can have a better understanding on why the Jacobian eigenvectors can reveal the memorization behavior of Simplified UNets, UNets and DiT.

3. *Novelty*: Although this work is of significance, the analysis perspective and method is still similar to [1]. It would be better to highlight the novelty and contributions of this work compared with [1], especially the differences on the analysis method and techniques. Theoretical analysis will be helpful to improve the novelty, as well as the quality.

[1] Zahra Kadkhodaie, Florentin Guth, Eero P Simoncelli, and Stéphane Mallat. Generalization in diffusion models arises from geometry-adaptive harmonic representation. In *ICLR*, 2024. 1, 2, 3, 4, 5, 10

questions:
1. Could the experimental results also include the simplified UNet and UNets in improved diffusion models? Hence the comparison can be more comprehensive and promising.

2. It would also enhance the qualitative studies in the experimental results to show the improved generalizability of DiT with local attention windows. The impact of the local attention window size can also be presented.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669918219

thread_type:
official_review

<thread_object>
creator: reviewer_vbjl

<content>
summary:
This paper investigates the inductive bias of diffusion transformers. Specifically, the authors start by looking at the difference in the generalization behavior between UNet and DiT. Visualizing the eigenvectors of the harmonic bases, they show that the eigenvectors of smaller eigenvalues of UNet display more interesting patterns than DiT, while the PSNR gap between the training and testing dataset is large for UNet when the number of training data is smaller. Motivated by this difference, they visualize the attention maps of DiT with different training sizes and find out that the attention maps exhibited a more obvious locality pattern when the training size increased. At last, they propose to restrain the attention window size to enhance this locality property of the attention maps during the training. Their results show considerable improvement in the FID score when the training size is smaller.

strengths:
The paper is well-presented and easy to follow. The authors conduct extensive experiments and show relative improvements in the image generalization quality on CelebA and ImageNet datasets.

Overall, I found the topic of looking into the generalization behavior of diffusion transformers very interesting.

weaknesses:
1. The focus in the 1st part of the paper seems to disconnect from the rest, especially the method and the experimental results. I don't understand how the proposed method relates to the difference between the patterns of the harmonic bases for UNet and DiT. I am wondering do the authors observe differences in the harmonic bases after they alter the window size of the transformers. 

2. The results are concerning. Though the current results show a considerable amount of improvement with a very limited size of the training data, the downgraded performance when scaling to a larger dataset makes it hard to believe that locality is the right reason behind the good generalization capability. 

Typos and the minors:

Line 190: $\sigma_t$ is undefined.

Line 202: I assume you are talking about PSNR instead of PNSR, right?

questions:
1. Eqn. 4: What does $\hat{x}_0^k$ mean? And what’s the intuition behind setting $K=300$?

2. What is the training time for different training sizes (i.e., N)? What's the norm of the attention maps when you increase the training time? Should this be considered as an impact on the emergence of the locality?

3. If you conduct the experiments in the latent space, have you observed similar results? Related to this, if you conduct the results on a more complex dataset where long-range correlations between objects matter(e.g., COCO-2017), what's your observation?

4. For Figure 4, can you show the results with different thresholds of the colormaps?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732344978388

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We thank all reviewers for their valuable comments. We are encouraged that reviewers recognize the following merits of this paper:
* The studied topic is interesting, significant, and highly-relevant (Reviewer `VbJL`, `oYGR`, `pa74`, `MCon`)
* The paper is of high-quality and well-presented (Reviewer `VbJL`, `oYGR`)
* The analysis is coherent, comprehensive, and interesting (Reviewer `oYGR`, `pa74`)
* The experimental verification is extensive and sufficient (Reviewer `VbJL`, `pa74`)

We address the raised questions one-by-one below. Please note we also revised the paper to include new experimental results and visualization in response to reviewers’ questions and comments. Newly added and modified text, figures, and tables are highlighted with **red** color.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733294151902

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Author Response Summary

comment:
Dear Reviewers, ACs, SACs, PCs:

Thanks a lot for your valuable comments and active participation in the discussion. We think that all of the questions raised by reviewers have been thoroughly answered. Below we restate our contribution before summarizing how we addressed the reviewers’ questions.

This paper studies the inductive biases of diffusion transformers (DiTs) that enable their generalization. Through an empirical analysis, we identify the locality of attention maps as a key inductive bias, which contributes to the generalization of a DiT. To verify this inductive bias, we present experimental results by incorporating local attention into a DiT. We show that enhancing the locality of attention can effectively modify a DiT’s generalization. We added a theoretical explanation showing that the attention locality bias is connected with the low sensitivity bias studied by prior work.

During the discussion period, we address the questions of all reviewers as summarized below:
* Reviewer `VbJL`:
  * **Differences in the harmonic bases when altering the attention window size**: We added Fig. 4 to the revised Appendix F. It shows the harmonic bases after applying local attention.
  * **New experiments with MSCOCO and latent space DiT**: We analyzed and presented the results on MSCOCO in Tabs. 2 and 3 of the revised Appendix E.1. We included experimental results for a latent diffusion model in Tab. 4 of the revised Appendix E.2, using CelebA and MSCOCO data.
  * **Possible misunderstanding of the term “generalization”**: We provided an explanation to note that the generalization we study is not identical to the data scaling capability of a DiT.
  * **New experiments with a training set of $10^6$ images**: We provided results with $10^6$ images using ImageNet and LSUN Bedroom data. We explained why we think using $10^5$ images is sufficient to verify the discovered inductive bias. Results with $10^6$ images corroborate this.
* Reviewer `oYGR`:
  * **Details about the Jacobian eigenvector computation**: We moved details about the Jacobian eigenvector computation from the Appendix to Sec. 2 of the paper.
  * **New experiments with UNet**: We included experimental results for the UNet in Tabs. 2, 3, and 4 of the revised Appendix E.
  * **Qualitative results**: We provided qualitative results in Fig. 8 of the revised Appendix F.
* Reviewer `pa74`: 
  * **Details about the Jacobian eigenvector computation**: We provided an explanation in the response and moved the details about the Jacobian eigenvector computation from the Appendix to Sec. 2 of the paper.
  * **Details about token dimension**: We provided an explanation in the response.
* Reviewer `MCon`: 
  * **Effectiveness of local attention compared with simply reducing model complexity**: We provided an explanation noting that reducing the model complexity does not always lead to the improved generalization. We added new experiments with a smaller DiT-XXS/1 in Tabs. 2 and 3 of Appendix E, showing that reducing model complexity is not as effective as local attention.
  * **The contribution of this paper considering that local attention has been widely used**: We provided an explanation stating that our goal isn’t to improve a DiT’s performance via local attention. In contrast, we recognize attention locality as a key inductive bias of DiT. Local attention serves as a way to encourage attention locality of a DiT. We solely apply local attention to verify our discovered inductive bias rather than as a contribution of this work.
  * **New experiments with MSCOCO**: We included experimental results using MSCOCO data in Tabs. 2 and 3 of the revised Appendix E.1.
  * **Qualitative results**: We provided qualitative results in Fig. 8 of the revised Appendix F.
  * **Theoretical explanation for why local attention improves generalization of a DiT**: We provided a theoretical explanation in Appendix H linking the attention locality bias with the low sensitivity bias, flat minima, and generalization. We also provided a line-by-line explanation of the theoretical derivation.

In summary, we think we have comprehensively addressed all reviewers' comments. We are optimistic that our work provides valuable insights for generative modeling, facilitating understanding of the behavior of a DiT. Lastly, we thank everyone once again for their time and feedback, which helped us to improve the paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores the generalization behaviour of Diffusion Transformers, and claims that attention locality is the inductive bias, as opposed to UNet-based diffusion models which were said to have harmonic bases as inductive bias in prior work. Local attention windows are proposed as a way to improve generalizability.

However, reviewers felt that the experimental evaluation was not sufficient to robustly back up these claims. In particular, there can be alternative explanations for the beneficial effect of local attention windows on small datasets which were not ruled out by the authors. For example a more mundane explanation that they simply decrease the model complexity improving generalizability, rather than having a beneficial impact on the inductive bias. The authors have failed to consistently apply the scientific method by making claims without sufficient evidence and considering alternatives, and hence I am recommending rejection.