SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers

Q2. Difference usage of Decoder-only LLM

A2: Both ELLA[3] and LLM4GEN[4] are early explorations of using LLM as a text encoder, we will cite them in our paper. However, they only simply replace T5/Clip with LLM and lack research on stimulating LLM's reasoning capabilities. LiDiT[5] attempts to promote the reasoning ability of LLM through simple human instruction, but it is a relatively simple attempt. Additionally, LiDiT is also necessary to fine-tune the last few layers of the LLM model, which increases the training cost and pipeline complexity. Our CHI in SANA designs more complex prompt templates and in-context-learning techniques, which can fully stimulate the high-order reasoning ability of LLM, extract better text embedding, and thus improve semantic alignment capabilities.

We added the ablation study below to compare the effects of LiDiT and SANA's CHI and found that CHI consistently outperformed LiDiT.

Complex Human Instruction(SANA) vs Simple Human Instruction(SHI)(LiDiT)

Q3 Discuss Potential Misuse of SANA

Good question, this is a common question for all image generative AI models. To enhance safety, we have equipped SANA together with a safety check model (e.g. ShieldGemma-2B[6]). Specifically, the user prompt will first be sent to the safety check model to determine whether it contains NSFW(not safe for work) content. If the user prompt does not contain NSFW, it will continue to be sent to SANA to generate an image. If the user prompt contains NSFW content, the request will be rejected. After extensive testing, we found that ShieldGemma can perfectly filter out NSFW prompts entered by users under strict thresholds and our pipeline will not create harmful AI-generated content. We will add this part to our final paper.

[1] Chen et. al. Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models. 2024

[2] Liu et. al. LinFusion: 1 GPU, 1 Minute, 16K Image. 2024

[3] Hu et. al. ELLA: Equip Diffusion Models with LLM for Enhanced Semantic Alignment. 2024

[4] Liu et. al. LLM4GEN: Leveraging Semantic Representation of LLMs for Text-to-Image Generation. 2024

[5] Ma et. al. Exploring the Role of Large Language Models in Prompt Encoding for Diffusion Models. 2024

[6] ShieldGemma: Generative AI Content Moderation Based on Gemma https://huggingface.co/google/shieldgemma-2b

Thank you very much for taking the time to review and for your support. We try our best to address your questions as follows.

Q1: Difference with previously explored methods

Thank you for your insightful comments regarding SANA combining several explored components, but SANA does not simply combine several existing components but makes significant modifications:

A1: 1. Linear Attention in LinFusion[2] vs. SANA:

• Differences in use: LinFusion employs linear attention mechanisms selectively and sparingly within its U-Net based structure (SDv2 and SDXL), mainly as a supplement to standard self-attention layers. This contrasts sharply with our approach in the DiT architecture, where linear attention is used comprehensively across all layers, enhancing computational efficiency in generative tasks.
• Foundational vs. distillation Focus: LinFusion’s use of linear attention is part of a distillation strategy. In contrast, our model is designed as a foundational generative model and is trained from scratch, where the global replacement of vanilla self-attention with linear attention is explored in depth. We discovered through extensive experimentation that straightforward replacement does not suffice; it requires a suite of supportive design tweaks, such as the Mix-FFN as discussed in Section 2.2, to achieve optimal performance.
• Technical implementation: The linear attention in LinFusion relies on matrix compression techniques akin to State Space Models (SSM), mapping all tokens to a reduced NxN low-rank state before updating. Our implementation, however, adheres more closely to direct O(N) attention computation, which is integral to maintaining the integrity of the transformer architecture. We aim to successfully apply the linear attention designs from the perception field to the generative field, which has not been effectively addressed in previous studies.

2. DC-AE[1]: First, DC-AE and SANA are concurrent works and both are submitted to ICLR 2025. Second, while higher compression ratio autoencoder is discussed in concurrent works DC-AE, they focused more on ImageNet data reconstruction and simple class-condition generation. However, SANA focus on efficient and high-quality Text-to-Image generation, which presents unique challenges and requirements not addressed by DA-AE paper, as shown below:
   • Directly replaced with DC-AE on our baseline PixArt-Sigma leading to unsatisfactory results. Our in-depth investigation found that letting DiT focus on denoising purely instead of compressing tokens can compensate for performance loss (e.g. patch size: 2 -> 1).
   • Simply enlarging the channel in latent space (F8C4->F32C32) will make training converge much slower, we instead design Linear Attention + Mix-FFN blocks for fast convergence
   • We design complex human instruction and learnable small-scale factor~(e.g. initialize it with 0.01) in Decoder-only LLMs to stabilize training and improve semantic alignment.

