Unveiling the Secret of AdaLN-Zero in Diffusion Transformer

Q1: The contribution of this paper is limited; the only significant advancement is the introduction of adaLN-Gaussian

A1: We respectfully disagree with the reviewer's comment.

Our paper is not merely about a method proposal but, more importantly, an analytical study. Specifically, besides the introduction of adaLN-Gaussian, our detailed analysis about adaLN-Zero also can not be ignored because it unveils the reason why adaLN-Zero outperforms adaLN, which enhances the community's understanding. Moreover, our analysis is the basis of the introduction of adaLN-Gaussian. It would be inappropriate to neglect the contribution of our analysis.

In contrast, our analysis is acknowledged by other reviewers. For example, it is acknowledged by Reviewer [p8jh] that we provide sufficient evidence to support the exploration inside the adaLN-zero. And Reviewer [KCHJ] also comments that our analysis provides some inspiring conclusions.

Q2: The performance gains are also limited.

A2: We may not agree with the reviewer's comment about the performance gains of our methods.

The improvement of our method is acknowledged by reviewer [ZpLc] who comments that our work is simple and effective, achieving remarkable performance.

If the reviewer thinks that our performance gains, e.g., 2.16 FID in 400K, are limited, how does the reviewer think about the improved performance of SiT[1] accepted by ECCV 2024? For example, SiT-XL/2 outperforms DiT-XL/2 by 2.3 FID in 400K in Table 1 in its paper (https://arxiv.org/pdf/2401.08740). (Our method and SiT both use the same structure of DiT.)

We sincerely appreciate the reviewer's time and effort, and we have no intention of offending the reviewer. However, to some extent, these comments seem to be unreasonable.

We hope that the reviewer could reevaluate the contributions of our work. As acknowledged by reviewer [ZpLc], our work is simple and effective (with only one line code replaced). Moreover, our method in Table 4 of our paper has shown great generalization to various DiT variants and DiT-based models.

[1] Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. ECCV2024

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

A3: We kindly remind the reviewer that in this paper we do not alter the learning algorithm and any structure of DiT model, but only initialize its condition mechanism differently. Therefore, theoretically, if there is no limit on training steps and excluding local optimum, the finally converged performance of adaLN-Zero would be similar to the performance of adaLN-Gaussian because the model capacity is the same. However, as reviewer [KCHJ] points out, different initialization methods mainly affect the convergence speed. Under the same training steps, our method allows DiT model to converge faster. For example, adaLN-Gaussian achieves an FID of 14.84 after 600K training steps, while adaLN-Zero requires 800K steps to reach a similar FID of 14.73. When further extending the training time, adaLN-Gaussian achieves an FID of 10.83 after 1.5M training steps while adaLN-Zero achieves requires around 2.4M steps to reach a similar FID of 10.67.

Q4: Conducting additional experiments on transformer-based models such as SiT and PixArt-alpha (text-to-image).

A4: We thank the reviewer's suggestion. SiT and PixArt-alpha are both very excellent work in the community. Due to the limited time, high GPU computing demand (64 V100 & 26 days), and large amount of internal data used in PixArt-alpha, we select to perform additional experiments on transformer-based SiT. We use the best-performing SiT-XL/2 training on ImageNet1K 256x256 for 50K. adaLN-Zero produces 71.90 for FID, while adaLN-Gaussian yields 67.15 FID, significantly outperforming adaLN-Zero and demonstrating the effectiveness of our method. Moreover, we also perform experiments using DiT-XL/2 training on other datasets including Tiny ImageNet, AFHQ, and CelebA-HQ for 50K. We report all the results below. These results further show the effectiveness and generalization of adaLN-Gaussian.

Q5: How do you summarize the differences between AdaLN-Zero and adaLN as the three points?

A5: Essentially, our summarized three points reflect the roles of the two additional steps of adaLN-Zero compared to adaLN: 1) introducing scaling element $\alpha$ 2) zero-initializing corresponding linear layers.

