本文的共同一作為墨爾本大學的胡冬庭和香港科技大學的陳捷潤和黃悉偈，完成于在 Snap 研究院 Creative Vision 團隊實習期間。主要指導老師為任健、徐炎武和 Anil Kag，他們均來自 Snap Creative Vision 團隊。該團隊的主要研究方向包括 Efficient AI 和圖像/視頻/三維生成模型。

近些年來，以 Stable Diffusion 為代表的擴散模型為文生圖（T2I）任務樹立了新的標準，PixArt，LUMINA，Hunyuan-DiT 以及 Sana 等工作進一步提高了圖像生成的質(zhì)量和效率。然而，目前的這些文生圖（T2I）擴散模型受限于模型尺寸和運行時間，仍然很難直接部署到移動設備上。

盡管以量化 / 剪枝為代表的模型壓縮技術可以解決一部分問題，但直接從頭訓練一個輕量化可以部署在移動設備上的高效高質(zhì)文生圖模型仍然是巨大的挑戰(zhàn)。

最近，來自 Snap 研究院的 Creative Vision 研究團隊提出了 SnapGen，從頭訓練了一個僅有 379M 參數(shù)的文生圖模型，并且在 iPhone 16 Pro-Max 上僅需 1.4s 就可以生成超高質(zhì)量的 1024x1024 圖片。

和 SOTA 模型 SDXL, SD3, SD3.5, PixArt-α 等相比，SnapGen 有著同等或更好的指令跟隨能力以及圖像生成質(zhì)感。

在多個定量測試基準和人類偏好測試中，SnapGen 同樣顯著超過了擁有更多參數(shù)量的模型，在 GenEval 指標上達到 0.66，并且在美學和文字 - 圖像一致性等方面接近 SD3-Medium 以及 SD3.5-Large。在少步數(shù)生成的情景下， SnapGen 也同時保持了其生成質(zhì)量， 在 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）進行了全面優(yōu)化，從而獲得資源使用和性能之間的最佳權衡。

與以往專注于預訓練擴散模型的優(yōu)化和壓縮的研究不同，SnapGen 從整體架構和微觀算子設計同時入手，提出了一種高效的模型結構，在顯著降低模型參數(shù)和計算復雜度的同時，仍能保持高質(zhì)量的生成效果。

多級知識蒸餾（Multi-level Knowledge Distillation）

為了對齊最先進的文生圖模型 SD3.5-Large，SnapGen 使用 Rectified Flows 目標進行訓練，從而可以直接使用 SD3.5 系列作為知識蒸餾的教師模型。

與已有的工作在相同架構類型下進行蒸餾不同，SnapGen 使用 DiT 教師模型跨架構蒸餾 UNet 學生模型，并且提出了一種先進的多級知識蒸餾框架，分別在輸出和特征維度進行教師與學生模型的對齊。

為了解決不同時間步上蒸餾損失函數(shù)尺度不一的問題，作者們提出了時間步感知的縮放（timestep-aware scaling）操作，這種操作顯著加速了知識蒸餾的收斂并增強了 SnapGen 學生模型的生成能力。

步數(shù)蒸餾（Step Distillation）

為了進一步減少模型的推理時間，作者們考慮使用了一種基于 LADD 的少步數(shù)蒸餾。在步數(shù)蒸餾算法中，可以直接進行 4 步推理的 SD3.5-Large-Turbo 被用來作為教師模型和判別器的特征提取。蒸餾過后的模型具有和 28 步相當?shù)?4/8 步生成能力。

SnapGen 可視化

和 SOTA 模型 SDXL, SD3, SD3.5, PixArt-α 等相比，SnapGen 參數(shù)量最小，也是唯一可以直接部署到移動端的模型，同時有著接近或者更強的高像素圖像生成能力。

下面的視頻更加直觀地展示 SnapGen 在移動端設備上的文生圖效率與質(zhì)量，在 iPhone 16 Pro Max 上僅需 1.4s 就可以生成超高質(zhì)量的 1024x1024 圖片。

這里展示了更多的 1024x1024 圖像生成結果，SnapGen 在具有挑戰(zhàn)性的文字，人物肢體，特定風格和概念生成中均表現(xiàn)出色。

SnapGen 方法細節(jié)

高效的網(wǎng)絡結構

在去噪模型結構的選擇上，擴散模型主要分為 UNet 和純 Transformer 兩大流派。

盡管純 Transformer 架構（如 DiT）在大規(guī)模數(shù)據(jù)和算力支持下展現(xiàn)了廣闊前景，亞馬遜 AWS AI Lab 的一項研究（On the Scalability of Diffusion-based Text-to-Image Generation）表明，UNet 架構（尤其是 SDXL 架構）在相同參數(shù)量下表現(xiàn)出更高的性能、更低的算力需求以及更快的收斂速度。

基于此，Snap 團隊調(diào)整 SDXL 中 UNet 架構的深度和寬度，并探索了如下圖所示的一系列架構優(yōu)化，包括移除高分辨率自注意力（SA）層、使用寬度擴展后的深度可分卷積（SepConv）替代常規(guī)卷積（Conv）、降低全連接層（FFN）的中間通道維度、更早注入文字等條件信息，以及優(yōu)化自注意力（SA）與交叉注意力（CA）算子（如將多頭自注意力 MHSA 替換為多查詢注意力 MQA、對查詢和鍵值應用 RMSNorm 歸一化、插入旋轉(zhuǎn)式位置編碼 RoPE）。

雖然部分方法已在其他工作（如谷歌的 MobileDiffusion 模型）中有所提及，但這些研究往往缺乏對改動前后模型性能的全面量化評估。

