U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers

Authors: yuchuan  tian
zhijun  tu
hanting  chen
jie  hu
chao  xu
yunhe  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 19:26:47

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.

Main paper: 
 introduction: 
 Thanks to the attention mechanism that establishes long-range spatial dependencies, Transformers [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] , et cetera.Interestingly, DiTs have discarded the U-Net architecture [30] that is universally applied in manifold previous works, either in pixel [20; 13] or latent space [29] . The use of isotropic (i.e. standard transformer; a plain stack of transformer blocks) architectures in DiTs is indeed successful, as scaledup 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 Figure 1 : Comparing U-DiTs with DiTs and their improvements. We plot FID-50K versus denoiser GFLOPs (in log scale) after 400K training steps. U-DiTs could achieve better performance than its counterparts. DiT-UNets are merely comparable to DiTs at similar computation costs. From this toy experiment, it is inferred that the inductive bias of U-Net is not fully leveraged when U-Nets and plain transformer blocks are simply combined.Hence, we rethink the self-attention mechanism in DiT-UNet. The backbone in a latent U-Net denoiser provides a feature where low-frequency components dominate [31] . The discovery implies the existence of redundancies in backbone features: the attention module in the U-Net diffuser should highlight low-frequency domains. As previous theories praised downsampling for filtering high-frequency noises in diffusion [39] , we seek to leverage this natural low-pass filter by performing token downsampling on the features for self-attention. Unlike previous transformer works [17; 44; 32] that downsample key-value pairs only, we radically downsample the query-key-value tuple altogether, such that self-attention is performed among downsampled latent tokens. It is surprising that when we incorporate self-attention with downsampled tokens into DiT-UNet, better results are achieved on latent U-Net diffusers with a significant reduction of computation.Based on this discovery, we scale U-Nets with downsampled self-attention up and propose a series of State-of-the-Art 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. 1 & Fig. 2 , U-DiTs could outperform DiTs by large margins. Amazingly, the proposed U-DiT model could perform better than DiT-XL/2 which is 6 times larger in terms of FLOPs.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The abstract demonstrates our motivation, the proposed ideas and a brief summary of experiment results.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 be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
