Lumina-T2X: Scalable Flow-based Large Diffusion Transformer for Flexible Resolution Generation

We thank reviewer 5 (LZg1) for acknowledging the contribution of our paper and providing thoughtful comments. Please see our response to the feedback below.

Q1: Novelty.

Please refer to the global reply for more details. We acknowledge that our framework incorporates existing techniques. Our objective is to demonstrate their effectiveness in enhancing training stability and efficiency when developing scalable diffusion transformers across various modalities. By leveraging these design choices, we scale the diffusion transformer to 7B parameters and showcase new capabilities, such as RoPE-based resolution extrapolation, high-resolution editing, compositional generation, and style-consistent generation.

Q2: Evaluation on more modalities.

Thank you for your suggestion. As clarified in our contribution above, our experiments on other modalities are extensions of the Lumina architecture, so we provided preliminary exploration in the appendix. Additionally, we included quantitative results for audio in Table 5 of the appendix, and further quantitative results for text-to-image generation are provided in our response to Reviewer 4 (ozQQ).

Q2: Scalability.

Initially, we found that naive scaling the Diffusion Transformer (such as DiT and SiT) failed due to training instability. To address this, we introduced a series of architectural modifications that successfully enabled the stable training of 600M-7B parameter Diffusion Transformers across various modalities. Besides, our ablation studies on ImageNet (Figure 3) also demonstrate the effectiveness of scaling diffusion transformers.

Q3: Typo.

Thank you for pointing out this typo. We have corrected it.

We thank reviewer 4 (ozQQ) for acknowledging the contribution of our paper and providing thoughtful comments. Please see our response to the feedback below.

Q1: Novelty.

Please refer to the global reply for more details. We acknowledge that our framework incorporates existing techniques. Our objective is to demonstrate their effectiveness in enhancing training stability and efficiency when developing scalable diffusion transformers across various modalities. By leveraging these design choices, we scale the diffusion transformer to 7B parameters and showcase new capabilities, such as RoPE-based resolution extrapolation, high-resolution editing, compositional generation, and style-consistent generation.

Q2: More comparison with sota methods

Thank you for your insightful feedback. We would like to clarify that Flag-DiT demonstrates better performance on metrics such as FID and CLIP-Score, surpassing the comparison methods. Additionally, our ablation studies confirm that Flag-DiT achieves significantly faster convergence compared to DiT and SiT of the same size.

As mentioned in our above response, the focus of our paper is on exploring a scalable architecture for DiT, which is compatible with many other methods. For instance, we conducted experiments incorporating MoE into the Flag-DiT architecture. Specifically, we implemented both time-centric and token-centric router versions of DiT-MoE, training them on the ImageNet-256 benchmark for 700k steps. The results, detailed in the table below, show improvements in FID and other metrics compared to the baseline, demonstrating the potential of combining Flag-DiT with more advanced techniques.

Besides, since ImageNet serves primarily as a proxy task for generation and conventional metrics may not accurately reflect the generation quality, we have conducted more comprehensive comparative experiments and visualizations on the more complex text-to-image task. Some generated images are visualized in Figure 7-8. Besides, we design an AI preference study to evaluate Lumina-T2I against other text-to-image models, following PixArt [1]. Specifically, we employ GPT-4o, the SoTA multimodal LLM exhibiting strong alignment with human preference, as our evaluator to vote based on image quality and text-image alignment. As shown in the following table, Lumina-T2I demonstrates competitive performance with advanced text-to-image models including PixArt and SD3. Note that SD3 uses over 1B text-image pairs, which is ~100x greater than our models. Besides, our model also uses less than 1/3 training compute of PixArt-Sigma, which is already a training-efficiencient model. However, we have to admit that Lumina-T2I still underperforms these SoTA models in terms of text-image alignment or compositional generation, due to inadequate data and training.

