AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
本文的共同一作为墨尔本大学的胡冬庭和香港科技大学的陈捷润和黄悉偈，完成于在 Snap 研究院 Creative Vision 团队实习期间。主要指导老师为任健、徐炎武和 Anil Kag，他们均来自 Snap Creative Vision 团队。该团队的主要研究方向包括 Efficient AI 和图像/视频/三维生成模型。
近些年来，以 Stable Diffusion 为代表的扩散模型为文生图（T2I）任务树立了新的标准，PixArt，LUMINA，Hunyuan-DiT 以及 Sana 等工作进一步提高了图像生成的质量和效率。然而，目前的这些文生图（T2I）扩散模型受限于模型尺寸和运行时间，仍然很难直接部署到移动设备上。
尽管以量化 / 剪枝为代表的模型压缩技术可以解决一部分问题，但直接从头训练一个轻量化可以部署在移动设备上的高效高质文生图模型仍然是巨大的挑战。
最近，来自 Snap 研究院的 Creative Vision 研究团队提出了 SnapGen，从头训练了一个仅有 379M 参数的文生图模型，并且在iPhone 16 Pro-Max 上仅需 1.4s 就可以生成超高质量的 1024x1024 图片。
和 SOTA 模型 SDXL, SD3, SD3.5, PixArt-α 等相比，SnapGen 有着同等或更好的指令跟随能力以及图像生成质感。
在多个定量测试基准和人类偏好测试中，SnapGen 同样显著超过了拥有更多参数量的模型，在 GenEval 指标上达到 0.66，并且在美学和文字 - 图像一致性等方面接近 SD3-Medium 以及 SD3.5-Large。在少步数生成的情景下， SnapGen 也同时保持了其生成质量， 在 GenEval 指标上达到 0.63（8 步）和 0.61（4 步）。
- 论文标题：SnapGen: Taming High-Resolution Text-to-Image Models for Mobile Devices with Efficient Architectures and Training
- 论文链接：https://arxiv.org/abs/2412.09619
- 项目主页：https://snap-research.github.io/snapgen/
SnapGen 方法简介
高效的模型结构
作者们对去噪 UNet 和图像解码器（AE decoder）进行了全面优化，从而获得资源使用和性能之间的最佳权衡。
与以往专注于预训练扩散模型的优化和压缩的研究不同，SnapGen 从整体架构和微观算子设计同时入手，提出了一种高效的模型结构，在显著降低模型参数和计算复杂度的同时，仍能保持高质量的生成效果。
多级知识蒸馏（Multi-level Knowledge Distillation）
为了对齐最先进的文生图模型 SD3.5-Large，SnapGen 使用 Rectified Flows 目标进行训练，从而可以直接使用 SD3.5 系列作为知识蒸馏的教师模型。
与已有的工作在相同架构类型下进行蒸馏不同，SnapGen 使用 DiT 教师模型跨架构蒸馏 UNet 学生模型，并且提出了一种先进的多级知识蒸馏框架，分别在输出和特征维度进行教师与学生模型的对齐。
为了解决不同时间步上蒸馏损失函数尺度不一的问题，作者们提出了时间步感知的缩放（timestep-aware scaling）操作，这种操作显著加速了知识蒸馏的收敛并增强了 SnapGen 学生模型的生成能力。
步数蒸馏（Step Distillation）
为了进一步减少模型的推理时间，作者们考虑使用了一种基于 LADD 的少步数蒸馏。在步数蒸馏算法中，可以直接进行 4 步推理的 SD3.5-Large-Turbo 被用来作为教师模型和判别器的特征提取。蒸馏过后的模型具有和 28 步相当的 4/8 步生成能力。
SnapGen 可视化
和 SOTA 模型 SDXL, SD3, SD3.5, PixArt-α 等相比，SnapGen 参数量最小，也是唯一可以直接部署到移动端的模型，同时有着接近或者更强的高像素图像生成能力。
下面的视频更加直观地展示 SnapGen 在移动端设备上的文生图效率与质量，在 iPhone 16 Pro Max 上仅需 1.4s 就可以生成超高质量的 1024x1024 图片。
这里展示了更多的 1024x1024 图像生成结果，SnapGen 在具有挑战性的文字，人物肢体，特定风格和概念生成中均表现出色。
SnapGen 方法细节
高效的网络结构
在去噪模型结构的选择上，扩散模型主要分为 UNet 和纯 Transformer 两大流派。
尽管纯 Transformer 架构（如 DiT）在大规模数据和算力支持下展现了广阔前景，亚马逊 AWS AI Lab 的一项研究（On the Scalability of Diffusion-based Text-to-Image Generation）表明，UNet 架构（尤其是 SDXL 架构）在相同参数量下表现出更高的性能、更低的算力需求以及更快的收敛速度。
基于此，Snap 团队调整 SDXL 中 UNet 架构的深度和宽度，并探索了如下图所示的一系列架构优化，包括移除高分辨率自注意力（SA）层、使用宽度扩展后的深度可分卷积（SepConv）替代常规卷积（Conv）、降低全连接层（FFN）的中间通道维度、更早注入文字等条件信息，以及优化自注意力（SA）与交叉注意力（CA）算子（如将多头自注意力 MHSA 替换为多查询注意力 MQA、对查询和键值应用 RMSNorm 归一化、插入旋转式位置编码 RoPE）。
虽然部分方法已在其他工作（如谷歌的 MobileDiffusion 模型）中有所提及，但这些研究往往缺乏对改动前后模型性能的全面量化评估。
相较之下，SnapGen 在 ImageNet-1K 256 像素类条件图像生成任务中，通过生成质量指标（FID）、模型时延、计算量和参数规模的综合评估，验证了每项架构改动的合理性和有效性。
最终，SnapGen 在生成质量（FID 2.06）与现有模型（如 SiT-XL）相当的情况下，大幅降低了模型大小和计算量，展现出卓越的性能和资源效率。
除了去噪模型，图像解码器同样是一个重要的优化对象。
首先，相较于整体生成时间，图像解码器的推理时间不容忽视，尤其是在少步甚至单步去噪模型的情况下。此外，在部署到移动端生成高分辨率图像时，解码器常常会遇到显存不足的报错。
Snap 团队发现，现有的 SD3 图像解码器存在大量的参数和计算冗余。这主要是由于其潜在空间采用了 16 通道，而与 SDXL 使用的 4 通道相比，16 通道更容易实现图像重建，因此其网络结构在压缩和加速方面具有更大的潜力。
为了优化这一点，Snap 团队通过移除不必要的自注意力机制和冗余的 GroupNorm 归一化层，同时减小网络宽度等方式，成功实现了近乎无损的 36 倍参数压缩，并在移动端部署中实现了 54 倍的解码加速。
高效的训练以及高级知识蒸馏
SnapGen 采用 Rectified Flows 为目标优化模型训练， 与 SD3 和 SD3.5 等较大的模型保持一致。同时 SnapGen 利用多个文本编码器 (text encoders) 包括 CLIP 以及 Gemma2-2b， 在训练中使用 classifier-free guidanc 以实现不同硬件环境下的部署需求。
基础模型在从初始训练的情况下在 GenEval 上的表现为 0.61。得益于使用相同的训练目标，SnapGen 可以将最新的 SD3.5-Large 作为知识蒸馏的教师模型。
然而在知识蒸馏过程中，仍然有很多需要解决的挑战：教师模型（DiT）和学生模型（UNet）的异构性，蒸馏损失函数和 Rectified Flows 任务损失函数的尺度不一致，以及常常被研究人员忽视的不同时间步上去噪预测难度的差异。
为了解决上述的问题，Snap 团队提出了一种新颖的多级别知识蒸馏范式，并且进行了时间步感知的损失函数尺度缩放。在任务损失函数之外，SnapGen 的训练还使用了输出蒸馏损失函数与特征蒸馏损失函数。
和之前使用知识蒸馏的工作（LinFusion，BK-SDM）不同，SnapGen 不需要预设不同损失函数的尺度，而是根据不同时间步上的统计数据将这些损失函数缩放到同一个尺度，保证每部分对训练的贡献均等，这种操作也被作者称为时间步感知尺度缩放（timestep-aware scaling）。
实验表明，这种考虑时间步变化的尺度缩放可以有效加速训练，并且知识蒸馏后的模型在 GenEval 上的表现高达 0.66。
步数蒸馏提升推理速度
作者们基于 LADD 等 diffusion-GAN 混合结构对 SnapGen 进行步数蒸馏。使用 4 步模型 SD3.5-Large-Turbo 同时作为教师模型和判别器的特征提取器。
SnapGen 可以快速适应少步推理的需要，和 28 步的基础模型相比，4 步与 8 步的推理结果在视觉效果上没有显著差别，GenEval 上也仅仅只有 0.05/0.03 的差距。
实验结果
作者们在多个测试基准（GenEval，DPG-Bench，CLIP Score on COCO，ImageReward）上定量对比了 SnapGen 和大量现有的高分辨率文生图模型，涵盖了 PixArt 系列，Stable Diffusion 系列，Sana，LUMINA-Next，Playgroundv2/v2.5，IF-XL。尽管参数量最小且吞吐量最高，SnapGen 仍然在所有指标中均排在前列。
在人类偏好测试中，和 SD3-Medium，SD3.5-Large 相比，SnapGen 生成的图像具有更真实的质感，并且在美学和文字 - 图像一致性等方面接近；同时 SnapGen 在所有评测指标中均显著超越 SDXL。
经过步数蒸馏的少步模型同样具有优秀的高效文生图能力，对于基准模型而言，4/8 步的生成结果通常会比较模糊或者确实重要细节。
而 SnapGen 的 4/8 步生成仍然能保持接近 28 步的视觉效果，在 GenEval 等定量基准上也非常接近 28 步的 baseline。
通过设计高效的去噪模型架构，使用先进的知识蒸馏和少步数蒸馏算法，Snap 团队提出了能直接部署到手机上的 SnapGen 模型。SnapGen 仅有 379M 参数，仅用 1.4s 就可以在 iPhone 上生成 1024x1024 图片，却在生成质量方面超出大部分现有模型。
在文生图模型随着 scaling law 越来越大的今天，SnapGen 作为小尺寸高效率模型走出了一条不一样的路，也相信会给生成模型的研究带来启发。
