RTX 4090 笔记本 0.37 秒直出大片：英伟达联手 MIT 清华祭出 Sana 架构，速度秒杀 FLUX

一台 4090 笔记本，秒生 1K 质量高清图。英伟达联合 MIT 清华团队提出的 Sana 架构，得益于核心架构创新，具备了惊人的图像生成速度，而且最高能实现 4k 分辨率。

一台 16GB 的 4090 笔记本，仅需 0.37 秒，直接吐出 1024×1024 像素图片。

如此神速 AI 生图工具，竟是出自英伟达 MIT 清华全华人团队之笔！正如其名字一样，Sana 能以惊人速度合成高分辨率、高质量，且具有强文本-图像对其能力的模型。

而且，它还能高效生成高达 4096×4096 像素的图像。

Sana 的核心设计包含了以下几个要素：

• 深度压缩自编码器（AE）：传统自编码器只能将图像压缩 8 倍，全新 AE 可将图像压缩 32 倍，有效减少了潜在 token 的数量。

• 线性 DiT（Diffusion Transformer）：用「线性注意力」替换了 DiT 中所有的普通注意力，在高分辨率下更加高效，且不会牺牲质量。

• 基于仅解码器模型的文本编码器：用现代的仅解码器 SLM 替换 T5 作为文本编码器，并设计了复杂的人类指令，通过上下文学习来增强图像-文本对齐。

• 高效的训练和采样：提出 Flow-DPM-Solver 来减少采样步骤，并通过高效的标题标注和选择来加速收敛。

基于以上的算法创新，相较于领先扩散模型 Flux-12B，Sana-0.6B 不仅参数小 12 倍，重要的是吞吐量飙升 100 倍。

以后，低成本的内容创作，Sana 才堪称这一领域的王者。

效果一览

一只赛博猫，和一个带有「SANA」字样的霓虹灯牌。

一位站在山顶上的巫师，在夜空中施展魔法，形成了由彩色能量组成的「NV」字样。

在人物的生成方面，Sana 对小女孩面部的描绘可以说是非常地细致了。

下面来看个更复杂的：

一艘海盗船被困在宇宙漩涡星云中，通过模拟宇宙海滩旋涡的特效引擎渲染，呈现出令人惊叹的立体光效。场景中弥漫着壮丽的环境光和光污染，营造出电影般的氛围。整幅作品采用新艺术风格，由艺术家 SenseiJaye 创作的插画艺术，充满精致细节。

甚至，像下面这种超级复杂的提示，Sana 也能 get 到其中的关键信息，并生成相应的元素和风格。

Prompt：a stunning and luxurious bedroom carved into a rocky mountainside seamlessly blending nature with modern design with a plush earth-toned bed textured stone walls circular fireplace massive uniquely shaped window framing snow-capped mountains dense forests, tranquil mountain retreat offering breathtaking views of alpine landscape wooden floors soft rugs rustic sophisticated charm, cozy tranquil peaceful relaxing perfect escape unwind connect with nature, soothing intimate elegance modern design raw beauty of nature harmonious blend captivating view enchanting inviting space, soft ambient lighting warm hues indirect lighting natural daylight balanced inviting glow

顺便，团队还给经典梗图，生成了一个卡通版变体（右）。

设计细节

Sana 的核心组件，已在开头简要给出介绍。接下来，将更进一步展开它们实现的细节。

模型架构的细节，如下表所示。

- 深度压缩自编码器

研究人员引入的全新自编码器（AE），大幅将缩放因子提高至 32 倍。

过去，主流的 AE 将图像的长度和宽度，只能压缩 8 倍（AE-F8）。

与 AE-F8 相比，AE-F32 输出的潜在 token 数量减少了 16 倍，这对于高效训练和生成超高分辨率图像（如 4K 分辨率）至关重要。

- 高效线性 DiT（Diffusion Transformer）

原始 DiT 的自注意力计算复杂度为 O (N²)，在处理高分辨率图像时呈二次增长。

线性 DiT 在此替换了传统的二次注意力机制，将计算复杂度从 O (N²) 降低到 O (N)。

与此同时，研究人员还提出了 Mix-FFN，可以在多层感知器（MLP）中使用 3×3 深度卷积，增强了 token 的局部信息。

实验结果显示，线性注意力达到了与传统注意力相当的结果，在 4K 图像生成方面将延迟缩短了 1.7 倍。

此外，Mix-FFN 无需位置编码（NoPE）就能保持生成质量，成为首个不使用位置嵌入的 DiT。

- 基于仅解码器「小语言模型」的文本编码器

这里，研究人员使用了 Gemma（仅解码器 LLM）作为文本编码器，以增强对提示词的理解和推理能力。

尽管 T2I 生成模型多年来取得了显著进展，但大多数现有模型仍依赖 CLIP 或 T5 进行文本编码，这些模型往往缺乏强大的文本理解和指令跟随能力。

与 CLIP 或 T5 不同，Gemma 提供了更优的文本理解和指令跟随能力，由此解训练了不稳定的问题。

他们还设计了复杂人类指令（CHI），来利用 Gemma 强大指令跟随、上下文学习和推理能力，改善了图像-文本对齐。

在速度相近的情况下，Gemma-2B 模型比 T5-large 性能更好，与更大更慢的 T5-XXL 性能相当。

- 高效训练和推理策略

另外，研究人员还提出了一套自动标注和训练策略，以提高文本和图像之间的一致性。

首先，对于每张图像，利用多个视觉语言模型（VLM）生成重新描述。尽管这些 VLM 的能力各不相同，但它们的互补优势提高了描述的多样性。

