On Inductive Biases That Enable Generalization in Diffusion Transformers

## Abstract

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. Project page: dit-generalization.github.io/ .

## 1 Introduction

Diffusion models have achieved remarkable success in visual content generation. Their training approximates a distribution in a high-dimensional space from a limited number of training samples–a task that is highly challenging due to the curse of dimensionality. Nonetheless, recent diffusion models (Sohl-Dickstein et al., ; Song et al., ; Ho et al., ; Kadkhodaie & Simoncelli, ; Nichol & Dhariwal, ; Song et al., ) learn to generate high-quality images (Nichol et al., ; Dhariwal & Nichol, ; Saharia et al., ; Rombach et al., ; Chen et al., ; ) and even videos (Singer et al., ; Ho et al., ; Girdhar et al., ; Blattmann et al., ; OpenAI, ) using relatively few samples when compared to the underlying high-dimensional space. This indicates that diffusion models exhibit powerful inductive biases (Wilson & Izmailov, ; Goyal & Bengio, ; Griffiths et al., ) that promote effective generalization.
What exactly are these powerful inductive biases?
Understanding them is crucial for gaining deeper insights into the behavior of diffusion models and their remarkable generalization.

Recent work by Kadkhodaie et al. () on UNet-based diffusion models reveals that
the strong generalization of UNet-based denoisers is driven by inductive biases that can be expressed
via a set of geometry-adaptive harmonic bases (Mallat et al., ).
Their result is illustrated in Fig. 1 (a): the harmonic bases are extracted from a simplified one-channel UNet via the eigenvectors of the denoiser’s Jacobian matrix.
It is easy to extend the analysis of Kadkhodaie et al. () to show that similar harmonic bases are also observed in more complex and classic multi-channel UNets (Nichol & Dhariwal, ), as shown in Fig. 1 (b).
Given this observation, it is natural to ask: does the emergence of harmonic bases also occur in compelling recent transformer-based diffusion model backbones, e.g., diffusion transformers (DiTs) (Peebles & Xie, )?
To explore this possibility,
we perform an eigendecomposition of a DiT’s Jacobian matrix, following Kadkhodaie et al. (). To our surprise, as shown in Fig. 1 (c), a DiT trained in the pixel space does *not* exhibit geometry-adaptive harmonic bases, making it different from a UNet.
Building on these insights, a natural question arises: what are the inductive biases that enable the strong generalization of DiTs?

Answering this question is particularly important because of the growing adoption of DiTs in recent methods (Chen et al., ; Esser et al., ), partly for its observed performance at scale (Peebles & Xie, ).
In a new study in this paper, using the PSNR gap (Kadkhodaie et al., ) as a metric to evaluate the generalization of diffusion models, we confirm that a DiT indeed exhibits better generalization than a UNet with the same FLOPs.
Yet, as mentioned before, this observation alone doesn’t reveal the inductive biases which enable generalization.

The generalization mechanism of a DiT may differ from that
of UNet-based models, potentially due to the self-attention (Vaswani, ) dynamics which are pivotal in DiT models but not in UNets.
In a self-attention layer, the attention map, derived from the multiplication of query
and key
matrices, determines how the value matrix
obtained from input tensors influences output tensors.
To shed some light, we analyze the attention maps of a DiT and show that locality of the attention maps is closely tied to its generalization ability. Specifically, the attention maps of a DiT trained with insufficient images, i.e., with weak generalization, exhibit a more position-invariant pattern: the output tokens of a self-attention layer are largely influenced by a certain combination of input tensors, irrespective of their positions. In contrast, the attention maps of a DiT trained with sufficient images, which demonstrates strong generalization, exhibit
a sparse diagonal pattern. This indicates that each output token is primarily influenced by its neighboring input tokens.
This analysis provides insight into how the generalization ability of DiTs can be modified, if necessary, such as when only a small number of training images are available.

Restricting the attention window in self-attention layers should permit modifying a DiT’s generalization.
Indeed, we find that employing local attention windows (Beltagy et al., ; Hassani et al., ) is effective. A local attention window restricts the dependence of an output token on its nearby input tokens, thereby promoting the locality of attention maps. In addition, the placement of attention window restrictions within the DiT architecture and the effective size of attention windows are critical factors to steer a DiT’s generalization.
Our experiments show that placing attention window restrictions in the early attention layers of the DiT architecture yields the most benefit.
Experimental results on the CelebA (Liu et al., ), ImageNet (Deng et al., ), and LSUN (Yu et al., ) (bedroom, church, tower, bridge) datasets demonstrate that applying attention window restrictions improves generalization, as reflected by a reduced PSNR gap. We also observe an improved FID (Heusel et al., ) when training with insufficient data, confirming that DiT’s generalization can be successfully modified through attention window restrictions.

In summary, the contributions of this paper include the following: 1) We identify the locality of attention maps as a key inductive bias contributing to the generalization of a DiT, and 2)
we demonstrate how to control this inductive bias by incorporating local attention windows into a DiT. Enhancing the locality in attention computations effectively modifies a DiT’s generalization, resulting in a lower PSNR gap and improved FID scores when insufficient training images are available.

## 2 Analyzing the Inductive Bias of Diffusion Models

Diffusion models are designed to map a Gaussian noise distribution to a dataset distribution. To achieve this, diffusion models take a noisy image $\bm{x}_{t}$, obtained by adding Gaussian noise $\bm{\epsilon}$ to a training sample $\bm{x}_{0}$ following a noise schedule depending on step $t$, and estimate noise $\bm{\epsilon}$.
The loss function of diffusion model training is as follows:

$$ $\mathcal{L}=\mathbb{E}_{\bm{x}_{0},\bm{\epsilon},t}\left[\left\|\bm{\epsilon}- \bm{\epsilon}_{\theta}(\bm{x}_{t},t)\right\|_{2}^{2}\right].$ (1) $$

In Eq. (1), $\bm{\epsilon}_{\theta}(\cdot)$ represents the backbone network with trainable parameters $\theta$, which plays a crucial role in diffusion model generalization and hence is our primary focus.
In this section, we first compare the generalization ability of a DiT (Peebles & Xie, ) and a UNet (Nichol & Dhariwal, ), two of the most popular diffusion model backbones. Subsequently, we investigate the inductive biases that drive their generalization.