preliminaries: 
 Vision Transformers. ViTs [15] have introduced a transformer backbone to vision tasks by patchifying the input and viewing an image as a sequence of patch tokens and have proved its effectiveness on large-scale image classification tasks. While ViTs adopt an isotropic architecture, some following works on vision transformers [37; 25; 19; 40] adopt a pyramid-like hierarchical architecture that gradually downsamples the feature. The pyramid architecture is proved highly effective in classification and other downstream tasks. Apart from architectural improvements, some other works [3; 41] focuses on improving the Feed-Forward Network module in transformers.Vision transformers are also mainstream backbones for denoising models. IPT [6] introduces an isotropic transformer backbone for denoising and other low-level tasks. Some later works [23; 22; 9] follow the isotropic convention, but other denoising works [38; 42] shift to U-Net backbones as their design. The pioneering work of U-ViT [1] and DiT [28] introduces full-transformer backbones to diffusion as denoisers. We propose to downsample the input features for self-attention. The downsampling operation could amazingly improve DiT-UNet with a huge cut on the amount of computation.SwishGLU. These transformer-based diffusion works agree on adopting isotropic architectures on latents, i.e. 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 high-resolution image synthesis, which is highly successful in manifold generative tasks. Some previous trials on diffusion transformers [4; 18; 11; 21] 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 transformer-oriented works on latent-space diffusion.Motivated by this, we are dedicated to investigating the potential of Transformer-backboned U-Net on latent-space diffusion. It is noteworthy that our goal is significantly different from U-ViT [1] : 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.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 are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it 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, a worse outcome might be that reviewers discover limitations that aren't 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 not penalize honesty concerning limitations. 
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 transformer-backboned U-Net denoisers in latent diffusion via a 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 a 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 2 is 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 small, cramped latent space (32× 32 for 256×256-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 to align the model with the DiT design, we repeatedly use plain DiT blocks in each stage. Each DiT block includes a 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 Improved U-Net-style DiT via token downsampling. In seeking to incorporate attention in transformers to diffusion U-Nets better, we 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 a 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 downsampling has been investigated under the diffusion scenario, which concludes that downsampling helps denoisers in diffusion as it automatically "discards those higher-frequency subspaces which are 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 [17; 44] have proposed methods to downsample key-value pairs for computation cost reduction. Recent work on acceleration of diffusion models [32; 7] also applies key-value downsampling on 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 a 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 diffusion-friendly 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 2× downsampled features by the downsampler (the downsampler design is detailed in Sec. 4.2). Then, the downsampled features are 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, we are not reducing or 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. 1, the substitution of downsampled self-attention to full-scale self-attention brings slight improvement in the Fréchet Inception Distance (FID) metric despite a 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 as 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 an input feature of size N × N and dimension d, we denote Q, K, V ∈ R N 2 ×d as mapped query-key-value tuples. The complexity of self-attention is analyzed as:X = AV O(N 4 D) s.t. A = Softmax QK T O(N 4 D).In the proposed self-attention on downsampled tokens, four sets of downsampled query-key-value tuples 4× (Q ↓2 , K ↓2 , V ↓2 ) ∈ R ( N 2 )2 ×d 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 is 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?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 proofs in the paper should be numbered and cross-referenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but 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 or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
scaling the model up: 
 Based on the discovery in our toy experiment, we propose a series of U-shaped DiTs (U-DiT) by applying the downsampled self-attention (proposed in 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 1e -4, weight decay 0, and global seed 0. The training is conducted with the training set of ImageNet 2012 [12] . We used 8 NVIDIA A100s (80G) to train U-DiT-B and U-DiT-L models. The training overhead is listed in the appendix.Apart from the self-attention on downsampling as introduced in the toy experiment (Section 3), we further introduce a series of modifications to U-DiTs, including cosine similarity attention [24; 22] , RoPE2D [34; 26; 10] , depthwise conv FFN [38; 3; 44] , and re-parametrization [14; 35] . The contribution of each modification is quantitatively evaluated in Sec. 9.Comparison with DiTs and their improvements. In order to validate the effectiveness of the proposed U-DiT models beyond simple toy experiments, we scale them up and compare them with DiTs [28] of larger sizes. For a fair comparison, we use the same sets of training hyperparameters as DiT; all models are trained for 400K iterations. The results on ImageNet 256×256 are shown in Tab. 2, where we scale U-DiTs to ∼ 6e9, ∼ 20e9, ∼ 80e9 FLOPs respectively and compare them with DiTs of similar computation costs, more details about the U-DiT architectures are shown in Tab. 8.It 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-DiTs over several competitive diffusion transformers [1; 28; 8; 18] . To present the advantage of our method better, we also include the performance of U-DiTs in an FID-50K versus FLOPs plot (Fig. 1 ). Apart from DiTs and U-DiTs, we also include other state-of-the-art methods: SiT [27] that proposes an interpolant framework for DiTs, and SiT-LLaMA [10] that combines state-of-the-art DiT backbone VisionLLaMA and SiT. The advantages of U-DiTs over other baselines are prominent in the plot. The results highlight the extraordinary scalability of the proposed U-DiT models. Table 7 : Ablations on the choice of downsampler. We have tried several downsampler designs, and it turns out that the parallel connection of a shortcut and a depthwise convolution is the best fit. We avoid using ordinary convolution (i.e. Conv.+PS) because channel-mixing is costly: conventional convolution-based downsamplers could double the amount of computation. The U-DiT with a conventional downsampler costs as many as 2.22G FLOPs in total.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 as 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-based downsamplers could double computation costs. As a compromise, we apply depthwise convolution instead. We also add a shortcut that short-circuits this depthwise convolution, which has proved crucial for better performance. The shortcut adds negligible computation cost to the model, and in fact, it could be removed during the inference stage with 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 to 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. 9, 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, we inspect other modifications that further refine U-DiTs (lower part of Tab. 9). Swin Transformer V2 [24] proposes a stronger variant of self-attention: instead of directly multiplying Q and K matrices, cosine similarities between queries and keys are used. We apply the design to our selfattention, which yields ∼ 2.5FIDs of improvement. RoPE [34] is a powerful positional embedding method, which has been widely applied in Large Language Models. Following the latest diffusion transformer works [26; 10], we inject 2-dimensional RoPE (RoPE2D) into queries and keys right before 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 [38; 3; 44] . As the measure is 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.Apart from the modifications that improve U-DiT, it is worth noting that vanilla U-DiTs (i.e. U-DiTs without any of the modifications mentioned above) are still competitive. According to Tab. Question: 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 of whether the code and data are provided or not)?Answer: [Yes] Justification: The paper fully discloses all the information needed to reproduce the main experimental results of the paper. 
conclusion: 
 In this paper, we lay emphasis on DiTs in U-Net architecture for latent-space generation. Though isotropic-architectured DiTs 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 on plain DiT-UNet, which is a straightforward combination of U-Net and DiT blocks, and try to reduce computation redundancy in the U-Net backbone. Inspired by previous wisdom on diffusion, we propose to downsample the visual tokens for self-attention and yield extraordinary results: the performance is further improved despite a huge cut on FLOPs. From this interesting discovery, we scale the U-Net architecture up and propose a series of U-shaped DiT models (U-DiTs). We have done various experiments to demonstrate the outstanding performance and scalability of our U-DiTs.Limitations. For lack of computation resources and tight schedule, at this time we could not further extend training iterations and scale the model size up to fully investigate the potential of U-DiTs. Broader Impacts. Due to 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 to identify and mitigate content that contravenes legal or ethical usages.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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 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 the results?Answer: [Yes] Justification: This paper has specified all the training and test details.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [NA] Justification: This is not relevant to this paper. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal 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.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and 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 is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes] Justification: This paper has described safeguards.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 means 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 of the asset is used and, if 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 a 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 be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA]  
 
 Appendix: 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 be desk rejected. The checklist should follow the 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] , [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 be published with the paper.The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it 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 we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.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 the questions and only use the provided macros for your answers.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a 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 or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to 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 closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. 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 type 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 each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).Justification: This research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


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

