Trained Quantization Thresholds (TQT)

阅读说明

个人论文阅读报告留档：介绍发表于 MLSys 2020 的论文 TQT（Trained Quantization Thresholds），一种基于 QAT 学习均匀、对称、per-tensor 量化截断阈值的方法。

论文信息

Title: Trained Quantization Thresholds for Accurate and Efficient Fixed-point Inference of Deep Neural Networks
Venue: MLSys 2020
Author(s): Sambhav R. Jain, Albert Gural, Michael Wu, Chris H. Dick
Institution(s): Xilinx Inc, Stanford University
Link: arXiv:1903.08066
Project: Graffitist、Vitis-AI Quantizer

引言#

本文提出 Trained Quantization Thresholds（TQT）¹：一种基于量化感知训练（Quantization-Aware Training, QAT），学习均匀（Uniform）、对称（Symmetric）、张量级（per-tensor）量化器截断阈值的方法。TQT 在训练时使用直通估计器（Straight-Through Estimator, STE）²，并将量化步长限制为 2 的幂次（Power-of-Two, PoT）。采用 PoT 步长后，推理阶段可用整型加法与移位完成缩放，便于硬件部署。作者对 TQT 的鲁棒性做了数学分析，并在 ImageNet（ILSVRC 2012）图像分类任务上对多个 CNN 进行了实验。结果表明：在 MobileNet 等以往较难量化的网络上，TQT 往往只需少于 5 个 epoch 的再训练，即可达到接近浮点模型的精度。

本文主要关注以下三个问题：

此前大量工作依赖后训练量化（Post-Training Quantization, PTQ）的静态校准：校准得到的截断阈值在之后保持固定，相对 QAT 学到的阈值，泛化往往更弱。
PTQ 中按 per-tensor / per-channel 定义的量化误差指标（如 KLD、MSE）多属经验设计，与最终任务损失是否相关仍缺乏严格理论保证。相较优化人工指标的 PTQ，基于训练的方法通常更有效，也更易从优化目标上解释。
Google 的 Integer-Only Inference³ 方案中，比例系数 $M$ 为 dyadic 形式而非 PoT，仍可能带来高位整型存储与乘法开销。更理想的全整型推理可将量化步长限制为 PoT，从而用移位完全替代步长乘法，更利于硬件加速。

相关工作主要包括 STE²、Integer-Arithmetic-Only Inference³、PACT⁴，以及同期的 LSQ⁵ 等。

方法#

量化器约束#

TQT 采用的均匀量化公式如下：

r = s \cdot (q - z)

其中 $s$ 为量化步长， $z$ 为量化零点， $q$ 为量化后的整型数值。为简化运算，TQT 采用对称量化并令 $z=0$ ，从而省略含零点的展开项：

r = s \cdot q

同时，TQT 要求量化步长 $s$ 为 PoT 形式，即 $s=2^{-f}$ ，其中 $f$ 为 fractional length（可理解为定点小数点位置），是一个有符号整数。这样计算 $r = s \cdot q$ 时可用移位代替乘法。

前向传播#

TQT 的量化前向包含缩放（Scaling）、取整（Rounding）、截断（Saturation）、反量化（De-quantization）四步，量化函数为：

q(x;s)=\text{clip} \left( \left\lfloor \frac{x}{s} \right\rceil ;n,p \right)\cdot s

其中 $q(x;s)$ 为量化函数， $x$ 为输入张量， $s$ 为量化步长， $\text{clip}$ 为区间截断， $n$ 与 $p$ 分别为区间下界与上界。量化参数取值取决于输入是否有符号：

有符号（signed）：

n=-2^{b-1}, \quad p=2^{b-1}-1, \quad s=\frac{2^{\left\lceil \log_2 t \right\rceil}}{2^{b-1}}

无符号（unsigned）：

n=0, \quad p=2^b-1, \quad s=\frac{2^{\left\lceil \log_2 t \right\rceil}}{2^b}

其中 $t$ 为原始量化截断阈值；实际使用的截断范围取不小于 $t$ 的最小 PoT 值 $2^{\left\lceil \log_2 t \right\rceil}$ 。

反向传播#

训练中同时优化输入 $x$ 与量化步长 $s$ 。为便于写出梯度，可将前向中的 $\text{clip}$ 改写为分段形式：

q(x;s)= \begin{cases} \left\lfloor \dfrac{x}{s} \right\rceil \cdot s & \text{if } n\leq\left\lfloor \dfrac{x}{s} \right\rceil \leq p, \\[0.5em] n\cdot s & \text{if } \left\lfloor \dfrac{x}{s} \right\rceil < n, \\[0.5em] p\cdot s & \text{if } \left\lfloor \dfrac{x}{s} \right\rceil > p. \end{cases}

反向传播时，对不可微的取整运算采用 STE²：当 $\lceil x \rceil \neq x$ 、 $\lfloor x \rfloor \neq x$ 、 $\left\lfloor x \right\rceil \neq x$ 时，令

