Source-linked AI summary

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

Enze Xie, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, Song Han

arXiv:2410.10629v3cs.CV

TL;DR

High-resolution text-to-image generation remains expensive as diffusion models grow larger. Sana combines deep latent compression, linear DiT, decoder-only LLM text encoding, and efficient captioning and flow-based sampling; Sana-0.6B delivers competitive results with over 100× higher 4K throughput than state-of-the-art methods. The system supports generation up to 4096×4096 and can run efficiently on edge hardware.

  • Problem

    Increasing diffusion-model scale raises training and inference costs, motivating high-quality, high-resolution text-to-image generation that is computationally efficient.

  • Method

    Sana combines a 32× compression autoencoder, linear attention in DiT, a decoder-only Gemma text encoder with complex human instructions, automatic caption selection, and Flow-DPM-Solver.

  • Results

    Sana-0.6B generates images up to 4096×4096 with competitive results and over 100× higher throughput than state-of-the-art methods for 4K generation.

  • Takeaways & Limitations

    Sana provides an efficient pipeline for high-resolution image generation, including deployment on edge devices for real-time generation.

  • Takeaways & Limitations

    Sana cannot fully guarantee the safety and controllability of generated content, and challenges remain for text rendering and generating faces and hands.

Abstract

from arXiv · show

We introduce Sana, a text-to-image framework that can efficiently generate images up to 4096$\times$4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU. Core designs include: (1) Deep compression autoencoder: unlike traditional AEs, which compress images only 8$\times$, we trained an AE that can compress images 32$\times$, effectively reducing the number of latent tokens. (2) Linear DiT: we replace all vanilla attention in DiT with linear attention, which is more efficient at high resolutions without sacrificing quality. (3) Decoder-only text encoder: we replaced T5 with modern decoder-only small LLM as the text encoder and designed complex human instruction with in-context learning to enhance the image-text alignment. (4) Efficient training and sampling: we propose Flow-DPM-Solver to reduce sampling steps, with efficient caption labeling and selection to accelerate convergence. As a result, Sana-0.6B is very competitive with modern giant diffusion model (e.g. Flux-12B), being 20 times smaller and 100+ times faster in measured throughput. Moreover, Sana-0.6B can be deployed on a 16GB laptop GPU, taking less than 1 second to generate a 1024$\times$1024 resolution image. Sana enables content creation at low cost. Code and model will be publicly released.

1 INTRODUCTION

Sana targets efficient, high-quality text-to-image generation at up to 4096×4096 resolution despite the rising cost of large diffusion models. It combines compressed latents, linear DiT attention, decoder-only LLM text encoding, and efficient training and inference strategies.

  • Motivation: Sana addresses the challenge of making high-quality, high-resolution image generation computationally efficient as diffusion models become increasingly costly to train and run.Industry models have grown from 0.6B to 24B parameters, increasing training and inference costs.
  • Core designs: Sana’s AE-F32 compresses images by a factor of 32 and outputs 16× fewer latent tokens than the mainstream AE-F8.The reduced token count supports efficient training and ultra-high-resolution generation, including 4K images.
  • Core designs: Linear DiT replaces quadratic self-attention with linear attention, reducing complexity from O(N^2) to O(N) while reporting comparable quality and 1.7× acceleration at 4K.Mix-FFN aggregates local token information with 3×3 depth-wise convolution, and removing positional embeddings produced no quality loss.
  • Core designs: A decoder-only Gemma text encoder uses complex human instructions, in-context learning, and reasoning capabilities to improve image-text alignment.The design also addresses training instability from directly adopting an LLM as a text encoder.
  • Efficient training and inference: Caption relabeling with multiple VLMs, CLIP-score-based caption selection, and Flow-DPM-Solver improve training consistency and reduce sampling steps from 28–50 to 14–20.The solver is reported to achieve better results than the widely used Flow-Euler-Solver.
  • Results: Sana-0.6B achieves over 100× higher throughput than FLUX for 4K generation and 40× higher throughput for 1K generation while maintaining competitive benchmark results.The model is quantized and deployed on an edge device, generating a 1024×1024 image in 0.37 seconds on a customer-grade 4090 GPU.

2 METHODS

