首页 时事 民生 政务 教育 文化 科技 财富 体娱 健康 情感
更多
旅行 百科 职场 楼市 企业 乐活 学术 汽车 时尚 创业 美食 幽默 美体 文摘

终于等来能塞进手机的文生图模型！十分之一体量，SnapGen实现百分百的效果

科技   2024-12-27 21:01   江苏  

点击蓝字


关注我们

关注并星标

从此不迷路

芒果智能科技

 

  • 论文标题：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/ 


芒果智能科技专栏

Column of Computer Vision Institute

本文的共同一作为墨尔本大学的胡冬庭和香港科技大学的陈捷润和黄悉偈，完成于在 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 方法简介
高效的模型结构
作者们对去噪 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 作为小尺寸高效率模型走出了一条不一样的路，也相信会给生成模型的研究带来启发。

END



转载请联系本公众号获得授权


计算机视觉研究院学习群等你加入！


ABOUT

计算机视觉研究院

计算机视觉研究院主要涉及深度学习领域，主要致力于目标检测、目标跟踪、图像分割、OCR、模型量化、模型部署等研究方向。研究院每日分享最新的论文算法新框架，提供论文一键下载，并分享实战项目。研究院主要着重”技术研究“和“实践落地”。研究院会针对不同领域分享实践过程，让大家真正体会摆脱理论的真实场景，培养爱动手编程爱动脑思考的习惯！


 往期推荐 

🔗
计算机视觉研究院
计算机视觉研究院主要涉及AI研究和落地实践，主要致力于目标检测、目标跟踪、图像分割、OCR、模型量化、模型部署等研究方向。研究院每日分享最新的论文算法新框架，提供论文一键下载，并分享实战项目。研究院主要着重”技术研究“和“实践落地”！
 最新文章
被导师放养，后果可能很严重。。。
9h打通顶会壁垒！实现1天内复现任意顶会文章！
利用先进技术保家卫国：深度学习进行小目标检测（适合初学者入门）
首个全量化Vision Transformer的方法FQ-ViT，AI大模型落地不远了！
代码开源 | 基于YOLO的16种目标检测项目，人体/交通/医疗/工业等方向超全
无人机小目标检测的轻量级、精确的类YOLO网络
视觉算法岗工资正式进入4开头时代？！
昼夜无监督域自适应目标检测（附原代码）
无人机+AI：光伏巡检自动化解决方案
「无人机+AI」总结成为智能巡检好帮手
60k感知算法岗，考察风向变了？！
超越所有视觉方案！CVT-Occ：成本体积时间融合问鼎OCC！
无人机+AI：风机全自动巡检解决方案
性能提升：Facebook提出FP-NAS——搜索速度更快、分类精度更高、性能更好
即插即用，轻松涨点！把大牛的模块缝合到自己的paper里
「源神」稚晖君又双叒叕开源，这一次机器人直接进入人类生活！
GitHub 7.5k star量，各种视觉Transformer的PyTorch实现合集整理好了
优于FCOS：在One-Stage和Anchor-Free目标检测中以最小的成本实现最小的错位（代码待开源）
发SCI一般奖励多少钱，听说重庆大学60万？
用Kaggle比赛项目金牌发了顶会！
GPT理解的CV：基于Yolov5的半监督目标检测
首个全量化Vision Transformer的方法FQ-ViT，AI大模型落地不远了！
2024新技术：远距离的小目标也可以准确检测
Drone-YOLO：一种有效的无人机图像目标检测
2PCNet：昼夜无监督域自适应目标检测（附原代码）
CRAS-YOLO：多类别船舶检测与分类模型
解锁自动驾驶与具身智能技术栈！各大优质社区汇总
YOLO-S：小目标检测的轻量级、精确的类YOLO网络
旋转角度目标检测的重要性！！！（附源论文下载）
大改Yolo框架 | 能源消耗极低的目标检测新框架（附论文下载）
智慧建筑：基于YOLOv7的建筑外墙缺陷检测
2PCNet：昼夜无监督域自适应目标检测（附原代码）
CRAS-YOLO：多类别船舶检测与分类模型
解锁自动驾驶与具身智能技术栈！各大优质社区汇总
终于等来能塞进手机的文生图模型！十分之一体量，SnapGen实现百分百的效果
2024全球无人机市场洞察报告
YoloV：视频中目标实时检测依然很棒（附源代码下载）
YOLO-SLD：一种改进的注意机制YOLO车牌检测
Consistent-Teacher：半监督目标检测超强SOTA
一种有效的无人机图像目标检测（Drone-YOLO）
摸鱼一年半，摸出了一篇顶会论文！
突破数据瓶颈！交大研发电脑智能体，让 AI 替你熬夜做 PPT
YOLO跌落神坛？Mamba YOLO干翻YOLO系列模型！
微调时无需泄露数据或权重，这篇AAAI 2025论文提出的ScaleOT竟能保护隐私
YOLO-SLD：一种改进的注意机制YOLO车牌检测
“计算机视觉研究院”商务合作
ResNet超强变体：京东AI新开源的计算机视觉模块！（附源代码）
自己觉得挺有意思的目标检测框架，分享给大家（源码论文都有）
更快、更灵活的Transformer图像去雾网络
改进的检测算法：用于高分辨率光学遥感图像目标检测
 分类
时事
民生
政务
教育
文化
科技
财富
体娱
健康
情感
旅行
百科
职场
楼市
企业
乐活
学术
汽车
时尚
创业
美食
幽默
美体
文摘
 原创标签
时事 社会 财经 军事 教育 体育 科技 汽车 科学 房产 搞笑 综艺 明星 音乐 动漫 游戏 时尚 健康 旅游 美食 生活 摄影 宠物 职场 育儿 情感 小说 曲艺 文化 历史 三农 文学 娱乐 电影 视频 图片 新闻 宗教 电视剧 纪录片 广告创意 壁纸头像 心灵鸡汤 星座命理 教育培训 艺术文化 金融财经 健康医疗 美妆时尚 餐饮美食 母婴育儿 社会新闻 工业农业 时事政治 星座占卜 幽默笑话 独立短篇 连载作品 文化历史 科技互联网
 发布位置
广东 北京 山东 江苏 河南 浙江 山西 福建 河北 上海 四川 陕西 湖南 安徽 湖北 内蒙古 江西 云南 广西 甘肃 辽宁 黑龙江 贵州 新疆 重庆 吉林 天津 海南 青海 宁夏 西藏 香港 澳门 台湾 美国 加拿大 澳大利亚 日本 新加坡 英国 西班牙 新西兰 韩国 泰国 法国 德国 意大利 缅甸 菲律宾 马来西亚 越南 荷兰 柬埔寨 俄罗斯 巴西 智利 卢森堡 芬兰 瑞典 比利时 瑞士 土耳其 斐济 挪威 朝鲜 尼日利亚 阿根廷 匈牙利 爱尔兰 印度 老挝 葡萄牙 乌克兰 印度尼西亚 哈萨克斯坦 塔吉克斯坦 希腊 南非 蒙古 奥地利 肯尼亚 加纳 丹麦 津巴布韦 埃及 坦桑尼亚 捷克 阿联酋 安哥拉