SVDQuant：通过低秩分量吸收异常值的4比特扩散模型

摘要

图1：SVDQuant是一种4比特权重和激活的后训练量化技术，能够很好地保持视觉保真度。在12B FLUX.1-dev上，与BF16模型相比，它实现了3.6倍的内存减少。通过消除CPU卸载，在16GB笔记本电脑4090 GPU上，它比16比特模型提供了8.7倍的加速，比NF4 W4A16基线快3倍。在PixArt-Σ上，它在视觉质量上显著优于其他W4A4甚至W4A8基线。“E2E”表示包括文本编码器和VAE解码器的端到端延迟。

图2：LLMs和扩散模型的计算量与参数量的对比。LLMs的计算量以512个上下文和256个输出标记为基准，而扩散模型的计算量则为单步计算。虚线表示趋势。

扩散模型在生成高质量图像方面已被证明非常有效。然而，随着这些模型规模的扩大，它们需要更多的内存并遭受更高的延迟，给部署带来了重大挑战。在这项工作中，我们旨在通过将权重和激活量化到4比特来加速扩散模型。在这种激进的水平上，权重和激活都非常敏感，大型语言模型的传统后训练量化方法（如平滑）变得不足。为了克服这一限制，我们提出了_SVDQuant_，一种新的4比特量化范式。与平滑方法在权重和激活之间重新分配异常值不同，我们的方法通过低秩分支_吸收_这些异常值。我们首先通过将异常值从激活转移到权重来巩固它们，然后使用高精度的低秩分支通过奇异值分解（SVD）接收权重异常值。这一过程缓解了两边的量化难度。然而，天真地独立运行低秩分支会由于激活的额外数据移动而带来显著的开销，抵消了量化加速。为了解决这个问题，我们共同设计了一个推理引擎_Nunchaku_，它将低秩分支的内核与低比特分支的内核融合，以减少冗余的内存访问。它还可以无缝支持现成的低秩适配器（LoRAs），而无需重新量化。在SDXL、PixArt-ΣΣ和FLUX.1上的广泛实验验证了SVDQuant在保持图像质量方面的有效性。我们将12B FLUX.1模型的内存使用减少了3.5倍，在16GB笔记本电脑4090 GPU上比4比特权重量化基线实现了3.0倍的加速，为PC上的更多交互式应用铺平了道路。我们的量化库1和推理引擎2是开源的。

Footnote 1: 量化库：github.com/mit-han-lab/deepcompressor

Footnote 2: 推理引擎：github.com/mit-han-lab/nunchaku

1 引言

扩散模型在生成高质量图像方面展示了显著的能力（Ho et al., 2020），最近的进展进一步增强了用户对生成过程的控制。在大量数据上训练的这些模型可以从简单的文本提示中创建令人惊叹的图像，解锁了多样化的图像编辑和合成应用（Meng et al., ; Ruiz et al., 2023; Zhang et al., 2023）。

为了追求更高的图像质量和更精确的文本到图像对齐，研究人员正在不断扩大扩散模型的规模。如图2所示，稳定扩散（SD）（Rombach et al., 2022）1.4只有8亿参数，而SDXL（Podell et al., 2024）将其扩展到26亿参数。AuraFlow v0.1（fal.ai, 2024）进一步扩展到60亿参数，最新的模型FLUX.1（Black-Forest-Labs, 2024）将边界推至120亿参数。与大型语言模型（LLMs）相比，扩散模型的计算密集度显著更高。它们的计算成本3随着模型规模的增加而更快速地增长，给实际模型部署带来了内存和延迟障碍，特别是对于需要低延迟的交互式用例。

Footnote 3: 计算成本通过乘积累加操作（MACs）的数量来衡量。1 MAC=2 FLOPs。