Thank the authors for the detail answers to my questions. Most of my concerns, especially about novelty, are well-addressed and I have to admit that I was wrong about DC-AE as I did not check that it was indeed a concurrent work. Also, is the Fig. 13 the description of how the authors perform CHI? The author should inlclude a brief paragraph saying that figure 13 is a description of CHI pipeline, so that it is easier for readers to search. Based on the rebuttal, I've increased my score to 8. Below I have two additional remarks (these are merely further suggestions and have no impact on my scoring):

1. regarding to the question about token limit of CLIP from Reviewer 9t32. The authors can consider to use AltCLIP [6] to evaluate CLIPScore as its token limit is 514 which I believe is long enough.

2. And the question raised by Reviewer S4E6 about evaluation for 4K is interesting. What evaluation protocol that the authors think is suitable for 4K ? Because there is an increasing number of work tackles the 4K image synthesis task.

Reference: [6] https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltCLIP

Thank you very much for taking the time to review and for your support. We try our best to address your questions as follows.

Q1: Details of linear attention modeling

A1: Our linear attention maintains the global receptive field and long-range dependency capabilities of traditional vanilla self-attention. The primary modification involves replacing the softmax similarity computation with a ReLU-based similarity function[1,2]. This change allows us to leverage the associative property of matrix multiplication, significantly reducing computational complexity and memory usage from quadratic to linear, without compromising the mechanism’s effectiveness. Specifically:

As discussed in EfficientViT[2], given input $x \in \mathbb{R}^{N \times f}$, traditional softmax attention can be expressed as:

where $Q = xW_Q$, $K = xW_K$, $V = xW_V$ are projections, and $Sim(Q, K) = \frac{\exp(QK^T)}{\sqrt{d}}$ defines the softmax-based similarity.

To reduce computational complexity from quadratic to linear, we use ReLU linear attention:

which modifies the output calculation to:

This formulation enables using the associative property of matrix multiplication, simplifying the calculation of $O_i$ to:

As shown in Eq. (3), we only need to compute $(\sum_{j=1}^{N} ReLU(K_j)^T V_j) \in \mathbb{R}^{d \times d}$ and $(\sum_{j=1}^{N} ReLU(K_j)^T) \in \mathbb{R}^{d \times 1}$ once, then can reuse them for each query, thereby only requires $\mathcal{O}(N)$ computational cost and $\mathcal{O}(N)$ memory.

Besides, We discovered through extensive experimentation that straightforward replacement of vanilla self-attention to linear attention does not suffice; it requires a suite of supportive design tweaks, such as the Mix-FFN as discussed in Section 2.2, to achieve optimal performance.

[1] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and Fran ̧ cois Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156–5165. PMLR, 2020.

[2] Han Cai, Junyan Li, Muyan Hu, Chuang Gan, and Song Han. Efficientvit: Lightweight multi-scale attention for high-resolution dense prediction. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 17302–17313, 2023.

Q2: Should include metrics e.g. FID on 4K image generation?

A2: Thanks for the insightful asking. Currently, the datasets primarily utilized, such as COCO and MJHQ, consist of images up to 1K resolution, and there are no established versions of the 4K version of InceptionNet, which is typically used for FID calculations. Evaluating on 4K would not be particularly meaningful due to the lack of comparative methods and suitable benchmarks. Additionally, there is no standardized approach for conducting FID analysis at 4K resolution. Therefore, at this stage, our focus is on exploratory analysis with the available resolutions, primarily assessing visual quality rather than striving for high-resolution FID metrics.

First and foremost, I would like to express my gratitude to the authors for their responses, which have addressed the concerns I previously raised. After re-examining the paper and the authors' latest replies, I have the following suggestions and questions:

1. I noticed that the authors have not updated the experimental results from their responses into the appendix of the paper. I believe these experimental findings are highly significant and valuable for reference, and they should be included in the paper, at the very least in the appendix.

2. The authors have emphasized the use of ClipScore. Upon reviewing the relevant section in the paper's appendix, I noticed that the authors did not explain the implementation or calculation method of ClipScore. I am curious about how ClipScore, which is based on CLIP and is known to have a text length limitation of approximately 77 tokens, is used to assess the quality of captions of such length.Regarding the Experimental Results and the Use of ClipScore in the Paper

3. I believe some relevant works should be cited, such as ultra-high-resolution generation: UltraPixel: Advancing Ultra-High-Resolution Image Synthesis to New Peaks (NeurIPS 2024) and efficient generation work: Stable Cascade.