\frac{\partial \lceil x \rceil}{\partial x} = \frac{\partial \lfloor x \rfloor}{\partial x} = \frac{\partial \left\lfloor x \right\rceil}{\partial x} = 1.

由此得到 $q$ 关于 $s$ 的梯度：

\nabla_s q(x;s)= \begin{cases} \left\lfloor \dfrac{x}{s} \right\rceil - \dfrac{x}{s} & \text{if } n\leq\left\lfloor \dfrac{x}{s} \right\rceil \leq p, \\[0.5em] n & \text{if } \left\lfloor \dfrac{x}{s} \right\rceil < n, \\[0.5em] p & \text{if } \left\lfloor \dfrac{x}{s} \right\rceil > p. \end{cases}

注意到 $\nabla_{\log_2 t} s = s \ln 2$ ，上式可写为：

\nabla _{\log_2 t} q(x;s)=s \ln 2 \cdot \begin{cases} \left\lfloor \dfrac{x}{s} \right\rceil - \dfrac{x}{s} & \text{if } n\leq\left\lfloor \dfrac{x}{s} \right\rceil \leq p, \\[0.5em] n & \text{if } \left\lfloor \dfrac{x}{s} \right\rceil < n, \\[0.5em] p & \text{if } \left\lfloor \dfrac{x}{s} \right\rceil > p. \end{cases}

同理， $x$ 的梯度为：

\nabla _x q(x;s)= \begin{cases} 1 & \text{if } n \leq \left\lfloor \dfrac{x}{s} \right\rceil \leq p, \\[0.5em] 0 & \text{otherwise}. \end{cases}

梯度含义解读#

上面给出了量化器 $q$ 对 $x$ 与 $\log_2 t$ 的梯度。为直观理解 QAT 中的梯度行为，考虑用最小二乘损失

L = \frac{(q(x;s) - x)^2}{2}

优化一个 TQT 量化器。此时：

\begin{align*} \nabla _{\log_2 t} L &= (q(x;s) - x) \cdot \nabla _{\log_2 t} q(x;s) \\ \nabla _x L &= (q(x;s) - x) \cdot (\nabla _x q(x;s) - 1) \end{align*}

图 1 展示了 $b=3,\,t=1.0$ 时有符号与无符号情况下的函数值及各参数梯度曲线：第一行为量化函数，第二行为损失函数。由于取整 $\left\lfloor\cdot\right\rceil$ 与截断 $\text{clip}$ 共同作用，输入 $x$ 的有效范围实际为 $[s\cdot(n-0.5),\,s\cdot(p+0.5)]$ 。从图 1 可以看出：

当 $x$ 落在有效范围内时， $L$ 对 $\log_2 t$ 的梯度非负，对 $x$ 的梯度为 $0$ ；
当 $x$ 落在有效范围外时， $L$ 对 $\log_2 t$ 的梯度恒负，对 $x$ 的梯度左负右正。

图 2 展示了输入 $x$ 呈钟型分布（如高斯）时， $\log_2 t$ 梯度的典型情形：

左图：绝大多数 $x$ 落在有效范围内，对 $\log_2 t$ 的累计梯度为正，梯度下降使 $\log_2 t$ 减小并向分布中心靠近；
中图：落在范围外的样本对 $\log_2 t$ 产生负累计梯度，使阈值向外扩大；
右图：收敛时正负梯度相互抵消，此时截断阈值为 PoT，并能在精度损失与截断范围之间取得较好折中。

TQT 在 QAT 中也会同时优化输入 $x$ 。与阈值梯度类似：

范围内： $L$ 对 $x$ 的梯度为 $0$ ，范围内输入保持不变；
范围外： $L$ 对 $x$ 的梯度左负右正，使离群点向分布中心收缩。

与其它截断阈值梯度方法的比较#

与 TensorFlow FakeQuant 的比较#

图 3 FakeQuant 在 b=3, n=-1.125, p=0.875 时的函数与梯度

TensorFlow 伪量化模块 FakeQuant 参考的量化函数为：

q(x;n,p)=\left\lfloor \frac{\text{clip}(x;n,p)-n}{\frac{p-n}{2^b-1}} \right\rceil \cdot \frac{p-n}{2^b-1} + n

图 3 给出了 $b=3,\,n=-1.125,\,p=0.875$ 时 FakeQuant 的函数与梯度。与 TQT 不同，FakeQuant 直接学习截断端点 $n$ 与 $p$ ，反向时将取整近似为恒等映射。由图可见， $n$ 与 $p$ 的梯度恒为正，梯度下降只能使边界向外扩张，结果更接近 PTQ 的 min/max，难以在精度与截断范围之间主动折中。

与 PACT 的比较#

PACT⁴（PArameterized Clipping acTivation）通过可学习的 Clipped ReLU 将截断阈值引入 QAT：