随着摩尔定律的放缓，硬件供应商正在转向低精度推理以维持性能改进。例如，NVIDIA的Blackwell Tensor Cores引入了一种新的4比特浮点（FP4）精度，性能比FP8翻倍（NVIDIA, 2024）。因此，使用4比特推理来加速扩散模型是吸引人的。在LLMs领域，研究人员利用量化来压缩模型大小并提升推理速度（Dettmers et al., 2022; Xiao et al., 2023）。然而，与LLMs不同——其中延迟主要受加载模型权重的限制，特别是在小批量情况下——扩散模型即使在单批量情况下也是计算密集型的。因此，仅量化权重无法加速扩散模型。为了实现加速，权重和激活必须量化到相同的比特宽度；否则，较低精度的部分将在计算过程中向上转换，抵消潜在的性能提升。

SVDQuant可以量化各种文本到图像扩散架构，包括UNet（Ho et al., 2020; Ronneberger et al., 2015）和DiT（Peebles & Xie, 2023）骨干，同时保持视觉质量。它支持INT4和FP4数据类型，并能无缝集成预训练的低秩适配器（LoRA）（Hsu et al., 2022）而无需重新量化。据我们所知，我们是第一个成功将4比特后训练量化应用于扩散模型的权重和激活，并在NVIDIA GPU上实现测量加速的。在最新的12B FLUX.1上，我们大大保留了图像质量，并将原始BF16模型的内存占用减少了3.5倍，在16GB笔记本电脑RTX4090 GPU上比NF4权重量化基线实现了3.0倍的加速。参见图1中的视觉示例。

图3：SVDQuant概述。(a) 最初，激活X和权重W都包含异常值，使得4比特量化具有挑战性。(b) 我们将异常值从激活迁移到权重，得到更新后的激活Xˆ和权重Wˆ。虽然Xˆ变得更容易量化，但Wˆ现在变得更加困难。(c) SVDQuant进一步将Wˆ分解为低秩分量L1L2和残差Wˆ − L1L2，通过SVD实现。因此，量化难度通过运行在16比特精度的低秩分支得到缓解。

2 相关工作

扩散模型。 扩散模型（Sohl-Dickstein et al., 2015; Ho et al., 2020）已成为一类强大的生成模型，通过建模数据分布的迭代去噪过程生成高质量样本。最近的文本到图像扩散模型（Balaji et al., 2022; Rombach et al., 2022; Podell et al., 2024）已经彻底改变了内容生成。研究人员进一步从基于卷积的UNet架构（Ronneberger et al., 2015; Ho et al., 2020）转向transformer（如DiT（Peebles & Xie, 2023）和U-ViT（Bao et al., 2023））并扩大模型规模（Esser et al., 2024）。然而，扩散模型由于其长时间的去噪序列和密集的计算而遭受极慢的推理速度。为了解决这个问题，各种方法被提出，包括少步采样器（Zhang & Chen, 2022; Zhang et al., 2022; Lu et al., 2022）或从预训练模型中蒸馏少步模型（Salimans & Ho, 2021; Meng et al., ; Song et al., 2023; Luo et al., 2023; Sauer et al., 2023; Yin et al., , ; Kang et al., 2024）。另一类工作选择通过高效架构设计（Li et al., ; Li et al., ; Li et al., ; Liu et al., ）、量化（Shang et al., 2023; Li et al., ）、稀疏推理（Li et al., 2022; Ma et al., , ）和分布式推理（Li et al., ; Wang et al., ; Chen et al., ）来优化或加速计算。本文专注于将扩散模型量化到4比特以减少计算复杂性。我们的方法也可以应用于少步扩散模型以进一步减少延迟（见第5.2节）。

