U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers
Authors: Yuchuan Tian
Zhijun Tu
Hanting Chen
Jie Hu
Chao Xu
Yunhe Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion Transformers; U-Net; Architecture
Primary area: diffusion_based_models

Abstract: Diffusion Transformers (DiTs) introduce the transformer architecture to diffusion tasks for latent-space image generation. With an isotropic architecture that chains a series of transformer blocks, DiTs demonstrate competitive performance and good scalability; but meanwhile, the abandonment of U-Net by DiTs and their following improvements is worth rethinking. To this end, we conduct a simple toy experiment by comparing a U-Net architectured DiT with an isotropic one. It turns out that the U-Net architecture only gain a slight advantage amid the U-Net inductive bias, indicating potential redundancies within the U-Net-style DiT. Inspired by the discovery that U-Net backbone features are low-frequency-dominated, we perform token downsampling on the query-key-value tuple for self-attention and bring further improvements despite a considerable amount of reduction in computation. Based on self-attention with downsampled tokens, we propose a series of U-shaped DiTs (U-DiTs) in the paper and conduct extensive experiments to demonstrate the extraordinary performance of U-DiT models. The proposed U-DiT could outperform DiT-XL with only 1/6 of its computation cost. Codes are available at https://github.com/YuchuanTian/U-DiT.

U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers
Yuchuan Tian',, Zhijun Tu Hanting Chen', Jie Hu', Chao Xu', Yunhe Wang" 1 State Key Lab of General AI, School of Intelligence Science and Technology. Peking University.
Huawei Noah's Ark Lab.
tianyclstu. man.eddu.com cn, fzhijun.tu, chenhanting, hujie23, yunhe. wang)@huawei.... com xuchaoßcis. pku. eku.edu.cc
Abstract
Diffusion Transformers (DiTs) introduce the transformer architecture to diffusion tasks for latent-space image generation. With isotropic architecture that chains a series transformer blocks, DiTs demonstrate competitive performance and good scalability; but meanwhile, the abandonment of U-Net by DīTs and their following improvements is worth rethinking. To this end, we conduct simple toy experiment by comparing U-Net architectured DiT with an isotropic one. turns out that the U-Net architecture only gain a slight advantage amid the U-Net inductive bias, indicating potential redundancies within the U-Net-style DiT. Inspired by the discovery that U-Net backbone features are low-freeu frequency-domminet token downsampling the query-key-value tuple for self-atter self-attention that bring further improvements despite considerable amount of reduction in computation. Based self-attention with downsampled tokens, we propose series of U-shaped DiTs (U-DiTs) in the paper and conduct extensive experiments to demonstrate the extraordinary performance of U-DiT models. The proposed U-DiT could outperform DiT-XL/2 with only 1/6 of its computation cost. Codes are available at /ttps://githii comm//uchaanTaan//uu
we perform
Introduction
Thanks to the attention mechanism that establishes long-range spatial dependencies, Transform- ers (36] are proved highly effective on various vision tasks including image classification [15]. object detection [5), segmentation 43), and image restoration (6). DiTs 28 introduce full transformer backbones to diffusion, which demonstrate outstanding performance and scalability on image-space and latent-space generation tasks. Recent follow-up works have demonstrated the promising prospect of diffusion transformers by extending their applications to flexible-resolution image generation (26). realistic video generation [2] cetera.
Interestingly, DiTs have discarded the U-Net architecture that is universally applied in manifold
previous works, either in pixel 220 [3] or latent space 29 The use of isotropic (i.e. standard transformer; plain stack of transformer blocks) architectures DiT' is indeed successful, as sealed- up DiT models achieve supreme performance. However, the abandonment of the widely-applied U-Net architecture by DiTs and their improvements [18] 10 26 on latent-space image generation tasks triggers our curiosity, because the U-Net inductive bias is always believed to help denoising. Hence, we rethink deploying DiTs on a canonical U-Net architecture.
In order to experiment with the combination of U-Net with DiT, we first propose a naive DIT in U-Net style (DiT-UNet) and compare it with an isotropic DIT of similar size. Results turn out that
'Equal Contribution.
*Corresponding Author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

