On improved Conditioning Mechanisms and Pre-training Strategies for Diffusion Models

We appreciate the reviewer's feedback and are glad they found our paper a good advancement to the state of the art. In the following, we respond to the points raised:

• Lack of small-scale experiments. The objective of our work is to increase understanding of the design choices in state-of-the-art diffusion models, and propose improvements upon those. However, SOTA diffusion models require sufficiently large scale pre-training to be highly performant. Under this premise, we believe that testing on smaller scales is not necessarily conducive to our objective, since convergence behavior can differ significantly at smaller scales, e.g., due to overfitting of the data-hungry transformers. Also, in the literature, fundamental diffusion studies have been already tested at small-scale, e.g., Nichol et al (2021) and Karras et al (2022 and 2023). $Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." (2021); Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." (2022); Karras, Tero, et al. "Analyzing and improving the training dynamics of diffusion models." (2023)$

• Instability issues with certain models not detailed (e.g., DiT with cross-attention). As mentioned in our paper at L84-90 and L239-241, these instabilities that we report have been already studied and solved in the SD3 paper. They come from the keys and queries growing larger in amplitude, causing instabilities later on during training. They were solved by using RMSnorm on the keys and queries before every attention operation.

• Ablations are not comprehensive. We are not sure about which ablations the reviewer would expect us to report. However, for an overall ablation we have added Table R1 in the attached pdf. The added table should better isolate the contribution of each of our improvements. If this would not be enough for the reviewer, we invite them to specify which components they would like to see ablated.

• Limitation of our work. We refer the reviewer to the limitations section in Appendix C of our paper. In addition, we expand on the limitation section, focusing on failure cases of our method:

  • We found our method to provide improvements consistently across the settings tested. However, there are cases where these improvements can be less pronounced. For example, noisy replicates for the text embeddings can become less pertinent if the model is trained exclusively with long prompts.
  • Also, while low resolution pretraining with local crops on ImageNet resulted in better FID at 512 resolution (see Table 5c), it might not be necessary if pretraining on much larger datasets (~300M samples or more which we did not experiment in this work).
  • Similarly, flip conditioning is only pertinent if the training dataset contains position sensitive information (left vs. right in captions, or rendered text in images), otherwise this condition will not provide any useful signal.

• Extension to audio or text generation. While we find possible extensions to other data modalities intriguing as future works, we believe they are out of the scope of the current work.

We appreciate the thorough feedback. In the following we address it:

• Design choices

  • Noisy replicate padding has minor impact. In addition to slightly improving the FID and CLIPScore (0.54 and 0.58, respectively), the added noisy replicate can serve as a regularizer with no added cost during training. The intuition is that it disincentivizes the model from ignoring the later copies in the attention, as it would happen with zero-padding (see Figure R3 in the attached pdf). We added Figure R2 to depict this intuition. Moreover, the local variations caused by the added noise acts as a sort of data augmentation in the text encoder’s embedding space, intuitively robustifying the model to small variations in the semantics.
  • Cosine scheduler for conditioning disentaglement. We chose the power cosine schedule because it provides an easy-to-tune function, in the form of one hyperparameter $\\alpha$ (decay rate of the conditioning), and generates a smooth transition, which we intuited to be more inductive for learning than using a step function. To prove our choice and clarify the reviewer's doubt, we have run a comparison against step and linear schedulers, see Table R3. We observe superior performance provided by the cosine scheduler as per FID, LPIPS, and LPIPS at high resolution (LPIPS/HR). We recall that LPIPS is measured across 10 generations obtained with the same random seed, prompt, and initial noise, but different size conditioning (for HR, we exclude sizes smaller than the target resolution); then we report the average over 10K prompts.

• Increased number of hyper-parameters. While we introduced some new hyperparameters, the model shows improved performance for a significant range of their values. For example, for all $\\alpha$ tested in the control scheduler, we obtained lower LPIPS (see Table 2a) and practically removed semantic drift across different size conditionings (see Figure 2 and Table 2b). The control scale, $\\gamma\_c$, also shows consistent performance across datasets for the same values (range $0, 2$ , see Figure 7 right). Moreover, when possible, e.g., for the guidance scale, we included theoretical analysis to guide the transfer of these hyperparameters across resolutions (see Eq.4 and its derivation in Appendix F), thereby restricting the search/tuning space.

• Trivial results in Table 2b. We agree that it is expected to obtain the best FID (distribution overlapping) when conditioning on the original image size. However, this result serves as a reference point to understand how much the choice of the size conditioning affects the distribution of generated images. Indeed, the purpose of Table 2b is to highlight the larger FID increase (i.e., distribution shifting) in the baseline vs ours. The results obtained with our control conditioning suggest that better FID numbers can be obtained without the need of careful tuning of the size conditions for every prompt. We will clarify the reading of Table 2b when updating our paper.

• Class embedding position in control conditioning. Moving class embeddings into control conditions results in the vanilla DiT architecture, while moving the control conditioning to the cross-attention in mmDiT results in an architecture similar to Hatamizadeh et al (2023), which we found to underperform with respect to the baseline. $Hatamizadeh, Ali, et al. "Diffit: Diffusion vision transformers for image generation." (2023)$

