On Inductive Biases That Enable Generalization in Diffusion Transformers
Authors: Jie An
De Wang
Pengsheng Guo
Jiebo Luo
Alex Schwing
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion Model
Diffusion Transformer
Generalization
Inductive Bias
Primary area: deep_learning
TL;DR: The generalization of a DiT is influenced by the inductive bias of attention locality rather than harmonic bases like UNet. Using attention window restrictions can modify its generalization ability.

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.

On Inductive Biases That Enable Generalization of Diffusion Transformers
Jie An De Wang', Pengsheng Guo', Jiebo Luo?, Alexander G. Schwing'
Apple, University of Rochester [jan6, jluo)[cs rochester edu [de_wange peweanggenggennggeee guo, ag_ ag wew.swhwnn@@aappe..oom
Abstract
Recent work studying the generalization of diffusion models with locally linear UNet-based denoisers reveals inductive biases that be expressed via geometry- adaptive harmonic bases. For such locally linear UNets, these geometry-adaptivy harmonic bases be conveniently visualized through the sigen-decompositioon of a UNet's Jacobian matrix. In practice, however, more recent denoising networks are often transformer-basee 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, 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 DiT can improve both generalization and generation quality when less training data is available. Source code available at hattps://gitae http/DiT-GGeeralizzee 2at.ra/iiaaaaaoo//ooeeeeeaaaaaa
Introduction
Diffusion models have achieved remarkable success in visual content generation. Their training involves approximating a distribution in high-dimensional space from a limited number of training samples-a task that demanding due the curse of dimensionality. Nonetheless, recent diffusion
models (Song et al.] 2020. Ho et al. 2020) learn to generate high-quality images (Nichol al.
2021 Dhariwal
Nichol 2021 I Saharia et al. 2022 Rombach ct al. 2022 Chen et al. 2023
2024a and even videos/audio (Singer et al. 2022 Ho et al. 2022 Girdhar et al. 2023 Blattmann
et al. 2023 OpenAl 2024 Cheng et al. 2025) using relatively few samples when compared to the dimensionality of the underlying space. This indicates that diffusion models exhibit powerful
inductive biases (Wilson & Izmailov 2020, Goyal & Bengio 2022) Grifitts et al.) (024) that promote
effective generalization. What exactly are these inductive biases? Answering this question is crucial for understanding the behavior of diffusion models and their generalization.
Recent work by Kadkhodaie et al. (004)) locally linear single-channel UNet-based diffusion models reveals that the strong generalization of UNet-based denoisers is driven by inductive biases that can be expressed via set of ceometry-adaptic harmonic bases (Mallat al.] 2020). For a UNet that has been modified to be locally linear, such harmonic bases can be extracted via the eigenvectors
Work done during internship at Apple.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Clean
1.731
Naisy
Noiey
Auu
Apza
Appe
The Largeot Elgenwolue
Elpansectors
(a) Locally Linear UNet Kadkhodaie et al. 2024p
(b) Hybrid UNet
(c) DiT
Nichol & Dhariwal 2021
Peebles & Xie 2023
Figure Jacobian eigenvectors of (a) locally linear channeanee UNet. (b) the hybrid UNet
introduced in improved diffusion (Nichol & Dhariwal, 2021). and (c) DiT (Peebles & Xie 2023). Kadkhodaie et al. (2024) find that the generalization
UNet-based diffusion models is driven by
geometry adaptive harmonic bases (a), which exhibit oscillatory patterns whose frequency increases as the eigenvalue 0g decreases. For hybrid UNets (Nichol & Dhariwal 202T), due the inclusion of non- linear operations such softmax in transformer blocks and normalization layers in both transformer and convolutional layers, the harmonic bases extracted from their split-channel eigenvectors (b) do not adapt well to the input geometry, though oscillatory patterns still persist. In contrast, the harmonic bases completely disappear in a DiT (Peebles & Xie 2023) as shown in (c), indicating that the bigen-dccompositine motivates us to investigate the inductive biases a DiT that enable its generalization. The RGB channels of the split-channel eigenvectors outlined with red,, ggreen, and blue boxes, respectively. All models operate directly the pixel space without the patchify operation.
analysis is longer valid for transformer based DiTs. This observation
of the denoiser's Jacobian matrix, as shown in Fig. Ta). Extending the eigen-decomposition analysis of Kaakhodaie et al. (2024) to more complex, classic multi-channel UNets shows that geometry- adaptive harmonic bases become harder to observe. As illustrated in Fig. [[.). these eigenvectors do not adapt well to the input geometry, although oscillatory patterns whose frequencies increase as the eigenvalues A decrease still exist. This because modern UNets (Nichol Dhariwal 2021) adopt hybrid architectures that incorporate several transformer layers, where the softmax in attention and normalization layers, present both convolutional and transformer blocks, degrade the network's 's local linearity. For DiT. the eigenvectors of its Jacobian matrix show geometry-adaptis harmonic bases, as shown in Fig. [tc. This does not necessarily imply that a DiT fails to capture adaptivv harmonic structures, but rather suggests that the eigen- ggge-ddccompoittonn analysis of Kadkhodaie et al. (2024) isn't 'tn't applicable for probing the inductive bias DiT models. Motivated by this issue, seek alternative approaches address the question: what inductive biases enable the strong generalization ability DiTs?
Answering this question is particularly important because of the recent growing adoption of DiTss (Chen al.) 2024b; Esser et al., 2024, Cheng et al.] 2025; Chen al., 2025), partly for its observed performance at scale (Peebles & Xie 2023). In a new study in this paper, using the PSNR gap (Kadkhodaie et al., 2024) as a metric to evaluate the generalization of diffusion models, confirm that DiT indeed exhibits better generalization than UNet with the same FLOPs. Yet, this observation alone doesn't reveal the inductive biases which enable generalization.
The generalization mechanism of a DiT can be determined by inductive biases introduced by the diffusion model theory, training objectives, target optimal score functions, and network architectures. Prior works Zhang et al. 2024 Niedoba et al., 2024 Li et al. 2024 Wang & Vastola 2024) reveal that the inductive biases enabled by diffusion model theory, training objectives, and target optimal score functions can be similar between a UNet and DIT. However, the inductive bias driven by the model architecture, differs between a UNet and DIT, potentially due to the self-attention (Vaswani 2017) dynamics which are pivotal in DiT models but not in UNets. In a self-attention layer, the

---

PSNR Gap Timestep
3060, PSTR 20.39
1000 PSW (OOOL des PURK
1950 doy Cec: 55
10E6S PORR 27-7
1933, Ary copi 3.50
Outpur PSNAA
12 x
Standard Devistion Range
UNet
Ourput PSWR
2
3
Diftuston
Diffusion
PSNR Timescep
to 1000, Aug P2S0: 24.30
1o00, Aop PSN 54
5030 dey Cas: 7.49 Avy Con ta
Oupput PSNA
Standord Deviation Range
DiT
Ingut Diffusion Stop.t
Difusion Step:
Input Difusion Stop::
(a) PSNR Curves of Training and Testing Sets
(b) PSNR Gap
Figure 2: PSNR (a) and PSNR gap (b) comparison between a UNet and DiT with the same FLOPs for different training image quantities (N). All curves are averaged over three training runs using different dataset shuffles. The standard deviations, illustrated by the curve shadows in the zoomed-in windows, are negligible, indicating minimal variation.
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 find that locality of the attention maps closely tied to its generalization ability. Specifically, the attention maps of a DiT trained with insufficient images, i.e., a DiT with weak generalization, exhibit more -ivarrannnn pattern, especially in early layers: the output tokens of self-attention layer are largely influenced by 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.
If the above finding is true, restricting the attention window in self-attention layers should permit us to modify a DiT's s generalization. Indeed, we find that employing local attention windows (Beltagy eet al. 2020; Hassani et al. 2023) 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 DiT's generalization. Our experiments show that placing attention window restrictions in the early attention layers of the DiT architecture has most impact. Results on CelebA (Liu et al., 2015), ImageNet (Deng et al., 2009), MSCOCO (Lin et al. 2014), and LSUN (Yu et al. 2015) (bedroom, church, tower, bridge) data reveal that applying attention window restrictions modifies generalization, as reflected by a reduced PSNR gap. We also observe improved FID (Heusel et al. 2017). Inception Score (IS) (Barratt & Sharmal 2018). and FD-DINOv2 (Qqua et al. (223)) when training with insufficient data, confirming that a 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 key inductive bias contributing to the generalization of DIT, and 2) we demonstrate how to control this inductive bias by incorporating local attention windows into early layers of DIT. Enhancing the locality in attention computations effectively modifies a DiT's generalization, resulting lower PSNR gap and improved FID, IS, and FD-DINOv2 scores when insufficient training images available for training.
3

