Unveiling the Secret of AdaLN-Zero in Diffusion Transformer

Authors: Jie Zhu, Mingyu Ding, Boqiang Duan, Leye Wang, Jingdong Wang
Primary area: generative models
Keywords: Diffusion transformer, zero-initialization, image generation

Abstract:
Diffusion transformer (DiT), a rapidly emerging architecture for image generation, has gained much attention. However, despite ongoing efforts to improve its performance, the understanding of DiT remains superficial. In this work, we delve into and investigate a critical conditioning mechanism within DiT, adaLN-Zero, which achieves superior performance compared to adaLN. Our work studies three potential elements driving this performance, including an SE-like structure, zero-initialization, and a “gradual” update order, among which zero-initialization is proved to be the most influential. Building on this insight, we heuristically leverage Gaussian distributions to initialize each condition modulation, termed adaLN-Gaussian, leading to more stable and effective training. Extensive experiments following DiT on ImageNet1K demonstrate the effectiveness and generalization of adaLN-Gaussian, e.g., a notable improvement of 2.16% in FID score over adaLN-Zero.

Decision: reject
Reviewer scores: 5, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1729689775723

thread_type:
official_review

<thread_object>
creator: reviewer_nddd

<content>
summary:
This paper investigates AdaLN-Zero within the DiT architecture and proposes three potential reasons for the performance difference between AdaLN-Zero and AdaLN: the incorporation of an SE-like structure, the use of an effective zero-initialized value, and a gradual weight update order. The authors conduct a series of analyses, ultimately concluding that the zero initialization is the most significant factor. Building on these findings, they introduce a method that leverages Gaussian distributions to initialize each condition modulation, referred to as adaLN-Gaussian. Extensive experiments demonstrate an improvement of 2.16% in FID on ImageNet1K.

strengths:
The presentation is clear, three potential reasons for the performance difference between AdaLN-Zero and AdaLN are studied in detail, and the designed experiment is reasonable.

weaknesses:
1. The contribution of this paper is limited; the only significant advancement is the introduction of adaLN-Gaussian, which leverages Gaussian distributions to initialize each condition modulation in the DiT architecture. The performance gains are also limited.
2. The experiments conducted are insufficient. Training for 800k iterations may not be adequate for the convergence of the DiT model. Given the 2% improvement in the current results at 400k iterations, it raises doubts about whether Gaussian initialization will outperform zero initialization in terms of final performance. Furthermore, to demonstrate the broader applicability of this method, I recommend conducting additional experiments on transformer-based models such as SiT[1] and PixArt-alpha[2] (text-to-image).

[1] Ma, Nanye, et al. "Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers." arXiv preprint arXiv:2401.08740 (2024).
[2] Chen, Junsong, et al. "Pixart-$\alpha $: Fast training of diffusion transformer for photorealistic text-to-image synthesis." arXiv preprint arXiv:2310.00426 (2023).

questions:
1. How do you summarize the differences between AdaLN-Zero and adaLN as the three points?
2. In Fig. 4, it appears that regardless of the type of initialization used, the weight distribution converges to a similar state after a certain number of training iterations. This raises the question of why the choice of initialization has such a significant impact on performance.
3. A few minor suggestions: 1) Simplify Figure 6 by merging the subgraphs to enhance information density. 2) Rotate Figure 2 to improve the visibility of the text.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730293578339

thread_type:
official_review

<thread_object>
creator: reviewer_kchj

<content>
summary:
This paper investigates adaLN-Zero, a key component of DiT. They find three key factors that contribute to the superior performance of adaLN-Zero: an SE-like structure, a good zero-initialized value, and a gradual weight update order. The second one plays the most important role. Finally, they propose adaLN-Gaussian, which achieves better results.

strengths:
1. This paper is easy to follow, and the experiments are relatively sufficient.
2. This paper provides some inspiring conclusions.

weaknesses:
1. Part of the description needs further explanation, e.g. line 208：They also zeros out weights of all modulations including Wγ1,Wβ1, Wγ2, and Wβ2 in a block, rendering γ1, β1, γ2, and β2 zero.
2. Some of the inferences are intuitive, and it would be better if there were more rigorous analysis. e.g. line 448：Moreover, this moment should be neither too late, ... ,nor too early, as there may be minimal difference from zero-initialization.

questions:
1. From line 129~130, we know that all alpha are initialized to 0. Can you expalin line 209: why are gama1, beta1, gama2, beta2 also zero?
2. In Table 4, if there is no limit on steps, will the results of adaLN-Zero catch up with adaLN-Gaussian when training for more steps? Do you have any relevant experimental results? Maybe different initialization methods mainly affect the convergence speed.
3. In addition to adaLN-Zero, have you analyzed the distribution of parameters of other parts of DiT? Do they also end up being Gaussian? How were these parameters initialized?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730358969578

