[论文解读] SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers

Enze Xie, Jun Song Chen+9 more|arXiv (Cornell University)|Oct 14, 2024

Advanced Vision and Imaging被引用 4

一句话总结

Sana 提出了一种高度高效的文本到图像扩散框架，通过使用 32× 深度压缩自编码器、DiT 架构中的线性注意力机制、基于上下文指令微调的仅解码器 LLM（Gemma）进行文本编码，以及基于流的 DPM 求解器，实现了前所未有的速度，能够在 4096×4096 分辨率下生成具有最先进质量的图像。Sana-0.6B 的推理速度比 SOTA 模型（如 FLUX）快 100 倍，参数量仅为后者的 1/20，且可部署于 16GB 显存的笔记本电脑 GPU 上。

ABSTRACT

We introduce Sana, a text-to-image framework that can efficiently generate images up to 4096$ imes$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$ imes$, we trained an AE that can compress images 32$ imes$, 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$ imes$1024 resolution image. Sana enables content creation at low cost. Code and model will be publicly released.

研究动机与目标

• 解决现有高分辨率文本到图像扩散模型（尤其是 4K 分辨率下）计算成本高、推理速度慢的问题。
• 实现在消费级硬件（如 16GB 显存笔记本电脑 GPU）上高效、实时推理高保真 4K 图像。
• 在不牺牲图像质量或文本-图像对齐性能的前提下，减小模型尺寸并降低推理延迟。
• 通过利用仅解码器 LLM 的指令遵循与推理能力，克服传统文本编码器（如 CLIP、T5）的局限性。
• 构建一个可扩展的端到端训练与推理流程，支持在极低硬件要求下实现超高分辨率图像合成。

提出的方法

• 提出一种深度压缩自编码器（AE-F32），将图像空间维度压缩 32 倍，与标准 8× 压缩自编码器相比，潜变量 token 数量减少 16 倍。
• 将 DiT 中所有标准自注意力模块替换为线性注意力机制，将计算复杂度从 O(N²) 降低至 O(N)，从而实现高分辨率生成的高效计算。
• 引入 Mix-FFN，一种嵌入 MLP 模块的 3×3 深度可分离卷积，实现局部 token 聚合，消除对位置编码（NoPE）的需求。
• 将 T5 替换为小型仅解码器 LLM（Gemma）作为文本编码器，并采用复杂人类指令（CHI）提示与上下文学习，以提升指令遵循能力与图像-文本对齐性能。
• 设计了基于多 VLM 的图像字幕生成流水线，结合 clipscore 的动态字幕选择机制，以提升训练数据质量与模型收敛性。
• 提出 Flow-DPM-Solver，一种基于流的采样方法，将推理步数从 28–50 减少至 14–20，同时提升生成质量并降低延迟。

实验结果

研究问题

• RQ1 能否通过紧凑架构实现在 4K（4096×4096）分辨率下高质量、低延迟的文本到图像扩散模型生成？
• RQ2 线性注意力能否在不损失质量的前提下替代 DiT 中的标准自注意力机制，从而实现高效的高分辨率推理？
• RQ3 当使用复杂人类指令进行提示时，小型仅解码器 LLM（如 Gemma）能否在文本-图像对齐方面超越传统 T5 或 CLIP 编码器？
• RQ4 通过先进自编码、高效注意力与优化采样方法的结合，能否实现相比 SOTA 模型推理时间减少超过 100 倍？
• RQ5 是否可行将高分辨率扩散模型部署于 16GB 显存的笔记本电脑 GPU 上，并实现 1024×1024 图像生成的亚秒级推理？

主要发现

• Sana-0.6B 在 16GB 显存笔记本电脑 GPU 上生成 1024×1024 图像的时间低于 1 秒，展现出实时推理能力。
• Sana-0.6B 在 4096×4096 分辨率下推理速度比 FLUX 快 106 倍，将 A100 GPU 上的延迟从 469 秒降低至 9.6 秒。
• Sana-0.6B 在 1024×1024 图像生成中取得 FID 5.81 和 CLIP 分数 28.36，优于 SDXL 和 PixArt-Σ 等模型在两项指标上的表现。
• 使用 Gemma-2B 作为文本编码器时，取得 FID 5.9 和 CLIP 分数 26.8，性能与 T5-Large 和 T5-XXL 相当或更优，且延迟更低或相当。
• Flow-DPM-Solver 将采样步数从 28–50 减少至 14–20，提升推理速度，同时保持或改善生成质量。
• 结合 Mix-FFN 的 NoPE（无位置编码）设计可有效建模序列依赖关系，无需位置编码，简化了网络结构并提升了训练稳定性。