扩散模型已经被证明在生成高质量图像方面非常有效。然而，随着这些模型的增大，它们需要更多的内存，并且延迟更高，这给部署带来了巨大的挑战。在这项工作中，作者的目标是通过将扩散模型的权重和激活量化到 4 位来加速扩散模型。在如此激进的 Level 上，权重和激活都是高度敏感的，而传统的用于大型语言模型（如平滑）的课后量化方法变得不适用。
为了克服这个限制，作者提出了一种新的 4 位量化范式，称为 SVDQuant。与平滑方法不同，作者的方法使用低秩分支来吸收这些异常值。作者首先将异常值从激活转移到权重，然后使用高精度低秩分支通过奇异值分解（SVD）吸收权重异常值。
这个过程使得两边的量化都变得更容易。然而，独立运行低秩分支会因为额外的数据移动而导致显著的开销，抵消了量化速度的提升。
为了应对这个问题，作者设计了一个推理引擎 Nunchaku，将低秩分支的核函数与低位分支的核函数融合，以切断冗余内存访问。它也可以无缝支持即插即用的低秩 Adapter （LoRA），无需重新量化。
在 SDXL、PixArt-Σ 和 FLUX.1 等大量实验中，作者验证了 SVDQuant 在保持图像质量方面的有效性。作者将 12B 的 FLUX.1 模型的内存使用降低了 3.5 倍，实现了在 16GB 笔记本电脑 4090 GPU 上的 4 位仅量化基准的 3.0 倍速度提升，为在 PCS 上实现更多交互式应用铺平了道路。作者的量化库和推理引擎都是开源的。
1 Introduction
扩散模型在生成高质量图像（Ho 等人，2020 年）方面显示出惊人的能力，而最近的进展进一步增强了用户对生成过程的控制。在大量数据上进行训练后，这些模型可以从简单的文本 Prompt 中创建出令人惊叹的图像，解锁了各种图像编辑和合成应用。
为了追求更高的图像质量和更精确的文字到图像对齐，研究行人越来越倾向于使用扩散模型。如图 2 所示，Stable Diffusion (SD) (Rombach 等人，2022 年) 1.4 只有 800M 参数，而 SDXL (Podell 等人，2024 年) 将此扩展到 2.6B 参数。AuraFlow v0.1 (fal.ai, 2024 年) 进一步扩展到 6B 参数，最新的模型 FLUX.1 (Black-Forest-Labs, 2024 年) 将参数增加到 12B。与大型语言模型（LLMs）相比，扩散模型的计算成本显著更高。随着模型大小的增加，计算成本增长得更快，这给实际模型的部署带来了内存和延迟上的巨大障碍，尤其是对于需要低延迟的交互式使用场景。
随着摩尔定律的放缓，硬件供应商开始转向低精度推理以维持性能提升。例如，英伟达的 Blackwell 张量核心引入了新的 4 位浮点数（FPA）精度，与 FP8（英伟达，2024 年）相比性能翻倍（英伟达，2022 年）。因此，使用 4 位推理加速扩散模型具有吸引力。在 LLM 领域，研究行人利用量化压缩模型大小并提高推理速度（Dettmers 等人，2022 年；Xiao 等人，2023 年）。然而，与 LLM 不同，扩散模型在计算上受到重约束，即使单个批次也是如此。因此，仅对权重进行量化压缩不能加速扩散模型。为了实现加速，权重和激活都需要量化到相同的位宽；否则，在计算过程中低精度一侧会被向上转换，抵消潜在的性能提升。
在本工作中，作者关注将扩散模型的权重和激活值量化到 4 位。这种具有挑战性和侵略性的方案通常会导致严重的质量退化。现有方法如平滑（Xiao 等人，2023；Lin 等人，2024），试图在权重和激活值之间传递异常值，由于两者都高度易受异常值影响，因此其效果不佳。为了解决这个问题，作者提出了一种新的通用量化范式_SVDQuant_.作者的核心思想是引入一个低成本的分枝来吸收两边的异常值。为了实现这一目标，如图 3 所示，作者首先通过平滑将异常值从激活值 X 迁移到权重 W。然后，作者对更新后的权重 W 应用奇异值分解（SVD），将其分解为一个低秩分枝 L1L2 和一个残差 W — L1L2。低秩分枝在 16 位运行，允许作者仅将残差量化到 4 位，从而显著降低了异常值和幅度。然而，单独运行低秩分枝会导致大量的内存访问开销，抵消了 4 位推理速度的提升。为了克服这个问题，作者专门设计了一个推理引擎_Nunchaku_，将低秩分枝的计算融合到 4 位量化和计算核心中。这种设计使作者能够实现即使有额外的分支，也能测量到的推理速度提升。
SVDQuant 可以量化各种文本到图像的扩散架构，包括 UNet（Ho 等人，2020 年；Ronneberger 等人，2015 年）和 DiT（Peebles 和 Xie，2023 年） Backbone ，同时保持视觉质量。它支持 INT4 和 FP4 数据类型，并与预训练的低秩 Adapter （LoRA）（Hsu 等人，2022 年）无缝集成，无需重新量化。据作者所知，作者是首次成功地将 4 位 PTQ 应用于扩散模型的权重和激活，并在 NVIDIA GPU 上实现了测量速度的提高。在最新的 12B FLUX.1 中，作者很大程度上保留了图像质量，并将原始 BF16 模型的内存占用降低到原来的 3.5 倍，同时在 16GB 笔记本级 RTX4090 GPU 上实现了 3 倍的加速，相对于 NF4 仅量化的 Baseline 。请参阅图 1 以查看视觉示例。
2 Related Work
扩散模型
扩散模型（Sohl-Dickstein 等人，2015 年；Ho 等人，2020 年）已成为一种强大的生成模型类别，以其通过迭代去噪过程模拟数据分布来生成高质量样本的能力而著称。最近，在文本到图像扩散模型（Balaji 等人，2022 年；Rombach 等人，2022 年；Podell 等人，2024 年）方面取得了突破性进展，已经彻底改变了内容生成。研究行人进一步从基于卷积的 UNet 架构（Ronneberger 等人，2015 年；Ho 等人，2020 年）转向了 Transformer （例如，DiT（Peebles 和 Xie，2023 年）和 U-ViT（Bao 等人，2023 年））并将模型规模扩大（Esser 等人，2024 年）。然而，由于扩散模型具有过长的去噪序列和强大的计算，推理速度极其缓慢。为了解决这个问题，已经提出了各种方法，包括少步采样器（Zhang 和 Chen，2022 年；Zhang 等人，2022 年；Lu 等人，2022 年）或从预训练模型中提取更少的步长模型（Salimans 和 Ho，2021 年；Meng 等人，2022 年；Song 等人，2023 年；Luo 等人，2023 年；Sauer 等人，2023 年；Yin 等人，2024 年；Xing 等人，2024 年）。另一条工作线则选择通过高效架构设计（Li 等人，2023b；2020 年；Cai 等人，2024 年；Liu 等人，2024a），量化（Shang 等人，2023 年；Li 等人，2023a），稀疏推理（Li 等人，2022 年；Ma 等人，2024b；A），和分布式推理（Li 等人，2024b；Wang 等人，2024c；Chen 等人，2024b）来优化或加速计算。本工作专注于将扩散模型量化到 4 位，以减少计算复杂性。作者的方法还可以应用于少步扩散模型，以进一步减少延迟（参见第 5.2 节）。
量化。量化已被认为是减少语言模型模型大小并加速推理的有效方法（Dettmers 等人，2022 年；Frantar 等人，2023 年；Xiao 等人，2023 年；Lin 等人，2024b；a；Kim 等人，2024 年；Zhao 等人，2024 年）。对于扩散模型，Q-Diffusion（Li 等人，2023 年 a）和 PTQ4DM（Shang 等人，2023 年）首先实现了 8 位量化。后续的工作通过诸如敏感性分析（Yang 等人，2023 年）和时间步长感知量化（He 等人，2023 年；Huang 等人，2024 年；Liu 等人，2024b；Wang 等人，2024 年 a）等方法改进了这些技术。一些最近的工作扩展了这些设置，应用于文本到图像模型（Tang 等人，2023 年；Zhao 等人，2024c），Dilated Convolution（Wu 等人，2024 年），量化感知训练（He 等人，2024 年；Zheng 等人，2024 年；Wang 等人，2024b；Sui 等人，2024 年），视频生成（Zhao 等人，2024b），以及不同的数据类型（Liu 和 Zhang，2024 年）。在这些工作中，只有 MixDQ（Zhao 等人，2024c）和 ViDiT-Q（Zhao 等人，2024d）实现了低位推理引擎，并在 GPU 上报告了 8 位速度提升。在本工作中，作者进一步将扩散模型量化到 4 位，支持整数或浮点数据类型，兼容 UNet Backbone 网络（Ho 等人，2020 年）和最近的 Dilated Convolution（Peebles 和 Xie，2023 年）架构。作者的自定义推理引擎 Nunchaku 进一步确保了硬件加速。此外，在将 LoRA 应用于模型时，现有方法需要将 LoRA 分支合并到主分支并重新量化模型以避免 LoRA 分支中的巨大内存访问开销。Nunchaku 通过 Kernel 融合切断了这种开销，使低秩分支作为单独的分支高效运行，无需重新量化。
低秩分解在深度学习领域受到了广泛关注，因为它可以提高计算和内存效率（Hu 等人，2022 年；赵等人，2024 年；Jaiswal 等人，2024 年）。虽然直接将这种方法应用于模型权重可以降低计算和内存需求（Hsu 等人，2022 年；于等人，2023 年；李等人，2023 年 c），但通常会导致性能下降。相反，Yao 等人（2023 年）将低秩分解与量化相结合用于模型压缩，使用低秩分支来补偿量化误差。低秩自适应（LoRA）（Hu 等人，2022 年）通过使用低秩矩阵来调整一组预训练权重的子集以实现有效的微调，从而开辟了另一项重要研究。这一方法催生了许多进展（Dettmers 等人，2023 年；郭等人，2024 年；李等人，2024 年 c；何等人，2024 年；许等人，2024 年；Xu 等人，2024 年 b），它们将量化模型与低秩 Adapter 相结合，在模型微调期间降低内存使用。然而，作者的工作有两个主要区别。首先，作者的目标不同，因为作者希望通过量化来加速模型推理，而以前的工作主要关注模型压缩或有效的微调。因此，他们主要考虑权重量化，导致没有加速。其次，如作者在实验（图 6 和第 5.2 节中的消融研究）中所见，直接应用这些方法不仅会降低图像质量，还会引入显著的额外开销。相比之下，作者的方法由于对权重和激活的联合量化以及推理引擎 Nunchaku 通过将低秩分支核融合到低位计算中而获得了更好的性能。
3 Quantization Preliminary
量化是加速网络线性层的有效方法。给定张量 X，量化过程定义如下：
在这项工作中，作者专注于 W4A4 量化加速，其中权值和激活值中的异常值给加速带来了巨大障碍。传统方法来抑制这些异常值包括量化感知训练（QAT）（He 等人，2024 年）和旋转（Ashkboos 等人，2024 年；Liu 等人，2024 年；Lin 等人，2024b 年）。QAT 需要大量的计算资源，尤其是对于参数超过 10 亿的大型模型（例如，FLUX.1）。旋转不适用，因为扩散模型中使用了自适应归一化层（Peebles 和 Xie，2023 年）。运行时生成的归一化权重阻止了旋转矩阵与投影层权值的离线集成。因此，在线旋转激活和权值会导致显著的运行时开销。
4 Method
在本节中，作者首先提出问题并讨论量化误差来源。接下来，作者提出 SVDQuant，一种用于扩散模型的新的 W4A4 量化范式。作者的关键思想是引入一个额外的低秩分支，可以吸收权重和激活的量化困难。最后，作者提供一个与 Kernel 融合的协同设计推理引擎 Nunchaku，以最小化 4 位模型中低秩分支的开销。
Problem Formulation
命题 4.1（误差分解）。：量化误差可以分解如下：
SVDQuant: Absorbing Outliers via Low-Rank Branch
从等式 5 中，量化误差可以表示为
Nunchaku: Fusing Low-Rank and Low-Bit Branch Kernels
尽管低秩分支在理论上的计算成本可以忽略不计，但将其作为单独的分支运行将产生显著的延迟开销--大约是 4 位分支延迟的 50%，如图 6(a)所示。这是因为，对于一个小秩 r，尽管计算成本大大降低，输入和输出激活的数据大小保持不变，将 Bottleneck 从计算转移到内存访问。情况变得更糟，尤其是当激活无法放入 GPU L2 缓存。
例如，低秩分支的 QKV 投影的向上投影由于其输出超过可用的 L2 缓存而变得非常慢，导致额外的加载和存储操作到 DRAM。幸运的是，作者观察到低秩分支的向下投影 L1 与低位分支的量化核共享相同的输入，而向上投影 L2 与 4 位计算核共享相同的输出，如图 6(b)所示。
通过将向下投影与量化核融合和将向上投影与 4 位计算核融合，低秩分支可以与低位分支共享激活，消除额外的内存访问，并将 Kernel 调用数量减半。因此，低秩分支只增加了 5%~10%的延迟，使其几乎免费。
5 Experiments
Setups
模型。作者使用以下文本到图像模型来评估作者的方法，包括 UNet 和 Dilated Backbone ：
FLUX.1 是基于 DiT 的开源扩散模型中的最先进版本。它包括 19 个联合注意力块 和 38 个并行注意力块 ，总共拥有 120 亿个参数。作者在 50 步指导蒸馏（FLUX.1-dev）和 4 步时间步蒸馏的变体上进行评估。
PixArt-Σ(Chen et al., 2024a) 是另一个基于 DiT 的模型。它没有使用联合注意力，而是堆叠了包含自注意力、交叉注意力和前馈层的 28 个注意力块，总共拥有 600M 个参数。作者在默认的 20 步设置上评估它。
稳定扩散 XL（SDXL）是一种广泛使用的基于 UNet 的模型，参数数量为 2.6 亿（Podell 等人，2024 年）。该模型可以预测三种分辨率的噪声。最高分辨率阶段完全由 ResBlocks（He 等人，2016 年）处理，而其他两个阶段则共同使用 ResBlocks 和注意力层。与 PixArt-Σ 类似，SDXL 也采用了交叉注意力层进行文本条件化。作者在 30 步设置下评估它，同时包括其 4 步浓缩变体，SDXL-Turbo（Sauer 等人，2023 年）。
数据集。 遵循先前的研究,作者在 COCO Captions 2024（Chen 等人，2015）中随机选择了 Prompt 语进行校准。为了评估作者的方法的一般化能力，作者采用了两种具有不同风格的 Prompt 集作为基准：
MHQ-30K（李等，2024a）包括来自 Midjourney 的 30K 个样本，分为 10 个常见类别，每个类别 3K 个样本。作者从这个数据集中均匀选择 5K 个 Prompt 来评估模型在艺术图像生成方面的性能。
密集标注图像（DCI）(Urbanek 等，2024)是一个包含约 8000 张带有详细人机标注的图像的的数据集，平均每张图像有 1000 多字。在作者的实验中，作者使用摘要版（sDCI），其中标注被压缩为使用大型语言模型（LLMs）的 77 个 Token ，以容纳扩散模型。同样，作者随机选择 5000 个 Prompt 以进行现实图像生成的有效评估。
基准。作者将 SVDQuant 与以下 PTQ（PTQ）方法进行比较：
4-bit NormalFloat (NF4)是一种用于仅量化权重（Dettmers 等人，2023 年）的数据类型。它假设权重遵循正态分布，并且是信息论上最优的 4 位表示。作者使用社区量化的 NF4 FLUX.1 模型（Lllyasviel，2024）作为基准。
ViDiT-Q (赵等，2024b) 通过每个 token 的量化和平滑（肖等，2023）来缓解不同批次和 token 之间的异常，并在 PixArt-上实现了无损的 8 位量化。
MixDQ (赵等，2024c) 在文本嵌入的句子开头的 Token 中识别出异常值，并用 16 位预计算进行保护。这种方法在 SDXL-Turbo 上实现了高达 W4A8 的量化和几乎不降低性能。
TensorRT 包含一个行业级的 PTQ 工具包，可以将扩散模型量化到 8 位。它使用平滑技术，仅在选定的时间步长范围内使用百分位数方案校准激活值。
评价指标参考之前的研究（Li 等人，2022 年；2024b 年）（16 位原模型产生的结果），作者主要评估图像质量和相似度。在图像质量评估中，作者使用弗雷歇斯 inception 距离（FID，越低越好）（Heusel 等人，2017 年）来衡量生成图像和真实图像之间的分布距离。此外，作者使用图像奖励（越高越好）（Xu 等人，2024a 年）来近似人类对生成图像的评分（LPIPS，越低越好）（Zhang 等人，2018 年）用于测量 16 位原模型的图像的感知相似性（PSNR，越高越好）。请参阅作者的附录 B.1 以获取更多评价指标（CLIP IQA（Wang 等人，2023 年）；CLIP 得分（Hessel 等人，2021 年）和 SSIM4）。
实现细节对于 8 位设置，作者使用按词动态激活量化（per-token dynamic activation quantization）和按通道权重量化（per-channel weight quantization）以及低秩分支（low-rank branch）的秩为 16。对于 4 位设置，作者采用对激活和权重都进行按组对称量化（per-group symmetric quantization）的方法，同时配合低秩分支（low-rank branch）的秩为 32。INT4 量化使用组大小为 64，量级为 16 位的标度。FP4 量化使用组大小为 32，量级为 FP8 标度（Rouhani 等人，2023）。对于 FLUX.1 模型，自适应归一化线性层中的输入保持为 16 位（即 W4A16）。对于其他模型，交叉注意力中 Key 和 Value 投影在 16 位保留，因为它们的延迟仅占总运行时间的不到 5%。
Results
作者报告了各种模型和精度 Level 下的定量质量结果，并在图 7 中展示了相应的 4 位定性比较。在所有模型中，作者的 8 位结果可以完美地反映 16 位结果，实现了比所有其他 8 位 Baseline 更高的 PSNR（27）。在 FLUX.1-dev 上，作者的 INT8 PSNR 甚至达到了 27 分（MJHQ）。
对于 4 位量化的 FLUX.1，作者的 SVDQuant 在图像奖励方面超过了 NF4 W4A16 Baseline 。在 schnell 变体中，作者的图像奖励甚至超过了原始 BF16 模型，这表明了人类偏好的强烈性。在 PixArt-上，虽然作者的 INT4 图像奖励略有下降，但作者的 FP4 模型比 FP16 模型获得了更高的分数。这可能是由于 PixArt-Σ 的小模型大小（600M 参数）使其已经非常紧凑，并得益于较小的组大小。值得注意的是，作者的 INT4 和 FP4 结果在所有指标上明显优于 ViDiT-Q 的 W4A8 结果。对于基于 UNet 的模型，在 SDXL-Turbo 上，作者的 4 位模型显著超过了 MixDQ 的 W4A8 结果，且作者的 FID 分数与 FP16 模型相当，表明性能没有损失。在 SDXL 上，作者的 INT4 和 FP4 结果达到了与 TensorRT 的 W8A8 性能相媲美的质量，这代表了 8 位的最优水平。如图 15 所示，在附录中的图 15，作者的视觉质量仅略有下降。
内存节省与加速。 在图 8 中，作者报告了 FLUX.1 的模型大小、内存节省和加速情况。作者的 INT4 量化将原始 Transformer 大小从 22.2 GiB 减少到 6.1 GiB，包括低秩分支的 0.3 GiB 开销，导致总体 3.6 倍减少。由于权重和激活值都进行量化，与仅对 NF4 权重进行量化的 Transformer 相比，作者的推理引擎 Nunchaku 甚至节省了更多的内存占用，并在桌面级和笔记本级 NVIDIA RTX 4090 GPU 上提供了 3.0 倍加速。值得注意的是，原始 BF16 模型在 16GB 笔记本电脑上的每一层都需要 CPU 卸载，而作者的 INT4 模型完全可以在 GPU 内存中运行，通过避免卸载实现了 10.1 倍加速。作者预计在 NVIDIA 下一代 Blackwell GPU 上的 FP4 量化模型将实现更大的加速，因为它们天生支持组量化，无需专门的 GPU Kernel 。
将 LoRA 集成到 Nunchaku 中。之前的量化方法需要将 LoRA 分支与低秩分支融合，并在融合 LoRA 时重新对模型进行量化。然而，作者的 Nunchaku 消除了冗余的内存访问，允许添加独立的 LoRA 分支。在实际应用中，作者可以通过稍微提高秩的方式，将 LoRA 分支融合到作者的低秩分支中，从而进一步提高效率。在图 9 中，作者展示了将五种不同风格（真实感、奇异天空插图、动漫、儿童素描和编织艺术）的 LoRAs 应用到作者的 INT4 FLUX.1-dev 模型中的视觉示例。作者的 INT4 模型成功适应了每种风格，同时保持了 16 位版本图像质量。更多视觉示例请参见附录 B.2。
在图 10 中，作者展示了 SVDQuan 在 PixArt-Σ 上的几个消融实验。首先，仅使用 SVD 和 Naive 量化在 4 位设置下表现不佳，导致图像质量严重降低。尽管对量化进行平滑处理相对于 Naive 量化稍微提高了图像质量，但整体结果仍然不满意。LoRC（Yao 等人，2023 年）引入了一个低秩分支来补偿量化误差，但这种方法并非最优。量化误差呈现平滑的奇异值分布。因此，低秩补偿未能有效地减小这些误差，如第 4.2 节所述。相比之下，作者首先分解权重，仅量化残差。如图 5 所示，前几个奇异值显著大于其余的，使作者能够将它们转移到低秩分支，从而有效地减小权重大小。最后，平滑处理整合了离群点，进一步使低秩分支吸收激活中的离群点，从而显著提高图像质量。
权衡增加秩的利弊。图 11 展示了在 PixArt-Σ 上使用 SVDQuant 的不同秩 r 的结果。将秩从 16 增加到 64 显著提高了图像质量，但增加了参数和延迟开销。在作者的实验中，作者选择秩为 32，它提供了较好的质量，而开销较小。
6 Conclusion & Discussion
在这项工作中，作者提出了一种新颖的 4 位 PTQ 范式 SVDQuant，用于扩散模型。它采用低秩分支来吸收权值和激活中的异常值，从而简化量化的过程。作者的推理引擎 Nunchaku 进一步将低秩和低位分支核融合，减少内存使用并切断冗余数据移动开销。
大量实验表明，SVDQuant 保持了图像质量。Nunchaku 进一步实现了在原始 16 位模型上的内存使用减少 3.5 倍，并在 NVIDIA RTX-4090 笔记本上的仅量化权重实现 3.0 倍速度提升。这一进步使得大型扩散模型在边缘设备上的有效部署成为可能，解锁了交互式 AI 应用的更广泛潜力。
局限性。 在本研究中，作者没有报告作者的 FP4 模型的加速情况。这是因为作者没有访问 Blackwell GPU，它们原生支持组量化所需的准确性和微缩。在 Blackwell 硬件上，作者预计与 4090 GPU 上的 INT4 结果相比会有更大的加速。
参考文献
[0]. SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models.
END
作者：小书童
来源：集智书童
推荐阅读
- 【ml-engineering 翻译系列】计算加速器之 gpu
- 实时 4K 图像修复：TSFormer 以 3.38M 参数达 40fps 高性能，推动技术落地
- 图解 OpenRLHF 中基于 Ray 的分布式训练流程
- 超越 YOLOv10 和 YOLOv7，专为大规模高分辨率图像处理设计 ！
欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