量化。 量化已被认为是LLMs减少模型大小和加速推理的有效方法（Dettmers et al., 2022; Frantar et al., 2023; Xiao et al., 2023; Lin et al., , ; Kim et al., 2024; Zhao et al., ）。对于扩散模型，Q-Diffusion（Li et al., ）和PTQ4DM（Shang et al., 2023）首次实现了8比特量化。后续工作通过敏感性分析（Yang et al., 2023）和时间步长感知量化（He et al., 2023; Huang et al., 2024; Liu et al., ; Wang et al., ）等方法改进了这些技术。一些最近的工作将设置扩展到文本到图像模型（Tang et al., 2023; Zhao et al., ）、DiT骨干（Wu et al., 2024）、量化感知训练（He et al., 2024; Zheng et al., 2024; Wang et al., ; Sui et al., 2024）、视频生成（Zhao et al., ）和不同数据类型（Liu & Zhang, 2024）。在这些工作中，只有MixDQ（Zhao et al., ）和ViDiT-Q（Zhao et al., ）实现了低比特推理引擎并在GPU上报告了测量的8比特加速。在这项工作中，我们将边界进一步推至4比特量化扩散模型，支持整数或浮点数据类型，兼容UNet骨干（Ho et al., 2020）和最近的DiT架构（Peebles & Xie, 2023）。我们共同设计的推理引擎Nunchaku进一步确保了硬件加速。此外，当将LoRA应用于模型时，现有方法需要将LoRA分支融合到主分支并重新量化模型，以避免LoRA分支中的巨大内存访问开销。Nunchaku通过内核融合切断了这种开销，允许低秩分支作为单独分支高效运行，无需重新量化。

低秩分解。 低秩分解在深度学习中引起了广泛关注，以提高计算和内存效率（Hu et al., 2022; Zhao et al., ; Jaiswal et al., 2024）。虽然直接将这种方法应用于模型权重可以减少计算和内存需求（Hsu et al., 2022; Yuan et al., 2023; Li et al., ），但通常会导致性能下降。相反，Yao et al.（2023）将其与量化结合用于模型压缩，使用低秩分支来补偿量化误差。低秩适配（LoRA）（Hu et al., 2022）引入了另一条重要研究线，使用低秩矩阵调整预训练权重子集以实现高效微调。这引发了众多进展（Dettmers et al., 2023; Guo et al., 2024; Li et al., ; He et al., 2024; Xu et al., ），将量化模型与低秩适配器结合以减少模型微调期间的内存使用。然而，我们的工作在两个主要方面有所不同。首先，我们的目标不同，我们旨在通过量化加速模型推理，而之前的工作主要关注模型压缩或高效微调。因此，他们主要考虑仅权重量化，导致没有加速。其次，如我们的实验所示（图6和第5.2节的消融研究），直接应用这些方法不仅会降低图像质量，还会引入显著的开销。相比之下，我们的方法由于权重和激活的联合量化以及推理引擎Nunchaku最小化了开销，因此性能更好。

3 量化预备知识

量化是加速网络中线性层的有效方法。给定一个张量X，量化过程定义为：

在这项工作中，我们专注于W4A4量化以实现加速，其中权重和激活中的异常值构成了重大障碍。传统方法抑制这些异常值包括量化感知训练（QAT）（He et al., 2024）和旋转（Ashkboos et al., 2024; Liu et al., ; Lin et al., ）。QAT需要大量计算资源，特别是对于超过100亿参数的模型（如FLUX.1）。旋转由于扩散模型中使用了自适应归一化层（Peebles and Xie, 2023）而不可行。运行时生成的归一化权重排除了旋转矩阵与投影层权重的离线集成。因此，在线旋转激活和权重会带来显著的运行时开销。

4 方法

在本节中，我们首先阐述我们的问题并讨论量化误差来自何处。接下来，我们介绍SVDQuant，一种用于扩散模型的新W4A4量化范式。我们的关键思想是引入一个额外的低秩分支，可以吸收权重和激活中的量化难度。最后，我们提供了一个共同设计的推理引擎Nunchaku，通过内核融合最小化4比特模型中低秩分支的开销。

