I-ViT: 用于高效视觉Transformer推理的纯整数量化

Li Z, Gu Q. I-vit: Integer-only quantization for efficient vision transformer inference[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 17065-17075.

1. 引言

视觉Transformer（ViTs）在计算机视觉任务中展现出了卓越的性能，但其庞大的内存占用和计算开销严重阻碍了在资源受限的边缘设备上的部署。模型量化作为一种有效的压缩方法，通过降低权重和激活参数的表示精度来提升模型效率。现有的整数推理管线基于卷积神经网络（CNNs）的齐次性条件设计，然而这一条件并不适用于ViTs中的非线性组件，使得ViTs的纯整数推理成为一个亟待解决的开放性问题。

本文提出I-ViT，这是首个专门为ViTs设计的纯整数量化方案。I-ViT使得ViTs能够在整个推理过程中仅使用整数算术和位移操作，完全避免浮点运算。线性操作（如MatMul和Dense）遵循带有二进制算术的整数管线，而非线性操作（如Softmax、GELU和LayerNorm）则通过我们提出的轻量级整数算术方法进行近似。

2. 方法论框架

2.1 整体架构

图1描述：该图对比了FasterTransformer和I-ViT在处理非线性操作时的计算流程。图1(a)显示FasterTransformer需要在Softmax、GELU和LayerNorm操作前进行反量化（Dequantization）转换到FP32精度，执行浮点运算后再量化回INT8。而图1(b)展示了I-ViT的创新之处——通过Shiftmax、ShiftGELU和I-LayerNorm实现纯整数运算，整个过程保持在INT8精度，避免了浮点运算和精度转换的开销。

ViTs的主体结构是一系列堆叠的块，每个块包含多头自注意力（MSA）模块和多层感知器（MLP）模块：

$$\hat{X} = \text{MSA}(\text{LayerNorm}(X)) + X$$

$$Y = \text{MLP}(\text{LayerNorm}(\hat{X})) + \hat{X}$$

MSA模块通过计算全局注意力学习patch间的表示：

$$\text{MSA}(X) = \text{Concat}(\text{Attn}_1, \text{Attn}_2, ..., \text{Attn}_h)W^O$$

其中每个注意力头计算为：

$$\text{Attn}_i = \text{Softmax}\left(\frac{Q_i \cdot K_i^T}{\sqrt{d}}\right)V_i$$

这里$Q_i = XW_i^Q$，$K_i = XW_i^K$，$V_i = XW_i^V$是通过线性投影获得的查询、键和值矩阵。

图3描述：该图展示了I-ViT的完整计算流程。整个计算图完全使用整数算术实现，其中线性的MatMul和Dense操作遵循二进制算术管线，而提出的Shiftmax、ShiftGELU和I-LayerNorm完成非线性操作。图中清晰标注了数据流的精度：除了特别标记的INT32位置，其余数据流均为INT8精度。重量化（Requantization）步骤通过$I_A = (I_A' \cdot b) \gg c$实现，其中$b/2^c = DN(S_Q \cdot S_K / S_A)$。

2.2 量化方案

我们采用最简单的对称均匀量化：

$$I = \left\lfloor \frac{\text{clip}(R, -m, m)}{S} \right\rceil, \quad \text{where} \quad S = \frac{2m}{2^k - 1}$$

其中$R$和$I$分别表示浮点值和量化后的整数值，$S$是缩放因子，$m$是通过min-max方法确定的截断值，$k$是量化位精度，$\lfloor \cdot \rceil$是舍入操作。

3. 线性操作的二进制算术

对于线性操作，当输入为$Q = (S_Q, I_Q)$和$K = (S_K, I_K)$时，输出计算过程如下：

$$A' = S_{A'} \cdot I_{A'} = S_Q \cdot S_K \cdot (I_Q * I_K^T)$$

其中$I_{A'} = I_Q * I_K^T$执行纯整数算术。根据实际硬件实现原则（如DP4A），当输入$I_Q$和$I_K$为INT8类型时，输出$I_{A'}$为INT32类型。因此需要将$I_{A'}$重量化为INT8类型：

$$I_A = \left\lfloor \frac{S_{A'} \cdot I_{A'}}{S_A} \right\rceil = \left\lfloor \frac{S_Q \cdot S_K}{S_A} \cdot (I_Q * I_K^T) \right\rceil$$

为避免浮点运算，将缩放转换为二进制数（Dyadic Number）：

$$DN\left(\frac{S_Q \cdot S_K}{S_A}\right) = \frac{b}{2^c}$$

其中$b$和$c$都是正整数。这样缩放可以高效地通过整数乘法和位移完成：

$$I_A = (b \cdot (I_Q * I_K^T)) \gg c$$

4. Shiftmax: 整数化Softmax

4.1 理论推导

Softmax在ViTs中将注意力分数转换为概率：

$$\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^d e^{x_j}} = \frac{e^{S_{x_i} \cdot I_{x_i}}}{\sum_{j=1}^d e^{S_{x_j} \cdot I_{x_j}}}$$

由于非线性特性，Softmax无法直接遵循二进制算术管线。为解决这个问题，我们首先平滑数据分布并防止溢出：

$$\text{Softmax}(x_i) = \frac{e^{S_{\Delta_i} \cdot I_{\Delta_i}}}{\sum_j e^{S_{\Delta_j} \cdot I_{\Delta_j}}} = \frac{e^{S_{x_i} \cdot (I_{x_i} - I_{\max})}}{\sum_j e^{S_{x_j} \cdot (I_{x_j} - I_{\max})}}$$

其中$I_{\max} = \max\{I_{x_1}, I_{x_2}, ..., I_{x_d}\}$，$I_{\Delta_i} = I_{x_i} - I_{\max} \leq 0$。

4.2 基数转换与近似

为充分利用硬件移位器，我们将指数基数从$e$转换为2。由于$\log_2 e$可以近似为二进制$(1.0111)_b$，浮点乘法可通过整数移位实现：

$$e^{S_\Delta \cdot I_\Delta} = 2^{S_\Delta \cdot (I_\Delta \cdot \log_2 e)} \approx 2^{S_\Delta \cdot (I_\Delta + (I_\Delta \gg 1) - (I_\Delta \gg 4))}$$

记幂次项为$S_\Delta \cdot I_p$，其中$I_p = I_\Delta + (I_\Delta \gg 1) - (I_\Delta \gg 4)$。由于$S_\Delta \cdot I_p$不保证是整数，我们将其分解：

$$2^{S_\Delta \cdot I_p} = 2^{(-q) + S_\Delta \cdot (-r)} = 2^{S_\Delta \cdot (-r)} \gg q$$

其中$S_\Delta \cdot (-r) \in (-1, 0]$是小数部分，$q$和$r$都是正整数。

4.3 线性近似与整数除法

对于小数部分$2^{S_\Delta \cdot (-r)}$，我们使用线性函数近似：

$$2^{S_\Delta \cdot (-r)} \approx \frac{S_\Delta \cdot (-r)}{2} + 1 = S_\Delta \cdot [((-r) \gg 1) + I_0]$$

其中$I_0 = \lfloor 1/S_\Delta \rceil$。完成分子近似后，即$S_\Delta \cdot I_{\exp} \approx e^{S_\Delta \cdot I_\Delta}$，通过分数化简可以去除$S_\Delta$。最终通过整数除法得到输出：

$$I_{out_i} = \text{IntDiv}(I_{\exp_i}, \sum_{j} I_{\exp_j}, k_{out}) = \left(\left\lfloor \frac{2^M}{\sum_j I_{\exp_j}} \right\rfloor \cdot I_{\exp_i}\right) \gg (M - (k_{out} - 1))$$

$$S_{out_i} = 1/2^{k_{out}-1}$$

5. ShiftGELU: 整数化GELU

5.1 GELU近似

GELU激活函数可以通过sigmoid函数近似：

$$\text{GELU}(x) = x \cdot \frac{1}{\sqrt{2\pi}} \int_{-\infty}^x e^{-t^2/2}dt \approx x \cdot \sigma(1.702x) = S_x \cdot I_x \cdot \sigma(S_x \cdot 1.702I_x)$$

首先，1.702可以近似为二进制$(1.1011)_b$，因此$1.702I_x$可通过整数移位实现：

$$I_p = I_x + (I_x \gg 1) + (I_x \gg 3) + (I_x \gg 4)$$

5.2 Sigmoid的整数化

我们等价变换sigmoid函数：

$$\sigma(S_x \cdot I_p) = \frac{1}{1 + e^{-S_x \cdot I_p}} = \frac{e^{S_x \cdot I_p}}{e^{S_x \cdot I_p} + 1}$$

进一步变换得到：

$$\sigma(S_x \cdot I_p) = \frac{e^{S_x \cdot (I_p - I_{\max})}}{e^{S_x \cdot (I_p - I_{\max})} + e^{S_x \cdot (-I_{\max})}}$$

这里分子与Shiftmax中的分子形式完全相同，因此两者实现一致。完成整数除法后，再乘以$S_x \cdot I_x$即得到GELU的整数近似。

6. I-LayerNorm: 整数化层归一化

LayerNorm在隐藏特征维度上归一化输入：

$$\text{LayerNorm}(x) = \frac{x - \text{Mean}(x)}{\sqrt{\text{Var}(x)}} \cdot \gamma + \beta$$

其中$\gamma$和$\beta$是线性仿射因子。与BatchNorm不同，LayerNorm需要在推理阶段动态计算统计量。虽然整数算术单元可以直接计算均值和方差，但不支持平方根运算。

我们通过位移改进了轻量级整数迭代方法：

$$I_{i+1} = \frac{I_i + \lfloor \text{Var}(I_x)/I_i \rfloor}{2} = (I_i + \lfloor \text{Var}(x)/I_i \rfloor) \gg 1$$

其中$I_i$是第$i$次迭代的结果，$I_0$初始化为$2^{\lfloor \text{bit}(\text{Var}(I_x))/2 \rfloor}$。朴素的停止准则是$I_{i+1} \geq I_i$，但这无法保证恒定延迟。实验发现10次迭代可实现大部分收敛，因此我们将停止准则修改为迭代次数以便于硬件实现。

7. 实验结果与分析

7.1 准确率评估

图2描述：该图展示了I-ViT和FP基线在DeiT和Swin模型上的准确率-速度曲线。横轴表示每秒处理的图像数（Images/s），纵轴表示Top-1准确率（%）。可以看到，整数化的DeiT和Swin（实线）相比FP基线（虚线）在保持相似甚至略高准确率的同时，实现了3.72～4.11倍的显著加速。这充分证明了I-ViT在精度和速度之间取得了优秀的平衡。

表1展示了详细的准确率和延迟结果。I-ViT在各种模型上都保持了与FP32基线相当的准确率。例如，DeiT-S通过I-ViT量化后达到80.12%的Top-1准确率，比FP32基线（79.85%）还高0.27%。相比之下，FasterTransformer由于保留了浮点运算，加速效果有限（仅3.53倍）；I-BERT由于近似不匹配，准确率下降明显（-0.74%）。

7.2 延迟评估

图4描述：该图展示了DeiT-S在不同批大小（8、16、32、64）下的延迟结果。纵轴使用对数刻度表示延迟（ms）。可以观察到，I-ViT（紫色柱）在所有批大小下都保持了相对于基线（蓝色柱）、FasterTransformer（橙色柱）和I-BERT（绿色柱）的恒定加速比例。这表明I-ViT对批大小具有良好的鲁棒性，能够在不同的部署场景下保持稳定的加速效果。

硬件实现细节：我们使用TVM在RTX 2080Ti GPU上部署I-ViT。尽管GPU不是纯整数硬件，但依靠DP4A指令，I-ViT可以在Turing Tensor Cores上执行高效的纯整数推理。值得注意的是，TVM的软件支持和Turing Tensor Cores的硬件支持并非最优——批大小增加时并未实现完全并行化。因此，在专用硬件（如FPGA）上部署I-ViT将进一步提升加速潜力。

7.3 消融研究

表2的消融研究揭示了各组件的重要性。将Shiftmax和ShiftGELU替换为I-BERT的二阶多项式近似会导致准确率下降和延迟增加。在DeiT-B上，多项式GELU导致0.86%的准确率损失和0.17ms的延迟增加。这是因为多项式GELU仅对特定区间进行近似，不适用于ViTs的全定义域。

使用L1 LayerNorm虽然略微减少了延迟（-0.02ms），但导致了严重的准确率下降（DeiT-B上-2.49%，Swin-S上-3.32%）。这说明我们提出的I-LayerNorm在保持准确率方面的重要性。

LIS（Log-Int-Softmax）方法基于I-BERT并添加了对数函数，但同样遇到了数据分布不匹配的问题，准确率和延迟表现都不如Shiftmax。

8. 结论

I-ViT首次实现了ViTs的完整计算图量化，使整个推理过程仅需整数运算和位移操作，完全避免了浮点运算。通过提出的Shiftmax、ShiftGELU和I-LayerNorm，非线性操作得以高效地用整数算术实现。实验结果表明，I-ViT在保持甚至略微提升准确率的同时，实现了3.72～4.11倍的推理加速，为ViTs在资源受限设备上的部署提供了实用的解决方案。

附录：数学推导

A. Shiftmax完整推导

A.1 指数函数的二进制近似

从指数函数的换底公式开始：

$$e^x = 2^{x \cdot \log_2 e}$$

其中$\log_2 e = 1.4426950408889...$。为了使用整数运算，我们需要找到$\log_2 e$的二进制近似。通过分析：

$$\log_2 e \approx 1 + \frac{1}{2} - \frac{1}{16} = 1.4375$$

误差为$|1.4426950408889 - 1.4375| = 0.0051950408889 < 0.006$，相对误差小于0.4%。

因此：

$$e^{S_\Delta \cdot I_\Delta} \approx 2^{S_\Delta \cdot I_\Delta \cdot (1 + 1/2 - 1/16)} = 2^{S_\Delta \cdot (I_\Delta + I_\Delta/2 - I_\Delta/16)}$$

用位移表示：

$$I_p = I_\Delta + (I_\Delta \gg 1) - (I_\Delta \gg 4)$$

A.2 幂次分解的数学基础

对于$2^{S_\Delta \cdot I_p}$，其中$S_\Delta \cdot I_p$不是整数，我们需要分解为整数和小数部分。设：

$$S_\Delta \cdot I_p = -q + f$$

其中$q = \lfloor -S_\Delta \cdot I_p \rfloor$是正整数（因为$I_p \leq 0$），$f \in [0, 1)$是小数部分。

为了计算方便，我们重新参数化：

$$S_\Delta \cdot I_p = -q + S_\Delta \cdot (-r)$$

其中$r = -f/S_\Delta$，满足$S_\Delta \cdot (-r) \in (-1, 0]$。

因此：

$$2^{S_\Delta \cdot I_p} = 2^{-q} \cdot 2^{S_\Delta \cdot (-r)} = \frac{2^{S_\Delta \cdot (-r)}}{2^q}$$

A.3 小数幂次的Taylor展开

对于$2^{S_\Delta \cdot (-r)}$，其中$S_\Delta \cdot (-r) \in (-1, 0]$，使用Taylor展开：

$$2^x = 1 + x\ln 2 + \frac{(x\ln 2)^2}{2!} + \frac{(x\ln 2)^3}{3!} + ...$$

对于$x \in (-1, 0]$，保留前两项：

$$2^x \approx 1 + x\ln 2 = 1 + 0.693147x$$

为了使用整数运算，我们近似$\ln 2 \approx 1/2$（误差约为27.7%，但在小区间内可接受）：

$$2^{S_\Delta \cdot (-r)} \approx 1 + \frac{S_\Delta \cdot (-r)}{2}$$

重新整理：

$$2^{S_\Delta \cdot (-r)} \approx S_\Delta \cdot \left[\frac{(-r)}{2} + \frac{1}{S_\Delta}\right] = S_\Delta \cdot [((-r) \gg 1) + I_0]$$

其中$I_0 = \lfloor 1/S_\Delta \rceil$。

B. ShiftGELU的Sigmoid推导

B.1 Sigmoid函数的等价变换

从sigmoid函数的定义开始：

$$\sigma(x) = \frac{1}{1 + e^{-x}}$$

分子分母同时乘以$e^x$：

$$\sigma(x) = \frac{e^x}{e^x + 1}$$

为了防止溢出，引入偏移量$x_{\max}$：

$$\sigma(x) = \frac{e^{x-x_{\max}}}{e^{x-x_{\max}} + e^{-x_{\max}}}$$

这个形式的分子与Softmax中单个元素的分子完全相同，可以复用Shiftmax的实现。

B.2 GELU近似的误差分析

GELU的精确定义：

$$\text{GELU}(x) = x \cdot \Phi(x) = x \cdot \frac{1}{\sqrt{2\pi}} \int_{-\infty}^x e^{-t^2/2}dt$$

使用sigmoid近似：

$$\text{GELU}(x) \approx x \cdot \sigma(\alpha x)$$

其中$\alpha = \sqrt{2/\pi} \cdot \text{exp}(1/\pi) \approx 1.702$。

这个近似的最大绝对误差约为0.02，相对误差在$|x| < 3$的范围内小于2%。

C. I-LayerNorm的收敛性分析

C.1 Newton-Raphson方法

求平方根$\sqrt{v}$等价于求解方程$f(x) = x^2 - v = 0$。

Newton-Raphson迭代公式：

$$x_{i+1} = x_i - \frac{f(x_i)}{f'(x_i)} = x_i - \frac{x_i^2 - v}{2x_i} = \frac{x_i + v/x_i}{2}$$

这正是我们使用的迭代公式。

C.2 收敛速度

Newton-Raphson方法具有二次收敛性。设$\epsilon_i = |x_i - \sqrt{v}|$是第$i$次迭代的误差，则：

$$\epsilon_{i+1} \leq \frac{\epsilon_i^2}{2\sqrt{v}}$$

这意味着每次迭代误差的有效数字位数大约翻倍。

C.3 初始值选择

我们选择$I_0 = 2^{\lfloor \text{bit}(\text{Var}(I_x))/2 \rfloor}$作为初始值。这保证了：

$$\frac{\sqrt{\text{Var}(I_x)}}{2} \leq I_0 \leq 2\sqrt{\text{Var}(I_x)}$$

初始相对误差最多为2倍，经过10次迭代后，相对误差小于$2^{-1024}$，远超过INT8精度要求。

D. 量化误差分析

D.1 量化噪声模型

均匀量化引入的噪声可以建模为：

$$e = R - S \cdot I$$

其中$e$近似服从均匀分布$U(-S/2, S/2)$。

量化噪声的方差：

$$\text{Var}(e) = \frac{S^2}{12}$$

D.2 误差传播

对于线性操作$Y = WX$，量化误差传播为：

$$e_Y = W \cdot e_X + e_W \cdot X + e_W \cdot e_X$$

忽略二阶项$e_W \cdot e_X$，误差方差：

$$\text{Var}(e_Y) \approx \|W\|_F^2 \cdot \text{Var}(e_X) + \|X\|_F^2 \cdot \text{Var}(e_W)$$

这表明误差随层数线性累积，但通过合理的缩放因子选择可以控制在可接受范围内。