本文目录

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)。通过上式, 基本上能够得到一个表示权重中的最大值与最小值关系的指标 。本文的目标是用低误差表示该区间内的所有元素值。另一方面, 对于给定的指数和尾宽分别为