The first step changes the structure of DiT model. Through our closer observation of the formed structure, we find that the modified structure shares similarities with the SE module (The first point).

Our last two points are derived from the second step. First of all, we know that initialization has a basic role, namely, determining the initial value of the module weight (The second point). Particularly, for zero initialization, as previous work[2][3] suggested, it additionally can nullify certain output pathways at the beginning of training, thereby causing the signals to propagate through identity shortcuts. A direct way to reveal the influence of this shortcut behavior on model optimization is to analyze through the lens of gradient update (The third point).

Finally, our step-by-step and decoupling experiments and analysis show that these three points are all related to the improved performance of DiT model.

[2] Accurate, large minibatch sgd: Training imagenet in 1 hour. Arxiv2017

[3] Scalable Diffusion Models with Transformers. ICCV2023

Q6: In Fig. 4, regardless of the type of initialization used, the weight distribution converges to a similar state after a certain number of training iterations. Why the choice of initialization has such a significant impact on performance?

A6: We think there may be two reasons. Firstly, the results presented in Figure 5 of the DiT paper (https://arxiv.org/pdf/2212.09748), such as the significant performance improvement of adaLN-Zero over adaLN in FID, indicates that the condition mechanism is essential for achieving superior metric outcomes. Hence, leveraging a more appropriate initialization allows condition mechanism to learn better and thereby help DiT model to obtain superior outcomes.

Second, though the weight distribution converges to a similar state after a certain number of training iterations, e.g. 400K in Fig 4 where their distributions are very close to each other, the element-wise discrepancy may be still large. We average the absolute values ​​of the differences between each element in all weight matrices corresponding to adaLN-Zero and adaLN-Mix. The averaged element-level value is 0.05, which indicates that there is still a large element-level value shift in weight matrice. This may be the reason to the performance gap between adaLN-Zero and adaLN-Mix. We also calculate the average difference of the whole model weights between adaLN-Zero and adaLN-Mix. The result is 0.051, indicating there is also a large element-level value shift in the whole model.

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

A7: We thank the reviewer's suggestions. We follow the reviewer's suggestion to rotate Figure 2, which indeed significantly improves the visibility of the text. As for simplifying Figure 6 by merging the subgraphs, we have tried to do so, e.g., by merging 4 subgraphs. However, while enhancing information density, the resulting figure is not clear to see the distribution of each block. This may not match the purpose of showing Figure 6 where we aim to demonstrate that $W^{L}_{alpha}$ in each block of DiT exhibits a Gaussian-like distribution.

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

A1: Why are gama1, beta1, gama2, beta2 also zero? The answer is that in adaLN and adaLN-Zero, as shown in Fig 2 of our paper, the weights of the linear layer that produces gama1, beta1, gama2, and beta2 are initialized to zero at the beginning. Sorry for the confusion. We have revised the corresponding part of the paper to make it clearer.

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

A2: We thank the reviewer for the constructive suggestions. To make the inference more rigorous, we analyze two representative std settings including $std=0.05$ and $std=0.0005$, which correspond to a late moment (a large std) and an early moment (a small std), respectively.

We first illustrate their weight distributions of $W_{\alpha}$ in the conditioning mechanism and compare them with that of adaLN-Zero and adaLN-Gaussian ($std=0.001$). The results are shown in Fig. 21 of our revised paper in Appendix.9. It can be seen that a large std $std=0.05$ presents a relatively uncompact distribution and exhibits a significant discrepancy in distribution shape compared to the rest settings. This result indicates that a large std may be incompatible with other parameters, resulting in a slow speed of convergence and a poor performance. Moreover, we consider this a step further. Theoretically, if we further increase the std value, it would become close to the default initialization in adaLN-Step1 (xavier_uniform) while the performance of adaLN-Step1 is also bad.

For a small std $std=0.0005$, it can be seen that the distribution of $W_{\alpha}$ is quite similar to that of adaLN-Zero and adaLN-Gaussian ($std=0.001$). However, there still exists a slight discrepancy. To make this discrepancy clearer, we average the absolute values ​​of the differences between each element in $W_{\alpha}$ corresponding to $std=0.0005$ and adaLN-Zero, and $std=0.0005$ and adaLN-Gaussian. The element-wise averaged results are 0.0121 and 0.0124, respectively. By comparing the results (0.0121 < 0.0124), it is shown that small std leads to weights relatively closer to that of zero-initialization (adaLN-Zero). And, to some extent, the corresponding performance also proves it where $std = 0.0005$ produces 80.68 for FID, closer to adaLN-Zero (78.99) compared to adaLN-Gaussian (76.21).

We thank the reviewer again and have revised the corresponding part in our paper (as highlighted in blue), which makes our paper clearer and reasonable.

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

A3: Yes, if there is no limit on steps, the performance of adaLN-Zero could catch up with that of adaLN-Gaussian when training for more steps. For example, adaLN-Zero training for 800K produces 14.73 FID, catching up with adaLN-Gaussian training for 600K (14.84 FID).

And we agree with the reviewer's comment that different initialization methods mainly affect the convergence speed. Essentially, adaLN-Gaussian only changes the initialization strategy and does not alter any structure of DiT model, which means the model's capacity is the same. Therefore, theoretically, adaLN-Gaussian will not significantly influence the final converged performance but converge faster. For example, adaLN-Gaussian achieves an FID of 14.84 after 600K training steps, while adaLN-Zero requires 800K steps to reach a similar FID of 14.73. When further extending the training time, adaLN-Gaussian achieves an FID of 10.83 after 1.5M training steps while adaLN-Zero achieves requires around 2.4M steps to reach a similar FID of 10.67.

Q4: Have you analyzed the distribution of parameters of other parts of DiT? Do they also end up being Gaussian? How were these parameters initialized?

A4: Yes, we have analyzed the weight distribution of other parts of DiT including Attention and Mlp in DiT Block, PatchEmbed, LabelEmbedder, and TimestepEmbedder in Appendix A.7. In our experiments, we find that all of their weight distributions except PatchEmbed end up being Gaussian-like.

For the last question (How were these parameters initialized?) In DiT code, Attention and Mlp in DiT Block, and PatchEmbed are with Xavier uniform. LabelEmbedder and TimestepEmbedder are initialized with normal distribution.

Naturally, we can consider Gaussian initializations for these modules as well except PatchEmbed to accelerate training. For example, we could uniformly use Gaussian initialization for Attention and Mlp in DiT Block. We set the mean to 0 and use several choices for std such as 0.001, 0.01, 0.02, 0.03, and 0.04. We use DiT-XL-2 and train for 50K steps for simplicity. The results are shown below.

We see that the performance is inferior to the default initialization. Therefore, more precise hyperparameter tuning may be needed for these modules to further improve the performance, which we leave as future work.

We hope our response will clarify the reviewer's confusion and alleviate the concern. And we sincerely hope to obtain support from the reviewer.

Dear Reviewer KCHJ,

We hope that this letter finds you well.

As the end of this discussion is going to be close, we would greatly appreciate it if you could let us know whether our responses and revisions align with your expectations and have addressed your concerns. We may need to highlight our new experiment where adaLN-Gaussian achieves 2.27 FID in 5400K training steps, less than 7000K in adaLN-Zero which is the longest training steps in DiT paper. This result demonstrates that our method adaLN-Gaussian is also superior over adaLN-Zero in the scaled training for final convergence.

If there are any remaining issues where further clarification is needed, please don’t hesitate to let us know. We are more than willing to provide additional explanations.

Thank you once again for your time and effort. We look forward to hearing your thoughts and would greatly appreciate it if you could consider kindly raising your rating.

Sincerely,

Authors

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

A1: We thank the reviewer's comment and we are very willing to explain how we find such similarity. We primarily find the connections between them from three aspects: 1) overall structure; 2) module function; and 3): detailed mathematical formula.

First, from the view of the overall structure, the adaLN-Zero, and SE module both serve as a side pathway compared to the main path.

Then, from the view of the module function, scaling element $alpha$ in adaLN-Zero plays a similar role as the SE module, both of which aim to perform a channel-wise modulation operation. Hence, to achieve it, they may yield outputs of the same structure, e.g., vector, thereby sharing some similarity in output formulation.

Finally, from the view of the mathematical formula, though we may not directly find the similarity due to the existence of the SiLU function, a more detailed expandation allows us to find a close connection in the mathematical formula between the adaLN-Zero and SE module. With a slight adjustment according to experience, the similarity appears.

We thank the reviewer's question and have added the response as well as a figure of the structure of the SE module to our paper (as highlighted in blue), which makes the start point clearer.

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

A2: The reviewer may want to say Table 2 instead of Figure 2. Intuitively, since the weights of the conditional mechanisms we counted are Gaussian-like distributions, there should exist an optimal std hyperparameter when initializing these weights with Gaussian, and naturally, the values ​​on both sides of this hyperparameter are relatively unsuitable.

We follow the reviewer's advice and analyze this U-shaped trending by leveraging two representative settings, i.e., $std=0.0005$ and $std=0.05$, which are the two ends of U-shaped results.

We first illustrate their weight distributions of $W_{\alpha}$ in the conditioning mechanism and compare them with that of adaLN-Zero and adaLN-Gaussian ($std=0.001$). The results are shown in Fig. 21 of our revised paper in Appendix.9. We find that a large std $std=0.05$ presents a relatively uncompact distribution and exhibits a significant discrepancy in distribution shape compared to the rest settings. This result indicates that a large std may increase the difficulty of optimization, resulting in a slow speed of convergence. Moreover, we consider this a step further. Theoretically, if we further increase the std value, it would become close to the default initialization in adaLN-Step1 (xavier_uniform) while the performance of adaLN-Step1 is also bad.

For a small std $std=0.0005$, we find that the distribution of $W_{\alpha}$ is more compact and is quite similar to that of adaLN-Zero and adaLN-Gaussian ($std=0.001$). Though similar, there still exists a slight discrepancy. To make this discrepancy clearer, we average the absolute values ​​of the differences between each element in $W_{\alpha}$ corresponding to $std=0.0005$ and adaLN-Gaussian, and adaLN-Zero and adaLN-Gaussian. The results are 0.0124 and 0.0122, respectively, which indicates that there is still an element-level value shift in weight matrice. The reason that the performance of $std=0.0005$ is relatively bad in Table 2 may possibly be because $W_{\alpha}$ in $std=0.0005$ is farther away from that in adaLN-Gaussian compared to adaLN-Zero since 0.0124>0.0122. Note that though the difference is small, considering the large number of elements (about 74.3M), the influence may not be ignored.

We hope our response will clarify the reviewer's confusion. And we sincerely hope to obtain support from the reviewer.

Q1: Could the authors provide a more detailed explanation of why Gaussian distributions are used to initialize weights?

A1: Yes, we feel very sorry for the confusion and are very willing to explain the reason for using Gaussian distributions. The logic is that:

First of all, by comparing the differences between adaLN-Zero and adaLN, our analysis studies three elements that collectively drive the performance enhancement: 1) an SE-like structure, 2) zero-initialized value, and 3) a “gradual” update order. Though previous work thinks the manner of shortcuts plays a major role, our analysis finds that it is a good zero-initialized location itself that plays a significant role, which indicates that it is important to find a suitable initialization.