We also conducted a quantitative evaluation to validate the resolution extrapolation capability of Lumina-T2I. The results indicate that Lumina-T2I, equipped with NTK-aware Scaled RoPE, Time-shifting, and Proportional Attention, exhibits better extrapolation performance compared to PixArt and SD3.

[1] Chen, Junsong, et al. "Pixart-\sigma: Weak-to-strong training of diffusion transformer for 4k text-to-image generation." arXiv preprint arXiv:2403.04692 (2024).

We thank reviewer 3 (LFRv) for acknowledging the contribution of our paper and providing thoughtful comments. Please see our response to the feedback below.

Q1: Limited originality.

Please refer to the global reply for more details. We acknowledge that our framework incorporates existing techniques. Our objective is to demonstrate their effectiveness in enhancing training stability and efficiency when developing scalable diffusion transformers across various modalities. By leveraging these design choices, we scale the diffusion transformer to 7B parameters and showcase new capabilities, such as RoPE-based resolution extrapolation, high-resolution editing, compositional generation, and style-consistent generation.

Q2: Joint training of different modalities.

Thank you for your insightful suggestion regarding the potential for co-training an all-in-one foundation model. We would like to highlight that it is difficult to directly train such a powerful model without comprehensively exploring the algorithms, architectures, and design choices of diffusion transformers. Therefore, the key contribution of our paper is to explore the principle framework for building diffusion transformers in various modalities. Our proposed Lumina-T2X can be adapted to any modality with minimal changes as long as there is a well-defined continuous space.

To elaborate on the challenges mentioned in Appendix Section A, the latent space distributions differ significantly across modalities. This is unlike autoregressive models, which utilize a unified discrete token representation. For example, while joint training with images and videos can enhance visual quality, it may negatively affect the dynamic aspects of video generation. Furthermore, multiview and audio representations vary even more significantly. The disparity in data quantity—where high-quality image data is more abundant than video data, which in turn exceeds multiview data—poses additional challenges for joint training. We believe that addressing these challenges will be critical for future work on an end-to-end foundation model and we have incorporated the above discussion into our revised paper appendix.

Q3: Settings of image editing.

For image editing, we use the prompt of the target image directly during the sampling. We assume that the user provides a mask for the specific region or object they wish to edit. Alternatively, we can automate mask creation using models like SAM or DINO based on the semantics of the target prompt.

Q4: Model sizes in Table 1.

Thank you for your suggestion. In Table 1, we only highlighted the size of our models because they are significantly larger than the others listed. We also demonstrated in our ablation studies that the 600M Flag-DiT outperforms other models of the same size such as DiT and SiT.

We thank reviewer 2 (L7zd) for acknowledging the contribution of our paper and providing thoughtful comments. Please see our response to the feedback below.

Q1: Novelty.

Please refer to the global reply for more details. We acknowledge that our framework incorporates existing techniques. Our objective is to demonstrate their effectiveness in enhancing training stability and efficiency when developing scalable diffusion transformers across various modalities. By leveraging these design choices, we scale the diffusion transformer to 7B parameters and showcase new capabilities, such as RoPE-based resolution extrapolation, high-resolution editing, compositional generation, and style-consistent generation.

Q2: Video Quality.

Thank you for highlighting the limitations of our video results. Our paper primarily focuses on establishing a unified, simple, and scalable DiT architecture with extensive experiments on image generation, alongside preliminary validation in other modalities such as video and 3D. Compared to current state-of-the-art video generation models, Lumina utilizes significantly fewer computational resources, data, and tricks.

One of the biggest challenges in video generation is efficiently modeling the highly redundant video frames along the temporal dimension. Lumina follows the t2i settings and still uses the sdxl image VAE to compress each frame, which introduces redundancy in video sequences. Moreover, using a large amount of image data for joint training is a recently recognized technique in the community to significantly enhance video visual quality. The oversaturation issue you mentioned may be related to Lumina's use of video data alone for training and the high CFG scale during inference.