Sana combines aggressive latent compression, linear attention, and decoder-only text encoding to make high-resolution text-to-image generation more efficient while preserving generation quality and alignment.

  • Deep Compression Autoencoder: AE-F32C32P1 outperforms AE-F16C32P2 and AE-F8C16P4 in generation FID despite weaker reconstruction metrics.The three configurations encode 1024×1024 images into the same 32 × 32 token count.
  • Efficient Linear DiT: Linear DiT replaces quadratic self-attention O(N^2) with linear attention O(N), while Mix-FFN adds 3×3 depth-wise convolution for local token aggregation.The design retains DiT’s macro architecture and uses ReLU linear attention to improve high-resolution efficiency.
  • Efficient Linear DiT: Sana removes positional embeddings without performance loss, with 3×3 convolution providing implicit positional information.This NoPE design differs from earlier DiT methods using absolute, learnable, or rotary positional embeddings.
  • Text Encoder Design: Gemma-2-2B runs 6× faster than T5-XXL with comparable CLIP Score and FID, while complex human instruction further improves prompt understanding and alignment.CHI focuses Gemma’s output on extracting and enhancing prompt details and helps stabilize generation for short prompts.

3 EFFICIENT TRAINING/INFERENCE

Sana combines multi-caption training, flow-based objectives and sampling, and staged resolution training to improve efficiency, convergence, and high-resolution generation.

  • Caption labeling and selection: Four VLMs generate diverse captions for each image, while CLIP-score sampling gives higher-quality captions greater selection probability during training.The sampler uses a temperature-controlled probability based on each caption’s CLIP score.
  • Resolution training: Cascade resolution training starts at 512px and progressively fine-tunes at 1024px, 2K, and 4K instead of pre-training at 256px.The strategy relies on AE-F32C32P1 and is intended to improve high-resolution reconstruction and image-text alignment.
  • Flow-based training: Flow-based training uses data or velocity prediction rather than DDPM noise prediction, supporting faster convergence and improved performance.Rectified Flow predicts velocity, while EDM predicts data.
  • Flow-based inference: Flow-DPM-Solver adapts DPM-Solver++ to Rectified Flow by redefining time steps and using the model’s predicted velocity to derive data.Its time range is changed to [0, 1], with a modified scaling factor and time-step shift.
  • Flow-based inference: 14–20 sampling steps yield better convergence and performance with Flow-DPM-Solver, whereas Flow-Euler requires 28–50 steps and performs worse.The comparison is reported in Figure 8.

4 ON-DEVICE DEPLOYMENT

Sana’s on-device deployment combines INT8 quantization with fused CUDA kernels to accelerate 1024px generation while preserving image quality.

  • Quantized deployment: W8A8 quantization uses per-token symmetric INT8 activations and per-channel symmetric INT8 weights, while selected layers remain in full precision.Normalization, linear attention, and key-value projection layers are retained at full precision to preserve semantic similarity.
  • Inference optimization: Kernel fusion combines linear-attention products, QKV projection, GLU, quantization, and element-wise operations to reduce runtime overhead.The implementation uses a W8A8 GEMM kernel written in CUDA C++.
  • Deployment results: 2.4× speedup reduces optimized 1024px generation to 0.37 seconds on a laptop GPU while maintaining almost lossless image quality.Table 5 evaluates quality with CLIP-Score and ImageReward.

5 EXPERIMENTS

Experiments evaluate Sana’s architecture, efficiency, image quality, text-image alignment, and deployment-oriented performance against contemporary text-to-image models.

  • Model architecture: 590M parameters make Sana-0.6B comparable in scale to DiT-XL and PixArt-Σ while Sana-1.6B uses 1.6B parameters.Sana-1.6B has 20 layers and 2240 channels per layer.
  • Evaluation setup: Sana is evaluated with FID, CLIP Score, GenEval, DPG-Bench, and ImageReward across image quality and text-image alignment.FID and CLIP Score use MJHQ-30K; GenEval and DPG-Bench measure alignment.
  • Performance comparison: 5× faster throughput than PixArt-Σ at 512×512 accompanies significantly better FID, CLIP Score, GenEval, and DPG-Bench results for Sana-0.6.At 1024×1024, Sana is stronger than most models with fewer than 3B parameters.
  • Performance comparison: 39× faster throughput than FLUX-dev for Sana-0.6B and 23× faster for Sana-1.6B are reported with equivalent or slightly lower alignment results.The comparison uses DPG-Bench and GenEval at 1024×1024.
  • Block ablations: Replacing quadratic attention with linear attention lowers latency, while Mix-FFN compensates for the resulting performance loss at some efficiency cost.The block-design analysis is conducted at 1024×1024 resolution.
  • Qualitative comparison: Sana renders text accurately and produces image quality comparable to FLUX-dev in the reported visual comparisons.The comparison also reports inaccurate text understanding for SD3 and limited text rendering for PixArt-Σ.

6 RELATED WORK