---

2 Inductive Bias Analysis of Diffusion Models
Diffusion models are designed to map a Gaussian noise distribution to a dataset distribution. To achieve this, diffusion models can be formulated to estimate the noise c that was used to compute the corrupted image 3; by perturbing the training sample zo following noise schedule depending on step t. The loss function of diffusion model training hence reads as follows:
C = Eange.t [|ee (a 1)][]]
Here, ce(+) represents the backbone network with trainable parameters 0, which plays crucial role in diffusion model gencralization. In this section, we first compare the generalization ability of a DiT (Peebles & Xie 2023) and UNet (Nichol & Dhariwal 2021), two of the most popular diffusion model backbones. Subsequently, we investigate the inductive biases that drive their generalization.
2.1 Comparing DiT and UNet Generalization
We compare the generalization of pixel-space DiT and UNe using metric the PSNR gap proposed by Kadkhodaic et al. (2024). The PSNR gap at diffusion step t, denoted as Gap (t), is the zero-truncated difference between the training set PSNR and the testing set PSNR at step t:
Gap (t) max (PSNB (t) - PSNPS test (t).0),
(2)
where PSN (t) and PSNR (t) are obtained following Kadkhodaie et al. (024). To elaborate, given K images from either training or testing set, we first feed noisy images at step diffusion models and obtain the estimated noise é. Next, we compute the one-step denoising result 2o via
i = xz orê,
where or is defined by the diffusion model noise schedule. Finally, we derive the training and testing PSNRs at diffusion step as follows:
1 K
K Σ
M MSE (#t xx)
PSNR (t)
10- log
Here, E denotes the estimate of image k. obtained by using Eq. ( M denotes the intensity range
of æo. which is 2 since images are normalized [-1, 1]. K is set 300 following the PSNR gap computation of Kadkhodaie et al. (2024).
Turning to diffusion model backbones, prior work (Peebles & Xie 2023) has shown that DiT achieves better image generation quality than UNet with equivalent FLOPs. This prompts our curiosity to study whether DiT also demonstrate superiority in generalization when using the PSNR 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 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 substantial performance difference in distribution mapping given sufficient training data. Nevertheless, as shown in Fig. 2b), when trained with less data, e.g., N=103 and N=10 DiT has remarkably smaller PSNR gap than UNet, suggesting that 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 this generalization difference.
2.2 Digep-Decompoostion Analysis Cannot Explain DiT Generalization
Kadkhodaie et al. (2024) reveal that the generalization of locally linear one-channel UNet driven by the emergence of geometry-adaptive harmonic bases. These harmonic bases are obtained from the eigenvectors of a locally linear UNet's Jacobian matrix. This raises an important question: Do classic hybrid UNets and DîT's also possess harmonic bases that account for their generalization difference? Unfortunately, due to the use of nonlinear operations such as softmax in transformer blocks and normalization layers in both convolution and transformer layers, the cigen-decomposition analysis used by Kadkhodaie et al. 2024) fails to reveal meaningful insights about the inductive
www www.github..o www.gggttt....oopeee http:///ww.....aaa////////adddda

---

88.225
15.0
Dominated large
125
10.0
eigenwalues
6.0
1000
Goog
1000
1000
An 3.033
7.755
80-
50
Nokey
20
00000
0500
1000
X000
(a) UNet, FLOPs: 303.17G; Params: 109.55M
(h) DIT, FLOPs: 300.49G; Params: 14.27M
Figure 3: Jacobian eigenvector comparison between the hybrid UNet (Nichol & Dhariwal 2021 and DiT (Peebles & Xie 2023) with equivalent FLOPs. (a) The eigenvectors of hybrid UNet form harmonic bases that tend to memorize the training images when - but not adapt well to the input geometry, differing from the behavior observed by Kadkhodaie al. (022). In contrast, (b) the DiT's eigenvectors do not form harmonic bases at either N=10 or =10%. Overall, oegepddecompoiition analysis for both the hybrid UNet and DiT fails to reveal sufficient insight into the inductive biases underlying their generalization.
biases that explain the generalization difference between a UNet and DIT. To investigate this further, we follow Kadkhodaie et al. (024)) and perform an eigen-decomposition of the Jacobian matrices for three-channel hybrid UNet (Nichol & Dhariwal 2021) and DiT. Specifically, first feed noisy image x (xy, is omitted for readability) into a DiT and a UNet and obtain their Jacobian matrices:
Jacobian VEA
OR H.W. the 1EH
Each entry of the Jacobian represents the partial derivative of output pixel w.r.f. all input pixels. Next, we perform an cigen-decomposition of the Jacobian matrix and obtain the eigenvectors.
Fig. presents the eigenvalues and eigenvectors of hybrid UNet and DiT trained with 10 and 10a images, respectively. For UNet trained with a small dataset (e.g., N=10), the Jacobian eigenvectors corresponding to several large eigenvalues tend to memorize the geometry of the input image. The leading eigenvalues are sigmificantly larger than the rest, suggesting that the UNet trained with 10 images primarily memorizes training examples (Carlini al. 2023 Somepalli al. 2023). When the training set size increases to N=10", the UNet's eigenvectors exhibit oscillatory patterns whose frequency increases eigenvalues X decrease. However, these harmonic bases no longer adapt well to the geometry of the input image.
In contrast, shown in Fig. 3bb). the eigenvectors of DiT display random, sparse patterns regardless of the training dataset size. Unlike the UNet, the eigenvalue distribution of the DiT changes little between N=10 and 1010 and no harmonic bases emerge. Overall, the dego-odecompoiition of the Jacobian matrices for hybrid UNets and DiTs does not reveal the coometry--daattin harmonic bases observed by Kadkhodaie et al. (202)). This does not imply that hybrid UNets and DiTs aren t capable of forming such bases. rather indicates that the eigen-decommn gege-ddcompooition analysis is invalid for characterizing the inductive biases underlying DIT generalization. This observation motivates to seek alternatives to investigate the inductive biases that enable the generalization of DiTs.
2.3 How Does a DiT Generalize?
The generalization of a DiT may originate from the self-artention (Vaswani 2017) dynamics because of its pivotal role DIT. Could the attention maps of a DiT provide insights into its inductive biases? To shed light, we empirically compare the attention maps of DiTs with varying levels of
5

---

