本文目录

1 HQ-DiT：高效的 FP4 混合精度量化 DiT
(来自纽约大学)
1 HQ-DiT 论文解读
1.1 在 DiT 中使用浮点数量化
1.2 扩散模型简介
1.3 扩散模型中的激活值是如何分布的？
1.4 给激活值乘以 Hadamard 变换
1.5 权重量化的 FP 格式选择
1.6 实验结果

太长不看版

本文提出了一种对 Diffusion Transformer (DiT) 的混合浮点数量化技术，属于是后训练量化 (Post-Training Quantization)，把权重量化为 4-bit 的浮点数以完成推理。因为 DiT 这种模型最近在工业和学术领域获得了极大的关注，因为其良好的视觉生成能力优于以 U-Net 为代表的传统扩散模型。但是，DiT 也伴随着较高的参数计数和实现成本，严重限制了它们在手机等资源有限设备上的使用。与定点量化 (例如 INT8) 相比，浮点数量化，再加上本文提出的裁剪范围选择机制，这个机制可以自然地与 DiT 中的数据分布对齐，使得量化误差最小。本文的结果表面：DiT 可以实现极低的精度量化 (4 位)，对性能的影响可以忽略不计。

本文做了什么工作

1. 提出了一种为 DiT 设计的高效的 PTQ 方法，能使用 4-bit 浮点数 (FP4) 量化 DiT，并实现与全精度模型相当的性能。据作者所知，HQ-DiT 首次尝试使用 FP 数据格式量化 DiT。
2. 提出了一种新的算法，该算法可以根据数据分布自适应地选择最佳 FP 格式，能够节约计算开销。
3. HQ-DiT 使用 FP4 量化 DiT 中的权重和激活值，与全精度模型相比，加速 5.09 倍，内存节约 2.13 倍。HQ-DiT在低精度量化方面取得了 SOTA 的结果，FP4 模型在 Inception Score (IS) 和 Frechet Inception Distance (FID) 方面都优于全精度的 latent diffusion model (LDM)。

1 HQ-DiT：高效的 FP4 混合精度量化 DiT

论文名称：HQ-DiT: Efficient Diffusion Transformer with FP4 Hybrid Quantization

论文地址：

http://arxiv.org/pdf/2405.19751

1.1 在 DiT 中使用浮点数量化

Diffusion Transformers (DiTs[1]) 因其卓越的性能而备受学术界和工业界的关注，比较引人关注的应用是 OpenAI 的 SoRA[2]。最近的研究 Lumina-T2X[3]也已经证明了它在不同模式的令人印象深刻的生成能力。但是，迭代去噪步骤和巨大的计算需求会减慢其推理速度。而且，在每个 Denoising timestep 上，DiT 模型大量参数和复杂的网络结构仍会带来很大的计算负担，阻碍了它们在资源有限的实际场景中的适用性。

模型量化通过将权重和激活值压缩到低位表示，来减少内存和计算负担。在各种量化方法中，后训练量化 (Post-Training Quantization, PTQ) 是一种无训练的方法，用于快速有效的量化。与需要多轮微调的量化感知训练 (Quantization-Aware Training, QAT) 相比，PTQ 可以显著降低计算成本。这使得量化像 DiT 这样的大模型成为一个很有吸引力的解决方案。现有的对于扩散模型的 PTQ 方法主要采用定点数量化 (即 INT 量化)。但是这在低比特的情况下会带来比较大的量化误差。为了验证这点，作者将几个为大语言模型设计的量化方法用在 DiT 上，包括 SmoothQuant[4], FPQ[5], 和 GPTQ[6]。如图 1 所示，将权重和激活量化为 4 位精度会导致严重的性能下降，导致 FID 的增加高达 100。此外，由于搜索最优量化方案所需的校准过程，这些方法所涉及的量化过程也会产生较高的计算成本。

浮点数 (FP) 量化为整数 (INT) 量化提供了一种更灵活的替代方案。相比整数使用固定的 scaling factor，FP 包含指数，对不同范围的数据自适应。这种适应性使其能够保持不同量级的精度，使 FP 成为各种商业硬件平台的理想数据格式选择，如 Nvidia 的 Blackwell (H200)。

但是，与 INT 量化通过截断和舍入得到量化值不同，FP 量化的难点在于选择适当的组合：有多少位被分配给指数，有多少位被分配给尾数。如果选择没选好，可能会导致性能不佳。为了解决这个问题，作者基于 channel 数据分布来确定 FP 组合。与之前的 LLM-FP4[7]相比，本文方法通过定制每个 channel 中数据的独特特征来优化性能。由于较高的通道间方差和较低的通道内方差，量化激活值是一个挑战[8]。

1.2 扩散模型简介

扩散模型 (DM) 因其生成各种逼真图像的显著能力而备受关注。在 DM 前向过程, 给定输入图像 , 生成一系列高斯噪声并将其添加到 中, 从而产生一系列噪声样本 。

式中， 是控制每一步高斯噪声强度的超参数。

