終於等來能塞進手機的文生圖模型！十分之一體量，SnapGen實現百分百的效果
机器之心發表於2024-12-25
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/
作者們對去噪 UNet 和影像解碼器（AE decoder）進行了全面最佳化，從而獲得資源使用和效能之間的最佳權衡。與以往專注於預訓練擴散模型的最佳化和壓縮的研究不同，SnapGen 從整體架構和微觀運算元設計同時入手，提出了一種高效的模型結構，在顯著降低模型引數和計算複雜度的同時，仍能保持高質量的生成效果。多級知識蒸餾（Multi-level Knowledge Distillation）為了對齊最先進的文生圖模型 SD3.5-Large，SnapGen 使用 Rectified Flows 目標進行訓練，從而可以直接使用 SD3.5 系列作為知識蒸餾的教師模型。與已有的工作在相同架構型別下進行蒸餾不同，SnapGen 使用 DiT 教師模型跨架構蒸餾 UNet 學生模型，並且提出了一種先進的多級知識蒸餾框架，分別在輸出和特徵維度進行教師與學生模型的對齊。為了解決不同時間步上蒸餾損失函式尺度不一的問題，作者們提出了時間步感知的縮放（timestep-aware scaling）操作，這種操作顯著加速了知識蒸餾的收斂並增強了 SnapGen 學生模型的生成能力。為了進一步減少模型的推理時間，作者們考慮使用了一種基於 LADD 的少步數蒸餾。在步數蒸餾演算法中，可以直接進行 4 步推理的 SD3.5-Large-Turbo 被用來作為教師模型和判別器的特徵提取。蒸餾過後的模型具有和 28 步相當的 4/8 步生成能力。和 SOTA 模型 SDXL, SD3, SD3.5, PixArt-α 等相比，SnapGen 引數量最小，也是唯一可以直接部署到移動端的模型，同時有著接近或者更強的高畫素影像生成能力。下面的影片更加直觀地展示 SnapGen 在移動端裝置上的文生圖效率與質量，在 iPhone 16 Pro Max 上僅需 1.4s 就可以生成超高質量的 1024x1024 圖片。
這裡展示了更多的 1024x1024 影像生成結果，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 作為小尺寸高效率模型走出了一條不一樣的路，也相信會給生成模型的研究帶來啟發。