Related work spans diffusion-transformer architectures, text encoders, high-resolution generation, and deployment of diffusion models on devices.

  • Generative architectures: Diffusion Transformer and DiT-based text-to-image extensions form the architectural line of related work discussed by the paper.The paper places these methods within recent progress in generative model architecture.
  • Text encoders: Text encoders in prior work include CLIP, T5-XXL, and combinations of T5 and CLIP.The paper contrasts these prior choices with its decoder-only language-model approach elsewhere.
  • High resolution and deployment: Prior high-resolution and on-device efforts include PixArt-Σ and GigaGAN for 4K generation, plus diffusion-model deployment on mobile devices.PixArt-Σ directly generates 4K images, while GigaGAN uses super-resolution.

7 CONCLUSION

Sana is an efficient text-to-image pipeline that generates images up to 4096×4096 with competitive quality and substantially higher throughput. The paper also identifies safety, controllability, and difficult content-generation cases as limitations.

  • Sana combines deep compression, linear attention, decoder-only language encoding, automatic captioning, and flow-based sampling in one efficient pipeline.
  • 4096×4096 resolution and more than 100× higher throughput than state-of-the-art methods are reported while maintaining competitive generation results.
  • The authors do not fully guarantee the safety and controllability of generated content, and challenging cases include text rendering and faces and hands.

A FULL RELATED WORK

Related work traces efficient image generation from U-Net replacement and improved text encoders to on-device diffusion optimization. These developments motivate Sana’s focus on scalable architectures, language understanding, and deployment efficiency.

  • Diffusion Transformers replaced traditional U-Net architectures, while PixArt-α extended transformer-based diffusion to text-to-image generation.
  • Text encoders evolved from CLIP in latent diffusion models toward T5-XXL in Imagen, reflecting the importance of language understanding for image generation.
  • On-device diffusion research has investigated post-training quantization through calibration objectives and data-acquisition strategies.

B MORE IMPLEMENTATION DETAILS

The implementation details explain Sana’s flow-based training and sampling, efficient captioning, kernel acceleration, and high-resolution fine-tuning. Flow-DPM-Solver is designed to improve sampling stability and reduce the steps needed for high-quality images.

  • Flow-based Training: Flow matching and EDM use data or velocity prediction rather than DDPM noise prediction, supporting faster convergence and improved performance.
  • Flow-based Inference: Near t = T, noise prediction becomes unstable and accumulates sampling errors, whereas the data prediction model becomes approximately constant.
  • Flow-based Inference: Flow-DPM-Solver adapts DPM-Solver++ to Rectified Flow by transforming scaling factors, time steps, and model outputs for velocity-based sampling.
  • Flow-based Inference: Flow-DPM-Solver produces stable, high-quality images in 10–20 steps, whereas Flow-Euler typically requires 30–50 steps.
  • Multi-Caption Auto-labeling Pipeline: The multi-caption pipeline pairs each image with its original prompt and four VLM-generated captions whose complementary variations improve semantic alignment.
  • Triton Acceleration Training/Inference Detail: Triton kernel fusion accelerates training and inference by fusing activation, precision-conversion, padding, division, and projection operations.
  • 2K/4K fine-tuning: Positional encoding enables 2K and 4K fine-tuning, with the added encoding typically reaching convergence within 10K iterations.

C MORE RESULTS

Additional experiments evaluate reconstruction, block design, text encoding, language transfer, throughput, and high-resolution visual quality. They report competitive quality, increasing efficiency advantages with resolution, and successful Chinese and Emoji prompt transfer.

  • Autoencoder Comparison: DC-AE-F32C32 and SDXL’s VAE-F8C4 both produce reconstructions nearly indistinguishable from the original images.
  • Sana Block Ablation: Adding Mix-FFN compensates for the performance loss from replacing DiT self-attention with linear attention, while Triton fusion speeds execution without harming performance.
  • Complex Human Instruction: Complex human instruction makes Gemma-2 outputs focus more directly on understanding and enriching user prompts rather than producing conversational responses.
  • Zero-shot Language Transfer: Using Gemma-2, Sana generates corresponding images from Chinese and Emoji prompts despite training on English prompts only.
  • Text Encoder Speed: Gemma-2B has latency similar to T5-large while increasing model size, which the authors associate with improved capabilities and efficiency.
  • Diffusion Model Speed: Sana’s efficiency advantage becomes more significant as resolution increases across tests at 512, 1024, 2048, and 4096.
  • More Visualizations: The 4K visualizations contain more detail than 1K images, and supplementary materials show Sana deployed on a laptop.
Loading 2410.10629v3…