Title: On Inductive Biases That Enable Generalization in Diffusion Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion Model

Diffusion Transformer

Generalization

Inductive Bias
Abstract: Recent work studying the generalization of diffusion models with locally linear UNet-based denoisers reveals inductive biases that can be expressed via geometry-adaptive harmonic bases. For such locally linear UNets, these geometry-adaptive harmonic bases can be conveniently visualized through the eigen-decomposition of a UNet’s Jacobian matrix. In practice, however, more recent denoising networks are often transformer-based, e.g., the diffusion transformer (DiT). Due to the presence of nonlinear operations, similar eigen-decomposition analyses cannot be used to reveal the inductive biases of transformer-based denoisers. This motivates our search for alternative ways to explain the strong generalization ability observed in DiT models. Investigating a DiT’s pivotal attention modules, we find that locality of attention maps in a DiT’s early layers are closely associated with generalization. To verify this finding, we modify the generalization of a DiT by restricting its attention windows 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, MSCOCO, and LSUN data show that strengthening the inductive bias of a DiT can improve both generalization and generation quality when less training data is available. Source code is available at https://github.com/DiT-Generalization/DiT-Generalization.


Peer reviews:
comments: SUMMARY:
This paper studies the inductive bias of DiT models. It is found that the inductive bias of DiT models does not manifest as geometry-adaptive harmonic bases as UNet models do. Instead, the authors discover that it is related to the locality of early self-attention layers. Furthermore, abundant experiments validate this finding.

STRENGTHS AND WEAKNESSES:
Pros:
1. The authors are focusing on the newest trend of diffusion model advancement, studying the property of DiT models.
2. This study has a solid methodology, along with abundant experimental evidence and an interesting finding.
3. For certain scenarios where the training data is scarce, the finding might have some practical impact.

Minor Cons:
1. The architecture of actual commercial diffusion models is more complex than the simplified versions in this study. For example, Stable Diffusion XL model, as a UNet model, has many 10-layer ViTs integrated in the UNet, which may make it behave specially. In Stable Diffusion 3 and Flux, the attention blocks are upgraded into MMDiT fusion attention, which may also introduce new properties. I would recommend some future work on commercial diffusion models to extend the study to more complex architectures.

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

LIMITATIONS:
I recommend considering the special ViT-UNet-mixed architecture of Stable Diffusion XL as a possible future topic.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-04

comments: SUMMARY:
This paper studies the inductive biases of DiT for generalization. The research question is: do transformer-based denoising networks have inductive biases that can be expressed via geometry-adaptive harmonic bases like in UNet-based denoisers? The paper then shows that this is not the case, and proposes to restrict the attention windows in early layers, which can improve the generalization ability. The paper conducts experiments on CelebA. ImageNet, MSCOCO, and LSUN datasets, using DiT.

STRENGTHS AND WEAKNESSES:
strength:
1. the research question is specific and clearly stated.

2. the analysis in section 2.3 and the verification in section 3 is reasonable.

weakness:
1. 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.

2. The proposed method is not very easy to be adapted to many other Transformer-based Diffusion models. The effectiveness of the method is affected by window sizes and which layers being modified. If we target on other Transformer-based models, does it mean that we should do these analysis again and try to choose appropriate parameters?

QUESTIONS:
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?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-06

comments: SUMMARY:
This paper investigates the inductive biases that facilitate the generalization ability of DiT. The authors propose that locality of attention maps in a DiT’s early layers are closely associated with generalization. They use local attention windows in early layers of a DiT and observe an improvement in generalization. Experimental results on the CelebA, ImageNet, MSCOCO, and LSUN data show that strengthening the inductive bias can improve both generalization and generation quality when less training data is available.

STRENGTHS AND WEAKNESSES:
Strengths：
1. The paper proposes the locality of attention maps as a key inductive bias contributing to the generalization of a DiT. 
2. The paper demonstrates how to control this inductive bias by incorporating local attention windows into early layers of a DiT. 

Weaknesses:
The experimental part has limitations. The proposed approach lacks comprehensive comparisons with state-of-the-art methods and in-depth ablation studies to fully validate both generalization and generation quality. For DiT, there is not any generated image as experiment result to demonstrate and compare. That is not enough only using PSNR gap and FID to verify the effectiveness of the viewpoints.

