SnapGen++：释放扩散 Transformer 在边缘设备上的高效高保真图像生成潜力 SnapGen++: Unleashing Diffusion Transformers for Efficient High-Fidelity Image Generation on Edge Devices
首个在手机端实现 DiT 高质量图像生成的系统，0.4B 模型 1.8 秒生成 1K 图像
前置知识
Diffusion Transformer (DiT)
DiT 是将 Transformer 架构应用于扩散模型的范式，取代了传统 U-Net 作为去噪网络。它将图像的潜在表示（latent）分割成 patch 序列，通过 Transformer 的自注意力和前馈网络进行迭代去噪。DiT 以 PixArt-α、Stable Diffusion 3 和 Flux 为代表，展现出比 U-Net 更强的可扩展性和生成质量，但计算复杂度随 token 数量呈二次增长。
本文的核心目标是将 DiT 部署到移动设备上，理解 DiT 的架构特点和计算瓶颈是理解本文方法设计的关键。
Flow Matching
Flow Matching 是一种扩散模型训练范式，学习从噪声到数据的向量场（velocity field）。相比传统 DDPM 的 ε-prediction，Flow Matching 直接预测速度 v(x_t, t)，训练目标为最小化预测速度与真实速度之间的 MSE 损失。这种线性插值路径更稳定，采样效率更高。
本文采用 Flow Matching 作为基础训练框架，所有模型变体都基于这个范式训练，理解它才能理解后续的蒸馏方法。
Distribution Matching Distillation (DMD)
DMD 是一种步蒸馏方法，通过让学生的输出分布与教师的分布对齐来压缩采样步数。它使用一个 critic 模型估计学生的输出分布，然后计算与教师分布的 KL 散度作为蒸馏损失。相比渐进蒸馏或一致性模型，DMD 更稳定且能保持更高的图像质量。
本文在 DMD 基础上提出 K-DMD，加入了少步教师的知识转移，是实现 4 步高质量生成的关键技术。
稀疏注意力机制
稀疏注意力通过限制每个 token 的注意力范围来降低自注意力的 O(n²) 复杂度。常见方式包括局部窗口注意力（Neighborhood Attention）、块注意力（Blockwise Attention）和 KV 压缩（Key-Value Compression）。本文的 ASSA 将 KV 压缩和块邻域注意力组合成自适应的全局-局部注意力。
注意力机制的计算开销是 DiT 在高分辨率图像上部署的主要瓶颈，本文的核心创新就是设计高效的稀疏注意力。
弹性网络（Elastic Network）
弹性网络源自 Slimmable Networks 和 Once-for-All 的思想，通过在一个超网络中同时训练多个宽度不同的子网络，实现单一模型适配不同硬件。子网络共享主干参数，通过切片隐藏维度来获得不同容量的模型。训练时需要稳定的梯度聚合策略来避免子网络之间的干扰。
本文将这一思想首次应用于生成模型，使得同一个 DiT 能够部署在从低端手机到高端服务器的不同设备上。
研究动机
当前最先进的图像生成模型（如 Flux 12B、Qwen-Image 20B）基于 Diffusion Transformer 架构，在生成质量和编辑灵活性上远超传统 U-Net 模型，但这些模型动辄数十亿参数，需要服务器级 GPU 和定制 CUDA 内核才能运行。在移动端部署场景下，高分辨率（如 1024×1024）的自注意力操作需要处理 4096 个 token，导致 O(n²) 的计算复杂度常常引发内存溢出（OOM）错误。现有的移动端解决方案（如 SnapGen、MobileDiffusion）虽然能在手机上生成 512 像素或 1024 像素图像，但它们都基于 U-Net 架构，在生成质量和可扩展性上远远落后于 DiT 模型。此外，移动设备硬件异构性极大——从入门级安卓机到高端 iPhone，单一静态模型无法在所有设备上高效运行，导致开发碎片化。
本文的目标是本文的目标是设计一个高效 DiT 框架，在严格的计算和内存约束下实现 Transformer 级别的图像生成质量。具体而言：（1）在 iPhone 16 Pro Max 上以 1.8 秒内生成 1024×1024 高保真图像；（2）提供 0.3B、0.4B 和 1.6B 三个规格的模型，分别适配低端安卓、高端手机和服务器；（3）通过 4 步蒸馏实现接近实时的生成速度；（4）在主流 T2I 基准上达到与数倍参数量模型可比的性能。
与已有工作不同的是，本文的独特切入点在于「三个缺一不可的组件协同设计」。之前的工作要么只关注架构优化（如 SANA 的线性注意力）、要么只关注训练压缩（如 SD3.5-Flash 的蒸馏）、要么只关注特定硬件（如 SnapGen 的 U-Net），没有一个统一框架同时解决架构效率、多设备适配和快速推理三个问题。本文首次将自适应稀疏注意力、弹性训练和知识引导蒸馏三者结合，形成一个完整的端到端解决方案，填补了 DiT 在移动端部署的空白。
核心方法
本文的方法可以类比为「三层包装」。想象你要把一个大型餐厅（1.6B DiT）搬到手机上：第一层是「缩小厨房」——通过三阶段架构和自适应稀疏注意力降低每道菜的烹饪时间；第二层是「弹性菜单」——同一个厨房可以做出不同规格的菜，适配不同大小的餐桌（设备）；第三层是「师傅带徒弟」——通过知识蒸馏让小模型快速学会大模型的做菜技巧。技术路线上，整个系统以 1.6B 参数的 DiT 作为超网络，包含 Down-Middle-Up 三个阶段。Down 和 Up 阶段在高分辨率特征图上运行，使用本文提出的自适应稀疏自注意力（ASSA）替代全注意力；Middle 阶段在下采样后的特征图上运行，使用标准注意力。训练分为三个阶段：先在 256 分辨率上用 Flow Matching 预训练弹性模型，再在 1024 分辨率上进行知识蒸馏，最后用 K-DMD 进行 4 步蒸馏。
本文最核心的创新是自适应稀疏自注意力（Adaptive Sparse Self-Attention, ASSA）。它的本质思想是：在高分辨率特征图上，不需要每个 token 都与所有 token 计算注意力——全局结构只需要粗粒度的信息，局部细节只需要关注邻近区域。具体来说，ASSA 将注意力计算拆分为两个并行分支：（1）全局注意力分支通过对 key 和 value 进行 2×2 步长卷积压缩 token 数量（从 4096 压缩到 1024），让每个 query 能以极低代价获取全局上下文；（2）局部注意力分支将 token 网格分成 B 个不重叠的空间块，在每个块的邻域内独立计算注意力，保留精细的空间关系。最关键的是，这两个分支的融合权重是逐头、输入自适应的——通过一个 sigmoid 门控函数根据当前输入动态决定每个注意力头应该更依赖全局还是局部信息。这使得模型在处理天空等大面积区域时倾向于使用全局注意力，在处理纹理细节时倾向于使用局部注意力。
方法步骤详情
完整的方法分为以下步骤：（1）架构设计阶段：构建三阶段 DiT，Down 和 Up 阶段各使用 ASSA 层，Middle 阶段使用标准自注意力。Down 阶段包含 6 层，Middle 阶段 2 层，Up 阶段 10 层（比 Down 多以获得更好的上采样质量）。中间通过 2×2 窗口下采样和上采样连接。在 ASSA 中，key 和 value 首先经过 Conv2×2 步长卷积压缩到 1/4 大小用于全局注意力，同时原始 key/value 分成 B=16 个空间块用于局部注意力。两种注意力的结果通过输入自适应的权重 sigmoid(W·h) 线性插值融合。（2）弹性训练阶段：在每个训练迭代中，从 1.6B 超网络中均匀采样 0.3B 和 0.4B 子网络，三者同时计算 Flow Matching 损失。梯度通过自适应缩放聚合，确保不同宽度子网络的更新平衡。此外，每个子网络还接收来自超网络输出的蒸馏损失以稳定训练。（3）知识蒸馏阶段：使用 Qwen-Image（20B）作为云端教师，通过输出级速度预测匹配和特征级表征对齐进行蒸馏。损失函数结合 Flow Matching 损失、教师输出损失和特征损失，并采用时间步感知的缩放策略。（4）K-DMD 步蒸馏阶段：使用 Qwen-Image-Lightening 作为少步教师，结合 DMD 目标（通过 critic 模型估计学生输出分布与教师分布的 KL 散度）和少步教师的输出/特征损失，进行 10K 步迭代训练，实现 4 步高质量生成。
技术新颖性
本文的技术新颖性体现在三个层面。首先，ASSA 是首个将 KV 压缩与块邻域注意力结合并实现输入自适应融合的注意力机制。与 SANA 的线性注意力不同，ASSA 保留了 Transformer 的表达能力；与 PixArt-Σ 的固定 KV 压缩不同，ASSA 的融合权重是动态的；与传统的窗口注意力不同，ASSA 通过块邻域设计支持灵活的感受野大小。其次，弹性 DiT 框架首次将 Once-for-All 思想应用于生成模型，通过沿隐藏维度切片投影矩阵实现参数共享，同时隔离了 LayerNorm 和调制层等与宽度绑定的参数。这种设计使得 0.4B 子网络的训练几乎「免费」——它复用 1.6B 模型的参数，不增加额外训练成本。最后，K-DMD 在 DMD 基础上引入少步教师的知识转移，解决了 DMD 在小模型上需要精细调参且容易不收敛的问题，通过 LoRA 实现零额外内存开销的少步教师集成。
实验结果
本文在多个维度展示了全面的实验验证。在 DPG-Bench 上，0.4B 小模型达到 85.2 分，超过了 2B 参数的 SD3-Medium（85.1）、12B 参数的 Flux.1-dev（83.8）和 1.6B 的 SANA（84.8）；1.6B 全模型更是达到 87.2 分，接近 20B 的 Qwen-Image（88.3）。在 GenEval 评估中，小模型达到 0.70，与 8.1B 的 SD3.5-Large（0.71）相当；全模型达到 0.76，大幅超越所有同规模模型。在 T2I-CompBench 上，小模型整体得分 0.506，全模型 0.536，均超过大多数基线。在延迟方面，0.3B tiny 模型在 iPhone 16 Pro Max 上单步 280ms，4 步生成仅 1.2 秒；0.4B small 模型单步 360ms，4 步生成 1.8 秒；量化后的 1.6B full 模型单步 1580ms，4 步生成 6.7 秒。人类评估（Parti prompts 用户研究）显示，全模型在保真度和写实性上全面超越所有基线，在图文对齐上与 SD3-Medium 竞争激烈。4 步蒸馏后的模型虽然在 DPG 和 GenEval 分数上有轻微下降（如小模型从 85.2/0.70 降至 82.7/0.69），但视觉质量几乎无损，证明了 K-DMD 的有效性。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| DPG-Bench | Overall Score | 85.2 (0.4B) | Flux.1-dev (12B): 83.8 | +1.4 分，参数量仅 1/30 |
| DPG-Bench | Overall Score | 87.2 (1.6B) | SD3.5-Large (8.1B): 85.6 | +1.6 分，参数量仅 1/5 |
| GenEval | Overall Score | 0.70 (0.4B) | SD3-Medium (2B): 0.62 | +0.08，参数量仅 1/5 |
| GenEval | Overall Score | 0.76 (1.6B) | Flux.1-dev (12B): 0.66 | +0.10，参数量仅 1/7.5 |
| T2I-CompBench | Overall Score | 0.506 (0.4B) | SANA (1.6B): 0.476 | +0.030，参数量仅 1/4 |
| CLIP Score (COCO) | CLIP-S | 0.332 (0.4B) | SnapGen (0.4B): 0.332 | 持平，架构从 U-Net 升级为 DiT |
| iPhone 延迟 | 4-step Gen. Time | 1.8s (0.4B) | SnapGen (0.4B): 约 2s | 更快且质量更高 |
局限与改进
本文存在以下几个值得关注的局限性。首先，论文主要在 ImageNet-1K 上进行消融实验，而最终的 T2I 模型使用的是闭源的 Qwen-Image 作为教师模型，这意味着完整的训练流程无法被独立复现——读者无法获得相同的教师模型权重。其次，虽然论文声称模型可以在低端安卓设备上运行，但所有延迟数据均在 iPhone 16 Pro Max 上测量，缺乏安卓设备的实测数据。第三，K-DMD 蒸馏后的 4 步模型虽然在 DPG 和 GenEval 分数上与 28 步模型接近，但论文未展示在复杂组合提示（如涉及多物体空间关系的提示）上的具体表现差异。第四，弹性训练虽然节省了参数存储，但训练时仍需同时计算三个子网络的前向和反向传播，实际训练成本并未显著降低。最后，论文缺乏与 SANA-Sprint 等最新的步蒸馏方法的直接对比，使得 K-DMD 的相对优势不够清晰。
独立分析的弱点
从独立分析的角度，本文存在以下弱点及其改进方向。第一，ASSA 的块邻域注意力虽然降低了计算复杂度，但块之间的边界效应可能导致相邻区域的语义不一致。在高分辨率图像中，物体可能横跨两个注意力块的边界，此时局部注意力无法捕捉完整的物体结构。改进方向可以考虑引入重叠块或跨块信息交换机制，例如使用滑动窗口而非固定分块。第二，弹性训练中的子网络采样策略是均匀的，但实际上不同设备的使用频率和重要性不同。改进方向是根据目标设备的市场份额或推理负载动态调整采样概率，实现「按需训练」。第三，模型使用 TinyCLIP 和 Gemma3-4b-it 两个文本编码器的拼接作为条件输入，这增加了部署复杂度和推理延迟。改进方向是探索单一更高效的文本编码器，或对现有双编码器进行联合蒸馏压缩。第四，论文中提到的 0.3B tiny 模型虽然延迟最低（280ms/步），但在 GenEval 的颜色属性上（0.85）明显弱于小模型（0.89），说明过度压缩会导致细节生成能力下降。改进方向可以探索混合精度策略，对关键层保留更高精度。
未来方向
基于本文的成果，未来可以从以下几个方向延伸。首先，当前的弹性框架只支持宽度维度的缩放，未来可以扩展到深度维度（层数）和分辨率维度的联合缩放，实现更细粒度的资源适配。其次，ASSA 的设计可以推广到视频生成领域——视频 DiT 面临的时空注意力复杂度问题更加严重，自适应稀疏注意力的思想可以直接应用。第三，K-DMD 的少步教师集成方式（通过 LoRA）可以进一步探索为「教师池」机制，维护多个不同特化的少步教师，根据输入内容动态选择最合适的教师进行指导。第四，当前模型只支持文本到图像生成，未来可以将弹性 DiT 框架扩展到图像编辑、图像修复、可控生成等任务，验证架构的通用性。最后，论文展示了 4-bit 量化的 1.6B 模型可以在手机上运行，未来可以探索更低比特（如 2-bit）的量化方案，结合本文的弹性框架实现更极致的压缩。
复现评估
本文的复现存在较大挑战。从开源情况看，论文提供了项目页面（https://snap-research.github.io/snapgenplusplus/2026），但未明确说明代码和模型权重是否会开源。训练数据方面，模型在大规模内部数据集上训练，使用了 Qwen-Image（20B）作为教师模型，这些都是闭源资源。算力需求上，论文使用 32 个节点、每节点 8 张 A100 GPU（80GB）进行训练，这是普通实验室难以承受的。弹性预训练阶段 400K 迭代加上知识蒸馏阶段 100K 迭代，再加 K-DMD 的 10K 迭代，总计约 510K 迭代的计算量。对于希望复现核心思想（如 ASSA 机制）的研究者，可以在较小规模（如 ImageNet-1K 256×256）上验证消融实验，这部分的计算需求相对可控。但如果想复现完整的 T2I 系统，几乎必须依赖 Snap 的内部资源。建议关注项目的开源动态，如果代码发布，可以在其基础上进行微调或改进。
论文图表