Thank you very much for taking the time to review and for your support. We try our best to address your questions as follows.

Q1 and Q4: Definition of Gemma2-2B-IT and adding more citations

A1 & A4: Thanks a lot for your noticing. Gemma[1] is a family of lightweight, state-of-the-art open models from Google, built from the same research and technology used to create the Gemini models. The specific Gemma we use in Sana is the Gemma-2-2B-IT version. "IT" here represents the model is Instruct-Tuned for better prompt following and chat ability, experiments show that Instrucion-tuned versions are superior to non-Instruction-Tuned models, as shown in Table 9 in the revised version. The related papers mentioned by reviewers are also cited in the Appendix.

Q2: More comprehensive comparison between Gemma and T5 series

A2: Thank you for your insightful comments regarding the comparative evaluation of the Gemma2 and T5-XXL models as depicted in Table 9. As requested, we expanded our analysis beyond FID and CLIP-Score to include a comprehensive set of performance metrics from both GenEval and HPSv2 benchmarks, providing a broader perspective on each model’s capabilities.

Table 9: The extended results, as shown in the updated table below, highlight the comparative performance after 50K training steps on an A100 GPU with FP16 precision

These results confirm that the Gemma2-2B-IT model not only matches the performance of the significantly larger and slower T5-XXL model in all metrics but does so with considerably higher efficiency. In addition, the Gemma-2B models outperform the T5-Large while operating at a similar speed, offering a more efficient solution.

Furthermore, Gemma2 offers additional advantages including multilingual capabilities (zero-shot supporting Chinese and emoji), as detailed in Figure 13 of the Appendix, and features advanced in-context learning, illustrated in Figure 12. It is also optimized for rapid inference on consumer-grade GPUs with less than 12GB of memory, a capability demonstrated in the appendix video. These attributes make the Gemma2 series particularly suitable for applications requiring efficient, high-quality multilingual text processing on commonly available hardware.

Q3: Why is Sana better than Lumina-Next with the same Gemma?

A3: Thank you for your question regarding the comparative performance of Gemma2 versus Lumina-Next. It’s important to note that the language model is just one component of our approach, and there are several distinct aspects in which Gemma2 diverges from Lumina-Next:

• 1. We use Gemma2-2B-Instruct, rather than Gemma1 used in Lumina-Next.
• 2. Model Architecture: Sana features unique design elements such as a 32x compressed VAE and linear attention mechanisms, enhancing both its performance and efficiency, which is significantly different from Lumina-Next.
• 3. Data Construction and Quality: We have enriched our dataset by integrating captions from multiple visual language models (VLMs) and employing a temperature-based caption selection strategy using clip-score. These methods not only diversify the dataset but also significantly enhance the model’s semantic alignment capabilities, as extensively discussed in Section 3.1.

In summary, the above differences underpin the superior performance of Sana over Lumina-Next in key metrics, demonstrating its advanced semantic understanding capabilities.

Thank you very much for your insightful questions and valuable suggestions.

Here, we try to address your concern on the caption selecting strategy based on CLIP-Score.

A6: We primarily uses the CLIP model to score the caption as the selective strategy during training. The 77-token limitation of the original CLIP model may seem restrictive. However, after extensive observation, we found that the most valuable content typically appears in the first half of the caption, such as "The image captures $content of image\$...", that is, 77-token length is typically sufficient to encapsulate essential semantic information for our purposes. Consequently, we believe that the CLIP-Score based on 77 tokens still offers valuable insights and relevance for assessing caption quality in our context, which will benefit for performance rather than harmful.

Additionally, we conducted an analysis of the effects of employing different sampling strategies. As illustrated in the figure (Caption-selecting-statistic.jpg) in the updated supplementary material zip file, using a ‘random’ temperature involves selecting a training sample randomly from various VLM captions. A temperature of 0.1 corresponds to sampling based on Clip-Score with this specific temperature setting, which moderates selection intensity, whereas a temperature of 0 always selects the caption with the highest Clip-Score.

Upon analyzing 50,000 samples and the frequency with which three different captions were sampled, we observed that the strategies with temperatures of 0.1 and 0 consistently outperform the random selection approach. However, exclusively using the highest Clip-Score caption can lead to the issues noted by reviewers, where Clip-Score may not perfectly assess the image-text alignment. Consequently, we opted for a more balanced approach by setting the temperature to 0.1, aiming to select captions that more accurately reflect the semantic content of the images from a holistic dataset perspective.

