字节开源Hybrid SD，端云协同推理实现文生图端侧部署

最新推荐文章于 2024-11-03 06:26:00 发布

PaperWeekly

最新推荐文章于 2024-11-03 06:26:00 发布

阅读量198

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzIwMTc4ODE0Mw==&mid=2247686091&idx=2&sn=a6af0ee14be3aad50285ffb5bd3da58b&chksm=9779054fa7fb0f9e17acc0217af3daf22e3fdd96c18177f47f4d3e65360fe2299d0fc3470fce&scene=126&sessionid=0

版权

©PaperWeekly 原创 · 作者 | 刘松伟

单位 | 字节跳动

研究方向 | 大模型推理、模型优化

Hybrid SD 是面向扩散模型（Stable Diffusion Models, SDMs）的端云协同混合推理工作，主要解决 SDMs 推理中存在的效果与成本的挑战：1）大模型在语义理解方面具备强大能力，但是计算开销大，难以在端侧部署；2）轻量化的小模型可以部署在端侧，但是效果较大模型有较大差距。

论文标题：

Hybrid SD: Edge-Cloud Collaborative Inference for Stable Diffusion Models

文章地址：

https://arxiv.org/pdf/2408.06646

针对这两个挑战，在推理框架层面，我们推出 Hybrid SD 端云协同混合推理框架，巧妙地融合了大型模型的强大语义理解能力与小型模型的高效重建优势，实现了效果与成本的权衡。在端侧模型层面，我们通过模型优化产出了可以端侧部署的轻量化 U-Net 和 VAE，有效降低端侧设备的模型尺寸，降低了总的推理计算量。

▲ 图 1: Hybrid SD端云协作混合推理。在 SD Models（25步）推理中，Hybrid SD 在保证推理效果的前提下大幅降低了推理成本（FLOPs）。在加速后的一致性模型（Latent Consistency Models, LCM）中，Hybrid SD 取得了进一步的效果收益，在 FID 指标上，Ours-Tiny-LCM（k=4）相比于单独的小模型推理 Ours-Tiny-LCM（k=0）从 23.42 下降到 16.19。同时将 VAE 替换为我们的轻量化 VAE 后，计算量 FLOPs 进一步从 12.8 下降到 8.1，降低了在端侧推理的成本。

引言

近年来，扩散模型（Stable Diffusion Models, SDMs）在文生图任务中表现出了卓越的效果，这些模型能够精准地解析用户输入的文本描述，并生成与之高度匹配的图像。不仅极大地推动了以人工智能为驱动力的内容生成技术（AIGC）的发展，也为数字艺术、游戏设计、影视制作等多个领域带来了革命性的变化。

然而，受限于其庞大的模型规模及对计算资源的海量需求，当前大多数 SDMs 只能部署在昂贵的云端服务器上。但随着生成式 AI 模型使用量和复杂性的不断增长，仅在云端进行推理并不划算。因为数据中心基础设施成本，包括硬件、场地、能耗、运营、额外带宽和网络传输的成本将持续增加。

基于 AWS 的公开数据，部署一个 SDXL 模型，假设每天推理时长为 8 小时，每个月工作 20 天，共生成 72000 张图像，预计需要每月花费 310.45 美元。除了成本的开销之外，云端的模型推理对我们赖以生存的环境的影响也不容小觑，论文测量了 SD1.5 单次推理所消耗的能量为 1.38×10^-3kWh（约等于一辆电动汽车行驶 8 米）。

如表 2 和图 2 所示，FU1 和 FU2 分别表示用户视角和服务视角测量了 SD 服务产生的环境影响，推理在整个生命周期里占了约 50% 的 GWP 和 PE，即一年排放了 180 吨二氧化碳（相当于一个人生活 30 年的碳排），消耗 1.24 GW 的能量（相当于燃烧 106.6 吨油的能量）。

▲ ADP：矿物和金属消耗；GWP：碳排放；PE：能耗

一种减轻云基础设施的压力和开支的方式是利用终端设备执行推理，终端设备能够以很低的能耗运行 SD 模型，同时帮助云服务提供商降低数据中心的能耗，实现环境和可持续发展目标。尽管市场上已出现一些针对边缘设备的轻量级 SDMs 模型，它们在减少计算负担的同时，往往不得不在画面质量、语义精确度方面做出妥协。

为了解决这一问题，我们提出了一种创新的 Hybrid SD 端云协同推理框架，巧妙地融合了大型模型的强大语义理解能力与小型模型的高效重建优势。该框架利用了去噪过程中各阶段的信息处理特性，通过分阶段的方式，确保在不同阶段的任务分配中，既保留了大模型在语义信息理解上的深度和广度，又发挥了小型模型在图像重建方面的灵活性和效率。