Further, we find that, from the perspective of weight distribution, though weights in the conditioning mechanism are zero-initialized, after a certain number of training steps, they transition from zero distributions to Gaussian-like distributions. Hence, inspired by this, our insight is that we can expedite this distribution shift by directly initializing the weights via a suitable Gaussian distribution to potentially accelerate training.

We have revised the corresponding part of the paper (as highlighted in blue) to make it clearer.

Q2: More extensive and general experiments

A2: We thank the reviewer for the kind and constructive suggestions. To further show the effectiveness of adaLN-Gaussian, we add more experiments on other datasets including Tinyimagenet, AFHQ, and CelebA-HQ using the best-performing DiT-XL/2 with 50K training steps while keeping all training settings. Moreover, we also use another DiT-based model SiT-XL/2 [1] training on ImageNet1K 256x256 for 50K to further show the effectiveness and generalization of adaLN-Gaussian. We report all the results below. These results show that adaLN-Gaussian consistently outperforms adaLN-Zero, demonstrating the effectiveness of our method.

We hope our response will clarify the reviewer's confusion and alleviate the concern. And we sincerely hope to obtain support from the reviewer.

[1] Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. ECCV2024

Dear Reviewer ZpLc:

Thank you for your comment; we are delighted to engage in this discussion! We apologize for the absence of a theoretical analysis. Our experiments are still ongoing and require additional time. We will share the results and provide further evidence promptly within the discussion period.