**Figure 2: The PSNR (a) and PSNR gap (b) comparisons between a UNet and a DiT with the same FLOPs for different training image quantities ($N$). When $N{=}10^{5}$, both DiT and UNet show small PSNR gaps between the training and testing sets. Nevertheless, when $N{=}10^{3}$ and $N{=}10^{4}$, a DiT exhibits smaller PSNR gaps compared to a UNet, indicating a better generalization ability under insufficient training data. All PSNR and PSNR gap curves are averaged over three models trained on different dataset shuffles. The standard deviations, illustrated by the curve shadows in the zoomed-in windows, are negligible, indicating minimal variation.**
| UNet |  |  |
| --- | --- | --- |
| DiT |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

### 2.1 Comparing DiT and UNet Generalization

We compare the generalization of pixel-space DiT and UNet(^1^11[https://github.com/openai/improved-diffusion](https://github.com/openai/improved-diffusion)) using as a metric the PSNR gap proposed by Kadkhodaie et al. ().
The PSNR gap is the zero-truncated difference between the training set PSNR and the testing set PSNR at a diffusion step $t$:

$$ $\mathrm{PSNR\ Gap}\left(t\right)=\max\left(\mathrm{PSNR}_{\mathrm{train}}\left (t\right)-\mathrm{PSNR}_{\mathrm{test}}\left(t\right),0\right),$ (2) $$

where $\mathrm{PSNR}_{\mathrm{train}}\left(t\right)$ and $\mathrm{PSNR}_{\mathrm{test}}\left(t\right)$ are obtained following Kadkhodaie et al. (). To elaborate, given $K$ images from either training or testing set, we first feed noisy images at step $t$ to diffusion models and obtain the estimated noise $\hat{\bm{\epsilon}}$. Next, we get the one-step denoising result $\hat{\bm{x}}_{0}$ via

$$ $\hat{\bm{x}}_{0}=\bm{x}_{t}-\sigma_{t}\hat{\bm{\epsilon}}.$ (3) $$

Finally, we derive the training and testing PSNRs at diffusion step $t$ as follows:

$$ $\mathrm{PSNR}_{\mathrm{train/test}}\left(t\right)=10\cdot\left(\log(M^{2})- \log\left(\frac{1}{K}\sum_{k=1}^{K}\mathrm{MSE}\left(\hat{\bm{x}}_{0}^{k},\bm{ x}_{0}^{k}\right)\right)\right).$ (4) $$

Here, $M$ denotes the intensity range of $\bm{x}_{0}$, which is set to $2$ since $\bm{x}_{0}$ is normalized to $\left[-1,1\right]$. $K$ is set to $300$ following the PSNR gap computation of Kadkhodaie et al. ().

Turning to diffusion model backbones, prior work (Peebles & Xie, ) has shown that a DiT achieves better image generation quality than a UNet with equivalent FLOPs. This advantage of DiT prompts our curiosity to study whether DiT can also demonstrate superiority in generalization, using the PNSR gap as a metric.
Fig. 2 compares the PSNR and PSNR gap of a UNet and a DiT.
Interestingly, when the number of training images is sufficient for the model size, e.g., $N{=}10^{5}$, the training and testing PSNR curves of both DiT and UNet are nearly identical, and their PSNR gaps remain small. This indicates that DiT and UNet have no substantial performance difference in distribution mapping given sufficient training data.
Nevertheless, as shown in Fig. 2 (b), when trained with less data, e.g., $N{=}10^{3}$ and $N{=}10^{4}$, a DiT has a remarkably smaller PSNR gap than a UNet, suggesting that a DiT has a better generalization ability than a UNet.
This discrepancy of the PSNR gap motivates us to explore the underlying inductive biases that contribute to the generalization difference between a DiT and a UNet.

Figure: Figure 3: Jacobian eigenvector comparison between UNet (Nichol & Dhariwal, ) and DiT (Peebles & Xie, ) with equivalent FLOPs. (a) The eigenvectors of a UNet tend to memorize the training images when $N{=}10$ and drive the generalization througth harmonic bases (Kadkhodaie et al., ) when $N{=}10^{5}$. In contrast, (b) the DiT’s eigenvectors exhibit neither the memorization effect at $N{=}10$ nor harmonic bases at $N{=}10^{5}$.
Refer to caption: x8.png

### 2.2 DiT Does Not Have Geometry-Adaptive Harmonic Bases

Kadkhodaie et al. () reveal that the generalization of a simplified one-channel UNet is driven by the emergence of geometry-adaptive harmonic bases. These harmonic bases are obtained from the eigenvectors of a UNet’s Jacobian matrix.
This raises an important question: can the potential difference in harmonic bases between a DiT and a UNet account for their generalization differences?
To address this, we follow the approach of Kadkhodaie et al. () and perform an eigendecomposition of the Jacobian matrices for a three-channel classic UNet (Nichol & Dhariwal, ) and a DiT.

Fig. 3 presents the eigenvalues and eigenvectors of a UNet and a DiT trained with $10$ and $10^{5}$ images, respectively.
For a UNet which is trained with a small number of images, e.g., $N{=}10$, the Jacobian eigenvectors corresponding to several large eigenvalues tend to memorize the geometry of the input image. Notably, the leading eigenvalues are significantly larger than the rest, indicating that the UNet trained with $10$ images is governed by memorization of the training images (Carlini et al., ; Somepalli et al., ). In contrast, when the training set size is increased to $N{=}10^{5}$, the UNet’s eigenvectors show the geometry-adaptive harmonic bases similar to the ones reported by Kadkhodaie et al. (): oscillating patterns which increase in frequency as eigenvalues $\lambda_{k}$ decrease.
This clear transition from memorizing to generalizing, observed as $N$ increases, indicates that harmonic bases play a key role in driving the generalization of a UNet.

In contrast, harmonic bases do not appear to be the driving factor behind a DiT’s generalization. As shown in Fig. 3 (b), the eigenvectors of the DiT do not exhibit the harmonic bases similar to the ones observed for the UNet. Instead, the DiT displays random sparse patterns regardless of the training dataset size. Additionally, the difference in the distribution of DiT’s eigenvalues between $N{=}10$ and $N{=}10^{5}$ is much less pronounced compared to that of the UNet.
Notably, unlike the UNet, the Jacobian eigenvectors of the DiT does not transition from memorization to generalization as the training dataset size increases, indicating that the driving factor of a DiT’s generalization is fundamentally different from a UNet.
This difference calls for a further study about what other inductive biases drive the generalization ability of a DiT?

Figure: Figure 4: Attention maps of DiTs trained with $10$, $10^{3}$, and $10^{5}$ images. All attention maps are linearly normalized to the range $\left[0,1\right]$, with a colormap applied to the interval $\left[0,0.1\right]$ for enhanced visualization. The top-right insets provide a zoomed-in view of the center patch of each attention map. As the number of training images increases, DiT’s generalization improves, and attention maps across all layers exhibit stronger locality. The pink boxes highlight the attention corresponding to a specific output token, obtained by reshaping a single row from the layer-$12$ attention map (original shape: $1{\times}(HW)$) into a matrix of shape $H{\times}W$. As $N$ increases from $10$ to $10^{5}$, the token attentions progressively concentrate around the region near the output token (highlighted with blue boxes).
Refer to caption: x12.png

### 2.3 How Does a DiT Generalize?

The generalization of a DiT may originate from the self-attention (Vaswani, ) dynamics because of its pivotal role in a DiT. Could the attention maps of a DiT provide insights into its inductive biases?
In light of this, we empirically compare the attention maps of DiTs with varying levels of generalization: three DiT models trained with $10$, $10^{3}$, and $10^{5}$ images, where a DiT trained with more images demonstrates stronger generalization.
Specifically, we extract and visualize the attention maps from the self-attention layers of these DiT models as follows,

$$ $\mathrm{Attention\ Map}=\mathrm{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right ),~{}~{}\{Q,K\}\in\mathbb{R}^{\left(HW\right)\times d},$ (5) $$

where $Q$ and $K$ represent the query and key matrices. $H$ and $W$ are the height and width of the input tensor, while $d$ denotes the dimension of a self-attention layer.
For better readability of the attention maps, we linearly normalize each attention map to the range of $\left[0,1\right]$ and apply a colormap to the interval $\left[0,0.1\right]$, i.e., values exceeding the upper bound are clipped at $0.1$.

Fig. 4 shows the attention maps of DiTs with varying levels of generalization on a randomly selected image. Empirically, we observe that the attention maps of a DiT’s self-attention layers remain highly consistent across different images. Further details are provided in Appendix B.
As the number of training images increases from $N{=}10$ to $N{=}10^{5}$, the attention maps of a DiT become increasingly concentrated along several diagonal lines. A closer inspection of the token attentions of a specific target token, i.e., a row in the attention map, shows that these diagonal patterns correspond to tokens near the target token, indicating that the generalization ability of a DiT is linked to the locality of its attention maps.
Delving deeper, can one modify the generalization of a DiT with this inductive bias? We explore this next.

Figure: Figure 5: Global and local attention maps: (a) global attention captures the relationship between the target token and any input token, whereas (b) local attention focuses only on tokens within a nearby window around the target.
Refer to caption: x19.png

## 3 Injecting Inductive Bias by Restricting Attention Windows

To verify that the locality of attention maps enables the generalization of a DiT, we hypothesize that it’s possible to adjust the inductive bias of a DiT by restricting attention windows.
To test this hypothesis, we set up baselines by adopting the diffusion model and DiT implementations from the official repository(^2^22[https://github.com/facebookresearch/DiT](https://github.com/facebookresearch/DiT)) of Peebles & Xie (). Specifically, we remove the auto-encoder and set the patchify size to $1{\times}1$, transforming it into a pixel-space DiT. This modification rules out irrelevant components and ensures more straightforward comparisons in downstream experiments. For model training, we use images of resolution $32{\times}32$, which is equivalent in dimensionality to $512{\times}512$ for a latent-space DiT with a patchify size of $2{\times}2$.

In the remainder of this section, we show that based on the PSNR gap, injecting local attention can effectively modify the generalization of a DiT, often accompanied by an FID change when insufficient training data is used. Furthermore, we discover that placing the attention window restrictions at different locations in a DiT and adjusting the effective attention window sizes allow for additional control over its generalization behavior.
Details w.r.t. experimental settings and more results are deferred to the Appendix A and D.

Figure: Figure 6: PSNR gap$\downarrow$ comparison between a standard DiT and a DiT equipped with local attention for two architectures: (a) DiT-XS/1 and (b) DiT-S/1. Incorporating local attention reduces the PSNR gap consistently across $N{=}10^{3}$, $N{=}10^{4}$, and $N{=}10^{5}$. This advantage is robust across six different datasets and both DiT backbones. In this setup, local attention with window sizes $\left(3,5,7,9,11,13\right)$ is applied to the first six layers of the DiT. Textured bars highlight the default DiT baselines.
Refer to caption: extracted/5952280/images/v2_CelebA_ditxs.png

**Table 1: FID$\downarrow$ comparison between a standard DiT and a DiT equipped with local attention. ^† indicates training with different random seeds, train-test splits, and doubled batch sizes. For the DiT-XS/1 and DiT-S/1 architectures, local attention reduces FID when the DiT’s generalization is not saturated ($N{=}10^{4}$). At $N{=}10^{5}$, local attention achieves comparable or marginally higher FID compared to the standard DiT. These findings are consistent across various datasets, random seeds, train-test splits, and batch sizes. In this setting, local attention with window sizes of $\left(3,5,7,9,11,13\right)$ is applied to the first six layers of the DiT, where both the placement and window size play a crucial role in determining a DiT’s FID result. Further details are provided in Sec. 3.2 and Sec. 3.3.**
| Model | CelebA | ImageNet | LSUN Church | LSUN Bedroom | LSUN Bridge | LSUN Tower |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $N{=}10^{4}$ | $N{=}10^{5}$ | $N{=}10^{4}$ | $N{=}10^{5}$ | $N{=}10^{4}$ | $N{=}10^{5}$ | $N{=}10^{4}$ | $N{=}10^{5}$ | $N{=}10^{4}$ | $N{=}10^{5}$ | $N{=}10^{4}$ | $N{=}10^{5}$ |  |
| DiT-XS/1 | $9.6932$ | $2.6303$ | $52.5650$ | $17.3114$ | $12.8842$ | $5.2927$ | $14.8354$ | $5.4066$ | $23.1771$ | $8.0791$ | $12.5532$ | $4.6619$ |
| w/ Local | $8.4580$ | $2.5469$ | $43.8687$ | $18.0671$ | $10.4794$ | $5.2672$ | $11.9566$ | $5.3542$ | $18.1470$ | $8.3546$ | $10.5644$ | $4.8041$ |
| $-12.74$% | $-3.17$% | $-16.54$% | $+4.37$% | $-18.66$% | $-0.97$% | $-19.40$% | $-0.97$% | $-21.70$% | $+3.41$% | $-15.84$% | $+3.05$% |  |
| DiT-XS/1^† | $10.5432$ | $2.5215$ | $36.8461$ | $20.1907$ | $13.4921$ | $3.9033$ | $15.6740$ | $4.8256$ | $22.0032$ | $7.7771$ | $13.8952$ | $4.1576$ |
| w/ Local^† | $8.4258$ | $2.4988$ | $31.4555$ | $20.3175$ | $10.2708$ | $4.5322$ | $11.2033$ | $5.0868$ | $17.8903$ | $7.7477$ | $10.1938$ | $4.6146$ |
| $-20.08$% | $-0.90$% | $-14.63$% | $+0.63$% | $-23.88$% | $+16.11$% | $-28.53$% | $+5.41$% | $-18.69$% | $-0.38$% | $-26.64$% | $+10.99$% |  |
| DiT-S/1 | $23.2496$ | $2.3278$ | $36.6378$ | $20.6101$ | $14.8826$ | $3.9390$ | $16.1094$ | $4.6086$ | $51.5729$ | $5.7950$ | $28.9727$ | $3.1897$ |
| w/ Local | $20.7768$ | $2.3321$ | $33.1807$ | $20.7972$ | $11.7540$ | $4.4097$ | $11.6833$ | $5.0519$ | $37.6523$ | $5.5825$ | $21.8068$ | $3.5586$ |
| $-10.64$% | $+0.18$% | $-9.44$% | $+0.91$% | $-21.02$% | $+11.95$% | $-27.48$% | $+9.62$% | $-26.99$% | $-3.67$% | $-24.73$% | $+11.57$% |  |
| DiT-S/1^† | $14.1763$ | $2.5061$ | $37.3477$ | $20.4165$ | $15.4509$ | $4.2317$ | $15.5820$ | $4.8336$ | $24.4374$ | $7.3170$ | $14.8695$ | $4.4495$ |
| w/ Local^† | $11.1046$ | $2.6598$ | $33.1323$ | $20.6006$ | $11.4956$ | $4.5546$ | $11.3673$ | $5.0552$ | $20.3403$ | $7.5565$ | $12.3236$ | $4.4927$ |
| $-21.67$% | $+6.13$% | $-11.29$% | $+0.90$% | $-25.60$% | $+7.63$% | $-27.05$% | $+4.58$% | $-16.77$% | $+3.27$% | $-17.12$% | $+0.97$% |  |

### 3.1 Attention Window Restriction

Local attention, initially proposed to enhance computational efficiency (Liu et al., ; Yang et al., ; Hatamizadeh et al., ; Hassani et al., ), is a straightforward yet effective way to modify a DiT’s generalization. Different from global attention which enables a target token to connect with all input tokens (Fig. 5 (a)), local attention only permits a target token to attend within a small nearby window. The resulting attention map structure is depicted in Fig. 5 (b).
Notably, a local attention constrains the attention map to a sparse activation pattern only along the diagonal direction, thereby enforcing locality of the attention map. The resulting attention map patterns produced by a local attention align well with the inductive bias that a DiT exhibits when observing a strong generalization ability, as illustrated in Fig. 4 (row $N{=}10^{5}$).

Using local attentions in a DiT can consistently improve its generalization (measured by PSNR gap) across different datasets and model sizes.
Specifically, we consider a DiT model with $12$ DiT blocks, and replace the first $6$ global attention layers with local attentions, whose window sizes range from $3{\times}3$ to $13{\times}13$ with a stride of $2$.
We train both the vanilla DiT and a DiT equipped with local attentions with $N{=}10^{3},10^{4}$ and $10^{5}$ images for the same $400$k training steps. Then we calculate the PSNR gap between the training and testing images for models trained with different amounts of images.
In Fig. 6, we show the PSNR gap comparison between a DiT with and without local attentions on CelebA, ImageNet, and LSUN (Church, Bedroom, Bridge, Tower) datasets, using baseline DiT models of two sizes (DiT-XS/1 and DiT-S/1). Notably, using local attentions reduces a DiT’s PSNR gap with different amounts of training images. Importantly, the advantage of local attention is robust across different training datasets and backbone sizes.

For a discriminative model, e.g., a classifier, better generalization typically leads to better model performance when the training dataset is insufficient.
Is this also the case for generative models like a DiT?
To investigate, we compare the FID between the default DiT and a DiT using local attentions. For each dataset, we compare FID values of models trained with $10^{4}$ and $10^{5}$ images: the former represents the case of insufficient training images while the later case refers to use of sufficient training data.
Tab. 1 shows the FID comparison among the same six datasets and two DiT backbones used when comparing PSNR gaps.
Improving the generalization via local attentions can indeed improve the FID when $N{=}10^{4}$, which is in line with observations from discriminative models. When $N{=}10^{5}$, using the presented approach of adding local attentions either results in comparable FID values or experiences a slight compromise. Interestingly, we find that modifying the placement and effective attention window size permits fine-grained control of a DiT’s generalization and generation quality. More discussions are in Sec. 3.2 and Sec. 3.3 below.

In light of Occam’s razor, reducing the model parameter count has been shown to be yet another possible strategy to inject an inductive bias. This differs from the attention window restrictions considered above, as local attentions reduce the FLOPs of a DiT without changing the model parameter count. In contrast, to inject an inductive bias by reducing the parameter count of a DiT, we explore sharing parameters of a DiT’s attention blocks as well as modifying a DiT’s attention layers to learn the coefficients of pre-computed offline PCA components. Neither of these methods shows as compelling improvements of the generalization (measured via the PSNR gap) as using local attentions. We provide more details regarding the considered techniques in the Appendix C.

Figure: Figure 7: PSNR gap$\downarrow$ comparison for different local attention placement patterns. We find that placing local attention in the early layers (head) results in a smaller PSNR gap compared to mixing local and global attention (mix) or applying local attention in the later layers (tail). The latter two configurations may even perform worse than the vanilla DiT.
Refer to caption: extracted/5952280/images/CelebA_ditv2v5v10.png

### 3.2 Placement of Attention Window Restriction

Given the same set of local attentions, placing them at different layers of a DiT leads to different results.
For local attention, we study three placement schemes: 1) placing local attentions on the early layers of a DiT, 2) interleaving local attentions with global attentions, and 3) placing local attentions on the tail layers of a DiT.
In Fig. 7, we compare the PSNR gap for the three aforementioned placement schemes on the CelebA and ImageNet datasets, using two distinct local attention configurations. Specifically, *Local* refers to a setting with $6$ attention layers, where the window sizes vary from $3{\times}3$ to $13{\times}13$ with a stride of $2$, which is consistent with the local attention configuration used in Fig. 6 and Tab. 1 above. Meanwhile, *Local^∗* represents a different configuration consisting of $9$ local attention layers, arranged as $\left(3^{*3},5^{*3},7^{*3}\right)$, where $i^{*j}$ indicates repeating a local attention layer with a $\left(i{\times}i\right)$ window $j$ times.

The results in Fig. 7 indicate that applying local attention in the early layers of a DiT consistently leads to a smaller PSNR gap across different training data sizes. Additionally, the FID results in Tab. 3.1 demonstrate that the first placement scheme generally improves FID when the training data is limited ($N{=}10^{4}$).
In contrast, interleaving local and global attention, or applying local attention to the tail layers, enhances the model’s data-fitting ability but often compromises generalization. These two placement schemes tend to improve FID when $N{=}10^{5}$, though this improvement comes at the cost of reduced FID when $N{=}10^{4}$, further supporting the generalization results as measured by the PSNR gap.

Figure: Figure 8: PSNR gap$\downarrow$ changes when the effective attention window size is kept constant, decreased, or increased. Reducing the window size results in a smaller PSNR gap, indicating improved generalization.
Refer to caption: extracted/5952280/images/CelebA_ditv9v12.png

### 3.3 Effective Attention Window Size

Adjusting the effective attention window size provides an additional mechanism to control the generalization of a DiT. Specifically, our analysis reveals that smaller attention windows lead to stronger generalization, while larger windows enhance data fitting, typically at the cost of generalization. Furthermore, maintaining the total attention window size but altering the distribution across local attentions generally preserves the overall behavior of a DiT. These observations are based on an empirical study using the CelebA and ImageNet datasets, involving three paired comparisons of local attention configurations. The PSNR gap and FID results are shown in Fig. 8 and Tab. 3.2, respectively.

Specifically, in the first comparison, we apply two configurations of local attentions with window sizes $\left(5,5,5,5,5,5\right)$ and $\left(3,3,5,5,7,7\right)$ to the first six layers of a DiT. We observe that altering the attention window size distribution, while keeping the total window size fixed, has a limited impact on a DiT’s generalization, as indicated by the similar PSNR gaps across $N{=}10^{3}$, $10^{4}$, and $10^{5}$. This similarity in generalization is further corroborated by their comparable FID values.
In the second and third comparisons, using the DiT-XS/1 configurations with *Local* and *Local^∗* attention settings, we find that reducing the attention window size enhances generalization, while increasing the window size diminishes it. This is evidenced by a decrease in the PSNR gap for smaller window sizes and an increase for larger ones. Furthermore, the improved generalization is associated with better FID values under comparably insufficient training data, and vice versa.

## 4 Related Work

#### Inductive Biases of Generative Models.

Current diffusion models (Sohl-Dickstein et al., ; Song et al., ; Ho et al., ; Kadkhodaie & Simoncelli, ; Nichol & Dhariwal, ; Song et al., ; An et al., ) exhibit strong generalization abilities (Zhang et al., ; Keskar et al., ; Griffiths et al., ; Wilson & Izmailov, ), relying on inductive biases (Mitchell, ; Goyal & Bengio, ). Prior to the emergence of diffusion models,  Zhao et al. () show that generative models like GANs (Goodfellow et al., ) and VAEs (Kingma, ) can generalize to novel attributes not presented in the training data.
This generalization ability of generative models is possibly due to the inductive biases (Zhang et al., ; Keskar et al., ) introduced by model design and training. Following this line of research, Kadkhodaie et al. () show that the generalization of diffusion models arises due to geometry-adaptive harmonic bases (Mallat et al., ). However, their work only studies the generalization of a simplified one-channel UNet. It remains unclear whether their study can be generalized to commonly used three-channel UNets (Nichol & Dhariwal, ) and more compelling DiTs (Peebles & Xie, ). This work fills this gap and reveals that a classic UNet still exhibits the harmonic bases but a DiT does not. Further studies show that a DiT’s generalization is associated with a different inductive bias: locality of attention maps.

#### Attention Window Restrictions.

Prior studies have shown that restricting attention windows through mechanisms such as local attention (Beltagy et al., ; Liu et al., ; Hassani et al., ), strided attention (Wang et al., ; Xia et al., ), and sliding attention (Pan et al., ), among others, can significantly improve the efficiency of attention computation (Yang et al., ; Hatamizadeh et al., ; Hassani et al., ; Wang et al., ). These techniques limit the attention scope, reducing computational complexity while retaining the model’s ability to capture important contextual information. However, our work explores a different direction by investigating how attention window restrictions, especially through local attention, affect the generalization properties of DiTs. We show that beyond efficiency gains, local attention can be used to modulate the model’s generalization by enforcing the inductive bias of locality within attention maps.

## 5 Conclusion

This paper investigates the inductive biases that facilitate the generalization ability of DiTs. For insufficient training data, we observe that DiTs achieve superior generalization, as measured by the PSNR gap, compared to UNets with equivalent FLOPs. However, unlike simplified and standard UNet-based diffusion models, DiTs do *not* exhibit geometry-adaptive harmonic bases. Motivated by this discrepancy, we explore alternative inductive biases and identify that a DiT’s generalization is instead influenced by the locality of its attention maps. Consequently, we effectively modulate the generalization behavior of DiTs by incorporating local attention layers. Specifically, we demonstrate that varying the placement of local attention layers and adjusting the effective attention window size enables fine-grained control of a DiT’s generalization and data-fitting capabilities. Enhancing a DiT’s generalization often leads to improved FID scores when trained with insufficient data.

## Appendix A Experimental Settings

### A.1 Model Definition

We verify the effectiveness of local attention in modifying the generalization of a DiT using $2$ DiT backbones and $10$ local attention variations. Tab. 4 provides more details about the DiT backbones and local attention configurations.

**Table 4: DiT architectures and local attention settings. In the column titled ‘DiT Blocks’, $G$ denotes global attention while a number $k$ represents a local attention of window size $k{\times}k$.**
| Model | DiT | Hidden | Num | Depth | Patch |
| --- | --- | --- | --- | --- | --- |
| Blocks | Size | Heads | Size |  |  |
| DiT-XS/1 | $\left(G,G,G,G,G,G,G,G,G,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local | $\left(3,5,7,9,11,13,G,G,G,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local (mix) | $\left(3,G,5,G,7,G,9,G,11,G,13,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local (tail) | $\left(G,G,G,G,G,G,3,5,7,9,11,13\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local (smaller win) | $\left(3,5,7,9,11,13,15,17,19,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local^∗ | $\left(3,3,3,5,5,5,7,7,7,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local^∗ (mix) | $\left(3,3,3,G,5,5,5,G,7,7,7,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local^∗ (tail) | $\left(G,G,G,3,3,3,5,5,5,7,7,7\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local^∗ (larger win) | $\left(5,5,5,7,7,7,9,9,9,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local Attn $(5^{*6})$ | $\left(5,5,5,5,5,5,G,G,G,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| w/ Local Attn $(3^{*2},5^{*2},7^{*2})$ | $\left(3,3,5,5,7,7,G,G,G,G,G,G\right)$ | $252$ | $4$ | $12$ | $1{\times}1$ |
| DiT-S/1 | $\left(G,G,G,G,G,G,G,G,G,G,G,G\right)$ | $384$ | $6$ | $12$ | $1{\times}1$ |
| w/ Local | $\left(3,5,7,9,11,13,G,G,G,G,G,G\right)$ | $384$ | $6$ | $12$ | $1{\times}1$ |

To elaborate, we adopt two DiT backbones, DiT-XS/1 and DiT-S/1, to verify the effectiveness of attention window restrictions in modifying the generalization of a DiT. Both models have $12$ DiT Blocks. We remove the auto-encoder and use a patch size of $1{\times}1$. Note, DiT-XS/1 has a hidden size of $252$ and uses $4$ attention heads. In contrast, DiT-S/1 has a hidden size of $384$ and uses $6$ attention heads. Regarding the local attention variations, the default setting *Local* combines $6$ local attentions of window size $\left(3,5,7,9,11,13\right)$ and $6$ global attentions. Meanwhile, *Local^∗* is a variant using $9$ local attentions of window size $\left(3,3,3,5,5,5,7,7,7\right)$. For both *Local* and *Local^∗* settings, we place local attentions at the heading layers of a DiT. We also study interleaving the local and global attentions as well as placing local attentions at tailing layers of a DiT, leading to (mix) and (tail) variants. Additionally, to study the effects of modifying the attention window size, we decrease the attention window size of the *Local* model and increase the attention window size of the *Local^∗* model, resulting in (smaller win) and (larger win) models in Tab. 4.

### A.2 Training and Sampling Setting

The implementation of the UNet(^3^33[https://github.com/openai/improved-diffusion](https://github.com/openai/improved-diffusion)) and the DiT(^4^44[https://github.com/facebookresearch/DiT](https://github.com/facebookresearch/DiT)) are based on the official repositories of Nichol & Dhariwal () and  Peebles & Xie (). Specifically, for the UNet, we use the architecture which has a 4-stage encoder with channel multipliers of $\left(1,2,3,4\right)$. For each stage, we include $3$ ResBlocks. At the end of each stage, the resolution of the input tensor are down-sampled by a factor of $2$. In the last stage, we use one layer of self-attention. The decoder mirrors the encoder layers and places them in the reverse order, replacing down-sampling layers with up-sampling ones. Between the encoder and decoder, there are $2$ ResBlocks and $1$ self-attention layer by default. The default skip connections are used between the encoder and decoder at the same resolution. Consequently, the UNet has $303.17$G FLOPs and $109.55$M parameters. The FLOPs of this UNet are nearly identical to the DiT-XS/1 model in Tab. 4.

All DiT and UNet models are trained with the same hyper-parameter settings. Concretely, we train each model in the pixel-space using a resolution of $32{\times}32$. All networks are using the same diffusion algorithm: diffusion steps of $1000$ in training and $250$ in sampling, and predicting the added noise and sigma simultaneously. To train a network, we use the random seed $43$, learning rate $1e^{-4}$ and an overall batch size of $64$. All networks are trained with $8$ or $4$ A100/H100 GPUs, using the EMA checkpoint at train step $400k$ with EMA decay $0.9999$. For each dataset, we first randomly shuffle the whole dataset. Then we choose the last $N{=}10,10^{3},10^{4}$ and $10^{5}$ images as the training set. The train-test split of a dataset is kept consistent for different networks. When computing FID values for the model trained with $N{=}10^{4}$ and $10^{5}$ images, we randomly select $M{=}\min\{N,50k\}$ of the training images as the reference set. In Tab. 1 of the main manuscript, we present the results of DiT-XS/1^† and DiT-S/1^† models, with and without local attentions. For the ^† models, we use a different dataset shuffling, change the random seed to $143$, and double the training batch size to $128$. Notably, using local attention on both normal and ^† models can successfully modify the generalization of a DiT, confirming the effectiveness of the locality as the inductive bias of a DiT.

### A.3 PSNR Computation

We compute the PSNR based on a training or testing subset of $300$ images following Kadkhodaie et al. (). For each image, we re-space the diffusion steps from $1000$ to $50$, and compute the train and test PSNR on each step. Specifically, we first perform the noising step of the diffusion model to get the noisy image at a diffusion step $t$. Next, we feed the noisy image into the diffusion model backbone and get the estimation of the added noise, which is then used to recover the clean image from the training or testing subset, i.e., performing a one-step denoising. The final PSNR at step $t$ is obtained using the estimated clean image and the ground truth. Consequently, the PSNR value can estimate a diffusion model’s accuracy at each diffusion step. Therefore, the PSNR gap between the training and testing subsets can measure a diffusion model’s generalization: when a diffusion model has good generalization, its prediction accuracy should be comparably between the training and testing set, resulting in a small PSNR gap.

### A.4 Eigenvector of Jacobian Matrix

We follow Kadkhodaie et al. () to obtain the geometry-adaptive harmonic bases through eigendecomposition of a diffusion model’s Jacobian matrix.
Specifically, we first feed a noisy image into a DiT or a UNet and obtain their Jacobian matrices, where each entry of the Jacobian

$$ $\nabla\bm{\epsilon}_{\theta}=\begin{bmatrix}\frac{\partial\hat{{\epsilon}}_{1} }{\partial x_{1}}&\frac{\partial\hat{{\epsilon}}_{1}}{\partial x_{2}}&\cdots& \frac{\partial\hat{{\epsilon}}_{1}}{\partial x_{L}}\[4.30554pt] \frac{\partial\hat{{\epsilon}}_{2}}{\partial x_{1}}&\frac{\partial\hat{{ \epsilon}}_{2}}{\partial x_{2}}&\cdots&\frac{\partial\hat{{\epsilon}}_{2}}{ \partial x_{L}}\[4.30554pt] \vdots&\vdots&\ddots&\vdots\[4.30554pt] \frac{\partial\hat{{\epsilon}}_{L}}{\partial x_{1}}&\frac{\partial\hat{{ \epsilon}}_{L}}{\partial x_{2}}&\cdots&\frac{\partial\hat{{\epsilon}}_{L}}{ \partial x_{L}}\end{bmatrix},\quad\hat{\bm{\epsilon}}=\bm{\epsilon}_{\theta}( \bm{x},t),~{}~{}~{}\bm{x},\hat{\bm{\epsilon}}\in\mathbb{R}^{L\times d},~{}~{}~{}$ (6) $$

represents the partial derivative of an output pixel w.r.t. an input pixel.
Here, $L=H\times W$ denotes the overall number of pixels for the UNet and the DiT.
Next, we perform the eigendecomposition on $\nabla\mathbf{\epsilon}_{\theta}$ to obtain its eigenvectors and eigenvalues. The eigenvectors are considered to be the harmonic bases if it exists.

## Appendix B Attention Map Consistency

Figure: Figure 9: Attention maps of the $1^{\rm st}$ DiT Block produced by feeding $6$ random images to DiT-XS/1 models.
Refer to caption: x23.png

Figure: Figure 10: Attention maps of the $6^{\rm th}$ DiT Block produced by feeding $6$ random images to DiT-XS/1 models.
Refer to caption: x24.png

Figure: Figure 11: Attention maps of the $12^{\rm th}$ DiT Block produced by feeding $6$ random images to DiT-XS/1 models.
Refer to caption: x25.png

To verify the robustness of the discovered inductive bias of a DiT, i.e., the locality of attention maps, we obtain the attention maps corresponding to distinct input images and compare them visually. Specifically, we show the attention maps of the $1^{\rm st}$, $6^{\rm th}$, and $12^{th}$ self-attention layers in Fig. 9, Fig. 10, and Fig. 11, respectively, using randomly selected $6$ input images from the CelebA (Liu et al., ) dataset. In these figures, from top to bottom, each row is the attention map of a DiT model trained with $N{=}10,10^{3},10^{4}$, and $10^{5}$ images. Meanwhile, each column is related to an input image. Interestingly, we find that the attention maps of a DiT’s self-attention layers demonstrate a consistent pattern among different input images, suggesting that the attention maps of a DiT, after training, are part of its inductive biases rather than being mostly governed by a specific input image.

To verify the robustness of the discovered inductive bias of a DiT, i.e., the locality of attention maps, we obtain the attention maps corresponding to distinct input images and compare them visually. Specifically, we show the attention maps of the $1^{\rm st}$, $6^{\rm th}$, and $12^{th}$ self-attention layers in Fig. 9, Fig. 10, and Fig. 11, respectively, using randomly selected $6$ input images from the CelebA (Liu et al., ) dataset. In these figures, from top to bottom, each row is the attention map of a DiT model trained with $N{=}10,10^{3},10^{4}$, and $10^{5}$ images. Meanwhile, each column is related to an input image. For a better visualization, we use the logarithm normalization on attention maps before applying a colormap. For the same DiT Block, attention maps of different images demonstrate a similar pattern. Interestingly, we find that the attention maps of a DiT’s self-attention layers demonstrate a consistent pattern among different input images, suggesting that the attention maps of a DiT, after training, are part of its inductive biases rather than being mostly governed by a specific input image.

## Appendix C Reducing Parameters of a DiT

**Figure 12: The PSNR (a) and PSNR gap (b) comparison. Taking DiT-XS/1 as the baseline (row 1), using local attention (row 2) can achieve decent PSNR gap improvement. In contrast, using parameter reduction approaches: weight sharing (row 3) and PCA (row 4), hardly achieve a PSNR gap improvement (row 3) or even make it worse then the baseline (row 4).**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| w/ Local |  |  |
| w/ Weight Sharing |  |  |
| w/ PCA |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

**Table 5: FID$\downarrow$ comparison between reducing DiT FLOPs and parameter size. Using parameter sharing is not as effective in reducing FID as the local attention. Using PCA will make the FID worse when $N{=}10^{4}$. Both are not as effective as using local attention in improving the FID with $10^{4}$ training images.**
| Model | CelebA |  |
| --- | --- | --- |
| $N{=}10^{4}$ | $N{=}10^{5}$ |  |
| DiT-XS/1 | $9.6932$ | $2.6303$ |
| w/ Local | $8.4580$ | $2.5469$ |
| $12.74$% | $3.17$% |  |
| w/ Weight Sharing | $8.7819$ | $2.5802$ |
| $9.40$% | $1.90$% |  |
| w/ PCA | $11.3872$ | $2.5482$ |
| $+17.48$% | $3.12$% |  |

Figure: Figure 13: Principal components extracted from attention maps of different DiT Blocks. Based on a DiT-XS/1 model trained with $N{=}10^{5}$ data from ImageNet, we perform PCA on attention maps of its first three layers, using $2048$ images, resulting in a total of $24576$ attention maps. For a better visualization, we adopt the logarithm normalization to principal components before applying colormaps.
Refer to caption: x26.png

We study two parameter reduction approches for a DiT: parameter sharing and composing attention maps with PCA.

#### Parameter Sharing.

For this approach, the $2^{\rm nd}$ and $4^{\rm th}$ DiT Blocks reuse the parameters of the $1^{\rm st}$ and $3^{\rm rd}$ DiT Blocks, respectively, leading to a DiT model with the same FLOPs but fewer parameters. Fig. 12 (row 3) shows the PSNR and PSNR gap of the DiT using this parameter sharing approach. Meanwhile, Tab. 5 demonstrates the FID of the same model trained with $N{=}10^{4}$ and $10^{5}$ images. Notably, sharing parameters in a DiT can reduce the PSNR gap, resulting in FID improvement when $N{=}10^{4}$. However, it is not as effective as using local attention in modifying the generalization of a DiT.

#### Composing Attention Maps with PCA.

Another parameter reduction approach we explore is composition of attention maps of a DiT with PCA. To elaborate, we first collect the attention maps of $2048$ images. Particularly, we noise each image with $10,25$ and $40$ steps and obtain the attention maps of all attention heads corresponding to these noisy images, where the total diffusion step is re-spaced to $50$. Taking the DiT-XS/1 model as an example, we collect a total of $24,576=2048\ (\rm images){\times}3\ (\rm diffusion\ steps){\times}4\ (\rm
attention
\ heads)$ attention maps. We use the DiT-XS/1 model trained with the ImageNet (Deng et al., ) dataset and collect attention maps of a DiT’s first three self-attention layers using randomly selected $2048$ images from the testing set of the same dataset. Next, we compute the principal components of each self-attention layer from the corresponding attention maps. We use the low rank PCA function(^5^55[https://pytorch.org/docs/stable/generated/torch.pca_lowrank.html](https://pytorch.org/docs/stable/generated/torch.pca_lowrank.html)) of PyTorch and obtain the first $50$ principal components, where each principal component has the same size as the attention map.
Fig. 13 shows the principal components and the corresponding coefficients for the first two DiT Blocks. Notably, we find that PCA is effective in capturing the dominant diagonal patterns that indicate the locality in a DiT’s attention map. Finally, we use the principal components of the attention maps to reduce the parameters of a DiT. Concretely, we replace the two MLP layers that map the input tensor to query and key matrices by a smaller MLP mapping the input tensor to $50$ coefficients for each principal component (PC). Then the new attention map is obtained as follows:

$$ $\mathrm{Attention\ Map}=\mathrm{Coefficients}\odot\mathrm{PCs}+\delta,$ (7) $$

where $\odot$ denotes matrix multiplication while $\delta=\frac{1.0}{1024}$ is used to ensure that the attention weights for a specific token sum to $1$. We replace the normal attention maps of a DiT’s first three self-attention layers with the attention maps composed with principal components following Eq. (7).
According to the PSNR and PSNR gap comparison in Fig. 12 (row 4) as well as the FID comparison in Tab. 5, reducing parameters of DiT by composing its first three attention maps with principal components cannot reduce a DiT’s PSNR gap, leading to a worse FID when $N{=}10^{4}$.

## Appendix D Raw PSNR and PSNR Gap Curves

We show the raw PSNR and PSNR gap curves corresponding to models in Fig. 6 of the main manuscript in Fig. 14 to Fig. 19. Besides the PSNR gap change after using local attentions as discussed in Sec. 3.1 of the main manuscript, given sufficient images ($N{=}10^{5}$) of a specific dataset, the training PSNR curves and the average PSNR are very similar between a DiT with and without using local attentions, suggesting that using local attention can mostly maintain a DiT’s dataset fitting ability. This is likely because all the important information of a DiT’s attention map is retained.

**Figure 14: The PSNR (a) and PSNR gap (b) comparison on the CelebA (Liu et al., ) dataset.**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| DiT-XS/1 w/ Local |  |  |
| DiT-S/1 |  |  |
| DiT-S/1 w/ Local |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

**Figure 15: The PSNR (a) and PSNR gap (b) comparison on the ImageNet (Deng et al., ) dataset.**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| DiT-XS/1 w/ Local |  |  |
| DiT-S/1 |  |  |
| DiT-S/1 w/ Local |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

**Figure 16: The PSNR (a) and PSNR gap (b) comparison on the LSUN (Church) (Yu et al., ) dataset.**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| DiT-XS/1 w/ Local |  |  |
| DiT-S/1 |  |  |
| DiT-S/1 w/ Local |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

**Figure 17: The PSNR (a) and PSNR gap (b) comparison on the LSUN (Bedroom) (Yu et al., ) dataset.**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| DiT-XS/1 w/ Local |  |  |
| DiT-S/1 |  |  |
| DiT-S/1 w/ Local |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

**Figure 18: The PSNR (a) and PSNR gap (b) comparison on the LSUN (Bridge) (Yu et al., ) dataset.**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| DiT-XS/1 w/ Local |  |  |
| DiT-S/1 |  |  |
| DiT-S/1 w/ Local |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

**Figure 19: The PSNR (a) and PSNR gap (b) comparison on the LSUN (Tower) (Yu et al., ) dataset.**
| DiT-XS/1 |  |  |
| --- | --- | --- |
| DiT-XS/1 w/ Local |  |  |
| DiT-S/1 |  |  |
| DiT-S/1 w/ Local |  |  |
|  | (a) PSNR Curves of Training and Testing Sets | (b) PSNR Gap |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
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?
soundness: 2
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.
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.
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?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
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.
soundness: 3
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.
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.
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

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 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
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).
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

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the Weaknesses.
soundness: 2
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.
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.
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.

Meta-review: 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.

Decision: Reject