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

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.

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. (2) The paper provides extensive experimental results showing U-DiTs achieving better performance metrics than larger, more computationally intensive models. (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. (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

(1) How do you get the GFLOPs as illustrated in the paper? How to estimate the GPU hours according to such GFLOPs?

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

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.

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

1. Why do the tricks fail when the model is scaled up? 2. Please provide a comparison of parameter counts.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

n/a

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.

Strengths

- Experimental results at multiple model scales show the proposed architecture obtains better performance than the original DiT architecture, at the same or lower computational cost - The experiments suggest token downsampling for attention shows promising results and could be a worthile avenue for improvement.

Weaknesses

- 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

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

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

3

Limitations

Yes

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.

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

- if the downsampling factor is 2, how much layers will this downsampled design support?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

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.

Additional Comments on "why radically reduced tokens would gain better performance"

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.

Follow-up Question

Thank you for the detailed response. I have a follow-up question. What GPU that was used for your model training?

Response to JEXU

Sorry for not making the config clear. We use 80G A100. We will add this detail to the manuscript in the next revision.

Thank you for your response, and my remaining concerns.

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.