在反向过程中，给定一个随机采样的高斯噪声 。合成图像通过以下过程逐步生成：

式中, 和 通过下式来定义:

上式中， 表示使用 DNN 生成的预测噪声。DiT 模型用 Transformer 替换经常使用的 U-Net。DiT 在视觉生成能力方面被认为超过了传统的扩散模型。

Classifier-free Guidance

在 DiT 中，用户可以提供一个额外的类标签 cc 作为图像生成的指导。在这种情况下，反向过程变为：

Classifier-free Guidance[9]使用隐式分类器替换显式分类器，调整引导权重来控制生成图像的真实性和平衡多样性。根据贝叶斯公式，分类器的梯度可以表示为：

CFG 可以控制生成样本的真实性和多样性的平衡，广泛应用于 DALL·E 等生成模型。

1.3 扩散模型中的激活值是如何分布的？

为了了解 DiT 中输入激活的分布，作者在 50 个 Denosing Step 中收集了 DiT Block 的输入激活值，如图 2 所示。然后在特定时间步分析激活矩阵的分布，如图 3 所示。

与 Q-Diffusion[10]观察到的模式类似，Outlier 存在于整个 channel 级别。具体而言，激活中的异常值规模大约是激活值其余部分的 100 倍。因此，每个 token 量化时不可避免地会引入大量错误。鉴于激活在不同 channel 中表现出高方差，但 channel 内的方差较低，Outlier 会被限制在有限数量的 channel 中。因此，作者认为 "Per-Channel Quantization" + 最小化异常值可能可以实现较低的量化误差。

1.4 给激活值乘以 Hadamard 变换

之前的工作 已经证明 Hadamard 变换可以应用于消除数据中存在的异常值。本文引入随机 Hadamard 变换来消除 DiT 的输入激活 中存在的异常值, 方法是将其与正交 Hadamard 矩阵 相乘，其中 是 DiT 的 Embedding Dimension， 满足 。矩阵 就会表现出更平滑的分布, 即：大多数异常值被消除, 如图 3 的右侧所示。为了保持 DiT Block 内线性层的数学等价性, 需要将每个 DiT Block 的 Self-Attention 层和 FFN 层中的相应权重矩阵应用 Hadamard 矩阵。

为了数学上的等价性, Query, Key 和 Value 矩阵乘以 : 。这个过程在推理之前进行, 不会产生额外的在线计算成本。这样子一来，输出就不会因为 Hadamard 矩阵而发生改变： 。而且, 把输入激活值从 修改为了 ，相当于我们去量化后者, 量化误差要小得多, 如图 3 右侧所示。

然后, Attention 和 Value 的乘积这里还需要再量化一下, SA 的输出这里可以表示为 , 这里的 不是 Value 矩阵, 而是应该写成 ,这里 是 head 的数量。这里面的 才是 Value 矩阵。

那么现在要把 量化之后, 再与 相乘, 那么会产生上述的 Outlier 的问题。因此, 就有了下面的操作:

通过引入这两个 矩阵, 可以克服 Outlier 对量化引入的问题: 可以有效减轻 内的异常值, 而不需要任何额外的计算开销。但是要注意这里的 , 其中 是 Identity 矩阵。

此外, 在 FFN 模块里面, 也是用的 Hadamard 矩阵。这里面包含两个线性映射矩阵 和 还有一个 GELU 激活函数。第一个线性矩阵可以 , 因为输入 会变为 。该操作可以离线执行, 在推理过程中没有任何额外的成本。第1个线性层 的输出还包含异常值, 需要额外的 Hadamard 变换来去除异常值。由于 GELU 激活函数的存在, Hadamard 变换需要在线执行。为了减轻这种成本, 本文设计了一个高效的 Hadamard 变换, 其计算复杂度如下:

如果 的维度 是 2 的幂, Walsh-Hadamard 变换时间复杂度计算 的时间复杂度是 。如果 不是 2 的幂, 通过利用 Kronecker 乘积的性质, 时间复杂度可以变为 。因此可以通过 Kronecker 方法构建一个新的 Hadamard 矩阵（比如 , 这里 为 为 为 28)。

1.5 权重量化的 FP 格式选择

权重量化的目的是找到满足下式的权重：

选择合适的 FP 组合进行权重量化是至关重要的，因为指数和尾位宽的选择不当会导致显著的量化误差。

选择 FP 数据格式的简单方法涉及详尽地搜索每个可能的组合，但这会导致很高的计算开销。作者提出了一种简单高效的 FP 格式选择方法。本文方法基于一个简单的事实，即对于具有固定总位宽的 FP 格式，更多的指数位允许更大的可表示的值的范围，如下图5所示。但是，较大的指数位宽 (例如 E2M1) 也会导致表示数据的分布不均匀，使得长尾现象更加明显。比如下图 E2M1 的最大值表示 12 (bias=0 的情况下)，E1M2 的最大值表示 3.5 (bias=0 的情况下)。

对于 DiT 模型中权重矩阵，可以通过以下指标分析它们的数据分布来确定最佳 FP 格式：