thread_type:
official_review

<thread_object>
creator: reviewer_p8jh

<content>
summary:
This work focuses on an overlooked part in the diffusion transformer - the zero initialization of adaptive LayerNorm.
Authors start from the similarity between SENet with DiT's adaLN, and develop some variants from adaLN.
After several discussions about the gradient update over the weight of adaLN and the summarization of the benefits from adaLN-zero, authors  propose to leverage Gaussian distributions to initialize the adaLN.
The exps over IN-1K and DiT-based backbones are sufficiently perferformed.

strengths:
1. The motivation is very good (the overlooked part in DiT) and the issue in this work in quite interesting.

2. Clear clarity and good presentation.

3. Provide sifficient evidence to support the exploration inside the adaLN-zero, not only the gradient update but also the varients of adaLN.

4. Good and fair experimentss to support the proposed method. Authors provide relatively big scale exps on IN-1K 512X512 which is expensive, and the backbones are mainly based on large scale DiT. Besides authors also evaluate the effectiveness on other DiT-based backbones.

weaknesses:
1. The start point is very unclear for me. How can you find the similarity between adaLN and SE archtecture？ 
I consider that this start point should be better explained and provide a structure figure of SE archtecture.

questions:
1. How can you find the similarity between adaLN and SE archtecture？ This is an intersting point. Hope that authors can provide some principles but not intuitions.

2. As shown in Figure2, the performance of Gaussian initialization is a U-shaped trending. Could you please some analysis about this trending? Why the large std can brings a relatively bad results.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730714005949

thread_type:
official_review

<thread_object>
creator: reviewer_zplc

<content>
summary:
This paper investigates three mechanisms of adaLN-Zero, including the SE-like network structure, zero-initialization, and the weight update order. Based on the analysis of adaLN-Zero, this work proposes an improved initialization method, adaLN-Gaussian, which utilizes Gaussian distributions to initialize the weights of each condition modulation.

strengths:
The proposed adaLN-Gaussian initialization is simple and effective, achieving remarkable performance.

weaknesses:
1. The overall logic of the paper is somewhat disorganized, and there is a logical gap between the analysis of adaLN-Zero and adaLN-Gaussian. Could the authors provide a more detailed explanation of why Gaussian distributions are used to initialize weights?
2. This work lacks mathematical analysis; all conclusions are drawn from experimental results and statistics, and the authors only conducted experiments using the ImageNet dataset. I think more extensive and general experiments are needed to validate the effectiveness of adaLN-Gaussian.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732111959224

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely appreciate all reviewers’ time and efforts in reviewing our paper. We are glad to find that reviewers recognized our contributions:

- **Motivation.** Very good motivation and quite interesting work [p8jh]; Easy to follow [p8jh]; Clear presentation[nDDD]
- **Method.** Simple and effective[ZpLc]; Remarkable performance[ZpLc]
- **Experiments and analysis.** Good, fair, reasonable, and sufficient experiments [p8jh, KCHJ, nDDD]; Inspiring conclusions[KCHJ]; Detailed analysis [nDDD]

And we also thank all reviewers for their insightful and constructive suggestions, which helps a lot in further improving our paper.

In addition to the pointwise responses below, we summarize some supporting experiments and analysis added in the rebuttal according to reviewers’ suggestions.

### New experiments and analysis:

- Adding more experiments on other datasets including Tinyimagenet, AFHQ, and CelebA-HQ, and on another model SiT[1]  to show the effectiveness of our method
- Providing more analysis such as about a U-shaped  trending for Gaussian initialization in Table 2
- Using Gaussian initializations for other module parts of DiT
- Longer training steps for DiT-XL/2, e.g., 5.4M, to achieve 2.27, faster than adaLN-Zero which uses 7M steps, the longest steps in DiT paper

The additional experiments and modifications to the language have been delivered in our paper and reflected in the revised version. We hope our pointwise responses below could clarify all reviewers’ confusion and alleviate all concerns. We thank all reviewers’ time again.

[1] Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. ECCV2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1733295997880

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A concise summary to assist the AC in evaluating our work

comment:
## The contributions of this work

This work starts from an important but often overlooked module in DiT called adaLN-Zero. Initially introduced in the DiT paper, adaLN-Zero demonstrated superior performance compared to its counterpart, adaLN, which piqued our interest and prompted us to investigate further. Through step-by-step analysis, we provide some interesting conclusions and propose a simple yet effective refinement with one line code replaced, showing a promising pathway for future generative models. Our contributions can be summarized as follows:

- We study three key factors that collectively contribute to the superior performance of adaLN-Zero: an SE-like structure, a good zero-initialized value, and a gradual weight update order. Among them, we find that the a good zero-initialized value plays the most pivotal role.
- Based on the distribution variation of condition modulation weights, we heuristically leverage Gaussian distributions to initialize each condition modulation, termed **adaLN-Gaussian**.
- We conduct comprehensive experiments across various settings to demonstrate adaLN-Gaussian’s effectiveness and generalization including 1) different datasets (ImageNet1K, Tiny Imagnet, AFHQ, and Celeb-HQ),  2) DiT variants (DiT-B, DiT-L, and DiT-XL), and 3) DiT based models (VisionLlama, U-DiT, and SiT). We also involve 4) different training durations. Particularly, our method takes 5400K training steps to achieve 2.27 FID, faster than adaLN-Zero which uses 7000K steps, the longest steps in DiT paper.

## The concerns from reviewers in discussion and how we alleviate them

During the past discussions, the reviewers raised several concerns regarding our paper and our responses. We summarize them and show how we alleviate it as follows:

### Concerns from Reviewer ZpLc

**Concern 1: Lack of Mathematical Analysis**

* The reviewer pointed out that the mathematical analysis they requested was missing.
* **How we alleviate it** :  We provide the mathematical analysis from two aspects.  1) Inspired by previous study[1], we investigate the variance of the response of the scale element  $\alpha$ in a condition modulation under different initialization strategies including zero-initialization (adaLN-Zero), gaussian initialization (adaLN-Gaussian), and the default xavier uniform initialization (adaLN-Mix),  analyzing the relationships between the variance and the performance of different initialization strategies. 2) We use the nearest neighbor nonparametric estimation method to demonstrate that adaLN-Zero indeed forms a Gaussian-like distribution.
  
  [1] X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks.

**Concern 2: Unclear Necessity of Gaussian Distribution for Weight Initialization**

* The reviewer was unclear about the necessity of using Gaussian distribution for weight initialization.
* **How we alleviate it** : We explain two reasons for the necessity of using Gaussian distribution for weight initialization. First, the significant improvement of adaLN-Zero over adaLN shows the importance of condition mechanism for superior performance. Thus it is necessary to select a more suitable initialization. Secondly using Gaussian is motivated by our statistics results: we find that weights in the conditioning mechanism, though zero-initialized, always transition to a Gaussian-like distribution. Thus, using a suitable Gaussian initialization could better expedite this distribution shift.

**Concern 3: Insufficient Extended Experiments**

* The reviewer noted that the extended experiments with only 50K training steps were insufficient to demonstrate the effectiveness of adaLΝ - Gaussian over longer training periods.
* **How we alleviate it** : We present the results of extended training on different datasets to prove the long - term effectiveness of adaLΝ-Gaussian, e.g., an additional 50K for Tiny ImageNet, 150 K for ImageNet1K  (SiT-XL/2). We do not extend more training steps for Tiny Imagenet, AFHQ, and CelebA-HQ as our experiments show this will lead to overfitting, an increase in FID. We also show that our method takes 5400K to achieve 2.27 FID, faster than adaLN-Zero which uses 7000K steps, the longest steps in DiT paper.

### Concerns from Reviewer p8jh and nDDD (regarding scaled training)

They are both concerned about the scaled training of our method for DiT, namely whether our method could  outperform adaLN-Zero after a long training time and could be faster to achieve the finally converged performance (2.27 FID in DiT).

* **How we alleviate it** :  We report the FID results below which show the superiority of adaLN-Gaussian in the scaled training and it is faster than adaLN-Zero to achieve 2.27 FID, which uses 7000K steps, the longest steps in DiT paper.

| Dataset | ImagetNet (cfg=1)  |  ImagetNet (cfg=1.5) |  
|----------|----------|----------|
| adaLN-Zero | 10.67 (2352K) | 2.27 (7000K)   |
| adaLN-Gaussian | 10.09 (2300K) |  2.27 (5400K) |
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper analyzes the DiT architecture of image generation, and in particular the adaLN-Zero conditioning. It identies the zero initialization as critical for performance, and proposes an alternative Gaussian initialization, which is found to stabilize and improve training. 
Reviewers appreciate the presentation, and improved performance due to adaLN-Gaussian, 
There main concerns relate to the scope of the contribution, analysis and experiments.