【PTQ / NAS】LRP-QViT (arXiv 2024)

视觉Transformer（ViT）在各种视觉任务中展现出了卓越的性能。然而，ViT模型需要大量的计算和内存资源，这使得将其部署到资源受限的平台具有挑战性。量化是一种广泛采用的减少模型体积的方法，但大多数研究主要关注对整个网络进行等比特宽度量化，导致次优结果。尽管已有少量关于ViT混合精度量化（MPQ）的工作，但它们通常依赖基于搜索空间的方法，或以任意方式采用混合精度。在本文中，我们提出了LRP-QViT，这是一种基于可解释性的方法，用于根据分类过程中各层的重要性为其分配不同精度的比特数。具体而言，为了衡量每层在预测目标类别中的贡献度，我们采用逐层相关性传播（Layer-wise Relevance Propagation, LRP）方法。LRP在输出层分配局部相关性，并将其逐层向后传播，直到到达输入层。这些相关性分数被用作计算各层贡献度分数的指标。此外，我们引入了截断的逐通道量化方法，旨在消除LayerNorm激活后的异常值，以缓解严重的跨通道差异。为了验证和评估我们的方法，我们在多种数据集上对ViT、DeiT和Swin Transformer模型应用了LRP-QViT。实验结果表明，在4位和6位量化背景下，我们的固定比特数和混合比特数的训练后量化方法均优于现有模型。

问题: 统一比特量化无法兼顾不同层的量化敏感性差异，混合精度策略需要更好的指导机制。

核心创新:

模块	方法	作用
LRP 贡献评分	使用 Layer-wise Relevance Propagation 计算每层对分类的贡献	指导混合精度比特分配
CCLQ（Clipped Channel-wise Quantization）	对 post-LayerNorm 激活使用截断通道-wise 量化	缓解通道间方差

结果: 在 4-bit 和 6-bit 场景下，混合精度 LRP-QViT 优于 RepQ-ViT 等方法。

意义: 将可解释性 AI（Explainable AI）引入量化比特分配，提供了新的理论视角。

introduce

已有诸多研究利用后训练量化（PTQ）对视觉 Transformer（ViT）进行量化 [10, 23, 26, 48]。这些工作识别出若干限制量化模型性能的瓶颈组件，例如 LayerNorm、Softmax 和 GELU，并提出了带有改进措施的 PTQ 方案。

然而，大多数现有工作基于一个前提，即 ViT 中的所有层对模型输出的贡献是相等的，因此对整个网络采用统一的比特精度，从而导致次优解。固定比特量化迫使重要层与不重要层以相同的比特精度进行量化，错失了进一步减小模型尺寸并提升性能的机会。

混合精度量化（MPQ）通过允许不同层采用不同的比特精度来解决这一局限性，使得关键层可以使用比其他层更高的比特精度。以往针对 MPQ 的大部分工作聚焦于卷积神经网络（CNNs），并利用基于策略搜索的方法 [30, 31, 44] 或基于准则的方法 [11, 12] 来确定最优比特精度。

在本文中，我们提出一种面向视觉 Transformer 模型的混合比特量化的可解释性方法 [5, 6]，该方法基于各层对模型性能的贡献。我们利用层相关性传播（Layer Relevance Propagation, LRP）为所有层获取贡献得分，并据此指导针对不同层的混合精度比特分配策略。

在我们的框架中，我们采用并改进了 RepQViT [23] 框架以用于模型量化。具体而言，

• 针对 LayerNorm 之后的激活值，我们引入了截断的通道-wise 量化方法，以去除异常值并减轻推理阶段因通道间过度变化带来的影响。该截断操作通过调整 LayerNorm 的仿射因子及下一层的权重来实现。
• 针对 Softmax 之后的激活值，我们完全采用文献 [23] 中的修改方案：首先使用 log√2 量化器进行初始量化，以获得更高的精度表示；然后在推理过程中重新缩放至 log2 量化器，以实现更友好的量化过程。

结合改进后的量化模型与可解释的混合精度比特分配策略，我们的主要贡献总结如下：