Layer 2
Layer 4
Layer 6
Layer 8
Layer 10
Layer 12
Token Attention
0.10
N=10
Reshupe o (HW Hx
0.08
Better Locality
0.06
1-10
N-1
N
0.04
10
0.00
Shape of each attention map: (HW) x (HW)
Reshaped to: H x W
Figure 4: Attention maps of DiT's trained with 10, 10 and 10 images. The top-right insets provide 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 single row from the layer-12 attention map (original shape: 1x IxHHW) into matrix of shape H W. As N increases from 10 to 109, the token attentions progressively concentrate around the region near the output token (highlighted with blue boxes).
specific output token, obtained by reshaping
generalization: three DIT models trained with 10, 103, and 10 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 follows,
[OKT)
Attention Map Softmax
(6)
Vd
where {Q, K) E (HW))x represent the query and key matrices. H and W are the height and width of the input tensor, while denotes the dimension of self-attention layer. For better readability of the attention maps, we linearly normalize each attention map to the range of [o, 1] and apply colormap the interval [0, 0.1], ie., values exceeding the upper bound clipped 0.1.
Fig. shows the attention maps of DiTs with varying levels of generalization on a randomly selected image. Empirically, observe that the attention maps of a DiT's self-attention layer remain highly consistent across different images. Further details are provided in Appendix F As the number of training images increases from 1 to =0 the attention maps of DiT become increasingly concentrated along several diagonal lines, especially in early layers. A closer inspection of the attention values of specific target token, i.e., row in the attention map, shows that these diagonal patterns highlight spatially close locations, indicating that the generalization ability a DiT is linked the locality of its attention maps.
3 Verifying Attention Locality as a Bias by Restricting Attention
To verify attention locality an inductive bias, observed in Fig. map deviates from a pure identity attention. Specifically, for the attention map Attn E (MMx)) corresponding to target output token location (i,j), compute the deviation
M we assess how much an attention
Dev (i.j)
Σ (m,n)
(D())) Attn (m, (m, where D (i.j) (m,n) V + (n-
+
5)². (7)
MN
Eq. measures how much the attention map Attn deviates from the target token at location (i.j) (Wasserstein distance). We obtain the deviation for the whole attention map Alin by averaging the deviation for all target tokens. In the first row of Tab. [ we provide the deviation averaged over 300 random test images using DiT trained with 103. 10a, and 10 images. When increasing the number of training images (e.g., from 10 to 10'), the DiT tends to generalize better, which

---

Table I: Deviation, comparison between DIT' with and without local attention. In this setting, local attention with window sizes of (3, 5, 7,9. 9, 11, 13) is applied to the first six layers of the DIT. 1 xx and 5x5 denote the local kernel sizes from which the attention maps deviate.
DIT Layers
Layer
Layer 5
Layer
Train Set Size
10a
10*
10a
10a
10
10P
10a
10
10
DIT-XS/1 (1x1)
1.977
0.153
0.073
0.174
0.049
0.016
0.049
0.029
0.037
DIT-XS/1 (5x5)
0.274
0.016
0.010
0.075
0.055
0.033
0.070
0.054
0.040
w/ Local (1xl)
0.002
0.002
0.002
0.019
0.005
0.004
0.111
0.046
0.052
Attention Map
Token Map
Attention Map
Token Map
Q, E R x
Target
Token
Query
Key
MdV
Attended
OK
Softmax
Va
Skipped Token
(a) Global Attention
(b) Local Attention of Size (3x3)
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 nearby window around the target.
accompanied by a reduction of the deviation, especially in early layers. similar reduction is observed when measuring the deviation from a 5x5 local attention kernel, shown in the second row of Tab. I Based on this observation, we hypothesize that it is possible to adjust the inductive bias of a DiT by restricting attention windows of early layers. To test this, we set up baselines by of of & Xie (202)). Specifically, remove the auto-encoder and set the patchify size 1x1, transforming it into pixel-space DiT. This modification rules out irrelevant components and ensures more straightforward comparisons in downstream experiments. For model training, we images of resolution x2x32, which is equivalent in dimensionality to 512x512 for a latent-space DiT with a patchify size of 2x2.
adopting the DiT implementations from the official repository
In the remainder of this section, we show that based on the PSNR gap, injecting local attention in early layers can effectively modify a DiT generalization, often accompanied by an FID change when insufficient training data is used. Furthermore, we discover that placing the attention window restrictions at different locations a DiT and adjusting the effective attention window sizes allows for additional control over its generalization behavior. Details w.E. experimental settings, theoretical connections to other inductive biases, more quantitative, qualitative, as well as generation results, and limitations are deferred to the Appendices.
3.1
Attention Window Restriction
Local attention, initially proposed to enhance computational efficiency (Lu al. 2021: Yang et al.
2022 Hatamizadeh et al. 2023 Hassani et al. 2023). is straightforward yet effective way to modify DiT generalization. Different from global attention which enables a target token to connect with
all input tokens (Fig. a)), local attention only permits a target token to attend within a small nearby window. The resulting attention map structure is depicted in Fig. b)... Notably, local attention constrains the attention map to sparse activation pattern only along the diagonal direction, thereby enforcing locality of the attention map. The resulting attention map patterns produced by local attention align well with the inductive bias that a DiT exhibits when observing strong generalization ability, illustrated in Fig. (row --10").
Using local attentions in early layers a DiT can consistently improve its generalization (measured by PSNR gap) across different datasets and model sizes. Specifically, consider DiT model with

---

Table 2: PSNR gapl comparison between DiT with and without local attention for two architectures:
DiT-XS/1 and DiT-S/1. Local denotes applying local attention with window sizes (3, 5, 7, 9, 13) to the first six layers of the DIT.
Dataset
CelebA
ImageNet
MSCOCO
LSUN Church
LSUN Bedroom
LSUN Bridge
LSUN Tower
Train Set Size
10
10
10*
109
10*
10
10
109
10
10
10*
10
10
10
DIT-XS/I
0.80
0.01
1.08
0.05
0.60
0.13
0.38
0.00
0.70
0.26
0.52
0.03
0.52
0.00
0.57
0.01
0.74
0.05
0.41
0.13
0.21
0.00
0.52
0.26
0.31
0.03
0.33
0.00
Local
-0.20
-0.31
-0.19
-0.45
-0.26
-0.35
-0.37
DIT-S/I
2.30
0.02
0.65
0.05
0.72
0.13
0.61
0.00
0.94
0.26
1.74
0.03
1.97
0.00
1,73
0.02
0.43
0.05
0.54
0.13
0.36
0.00
0,64
0.26
1.26
0.03
1.34
0.00
Local
-.225
-0.34
0.18
-0.41
-0.32
-0.32
Table 3: FIDJ, comparison between DiT with and without local attention. highlighted in bold font.
The best results are
Model
CelebA
ImageNet
MSCOCO
LSUN Church
LSUN Bedroom
LSUN Bridge
LSUN Tower
Train Set Size
10*
10a
16*
10
10*
10a
10"
10a
10
10a
10*
10"
10*
10"
DIT-XS/1
9.69
2.63
52.57
17.31
28.35
12.97
12.88
4.38
14.84
5.41
23.18
8.08
12.55
4.66
w/ Looall
8,46
2.55
43.87
18.07
24.43
13.41
10.48
1.17
11.96
5.35
18.15
8.35
10.56
1.80
DIT-S/I
23.25
2.33
36.64
20.61
29. 25
13.78
14.88
3.94
16.11
4.61
51.57
5.80
28.97
3.19
w/ Local
20.78
2.33
33.18
20.80
27.11
13:16
11.75
4.41
11.68
5.05
37.65
5.88
21.81
3.56
12 DiT blocks, and replace the first global attention layers with local attentions, whose window sizes range from 3x3 to 13x 13 with a stride of 2. We train both the vanilla DiT and DiT equipped with local attentions with ==100, 10 and 10* images for the same 400k training steps. Then we calculate the PSNR gap between the training and testing images for models trained with different amounts of images. In Tab. 2 we show the PSNR gap comparison between a DiT with and without local attentions on CelebA, ImageNet, MSCOCO, and LSUN (Church, Bedroom, Bridge, Tower) datasets, using baseline DiT models of two sizes (DiT-XS/1 and DiT-S/I). 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 discriminative model, e-g., classifier, better generalization may lead better model perfor- mance when the training dataset is insufficient. Is this also the case for generative models like a DIT? To investigate, For each dataset, we compare FID values of models trained with 1 and 10s images: the former represents the case of insufficient training images while the later case refers use of sufficient training data. Tab. 3 shows the FID comparison among the same seven datasets and the two DiT backbones used when comparing PSNR gaps. Improving the generalization via local attentions can indeed improve the FID when ==10 When ==10², adding local attentions either results in comparable FID values or leads slight compromise because DiT trained with sufficient data can naturally develop a local attention pattern as shown in Fig. + So further encouraging attention locality expected have limited effect. However, offers the added benefit of reducing FLOPS with minimal performance loss. Both observations are in line with findings from discriminative models. Interestingly, we find that modifying the placement and effective attention window size permits to control a DiT's generalization and generation quality. More discussions are in Sec. 3.2 and Sec. 3.3 below. Going back to Tab. [ the third row shows the deviation when using local attention in early layers of a DIT. As expected, using local attention reduced the deviation of early layers. Interestingly, the deviation of the remaining layers without local attention increased. This shows that other factors beyond locality of attention are play. We leave identification and a study of those to future work.
compare the FID between the default DiT and a DiT using local attentions.
In light of Occam's razor, reducing the model parameter count has been shown 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 DīT, we explore sharing of the parameters of DiT's attention blocks as well as modifying DiT's attention layers to learn the coefficients of pre-computed offline PCA components. Neither of these methods showed as compelling improvements of the generalization (measured via the PSNR gap) as using local attention. We provide more details regarding the considered techniques in Appendix
8

---

Table 4: PSNR gapl and FID, comparison for different local attention placement patterns. The best results are highlighted in bold font.
PSNR Gap
CelebA
ImageNet
FID
CelebA
ImageNet
Train Set Size
10
10"
10a
10
10*
10
Train Set Size
10*
10
10"
10
DIT
7.49
0.80
0.01
7.77
1.08
0.05
DIT-XS/I
9.69
2.63
52.37
17.31
w/ Local (head)
6.56
0.57
0.01
6.76
0.74
0.05
w/ Local (head)
8.46
2.55
43.87
18.07
w/ Loccll (mix)
7.66
1.05
0.01
7.27
0.58
0.05
w/ Local (mix)
11.89
2.50
37 64
18.44
Local (tail)
9.05
1.46
0.05
(tail)
18.07
59.85
17.58
w/ Local" (head)
5.42
0.36
0.01
4.94
0.15
0.05
w/ Local" (head)
7.23
3.10
29.25
23.79
w/ Logal" (mix)
6.99
0.86
0.01
7.12
0.92
0.05
Local" (mix)
10.95
2.71
51.82
18.80
w/ Local" " (tail)
8.04
1.59
0.02
8.26
1.05
0.05
w/ Local" (tail)
17.04
3.04
49.64
22.17
Table 5: PSNR gapl and FID changes when the effective attention window size is kept constant, decreased, or increased. Best results are highlighted in bold font.
PSNR Gap
CelebA
ImageNet
FID
CelebA
ImageNet
Train
10a
10*
10a
10a
10
10
Train
10
10
10
16a
Local Attn (5*6)
7.19
1.05
0.02
6.55
0.69
0.05
Local Attn (8*8)
12.98
2.33
40.74
17.87
(3*2, 5 2=2)
7.00
1.01
0.02
6,55
0.05
(3"2 5 .7*2)
12.67
2.35
10.75
17.75
Local
6.56
0.57
0.01
6.76
0.74
0.05
Local
8.46
2.55
43.87
18.07
(smaller win size)
6.09
0.54
0.01
6,33
0.63
0.05
(smaller win size)
8.05
2.72
39.58
18.94
Local'
5.42
0.36
0.01
4.94
0.05
Local"
7.23
3.10
29.25
23.79
(larger win size)
5.92
0.16
0.01
6.15
0.56
0.05
(largr win size)
7.88
2.86
37.87
19,36
3.2 Placement of Attention Window Restriction
For local attention, we study three placement schemes: 1) using local attention in early layers of a DIT, 2) interleaving local attention with global attention, and 3) placing local attention the final layers of a DIT. In Tab. 4 we compare the PSNR gap for the three schemes on the CelebA and ImageNet data, using two distinct local attention configurations. Specifically, Local refers to setting with 6 attention layers, where the window sizes vary from 3x3 to 13x 13 with a stride of 2, which is consistent with the local attention configuration used in Tab. and Tab. above. Meanwhile, Local" represents a different configuration consisting of 9 local attention layers, arranged as (3a3, 5,3 7*3). where 2d indicates repeating local attention layer with (ixi) window j times.
The results in Tab. 4 indicate that applying local attention in the early layers of DiT leads smaller PSNR gap across different training data sizes, corroborating our carlier findings. Additionally, the FID results in Tab. f show that the first placement scheme generally improves FID when the training data is limited 10=1 In contrast, interleaving local and global attention, or applying local attention on the final layers, enhances the model's Sata-fitting ability but often compromises generalization. These two placement schemes tend to improve FID when ==10 at the cost of reduced FID when == further supporting the generalization results measured by the PSNR gap.
3.3 Effective Attention Window Size Analysis
Adjusting the effective attention window size provides an additional mechanism to control the gener- alization 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 DiT. These observations are based 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 Tab. 5
Specifically, in the first comparison, we apply two configurations of local attentions with window sizes (5, 5, 5, 5, 5, 5) and (3, 3. 5, 7, 7) 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 limited impact on a DiT's generalization, as indicated by the similar PSNR gaps across -=103, 10 and 10$ This similarity generalization is further corroborated by their comparable FID values. the second and third comparisons, using the DIT-XS/I configurations with Local and Local" attention settings,

