引言

生成式 AI 的繁荣为我们带来了诸如 Stable Diffusion XL (SDXL) 和 Stable Diffusion 3 (SD3) 等令人难以置信的工具。这些模型可以根据简单的文本提示生成照片级逼真的图像，但它们伴随着沉重的代价: 计算算力。通常，运行这些模型需要云服务器或高端游戏 PC 中那种庞大且耗能的 GPU。

对于普通用户来说，这意味着必须依赖云服务，从而带来延迟、订阅成本和数据隐私问题。在智能手机上本地运行这些模型一直是边缘计算的“圣杯”。虽然已经有人尝试通过压缩来缩小模型，但它们通常会导致低分辨率输出 (\(512 \times 512\) 像素) 或视觉质量的显著下降。

SnapGen 应运而生。

在一篇新论文中，来自 Snap 公司及其学术合作伙伴的研究人员提出了一套针对移动端生成问题的全面解决方案。他们开发了一种文生图 (T2I) 模型，该模型不仅体积小巧 (约 3.79 亿参数，相比之下 SDXL 为 26 亿) ，而且速度极快。

各种文生图模型在模型大小、移动设备兼容性和视觉输出质量方面的比较。

如上图所示，SnapGen 在实现 \(1024 \times 1024\) 分辨率的同时，达到了具有竞争力的视觉质量，并且完全兼容移动设备——这是 SDXL 和 SD3-Medium 等巨头难以在原生环境下做到的。本文将深入探讨使其成为可能的架构、训练配方和蒸馏技术。

背景: 效率瓶颈

要理解 SnapGen 的重要性，我们首先需要了解为什么当前的模型速度缓慢。大多数现代 T2I 系统都是扩散模型 (Diffusion Models) 。 它们的工作原理是在文本提示的引导下，通过逐步去除随机信号中的噪声，直到清晰的图像显现出来。

这个过程涉及两个主要的计算重担:

  1. 去噪骨干网络 (通常是 UNet 或 Transformer) : 它需要运行多次 (步数) 来去除噪声。由于注意力机制的存在，其计算成本很高 (二次复杂度) 。
  2. 解码器 (VAE) : 一旦在压缩的“潜空间”中完成去噪，解码器会将潜在表示扩展回像素空间。

现有的提高效率的方法通常涉及“剪枝” (切断网络的一部分) 或“量化” (降低数值精度) 。虽然有帮助，但这通常只是对为服务器而非手机设计的架构进行的权宜之计。SnapGen 团队采取了不同的方法: 他们从零开始设计了一个高效的架构，并创建了一个专门的训练流程，以实现越级挑战。

方法: 面向边缘端的架构设计

研究人员从三个角度攻克了这个问题: 重新设计的 UNet、微型解码器和复杂的训练方案。

1. 高效 UNet (The Efficient UNet)

SnapGen 的核心是其去噪骨干网络。团队从 SDXL 的 UNet 开始，将其精简为“更薄、更短”的基线。在此基础上，他们进行了一系列外科手术式的架构更改，以针对移动硬件 (特别是针对苹果 A 系列芯片的神经网络引擎) 进行优化。

高效 UNet 的架构演变。

以下是上图所示的分步演变过程:

  • (b) 在高分辨率阶段移除自注意力机制: 自注意力机制的计算量随分辨率呈二次方增长。通过仅在低分辨率瓶颈层保留它，并从高分辨率阶段移除，他们在不牺牲质量的情况下大幅降低了延迟。事实上，质量反而有所提高，这可能是因为高分辨率下的注意力有时会阻碍收敛。
  • (c) 深度可分离卷积: 标准卷积参数量很大。用深度可分离卷积 (SepConv) 代替它们可以显著减少参数。为了恢复因这种减少而损失的模型表达能力，他们扩展了通道宽度 (通用倒置瓶颈设计) 。
  • (d) 削减 FFN: 前馈网络通常会将维度扩展 \(4\times\)。将此扩展比率降低到 \(3\times\) 节省了计算量，且对输出影响极小。
  • (f) 先进的注意力和归一化: 他们采用了多查询注意力机制 (MQA) ，该机制在不同头之间共享键/值以节省内存，并实施了 RMSNorm 和 RoPE (旋转位置编码) 以获得更好的训练稳定性。

这些决策的影响量化如下。请注意，“Ours (Final)”在保持极低延迟和 FLOPs 的同时，实现了最低的 FID (越低代表质量越好) 。

各种设计选择的性能和效率比较。

2. 微型且快速的解码器

移动生成中一个令人惊讶的瓶颈是自动编码器 (VAE) 解码器。即使 UNet 生成潜在表示的速度很快，将其解码为 \(1024 \times 1024\) 的图像也可能导致移动端神经网络处理单元 (NPU) 出现内存溢出 (OOM) 错误。

标准的 SDXL 解码器大约有 5000 万参数。SnapGen 引入了一个仅有 138 万参数的微型解码器。

解码器架构比较。

他们通过完全移除解码器中的注意力层 (这会消耗巨大的峰值内存) 、使用更少的残差块以及利用深度可分离卷积来实现这一点。如下表所示，虽然 SDXL 和 SD3 解码器在 iPhone 的神经网络引擎上运行失败 (OOM) ，但 SnapGen 解码器仅需 174 毫秒即可运行。

解码器指标的性能比较。

进阶训练配方