进一步地，考虑到终端设备计算能力与存储空间的多样性，我们针对 SDMs 中的关键模块包括 U-Net 和 VAE 都进行了轻量化设计，以实现端侧模型的轻量化，从而在不牺牲生成质量的前提下，使 SDMs 模型更加适合在资源受限的环境中运行，同时降低云端部署 Hybrid SD 的成本。

我们还验证了 Hybrid SD 可以很好地兼容其他扩散模型推理加速方法，例如潜在一致性模型（Latent Consistency Models, LCM）。

本文提出的 Hybrid SD 端云协同框架，不仅为 SDMs 模型的云端-终端协作推理提供了一种新的解决方案，也为 AIGC 技术的进一步发展和应用开辟了新的可能性。通过这种框架，我们期待能够推动 AI 技术在更广泛的领域中发挥更大的作用，为用户带来更加丰富和个性化的创作体验。

方案设计

受到 T-GATE 的启发，SDMs 的去噪过程可以分成语义规划和保真度提升两个阶段。语义规划阶段通 CrossAttention 模块嵌入文本条件获得视觉语义信息，而保真度提升阶段则注重改善图像的生成细节。这表明，在去噪的早期阶段，噪声预测器在将条件信息编码到图像特征中起着关键作用；而在去噪的后期步骤中，噪声预测器主要聚焦于恢复视觉感知信息。

▲ 图1: 图源《TGATE: Cross-Attention Makes Inference Cumbersome》论文

2.1 端云协同推理

2.1.1 混合推理范式

通过整合不同大小模型的信息处理能力，我们提出了一种端云协同的混合推理范式 Hybrid SD。部署在云侧的大模型具备较强的语义信息理解能力，负责初始阶段进行语义规划的去噪过程，而部署在端侧的小模型具备足够的图像生成能力，负责接下来的保真度提升阶段。

如下图，对比了本文提出的 Hybrid SD 端云协同框架和传统推理方案的区别，其中（a）和（b）为传统方案仅使用单一大小的模型进行推理，（c）为本文提出的端云协同推理方案，结合大、小模型的能力协同推理。

▲ 图2：（a）大模型仅运行在云端（b）小模型仅运行在终端（c）大小模型端云协同推理

假设模型的去噪过程为，我们引入一个分割步数来确定大模型执行的时间步，即前步采用云侧大模型去噪，剩余的时间步使用端侧小模型来执行去噪。

我们展示了 Hybrid SD 的推理框架图，如下所示：

1. 当用户发起请求后，我们现在端侧将文本上传到云侧模型。如果是图生图任务，则将图像编码到潜在空间后，再上传到云侧模型；

2. 云侧大模型进行迭代去噪推理，从推理到时间步，此时云侧大模型已经通过其强大的语义理解能力，根据文本规划了图像的大致轮廓和布局信息；

3. Hybrid SD 将时间步的带噪样本和编码的文本信息传输给端侧设备，由端侧小模型完成后续的去噪任务和图像解码，充分利用小模型的高效性。

▲ 图3: Hybrid SD 推理系统框架

2.1.2 效率与成本分析

效率分析：我们首先分析本文提出的端云协同推理范式引入的数据传输效率问题。设传输时间为，数据大小为，网络带宽为，则有公式。

1. 对于传统的云侧推理范式，以标准模型 SD1.5 为例，一张 3×512×512的图片会通过网络传输返回端侧设备。假设图片是 8 比特的 PNG 格式，大小为 768KB，根据论文中报告的平均 WiFi 带宽为18.88 Mbps，传输需要约 0.33 秒。

2. 在我们提出的 Hybrid SD 端云协同推理中，有两个关键数据需要从云端传输到终端：潜在噪声编码和文本特征。以标准 SD 1.5 模型为例，潜在噪声编码的尺寸为 4×64×64，文本特征的尺寸为 77×768。在 FP16 精度下，累计的数据负载为 148 KB。传输 148KB 数据大约需要 0.06 秒，相比云侧推理带来的网络传输开销更小，同时这部分开销在去噪过程的整体耗时来说几乎可以忽略不计。

成本降低：根据前文的成本估算，我们进一步分析将部分的云端运算转移到终端处理可降低的成本，假设每个月上线 10 款应用，每个应用部署需要 3000 个实例（每个实例一个模型），若能将 1/2 的推理转移到用户终端处理，则每年可节约 5000 万美元（10300012*310.45/2）。

2.2 更小的模型