---

find that reducing the attention window size enhances generalization, while increasing the window size diminishes it. This is evidenced by decrease in the PSNR gap for smaller window sizes and an increase for larger ones. Furthermore, the improved generalization is associated with better FII) values under comparably insufficient training data, and vice versa.
Related Work
Inductive Biases of Generative Models. Current diffusion models (Sohl- et al. 2015
Song et al.
2020 Ho et al.,
2020 Kadkhodaic & Simoncelli 2020; Nichol & Dhariwal 2021; Song
et al..1.
2020 An al. (024)
exhibit strong generalization abilities (Zang et al.
|021 Keskar et al.
2016 Grithths et al. 2024, Wilson & Izmailov] 2020), relying on inductive biases (Mitchell) 1980,
Goyal & Bengio 2022). Prior to the emergence of diffusion models, Zhao al.
(2018) show that
generative models like GANs (Goodfellow et al. (0200) and VAES (Kingmal 2013) can generalize to novel attributes not presented in the training data. The generalization ability of generative models
is often attributed to inductive biases introduced by model architecture and training (Zhng et al.)
202T; Keskar et al., 2016). Kadkhodaic et al. (2024) link the generalization of diffusion models to adametty-adaptive harmonic bases (Mallat et al. (020)), but their analysis focuses on a simplified one-channel UNet. It remains unclear whether their findings extend to standard three-channel UNets (Ncholl & Dhariwal 202T) or DiTs (Peebles & Xie (023). This work addresses this gap: we show that UNets still exhibit harmonic bases, whereas DiT's do not. Instead, DiTs generalize through a different inductive bias attention locality. In contrast to Zhang et al. (2024), who argue that diffusion models converge to the optimal score function largely independent of architecture, we focus on the architectural inductive biases that influence the diffusion model generalization. Recent works (Wang& & Vastola 2024| Li et al.
2024) examine the linearity of score functions but do not address architectural biases. Niedoba et al. (2024) observe that diffusion models resemble patch-based
denoisers; our discovery of attention locality in DiTs offers an explanation for this behavior.
Attention Window Restrictions. Restricting attention windows through mechanisms such as local
attention (Beltagy et al. 2020; Liu et al. [02TT Hassani al. (023). strided attention (Wang et al.)
2021 Xia et al. 2022) and sliding attention (Pan et al.. 2023). among others, can sigmificanly
improve the efficiency of attention computation (Yang et al. 2022 Hatamizadeh et al. 2023 Hassani
et al. 2023 Apple 2024). These techniques limit the attention scope, reducing computational
complexity while retaining the model's ability to capture important contextual information. However, our work reveals another use for controlling the locality of attention. 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. We think this is particularly important for science domains where data for training generative models is less abundant.
5 Conclusion
This paper investigates the inductive biases that facilitate the generalization ability of DiT's For insufficient training data, we observe that DìT's achieve superior generalization, as measured by the PSNR gap, compared to UNets with equivalent FLOPs. However, the sigen-decomposition analysis that reveals geometry-adaptive harmonic bases as the key inductive bias of diffusion models based on locally linear UNets becomes invalid for classic hybrid UNets and DiTs due the presence of nonlinear operations. Therefore, we take an alternative approach to explore alternative inductive biases and identify that 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 DiT's generalization often leads to improved FID scores when trained with insufficient data. One limitation of this work is that our analysis focuses exclusively on DiTs. For future work, we consider it important and interesting to study the generalization behavior of hybrid models and conditional transformers (e.g., MMDiT modules), given their growing popularity in recent generative architectures.
10

---

Acknowledgement
We sincerely thank Zhongzheng Ren, Chen Chen, Byeongjoo Ahn, Saeed Khorram, and Aditya Sankar for their valuable discussions and insightful feedback. We are also deeply grateful to Alex Colburn, Qi Shan, and the Video Computer Vision team Apple for their generous support in providing the infrastructure and computational resources that made our experiments possible.
References
An, J., Yang, Z., Wang, J., L., Liu, Z., Wang. L., and Luo, J. Bring metric functions into diffusion models. arXiv preprint arXiv:2401.02414, 2024.
Apple. Deploying attention-basee baaeed vision transformers to apple neural engine, 2024.
Barratt, S. and Sharma, R. A note on the inception score. arXiv preprint
2018.
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.051 0500000005555.... 2020.
Blattmann, A., Dockhorn, T., Kulal, S., Mendelevitch, D., Kilian, M., Lorenz, D., Levi, Y., English, Z., Voleti, V., Letts, A., et al. Stable video diffusion: Scaling latent video diffusion models to large 2023.
datasets. arXiv preprint 1:xXi:2331.151/2
Carlini, N., Hayes, J., Nasr, M., Jagielski, M., Sehwag, V., Tramer, F., Balle. B., Ippolito, D., and Wallace, Extracting training data from diffusion models. 32nd USENIX Security Symposium, 2023.
Chen, C., Qian, R., Hu, W., Fu, T.JJ., Tong, J., Wang, X., Li, L., Zhang, B., Schwing, A., Liu, W., and Yang, Y. DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation. In arxiv.org/abs/2503.. 10618. 2025.
Chen, J., Yu, J., Ge, C., Yao, L., Xie, E., Wu, Y., Wang, Z., Kwok, J., Luo, P., Lu, H., et al. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.00426, 2023.
Chen, J., Ge, C., Xie, E., Wu, Y., Yao, L., Ren, X., Wang, Z., Luo, P., Lu, H., and Li, Z. Pixart-\sigmmn Weak-to-strong training of diffusion transformer for 4k text-to-image generation. arXiv preprint wrXiv:2403.04692, 2024a.
Chen, J., Wu, Y., Luo, S., Xie, E., Paul, S., Luo, P., Zhao, H., and Li, Z. Pixart- Fast and controllable image generation with latent consistency models. arXiv preprint arXiv:2401.05252. 2024b.
Cheng, H. K., Ishii, M., Hayakawa, A., Shibuya, T., Schwing, A., and Mitsufuji, Y. MMAudio: Taming Multimodal Joint Training for High-Quality Video-to-Audio Synthesis. In CVPR, 2025.
De Wolf, R. A brief introduction to fourier analysis on the boolean cube. Theory of Computing, 2008.
Deng, J., Dong, W., Socher, R., Li. L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. In NeurIPS, 2021.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, et al. Scaling rectified flow transformers for high-resolution image synthesis. In ICML, 2024.
Girdhar, R., Singh, M., Brown, A., Duval, Q., Azadi, S., Rambhatla, S. S., Shah, A., Yin, X., Parikh,
D.. and Misra, I. Emu video: Factorizing text-to-video generation by explicit image conditioning. arXiv preprint arXiv:2311.1070 2023.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 2020.
11

---

Goyal, A. and Bengio, Y. Inductive biases for deep learning of higher-level cognition. In Proceedings of the Royal Society A, 2022.
Griffiths, T. L., Zhu, J.-Q., Grant, E., and Thomas McCoy, R. Bayes the age of intelligent machines. Current Directions in Psychological Science, 2024.
Hahn, M. and Rofin, M. arXiv:2402.099663 2024.
Why are sensitive functions hard for transformers? arXiv preprint
Hassani, A., Walton, S., Li, J., Li, S., and Shi, H. Neighborhood attention transformer. In CVPR, 2023.
Hatamizadeh, A., Heinrich, G., Yin, H., Tao, A., Alvarez, J. M., Kautz, J., and Molchanov, P. Fastervit: Fast vision transformers with hierarchical attention. arXiv preprint arX6v:2306.06114 2023.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to local nash equilibrium. In NeurlPS, 2017.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. In NeurIPS, 2020.
Ho, J., Chan, W., Saharia, C., Whang, J., Gao, R., Gritsenko, A., Kingma, D. P., Poole, B., Norouzi, M., Fleet, J., et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02300: 2022.
Hron, J., Bahri, Y., Sohl-Dickstein,, J., and Novak, R. Infinite attention: Nngp and ntk for deep attention networks. In ICML, 2020.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. arXiv preprint arXiv:1912.02/17 2019.
Kadkhodaie, Z. and Simoncelli, E. P. Solving linear inverse problems using the prior implicit in denoiser. arXiv preprint arXiv:2007. 13440, 2020.
Kadkhodaie, Z., Guth, F., Simoncelli, P., and Mallat, S. Generalization diffusion models arises from geometry-adaptive harmonic representation. In ICLR, 2024.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, T. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1609.04336 2016.
Kingma, D. P. Auto-encoding variational bayes. arXiv preprint arXiv: 2013.
Li, X., Dai, Y., and Qu, Q Understanding generalizability of diffusion models requires rethinking the hidden gaussian structure. NeurIPS, 2024.
Lin, T.-Y., Maire, M., Belongie, S., Hays, J.. Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. Microsoft coco: Common objects in context. In ECCV, 2014.
Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In JCCV, 2015. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In CVPR, 2021.
Mallat, S., Zhang, S., and Rochette, G. Phase harmonic correlations and convolutional neural networks. Information and Inference: A Journal of the IMA, 2020.
Mitchell, T. M. The need for biases in learning generalizations, 1980.
Neyshabur, B., Bhojanapalli, S., McAllester, D., and Srebro, N. Exploring generalization in deep learning. NewrlPS, 2017.
Nichol, A., Dhariwal, P., Ramesh, A., Shyam, P., Mishkin, P., McGrew, B., Sutskever, I,, and Chen, M. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint 1xXiv:2112,10744 2021.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In ICML, 2021.
12

---

Niedoba, M., Zwartsenberg, B., Murphy, K., and Wood, F. Towards mechanistic explanation of diffusion model generalization. arXiv preprint arX1v:2411.19539, 2024.
OpenAl. Video generation models as world simulators, 2024.
Oquab, M., Darcet, T., Moutakanni, T., Vo, H., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza,
D., Massa, F., El-Nouby, A., et al. Dinov2: Learning robust visual features without supervision.
arXiv preprint urXiv:2304.0711
2023.
Pan, X., Ye, T., Xia, Z., Song, S., and Huang, G. Slide-transformer: Hierarchical vision transformer with local sclf-attention. CVPR, 2023.
Peebles, W. and Xie, S. Scalable diffusion models with transformers, In CVPR, 2023.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, 2022.
Saharia, C.,, Chan, W. Saxena, S., Li, L.,, Whang, J., Denton, E. L., Ghasemipour, K., Gontijo Lopes,
R., Karagol Ayan, B., Salimans, T., et al. Photorealistic text-to-image diffusion models with deep language understanding. In NeurIPS, 2022.
Singer, U., Polyak, A., Hayes, T., Yin, X., An, Zhang, S., Hu, Q. Yang, H., Ashual, O., Gafni, O., et al. Make-a-video: Text-to-video generation without text-video data. In ICLR, 2022.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML, 2015.
Somepalli, G., Singla, V., Goldblum, M., Geiping, J., and Goldstein, T. Diffusion art or digital forgery? investigating data replication in diffusion models. In CVPR, 2023.
Song. Y.. Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint 2rXiv:2011.13455 2020.
Vasudeva, B., Fu, D., Zhou, T., Kau, E., Huang, Y., and Sharan, V. Simplicity bias of transformers to learn low sensitivity functions. arXiv preprini arXiv:2403.06925, 2024.
Vaswani, A. Attention is all you need. In NeurlPS, 2017.
Wang, B. and Vastola, J. J. The unreasonable effectiveness of gaussian score approximation for diffusion models and its applications. arXiv preprint arXiv:2412.09770 2024.
Wang, W., Xie, E., Li, X., Fan, D.-P., Song, K., Liang, D., Lu, T., Luo, P., and Shao, L. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In ICCV, 2021.
Wilson, A. G. and Izmailov, P. Bayesian deep learning and probabilistic perspective of generalization. In NewrlPS, 2020.
Xia, Z.,, Pan, X., Song, S., L. E., and Huang, G. Vision transformer with deformable attention. In CVPR, 2022.
Yang, C., Qiao, S., Yu, Q., Yuan, X., Zhu, Y., Yuille, A., Adam, H., and Chen, L.-C. Moat: Alternating mobile convolution and attention brings strong vision models. In ICLR, 2022.
Yang, G. and Salman, H. A fine-grained spectral perspective on neural networks, arXiv preprint 1rXiv:1907110599 2019.
Yu, F., Zhang, Y.. Song, S., Seff, A., and Xiao, J. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:1506.03365. 2015.
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 2021.
Zhang. H., Zhou, J., Lu, Y., Guo, M.. Wang, P., Shen, L., and Qu, Q. The emergence of reproducibility and consistency in diffusion models. In ICML, 2024.
7hao, S., Ren, H., Yuan, A., Song, J., Goodman, N., and Ermon, Bias and generalization in deep generative models: An empirical study. In NeurlPS, 2018.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly convey that our work focuses on the inductive biases underlying DiT generalization, with our key contribution, the discovery of attention locality bias, explicitly highlighted.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We have clearly discussed the limitations of this work, including the lack of formal theoretical proof and the limited practical impact of the discovered attention locality bias in the large-data regime, in Appendix
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
14

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Appendix B we show that local attention promotes a simplicity bias, which reduces the model's s sensitivity to data perturbations. We further relate this reduced sensitivity to flatter minima, a well-established indicator of good generalization. All assumptions are consistent with prior work, and the proof is, to the best of our knowledge, correct.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
we provide all experimental settings, and our implementation is based publicly available DiT codebase. As stated in the abstract, we will release our code upon publication of the paper.
Justification: Appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
15

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: All datasets used in this paper are publicly available. Our implementation builds on the official DiT codebase, which is also publicly accessible. As stated in the abstract, we will release our code upon publication of the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https://nips.cor
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
hraps..ccpubllic/ui..ees http://gguuueeeeeeeueeuuiuiiaiia
ssinnooicciiy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have clearly presented all experimental settings in Appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For the PSNR and PSNR gap analysis in Fig. 2 we include standard deviation ranges to demonstrate the robustness our results. Although we not report error bars for other metries such FID, FD-DINOv2, and IS due to the high computational cost,
16

---

have thoroughly evaluated our model across seven diverse datasets using both pixel-space and latent-space diffusion models. The consistency of results across these settings further confirms the robustness of our findings.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (c.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Appendix A report that all models were trained using 4 or 8 A100/H100 GPUs, and all checkpoints were taken at 400k training steps.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We only publicly available datasets, ensure transparency by providing detailed experimental settings, and commit to releasing our code upon publication to support reproducibility and open science.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
17

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We have discussed the broader social impact of this work Appendix K work does not introduce any new unforeseen risks to the community.
This
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work is based entirely existing publicly available models and datasets. Our focus is on analyzing generalization to provide theoretical and empirical insights, rather than introducing new models methods that pose risk of misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
18

---

Answer: [Yes]
Justification: All assets are properly cited and their licenses are:
CelebA (non-commercial research license)
LSUN (non-commercial research license)
ImageNet (non-commercial research license)
MS COCO (Creative Commons CC-BY 4.0)
The official DiT codebase (Apache 2.0)
Our use is strictly non-commercial research, fully consistent with each license's 's terms.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssittcooeeeoo thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce any new assets. All experiments are conducted using existing public datasets and codebases.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
19

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This work does not involve human subjects study participants, and therefore does not require IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is required.
Answer: [NA]
Justification: We only used LLMs to assist with minor writing refinements.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.MM policy http:///earrrppeeccooooooa/aanna for what should or should not be described.
20


Reviews:
[{"review_id": "d63OTZAmPf", "reviewer": "Reviewer_XEv5", "summary": "This 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.", "questions": "- 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.", "limitations": "- No negative societal impacts. \n- Yes limitations have been addressed", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I think this paper is making a significant and valuable contribution regarding generalization in DiT models. However there are some analysis regarding using a first order approximation to analyze the network that seems incorrect to me, due to the inherent local nonlinearity of the model. Although the authors agreed the analysis was incorrect, and potentially misleading, they didn't make it clear that the paper would be revised to remove or replace such analysis from the paper. So I lowering my rating from 5 to 4 because I remain unsure if the incorrect conclusions will be included in the final revision or not. That said, I still regard the paper as an important contribution to the study of locality in transformers, which justifies the relatively high score.", "rebuttal": "Thanks a lot for your time and feedback. We are encouraged that our findings about DiT’s locality bias shown in Fig. 4 are recognized: A DiT can form a locality pattern similar to adaptive convolution when it achieves a good generalization, despite that access to global tokens is provided. We also appreciate the suggested related work by Raghu *et al.*, which echoes our finding but in ViTs. We will discuss this new related work in the revised paper.  We answer the remaining questions below.\n\n---\n* **Q1: DiT lacks local linearity due to the use of GeLU, Softmax, and group norm layers. Hence, the Jacobian analysis in Figs. 1 and 3 is not a valid interpretation tool for DiT’s mappings.**\n\n  We agree, a DiT is not locally linear. Note, we use the Jacobian to analyze the model’s sensitivity and representation behavior instead of linearly decomposing the full mapping. Therefore, we do not draw strong conclusions from Fig. 1 and Fig. 3 alone. Rather, we view them as suggestive empirical evidence that first-order Jacobian analysis may not suffice to study the inductive bias of a DiT. This motivates direct visualization of the attention maps of a DiT in Fig. 4, where we observe that attention patterns in DiT converge to localized structures. We’ll clarify in a revised version by adding this discussion and noting that this doesn’t rule out GAHBs.\n---\n* **Q2: The Jacobian analysis in Fig. 1 is not as informative due to the non-linearity in both DiT and UNet used by Nichol & Dhariwal 2021.**\n  \n  Thanks a lot for a great argument. To confirm, we use one attention layer at the last encoder stage and in the middle block in the UNet by Nichol & Dhariwal 2021, which is stated in line 40-42 of the appendix. We also use group norm and SiLU layers in the UNet. We think the amount of non-locality introduced by network operators is a reasonable and interesting explanation for the Jacobian eigenvector differences between the UNets and DiT in Fig. 1. We will add this discussion to Sec. 2.2 of the revised paper to provide more insight. Furthermore, for a DiT with the same FLOPs as a UNet, whether the amount of non-linearity is a deciding factor for their generalization difference is another interesting future direction. \n---\n* **Q3: Minor: Labeling the UNet in Kadkhodaie et al. (2024) with “simplified” is misleading. The UNet used by Nichol & Dhariwal 2021 should be labeled as augmented UNet.**\n  \n  Thanks for the valuable suggestion. Initially, we considered the UNet by Nichol & Dhariwal 2021 as the \"standard\" UNet as it has been widely used. For clarity, in the revised paper, we will rename the UNet used by Kadkhodaie et al. (2024) as the locally-linear UNet and the UNet by Nichol & Dhariwal 2021 as a hybrid UNet, recognizing that it contains attention layers.\n---\n* **Q4: To compute the average PSNR in Eq 4, the sum should be outside of log (the two operations do not commute). i.e compute individual PSNR for each image and then take the average.**\n  \n  Thank you for pointing this out. Upon reviewing our code, we confirmed that the reported PSNR values are computed by first calculating the PSNR for each image individually and then averaging the results. We will revise Eq. (4) to reflect this accordingly.\n---\n* **Q5: Suggestion: replacing the deviation from identity in eq 6 with deviation from convolutional kernel with a certain kernel size.**\n  \n  Thanks for a valuable suggestion. Below we present the deviation score *w.r.t.* a convolution kernel of $5\\times5$ for a DiT-XS/1 trained with CelebA data using $10^3$, $10^4$, and $10^5$ images, respectively.\n\n  | Train Data Num | $10^3$ | $10^4$ | $10^5$ |\n  | --- | --- | --- | ---- |\n  | Layer 1| 0.274 | 0.016 | 0.010 | \n  | Layer 5| 0.075 | 0.055 | 0.033 | \n  | Layer 9| 0.070 | 0.054 | 0.040 |\n\n  The results above also show that the deviation score *w.r.t.* a convolution kernel decreases when a DiT has better generalization, *i.e.*, when it is trained with more images. We will add the new deviation score results into the revised appendix.\n---\n* **Q6: Do you observe strong generalization (sampling the same image by two networks trained on disjoint datasets) in DiT?**\n  \n  Yes, we did. We will add images generated by two DiTs trained with disjoint data into the revised appendix."}, {"review_id": "Q6YiuyMuVd", "reviewer": "Reviewer_VB1R", "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.", "questions": "1. The experiments are not sufficient to support the proposed viewpoints.\n2. Are there some generated images to demonstrate?\n3. Are there some comparison results with SOTAs?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths：\n1. The paper proposes the locality of attention maps as a key inductive bias contributing to the generalization of a DiT. \n2. The paper demonstrates how to control this inductive bias by incorporating local attention windows into early layers of a DiT. \n\nWeaknesses:\nThe 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.", "was_revised": true, "final_justification": "I appreciate the authors’ detailed response, which has clarified part of my concerns. I remain some cautious about the novelty and effectiveness. I increase my score to Borderline reject.", "rebuttal": "Thanks a lot for your time, feedback, and highlighting that this paper reveals the locality of attention maps as an inductive bias of a DiT, which is verified by our experiments incorporating local attention windows into early layers of a DiT. We answer the remaining questions below.\n\n---\n* **Q1: More comprehensive comparisons with state-of-the-art methods and in-depth ablation studies to fully validate both generalization and generation quality.**\n\n  This paper aims to analyze the inductive bias of a DiT that contributes to its generalization, as stated in line 81-86. Based on experiments reported in Sec. 2, we hypothesize that attention locality is an important indicator of generalization. Our experiments hence modify attention patterns by enforcing attention locality while assessing the resulting impact on the generalization of a DiT. This should not be interpreted as a new method, but rather as studies to verify that the hypothesized attention locality bias is indeed at play. To our best knowledge, there are no state-of-the-art methods to compare to which aim to control the generalization of a DiT. Nonetheless, in Appendix I, we compare the effect of encouraging attention locality with a known way of improving generalization – reducing trainable parameters. As shown in Tab. 11 of Appendix I, two parameter reduction approaches, i.e., using parameter sharing and composing attention maps with pre-defined PCAs, are not as effective in explaining a DiT’s generalization as the local attention study. Further, in Tab. 4, we show that encouraging attention locality at the first few layers of a DiT can improve its generalization. Encouraging attention locality at the last few layers or interleaving global and local attention do not. In Tab. 5, we show that use of a smaller attention window size can lead to a larger generalization improvement while a larger window size leads to a smaller improvement. All these ablation studies verify that encouraging attention locality impacts a DiT’s generalization, corroborating our hypothesis.\n---\n* **Q2: For DiT, there is no generated image to demonstrate and compare.**\n\n  We present qualitative results in Fig. 2 of Appendix D. We find that DiTs trained with $10^4$ and $10^5$ images while using local attention produce images that are more like each other than DiTs trained without using local attention. Further, in Tab. 9 of Appendix D, we quantitatively verify that the average pixel intensity between generated images of models trained with $10^4$ and $10^5$ images is smaller if local attention is encouraged. This phenomenon shows that applying local attention to a DiT trained with $10^4$ images can make it generate images closer to a well-generalized model trained with $10^5$ images, indicating generalization improvement.\n---\n* **Q3: Only using PSNR gap and FID to verify the effectiveness of the viewpoints is not enough.**\n\n  We consider five evaluation metrics: PSNR gap (Tab. 2), FID (Tab. 3), IS (Tab. 4 in Appendix C.3), FD-Dinov2 (Tab. 5 in Appendix C.4), and Cosine Similarity to the nearest training image (Tab. 2 in Appendix C.2). Prior works studying diffusion model generalization [1,2] only use PSNR gap and Cosine Similarity as evaluation metrics. We hence believe the metrics used in our evaluation extend prior work and comprehensively support our claim. \n---\n  [1] Generalization in diffusion models arises from geometry-adaptive harmonic representation. In ICLR, 2024.\n  \n  [2] The emergence of reproducibility and consistency in diffusion models. In ICML, 2024."}, {"review_id": "qBoW8DaUrm", "reviewer": "Reviewer_jHWX", "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.", "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", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "strength:\n1. the research question is specific and clearly stated.\n\n2. the analysis in section 2.3 and the verification in section 3 is reasonable.\n\nweakness:\n1. 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.\n\n2. 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?", "was_revised": true, "final_justification": "Thanks the authors for rebuttal. My concerns on performance have been addressed. However, in my opinion the scope and significance of this work is still a bit limited and specific to certain DiT models. Thus I keep my score unchanged.", "rebuttal": "Thanks a lot for your time, feedback, and highlighting that our research question is specific and clearly stated while our analysis in Sec. 2.3 and verification in Sec. 3 are reasonable. We answer the remaining questions below.\n\n---\n* **Q1: It seems that the result is not very promising. As shown in Table 3, using the proposed local attention can be harmful to the performance in the $10^5$ setting. Although it can improve the performance in the $10^4$ setting, the FID is still high. This makes the applicability of this work a bit questionable.**\n\n  This paper aims to analyze the inductive bias of a DiT that contributes to its generalization, as stated in line 81-86. Based on experiments reported in Sec. 2, we hypothesize that attention locality is an important indicator of generalization. Our experiments hence modify attention patterns by enforcing attention locality while assessing the resulting impact on the generalization of a DiT (FID). This should not be interpreted as a new method, but rather as studies to verify that the hypothesized attention locality bias is indeed at play. Regarding the FID results shown in Tab. 3, when trained with $10^5$ images, as shown in Fig. 4 (N=$10^5$ row), attention locality patterns emerge naturally when training. Hence, it is expected that additionally encouraging attention locality doesn’t lead to any changes in a DiT’s generalization as validated via the PSNR gap in Tab. 2. For FID, we observe a minimal $4.33$% performance fluctuation averaged over seven datasets. Considering that using local attention offers the added benefit of reducing FLOPs, we think it is inaccurate to say local attention is harmful when training a DiT with $10^5$ samples. When using $10^4$ training samples, our experiments consistently show that introducing local attention in early DiT layers improves generalization (PSNR gap in Tab. 2 and Cosine Similarity in Tab. 3 of appendix) and sample quality (FID in Tab. 3, in appendix: IS in Tab. 4, FD-Dinov2 in Tab. 6). In Tab. 3, we observe a consistent $17.59$% FID improvement averaged over seven datasets, which we think is significant compared with the $4.33$% fluctuation when $N=10^5$.\n---\n* **Q2: The proposed method may not be easy to be adapted to other Transformer-based diffusion models, given sensitivity to window size and layer selection.**\n\n  Our experimental results in Tab. 2 and 3 consider two DiTs of different sizes and seven datasets. Using local attention at beginning layers of a DiT with gradually increasing attention window size consistently improves the model generalization. In Tab. 8 of the appendix, we apply the same local attention setting for a latent-space DiT, which also demonstrates reasonable generalization and performance improvement when trained with $N=10^4$ samples, as measured by PSNR gap and FID, respectively. Given these consistent improvements, we think our default setting would work reasonably well for most DiTs. Despite this, a careful analysis of the generalization behaviour is still beneficial for a specific DiT architecture, which we consider one of the key takeaways of this paper.\n---\n* **Q3: It might be a bit narrow if the paper only focuses on geometry-adaptive harmonic bases as in UNet-based denoisers. Is it possible to study the generalization of DiT through a more general exploration?**\n\n  This paper studies the generalization of DiTs from different perspectives. In Sec 2.2, we use a Jacobian analysis to study the representation behaviour of DiTs. Beyond that, in Sec. 2.3, we demonstrate the attention locality bias that is associated with a DiT’s generalization, which is verified by an improved PSNR gap when $N=10^4$ after injecting local attention into a DiT (Sec. 3). Further, in Appendix B, we provide a theoretical analysis connecting the discovered attention locality bias with the simplicity bias (Appendix B.1), and subsequently, low sensitivity bias (Appendix B.2), which indicates good generalization of a Transformer (Appendix B.3). Despite these, we recognize a more general exploration of a DiT’s generalization, an important and interesting future work."}, {"review_id": "f3OOeUDMGJ", "reviewer": "Reviewer_dzs6", "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.", "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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Pros:\n1. The authors are focusing on the newest trend of diffusion model advancement, studying the property of DiT models.\n2. This study has a solid methodology, along with abundant experimental evidence and an interesting finding.\n3. For certain scenarios where the training data is scarce, the finding might have some practical impact.\n\nMinor Cons:\n1. 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.", "was_revised": true, "final_justification": "I have read the response and the discussion between the authors and other reviewers. It seems there is very mixed opinion. Nevertheless, I insist that this paper as an exploratory study can reveal some interesting properties of DiT models, even though it as a practical method for improving DiT generation might not be that well applicable. I will keep my rating not changed as it is already high enough.", "rebuttal": "Thank you very much for your time, thoughtful feedback, and for highlighting the strengths of our work, specifically our focus on timely advances in diffusion models, the solid methodology, the well-supported experimental results, and the interesting findings. We also appreciate your recognition of the potential practical value of our method in low-data regimes. We address the remaining questions below.\n\n---\n* **Q1: Suggestion to extend the analysis to more complex commercial diffusion architectures (e.g., SDXL, SD3, Flux) which include hybrid UNet and ViTs and conditional attention blocks.**\n\n  Thanks for the valuable suggestion. As stated in line 40-42 of the appendix, the UNet architecture used in our experiments contains one attention layer at the last encoder stage and one attention layer at the middle block, which is a hybrid model combining convolution and two attention layers. In general, we agree that studying the generalization behaviour of hybrid models and conditional transformers (*e.g.*, MMDiT modules) are important and interesting due to the increasing popularity of these architectures.\n---\n* **Q2: What could happen if the attention window is somehow enlarged in DiTs trained with sufficient data? Could this modification make them generalize worse?**\n\n  For attention layers of a normal DiT, each token has full access to all other tokens, meaning that the attention window size is already at the maximum and cannot be enlarged. On the other hand, a DiT trained with sufficient data already forms a locality pattern in its attention layers as shown in Fig. 4 ($N=10^5$ row). Therefore, we think further enlarging the attention window would have very limited impact on a DiT’s generalization and performance since newly added non-local tokens in attention don’t contribute much.\n---\n* **Q3: I recommend considering the special ViT-UNet-mixed architecture of Stable Diffusion XL as a possible future topic.**\n\n  Thank you for a valuable suggestion. We agree and will note these future directions in the revised paper."}]

Author rebuttal:

Meta-review: This paper studies the inductive biases that are responsible for generalization in diffusion models. In particular, their main finding is that DiT's trained on larger datasets exhibit local attention patterns (i.e. they attend to a local neighbourhood), while those trained on smaller datasets do not have the locality bias and hence exhibit worse generalization performance. They demonstrate this by showing that enforcing locality in the smaller datasets leads to improved performance. There was large variance in the review scores, but I agree with Reviewer XEv5 and Reviewer dzs6 that the extensive experimentation & findings are interesting and worthy of acceptance. The authors should take on Reviewer XEv5's suggestions regarding softening the claims around the Jacobian analysis (as discussed in the rebuttal period).
Decision: Accept (poster)