Further, to demonstrate that adaLN-Zero indeed forms a Gaussian-like distribution, we employ KL-Divergence to measure the distance between its distribution and a true Gaussian. Specifically, we use the weights of adaLN-Zero at 50K steps to compute its mean and standard deviation. These parameters are then used to initialize a Gaussian distribution, from which we sample the same number of points as adaLN-Zero. Finally, the KL distance between the two sets of sampled points is calculated using the nearest neighbor nonparametric estimation method, as detailed below.

$D_{\text{KL}}(P \| Q) \approx \frac{1}{n} \sum_{i=1}^{n} \log \frac{\rho_i}{\nu_i} + \log \frac{m}{n-1}$

$m$ and $n$ are the number of sample points. $\rho_i$ represents the nearest neighbor distance of point $x_{i}$ in $P$. $\nu_i$ is similar.

The calculated KL-Div is 0.065. Typically, the closer the two distributions are, the smaller the KL-Div is. If the two distributions are the same, the KL-Div is 0. Therefore, the calculated result demonstrates that adaLN-Zero indeed formulates Gaussian-like distribution.

For the confusion of the necessity of using Gaussian distribution for weight initialization, we feel sorry for this. We think there may be two reasons for this necessity. Firstly, the results in Figure 5 of the DiT paper (https://arxiv.org/pdf/2212.09748), e.g., the significant improvement of adaLN-Zero over adaLN in FID, indicate that the condition mechanism is essential for achieving superior metric outcomes. Therefore, selecting a suitable initialization for the condition mechanism could be necessary. Simultaneously, this is also an easily achieved and low-cost method to help with training. 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. Therefore, using a suitable Gaussian initialization could better expedite this distribution shift.

Further, we follow your advice to extend our training steps. However, we may need to explain that in fact 50K iterations (batch size=256) is relatively suitable for Tiny Imagenet, AFHQ, and CelebA-HQ because the number of their images is small, only 100K, 14K, and 30K, respectively, compared to ImageNet1K. We feel sorry that we fail to update these information timely and sorry for the confusion this training step setting brings to you.

But we are willing to further extend the training steps. We extend an additional 50K for Tiny Imagenet and 150K for ImageNet1K (SiT-XL/2). The results are presented below where our method adaLN-Gaussian still outperforms adaLN-Zero when the training step is longer.

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. For example, an additional 2K training steps for CelebA-HQ increases the FID from 7.54 to 11.20.

In addition, we also extend the training steps of DiT to further show the effectiveness of our method under longer training time. The results are shown below. We find that adaLN-Gaussian produces 10.09 for FID in 2300K training steps, outperforming adaLN-Zero in 2352K training steps (10.67 FID in its paper). Moreover, using the same cfg=1.5 as adaLN-Zero, adaLN-Gaussian achieves 2.27 FID in 5400K training steps, less than 7000K in adaLN-Zero which is the longest training steps in DiT paper. These results further demonstrate the superiority of adaLN-Gaussian over adaLN-Zero.

We hope our response above could address your concerns and help you reassess our work. We would greatly appreciate it if you could consider kindly raising your rating! Thank you once again for your patience.