• 我们提出了 LRP-QViT，一种面向视觉 Transformer 模型量化的混合精度框架。
• LRP-QViT 利用层相关性传播为每层分配贡献得分，并据此指导比特分配。
• 我们额外引入了针对 LayerNorm 后激活值的截断通道-wise 量化方法，该方法能够去除异常值，并提升固定比特和混合精度量化的性能。
• 我们在 ViT、DeiT 和 Swin 模型上的实验结果表明，LRP-QViT 在分类、检测和分割基准任务上均优于固定比特量化方法。

相关研究

2.3. 变换器的可解释性
以往关于可解释性的研究主要集中在卷积神经网络（CNNs）上，主要侧重于梯度方法 [36, 37] 或归因方法 [2, 14]。然而，随着视觉变换器架构的发展，针对变换器可解释性的研究显著增加。

在“注意力展开”（Attention Rollout）方法 [1] 中，不同层的注意力分数被线性组合，但这种方法难以区分正负贡献。

逐层相关性传播（Layer-wise Relevance Propagation, LRP）方法从预测类别反向传播相关性至输入图像。已有若干研究将 LRP 应用于Transformer[2, 40]。然而，许多研究忽略了跨所有层的注意力传播，并忽略了执行两个激活图混合的网络部分，例如跳跃连接和矩阵乘法。此外，大多数工作并未明确区分各层对模型决策的正负贡献。由于缺乏这种区分，正负贡献被混合在一起，导致相关性分数高于必要水平。

在文献 [6] 中，相关性与梯度信息被用于迭代地去除负贡献，并计算变换器模型中每层各注意力头的准确相关性分数。在本工作中，我们采用文献 [6] 中提到的相关性分数计算方法，并将该方法扩展至计算其他层的相关性分数，例如 qkv 层、线性投影层、矩阵乘法层以及全连接层。

3.3. Layer-Wise Relevance Propagation

逐层相关性传播（Layer-wise Relevance Propagation ）是一种可解释的人工智能方法，能够揭示每个特征对模型输出的贡献。该方法首先为输出层分配一个相关性分数，然后通过反向传播将该分数逐层传递至网络各层，直至输入特征。这些相关性分数作为计算各层贡献度的指标。我们利用这些贡献度分数为后训练量化（PTQ）分配混合精度位宽。

3.3.1 相关性与梯度

设 $C$ 为分类头（classification head）的类别数量，$c\in 1\dots |C|$ 为待可视化的类别。我们针对类别 $c$ 传播相关性（Relevance, $R$）和梯度（Gradients, $\nabla$）。设 $x^{(n)}$ 为网络第 $n$ 层 $L^n$ 的输入，其中 $n\in [1\dots N]$，$x^N$ 为网络输入，$x^1$ 为网络输出。

对于类别 $t$ 的分类器输出 $y$，即 $y_t$，其相对于输入 $x_j^{(n)}$ 的梯度由下式给出：

$$\nabla x_j^{(n)}=\frac{\partial y_t}{\partial x_j^{(n)}}=\sum _i\frac{\partial y_t}{\partial x_i^{(n-1)}}\frac{\partial x_i^{(n-1)}}{\partial x_j^{(n)}}(18)$$

其中，索引 $j$ 对应 $x^{(n)}$ 中的元素，索引 $i$ 对应 $x^{(n-1)}$ 中的元素。

设 $L_j^n(X,W)$ 为张量 $X$ 和 $W$ 上的层运算。相关性传播遵循通用的深度泰勒分解（Deep Taylor Decomposition）[33]，如下所示：

$$R_j^n=G(X,W,R^{n-1})=\sum _i\frac{X_j\frac{\partial L_j^n(X,W)}{\partial X_j}}{{\sum }_{j^{\prime }}{L}_{j^{\prime }}^n(X,W)}{R}_i^{n-1}(19)$$

Transformer 模块包含 GELU [16]，其输出正值和负值。为了保持守恒规则，即所有层的相关性之和始终相同，我们移除所有具有负相关性的元素。我们构建一个索引子集 $p=\{(i,j)|x_j{w}_{ji}\ge 0\}$，从而得到以下传播公式：

