Paper: PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis (ICLR-2024)

Reviewer comment (Evaluation):
The GPU days reported in A.5 contradict the main table; Summing the last column and using their 2.2 scale factor for converting from v100 to a100 days yields 753 A100 days vs. the 675 reported in the main table. It's a relatively small difference (11%) but concerning that the reported numbers don't align.

Author response:
We sincerely apologize for the statistical errors in our paper due to the large amount of data information when we wrote this paper. We have fixed all these errors, including the misinterpretation of data referenced from GigaGAN [1] in the revised version. The updates involve GPU days and the training data required by other methods like Imagen [2]. The updated data shows that Pixart uses 5% data volume compared to SDv1.5 and costs less than 2% training time compared to RAPHAEL. Compared to RAPHAEL, our training costs are only 1%, saving approximately \\$3,000,000 (PixArt-α’s  \\$28,400 v.s. RAPHAEL’s \\$3,080,000). All the conclusion **consistently demonstrates the efficiency of PixArt**.

Moreover, Nvidia's tech report [8] illustrates the acceleration ratio between A100 and V100 spec in the first Table below, where the V100->A100 speedup for the Transformer BERT is **5x** v.s. U-Net **2.2x**. We have conservatively estimated the lowest speedup for PixArt in the original paper (we adopted the U-Net speedup ratio). If we consider BERT speedup since PixArt also uses pure Transformer architecture, same as BERT, the training time of PixArt would significantly reduce to **1656 V100 GPU days / 5** = **332 A100 GPU days**.

|         | ResNet-50 v1.5 | U-Net Medical | BERT    |
| ------- |:--------------:|:-------------:|:-------:|
| A100    | 4.1 hours      | 3 minutes     | 5 Days  |
| V100    | 9.7 hours      | 7 minutes     | 24 Days |
| Speedup | 2.4x           | **2.2x**      | **5x**  |

| Method   | #Images |     GPU days      |
| -------- |:-------:|:-----------------:|
| DALL·E   |  250M   |         -         |
| GLIDE    |  250M   |         -         |
| LDM      |  400M   |         -         |
| DALL·E 2 |  650M   |    41,667 A100    |
| SDv1.5   |  5000M  |    6,250 A100     |
| GigaGAN  |  2700M  |    4,783 A100     |
| Imagen   |  860M   |    7,132 A100     |
| RAPHAEL  | 5000M+  |    60,000 A100    |
| PIXART-α |   25M   |  753 A100 (2.2x)  |
| PIXART-α |   25M   | **332 A100** (5x) |

Reviewer comment (Experiments):
There are 3 primary improvements (dataset curation, multi-stage training, and unet->transformers), however there are no ablations quantifying how much of an effect each one has on model performance. While the results and final model are good, a lack of ablations limits the research contribution of the paper.

Author response:
Thanks for this valuable suggestion. Systematically ablation studies of Text-to-Image (T2I) models from model architecture, data curation, and training are computationally expensive and would demand extensive resources, which exceed our current capabilities. However, we still try our best to give some insights and analysis of the above aspects:  


1. Data Design: 
* Regarding the selection and re-annotation of LLaVA data, we discussed this in Section 2.4 via the total number of concepts. Additionally, we found that over 90% of the nouns in the LAION concept statistics appear less than 10 times throughout the dataset. Consequently, we believe that the large number of infrequently occurring nouns may be the main reason for prolonging the training time. This is where we attribute the primary cause for rapid training; merely replacing UNet with Transformer wouldn't be enough to accelerate training speed by an order of magnitude.
* We are continuing to explore the lower limit for the data required in stage 3. 
In stage 3 training, we found that increasing the data volume only marginally improves the model's capabilities, whereas decreasing the data volume has a progressively larger impact. 
We increased the data scale from 14M to 30M but did not observe significant improvements. We also decreased data to 1.4M, 700K, and 140K. The model's performance was relatively stable at 1.4M and 700K, but there was noticeable quality degradation at 140K. Due to time constraints, our current conclusion is that training with 1.4M data points in stage 3 achieves results comparable to 14M. However, volumes below 700K may lead to overfitting, as shown in the below table. The properties of the data should match the VN/DN analysis mentioned in Section 2.4. 