图4：PixArt-Σ中输入和权重的示例值分布。λ是平滑因子。红色表示异常值。最初，输入X和权重W都包含显著的异常值。经过平滑处理后，Xˆ的范围缩小，异常值大大减少，而Wˆ显示出更多的异常值。一旦通过SVD低秩分支L1L2减去，残差R的范围变窄，并且没有异常值。

问题公式化

SVDQuant：通过低秩分支吸收异常值

图6：(a) 天真地运行秩为32的低秩分支将引入57%的延迟开销，这是由于在向下投影中额外读取16比特输入和在向上投影中额外写入16比特输出。我们的Nunchaku引擎通过内核融合优化了这一开销。(b) 向下投影和量化内核使用相同的输入，而向上投影和4比特计算内核共享相同的输出。为了减少数据移动开销，我们将前两个和后两个内核融合在一起。

Footnote 8: https://en.wikipedia.org/wiki/Low-rank_approximation

Nunchaku：融合低秩和低比特分支内核

尽管低秩分支在理论上引入的计算可以忽略不计，但将其作为单独分支运行会带来显著的延迟开销——约占4比特分支延迟的50%，如图6(a)所示。这是因为对于小秩r，尽管计算成本大大降低，输入和输出激活的数据大小保持不变，将瓶颈从计算转移到内存访问。这种情况恶化，特别是在激活无法适应GPU L2缓存时。例如，QKY投影中的低秩分支的上投影要慢得多，因为其输出超过了可用的L2缓存，导致额外的加载和存储操作到DRAM。幸运的是，我们观察到低秩分支的下投影L1L1​与低比特分支中的量化内核共享相同的输入，而上投影L2与4比特计算内核共享相同的输出，如图6(b)所示。通过将下投影与量化内核融合，上投影与4比特计算内核融合，低秩分支可以与低比特分支共享激活，消除额外的内存访问，并将内核调用数量减半。因此，低秩分支仅增加了5∼∼10%的延迟，几乎无成本。

5 实验

设置

表1：不同模型之间的定量质量比较。IR表示ImageReward。我们的8比特结果与16比特模型的质量非常接近。此外，我们的4比特结果优于其他4比特基线，有效地保持了16比特模型的视觉质量。

模型。 我们使用以下文本到图像模型进行基准测试，包括UNet（Ronneberger et al., 2015; Ho et al., 2020）和DiT（Peebles & Xie, 2023）骨干：

• *FLUX.1（Black-Forest-Labs, 2024）是基于DiT的最先进的开源扩散模型。它由19个联合注意力块（Esser et al., 2024）和38个并行注意力块（Dehghani et al., 2023）组成，总计120亿参数。我们在50步指导蒸馏（FLUX.1-dev）和4步时间步长蒸馏（FLUX.1-schnell）变体上进行评估。

• *PixArt-ΣΣ(Chen et al., )是另一个基于DiT的模型。它堆叠了28个由自注意力、交叉注意力和前馈层组成的注意力块，总计6亿参数。我们在默认的20步设置上进行评估。

• *稳定扩散XL（SDXL）是一个广泛使用的基于UNet的模型，具有26亿参数（Podell et al., 2024）。它以三种分辨率预测噪声。最高分辨率阶段完全由ResBlocks（He et al., 2016）处理，而其他两个阶段联合使用ResBlocks和注意力层。与PixArt-ΣΣ类似，SDXL使用交叉注意力层进行文本条件化。我们在30步设置上进行评估，以及其4步蒸馏变体SDXL-Turbo（Sauer et al., 2023）。

数据集。 按照之前的工作（Li et al., ; Zhao et al., ），我们随机抽取COCO Captions 2024（Chen et al., 2015）中的提示进行校准。为了评估我们方法的泛化能力，我们采用了两个不同风格的提示集进行基准测试：

• *MJHQ-30K（Li et al., ）包含来自Midjourney的30K样本，分为10个常见类别，每个类别3K样本。我们从中均匀选择5K提示来评估模型在艺术图像生成上的性能。