60
DNT-8V2
U-DIT8
SIT
DIT-UZ
U-DIT-L
SIT-LLAMA
DIT-UU2
U-DIT (0urs)
50
40
FID.50x
FID-50K
20
20
10-
10"
200
600
800
Transforer GFLOPs
Training lberations (K)
Figure I: Comparing U-DiTs with DiTs and their improvements. We plot FID-50K ver- sus denoiser GFLOPs (in log scale) after 400K training steps. U-DiTs could achieve better per- formance than its counterparts.
Figure 2: The performance of U-DiTs and DiTs of various size. U-DiTs perform con- sistently better than DiTs with the increase of training steps. The marker size represents the computation cost of the model qualitatively.
DiT-UNets are merely comparable to DiTs similar computation costs. From this toy experiment, it inferred that the inductive bias of U-Net is not fully leveraged when U-Nets and plain transformer blocks simply combined.
Hence, we rethink the self-attention mechanism in DiT-UNet. The backbone in a latent U-Net denoiser provides feature where low-frequency components dominate [31]. The discovery implies the existence of redundancies backbone features: the attention module the U-Net diffuser should highlight low-frequency domains. As previous theories praised downsampling for filtering gigh-freuuency noises in diffusion 33). we seek to leverage this natural low-pass filter by performing token downsampling the features for self-attention. Unlike previous transformer works [17]44]] that downsample key-value pairs only, we radically downsample the query-key-ae key-value tuple altogether, such that self-attention is performed among downsampled latent tokens. is surprising that when incorporate self-attention with downsampled tokens into DiT-UNet, better results are achieved on latent U-Net diffusers with significant reduction of computation.
Based on this discovery, we scale U-Nets with downsampled self-attention up and propose series of State-of-the-Al U-shaped Diffusion Transformers (U-DiTs). We conduct manifold experiments to verify the outstanding performance and scalability of our U-DiT models over isotropic DiTs. As shown in Fig. [& & Fig. U-DiT's could outperform DiTss by large margins. Amazingly, the proposed U-DiT model could perform better than DiT-XL/2 which times larger in terms FLOPs.
Preliminaries
Vision Transformers. VITs 15| have introduced a transformer backbone to vision tasks by patchify- ing the input and viewing image sequence of patch tokens and have proved its effectiveness on large-scale image classification tasks. While ViTs adopt an isotropic architecture, some following adopt pyramid-lit hierarchical architecture that gradually downsamples the feature. The pyramid architecture is proved highly effective in classifica- tion and other downstream tasks. Apart from architectural improvements, some other works 41 focuses on improving the Feed-Forward Network module transformers.
works on vision transformers (37 25
19
40
Vision transformers also mainstream backbones for denoising models. IPT 6 introduces isotropic transformer backbone for denoising and other low-level tasks. Some later works 23 follow the isotropic convention, but other denoising works [38 22 shift to U-Net backbones their
22
design. The pioneering work of U-ViT and DiT 28 introduces full-transformer backbones to diffusion as denoisers.
Recent Advancements Diffusion Transformers. Following DiTs, some works investigate training and diffusion [16] 27 strategies of Diffusion Transformers. Other works focus on the
design of the DiT backhone. DiffIT [8] introduces a new fusion method for conditions; FiT (6
and Visi [10] strengthens DiT by introducing LLM tricks including RoPE2D [4 and

---

Nose 32<32x4
Σ
Naise
Noise
Σ
32H32+4
32x32-4
32H3224
32x4
32+32+4
Lineer and Reshage
Cinse and Reshape
+
Linse and Reshago
Laser
Layer Nort
FFN
Layer
Layer Nor
Transfomer Block
Layver Nora
Tronsforer
Torsferee Block
Trarffoma Block
Transfomtar
Tranformr Black
MBSA
Tinaslomer Block
Thansforee Blook
Inaefoomer Back
Layer Noms
Transfomer Block
Zmmp
Trazsforrye Block
Tranfommr Block
Traafomrer Block
Layer Norm
Thansomer Block
Embbed
Embbod
Embee
Notsed Latrot
Notsec. Latest
Natsed 32x32x44
32-32x4
32x32kx4
(a) DiT
(b) DiT- UNet
(c) U-DiT (Owss
Figure 3: The evolution from the DiT to the proposed U-DiT. Left (a): the original DiT, which uses isotropic architecture. Middle (b): DiT-UNet, which is a plain U-Net-style DiT. We try this as a simple combination of DiT and U-Net in the toy experiment. Right (c): the proposed U-DiT. We propose downsample the input features for self-attention The downsampling operation could amazingly improve DiT-UNet with a huge cut on the amount of computation.
SwishGLUU These transformer-bast based diffusion works agree on adopting isotropic architectures on latents, Lee the latent feature space is not downsampled throughout the whole diffusion model. The authors of DiT [28] even regard the inductive bias of U-Net as "not crucial".
U-Nets for Diffusion. From canonical works (20 33 13 29], the design philosophy of U-Net (30 is generally accepted in diffusion. Specifically, Stable Diffusion (29] uses a U-Net-based denoiser on the compressed latent space for resolutoon image synthesis, which is highly successful in manifold generative tasks. Some previous trials on diffusion transformers 4 [8] 11 (1
also adopt U-Net on pixel-space generation tasks; but strangely, they shifted to isotropic DiT-like structures for latent-space diffusion. Despite its popularity in pixel-space diffusion, the U-Net architecture is not widely accepted in recent transtormer-oriennedd works on latent-space diffusion.
Motivated by this, we are dedicated to investigating the potential of saansformer-backhoon U-Net on latent-space diffusion. It is noteworthy that our goal is significantly different from U-ViT Mī: U-ViT is an isotropic transformer architecture with shortcuts, but our work resort to true U-Net architectures that involves multiple stages of feature-map downsampling and upsampling.
3 Investigating U-Net DiTs in Latent
As is recapped, the U-Net architecture is widely adopted in diffusion applications; theoretical evaluations on U-Net denoisers also reveal their advantage, as downsampling U-Net stage transitions could filter noises that dominate high frequencies [39. The unprecedented desertion of isotropic architectures for latent diffusion transformers is thus counter-intuitive. We are rethinking and elucidating the potentials of rranforrer-backkoned U-Net denoisers in latent diffusion via toy experiment.
A canonical U-Net-style DIT. To start with, we propose a naive Transformer backboned U-Net denoiser named DIT-UNet by embedding DiT blocks into canonical U-Net architecture. Following previous U-Net designs, The DiT-UNet consists of an encoder and a decoder with an equal number of stages. When the encoder processes the input image by downsampling the image as stage-level amounts, the decoder scales up the encoded image from the most compressed stage to input size. At each encoder stage transition, spatial downsampling by the factor of performed while the feature dimension is doubled as well. Skip connections are provided at each stage transition. The skipped feature is concatenated and fused with the upsampled output from the previous decoder stage, replenishing information loss to decoders brought by feature downsampling. Considering the
3

---

small, cramped latent space (32x 32 for 256x256-sized generation), we designate 3 stages in total, i.e. the feature is downsampled two times and subsequently recovered to its original size. In order to fit time and condition embeddings for various feature dimensions across multiscale stages, we use independent embedders for respective stages. In addition, we avoid patchifying the latent, as the U-Net architecture itself downsamples the latent space and there is no need for further spatial compression.
Via toy experiments, we compare the proposed U-Net-style DiT with the original DiT that adopts an isotropic architecture. In order align the model with the DiT design, we repeatedly use plain DiT blocks in each stage. Each DiT block includes self-attention module as the token mixer and a two-layer feed-forward network as the channel mixer. We conduct the experiment by training the U-Net-Style DiT for 400K iterations and compare it with DiT-S/4 which is comparable in size. All training hyperparameters are kept unchanged. occurs that the U-Net style DiT only gains a limited advantage over the original isotropic DiT. The inductive bias of U-Net is insufficiently utilized.
ImageNet 256x256
Model
GFLOPs
FID
sFID
ISt
Precision
Recallf
DiT-S/4 DiT-UNet
1.41
97.85
21.19
13.27
0.26
0.41
1,40
93.48
20.41
14.20
0.27
0.42
DrT-UNet+Key-
Value Value Downsampling
0.91
94.38
23.21
14.32
0.27
0.40
DiT-UNet+Token Downsampling (Ours)
0.90
89.43
21.36
15.13
0.29
0.44
Table 1: Toy experiments on U-Net-style DiTs. The naive DiT-UNet performs slightly better than the isotropic DiT-S/4; but interestingly. when we apply token downsampling for self-attention, the DiT-UNet performs better with fewer costs.
Improved U-Net-style DiT via token downsampling, In seeking to incorporate attention in trans- formers to diffusion U-Nets better, review the role of the U-Net backbone as the diffusion denoiser. A recent work on latent diffusion models (31] conducted frequency analysis on intermediate features from the U-Net backbone, and concluded that energy concentrates at the low-frequency domain. This frequency-domain discovery hints at potential redundancies in the backbone: the U-Net backbone should highlight the coarse object from global perspective rather than the high-frequency details.
Naturally, we resort to attention with downsampled tokens. The operation of downsampling is a natural low-pass filter that discards high-frequency components. The low-pass feature of down- sampling has been investigated under the diffusion scenario, which concludes that downsampling helps denoisers in diffusion automatically "discards those higher- frequency subspaces which dominated by noise" (39]. Hence, we opt to downsample tokens for attention.
In fact, attention to downsampled tokens is not new. Previous works regarding vision transformers [7 44) have proposed methods to downsample key-value pairs for computation cost reduction. Recent work acceleration of diffusion models (32 7
also applies key-value downsampling
Stable
Diffusion models. But these works maintain the number of queries, and thus the downsampling operation is not completely performed. Besides, these downsampling measures usually involves reduction of tensor size, which could result in a significant loss in information.
Different from these works, we propose a simple yet radical token downsampling method for DiT- UNets: we downsample queries, keys, and values at the same time for ceffusion-friennde self-attention, but meanwhile we keep the overall tensor size to avoid information loss. The procedure is detailed as follows: the feature-map input is first converted into four 2x downsampled features by the downsampler (the downsampler design is detailed in Sec. 4.2) Then, the downsampled features mapped to Q. K, V for self-attention. Self-attention is performed within each downsampled feature. After the attention operation, the downsampled tokens are spatially merged as a unity to recover the original number of tokens. Notably, the feature dimension is kept intact during the whole process. Unlike U-Net downsampling, are not reducing increasing the number of elements in the feature during the downsampling process. Rather, we send four downsampled tokens into self-attention in a parallel manner.
Self-attention with downsampled tokens does help DiT-UNets on the task of latent diffusion. As shown in Tab. I the substitution of downsampled self-attention to full-scale self-attention brings

---

slight improvement in the Fréchet Inception Distance (FID) metric despite significant reduction in FLOPs.
Complexity analysis. Apart from the performance benefits, we are aware that adopting downsampled self-attention in the U-Shaped DiT could save much as 1/3 of the model's overall computation cost. We conduct a brief computation complexity analysis on the self-attention mechanism to explain where the savings come from.
Given input feature of size N N and dimension d, denote Q, K,V E R mapped query-key-value tuples. The complexity of self-attention is analyzed as:
== AV O(NDD)
S.t.
A= Softmax (QK*) O(NDD
In the proposed self-attention on downsampled tokens, four sets of downsampled query-key-value tuples 4 x (Qu2, K2 Vi2) )xd performs self-attention respectively. While each self-attention operation costs only 1/16 of full-scale self-attention, the total cost for downsampled self-attention is 1/4 of full-scale self-attention. 3/4 of the self-attention computation saved via token downsampling.
In a nutshell, we show from toy experiments that the redundancy of DiT-UNet is reduced by downsampling the tokens for self-attention.
4 Scaling the Model Up
Based on the discovery in our toy experiment, we propose series of U-shaped DiTs (U-DiT) by applying the downsampled self-attention (proposed Sec. 3 and scaling U-Net-Style DiT up.
Settings. We adopt the training setting of DiT. The same VAE (i.e. sd-vae-ft-ema) for latent diffusion models [29] and the AdamW optimizer is adopted. The training hyperparameters are kept unchanged. including global batch size 256, learning rate le 4, weight decay 0, and global seed O. The training conducted with the training set of ImageNet 2012 [1]. We used NVIDIA A A100s (80G) train U-DiT-B and U-DiT-L models. The training overhead is listed in the appendix.
Apart from the self-attention on downsampling introduced the toy experiment (Section 3), further introduce a series of modifications to U-DiTs, including cosine similarity attention 24 22]
RoPE2D 34 26 10), depthwise conv FFN 38 44 and re- re-parametrization 14 35 The
contribution of each modification is quantitatively evaluated in Sec.
4.1 U-DiT at Larger Scales
Comparison with DīTs and their improvements. In order to validate the effectiveness of the proposed U-DiT models beyond simple toy experiments, we scale them and compare them with DIT's 28 of larger sizes. For fair comparison, we use the same sets of training hyperparameters as DiT; all models are trained for 400K iterations. The results on ImageNet 256x256 are shown in Tab. 2 where we scale U-DiTs to 6c9, 20c9, 80c9 FLOPs respectively and compare them with DiTs of similar computation costs, more details about the U-DiT architectures are shown in Tab.
could be concluded from Tab. 2 that all U-DiT models could outcompete their isotropic counterparts by considerable margins. Specifically, U-DiT-S and U-DiT-B could outperform DiTs of comparable size by 30 FIDs; U-DiT-L could outperform DiT-XL/2 by ~ 10 FIDs. It is shocking that U-DiT-B could outcompete DiT-XL/2 with only 1/6 of the computation costs. In Tab. 3 we further demonstrate the advantage of U-DiTss over several competitive diffusion transformers the advantage of our method better, we also include the performance of U-DiTs an FID-50K versus
28 g 18). . To present
FLOPs plot (Fig. []. Apart from DiTs and U-DiTs, we also include other state-of offte-of-rn the-art methods:
SiT (27 that proposes an interpolant framework for DiTs, and SiT-LLaMA [10 that combines state-of-the-ae DiT backbone VisionI. LAMA and SiT. The advantages of U-DiT's over other baselines are prominent the plot. The results highlight the extraordinary scalability of the proposed U-DiT
models.
5

---

ImageNet 256x256
Model
FLOPs(G)
FID
sFIDL
1St
Precisionf
Recallt
DIT-S/2
68.40
DIT-S/2"
6.07
67.40
11.93
20.44
0.368 0.543
0.559
U-DiT-S (Ours)
6.04
31,51
8.97
51.62
0.633
DIT-L/4 28
19.70
45.64
DIT-L/4"
19.70
46.10
9.17
31.05
0,472
0.612
DIT-B/2 28
23.01
43.47
DIT-B/2
23.02
42.84
8.24
33.66
0.491
0.629
U-DiT-B (Ours)
22.22
16.64
6.33
85.15
0.642
0.639
DIT-L/2 28
80.71
23.33
DiT-L/2
80.75
23.27
6.35
59.63
0.611
0.635
DIT-XL/2 28
118.64
19.47
DIT-XL/22
118.68
20.05
6.25
66.74
0.632
0.629
U-DiT-L (Ours)
85.00
10.08
5.21
112.44
0.702
0.631
Table Comparing U-DiTs against DiT' on ImageNet 256 256 generation. Experiments with supermarks are replicated according to the official code of DIT. We compare models trained for 400K iterations with the standard training hyperparameters of DIT. The performance of U-DiTs is outstanding: U-DiT-B could beat DiT-XL/2 with only 1/6 of inference FLOPs; U-DiT-L could outcompete DiT-XL/2 by FIDs.
ImageNet 256x256
Model
FLOPs (G)
FIDJ
sFID
ISt
Precision
Recall
U-VIT-L
11
76.4
21.22
6.10
67.64
0.615
0.633
U-ViT-XL"
113.0
18.35
5.75
76.59
0.632
0.630
DŁT-XL/2
28
118.7
20.05
6.25
66.74
0.632
0.629
PixArt-a-XL/2"
8
118.4
24.75
6.08
52.24
0.612
0.613
DifAT-XL/2- [8]
118.5
36.86
6.53
35.39
0.540
0.613
U-DiT-B (Ours)
22.2
16.64
6.33
85.15
0.642
0.639
U-DiT-L (Ours)
85.0
10.08
5.21
112.44
0.702
0.631
Table 3: Comparing U-DiTs against competitive diffusion architectures on ImageNet x56x 256 generation. Since different architectures use different training settings, we align them under the official -oOK-iteration setting of DiT for fair comparison. The proposed U-DiT series could outperform these models by large margins fewer FLOPs. Experiments with supermarks include necessary modifications of the original work (detailed in the appendix).
U-DiTs are also performant in generation scenarios with classifier-free guidance. In Tab. we compare U-DiTs with DiTs at efg = 1.5. For a fair comparison, we train U-DiTs and DiTs for 400K iterations under identical settings.
Extended training steps. We evacuate the potentials of U-DITs by extending training steps to Million. Fig. further demonstrate that the advantage of U-DiTs is consistent all training steps. As training steps gradually goes up to Million, the performance U-DiTs is improving (Tab. We
ImageNet 256x256
Model
Cig-Scale
FLOPs(G)
FIDJ
sFID
ISt
Precisiont
Recallf
DIT-L/2"
1.5
80.75
7.53
4.78
134.69
0.780
0.532
DIT-XL/2"
1.5
118.68
6.24
4.66
150.10
0.794
0.514
U-DiT-B
1.5
22.22
4.26
4.74
199.18
0.825
0.507
U-DiT-L
1.5
85.00
3.37
4.49
246.03
0.862
0.502
Table 4: Generation performance with classifier-free guidance. We measure the performance of U-DiTs and DiTs at 400K training steps with cfg = 1.5. Experiments with supermark are replicated according to the official code of DIT. U-DiTs are also performant on conditional generation.

---

U-DIT-B
U-DIT-L
200K
400K
600K
800K
200K
400K
600K
800K
Figure 4: Quality improvements of generated samples as training continues. We sample from U-DiT models trained for different numbers of iterations on ImageNet 256x256. More training does improve generation quality. Best viewed on screen.
visualize the process where the image quality is gradually getting better (Fig Notably, U-DiT-L at only 600K training steps could outperform DiT-XL/2 at 7M training steps without classifier guidance. As additionally shown in Fig. U-DiT models could conditionally generate authentic images at merely 1M iterations.
Larger image size. We additionally compare the generation performance of U-DiT-B and DIT-XL/2 on ImageNet 512 x 512 under exactly the same training setting. As shown in Tab. 66] U-DiT-B could still outcompete DiT-XL/22 that is approximately 5 times larger in FLOPS.
ImageNet 256x 256
Model
Training Steps
FIDJ
sFID
ISt
Precisiont
Recallf
DFT-XL/2
7M
9.62
U-DiT-B
200K
23.23
6.84
64.42
0.610
0.621
U-DiT-B
400K
16.64
6.33
85.15
0.642
0.639
U-DiT-B
600K
14.51
6.30
94.56
0.652
0.643
U-DiT-B
800K
13.53
6.27
98.99
0.654
0.645
U-DIT-B
1M
12.87
6.33
103.79
0.661
0.653
U-DiT-L
200K
15.26
5.60
86.01
0.685
0.615
U-DiT-L
400K
10.08
5.21
112.44
0.702
0.631
U-DiT-L
600K
8.71
5.17
122.45
0.705
0.645
U-DiT-I
800K
7.96
5.21
131.35
0.705
0.648
U-DiT-I
1M
7.54
5.27
135.49
0.706
0.659
Table The performance of U-DiT-B and U-DiT-L models with respect to training iterations.
The unconditional generation performance of both models ImageNet 256x256 consistently improves as training goes on, where U-DiT-L at 600K steps strikingly beats DiT-XL/2 at 7M steps.

---

ImageNet 512x512
Model
FLOPs (G)
FID
sFIDL
IS1
Precision
Recallf
DIT-XL122
524.7
20.94
6.78
66.30
0.745
0.581
U-DiT-B
106.7
15,39
6.86
92.73
0.756
0.605
Table 6: Comparing U-DiTs against DiTs on ImageNet 512 x 512 generation. Experiments with supermark are replicated according to the official code of DiT. We compare models trained for 400K iterations with the standard training hyperparameters DĪT.
ImageNet 256
Model
FLOPs(G)
FIDJ
sFID
1ST
Precision
Recall
Pixel Shuffle (PS) Depthwise (DW) Conv. + PS
0.89
96.15
23.90
13.93
0.272
0.389
0.91
89.87
20.99
14.92
0.288
0.419
DW Conv. II Shorteut PS
0.91
89,43
21.36
15.13
0.291
0.436
Table 7: Ablations on the choice of downsampler. We have tried several downsampler designs, and turns out that the parallel connection of shortcut and depthwise convolution is the best fit. We avoid using ordinary convolution (i.e. Conv.+PS) because channel-mixin -ixing is costly: conventional convolution- downsamplers could double the amount of computation. The U-DiT with conventional downsampler costs as many as 2.22G FLOPs in total.
4.2 Ablations
The design of downsampler. The downsampling operation in the proposed U-DiT transforms a complete feature into multiple spatially downsampled features. Based on previous wisdom, we figured out that previous works either directly perform pixel shuffling, or apply a convolution layer before pixel shuffling. While we hold that it is much too rigid to shuffle pixels directly downsampling, applying convolution is hardly affordable in terms of computation costs. Specifically, ordinary convolutions are costly as extensive dense connections on the channel dimension are involved: using convolution- baseod downsamplers could double computation costs. As a compromise, we apply depthwise convolution instead. We also add a shorteut that short-circuits this depthwise convolution, which has proved crucial for better performance. The shorteut adds negligible computation cost to the model, and in fact, could removed during the inference stage with re- re-parameterization tricks. The results are shown in Tab. 7
The contribution of each individual modification. In this part, we start from a plain U-Net-style DiT (DiT-UNet) and evaluate the contribution of individual components. Firstly, we inspect the advantage of downsampled self-attention. Recapping the toy experiment results in Sec. 3 replacing the full-scale self-attention with downsampled self-attention would result in an improvement in FID and 1/3 reduction in FLOPs. In order evaluate the improvement of downsampling via model performance, we also design a slim version of DiT-UNet (i.e. DIT-UNet (Slim).. The DiT-UNet (Slim) serves as a full-scale self-attention baseline that spends approximately the same amount (~ 0.9GFLOPs) of computation as our U-DiT. As shown in the upper part of Tab. ( by comparing U-DiT against DiT-UNet (Slim), it turns out that downsampling tokens in DiT-UNet could bring a performance improvement of 18FIDs.
Next, inspect other modifications that further refine U-DiTs (lower part of Tab. 9.. Swin Trans- former V2 [24] proposes a stronger variant of self-attention: instead of directly muluplying Q and K matrices, cosine similarities between queries and keys are used. We apply the design to our self- attention, which yields ~ 2. 2.5FIDs of improvement. RoPE [34] is powerful positional embedding method, which has been widely applied in Large Language Models. Following the latest diffusion transformer works [26] [0]. we inject 2-dimensional RoPE (RoPE2D) into queries and keys right before self self-attention. The introduction of RoPE2D improves performance by 2.5FIDs. Some recent transformer works strengthen MLP by inserting a depthwise convolution layer between two linear mappings [3
As the measure proved effective in these works, we borrow it to our U-DiT model, improving 5FIDs. As re-parametrization during training [14] could improve model performance, we apply the trick to FFN (35) and bring an additional improvement of 3.5FIDs. Above all, based on these components, the proposed U-DiTs are further improved.
8

---

Apart from the modifications that improve U-DiT, it is worth noting that vanilla U-DiTs (i.e. U-DiTs without any the modifications mentioned above) are still competitive. According to Tab. I0 U-DiT-L could still achieve ~ 8FIDs of advantage over DIT-XI/2.
vanilla
Model
Params (M)
FLOPs (G)
Channel
Head Number
Encoder-Decolee
U-DiT-S
52.05
6.04
96
4
12, 5, 8, 5, 21 [2, 5, 21
U-DiT-B
204.43
22.22
192
8
U-DIT-L 810.19
85.00
384
16
[2, 2.
Table 8: Configurations of U-DiTs architecture with different model sizes. Channel represents the initial output channel number first layer. Encoder-Decode denotes the transformer block number of encoder and decoder module.
ImageNet 256x256
Model
FLOPs(G)
FID
sFID
IS1
Precisiont
Recall
DiT-UNet (Slim)
0.92
107.00
24.66
11.95
0.230
0.315
DIT-UNet
1,40
93.48
20.41
14.20
0.274
0.415
U-DiT-T (DiT-UNe++Dowwampling)
0.91
89.43
21.36
15.13
0.291
0.436
U-DiT-T (iCos.Sim.)
0.91
86.96
19.98
15.63
0.299
0.450
U-DIT-T (+RoPE2D)
0,91
84.64
19.38
16.19
0.306
0.454
U-DiT-T (+DWconv FFN)
0.95
79.30
17.84
17.48
0.326
0.494
U-DiT-T (+Re-param.)
0.95
75.71
16.27
18.59
0.336
0.512
Table 9: Ablations on U-DiT components. Apart from the toy example in Sec. 3 we further validate
the effectiveness of downsampled by comparing the U-DiT with slimmed version of DiT-UNet at equal FLOPs. Results reveal that downsampling could bring ~ 18FIDs on DiT-UNet. Further modifications top of the U-DiT architecture could improve 5 FIDs each.
ImageNet 256x256
Model
FLOPs(G)
FID
SFID
ISt
Precisiont
Recall
U-DiT-S (Vanilla)
5.91
41.01
10.96
39.29
0.489
0.622
U-DiT-S (+All Mods)
6.04
31.51
8.97
51.62
0.543
0.633
U-DiT-B (Vanilla)
21.96
20.89
7.33
72.85
0.611
0.637
U-DiT-B (+All Mods)
22.22
16.64
6.33
85.15
0.642
0.639
U-DiT-L (Vanilla)
84.48
12.04
5.37
102.63
0.684
0.628
U-DiT-L (+AIl Mods)
85.00
10.08
5.21
112.44
0.702
0.631
Table 10: Comparison between vanilla U-DiTs and improved U-DiTs with all modifications. With negligible extra computational overhead, the proposed modifications could improve the performance of U-DiT; but it is worth noting that vanilla U-DiTs are powerful enough against DiTs.
5 Conclusion
In this paper, we lay emphasis on DiTs in U-Net architecture for latent-space generation. Though "sotropic-architecuree DiT' have proved their strong scalability and outstanding performance, the effectiveness of the U-Net inductive bias is neglected. Thus, we rethink DiTs in the U-Net style. We first conduct an investigation and DiT blocks, and try to reduce computation redundancy in the U-Net backbone. Inspired by previous wisdom diffusion, we propose to downsample the visual tokens for self-attention and yield extraordinary results: the performance further improved despite a huge cut on FLOPs. From this interesting discovery, scale the U-Net architecture up and propose series of U-shaped DiT models (U-DiTss We have done various experiments to demonstrate the outstanding performance and scalability of our U-DiTs.
plain DiT-UNet, which is straightforward combination of U-Net
Limitations. For lack of computation resources and tight schedule, this time we could not further extend training iterations and scale the model size up to fully investigate the potential of U-DiTs.

---

Figure 5: Generated samples by U-DiT-L at 1M iterations. It astonishing that U-DiT could achieve authentic visual quality at merely Million training steps. Best viewed on screen.
Broader Impacts. Due the biases in the training data set, the generated content may contain pornographic, racist, hate and violent information. But we emphasize that the potential for misuse is mitigated through vigilant application.
Discussion of Safeguards. For cautionous usage, we suggest an algorithm capable of checking generated images, in order identify and mitigate content that contravenes legal ethical usages.
Acknowledgement. This work is supported by the National Key R&D Program of China under grant No. 2022ZD0160302 and the National Natural Science Foundation of China under grant No. 62276007. We gratefully acknowledge the support of MindSpore, CANN, and Ascend Al Processor used for this research.
References
[1] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: vit backbone for diffuuson models. In IEEE/CVF Conference Computer Vision and Pattern Recognition, CVPR 2023, Vancouver, BC, Canada, June 17-24, 2023, pages 22669-22679. IEEE, 2023.
[2] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024.
[3] Han Cai, Junyan Li. Muyan Hu, Chuang Gan, and Song Han. Efficientvit: Multi-scale linear attention for high-resolution dense prediction. arXiv preprint arXiv:2205. 14756, 2022.
[4] He Cao, Jianan Wang, Tianhe Ren, Xianbiao Qi, Yihao Chen, Yuan Yao, and Lei Zhang. Exploring vision transformers diffusion learners. CoRR, 1377 2022.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-ent object detection with transformers. In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm, editors, Computer Vision ECCV 2020 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part I, volume 12346 of Lecture Notes in Computer Science, pages 213 229, Springer, 2020.
10

---

[6] Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu. Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu, and Wen Gao. Pre-trained image processing transformer. CoRR, abs/2012.00364, 2020.
[7] Junsong Chen, Chongjjan Ge. Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-E: Weak-to-strong training of diffusion transformer for 4k text-to-image generation. CoRR, abs/2440300462. 2024.
[8] Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-imase image synthesis, 2023.
[9] Xiangyu Chen, Xintao Wang. Wenlong Zhang, Xiangtao Kong, Yu Qiao, Jiantao Zhou, and Chao Dong. Hat: Hybrid attention transformer for image restoration. arXiv preprint arXiv:2309.052:: 2023.
[10] Xiangxiang Chu, Jianlin Su, Bo Zhang, and Chunhua Shen, Visionllama: A unified Ilama interface for vision tasks. CoRR, abs/2403.00522. 2024.
[11] Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Z. Kaplan, and Enrico Shippole. Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers. CoRR, abs/2401.11603 2024.
[12] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vixion and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami, Florida, USA, pages 248-255. IEEE Computer Society, 2009.
[13] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthesis. In Marc' Aurelio Ranzato, Alina Beygelzimer. Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan. editors, Advances in Neural Information Processing Systems 34; Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 8780-88 0794, 2021.
[14] Xiaohan Ding, Xiangyu Zhang, Ningning Ma, Jungong Han, Guiguang Ding, and Jian Sun. Repvgg: Making vgg-style convncts great again. In Proceedings of the IEEE/CVF conference on computer vision and puttern recugnition, pages 13733-13742, 2021.
[15] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. CoRR,
abs/2010.1 1929, 2020.
[16] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Mdtv2: Masked diffusion transformer is strong image synthesizer, 2024.
[17] Jianyuan Guo. Kai Han, Han Wu, Chang Xu, Yehui Tang, Chunjing Xu, and Yunhe Wang. CMT: convolutional neural networks meet vision transformers. CoRR, abs/2107.06263, 2021.
[18] Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. Diffit: Diffusion vision transformers for image generation. CoRR, abs/2312.021399 2023.
[19] Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon Oh. Rethinking spatial dimensions of vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pages 11916-11925, IEEE, 2021.
[20] Jonathan Ho Ajay Jain, and Pieter Abbcel. Denoising diffusion probabilistic models. CoRR, abs/2006. 1229, 2020.
[21] Emiel Hoogeboom. Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Homolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 13213- 13232. PMLR, 2023.
[22] Yawei Li, Yuchen Fan, Xiaoyu Xiang, Denis Demandolx, Rakesh Ranjan, Radu Timofte, and Lue Van Gool. Efficient and explicit modelling of image hicrarchies for image restoration. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pages 18278-18289, 2023.
[23] Jingyun Liang, Jiezhang Cao, Guolei Sun, Kai Zhang. Luc Van Gool, and Radu Timofte. Swinir: Image restoration using swin transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops, pages 1833-1844, October 2021.
11

---

[24] Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao. Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao. Zheng Zhang. Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12009-12019, 2022.
[25] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang. Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), pages 10012-10022, October 2021.
[26] Zeyu Lu, Zidong Wang. Di Huang, Chengyue Wu, Xihui Liu, Wanli Ouyang, and Lei Bai. Fit: Flexible vision transformer for diffusion model. CoRR, abs/2402. 2024.
[27] Nanye Ma. Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. CoRR, abs/2401.087 2024.
[28] William Peebles and Saining Xie. Scalable diffusion models with transformers, In IEEE/CVF International Conference on Computer Vision, ICCV 2023, Paris, France, October 1-6, 2023, pages 4172-4182, IEEE, 2023.
[29] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 10674-10685. IEEE, 2022.
[30] Olaf Ronneberger. Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegge, William M. Wells III, and Alejandro F. Frangi, editors, Medical Image Computing and Computer-Asiisted Intervention MICCAI 2015 18th International Conference Munich, Germany, October 5 9, 2015, Proceedings, Part III, volume 9351 of Lecture Notes in Computer Science, pages 234-241. Springer. 2015.
[31] Chenyang Si, Ziqi Huang, Yuming Jiang, and Ziwei Liu. Freeu: Free lunch diffusion u-net. CoRR, abs/2309.1 11497, 2023.
[32] Ethan Smith, Nayan Saxena, and Aninda Saha. Todo: Token downsampling for efficient generation of high-resolution images, 2024.
[33] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. CoRR, abs/2010.02502, 2020.
[34] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 108:127063, 2024.
[35] Zhijun Tu, Kunpeng Du, Hamting Chen, Hailing Wang, Wei Li, Jie Hu, and Yunhe Wang. 1pt-v2: Efficient image processing transformer using hierarchical attentions, 2024.
[36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, abs/1706.03766 2017.
[37] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: versatile backbone for dense prediction without convolutions. CoRR, abs/2102.12122. 2021.
[38] Zhendong Wang, Xiaodong Cun, Jianmin Bao, Wengang Zhou, Jianzhuang Liu, and Houqiang Li. Uformer: A general u-shaped transformer for image restoration. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 17683-17693, 2022.
[39] Christopher Williams, Fabian Falck, George Deligiannidis, Chris C. Holmes, Amnud Doucet, and Saifuddin Syed. A unified framework for u-net design and analysis. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 16, 2023, 2023.
[40] Yixing Xu, Chao Li, Dong Li, Xiao Sheng, Fan Jiang, Lu Tian, and Ashish Sirasao. Fdvit: Improve the hicrarchical architecture vision transformer. In Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), pages 5950-5960, October 2023.
[41] Yixing Xu, Chao Li. Dong Li, Xiao Sheng, Fan Jiang. Lu Tian, Ashish Sirasao, and Emad Barsoum. Enhancing vision transformer: Amplifying non- linearity in feedforward network module. In Forty-first International Conference on Machine Learning, 2024.
12

---

[42] Syed Waqas Zamir, Aditya Arora. Salman Khan. Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Restormer: Efficient transformer for high-resolution image restoration. In CVPR, 2022.
[43] Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yahiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip H. S. Tomr and Li Zhang. Rethinking semantic segmentation from sequence-to- sequence perspective with transformers. CoRR, abs/2012.15840, 2020.
[44] Yupeng Zhou, Zhen Li, Chun-Le Guo, Song Bai, Ming-Ming Cheng, and Qibin Hou. Srformer: Permuted self-attention for single image super resouutssolt.on. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12780-12791, 2023.
13

---

A Appendix supplemental material
A.1 Details about Downsampling
Given input tuple of (queries, keys, values) QKV (shape=(b, 3c, w)). we firsty conduct Pixel-UnShufll operation on QKV, and get four spatially downsampled QKV (shape=1 x (bs* hc,h/s. w/s)). Then we perform vanilla multi-head self-attention, and get four downsampled output (shape=1 x (bx h/s. w/))). Fi- nally, we merge the four downsampled outputs into unity via Pixel Sixel-Shufflin (shape=(b, c, h/s, w/s)). Through- out the process, that the entire upsampling and downsampling process was completely lossless: the feature maps have not gone through lossy downsampling like bicubic or bilinear downsampling.
not only significantly reduced the computational overhead of sclf-attention, but also ensured
A.2 Additional Experiment Details
Training Overhcad. We report the training speed in Table comparable to that of DIT-XL/2.
11
The training speed of vanilla U-DiT-L
ImageNet 256
Model
TS (Steps/Sec)
FID
sFID
ISt
Precision
Recallf
DIT-XL/2" 28
1.71
20.05
6,25
66.74
0.632
0.629
U-DiT-B (Vanilla)
3.14
20.89
7.33
72.85
0.611
0.637
U-DiT-L (Vanilla)
1.55
12.04
5.37
102.63
0.684
0.628
U-DiT-L (+All Mods)
0.84
10.08
5.21
112.44
0.702
0.631
Table 11:
The training overhead of D/T-X1 and U-DiTs. "TS" stands for training speed, measured
steps per second on NVIDIA A100 (80G).
Experiment Details Table 3 Since different diffuusion architectures use different settings, we are dedicated to comparing them under identical settings for fair comparison. We adopt the 400K -iteration training setting of DLT-XL/2 28). Here some further details regarding certain baselines:
U-VIT-XL: increase the depth U- V-T-L from to 30 in order match FLOPs DIT-XL/2. We encounter loss explosion while training U- U-VTT-H (133.25 GFLOPs) on the codebase of DiTs.
2. PixArt-a-XL/2: As the original model is text-to-image model, removed its cross attention module for texts.
3. Diffi-XL/2:: This model is not open-sourced at the moment of this publication. Since it is a variant of DiT-XI/2, replicated the time-dependent self-attention (TMSA) based on the codes of DiT. Unfortunately, the performance gets worse compared the original DiT-XI/2.
Additional Visual Results. Due to large file size, we are unable to provide visual results in the appendix.
Please refer to the supplementary materials for two high-quality visual result demos.
14

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics. and societal impact. Do not remove the checklist: The papers not including the checklist will desk rejected. The checklist should follow references and precede the (optional) supplemental material, The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant
information is Not Available.
Please provide a short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked "[Yes] is generally preferable to "[No] ". is perfectly acceptable to answer "[No] provided a proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset used"). In general, answering "[No] or "[NA] is not phrased in binary way. acknowledge that the answer write a justification elaborate.
the checklist as of the factors in their evaluation. While
grounds for rejection. While the questions often more nuanced, so please just use your best judgment evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
supporting
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract demonstrates motivation, the proposed ideas and brief summary of experiment results.
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 in 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 he expected to generalize other setrings.
It is fine to include aspirational goals motivation long it is clear that these goals attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The paper has discussed the limitations of the work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors encouraged create separate "Limitations" section in their paper.
15

---

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 sell-sppeiiiiaatinn. asympiotic 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be 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-tex system might not be used reliably to provide closed captions for online lectures because fails to 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 fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't I acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Question:
Answer: [NA]
Justification: The paper does not inlcude theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross-referencedd
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 if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper fully discloses all the information needed to reproduce the main experimental results of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless whether the code and data are provided or not.
the contribution is a 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 novel architecture, describing the architecture fully might suffice, or if the contribution is 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 the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of model checkpoint, or other means that appropriate to the research performed.
16

---

While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. a large language model), then there should either be a way to access this model for reproducing the results a way to reproduce the model (e.g.. with open-source datasct instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of booeddsouuce models, it may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other rescarchers 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The paper will provide open access to the data and code during camera ready period.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines hattpa://nips.cc/pple cc/pulicc/
for more details.
While we encourage the release of code and data, we understand that this might not be possihle, so "No" an acceptable answer. Papers cannot rejected simply for not including code, unless this is central to the contribution (e.g.. for a 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 hattpa://nipe http:///iipsccccppbp/e cc/public/ htteee//oodSutmiisssoooooooi for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental Setting/Detail
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer. etc.) necessary to understand results?
Answer: [Yes]
Justification: This paper has specified all the training and test details.
Guidelincs:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [NA]
Justification: This not relevant this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

The authors should answer "Yes" if the results are accompanied by enoor bars, confidence intervals, statistical significance tests, at least for the experiments that support the main claims the paper.
The factors of variability that the error bars are capturing should be elearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given
experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clcar whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error hars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors not verified.
asymmetric distributions, authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (c.g. negative error rates). crror bars are reported tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
The paper has indicated sufficient information on the computer resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the lype of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for cach of the individual experimental runs well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e-g., preliminary or failed experiments that didn' make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https::/neuriip: wtwppeuuipppccppppaappeeeeeeeeee hctppp/ccc///thiiiiiicddeeeeee
Answer:
[Yes]
Justification: This research conducted in the paper conform, in every respect, with the NeurIPS Code of Fthics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure preserve anonymity (e.g., if there is special consideration due laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: This paper has discussed both potential positive societal impacts and negative societal impacts of the work performed.
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.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance). fairness considerations (e.g., deploy- ment lechnologies that could make decisions that unfairly impact specific groups), privacy considcrations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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, it not needed point out that generic algorithm for optimizing neural networks could enable people 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 or unintentional) misuse of the technology.
If there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitor- ing misuse, mechanisms to 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 relcase of data or models that have a high risk for misuse (e.g. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
This paper has described safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by requiring 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 we 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?
Answer: [Yes]
Justification: The utilization of code, data and models in this paper is in accordance with the license and the terms.
Guidelines:
The answer NA mcans that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version the asset is used and, possible, include a URL..
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should he provided. For popular datasets, papprsswi thcode. com/ datasetaatetaaaaaaaaaa has curated licenses for some datasets. Their licensing guide can help determine the license a dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
this information not available online, the authors are encouraged reach out to the asset'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]
19

---

Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/codc/model as part of their sub- missions via structured templates. This includes details about training, license, limitations,
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, other lahor should be paid at least minimum wage country data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review,
20


Reviews:
[{"review_id": "z46DODTeDZ", "reviewer": "Reviewer_Df8i", "summary": "The authors conduct a simple toy experiment by comparing a U-Net architectured DiT with an isotropic one. They find that the U-Net architecture only gains a slight advantage, indicating potential redundancies within the U-Net-style DiT.  Inspired by the discovery that U-Net backbone features are low-frequency-dominated, they perform token downsampling on the query-key-value tuple for self-attention and observed performance improvements. \n\nBased on self-attention with downsampled tokens, the authors propose a series of U-shaped DiTs (U-DiTs) in the paper and conduct extensive experiments to demonstrate the good performance of U-DiT models.", "questions": "- if the downsampling factor is 2, how much layers will this downsampled design support?", "limitations": "The authors may not have to pursue even longer iterations and comparison with DiTs. More justification on why downsampling tokens would help is more important. Also extending the experiments other than ImageNet would be better.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The author verified that radical token downsampling method for DiT-UNets could save the overall computation cost compared to full-scale self-attention, while can still improve the performance.\n\n- the authors also performed scaling experiments to compare with various scaled DiTs.", "weaknesses": "- there is a lack of comparison with some noticeable existing works, e.g., the authors did not compare with U-ViT, which is pretty similar as the plain UNet based DiT. There is also no mentioning or comparison with HourglassDiT. \n- there is no comparison with reducing tokens by simply increasing the patch size.\n- the saved computation is not clear, line 141 says 1/3 is saved while 149 claims 3/4.\n- there is no justification or explanation why radically downsampling of K, V, Q is better than only k-v downsampling.\n- there is no explanation on why radically reduced tokens would gain better performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer Df8i,\n\nThank you very much for your suggestions. Here are our responses:\n\n**W1.** We omitted U-ViT and Hourglass DiT previously, because these models are mainly targeted at pixel-space generation, while our work is focused on latent space generation. To demonstrate the advantage of our model, we conducted experiments as follows: \n\n|                           | GFLOPs | FID       | sFID      | IS        | Precision | Recall   |\n| ------------------------- | ------ | --------- | --------- | --------- | --------- | -------- |\n| DiffiT | 118.5   | 36.862 | 6.533 | 35.391 | 0.540 | 0.613 |\n| PixArt-XL | 118.4   |   24.751 | 6.075 | 52.237 | 0.612 | 0.613 |\n| U-ViT-Large | 76.4   | 21.218 | 6.100 | 67.644 | 0.615 | 0.633 |\n| Hourglass-DiT * | 53.8    | 564.574 | 747.528 | 1.000 | 0.000 | 0.000 |\n| U-DiT-B (Ours) | **22.2** | 16.64     | 6.33     | 85.15      | 0.642     | **0.639** |\n| U-DiT-L (Ours) | 85.0     | **10.08** | **5.21** | **112.44** | **0.702** | 0.631     |\n\nP.S. H-DiT is intended for pixel-space generation instead of latent space generation (which is the usage of our U-DiT model). The officially provided H-DiT models are very small (because they are intended for High-resolution pixel space generation), so we scale them up for 12x (which is the maximal size we could scale). The training of the enormously large H-DiT is not stable.\n\n\n  Plus, we want to stress that **U-ViT is not a U-Net architecture**: it is an isotropic architecture with shortcuts that does not involve feature downsampling. Our U-DiT model, on the other hand, is a U-Net architecture with feature downsampling. Besides feature downsampling, we are further adding attention downsampling to the self-attention module.\n\n**W2.** Thank you very much for your suggestions. Following the setting of Table 1, we have conducted an experiment that uses patchsize 2 as follows:\n\n|                           | GFLOPs | FID       | sFID      | IS        | Precision | Recall   |\n| ------------------------- | ------ | --------- | --------- | --------- | --------- | -------- |\n| Token Patchification      | 0.88   | 129.12    | 34.02     | 9.30      | 0.17      | 0.21     |\n| Token Downsampling (Ours) | 0.90   | **89.43** | **21.36** | **15.13** | **0.29**  | **0.44** |\n\nResults reveal that patchsize=2 performs worse than our proposed downsampling method. \n\n**W3.** We apologize for the ambiguities. \"1/3\" in line 141 refers to FLOPs saved in total (on the entire model); \"3/4\" in line 149 refers to FLOPs saved within self-attention. Apart from self-attention, there are many other components in the model. We will clarify the meaning of these fractions in the next revision.\n\n**W4.** We have provided a brief explanation in line 121-123, and here is some further clarification: KV Compression keeps the number of queries intact (which corresponse to the number of output tokens), which means downsampling is not performed completely on the feature map; the effect of noise filtering is thus reduced. We have downsampled queries for self-attention.  Besides, KV Downsampling involves a lossy reduction of Key-Value pair tokens; our Token Downsampling measure does not involve lossy reduction of tensors.\n\n\n**Q1.** All multi-head self-attention (MHSA) layers in U-DiT support downsampling tokens by 2, and we did apply downsampling in all MHSA layers in actual practice.\n\n\nSincerely,\n\nAuthors", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OhEsBrfWsV", "reviewer": "Reviewer_AGQh", "summary": "This paper proposes a transformer architecture as backbone for\ndiffusion modeling that is based on the UNet. The paper shows that a\nvariation of the transformer with downsampling layers and skip\nconnections achieves better results than the DiT at a lower\ncompoutational cost.", "questions": "- There is a repeated use of the word isotropic, seemingly to refer to\n  the standard transformer. I wonder what is the justification for\n  this characterization.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Experimental results at multiple model scales show the proposed\n  architecture obtains better performance than the original DiT\n  architecture, at the same or lower computational cost\n- The experiments suggest token downsampling for attention shows\n  promising results and could be a worthile avenue for improvement.", "weaknesses": "- The novelty and technical contribution is limited. The proposed\n  method consists of a minor architectural modification to the U-ViT\n  architecture.\n- Experiments are limited to ImageNet, and the FID scores achieved are\n  far from state-of-the-art (FID<2).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer AGQh,\n\nThank you very much for your comments.  Here are our responses:\n\nW1.. **Novelty is limited.**\n\nOur work is not an improvement of U-ViT. The architecture of our U-DiT model is **completely different from U-ViT**: we adopt a U-Net architecture, while U-ViT is an isotropic architecture with skipping shortcuts. The major architectural difference is that U-DiT (Ours) is a encoder-decoder model, where each part have several stages where downsampling or upsampling is used as stage transition; U-ViT does **not involve any downsampling or upsampling**. The feature size is not changed throughout the whole model. \n\nWe have also demonstrated via experiments that our model is much better than theirs.\n\n|                           | GFLOPs | FID       | sFID      | IS        | Precision | Recall   |\n| ------------------------- | ------ | --------- | --------- | --------- | --------- | -------- |\n| U-ViT-Large | 76.4   | 21.218 | 6.100 | 67.644 | 0.615 | 0.633 |\n| U-DiT-B (Ours) | **22.2** | 16.64     | 6.33     | 85.15      | 0.642     | **0.639** |\n| U-DiT-L (Ours) | 85.0     | **10.08** | **5.21** | **112.44** | **0.702** | 0.631     |\n\nW2. **About limited dataset and not reaching FID<2.** \n\nThe reason we use ImageNet is that latest latent-space conditional generation models, like DiT, SiT, DiT-LLaMA, DiffiT all use ImageNet as the only benchmark. There are two versions of the ImageNet dataset: ImageNet-256 and ImageNet-512. As experiments in the paper are performed on ImageNet-256, we further add the performance of our model on ImageNet-512 to prove its robustness:\n\n|                           | GFLOPs | FID       | sFID      | IS        | Precision | Recall   |\n| ------------------------- | ------ | --------- | --------- | --------- | --------- | -------- |\n| DiT-XL | 524.7 | 20.94 | **6.78** | 66.30 | 0.645     | 0.581  |\n| U-DiT-B (Ours) | **106.7** | **15.39** | 6.86 | **92.73** | **0.756** | **0.605** |\n\nWe are unable to pursue SOTA of FID<2 because it requires extremely large computation costs that we could not afford. For instance, DiT-XL needs 7M iterations to reach an FID of 2. We have compared with a SOTA model DiffiT (that reaches FID 1.73) under exactly the same 400K iteration setting. This indicates that our model is among the SOTA models in latent image generation.\n\n|                           | GFLOPs | FID       | sFID      | IS        | Precision | Recall   |\n| ------------------------- | ------ | --------- | --------- | --------- | --------- | -------- |\n| DiffiT | 118.5   | 36.862 | 6.533 | 35.391 | 0.540 | 0.613 |\n| U-DiT-B (Ours) | **22.2** | 16.64     | 6.33     | 85.15      | 0.642     | **0.639** |\n| U-DiT-L (Ours) | 85.0     | **10.08** | **5.21** | **112.44** | **0.702** | 0.631     |\n\nQ1. **The meaning of term \"isotropic\":** Yes, it refers to standard transformer where a stack of transformer blocks are concatenated in series. No downsampling or upsampling of the tokens is required within the model.\n\n\nSincerely,\n\nAuthors", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hAcg59CxIC", "reviewer": "Reviewer_GsTx", "summary": "This paper introduces a U-shaped diffusion Transformer (U-DiT) model, inspired by the departure from U-Net in DiT. The authors aim to combine the strengths of U-Net and DiT to determine if the inductive bias of U-Net can enhance DiTs. Initially, a simple DiT-UNet model was developed, but it showed minimal improvement over DiTs. Consequently, the authors explored downsampling QKV values in self-attention, which not only enhanced the model's performance but also significantly reduced computational costs. Leveraging this downsampled self-attention mechanism, a series of U-DiT models were proposed, with experimental results demonstrating their effectiveness.", "questions": "1. Why do the tricks fail when the model is scaled up?\n2. Please provide a comparison of parameter counts.", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper introduces an interesting idea that downsampled self-attention can reduce redundancy in U-Net while achieving performance improvements rather than losses. The authors support their claims with extensive experiments, demonstrating the convincing performance advantages of U-DiTs.", "weaknesses": "1. The statement in line 68, \"the latent feature space is not downsampled,\" is incorrect for DiTs. DiTs reduce the spatial size of features by patchifying the latent features.\n2. The paper presents several techniques that are claimed to be effective; however, their benefits decrease as the model scales up. For instance, in Table 8, the combined use of four techniques only results in a 2FID improvement for U-DiT-L.\n3. While parameter count is a critical measure of model size, the paper does not provide a comparison of parameter counts with baseline models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer GsTx,\n\nThank you very much for your comments. Here are our responses:\n\n**W1:** Statement \"the latent feature space is not downsampled\" is misleading.\n\n**A1**: Thanks for your suggestions. We will add qualifiers to this statement and limit it to the intermediate features.\n\n**W2&Q1:** Why do the tricks fail when the model scaled up?\n\n**A2:** Thank you for your insightful comments. (1) According to OpenAI’s scaling laws, model performance tends to be log-linear as model size increases. This means that as models become larger, the performance gains from the proposed tricks in our study, may appear less pronounced in absolute terms. This phenomenon occurs because larger models quickly approach the performance ceiling for the given evaluation metrics. (2) However, it is important to note that the improvements introduced by our tricks remain significant even for the largest models. These models are already nearing the upper limits of performance, and the relative improvements achieved by our methods are still impactful. Thus, while the absolute gain may be reduced, the effectiveness of our tricks is evident and noteworthy.\n\n**W3&Q2:** Please provide a comparison of parameter counts.\n\n**A3:** Thank you for your suggestion. The comparison of parameters is shown in the following table. With less FLOPs and a little more parameters, our proposed U-DiT-L still supress DiT-XL/2 significantly.\n\n| Model     | FLOPs (G)| Params (M) | FID |\n|:------------:|:----------------:|:--------:|:----------------:|\n| **U-DiT-B (Ours)** | **22.22** | **204.42** |  **16.64** |\n| DiT-L/2   | 80.73 | 458.10 |  23.33 |\n| DiT-XL/2   | 118.66  | 675.13  | 19.47 |\n| **U-DiT-L (Ours)** | **85.00** | **810.19** |  **10.08** |\n\nTypically, U-shaped models tend to have a larger number of parameters but lower computational overhead, whereas isotropic architecture models generally have higher computational requirements but fewer parameters. It is quite challenging to simultaneously align the parameter and computational overhead between these two types of models, and addressing this will be a primary focus of our future work.\n\nSincerely,\n\nAuthors", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bu69Zt0Uvl", "reviewer": "Reviewer_JEXU", "summary": "The paper presents a novel approach to enhancing diffusion models for image generation using Transformers. It introduces U-shaped Diffusion Transformers (U-DiTs) that employ token downsampling within the self-attention mechanism of U-Net architectures, aiming to reduce computational cost while maintaining or improving image generation quality. Computation efficiency is eagerly needed for the visual foundation model. This paper demonstrates that U-DiTs outperform traditional isotropic diffusion transformer models with significantly lower computation requirements. The paper also provides detailed ablations and analyses.", "questions": "(1) How do you get the GFLOPs as illustrated in the paper? How to estimate the GPU hours according to such GFLOPs?", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The introduction of token downsampling in the self-attention process of U-Net style transformers is novel and addresses key efficiency issues in image generation tasks.\n\n(2) The paper provides extensive experimental results showing U-DiTs achieving better performance metrics than larger, more computationally intensive models.\n\n(3) The reduced computational demand of U-DiTs could make high-quality image generation more accessible and cost-effective.", "weaknesses": "(1) This paper needs to include important baselines. It fails to compare with other efficient DiT models, such as Pixart-sigma or others. SiT seems to be the only baseline method that takes flow as the objective. However, it does not employ a different architecture other than DiT, which should not be a very fair baseline to compare with U-DiTs.\n\n(2) This paper only used ImageNet as the training data. No large-scale experiments have been conducted on Laion, JourneyDB, or others to verify its capacity for text-to-image generation. This limits the potential impact of this work.\n\n(3) The operation of token downsampling is not clear. I am very confused about the \"After the attention operation, the downsampled tokens are spatially merged as a unity to recover the original number of tokens.\" Fig 3 is also not clear since the down-sampler seems to be a black box. This is one of the most crucial parts of this paper. I would highly recommend the authors add more details on token merging, downsampling, and upsampling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer JEXU,\n\nThank you very much for your review. Here are our responses:\n\n**W1: Fail to compare with other efficient DiT models, like PixArt-Alpha.**\n\nThanks for your advice. Here we provide a comparison with powerful baselines: PixArt-Alpha [1], as well as U-ViT [2] and DiffiT-XL [3]. We align all methods under the same setting as the standard setting (Table 2). The FLOPs statistics and metric results are shown in the Table below. Our method has a clear advantage over other methods.\n\n|                | GFLOPs   | FID       | sFID     | IS         | Precision | Recall    |\n| -------------- | -------- | --------- | -------- | ---------- | --------- | --------- |\n| DiffiT-XL      | 118.5    | 36.86     | 6.53     | 35.39      | 0.540     | 0.613     |\n| PixArt-Alpha   | 118.4    | 24.75     | 6.08     | 52.24      | 0.612     | 0.613     |\n| U-ViT-Large    | 76.4     | 21.22     | 6.10     | 67.64      | 0.615     | 0.633     |\n| U-DiT-B (Ours) | **22.2** | 16.64     | 6.33     | 85.15      | 0.642     | **0.639** |\n| U-DiT-L (Ours) | 85.0     | **10.08** | **5.21** | **112.44** | **0.702** | 0.631     |\n\n**W2: No large-scale experiments for text-to-image generation.**\n\nWe are sorry that text-to-image generation training are too resource-intensive for us to conduct. We carefully consider this as a future work in order to verify its potentials on text-to-image generation.\n\n**W3: The operation of token downsampling is not clear. Please add more details on token merging, downsampling, and upsampling.** \n\nThank you for your feedback. We appreciate your keen observations and will ensure to include these additional details in our revised manuscript. Actually, given a input tensor QKV_0 (shape=(b, 3c, h, w)), we firstly conduct PixelUnShuffle operation on QKV_0, and get four smaller feature QKV_1 (shape=(b$\\times$s^2, 3c, h/s, w/s)). Then we perform vanilla multi-head self-attention, and get the output Y (shape=(b$\\times$s^2, c, h/s, w/s)). Finally, we reshape the Y into the original shape (b, c, h, w). Throughout the process, we not only significantly reduced the computational overhead of self-attention, but also ensured that the entire upsampling and downsampling process was completely lossless. This is a key reason why our approach significantly outperforms the other methods listed in Table 1 of the paper.\n\n**Q1: How do you get the GFLOPs as illustrated in the paper? How to estimate the GPU hours according to such GFLOPs?**\n\nWe got the GFLOPs via the Python library \"torchprofile''[1], which is a \"a general and accurate MACs / FLOPs profiler for PyTorch models\". We are willing to provide training speed statistics in the table below. For fair comparison, we have removed all tricks to demonstrate the actual training time difference caused by the proposed architecture.\n\n|                   | GPU hours (400K) | FID       | sFID     | IS         | Precision | Recall    |\n| ----------------- | ---------------- | --------- | -------- | ---------- | --------- | --------- |\n| DiT-XL/2          | 519              | 20.05     | 6.25     | 66.74      | 0.632     | 0.629     |\n| U-DiT-L (Vanilla) | 573              | **12.04** | **5.37** | **102.63** | **0.684** | 0.628     |\n| U-DiT-B (Vanilla) | **283**          | 20.89     | 7.33     | 72.85      | 0.611     | **0.637** |\n\n\n\nSincerely,\n\nAuthors\n\n[1] \"PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis.\" ICLR 2024.\n\n[2] \"All are Worth Words: A ViT Backbone for Diffusion Models.\" CVPR 2023.\n\n[3] \"DiffiT: Diffusion Vision Transformers for Image Generation.\" ECCV 2024.\n\n[4] https://github.com/zhijian-liu/torchprofile", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper improves UNet based on downsampling layers and skip connections for diffusion modeling. The paper received diverging ratings with a common concern around limited novelty. It is indeed not the most original idea to add downsampling/upsampling and skip connections to a neural architecture. But the paper does show convincing results that prove token downsampling inside the transformer is better than patchification at the input for a DiT. Hence this paper could inform future efforts on improving the popular DiT architectures and training recipes. Other than that, the majority of reviewers found that their concerns were sufficiently addressed during discussion and recommend acceptance. The AC shares the majority opinion.

It is recommended to integrate all the discussions/comparisons to the final paper. Furthermore, it's key to address the remaining concern from reviewer AGQh on the performance and scope of experiments. Given that the performance is far from the SOTA frontier, the authors argue for the trade-off between FID and training cost/model size. Then the authors need to answer one important question in the final paper: would this approach allow achieving the highest sample quality, or would the proposed token downsampling limit the model capacity at some point?
Decision: Accept (poster)