设计一个小模型只是成功的一半。与拥有数十亿参数的大模型相比，小模型通常难以捕捉文本和图像之间的复杂关系。为了解决这个问题，研究人员利用了知识蒸馏 (Knowledge Distillation, KD) 。

流匹配与知识蒸馏

SnapGen 没有使用标准的扩散目标，而是使用了整流 (Rectified Flow) 技术，这种技术可以拉直噪声分布和数据分布之间的路径。

训练目标本质上是最小化模型预测的“速度” (\(v_\theta\)) 与目标方向之间的差异。

任务损失方程

为了给小模型增效，他们使用了一个巨大的教师模型 (SD3.5-Large) 。目标是让小的“学生”模型模仿教师。然而，盲目复制教师的输出效率很低，因为教师和学生拥有不同的架构 (DiT 与 UNet) 。

研究人员引入了多级知识蒸馏策略:

  1. 输出蒸馏: 学生尝试匹配教师的最终速度预测。
  2. 特征蒸馏: 他们将教师的内部特征投影以匹配学生的维度，迫使学生学习教师“思考”的方式，而不仅仅是它的输出结果。

多级知识蒸馏概览。

时间步感知缩放

这篇论文的一个重要见解是，任务的难度随噪声水平 (时间步 \(t\)) 的不同而变化。在扩散过程的中间阶段，任务相对“容易”，而在最开始和最后阶段，预测难度会激增。

标准的蒸馏对损失应用恒定的权重。SnapGen 引入了时间步感知缩放 (Timestep-Aware Scaling) 。

平均损失幅度图。

如上图所示，任务损失 (\(\mathcal{L}_{task}\)) 和蒸馏损失 (\(\mathcal{L}_{kd}\)) 的幅度在边界处 (\(t=0\) 和 \(t=1\)) 显著发散。研究人员制定了一种缩放机制，根据时间步动态平衡这些损失:

时间步感知缩放方程。

这确保了模型恰好在最有益的时候专注于教师的监督，从而实现更快的收敛和更好的对齐。

步数蒸馏: 加速推理

即使拥有快速的模型，标准的扩散过程也需要 20-50 步才能生成图像。要在手机上实现 2 秒内的运行速度，步数需要降至个位数 (4 到 8 步) 。

SnapGen 通过对抗性步数蒸馏 (Adversarial Step Distillation) 实现了这一点。

对抗性步数蒸馏概览。

他们将扩散过程视为 GAN (生成对抗网络) 。训练一个判别器来区分学生的输出 (从先前状态一步生成) 和“真实”的去噪分布。这种对抗性压力迫使模型更快地跳向清晰图像，从而消除了对数十个精细化步骤的需求。

对抗性目标方程

实验与结果

那么，SnapGen 与那些重量级模型相比表现如何呢？

定量基准测试

研究人员在 ImageNet 和各种 T2I 基准 (如 GenEval 和 DPG-Bench) 上评估了该模型。结果令人瞩目。尽管参数量只是其他模型的一小部分 (0.38B 对比 2.6B+) ，SnapGen 在文本对齐和审美质量方面仍优于 SDXL 和 Playground v2.5。

定量基准测试评估。

在 GenEval (一个针对对象对齐的基准测试) 中，SnapGen 得分为 0.66 , 击败了 SDXL 的 0.55 。 这表明来自 SD3.5 的知识蒸馏非常有效，将大模型的“智慧”转移到了紧凑的学生模型中。

视觉质量与人类评估

数据很有用，但视觉艺术是主观的。团队进行了一项用户研究，将 SnapGen 与 SDXL、SD3-Medium 及其教师模型 SD3.5-Large 进行了比较。

人类评估结果。

SnapGen (紫色条) 在美学、对齐度和真实感方面始终击败 SDXL 和 SD3-Medium。它甚至可以与其大得多的教师模型 SD3.5-Large 分庭抗礼，特别是在真实感方面。

少步数性能

当比较步数时，对抗性步数蒸馏的影响清晰可见。如下所示，8 步甚至 4 步的生成保持了高保真度和结构，而基础模型 (没有步数蒸馏) 在低步数下则表现不佳。

少步数生成的性能比较。

端侧演示

最后，论文兑现了其主要承诺: 移动端执行。在 iPhone 16 Pro-Max 上，该模型生成一张 \(1024 \times 1024\) 的图像大约需要 1.4 秒 。

iPhone 16 Pro-Max 上的演示。

结论

SnapGen 代表了我们部署生成式 AI 思维方式的重大转变。与其简单地缩小现有的服务器级模型，这项研究强调了架构-硬件协同设计的重要性。通过专门针对移动端限制构建架构 (移除高分辨率注意力层、优化解码器) 并将其与最先进的训练技术 (流匹配、多级蒸馏) 相结合，该团队实现了以前认为需要庞大 GPU 才能完成的任务。

对于学生和研究人员来说，SnapGen 提供了一个高效深度学习的蓝图:

  1. 不要忽视解码器: 在潜在扩散中，VAE 可能是一个隐形的瓶颈。
  2. 名师出高徒: 如果从强大的教师模型正确蒸馏，微型模型也能学习复杂的概念。
  3. 动态训练: 损失权重不应总是一成不变的；适应时间步的难度会产生更好的结果。

随着像 SnapGen 这样的模型日趋成熟，我们正迈向一个未来: 高保真的内容创作可以即时、离线地在你的掌心中完成。