原始 SD1.5 的 UNet 结构如下图所示：

▲ 图4: SD1.5 UNet

基于原始 SD1.5 UNet ，可以通过蒸馏得到不同规模的轻量级块移除模型（Small 和 Tiny）。具体而言，将原始模型作为教师模型，将块移除后的模型作为学生模型。通过对齐学生模型与教师模型在不同时间步下的中间特征及输出噪声分布的差异实现对学生模型的微调。

▲ 图5: BK-SDM-Small

▲ 图6: BK-SDM-Tiny

我们基于 Small 模型（包含 12 个 ResnetBlock、9 个 SelfAttention 和 9 个 CrossAttention）进一步分析了 UNet 中不同模块的通道冗余度。对于 UNet 的三类基本组成模块，我们保持整个 block 的输出通道数不变，可剪裁通道如下图所示，对于 ResnetBlock 我们仅裁剪第一个 3x3-conv 的输出通道，对于 Self-Attention 和 Cross-Attention 模块，我们对 head 数进行剪裁。

▲ 图7: UNet 不同通道冗余度分析

我们基于权重的 L1-Norm 从这些模块中直接移除了 50% 的通道数，并对最终生成的图像进行了可视化。如下图所示，在不对模型进行重训练的情况下，我们发现：

1. 对于某些层，通道数减半并未导致生成内容质量的明显下降。比如，对 CrossAttention 模块进行参数裁剪对输出图像的影响比较小，移除第 4 个 ResnetBlock 和第 8 个 Self-Attention Block 的参数所造成的变化也很小，表明这些层内存在参数冗余，而第 1 个 ResnetBlock 和第 4 个 Self-Attention Block 参数则比较关键；

2. 观察到每个 block 对整体输出质量的贡献有所不同，表明了每个 block 需要采用不同的剪枝率。

▲ 图8: 参数移除对生成图的影响（Resnet-1 表示第 1 个 ResnetBlock）

2.2.2 通道剪枝过程

为了使剪裁过程可量化，我们基于 LD-Pruner 计算了层的重要性分数，和分别表示 UNet 原始以及剪裁参数后的潜在噪声编码集合，该分数表示了移除当前层参数对潜在噪声的均值偏移和方差变化的影响。分数越高意味着剪裁当前层对模型效果的影响越大。

如图 8 所示，通过该重要性指标我们发现了和可视化结果一致的趋势，ResnetBlock 第一层和最后一层的参数很重要，第四层参数相对来说最不重要，CrossAttention 的参数相对 ResnetBlock 和 SelfAttention 来说重要性低一些。

我们基于计算出来的重要性分数来确定不同层的剪枝比率，引入两个排序阈值和，对重要性排序大于的 block 采用 0.25 的剪枝率，对重要性排序小于的 block 采用 0.75 的剪枝率，其余则采用 0.5 的剪枝率，通过搜索的方式确定满足预期参数量的阈值。以此生成更轻量的模型。

▲ 图9: 不同模块打分

确定剪枝模型后，我们通过知识蒸馏进一步微调轻量模型，以提升小模型的性能。我们采用了 BK-SDM 的蒸馏 loss，让小模型在输出层面和特征层面与大模型对齐。总的训练 loss 如下：

其中，代表扩散 loss（计算预测噪声和真实噪声的 MSE loss），表示小模型的 output 与大模型 output 之间的蒸馏 loss，表示小模型与大模型在中间特征层面的蒸馏 loss。

2.2.3 提升VAE的性能

在端侧设备的 SD 推理中，我们发现扩散模型中的 VAE decoder 也带来了很大的计算开销。随着扩散模型采样加速的方法越来越流行（例如 LCM），扩散模型中 U-Net 的推理步数越来越小，VAE 在整个推理流程中的耗时占比越来越大。

尽管有一些开源的轻量化小模型，如 TAESD，他们的效果和 SD 的 VAE 仍然有较大差距。我们重新训练了一个轻量化的 VAE，通过 LPIPs loss 和 GAN loss，我们的轻量化 VAE 取得了与 SD VAE 接近的效果，同时参数量上远远少于 SD VAE，更有利于在端侧设备部署。

我们在 V100 GPU 上测试了在 8 步 LCM 推理场景中，通过替换轻量化 VAE 带来的时间收益，如图 10 所示，轻量化 VAE 大幅减少了推理中图像解码的耗时，推理时间从 1146 ms 下降到 749 ms。

▲ 图10: 轻量化 VAE 在 LCM 8步推理中的时间收益

实验结果

3.1 量化指标