| Data Volumn | Iterations | Image size |  FID  | CLIP score |
|:----------- |:----------:|:----------:|:-----:|:----------:|
| 30M         |    120K    |    512     | 8.62  |   0.276    |
| 14M         |    100K    |    512     | 8.61  |   0.275    |
| 14M         |    15K     |    256     | 16.61 |   0.268    |
| 1.4M        |    15K     |    256     | 16.96 |   0.266    |
| 700K        |    15K     |    256     | 17.14 |   0.265    |
| 140K        |    15K     |    256     | 21.58 |   0.255    |

2. Multi-stage Training: 
* We tried to skip the first stage and directly train the text-to-image task on SAM. However, due to the significant difference in scale between image and text in cross-attention inputs, and the lack of a corresponding relationship, this usually led to training crashes often represented by NAN. Therefore, we believe that pixel alignment in the first stage is key. 
* The decoupling of stages 2 and 3 is partial because the pseudo-label quality of SAM cannot be guaranteed by current open-source VLMs, leading to illusions, bias, and other issues. Therefore, we only use this data for high-concept-density text-image alignment learning, corresponding to stage 2 pre-training. Fine-tuning with high-aesthetics data can improve the model's aesthetic quality. We have conducted experiments to skip stage 2, and directly train stage 3 based on ImageNet pretrain. Under the same training iterations, we observed that both the visualization results and the text prompt instruction following ability were noticeably inferior compared to using stage 2's SAM pretraining.
3. UNet->Transformer:
* In the original DiT paper, they already thoroughly compared Transformer and UNet architectures on class-condition image generation, validating transformers are easier to scale up and have higher performance upper bound. This is the core reason we chose DiT as the base diffusion model and extended it to the text-to-image model.
* Moreover, we chose Transformer architecture because it also has several future potentials, including MaskDiT [3], and Window/Sparse/Linear Attention [4~6]. We leave these improvements in the future work.

Reviewer comment (Reproducibility):
The authors do not discuss the training of the VAE; since this is an important part of "learning the pixel distribution of natural images", the training time and data quantity of the VAE should be included in the paper.

Author response:
Our attempt to train a VAE resulted in an approximate training duration of 25 hours, utilizing 64 V100 GPUs on the OpenImage dataset. Training VAE seems does not consume too much time. We treat the pre-trained VAE as a ready-made component of a model zoo, the same as the pre-trained CLIP/T5-XXL text model, and our total training process did not include the training of VAE and we are also not sure whether other methods e.g. SDXL considered the VAE training time. To ensure a fair comparison, we have temporarily excluded the VAE training time and data quantity. We clarify this point in Appendix A.5 in the revised version of our paper.

Reviewer comment (Evaluation):
My intuition says that training cost and CO2 emissions should be directly proportional to GPU hours, however these proportions are inconsistent. For example, Fig 2b shows a 14x increase in cost of Imagen over PixArt, while the ratio of reported GPU hours is 10. For GigaGAN, Fig2b shows a 9.4x increase, whereas Table 2 shows a 7x increase in GPU hours.

Author response:
Sorry to make you confuse about CO2 emissions and training costs. The CO2 consumption here is directly proportional to the **A100 GPU** days because we referred to [7], which only contains A100 CO2 emission Statistics.
About the training cost, PixArt is calculated based on the actual training time and price of the **V100 GPU** (refer to https://www.leadergpu.com/). Due to the price differential between the V100 and A100 specs., the ratios of CO2 emissions to training time and cost to training time differ slightly. We will clarify this in the revised version of our manuscript.

Reviewer comment (Experiments):
How expensive was the dataset collection? I expect running a big image-text model like LLaVA on millions of images required a lot of time, GPU ressources, and memory. Should this non-trivial step be in some way included to the analysis of costs and co2 emissions?

Author response:
We use LLAVA-7B to generate captions. LLaVA's annotation time on the SAM dataset is approximately 24 hours with 64 V100; it is negligible compared to training time. Following your suggestion, we clarify this point in Appendix A.5 in the revised version of our paper.

Reviewer comment (Evaluation):
Although using a model like LLava clearly improves VN/DN, it probably introduces new biases to the distribution of text prompts? Does this affect the scope of text prompts for which PixArt-$\alpha$ works well or poorly?

Author response:
LLaVA may introduce certain illusions and biases, but it does not harm the learning of concepts and text-image alignment. Therefore, we only use LLaVA-caption for text-image alignment pre-training. In stage 3, we finetune our internal data with GT text prompts (users given) to quickly adapt to user preference. Our main contribution is a proof of concept that using VLM models like LLaVA for captioning is reasonable and practical. As VLM rapidly evolves, some new VLMs, e.g. GPT-4V, can produce better auto-labeling captions and significantly contribute to T2I pre-training.

Reviewer comment (Reproducibility):
What are the plans of the authors regarding open-sourcing? Is it expected to be released 1) complete training code; 2) collected dataset; 3) "internal" dataset?