Rjn=Gp(x,w,p,Rn−1)=∑(i,j)∈pxjwji∑{j′∣(j′,i)∈p}xj′wj′iRin−1(20) R^n_j = G_p(x, w, p, R^{n-1}) = \sum_{(i,j) \in p} \frac{x_j w_{ji}}{\sum_{\{j' | (j', i) \in p\}} x_{j'} w_{j'i}} R^{n-1}_i \quad (20) Rjn​=Gp​(x,w,p,Rn−1)=(i,j)∈p∑​∑{j′∣(j′,i)∈p}​xj′​wj′i​xj​wji​​Rin−1​(20)

为了启动相关性传播，我们设置 $R^0=1_c$，其中 $1_c$ 是一个指示目标类别 $c$ 的独热向量（one-hot vector）。

3.3.2 层重要性评分

在层相关性传播（LRP）中，相关性和梯度从分类头开始，通过每个模块内的所有层，向后传播至输入 patch embedding 层。在此过程中，网络中的每一层都会学习一个相关性评分图。具体而言，对于图像样本 $t$，考虑 Transformer 模块 $L$ 中的注意力层 $A$，其层梯度记为 $\nabla A_t^L$，相关性记为 $R_{A,t}^L$，则该注意力层的相关性评分图 $A_{s,t}^L$ 定义为：

$$A_{s,t}^L=E_h{\left(\nabla A_t^L\odot R_{A,t}^L\right)}_{+}(21)$$

其中 $E_h$ 表示在多头部（multi-head）上的均值，且我们的分析仅关注梯度-相关性乘积的正值部分。为了定量测量第 $L$ 个模块的注意力层对样本图像输出的贡献评分，我们简单地取相关性评分图的均值。由于 LRP 方法产生的相关性评分图是特定于类别的，即为不同的图像样本生成不同的图。第 $L$ 个模块的注意力层的总体贡献评分通过对从 ImageNet1k 训练数据集中随机选择的 $T=50,000$ 张图像取平均来计算，如下所示：

$$C_A^L=\frac{1}{T}\sum _{t=1}^{T}E\left(A_{s,t}^L\right)(22)$$

相同的计算方法也用于所有其他量化层，包括 qkv 层、矩阵乘法层（matmul1 和 matmul2）、投影层以及全连接层（fc1 和 fc2），以计算它们对输出分类的贡献评分。任意模块中任意层的相对重要性评分是该层的贡献评分在所有贡献评分中的平均归一化值。第 $L$ 个模块中注意力层的相对层重要性评分表达式如下：

$$I_A^L=\frac{C_A^L}{{\sum }_{l=1}^L{\sum }_{u\in U_l}{C}_u^l}(23)$$

这里，$U_l$ 表示第 $l$ 个 Transformer 模块中的所有量化层。我们使用相对重要性评分 $I$ 来决定模型量化期间的比特宽度，为具有较高相对重要性评分的层分配更多的比特数。

3.3.3 混合精度比特分配

Transformer 架构的早期模块通常负责捕获输入数据中的低层特征和细节。这些特征对于网络有效学习表示至关重要。此外，早期模块对输入数据中的微小变化非常敏感。如果这些模块被过度量化，网络可能会丧失其表示能力，从而导致整体性能下降。因此，为了避免在处理的早期阶段丢失信息，我们在训练后量化期间为前两个模块的所有层分配更高的精度。为了保持模型大小，我们利用层重要性评分来识别重要层，并减少这些层的比特分配。

（这里是不是缺少混合精度的具体方案，我看有些方案是做和全精度结果比较，进行精度搜索（利用遗传算法等））

论文应该是没写完整混合方案。

5. 结论

本文提出了一种新颖的方法 LRP-QViT，用于视觉 Transformer 的后训练混合位量化。LRP-QViT 通过反向传播相关性与梯度来学习层重要性分数，并基于该重要性分数为混合精度量化选择最优的比特数。此外，我们引入了裁剪式通道量化方法，用于处理 LayerNorm 之后的激活值，该方法能够消除异常值并抑制通道间差异，从而提升模型性能。

（正文完）

---