QUESTIONS:
1. The experiments are not sufficient to support the proposed viewpoints.
2. Are there some generated images to demonstrate?
3. Are there some comparison results with SOTAs?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to understand and explain the inductive biases responsible for generalization in transformer-based diffusion models. The authors show that the underlying mapping is not harmonic, unlike UNet-based diffusion models. They propose that the locality of attention maps lead to generalization. The claim is supported with extensive experiments.\n\nSTRENGTHS AND WEAKNESSES:\n**Main strength:**\n- I found the results in figure 4 amazing! This figure shows more than mere locality of the attention maps in the generalization regime. The plots show adaptive convolutional layer maps! It seems that with enough training data, the attention block in the transformer converges to convolutional layer. And the effective kernel size increases as the layer number increases. This is consistent with results presented in Fig 2 which shows that denoising performance is quite comparable between convolutional model and transformer based model. I find this result quite fascinating since it challenges the commonplace belief that DiT is superior to conv nets due to its access to global dependencies. The authors show the opposite! That in the generalization regime, the local information is what's needed and used even when access to global dependencies are provided. \n\nConsider relating your results to this paper that is related but for classification: Maithra Raghu et al, Do Vision Transformers See Like Convolutional Neural Networks?, 2021 \n\n**Main weakness:**\n- The Jacobian analysis is not valid as a tool for understanding the mapping in the DiT, making Fig 1 and Figure 3 difficult to interpret. The Jacobian is a first-order Taylor approximation to the denoising mapping, so it can fully explain the mapping only when the network is *locally linear*. In Kadkhodaie et al 2024, the UNet is modified *to become  locally linear*, as a result the Jacobian times the input is exactly equal to the end to end network operation (important pieces are: using ReLU, removing the bias terms of the convolutions and also using batch-norm without mean). Previous work has shown these changes do not compromise the performance but makes the network locally linear, which lends it to eigen analysis of the Jacobian. The standard DiT architecture is far from locally linear, due to the use of non-linearities other than ReLu, like softmax, and LayerNorm or GroupNorm layers. So it is not clear what a first order approximation captures about the mapping! One would need the higher orders and the affine bias to get the full picture. So  DiT may or may not be arriving at geometry adaptive harmonic bases, but the presented experiments do not answer this question. However, given the results in Fig 4, my guess is that DiT is probably arriving at GAHBs since the attention layers seems to be converging to convolutional patterns in generalization regime.\n\n************************* \n**Other points:**\n- The eigen vectors associated with Nichole & Dhariwal 2021 shown in Fig 1 do not really show GAHBs despite the caption. If the original architecture is used, it contains group norm, nonlinearity other than Relu (I think Gelu) that makes the mapping locally non-linear, and I believe that's the reason.\n- More importantly, the architecture used in Nichole & Dhariwal 2021 contains **attention blocks** (to my knowledge), yet the eigen vectors are very different from panel (c) in Fig1. This is another evidence that the Jacobian analysis for locally non-linear mappings is not very informative. Just from the fig 1 results, I infer that DiT is more locally non-linear than Nichole & Dhariwal 2021, maybe due to softmax and other pieces I am not aware of. But the general point is that the fig is not reflecting the true bases the networks are arriving at. \n- Minor point: Labeling the UNet in Kadkhodaie et al 2024 as simplified might be misleading. The change from grayscale (one-channel input) to color (three-channel input) is not significant and doesn't impact the end to end mapping, due to the high correlation between RGB channels. In other words, the shape information is way more important and difficult to capture than color information. On the other hand, the network used in Nichol & Dhariwal 2021 could be labeled as augmented UNet since it has attention block t-conditioning, etc hence more significantly different from original UNet. \n- 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. \n- Suggestion: replacing the deviation from identity eq 6 with deviation from convolutional kernel with a certain kernel size. \n************************* \nI think overall the paper is quite significant in shedding light on the operation of attention blocks in generalization regime. So despite the erroneous Jacobian analysis, I think this paper can have an impact. I would increase my score if the comments and questions about the weakness (regarding Jacobian analysis) are addressed.\n\nQUESTIONS:\n- Do you observe the strong generalization (sampling the same image by two networks trained on disjoint dataset) in DiT? I understand this needs lots of compute, so I'm not asking to do the experiment during the rebuttal, but curious if you have done the experiments already.\n\nLIMITATIONS:\n- No negative societal impacts. \n- Yes limitations have been addressed", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}