Author response:
We have released training/inference code, checkpoints, and demos, which have had a certain impact on the AIGC community.

Reviewer comment (Experiments):
Given its smaller dataset, there are concerns about the generalizability and compositional capabilities of the model for a broader range of concepts.

Author response:
**W1: Concerns About OoD Generalization with Smaller Dataset.**

We believe there are no generalizability issues in PixArt. As detailed in Section 2.4, although our total training samples are fewer than LAION-5B, each sample contains much more information, leading to better overall information volume and density. Specifically, as shown in **Table 1**, the total number of valid distinct nouns (VN) in SAM and Internal data is 152K + 23K = **175K**, which is comparable with LAION's **210K**. Moreover, we empirically found that for some OoD prompts, such as surreal image generation, like the prompt in Figure 1 where *"a cactus with a happy face"*, SDXL fails to produce satisfactory images while PixArt's generated image is good. 

In summary, we can conclude that our dataset scale is big enough, and PixArt does not suffer from generalizability issues.

Reviewer comment (Experiments):
What led to the decision of utilizing ImageNet data as the pretraining source instead of a pretrained VQGAN, and how did this choice impact the model's performance and capabilities?

Author response:
**Q1: Using ImageNet-Pretrained DiT Over Pretrained VQGAN**

We apologize for the confusion regarding the use of ImageNet. We explain our choice in two aspects:

1. Image encoder/decoder choice: We chose VAE over VQGAN [1], following current open-source popular methods like SDXL [2], SD [3], and DiT [4]. The choice between VAE and VQGAN is not the primary focus of our paper, and using VQGAN as an alternative is also feasible.
2. Genrative model choice: We chose the Diffusion Transformer instead of the Auto-regressive Transformer in VQGAN, mainly because Denoising Diffusion Probabilistic Models (DDPMs) are the most mainstream and mature data modeling method for image generation, currently most of the advanced image generative models are based on DDPMs. Meanwhile, auto-regressive Transformer for image generation is also promising since some recent papers, e.g. CM3Leon (meta), also adopted this solution.

Reviewer comment (Experiments):
The model's performance raises the important question of determining the minimum number of images required for training while still ensuring generalizability and maintaining the compositional properties of Text-to-Image models. Can the authors shed some light on what is the minimum set of images required to train such a model from scratch? What should be the properties of such a dataset?

Author response:
**Q3: Insights about Minimum Scale and Properties of the Dataset**

That's an excellent question. We are continuing to explore the lower limit for the data required in stage 3. Below is our observation. 

In stage 3 training, we found that increasing the data volume only marginally improves the model's capabilities, whereas decreasing the data volume has a progressively larger impact.

Specifically, we increased the data scale from 14M to 30M but did not observe significant improvements. We also decreased data to 1.4M, 700K, and 140K. The model's performance was relatively stable at 1.4M and 700K, but there was noticeable quality degradation at 140K. Due to time constraints, our current conclusion is that training with 1.4M data points in stage 3 achieves results comparable to 14M. However, volumes below 700K may lead to overfitting, as shown in the below table. The properties of the data should match the VN/DN analysis mentioned in Section 2.4. 

| Data Volumn | Iterations | Image size | FID   | CLIP score |
|:-----------:|:----------:|:----------:|:-----:|:----------:|
| 30M         | 120K       | 512        | 8.62  | 0.276      |
| 14M         | 100K       | 512        | 8.61  | 0.275      |
| 14M         | 15K        | 256        | 16.61 | 0.268      |
| 1.4M        | 15K        | 256        | 16.96 | 0.266      |
| 700K        | 15K        | 256        | 17.14 | 0.265      |
| 140K        | 15K        | 256        | 21.58 | 0.255      |

Reviewer comment (Novelty):
The paper mostly is a combination of multiple ideas that exist in the literature, so "novelty" in the traditional sense is somewhat limited.

Author response:
We understand the 'traditional novelty' mentioned in the review. We will discuss more about the novelty here. 
1. **Data**. Our paper's primary contribution to the community lies in the auto-labeling approach we have developed, a topic that has yet to see much discussion thus far. Coupled with the analysis of concept density, as detailed in Section 2.4, we provide a fresh perspective on assessing the suitability of a dataset for diffusion training.
2. **Training process**. Decoupling the entire training process into multiple stages is another innovative finding in our work. This approach significantly accelerates the training process, unlike previous methods [1,2], which decoupled the training process into different resolutions. While these earlier methods enhance the model's forward pass speed, we contribute to faster learning.
3. **Model structure**. we design an adaptation technique to validate the feasibility of transfer from class-condition to Text-to-Image (T2I) field and make it outperform existing methods.