In the future, we will attempt to fine-tune a larger context-length CLIP model or use more advanced VLM to score the text-image pairs to make the data pipeline better.

Thank you so much for your kind words and strong support! We will keep moving to make Sana and the following projects better and better. And of course, the VQScore and the AltCLIP[1] mentioned by the reviewer Dipw24 will be definitely considered in our dataset filter pipeline in the future.

[1] https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltCLIP

Dear Authors,

Thank you for your detailed response to the concerns raised regarding the CLIP-Score caption selection strategy. I appreciate your clarification on the observation that "the most valuable content typically appears in the first half of the caption." While this point has been noted in previous work, your explanation does help to address the doubts about the use of CLIP-Score.

However, I believe that VQASCore (ECCV'24), which directly measures text-image alignment using the T5-xxl model, could be a more reliable metric for assessing caption quality. VQASCore offers a more accurate assessment of text-image alignment.

Therefore, I recommend considering the use of VQASCore or other advanced VLMs in future work to further optimize your data pipeline and enhance the performance and robustness of your model.

Thank you once again for your efforts and detailed response.

Best regards,

Thank you very much for taking the time to review and for your support. We try our best to address your questions as follows.

Q1: In detail comparison with other works e.g. PixArt-Sigma and Playgroundv3

A1: We acknowledge the call for a detailed ablation study and provide extensive comparisons throughout the paper to highlight Sana’s advancements over existing technologies such as PixArt-Sigma. Actually, our experimental design and enhancements are systematically executed based on PixArt-Sigma, following a component-by-component approach to demonstrate Sana’s incremental and innovative contributions. Specifically,

• VAE: Our high-compression autoencoder’s efficacy is demonstrated in Table 1 by comparing reconstruction abilities against PixArt-Sigma’s and SDXL’s VAE.
• Training schedule: Additionally, in Table 3, we discuss the superiority of our training schedule based on flow-matching over PixArt-Sigma’s DDPM approach.
• Captions: Table 4 highlights the benefits of our enhanced captioning approach with complementary VLMs and a temperature-based caption selection strategy using clip-score as outlined in Section 3.1, compared to PixArt-Sigma’s single caption strategy.
• Model designs: Furthermore, our model design’s efficiency and capability improvements are evidenced in Tables 8 and 12 through innovations like Linear Attention, MixFFN, and Kernel-Fusion over the foundational model designs, Flash Attn & FFN, used by PixArt-Sigma.

Regarding PlaygroundV3, direct comparisons are challenging due to differences in dataset compositions, unknown training scale, and the model’s size—24B for PlaygroundV3 vs. 1.6B/0.6B for Sana. These disparities make an apple-to-apple comparison difficult. Our primary focus is on efficiency, aiming to develop models that perform optimally on consumer-grade GPUs.

Q2: Could enlarge LLMs further improve the generation quality?

A2: Our results from Table 9 in the revised paper illustrate that larger models such as LLMs and DIT do indeed offer performance improvements. However, our focus remains on efficiency, particularly for deployment on consumer devices with 12/16GB of GPU memory. The next size up for Gemma is 9B, which would entail substantial GPU memory overhead. Training and inference of LLM models at bf16 precision require 18GB, which is impractical for consumer GPUs with less than 24GB for training and less than 16GB for inference. The costs associated with retraining a Gemma9B model are also prohibitively high. Therefore, we prioritize developing more efficiency-focused models that can deliver significant performance while remaining feasible for use on widely available hardware.

Dear Reviewer mBM3,

As the discussion period is nearing its conclusion, we kindly ask if you could review our response to ensure it addresses your concerns. Your feedback is greatly appreciated.

Thank you for your time!

Best,

Authors

In the revised paper,

1. We add GenEval and HPSv2 to further compare T5 vs Gemma in Line 281 and Table 9. Asked by Reviewer-9t32.
2. We add more related works in the paper, like, Stable Cascade, UltraPixel in Line 823. Asked by Reviewer-9t32
3. We discuss the CLIP text encoder's input token length limitation asked by Reviewer-9t32 in Line 1021. Asked by Reviewer-9t32
4. We add more details about Gemma2-2B-IT used in Sana in Line 1041. Asked by Reviewer-9t32
5. We add more comparison with LiDiT's human instruction in Line 1059,1075, and add ablation in Table 10. Asked by Reviewer-Dipw
6. We discuss the potential misuse of AI models in Line 1221. Asked by Reviewer-Dipw

Please see our detailed responses below each review for other questions and suggestions. We sincerely thank all reviewers for their thorough reviews and constructive feedback again.