相較之下，SnapGen 在 ImageNet-1K 256 像素類條件圖像生成任務中，通過生成質(zhì)量指標（FID）、模型時延、計算量和參數(shù)規(guī)模的綜合評估，驗證了每項架構改動的合理性和有效性。

最終，SnapGen 在生成質(zhì)量（FID 2.06）與現(xiàn)有模型（如 SiT-XL）相當?shù)那闆r下，大幅降低了模型大小和計算量，展現(xiàn)出卓越的性能和資源效率。

除了去噪模型，圖像解碼器同樣是一個重要的優(yōu)化對象。

首先，相較于整體生成時間，圖像解碼器的推理時間不容忽視，尤其是在少步甚至單步去噪模型的情況下。此外，在部署到移動端生成高分辨率圖像時，解碼器常常會遇到顯存不足的報錯。

Snap 團隊發(fā)現(xiàn)，現(xiàn)有的 SD3 圖像解碼器存在大量的參數(shù)和計算冗余。這主要是由于其潛在空間采用了 16 通道，而與 SDXL 使用的 4 通道相比，16 通道更容易實現(xiàn)圖像重建，因此其網(wǎng)絡結構在壓縮和加速方面具有更大的潛力。

為了優(yōu)化這一點，Snap 團隊通過移除不必要的自注意力機制和冗余的 GroupNorm 歸一化層，同時減小網(wǎng)絡寬度等方式，成功實現(xiàn)了近乎無損的 36 倍參數(shù)壓縮，并在移動端部署中實現(xiàn)了 54 倍的解碼加速。

高效的訓練以及高級知識蒸餾

SnapGen 采用 Rectified Flows 為目標優(yōu)化模型訓練， 與 SD3 和 SD3.5 等較大的模型保持一致。同時 SnapGen 利用多個文本編碼器 (text encoders) 包括 CLIP 以及 Gemma2-2b， 在訓練中使用 classifier-free guidanc 以實現(xiàn)不同硬件環(huán)境下的部署需求。

基礎模型在從初始訓練的情況下在 GenEval 上的表現(xiàn)為 0.61。得益于使用相同的訓練目標，SnapGen 可以將最新的 SD3.5-Large 作為知識蒸餾的教師模型。

然而在知識蒸餾過程中，仍然有很多需要解決的挑戰(zhàn)：教師模型（DiT）和學生模型（UNet）的異構性，蒸餾損失函數(shù)和 Rectified Flows 任務損失函數(shù)的尺度不一致，以及常常被研究人員忽視的不同時間步上去噪預測難度的差異。

為了解決上述的問題，Snap 團隊提出了一種新穎的多級別知識蒸餾范式，并且進行了時間步感知的損失函數(shù)尺度縮放。在任務損失函數(shù)之外，SnapGen 的訓練還使用了輸出蒸餾損失函數(shù)與特征蒸餾損失函數(shù)。

和之前使用知識蒸餾的工作（LinFusion，BK-SDM）不同，SnapGen 不需要預設不同損失函數(shù)的尺度，而是根據(jù)不同時間步上的統(tǒng)計數(shù)據(jù)將這些損失函數(shù)縮放到同一個尺度，保證每部分對訓練的貢獻均等，這種操作也被作者稱為時間步感知尺度縮放（timestep-aware scaling）。

實驗表明，這種考慮時間步變化的尺度縮放可以有效加速訓練，并且知識蒸餾后的模型在 GenEval 上的表現(xiàn)高達 0.66。

步數(shù)蒸餾提升推理速度

作者們基于 LADD 等 diffusion-GAN 混合結構對 SnapGen 進行步數(shù)蒸餾。使用 4 步模型 SD3.5-Large-Turbo 同時作為教師模型和判別器的特征提取器。

SnapGen 可以快速適應少步推理的需要，和 28 步的基礎模型相比，4 步與 8 步的推理結果在視覺效果上沒有顯著差別，GenEval 上也僅僅只有 0.05/0.03 的差距。

實驗結果

作者們在多個測試基準（GenEval，DPG-Bench，CLIP Score on COCO，ImageReward）上定量對比了 SnapGen 和大量現(xiàn)有的高分辨率文生圖模型，涵蓋了 PixArt 系列，Stable Diffusion 系列，Sana，LUMINA-Next，Playgroundv2/v2.5，IF-XL。盡管參數(shù)量最小且吞吐量最高，SnapGen 仍然在所有指標中均排在前列。

在人類偏好測試中，和 SD3-Medium，SD3.5-Large 相比，SnapGen 生成的圖像具有更真實的質(zhì)感，并且在美學和文字 - 圖像一致性等方面接近；同時 SnapGen 在所有評測指標中均顯著超越 SDXL。

經(jīng)過步數(shù)蒸餾的少步模型同樣具有優(yōu)秀的高效文生圖能力，對于基準模型而言，4/8 步的生成結果通常會比較模糊或者確實重要細節(jié)。

而 SnapGen 的 4/8 步生成仍然能保持接近 28 步的視覺效果，在 GenEval 等定量基準上也非常接近 28 步的 baseline。

通過設計高效的去噪模型架構，使用先進的知識蒸餾和少步數(shù)蒸餾算法，Snap 團隊提出了能直接部署到手機上的 SnapGen 模型。SnapGen 僅有 379M 參數(shù)，僅用 1.4s 就可以在 iPhone 上生成 1024x1024 圖片，卻在生成質(zhì)量方面超出大部分現(xiàn)有模型。

在文生圖模型隨著 scaling law 越來越大的今天，SnapGen 作為小尺寸高效率模型走出了一條不一樣的路，也相信會給生成模型的研究帶來啟發(fā)。