基于我们提出的剪枝方法，生成了参数量仅为 224.5M 的轻量化模型，客观指标上对比 323.38M 的 BK-SDM-Tiny 获得更小的 FID（17.05 vs. 16.71），这表明了轻量化模型也具备了相当的图像生成能力。

我们以 SD-v1.4 作为大模型，进行大、小模型混合推理的实验，当与 BK-SDM-Small、BK-SDM-Tiny 以及 OursTiny-224 结合时，在 FID 和 IS 指标上对比单一小模型推理展现出了更优的性能。

相比于 FID 为 17.05 的 BK-SDM-Tiny，混合 SD-v1.4 + OursTiny（k=10）将 FID 减少到 13.75，并在 CLIP 分数上与 SD-v1.4 相当，表明了生成图像与文本提示之间语义对齐的能力。

我们同样验证了混合推理在加速模型上的效果。我们训练了 SD-v1.4 和我们的小模型 OursTiny 的 LCM 加速模型。在推理时，在 8 步推理的实验中，混合推理在所有指标上相对单一小模型推理取得了更好的性能，将 FID 降低到 16.19，同时 IS 和 CLIP 分别提升到 31.76 和 0.2698。

相比大模型的推理，混合推理取得了相当的 FID，同时计算量 FLOPs 从 10.86 降低到 7.86。

基于我们训练的轻量化 VAE，我们开展了 VAE 单独重建图像和 VAE 接入扩散模型文生图的实验，结果如下表所示。在 VAE 图像重建任务上，我们的 VAE 取得了 5.47 的 FID，优于 TAESD，同时在参数量上远小于 SD-v1.4 VAE。在文生图的任务中，我们的 VAE 在取得了与 SD-v1.4 VAE 相当的性能。

3.2 主观效果

如下图所示，比较了不同模型采用单一模型推理和混合推理范式生成效果的区别，可以看到大、小模型的混合推理保留了大模型的语义信息，比如对于 “two sheep standing in the snow...” 这句文本，OursTiny 没有表达出 “two sheep” 的含义，而采用大、小模型混合推理的范式则保留了大模型的文本理解能力，使得生成效果更符合文本描述。可以看到混合推理生成的图片，和大模型生成的图片高度一致。

▲ 图 11: 大小模型混合推理效果

图 12 展示了使用更多真实人像训练集进行微调训练的大模型（Realistic-Vision-V5.1）和小模型（segmind-small-sd）进行混合推理的结果，探索了不同分割步数的影响，可以看到在的情况下，大模型已经主导了基于文本描述的生成方向。

▲ 图 12：人像的混合推理效果

我们对文本“Faceshot Portrait of pretty young（18-year-old）Caucasian wearing a high neck sweater” 在不同分割点下的去噪过程进行了可视化。从可视化结果可以看到，小模型在理解 pretty young（18-year-old）这个特征时存在一定困难，但随着大模型推理步数的增加，模型逐渐理解了文本内容，生成的图像保留了大模型的语义信息。

▲ 图 13: 人像的混合推理效果

我们同样将 Hybrid SD 推理框架拓展到更大的扩散模型 SDXL 上，图 7 给出了大模型（SDXL）和小模型（Koala-700M）混合推理的结果。从可视化的结果可以看出，大模型生成的图片在语义上会优于小模型（例如 “A colorful bird with colorful feather in a green tree”, 大模型准确生成了一只鸟的图片，小模型却错误地生成了两只鸟）。

我们的混合推理生成的图像在语义布局上和大模型保持一致，同时将 10 步后的推理分配给小模型完成，节约了计算资源。

▲ 图 14: 混合推理在SDXL上的效果

我们可视化了接入不同 VAE 在图像重建上的效果和文生图中的效果，如图 15 所示。第一行图展示了 VAE 直接重建的效果，可以看到我们的轻量化 VA 对细节的重建优于 TAESD，对于小人脸的五官重建得更好。第二行展示了 VAE 接入文生图任务中的效果，TAESD 重建的图片在人眼处出现模糊和噪点，而我们的 VAE 很好地重建了人眼，并且整体饱和度更高。

▲ 图 15: VAE 效果比较

3.3 端到端混合推理

我们分别在云侧、端侧、端云结合三种方式下的混合推理的计算量 FLOPs（T）和时间（ms），如下表所示。我们提出的混合推理实现了 19% 的 FLOPs 降低（15.8 -> 12.8），将原始的 SD VAE 替换为我们的轻量化 VAE 后，实现了 49%的 FLOPs 降低（15.8 -> 8.1）。通过端云协同混合推理，我们可以进一步将 2.7 T 的 FLOPs 转移到端侧设备中，有效降低了云侧的推理成本。