In summary, our paper presents novel findings, insights, and designs concerning data, the training process, and model structure. We believe our approach could inspire the community in future research endeavors.

Reviewer comment (Experiments):
The SAM dataset blurs human faces in their training, won't this cause problem in generation cases where generating a face (not closed up) is needed?

Author response:
The blurry faces in the SAM dataset indeed pose challenges in pretraining models to generate clear faces. However, we primarily utilize SAM for Text-Image Alignment learning, where high image quality is not crucial. In later high-quality fine-tuning phases, we have ample data with clear facial images, allowing for rapid rectification of such issues.

Reviewer comment (Reproducibility):
How does the model generate images with more extreme aspect ratios?

Author response:
Drawing inspiration from SDXL [1], we pre-define various aspect ratios ranging from 1:4 to 4:1 and train batches of images with the same aspect ratio. We also add aspect ratios (AR) as one additional condition to the model, allowing the model to be aware of different ARs. During inference, users can specify a specific AR and a corresponding noise map with the same AR will be initiated and input into the model. Technically, PixArt can easily generate images with arbitrary aspect ratios. However, most of the training data's aspect ratio ranges from 1:4 to 4:1, so we cannot guarantee the quality of the generated images with an extreme aspect ratio.

Reviewer comment (Experiments):
DALL-E 3 technical report mentions the pitfall of "overfitting" to automated captions, is this the case in PixArt model? If not, how is this mitigated?

Author response:
DALL-E 3 and PixArt follow different training processes. DALL-E 3 uses a mixture of ground truth (GT) and pseudo captions during training. In contrast, PixArt initially uses  LLaVA-generated pseudo captions for pretraining (stage 2), then finetuning with GT text-image pairs (stage 3). Considering the captioning quality of LLaVA and similar VLM models are still under enhancement, the automated captions are mainly used for image-text alignment learning. Moreover, in stage 3, the model also learns text-image alignment; the only difference is that the captions of stage 3 are provided by real users (Section 2.4). After stage 3 finetuning, the model effectively aligns with users' preferences/habits, enabling PixArt to avoid overfitting from pseudo captions.

Reviewer comment (Experiments):
Since the dataset size is smaller, does it have trouble producing named entities, such as celebrities?

Author response:
PixArt has a certain ability to generate images of well-known people such as Donald Trump and Elon Musk, but it can't guarantee accuracy for individuals not present in the dataset. We test and observe that even when the Stable Diffusion is trained on a 5B volume dataset, its ability to generate celebrity images remains limited. To ease this problem, we can add related datasets for training. e.g. CelebA [7], which encompasses an array of celebrity faces.

Reviewer comment (Experiments):
How critical is training DiT on ImageNet needed? While being able to start with an existing model is good it also limits the possibilities to explore different architectures.

Author response:
Pretraining on ImageNet does not limit model exploration. Instead, it enhances training stability. 
1. Directly training a random-initialized DiT on the SAM dataset led to frequent NaN issues because the models need to learn image denoising and text-image alignment simultaneously. 
2. We have validated the feasibility and efficiency of extending the class-conditioned model to the text-conditioned model; this part is orthogonal to the network architecture.

So it is conceptually simple to explore different architectures to pretrain on ImageNet and finetune on the SAM dataset without restrictions.

Reviewer comment (Evaluation):
The CLIP score of the CLIP-FID curve of Pixart seems worse than SD 1.5. Is there any reason for that?

Author response:
PixArt and SD1.5 exhibit comparable CLIP scores. 
1. Considering both FID and CLIP scores, the comparison between PixArt and SD1.5 is evenly matched, with PixArt demonstrating a higher CLIP score at the lowest FID point, as shown in Fig. 20.
2. The datasets used by SD and PixArt differ: the SD utilizes the LAION-5B dataset and filters the images with CLIP scores < 0.28 [8]. Training with the data already having high CLIP scores might present an advantage. 
3. We tested the ground truth image-text pair CLIP score on MSCOCO-3W, which is only **0.257**, and we found that both SD1.5 and PixArt exceed this score. Therefore, we believe this metric can not effectively reflect the true performance of the advanced T2I models.