• Comparing curves instead of specific steps for Table 5. We agree with the reviewer that comparing curves provides richer observations. For this reason, we plotted the experiments presented in Table 5 evaluating runs at intermediate checkpoints. We show the plot in Figure R1. The trend of the curves depicted in all the three panels, (a-c), validate the improved performance of our contributions, which also showcase faster convergence rates. Unfortunately, due to limited time and compute resources, we were not able to run the comparison in panel (b) and (c) until 200K iterations, but we believe these plots could be enough to solve the reviewer's concerns regarding Table 5. The runs are still ongoing, and we will update the plots for the final version of the manuscript.

• Clarify LPIPS in Table 2a. We describe the LPIPS computation at L215 of the main paper, plus we clarify it here above: “We recall that LPIPS is measured across 10 generations obtained with the same random seed, prompt, and initial noise, but different size conditioning (for HR, we exclude sizes smaller than the target resolution); then we report the average over 10K prompts.” We will add this extra clarification also in the main paper.

We thank the reviewer for their valuable feedback and for raising their score.
While we acknowledge that many of the contributions we propose are experimentally based, we would highlight that we made an effort to provide theoretical backing for our contributions when possible, see answer to reviewer rR4N.

Regarding the significance of the gains with each modification, there are two points that we would like to raise:

• Differently from works like EDM (Karras et al 2022 and 2023), we are not optimizing for one single metric (FID) but trying to improve different facets of our model (FID, CLIP, LPIPS).
• We compare our model with very strong baselines, e.g., SD3, which achieve very high scores to start with. Hence, the absolute gap with these methods is small, but the relative improvement is considerable, see Table 1 and R1 below. Notably, our improvements do not carry computational overhead.

We thank the reviewer for the valuable feedback. In the following we address it:

• Unorganized Experiment Result Presentation. We appreciate the reviewer's suggestion to adopt a table as in Karras et al (2023) to enhance the presentation of our contributions. We have added Table R1 in the attached pdf, which reports improvements in $256^2$ pre-training (panel (a)) and resolution transfer setting (panel (b)). Notably, rows 2-5 in panel (a) address concerns regarding control conditioning contribution, showing FID improvements from 1.81 to 1.59 on ImageNet and 7.54 to 7.32 on CC12M, along with a slight CLIPScore boost. Given the added value of this table, we consider integrating it into the main paper. Finally, note that we do not report improvements in dataset shift because a similar table is already present in the main paper, see Table 4.
• Method presentation / complex design space. We do agree with the reviewer that the design space might be complex, but we believe that this observation applies to diffusion models as they have many parameters that are left unexplored most of the time. Though, we acknowledge that a table like the one in Karras et al (2022) would clarify many doubts on our configuration and contributions. Hence, we added Table R2 to the attached pdf and to the appendix of the paper. The table compares our model vs. SD3 and SDXL, taking into account the configuration of: Sampling, Network and conditioning, High resolution adaptation, and Optimization.

We are very glad the reviewer was satisfied by our response, and they increased their score.

We appreciate the reviewer’s opinion. We see our paper as an in-depth study that: (i) analyzes different model design choices, (ii) provides additional understanding, and (iii) proposes improvements on pre-existing approaches. However, we politely disagree with the premise of the comment that our observations might not generalize due to the lack of theoretical background. All our contributions are validated through an extensive set of experiments at large scale showing good generalization. In particular, we provide results on three datasets of different scales and distributions and for different resolutions. We may remark that previous foundational diffusion models, e.g., LDM and DiT, were trained and tested primarily on Imagenet at 256 and 512 resolutions. We would also note that empirical validation, even when not grounded by theory, has led to generalization several times over the years of DL algorithms development, e.g., MAE (He et al 2022), DINO (Caron et al 2021), spatial pyramid pooling (He et al 2015), ViT (Touvron et al 2021, Dehghani et al 2023) and ResNet recipes (He, et al 2019, Kolesnikov et al 2020, Bello, et al 2021). Finally, we would like to highlight that some of our proposed improvements are grounded by theory, e.g., decoupled control conditioning leverages the classifier free guidance theory, while noise schedule scaling is backed by Eq.3 (with proof of derivation in Appendix F).

We hope this answer addresses the reviewer's point for lowering confidence. If not, we are happy to discuss any other point that could assist to that end.

Touvron, Hugo, et al. "Training data-efficient image transformers & distillation through attention." International conference on machine learning 2021; He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition 2022; He, Kaiming, et al. "Spatial pyramid pooling in deep convolutional networks for visual recognition." IEEE transactions on pattern analysis and machine intelligence. 2015; Kolesnikov, Alexander, et al. "Big transfer (bit): General visual representation learning." Computer Vision–ECCV 2020; Dehghani, Mostafa, et al. "Scaling vision transformers to 22 billion parameters." International Conference on Machine Learning. 2023; Caron, Mathilde, et al. "Emerging properties in self-supervised vision transformers." Proceedings of the IEEE/CVF international conference on computer vision. 2021\. Bello, Irwan, et al. "Revisiting resnets: Improved training and scaling strategies." Advances in Neural Information Processing Systems 34\. 2021; He, Tong, et al. "Bag of tricks for image classification with convolutional neural networks." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019\.