• *密集标注图像（DCI）（Urbanek et al., 2024）是一个包含约8K图像的数据集，具有详细的人工标注字幕，平均超过1,000字。在我们的实验中，我们使用总结版本（sDCI），其中字幕通过大型语言模型（LLMs）浓缩为77个标记，以适应扩散模型。同样，我们随机抽取5K提示进行高效评估现实图像生成。

基线。 我们将SVDQuant与以下后训练量化（PTQ）方法进行比较：

• *4比特NormalFloat（NF4）是一种权重量化数据类型（Dettmers et al., 2023）。它假设权重服从正态分布，是信息理论上最优的4比特表示。我们使用社区量化的NF4 FLUX.1模型（Lilyasviel, 2024）作为基线。

• *ViDiT-Q（Zhao et al., ）使用每标记量化和平滑（Xiao et al., 2023）来缓解不同批次和标记之间的异常值，并在PixArt-ΣΣ上实现了无损的8比特量化。

• *MixDQ（Zhao et al., ）识别文本嵌入的句首标记中的异常值，并使用16比特预计算保护它们。该方法在SDXL-Turbo上实现了高达W4A8量化，性能下降可忽略。

• *TensorRT包含一个行业级的PTQ工具包，将扩散模型量化到8比特。它使用平滑，并仅在校准选定的时间步长范围内使用百分位方案校准激活。

结果

图7：MJHQ上的定性视觉结果。Image Reward在整个数据集上计算。在FLUX.1模型上，我们的4比特模型优于NF4 W4A16基线，展示了优越的文本对齐和更接近16比特模型的相似性。例如，NF4误解了“恐龙风格”，生成了一个真实的恐龙。在PixArt-Σ和SDXL-Turbo上，我们的4比特结果展示了明显优于ViDiT-Q和MixDQ的W4A8结果的视觉质量。

质量结果。 我们在表1中报告了不同模型和精度水平下的定量质量结果，并在图7中展示了一些相应的4比特定性比较。在所有模型中，我们的8比特结果可以完美反映16比特结果，PSNR高于21，击败了所有其他8比特基线。在FLUX.1-dev上，我们的INT8 PSNR甚至在MJHQ上达到27。

对于4比特量化，在FLUX.1上，我们的SVDQuant在Image Reward方面优于NF4 W4A16基线。在schnell变体上，我们的Image Reward甚至超过了原始BF16模型，表明更强的人类偏好。在PixArt-ΣΣ上，虽然我们的INT4 Image Reward略有下降，但我们的FP4模型实现了比FP16模型更高的分数。这可能是因为PixArt-ΣΣ的模型规模较小（6亿参数），已经非常紧凑，受益于较小的组大小。值得注意的是，我们的INT4和FP4结果在所有指标上一致优于ViDiT-Q的W4A8结果。对于基于UNet的模型，在SDXL-Turbo上，我们的4比特模型显著优于MixDQ的W4A8结果，FID分数与FP16模型相当，表明性能无损失。在SDXL上，我们的INT4和FP4结果在质量上与TensorRT的W8A8性能相当，代表了8比特的最先进水平。如图15所示，我们的视觉质量仅显示出轻微的下降。

图8：SVDQuant将12B FLUX.1的模型大小减少了3.6倍。此外，我们的引擎Nunchaku进一步将16比特模型的内存使用减少了3.5倍，并在桌面和笔记本电脑NVIDIA RTX 4090 GPU上比NF4 W4A16基线实现了3.0倍的加速。值得注意的是，在笔记本电脑4090上，通过消除CPU卸载，它总共实现了10.1倍的加速。

图9：我们的INT4模型无缝集成现成的LoRA，无需重新量化。应用LoRA时，它与原始16比特FLUX.1-dev的图像质量相匹配。详见附录C中的文本提示。