式中, 是一个表示分位数的超参数 (作者使用 25)。通过上式, 基本上能够得到一个表示权重中的最大值与最小值关系的指标 。本文的目标是用低误差表示该区间内的所有元素值。另一方面, 对于给定的指数和尾宽分别为 和 的 FP 格式, 最大值和最小值之间的比率可计算如下:

其中, 最大值是 , 最小值是 min_val 。目标是找到最优的 FP 格式, 使 接近于 。

然后作者使用前面描述的 FP 组合搜索方法为每个 DiT Block 找到了最佳的 FP 组合。

激活值的量化

对激活值的量化使用简单的 MinMax 量化方法，如下图 6 所示。MinMax 量化根据激活通道的最大值设置 bias 参数，然后通过除以 scale 参数将全精度数据映射到量化的数据范围。

如下图 7 所示，HQ-DiT 应用于每个 DiT Block 内的所有全连接层。FP4 和 FP32 的计算分别对应蓝色和红色。

1.6 实验结果

本文的实验是 ImageNet 无条件图像生成和条件图像生成，分辨率包括 256×256 和 512×512。作者还在 DiT 上实现了其他复杂的量化方法，包括 SmoothQuant[4], FPQ[5], 和 GPTQ[6]。

本节中，作者展示了使用 HQ-DiT 的条件图像生成和无条件图像生成的结果。在无条件图像生成中，DiT 去噪过程中没有提供外部指导。对于条件图像生成，将特定的类标签作为额外的输入来指导 DiT 的去噪过程。

无条件图像生成结果

结果如图 7 所示。虽然 SmoothQuant 和 FPQ 对 8 位权重和 8 位激活是有效的 (W8A8)，但它们的性能随着 4 位权重和 4 位激活而显著下降。相比之下，HQ-DiT 表现出卓越的性能，在 W8A8 精度上实现了最高的 IS 和 FID。此外，与全精度 (FP32) 结果相比，FID 仅略有增加。本文方法即使在 W4A4 精度下也能保持执行无条件生成任务的能力。这种显著的改进突出了 HQ-DiT 减轻异常值的影响并为每层选择最佳数据格式的能力。

条件图像生成结果

结果如图 8 所示。ImageNet 256×256 上的结果表明，采用 2 种不同的 Classifier-Free Guidance Scale，即 1.5 和 4.0。同样，作者提供了 3 个指标的综合评估结果，包括 IS、FID 和 sFID。可以观察到，对于 4.0 的 cfg，与 FP32 DiT 相比，HQ-DiT 在 W4A8 的设置下，FID 下降了 6.98，sFID 下降了 2.2。与其他基线算法相比，HQ-DiT 也实现了最佳性能。

作者进一步评估了 DiT 在 ImageNet 512×512 的性能。时间步长的总数设置为 20，而其余设置与 ImageNet 256×256 的设置相同。根据图 8 所示的评估结果，本文 W4A4 模型达到了 370.69 的 IS，FID 相对于 32 位的基线模型提升了 6.82，同时在基线算法中获得了最佳性能。

消融实验结果

Hadamard 变换对图像质量的影响

作者首先研究了 Hadamard 变换对不同设置下量化模型的性能的影响。作者评估了使用或者不使用 Hadamard 变化的情况下，HQ-DiT 生成的图像的质量。在没有 Hadamard 变换的情况下，HQ-DiT 在 W4A4 上无法生成高质量的图像，只能得到 3.14 的 IS 和 200.85 的 FID。相反，通过引入 Hadamard 变换，本文方法实现了 47.13 的 IS 和9.94 的 sFID，性能水平与全精度模型相当。

FP 格式选择消融实验

此外，作者通过实验来分析 FP 格式选择的效果。作者固定整个 DiT 层的 FP 组合，将 HQ-DiT 与其他基线进行比较，得到 3 个基线，包括 E3M0、E2M1 和 E1M2。图 9 所示的结果表明，本文的混合 FP 格式实现了 437.13 的 IS，优于所有其他 FP4 格式。本文的 FID 为 9.25，接近 E1M2 的 FID 结果 (9.20)。这些结果表明本文的 FP 组合选择和混合 FP 量化方案可以有效地提高 DiT 的性能。

本文介绍了一种新颖的高效混合精度量化方法HQ-DiT，该方法使用4位浮点数对扩散变换器（DiT）进行量化，显著改善了其在推理过程中的性能和资源使用效率。

1. ^Scalable Diffusion Models with Transformers
2. ^Sora: Creating video from text
3. ^Lumina-t2x: Transforming text into any modality, resolution, and duration via flow-based large diffusion transformers
4. ^abSmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
5. ^abAFPQ: Asymmetric Floating Point Quantization for LLMs
6. ^abGPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
7. ^LLM-FP4: 4-Bit Floating-Point Quantized Transformers
8. ^LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale
9. ^Classifier-Free Diffusion Guidance
10. ^Q-Diffusion: Quantizing Diffusion Models
11. ^QuIP#: Even Better LLM Quantization with Hadamard Incoherence and Lattice Codebooks

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT