笔记：理解 KL 散度与 INT8 量化校准

在做模型量化时，我们经常面临这样一个问题：激活值的分布往往是长尾的------大多数值集中在某个较小范围内，而少量极端值散布在两侧。如果我们直接使用全局最小最大值来做 INT8 量化，量化步长会被这些 outlier 拉得很大，导致中间主体部分的精度严重受损。

解决方案很直接：截断。我们只保留一个合适的范围 [−T,T][-T, T][−T,T] 或 [0,T][0, T][0,T]，把尾巴剪掉。但问题来了：TTT 取多少最合适？

截太大了，分辨率不够；截太小了，丢失太多信息。我们需要一个自动化的方法，帮我们在所有候选阈值中选出最优解。

这就是 KL 散度在量化校准中的作用。

一、核心思路一句话

KL 校准要做的事情是：对于每个候选阈值，模拟量化过程对分布造成的影响，然后选择那些让量化前后分布差异最小的阈值。

这个"差异"就是用 KL 散度来衡量的。

二、KL 散度快速回顾

KL 散度衡量的是用分布 QQQ 去近似分布 PPP 时的信息损失。它的定义是：

DKL(P∥Q)=∑xP(x)log⁡P(x)Q(x)D_{KL}(P \parallel Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}DKL(P∥Q)=x∑P(x)logQ(x)P(x)

需要记住的几个特点：

不对称 ：DKL(P∥Q)≠DKL(Q∥P)D_{KL}(P \parallel Q) \neq D_{KL}(Q \parallel P)DKL(P∥Q)=DKL(Q∥P)，因为方向不同，含义不同
非负：DKL≥0D_{KL} \geq 0DKL≥0，当且仅当 P=QP = QP=Q 时取 0
物理意义 ：表示"用 QQQ 代替 PPP 所需额外付出的信息代价"

在量化场景中，PPP 是原始浮点分布（截断后的），QQQ 是量化器能表达的近似分布。KL 越小，说明量化对分布的破坏越小。

三、算法全流程

我们一步步拆解。假设：

原始激活值直方图有 NNN 个细粒度 Bins（例如 2048）
量化后的等级数是 MMM（对称量化非负侧通常为 128 或 127）
候选阈值对应的截断位置为 iii（保留前 iii 个 Bins）

我们从 i=Mi = Mi=M 开始遍历到 i=Ni = Ni=N，对每个 iii 都做一遍下面的操作。

第一步：构造参考分布 P(i)P^{(i)}P(i)

这是"截断后的真实分布"。当我们决定把阈值设在第 iii 个 Bin 时，超过这个位置的概率质量不会消失，而是被 clip 到边界上。所以操作是：

保留前 iii 个细粒度 Bin
把索引 [i,N−1][i, N-1][i,N−1] 的所有概率质量累加到第 i−1i-1i−1 个 Bin 上

数学上：

Pj(i)={hist[j],0≤j<i−1∑t=i−1N−1hist[t],j=i−1 P_j^{(i)} = \begin{cases} \text{hist}[j], & 0 \leq j < i-1 \\[4pt] \displaystyle \sum_{t=i-1}^{N-1} \text{hist}[t], & j = i-1 \end{cases} Pj(i)=⎩ ⎨ ⎧hist[j],t=i−1∑N−1hist[t],0≤j<i−1j=i−1

然后对整个分布做归一化，使得 ∑j=0i−1Pj(i)=1\sum_{j=0}^{i-1} P_j^{(i)} = 1∑j=0i−1Pj(i)=1。

这一步模拟的是量化器对超出范围的值做裁剪的行为。

第二步：合并（压缩）到 MMM 个粗粒度 Bins

这一步模拟的是：量化器只有 MMM 个可区分的离散等级 ，它的分辨率远低于原始直方图。因此我们需要把 iii 个细粒度 Bin 压缩成 MMM 个粗粒度 Bin。

分区方式是对 [0,i−1][0, i-1][0,i−1] 均匀划分为 MMM 段。第 kkk 个粗粒度 Bin 覆盖的细粒度 Bin 范围是：

startk=⌊k⋅iM⌋,endk=⌊(k+1)⋅iM⌋−1\text{start}_k = \left\lfloor \frac{k \cdot i}{M} \right\rfloor, \quad \text{end}_k = \left\lfloor \frac{(k+1) \cdot i}{M} \right\rfloor - 1startk=⌊Mk⋅i⌋,endk=⌊M(k+1)⋅i⌋−1

该段长度为：

Lk=endk−startk+1L_k = \text{end}_k - \text{start}_k + 1Lk=endk−startk+1

粗粒度分布为：

Qk=∑j=startkendkPj(i),k=0,1,...,M−1Q_k = \sum_{j = \text{start}_k}^{\text{end}_k} P_j^{(i)}, \quad k = 0, 1, \dots, M-1Qk=j=startk∑endkPj(i),k=0,1,...,M−1

第三步：展开回 iii 个细粒度 Bins

现在有一个问题：KL 散度要求两个分布维度相同。P(i)P^{(i)}P(i) 有 iii 个元素，但 QQQ 只有 MMM 个元素，不能直接比较。

所以我们需要把粗粒度分布展开回 iii 个细粒度 Bin，得到一个 Q~(i)\tilde{Q}^{(i)}Q~(i)。

展开的原则是均匀分配：每个粗粒度 Bin 的概率质量，平均分给它覆盖的所有细粒度 Bin。

对于第 jjj 个细粒度 Bin，找到它所属的粗粒度 Bin kkk：

Q~j(i)=QkLk,其中 j∈[startk,endk] \tilde{Q}_j^{(i)} = \frac{Q_k}{L_k}, \quad \text{其中 } j \in [\text{start}_k, \text{end}_k] Q~j(i)=LkQk,其中 j∈[startk,endk]

等价地，可以用分段函数直接写：

Q~j(i)=∑t=startkendkPt(i)Lk,其中 k=⌊j⋅Mi⌋ \tilde{Q}j^{(i)} = \frac{ \sum{t = \text{start}_k}^{\text{end}_k} P_t^{(i)} }{ L_k }, \quad \text{其中 } k = \left\lfloor \frac{j \cdot M}{i} \right\rfloor Q~j(i)=Lk∑t=startkendkPt(i),其中 k=⌊ij⋅M⌋

为什么要均匀分配？

因为量化之后，所有落在同一个量化等级内的原始值都变得不可区分。量化器"看"到的是：这一段区间的概率质量总共是多少，但不知道内部的具体结构。均匀分配正是对这种不可分辨性的建模。

第四步：计算 KL 散度

现在我们可以计算 KL 散度了：

DKL(i)=∑j=0i−1Pj(i)log⁡Pj(i)Q~j(i)D_{KL}^{(i)} = \sum_{j=0}^{i-1} P_j^{(i)} \log \frac{P_j^{(i)}}{\tilde{Q}_j^{(i)}}DKL(i)=j=0∑i−1Pj(i)logQ~j(i)Pj(i)

注意：当 Pj(i)>0P_j^{(i)} > 0Pj(i)>0 且 Q~j(i)=0\tilde{Q}_j^{(i)} = 0Q~j(i)=0 时，KL 会趋向无穷大。这在数值上是不稳定的，所以在展开时通常要小心处理零值 Bin。

第五步：选择最优阈值

对所有候选阈值 iii 重复上述步骤，得到一组 KL 散度值：

i∗=arg⁡min⁡iDKL(i)i^* = \arg \min_i D_{KL}^{(i)}i∗=argiminDKL(i)

然后将索引 i∗i^*i∗ 转换为实际的浮点数阈值 TTT。

四、一个关键改进：展开时的零值处理

直接均匀展开有一个潜在问题：如果某个粗粒度 Bin 覆盖的细粒度 Bin 中有一些原本是零，把这些零 Bin 也赋予非零概率，会人为地引入量化器并不具备的信息，同时可能导致 KL 计算不稳定。

一种常用的改进是：只对非零的细粒度 Bin 分配概率。

设第 kkk 个粗粒度 Bin 覆盖的细粒度 Bin 集合为 Sk\mathcal{S}_kSk，其中非零概率的 Bin 集合为 Sk+={j∈Sk∣Pj(i)>0}\mathcal{S}_k^+ = \{ j \in \mathcal{S}_k \mid P_j^{(i)} > 0 \}Sk+={j∈Sk∣Pj(i)>0}。则展开公式改为：

Q~j(i)={Qk∣Sk+∣,if j∈Sk+ and ∣Sk+∣>00,otherwise \tilde{Q}_j^{(i)} = \begin{cases} \displaystyle \frac{Q_k}{|\mathcal{S}_k^+|}, & \text{if } j \in \mathcal{S}_k^+ \text{ and } |\mathcal{S}_k^+| > 0 \\[6pt] 0, & \text{otherwise} \end{cases} Q~j(i)=⎩ ⎨ ⎧∣Sk+∣Qk,0,if j∈Sk+ and ∣Sk+∣>0otherwise

这个改进在许多工业级实现（如 TensorRT 的早期版本）中被采用。

五、完整示例

我们用一个极小的例子走一遍流程。

设截断后保留 i=5i = 5i=5 个细粒度 Bin，概率分布为：

P=[0.1,0.2,0.3,0.2,0.2],∑P=1.0P = [0.1, 0.2, 0.3, 0.2, 0.2], \quad \sum P = 1.0P=[0.1,0.2,0.3,0.2,0.2],∑P=1.0

量化等级数 M=3M = 3M=3。

合并

按公式计算各段起止：

k=0k = 0k=0：start=⌊0×5/3⌋=0\text{start} = \lfloor 0 \times 5/3 \rfloor = 0start=⌊0×5/3⌋=0，end=⌊1×5/3⌋−1=0\text{end} = \lfloor 1 \times 5/3 \rfloor - 1 = 0end=⌊1×5/3⌋−1=0 →\rightarrow→ L0=1L_0 = 1L0=1，Q0=0.1Q_0 = 0.1Q0=0.1
k=1k = 1k=1：start=⌊1×5/3⌋=1\text{start} = \lfloor 1 \times 5/3 \rfloor = 1start=⌊1×5/3⌋=1，end=⌊2×5/3⌋−1=2\text{end} = \lfloor 2 \times 5/3 \rfloor - 1 = 2end=⌊2×5/3⌋−1=2 →\rightarrow→ L1=2L_1 = 2L1=2，Q1=0.2+0.3=0.5Q_1 = 0.2 + 0.3 = 0.5Q1=0.2+0.3=0.5
k=2k = 2k=2：start=⌊2×5/3⌋=3\text{start} = \lfloor 2 \times 5/3 \rfloor = 3start=⌊2×5/3⌋=3，end=⌊3×5/3⌋−1=4\text{end} = \lfloor 3 \times 5/3 \rfloor - 1 = 4end=⌊3×5/3⌋−1=4 →\rightarrow→ L2=2L_2 = 2L2=2，Q2=0.2+0.2=0.4Q_2 = 0.2 + 0.2 = 0.4Q2=0.2+0.2=0.4

粗粒度分布：Q=[0.1,0.5,0.4]Q = [0.1, 0.5, 0.4]Q=[0.1,0.5,0.4]

展开

j=0→k=0j = 0 \rightarrow k = 0j=0→k=0，Q~0=0.1/1=0.1\tilde{Q}_0 = 0.1 / 1 = 0.1Q~0=0.1/1=0.1
j=1→k=1j = 1 \rightarrow k = 1j=1→k=1，Q~1=0.5/2=0.25\tilde{Q}_1 = 0.5 / 2 = 0.25Q~1=0.5/2=0.25
j=2→k=1j = 2 \rightarrow k = 1j=2→k=1，Q~2=0.5/2=0.25\tilde{Q}_2 = 0.5 / 2 = 0.25Q~2=0.5/2=0.25
j=3→k=2j = 3 \rightarrow k = 2j=3→k=2，Q~3=0.4/2=0.2\tilde{Q}_3 = 0.4 / 2 = 0.2Q~3=0.4/2=0.2
j=4→k=2j = 4 \rightarrow k = 2j=4→k=2，Q~4=0.4/2=0.2\tilde{Q}_4 = 0.4 / 2 = 0.2Q~4=0.4/2=0.2

展开后分布：Q~=[0.1,0.25,0.25,0.2,0.2]\tilde{Q} = [0.1, 0.25, 0.25, 0.2, 0.2]Q~=[0.1,0.25,0.25,0.2,0.2]

计算 KL

DKL=0.1log⁡0.10.1+0.2log⁡0.20.25+0.3log⁡0.30.25+0.2log⁡0.20.2+0.2log⁡0.20.2=0+0.2log⁡(0.8)+0.3log⁡(1.2)+0+0≈−0.0446+0.0547=0.0101 \begin{aligned} D_{KL} &= 0.1 \log\frac{0.1}{0.1} + 0.2 \log\frac{0.2}{0.25} + 0.3 \log\frac{0.3}{0.25} + 0.2 \log\frac{0.2}{0.2} + 0.2 \log\frac{0.2}{0.2} \\ &= 0 + 0.2 \log(0.8) + 0.3 \log(1.2) + 0 + 0 \\ &\approx -0.0446 + 0.0547 \\ &= 0.0101 \end{aligned} DKL=0.1log0.10.1+0.2log0.250.2+0.3log0.250.3+0.2log0.20.2+0.2log0.20.2=0+0.2log(0.8)+0.3log(1.2)+0+0≈−0.0446+0.0547=0.0101

这就是该候选阈值对应的 KL 散度值。

六、再次理解：为什么要"压缩再展开"

你可能会有疑问：为什么不能直接在原始细粒度分布和粗粒度分布之间算 KL？原因很简单------维度不一致。KL 散度要求两个分布定义在同一个事件空间上。所以我们必须把粗粒度分布"投影"回细粒度空间，才能逐项比较。

这个过程可以类比图像处理：

压缩：把高分辨率图像（2048 像素）缩小到低分辨率（128 像素）
展开：再把低分辨率图像放大回原来的尺寸（2048 像素）

放大回去的图像虽然尺寸和原图一样，但细节已经丢失了------它是一张"被量化后的近似图像"。KL 散度比较的就是原始图像和这张"量化重建图像"之间的差异。

这个差异，就是量化带来的信息损失。

七、与 Min-Max 校准的对比

方法	思路	优点	缺点
Min-Max	直接取最小最大值作为量化范围	简单、快速	对 outlier 敏感，长尾分布效果差
KL 校准	找一个使量化前后分布差异最小的截断点	抗 outlier，对分布形状更友好	计算更复杂，依赖直方图

对于大多数实际场景（尤其是长尾分布），KL 校准的表现通常会优于 Min-Max。

八、总结

KL 校准寻找最优阈值的过程可以归纳为以下几步：

截断：对每个候选阈值，将尾部概率折叠到边界 Bin，得到参考分布 PPP
量化模拟 ：把 PPP 压缩到 MMM 个粗粒度 Bin，模拟 INT8 的有限表达能力
重建：将粗粒度分布均匀展开回原始粒度，得到近似分布 Q~\tilde{Q}Q~
评估：计算 DKL(P∥Q~)D_{KL}(P \parallel \tilde{Q})DKL(P∥Q~)
选取：选择使 KL 最小的候选阈值

整个过程的核心在于：我们不是在比较单个数值的误差，而是在比较整个分布形状的失真程度。 这才是 KL 散度在量化校准中发挥作用的关键所在。

希望这份笔记能帮你把 KL 校准的原理彻底理清。如果你在实际实现中遇到细节问题，欢迎继续深入讨论。