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

Dear reviewer JEXU,

Thank you very much for your review. Here are our responses:

W1: Fail to compare with other efficient DiT models, like PixArt-Alpha.

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

W2: No large-scale experiments for text-to-image generation.

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

W3: The operation of token downsampling is not clear. Please add more details on token merging, downsampling, and upsampling.

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

Q1: How do you get the GFLOPs as illustrated in the paper? How to estimate the GPU hours according to such GFLOPs?

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

Sincerely,

Authors

[1] "PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis." ICLR 2024.

[2] "All are Worth Words: A ViT Backbone for Diffusion Models." CVPR 2023.

[3] "DiffiT: Diffusion Vision Transformers for Image Generation." ECCV 2024.

[4] https://github.com/zhijian-liu/torchprofile

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.

Dear reviewer GsTx,

Thank you very much for your comments. Here are our responses:

W1: Statement "the latent feature space is not downsampled" is misleading.

A1: Thanks for your suggestions. We will add qualifiers to this statement and limit it to the intermediate features.

W2&Q1: Why do the tricks fail when the model scaled up?

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.

W3&Q2: Please provide a comparison of parameter counts.

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.

Typically, 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.

Sincerely,

Authors

Dear reviewer AGQh,

Thank you very much for your comments. Here are our responses:

W1.. Novelty is limited.

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

We have also demonstrated via experiments that our model is much better than theirs.

W2. About limited dataset and not reaching FID<2.

The 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:

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

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

Sincerely,

Authors

Dear reviewer Df8i,

Thank you very much for your suggestions. Here are our responses:

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:

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

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.

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:

Results reveal that patchsize=2 performs worse than our proposed downsampling method.

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.

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.

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.

Sincerely,

Authors