点击
阅读原文
观看作者视频讲解
！
ECCV'24公布的中稿结果中，本届共有2395篇文章被接收，录用率低于20%。由
清华大学、无问芯穹、微软、加州大学圣芭芭拉分校和上海交通大学研究团队合作的《MixDQ: Memory-Efficient Few-Step Text-to-Image Diffusion Models with Metric-Decoupled Mixed Precision Quantization》中稿！
MixDQ：
一种面向少步扩散模型的混合比特量化方案。
-
MixDQ分析定位了“少步扩散模型量化”的独特问题，并提出针对性解决方案。
-
针对少步生成模型，在现有量化方案在W8A8损失严重的情况下，MixDQ在能够实现多方面指标（图像质量，文图吻合，人为偏好）无损的W8A8量化，W4A8无明显视觉损失。
-
实现了高效的INT8 GPU算子，以实现实际的显存与延迟优化，并将模型开源为Huggingface Pipeline，通过几行代码即可调用。
图注：《MixDQ: Memory-Efficient Few-Step Text-to-Image Diffusion Models with Metric-Decoupled Mixed Precision Quantization》（简称MixDQ）
论文标题：
MixDQ: Memory-Efficient Few-Step Text-to-Image Diffusion Models with Metric-Decoupled Mixed Precision Quantization
论文链接：
https://arxiv.org/abs/2405.17873
Project page:
https://a-suozhang.xyz/mixdq.github.io/
Huggingface Pipeline:
https://huggingface.co/nics-efc/MixDQ
图注：《MixDQ: Memory-Efficient Few-Step Text-to-Image Diffusion Models with Metric-Decoupled Mixed Precision Quantization》（简称MixDQ）
近年来，
扩散模型(Diffusion Model)
在视觉生成领域取得了显著的进展，Stable Diffusion模型能够依据文本信息生成高度拟真且美观的图像，“AI绘图”也在各领域成为了热门话题。然而，由于文生图大模型具有巨大的参数规模（Stable Diffusion XL: 3.5B, 35亿参数）与扩散模型循环迭代式的推理特点（单次生成图片需要对大模型进行数十次推理），其运行的
硬件资源消耗十分巨大
，对其实际应用带来了巨大挑战。
在扩散模型效率优化的算法研究中，“少步数”生成模型近期成为热点研究话题，通过减少扩散模型中的迭代步数，以减少计算代价，可实现相较原本方法十余倍的延迟加速。LCM
[1]
首先实现2~8步的图像生成，SDXL-turbo
[2]
模型更是将生成步数减少到了单步。少步数生成模型，
解决了模型的延迟痛点
，目前在RTX3090上推理一次单步生成模型仅需要不到1s，但是，
模型的显存开销仍然显著
。SDXL模型生成单张图片所需要的显存为约9GB（FP16模式，采用FlashAttention），超过了许多桌面级GPU的显存容量（如RTX4070，8GB）。
图注：SDXL模型生成单张图片所需显存超出桌面级GPU显存容量
低比特量化
是一种被广泛使用的减少模型计算存储开销的方法，通过将原本高精度浮点（FP32/FP16）的模型全权重与激活值 (Weight and Activation, 简称W&A)，转化为低比特定点数（INT8/INT4），可以显著减少模型显存开销与计算复杂度。为
解决扩散模型的显存瓶颈
，
来自清华大学电子工程系、无问芯穹、微软、加州大学圣芭芭拉分校和上海交通大学研究团队，提出了一种新颖的扩散模型低比特量化方法：《MixDQ: Memory-Efficient Few-Step Text-to-Image Diffusion Models with Metric-Decoupled Mixed Precision Quantization 》
。这项工作中，研究人员分析了少步数生成模型低比特量化的挑战，并提出了一种
面向少步扩散模型的混合比特自动化设计方案
：MixDQ，在现有量化方案在W8A8失败的情况下，实现了
无损的W5A8量化
。将扩散模型中U-Net的
显存开销降低3.4倍，端到端延迟提速约1.5x，让文生图大模型“更小更快”，能够在各种小存储终端设备上被应用起来。
少步数扩散模型显著减少了迭代步数，提升了执行效率。然而，这也使得它对量化更加敏感。如下图所示，应用同样的基线量化方案Q-Diffusion，30步的SDXL模型在能够生成正常的图片，然而针对更少步的SDXL-turbo，图片视觉效果损失明显，在1步时图像更是变得模糊且扭曲，有时还会出现完全崩溃的情况（如下图中带粉色噪声的北极熊）。
因此，
针对少步数扩散模型设计量化方法更具挑战。
除了图像质量的损失之外，文生图任务的另一重要评估指标为“文图吻合度”，即生成的图像是否如实的遵从了文本描述。如下图，量化后模型产生的图片都产生了与全精度模型较大的差异，有时甚至会完全违背文本描述。这指出了，
当前文生图模型量化方法对文图吻合度的保持有待改进。
图注：量化后模型生成的图片与全精度模型、文本描述的差异
本文对少步模型量化失败的原因进行了系统实验分析，并得到了以下结论：
（1）神经网络模型中不同层的量化
敏感性存在着高度差异
，量化实际上被少部分“极端敏感”的层所瓶颈。
（2）现有的量化敏感性分析方法，
未区分量化对“图像质量”与“图像内容”
，而导致准确度不足。
受上述观察所启发，本文提出采用混合精度量化以解决“量化被少数极端敏感层所瓶颈”的问题，通过给这些特定层赋予更高的位宽，可以在保持量化性能的前提下，实现模型其他部分的低比特量化，获得硬件资源节省。
针对“如何准确识别出量化敏感层”的问题：
（1）本文考虑了现有方案对量化敏感性估计的不足，并提出了一种“指标解耦”的方案，来分别对量化对图像质量与内容的影响进行分析，以准确得到量化的敏感性。
（2）将混合比特位宽的分配定义为了一个整数规划问题，并高效自动化完成求解。
（3）针对特别敏感的文本特征量化，本文识别出了量化被“句首令牌”（Begin-of-Sentence Token）离群值所影响的关键问题，并提出了一种BOS-aware的量化方案，有效解决了该问题。
采用以上方案，MixDQ可以在几乎无损的情况下实现W5A8的量化，可获得约3.4x的显存优化，与1.5x的延迟优化。
为了在具体硬件上验证MixDQ量化方案的实际效率优化，针对Nvidia GPU，基于CUTLASS库开发了低比特量化算子，实现了SDXL-turbo量化模型的端到端推理。经过GPU实测，在多种量化位宽配置方案下，
MixDQ能获得显著的显存与延迟优化结果
。优化后的MixDQ-SDXL-turbo模型已封装为Huggingface Pipeline，以供用户一键实例化并调用。
图注：MixDQ-SDXL-turbo模型封装及调用展示
具体来说，如下图所示，MixDQ的4/8比特权重量化，可实现1.87x与3.03x的权重显存占用量优化，原本需要5.2GB存储的SDXL模型权重，经过W4优化后仅需不足2G，整个推理流程的显存占用也从8GB减少至4.5GB。当权重与激活值均被量化至INT8时，针对可量化层，MixDQ的高效量化算子实现，相比FP16版本能获得1.97x的延迟加速，考虑量化与反量化操作的额外开销，仍能获得1.45x的端到端延迟优化。INT4的量化算子仍在开发中，预期可获得更高的端到端延迟优化。
图注：MixDQ在不同比特宽度配置对内存和延迟的影响
对比现有的其他部署优化工具的量化加速方案，
MixDQ是第一个实现了少步数生成模型的量化实际显存与延迟优化的方案