Reviews: review: paper_summary: paper_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. 

Based 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.
strength_weakness: strength_weakness: - 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. 
- there is no comparison with reducing tokens by simply increasing the patch size.
- the saved computation is not clear, line 141 says 1/3 is saved while 149 claims 3/4.
- there is no justification or explanation why radically downsampling of K, V, Q is better than only k-v downsampling.
- there is no explanation on why radically reduced tokens would gain better performance.
questions: questions: - if the downsampling factor is 2, how much layers will this downsampled design support?
limitations: 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.
paper_summary: paper_summary: This paper proposes a transformer architecture as backbone for
diffusion modeling that is based on the UNet. The paper shows that a
variation of the transformer with downsampling layers and skip
connections achieves better results than the DiT at a lower
compoutational cost.
strength_weakness: strength_weakness: - The novelty and technical contribution is limited. The proposed
  method consists of a minor architectural modification to the U-ViT
  architecture.
- Experiments are limited to ImageNet, and the FID scores achieved are
  far from state-of-the-art (FID<2).
questions: questions: - There is a repeated use of the word isotropic, seemingly to refer to
  the standard transformer. I wonder what is the justification for
  this characterization.
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. While parameter count is a critical measure of model size, the paper does not provide a comparison of parameter counts with baseline models.
questions: questions: 1. Why do the tricks fail when the model is scaled up?
2. Please provide a comparison of parameter counts.
limitations: limitations: n/a
paper_summary: paper_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.
strength_weakness: strength_weakness: (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.

(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.

(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.
questions: questions: (1) How do you get the GFLOPs as illustrated in the paper? How to estimate the GPU hours according to such GFLOPs?
limitations: limitations: N/A

Comments: title: Thank you to reviewer AGQh
Thank You and Further Responses
Thank you for your response, and my remaining concerns.
Response to JEXU
Thanks
Authors' Response to Follow-up Question
Follow-up Question
Additional Comments on "why radically reduced tokens would gain better performance"
comment: Dear reviewer AGQh,

At the very end of the discussion period, we want to express our thanks to reviewer AGQh for their help in improving our paper. May I kindly ask if there is any other concerns that we could help you? Thank you!
Thanks for your approval and thanks again for your helpful suggestions. Here are our responses:

### Q1: Limited improvement in Table 7.

The key is that we also need to take model **FLOPs** into consideration in comparison. In Table 7, though the proposed downsampling could improve a mere 4 FID, it reduces the FLOPs by **more than 1/3** of the DiT-UNet model. In order to evaluate the advantage of downsampling through FID improvement, we evaluated a smaller "DiT-UNet (Slim)" model that is comparable in FLOPs to U-DiT-T. The advantage measured under comparable FLOPs is around **28 FID**, which is nearly **twice** the improvement of all the tricks. Below we provide a copy of Table 7 (upper section) for easy reference:

|                                     | GFLOPs | FID       | sFID      | IS        | Precision | 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-UNet+Downsampling)** | 0.91   | **89.43** | 21.36     | **15.13** | **0.291** | **0.436** |

### Q2: Filtering out high-frequency noise is harmful.

We provide a further explanation as follows: estimating noise is equivalent to estimating clear image via simple transformation, which implies that perception of clear denoised signal is vital in the denoiser. Additionally, shortcuts passes most high-frequencies, while the backbone is low-frequency dominated [1]. Hence, filtering out high-frequencies in the backbone is not causing significant impacts.

### Q3: Terminology issues with "U-Net" and "Isotropic".

Thank you for suggestions on terminology. We agree that the term "U-Net" may cause misunderstanding. After referring to Ronneberger et al. [2], we hold that "U-Shaped Architecture/Network" is a better substitute. However, we found it hard to find a term substitute to "Isotropic". We will explain it as "a standard transformer architecture that does not involve any change in token size" in the next revision.

### Q4: Comparison to UViT (in Simple Diffusion [3]).

We hold that our method is different from UViT [3] as follows:

1. Task difference: UViT experiments are conducted on **pixel-space**; DiT and U-DiT experiments are conducted on **latent-space**.

2. Training setting difference: UViT is using **batch-size 2048** for 500K iterations on the ImageNet-256 benchmark; DiT and our U-DiT uses batch-size 256 (which is only **1/8** of the batch-size of UViT). The authors of [3] themselves claim that "the batch size is larger (2048) which does affect FID and IS performance considerably".

3. Model size difference: UViT has **2 Billion** parameters, which is more than **2 times** the size of the largest variant of U-DiT (U-DiT has only 810M).

4. Architectural difference: UViT **only** uses Transformer Block at the medium stage; it keeps using **conventional ResBlock** at the encoder-decoder stage (Fig. 7 in [3]). Our U-DiT model uses **Transformer Blocks across all stages**.

Above all, we hold that UViT is **not fairly comparable** to DiT and U-DiT, but we do thank the reviewer for providing a competitive diffusion architecture, and we will discuss it in Section 2: Preliminaries in the next revision. Additionally, we sincerely apologize for not being able to test the performance UViT on the same setting of DiT due to limited time left for discussion and closed-source UViT codes.

[1] Freeu: Free lunch in diffusion u-net. CVPR 2024.

[2] U-Net: Convolutional Networks for Biomedical Image Segmentation. MICCAI 2015.

[3] Simple diffusion: End-to-end diffusion for high resolution images. ICML 2023.
I thank the authors for their response to my comments and to the other
reviewers.

I still have some concerns about the paper:

- Table 7 shows the additional modifications  are providing an
  improvement of 15 FID points, and that the downsampling alone is
  actually not that significant (although such high FID scores provide a
  weak signal). This raises the question of how much of the
  improvement in Table 2 is actually due to the downsampling
  introduced in Section 3, which is the main focus of the paper.

- I am not fully convinced that filtering out high-frequency noise is
  always an advantage for diffusion, as claimed by the authors. In
  many diffusion settings one seeks to predict the noise, in which
  case this could be harmful. 

- The writing and terminology could be improved to clearly state which
  elements of the original U-Net are being referred to. When the
  authors say "we adopt a U-Net architecture", it can be confusing
  since the original U-Net is fully convolutional. As another example,
  I remain doubtful about the use of the word "isotropic", which
  should mean something that "is the same in every direction". However
  in this case it doesn't seem any particular latent dimension is
  treated differently. 
  
- Note there is the UVit architecture introduced in [1] which also
  includes downsampling/upsampling layers and signifcanlty outperforms
  DiT-XL (which is not a strong ImageNet baseline).


While I maintain these concerns, after reading the other reviewers I
realize the practical value this paper can have to the community in
providing directions to improve the transformer architecture for diffusion, and therefore raise
my score by 1 point.


[1] Hoogeboom, E., Heek, J. and Salimans, T., 2023, July. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning (pp. 13213-13232). PMLR.
Sorry for not making the config clear. We use 80G A100. We will add this detail to the manuscript in the next revision.
80G A100 or 40G A100?
Thank you very much for your encouraging feedback, and thank you very much for your help in improving our paper as well!
Thank you for your helpful suggestions, and sorry for not specifying the devices. We used 8 NVIDIA A100s for training. The GPU-hours statistics were all measured on NVIDIA A100. We promise to add the results of more competitive baselines, as well as the table of FLOPs/GPU-Hour statistics, to Section 4 in later revisions.
Thank you for the response and the additional experiments. In my view, the idea of combining U-Net with token downsampling for generation is very promising, as simple downsampling alone may not be sufficient for generative tasks. Besides, the FID@400K results look good. The authors also show the efficiency compared to other baselines.

After considering the comments from other reviewers, I decided to raise my score to 7 (Accept).
Thank you for the detailed response. I have a follow-up question. What GPU that was used for your model training?
Sorry for not making why "reduced tokens would gain better performance" clear. On one hand, low-frequency dominates the backbone of U-Net, and thus downsampling would cause little information loss; on the other hand, downsampling could filter out high-frequency noises (according to line 115-117) and thus being beneficial for diffusion. That's why reduced tokens would gain better performance.