此外，他们还提出了一种基于 clipscore 的训练策略，根据概率动态选择与图像对应的多个描述中具有高 clip 分数的描述。

实验表明，这种方法改善了训练收敛和文本-图像对齐能力。

此外，与广泛使用的 Flow-Euler-Solver 相比，团队提出的 Flow-DPM-Solver 将推理采样步骤从 28-50 步显著减少到 14-20 步，同时还能获得更优的结果。

整体性能

如下表 1 中，将 Sana 与当前最先进的文本生成图像扩散模型进行了比较。

对于 512×512 分辨率：- Sana-0.6 的吞吐量比具有相似模型大小的 PixArt-Σ 快 5 倍- 在 FID、Clip Score、GenEval 和 DPG-Bench 等方面，Sana-0.6 显著优于 PixArt-Σ

对于 1024×1024 分辨率：- Sana 比大多数参数量少于 3B 的模型性能强得多- 在推理延迟方面表现尤为出色

与最先进的大型模型 FLUX-dev 的比较：- 在 DPG-Bench 上，准确率相当- 在 GenEval 上，性能略低- 然而，Sana-0.6B 的吞吐量快 39 倍，Sana-1.6B 快 23 倍

Sana-0.6 吞吐量，要比当前最先进 4096x4096 图像生成方法 Flux，快 100 倍。

而在 1024×1024 分辨率下，Sana 的吞吐量要快 40 倍。

如下是，Sana-1.6B 与其他模型可视化性能比较。很显然，Sana 模型生成速度更快，质量更高。

终端设备部署

为了增强边缘部署，研究人员使用 8 位整数对模型进行量化。

而且，他们还在 CUDA C++ 中实现了 W8A8 GEMM 内核，并采用内核融合技术来减少不必要的激活加载和存储带来的开销，从而提高整体性能。

如下表 5 所示，研究人员在消费级 4090 上部署优化前后模型的结果比较。

在生成 1024x1024 图像方面，优化后模型实现了 2.4 倍加速，仅用 0.37 秒就生成了同等高质量图像。

作者介绍

Enze Xie（谢恩泽）

共同一作 Enze Xie 是 NVIDIA Research 的高级研究科学家，隶属于由麻省理工学院的 Song Han 教授领导的高效 AI 团队。此前，曾在华为诺亚方舟实验室（香港）AI 理论实验室担任高级研究员和生成式 AI 研究主管。

他于 2022 年在香港大学计算机科学系获得博士学位，导师是 Ping Luo 教授，联合导师是 Wenping Wang 教授。并于朋友 Wenhai Wang 密切合作。

在攻读博士学习期间，他与阿德莱德大学的 Chunhua Shen 教授、加州理工学院的 Anima Anandkumar 教授以及多伦多大学的 Sanja Fidler 教授共事。同时，还与 Facebook 和 NVIDIA 等业界的多位研究人员进行了合作。

他的研究方向是高效的 AIGC / LLM / VLM，并在实例级检测和自监督 / 半监督 / 弱监督学习领域做了一些工作 —— 开发了多个 CV 领域非常知名的算法，以及一个 2000 多星的自监督学习框架 OpenSelfSup（现名为 mmselfsup）。

- PolarMask（CVPR 2020 十大影响力论文排名第十）

- PVT（ICCV 2021 十大影响力论文排名第二）

- SegFormer（NeurIPS 2021 十大影响力论文排名第三）

- BEVFormer（ECCV 2022 十大影响力论文排名第六）

Junsong Chen

共同一作 Junsong Chen 是 NVIDIA Research 的研究实习生，由 Enze Xie 博士和 Song Han 教授指导。同时，他也是大连理工大学 IIAU 实验室的博士生，导师是 Huchuan Lu 教授。

他的研究领域是生成式 AI 和机器学习的交叉，特别是深度学习及其应用的算法与系统协同设计。

此前，他曾在香港大学担任研究助理，由 Ping Luo 教授的指导。

Song Han（韩松）

Song Han 是 MIT 电气工程与计算机科学系的副教授。此前，他在斯坦福大学获得博士学位。

他提出了包括剪枝和量化在内广泛用于高效 AI 计算的「深度压缩」技术，以及首次将权重稀疏性引入现代 AI 芯片的「高效推理引擎」——ISCA 50 年历史上引用次数最多的前五篇论文之一。

他开创了 TinyML 研究，将深度学习引入物联网设备，实现边缘学习。

他的团队在硬件感知神经架构搜索方面的工作使用户能够设计、优化、缩小和部署 AI 模型到资源受限的硬件设备，在多个 AI 顶会的低功耗计算机视觉比赛中获得第一名。

最近，团队在大语言模型量化 / 加速（SmoothQuant、AWQ、StreamingLLM）方面的工作，有效提高了 LLM 推理的效率，并被 NVIDIA TensorRT-LLM 采用。

Song Han 凭借着在「深度压缩」方面的贡献获得了 ICLR 和 FPGA 的最佳论文奖，并被 MIT Technology Review 评选为「35 岁以下科技创新 35 人」。与此同时，他在「加速机器学习的高效算法和硬件」方面的研究，则获得了 NSF CAREER 奖、IEEE「AIs 10 to Watch: The Future of AI」奖和斯隆研究奖学金。

他是 DeePhi（被 AMD 收购）的联合创始人，也是 OmniML（被 NVIDIA 收购）的联合创始人。

参考资料：

• https://nvlabs.github.io/Sana/

• https://www.linkedin.com/feed/update/urn:li:activity:7251843706310275072/