We are currently training a version of Lumina-T2V with joint image-video training based on a 3D VAE, and we have found that it significantly improves video quality under the same architecture.

Q3: Additional Quantitative metrics, e.g., user studies.

We agree that conventional metrics such as FID and CLIP-Score may not accurately reflect the generation quality. Conducting user studies during the rebuttal stage can be time-consuming and expensive, so we design an AI preference study to evaluate Lumina-T2I against other text-to-image models, following PixArt [1]. Specifically, we employ GPT-4o, the SoTA multimodal LLM exhibiting strong alignment with human preference, as our evaluator to vote based on image quality and text-image alignment.

As shown in the following table, Lumina-T2I demonstrates competitive performance with advanced text-to-image models including PixArt and SD3. Note that SD3 uses over 1B text-image pairs, which is ~100x greater than our models. Besides, our model also uses less than 1/3 training compute of PixArt-Sigma, which is already a training-efficient model.

However, we have to admit that Lumina-T2I still underperforms some SoTA models in terms of text-image alignment or compositional generation. In addition to the gap in data size and training compute, SD3 proposes the MMDiT architecture, which leverages an additional text branch and joint attention to refine T5 text embedding. In contrast, Lumina-T2I leverages cross-attention to inject causal LLaMA text embeddings. We believe that the text-image alignment can be further enhanced by adding a bidirectional transformer to refine the causal LLaMA text embeddings.

[1] Chen, Junsong, et al. "Pixart-\sigma: Weak-to-strong training of diffusion transformer for 4k text-to-image generation." arXiv preprint arXiv:2403.04692 (2024).

Q4: Audio demos.

Thanks for your suggestion. We have added some audio demos in the updated supplementary materials.

We thank reviewer 1 (X8ux) for acknowledging the contribution of our paper and providing thoughtful comments. Please see our response to the feedback below.

Q1: Further discussion on some design choices, e.g., RoPE vs 3D RoPE.

We thank the reviewer for raising this question. Current implementations of 2D/3D RoPE [1,2] consider only axial frequencies by simply splitting the positional embedding of the x/y/z-axes. This makes them functionally equivalent to our 1D RoPE with identifiers when representing spatial-temporal positions in images or videos. We argue that these designs, which introduce visual priors, are more suitable for building expert models tailored to specific visual tasks. Considering our ultimate goal of building a foundational generative model, we choose 1D RoPE for its simplicity, which is similar to the unified paradigm in autoregressive models [3], making our framework easier to generalize across various modalities and tasks.

[1] Yang, Zhuoyi, et al. "Cogvideox: Text-to-video diffusion models with an expert transformer." arXiv preprint arXiv:2408.06072 (2024).

[2] Lu, Jiasen, et al. "Unified-IO 2: Scaling Autoregressive Multimodal Models with Vision Language Audio and Action." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[3] Wang, Xinlong, et al. "Emu3: Next-token prediction is all you need." arXiv preprint arXiv:2409.18869 (2024).

Q2: Further discussion on alternative strategies such as DiT-MoE and insights into complex tasks that require scaling.

Thank you for your insightful feedback. We acknowledge the potential of MoE to enhance model scalability, as recognized by the LLM community. Our primary focus was to explore a scalable base architecture of DiT, which is orthogonal and compatible with MoE techniques.

To address your suggestion, we conducted experiments incorporating MoE into the Flag-DiT architecture. Specifically, we implemented both time-centric and token-centric router versions of DiT-MoE, training them on the ImageNet-256 benchmark for 700k steps. The results, detailed in the table below, show improvements in FID and other metrics compared to the baseline, demonstrating the potential of DiT-MoE.

Since text-to-image generation serves as a foundation for more complex tasks such as video and 3D generation, we anticipate that our exploration with Lumina in T2I will further validate the proposed architecture's effectiveness in future, more complex tasks.