y=\mathrm{PACT}(x)=0.5\big(|x|-|x-\alpha|+\alpha\big)= \begin{cases} 0, & x \in (-\infty, 0) \\ x, & x \in [0, \alpha) \\ \alpha, & x \in [\alpha, +\infty) \end{cases}

输出对截断阈值 $\alpha$ 的梯度经 STE 近似为：

\frac{\partial y_q}{\partial \alpha} = \begin{cases} 0, & x \in (-\infty, \alpha) \\ 1, & x \in [\alpha, +\infty) \end{cases}

在 PACT 中， $\alpha$ 的梯度仅取决于 $x$ 落在 $\alpha$ 的哪一侧：只有右侧样本会推动 $\alpha$ ，因而训练中 $\alpha$ 易趋向样本右端最大值。为抑制该现象，PACT 在损失中对 $\alpha$ 加入 L2 正则，但系数 $\lambda_{\alpha}$ 需人工设定，无法像 TQT 那样在训练中自动权衡截断与精度。

框架实现#

与本文配套的 Graffitist 基于 TensorFlow 实现了端到端的 TQT 量化流程，支持多种网络的计算图优化与量化。

Xilinx

graffitist

Waiting for api.github.com...

00K

Waiting...

计算图优化：包括 BN 融合、多重 concat 合并为单一 concat、将平均池化改写为深度卷积等。
量化模式：支持 PTQ 静态量化与基于 TQT 的 QAT 等模式。
层精度配置：可分别针对计算层、激活层、平均池化与 concat 等操作配置量化位宽。
融合内核：提供面向 CPU/GPU 预编译的融合量化内核，降低量化器内存占用，从而可用更大 batch size 加速训练。

实验#

此处仅展示主要结果；初始化与实现细节见原论文。

结果#

表 3 给出了 12 种网络在 ImageNet 分类任务上的验证准确率。wt 表示只再训练权重；wt,th 表示权重与截断阈值同时训练。所有再训练均不超过 5 个 epoch。

讨论#

来自 TQT 的观察#

QAT 精度通常高于 PTQ 静态量化，符合预期。
对 VGG、Inception、ResNet 等 INT8 相对容易的模型，固定阈值、只训权重往往已足够；再引入 TQT 未必带来额外收益。
对 MobileNet、DarkNet 等 INT8 较难的模型，同时训练权重与阈值可比只训权重高出约 $4\%$ ，甚至接近 FP32。
对更低比特的 INT4，只训权重已很难恢复精度，通常必须配合 TQT 才能稳住精度。

MobileNet 对比#

对较难量化的 MobileNet，已有工作表明：简单的对称、per-tensor PTQ 往往效果很差。原因之一是深度可分离卷积的权重分布不规则，且通道间动态范围差异大。使用 TQT（wt,th）做 QAT 后，精度可接近 FP32。与 Google QAT 等结果对比如表 1：尽管 TQT 对步长的约束更严（PoT），精度反而更高。

图 5 展示了 TQT 再训练前后 MobileNet v1 的权重分布，体现了截断范围与精度之间的折中。

图 6 INT8 / INT4 下各层 fractional length 的始末差异

图 6 给出 INT8 与 INT4 量化时各层 fractional length 的始末差异。INT8 下 fractional length 的变化往往更大：比特更充裕时方法更倾向扩大截断范围；比特更少时，扩大范围与保精度的折中更为尖锐。

总结#

本文介绍了一种通用的量化截断阈值学习方法 TQT，主要特点包括：

可将量化步长限制为 PoT，适配多数定点硬件；
学习过程显式体现截断范围与精度的折中，在 INT8 / INT4 上均可取得较高精度；
具有较好的鲁棒性与较快的收敛速度。

作者基于 TensorFlow 实现了 Graffitist，并在 ImageNet 上对多种经典模型做了量化实验，同时对若干实验现象进行了讨论。

附录#

附录内容简要如下。

A 仿射量化器的开销#

讨论量化器计算开销。A.1 分析量化乘法，说明省略零点 $z$ 可降低复杂度；A.2 比较浮点、dyadic 与 PoT 步长下的乘法：dyadic 见 Google Integer-Only Inference³，PoT 可用移位替代乘法。

B 对数阈值训练#

从数值稳定性、尺度不变性与收敛性讨论在 $\log_2 t$ 空间训练阈值：

B.1 数值稳定性（Numerical Stability）：相对直接优化 $t\in\mathbb{R}^+$ ，优化 $\log_2 t\in\mathbb{R}$ 更稳，且与 PoT 形式一致；
B.2 尺度不变性（Scale Invariance）：希望梯度尽量与取值量级解耦；
B.3 收敛性（Convergence）：并给出不同比特下 Adam 超参参考。

C Adam 收敛性#

进一步分析 TQT 在 Adam 上的收敛行为。

D 最优验证与平均验证#

解释结果筛选方式。MobileNetV1 与 VGG16 上，1000 次验证的平均准确率与最高准确率差距分别约为 $0.1\%$ 与 $0.2\%$ ，说明结果较为稳定可靠。

音乐

音乐

引言#