内存节省和加速。 在图8中，我们报告了FLUX.1的测量模型大小、内存节省和加速。我们的INT4量化将原始transformer大小从22.2 GiB减少到6.1 GiB，包括由于低秩分支导致的0.3 GiB开销，总体减少了3.6倍。由于权重和激活都被量化，与NF4权重量化变体相比，我们的推理引擎Nunchaku进一步节省了内存占用，并在桌面和笔记本电脑NVIDIA RTX 4090 GPU上实现了3.0倍的加速。值得注意的是，原始BF16模型在16GB笔记本电脑4090上需要逐层CPU卸载，而我们的INT4模型完全适应GPU内存，通过避免卸载实现了10.1倍的加速。我们预计FP4量化模型在NVIDIA的下一代Blackwell GPU上会有更大的加速，因为它们本身支持组量化的微缩放，无需专门的GPU内核。

与LoRA集成。 之前的量化方法在集成LoRA时需要融合LoRA分支并重新量化模型。相比之下，我们的Nunchaku消除了冗余的内存访问，允许添加一个单独的LoRA分支。在实践中，我们可以通过稍微增加秩将LoRA分支融合到我们的低秩分支中，进一步提高效率。在图9中，我们展示了一些将五种不同风格的LoRA（现实主义、吉卜力插画、动漫、儿童素描和纱线艺术）应用于我们的INT4 FLUX.1-dev模型的视觉示例。我们的INT4模型成功适应每种风格，同时保持16比特版本的图像质量。更多视觉示例见附录B.2。

消融研究。 在图10中，我们展示了PixArt-ΣΣ上SVDQuant的几项消融研究。首先，仅SVD和天真量化在4比特设置下表现不佳，导致图像质量严重下降。虽然对量化应用平滑略微提高了图像质量，但总体结果仍不令人满意。LoRC（Yao et al., 2023）引入了一个低秩分支来补偿量化误差，但这种方法次优。量化误差表现出平滑的奇异值分布。因此，低秩补偿无法有效缓解这些误差，如第4.2节所述。相反，我们首先分解权重并仅量化残差。如图5所示，前几个奇异值显著大于其余值，使我们能够将它们转移到低秩分支，这有效地减少了权重幅度。最后，平滑巩固了异常值，进一步使低秩分支能够吸收激活中的异常值，并显著提高图像质量。

增加秩的权衡。 图11展示了PixArt-ΣΣ上不同秩rr的SVDQuant结果。将秩从16增加到64显著提高了图像质量，但也增加了参数和延迟开销。在我们的实验中，我们选择秩为32，它在质量上提供了不错的折衷，开销较小。

图10：PixArt-Σ上SVDQuant的消融研究。低秩分支的秩为64。Image Reward在MJHQ的1K样本上测量。我们的结果显著优于其他方法，以较大的优势实现了最高的图像质量。

图11：增加SVDQuant中低秩分支的秩r可以提高图像质量，但也会导致更高的参数和延迟开销。

6 结论与讨论

在这项工作中，我们引入了一种新的4比特后训练量化范式，SVDQuant，用于扩散模型。它采用低秩分支来吸收权重和激活中的异常值，缓解量化过程。我们的推理引擎Nunchaku进一步融合了低秩和低比特分支内核，减少了内存使用并切断了冗余的数据移动开销。广泛的实验表明，SVDQuant保持了图像质量。Nunchaku进一步实现了比原始16比特模型减少3.5倍的内存使用，并在NVIDIA RTX-4090笔记本电脑上比权重量化实现了3.0倍的加速。这一进展使得大规模扩散模型能够在边缘设备上高效部署，解锁了更多交互式AI应用的潜力。

局限性。 在这项工作中，我们没有报告FP4模型的加速。这是因为我们无法访问原生支持精度和组量化微缩放的Blackwell GPU。在Blackwell硬件上，我们预计FP4模型比4090 GPU上的INT4结果有更大的加速。