Amortized-Precision Quantization for Early-Exit Vision Transformers¶
会议: NeurIPS 2026
arXiv: 2605.07317
OpenReview: https://openreview.net/forum?id=EYh3OlleoZ
论文: https://arxiv.org/abs/2605.07317
领域: 模型压缩
关键词: 视觉Transformer, 动态推理, 早期退出, 摊销精度量化, 双层优化
一句话总结¶
针对低比特量化噪声扰动早期退出决策导致过早退出掉点或过晚退出低效的矛盾,本文提出摊销精度量化(APQ)范式与 MAQEE 双层优化框架,联合优化退出阈值与层级比特分配,在维持全精度性能的同时降低高达 95% 的 BOPs。
研究背景与动机¶
视觉 Transformer(ViT)在图像分类、目标检测和语义分割等任务上表现出色,但其自注意力机制带来的巨大计算开销严重阻碍了在端侧边缘设备上的实际部署。为了压低计算资源消耗,业界普遍采用量化技术减少权重与激活值的位宽,从最初粗暴统一位宽的固定精度量化(FPQ),逐步演进到依据层敏感度分配精度的混合精度量化(MPQ)。然而,传统的量化方法在理论和算法设计上都隐含了一个关键假设:网络具有确定性的静态全深度执行路径,每个输入样本都必须完整穿越全部层结构。
与此并行发展的是动态推理范式,尤其是早期退出(Early Exiting, EE)机制。EE 通过在 ViT 中间层挂载轻量级退出分类头,允许简单样本在置信度达标时提前终止推理,仅让困难样本深入更深层,从而在输入自适应的维度上大幅削减计算延迟。当这两种最具代表性的加速机制被直接叠加时,系统展现出严重的脆弱性与互斥性:传统 MPQ 是基于静态全深度敏感度校准的,而 EE 带来了随机且由输入决定的动态执行深度。量化引入的数值噪声极易扰动中间层退出头的置信度估算,诱发“过早退出”(premature exit,浅层置信度被噪声虚高而过早截止,导致准确率断崖式下滑高达 50%)或“过晚退出”(delayed exit,置信度被噪声打压导致无法早退,丧失加速收益)。反之,直接在已量化模型上挂载并训练退出头也因低比特权重的表示自由度受限而难以收敛。
这表明静态精度分配与动态随机深度之间存在本质上的目标失配。在动态推理路径中,浅层与高频激活层对整网量化误差的传播与退出决策起着决定性作用,而极少被访问的深层即使理论敏感度高,其全局影响也会被低激活概率大幅稀释。本文的核心 idea 是:将量化精度视为受动态利用率加权的摊销资源,提出摊销精度量化(APQ)范式,并通过显式建模决策不确定性与表征漂移风险的 MAQEE 双层优化框架,联合求解输入自适应的退出阈值与层级位宽分配。
方法详解¶
整体框架¶
本文构建的 MAQEE(Mutual Adaptive Quantization with Early Exiting)旨在解决摊销精度量化(APQ)这一 NP-hard 的离散双层优化问题。整体推理与优化系统包含两大紧密耦合的交互环路:外层循环负责优化各退出头的置信度阈值 \(\bm{\phi}\),控制样本的层级利用率分布;内层循环则在当前利用率与风险指导下自适应调整各层的位宽分配 \(\bm{b}\),再配合轻量级自蒸馏机制对量化造成的特征扭曲与置信度漂移进行参数微调。
整个框架通过综合考量两类决策层风险(性能差距风险 PGR、边界敏感性风险 BSR)与两类表征层风险(量化引起漂移 QID、逆量化信噪比 \(\text{SQNR}^{-1}\)),量化每一层在动态路径下的综合不稳定性得分,指导位宽在层间的最优再分配。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入样本 x 与预训练 ViT"] --> B["摊销精度量化目标建模<br/>层利用率 u_l 结合 BOPs 计算摊销代价"]
B --> C["四元风险量化度量<br/>决策级 PGR+BSR 与表征级 QID+SQNR^-1"]
C --> D["双层交替优化<br/>外层搜索阈值 φ 与内层贪心调整位宽 b"]
D --> E["轻量级自蒸馏微调<br/>全精度教师软标签约束对齐输出表征"]
E --> F["输出:Pareto 最优的低比特早期退出模型"]
关键设计¶
1. 摊销精度量化目标:面向随机停止深度的利用率感知建模
传统的混合精度量化在评估计算成本和层敏感度时,假设所有样本都会无差别地经过全网 \(L\) 个层。而在早期退出网络中,执行深度 \(K_{\bm{\phi}}(x)\) 是一个依赖于输入样本 \(x\) 和阈值策略 \(\bm{\phi}=\{\phi_\ell\}_{\ell=1}^L\) 的离散随机变量。为此,本文定义层利用率指示函数 \(u_\ell(x,\bm{\phi})=\mathbf{1}\{x\text{ 到达层 }\ell\mid\bm{\phi}\}\) 及其在数据分布上的期望利用率 \(u_\ell(\bm{\phi})=\mathbb{E}_{x\sim\mathcal{D}}[u_\ell(x,\bm{\phi})]=\Pr(K_{\bm{\phi}}(x)\ge \ell)\)。摊销精度量化(APQ)将全网期望计算复杂度度量为利用率加权的位操作数(BOPs):
其中 \(c_\ell(b_\ell)\) 为第 \(\ell\) 层在位宽 \(b_\ell\) 下的单次执行 BOPs。理论推导(定理 3)表明,静态量化仅依据静态敏感度 \(S_\ell\) 排序,若存在两层满足 \(S_i > S_j\) 但 \(u_i(\bm{\phi})S_i < u_j(\bm{\phi})S_j\),静态 MPQ 就会错误地将宝贵的高比特分配给访问概率极低的层 \(i\),导致次优的精度与计算开销折角。APQ 目标通过动态利用率与位宽的显式耦合,奠定了联合优化的理论基石。
2. 四元风险量化体系:解耦决策不确定性与表征畸变
量化噪声对动态推理的破坏主要源于两个维度:一是使退出置信度跨过阈值边界导致错误退出决策(定理 1 揭示的误退出上界),二是特征在沿随机路径向下传递时产生的误差累积与值域缩放畸变(定理 2 揭示的路径累积界)。MAQEE 相应设计了四个归一化的细粒度风险分量: - 边界敏感性风险(BSR):\(\text{BSR}_\ell=\mathbb{E}_{x}[\mathbf{1}\{|\rho_\ell(x)-\phi_\ell|\le \tau_\ell\}]\),统计全精度置信度落在阈值容限 \(\tau_\ell\) 狭窄边界内的样本比例,边界越密集代表决策越易被微小量化噪声扰动颠覆; - 性能差距风险(PGR):\(\text{PGR}_\ell=\mathbb{E}_{(x,y)}[\mathcal{L}(\hat{y}_\ell(x),y)-\mathcal{L}(\hat{y}_L(x),y)]\),衡量提前终止于第 \(\ell\) 层相较于最终层全深度推理所承受的任务性能损失; - 量化诱导漂移(QID):\(\text{QID}_\ell=\mathbb{E}_x[(\frac{\max_i \tilde{h}_\ell(x)_i - \min_i \tilde{h}_\ell(x)_i}{\max_i h_\ell(x)_i - \min_i h_\ell(x)_i}-1)^2]\),捕捉量化前后隐藏层表征动态值域分布的形变程度; - 逆量化信噪比(\(\text{SQNR}^{-1}\)):\(\text{SQNR}_\ell^{-1}=\mathbb{E}_x[\frac{\|h_\ell(x)-\tilde{h}_\ell(x)\|^2}{\|h_\ell(x)\|^2}]\),衡量点对点的绝对噪声能量比例。
这四个分量通过超参数 \(\alpha\) 组合为统一的层级综合不稳定性评分 \(R_\ell = \alpha(\text{PGR}_\ell + \text{BSR}_\ell) + (1-\alpha)(\text{QID}_\ell + \text{SQNR}_\ell^{-1})\)。
3. 双层优化与自蒸馏恢复:交替坐标搜索与自适应位宽重分配
由于直接求解非凸、离散的 APQ 问题属于 NP-hard(定理 5),MAQEE 采用高效的双层坐标交替迭代算法,并在循环中融入热启动(warm-start)与轻量级自蒸馏: - 外层阈值搜索:在固定位宽分配 \(\bm{b}\) 下,目标是最小化累积风险与摊销 BOPs 之和 \(J_{\text{out}}(\bm{\phi})=\sum_\ell R_\ell + \lambda_{\text{outer}}\text{BOPs}(\bm{b};\bm{\phi})\)。利用在校准集单次前向缓存的置信度分布,在极小的网格邻域内坐标搜索最优阈值 \(\bm{\phi}^\star\),随后重新评估各层利用率 \(u_\ell(\bm{\phi})\); - 内层位宽重分配:在给定退出策略下,计算每一层增加精度的边际效用增益优先级:
其中 \(\Delta R_\ell\) 为提高该层位宽带来的风险降幅,\(\Delta \text{BOPs}_\ell\) 为对应的摊销算力增加量。算法每次从可降比特层集中选出优先级最低的层削减 1 个比特,转赠给可增比特层集中优先级最高的层,优先保障高频访问且敏感的层; - 自蒸馏恢复:在每轮位宽微调后,冻结全精度模型作为教师网络输出软标签,通过交叉熵与 KL 散度联合微调可学习参数 \(\theta\):\(\mathcal{L}_{\text{rec}}(\theta)=\mathcal{L}(f_{\theta,\bm{b},\bm{\phi}}(x),y)+T^2\mathcal{L}_{\text{KL}}(p(x;f_\theta,T)\|p(x;f_{\theta,\bm{b},\bm{\phi}},T))\),快速拉平表征漂移并稳定边界置信度。
实验关键数据¶
主实验¶
实验在图像分类(CIFAR-100、ImageNet-1K)、语义分割(ADE20K/SceneParse150)和目标检测(MS COCO)上系统评测了 DeiT、ViT 与 Swin 三大架构。基准设置为 W4A4(平均 4 比特权重与激活)。评测协议包含“受控性能”(达到不加早退的纯量化模型基准线所需的退出层 \(\hat{L}\) 与 \(\widehat{\text{BOPs}}\))和“标准配置”(默认阈值下的平均退出层 \(\bar{L}\)、\(\overline{\text{BOPs}}\) 与实际精度)。
表 1:图像分类主实验对比(ImageNet-1K,W4A4 设置)
| 模型架构 | 量化方案 | 早退方法 | 受控退出层 \(\hat{L}\downarrow\) | 受控 \(\widehat{\text{BOPs}}\downarrow\) | 标准退出层 \(\bar{L}\downarrow\) | 标准 \(\overline{\text{BOPs}}\downarrow\) | Top-1 准确率 (%) \(\uparrow\) |
|---|---|---|---|---|---|---|---|
| Swin (FP32) | - | - | / | / | 24 | 1.39E13 | 83.50 |
| Swin (FP32) | - | ViT-EE | 16.62 | 1.23E13 | 18.23 | 1.35E13 | 83.68 |
| Swin (FP32) | - | LGViT | 11.13 | 8.83E12 | 14.20 | 1.10E13 | 82.70 |
| Swin (FP4/4) | RepQ | ViT-EE | N/A | N/A | 20.16 | 2.27E11 | 74.75 |
| Swin (FP4/4) | ERQ | ViT-EE | 23.46 | 2.65E11 | 20.96 | 2.31E11 | 75.45 |
| Swin (MP4/4) | MPTQ | ViT-EE | N/A | N/A | 20.99 | 2.39E11 | 71.55 |
| Swin (MP4/4) | MAQEE | 自适应 | 17.60 | 2.09E11 | 18.75 | 2.12E11 | 80.78 |
| ViT-B (FP32) | - | - | / | / | 12 | 1.80E13 | 83.10 |
| ViT-B (FP32) | - | ViT-EE | 7.65 | 1.90E13 | 8.68 | 2.05E13 | 79.60 |
| ViT-B (FP4/4) | RepQ | ViT-EE | N/A | N/A | 10.83 | 3.71E11 | 74.67 |
| ViT-B (FP4/4) | ERQ | ViT-EE | N/A | N/A | 9.88 | 3.49E11 | 77.35 |
| ViT-B (MP4/4) | MPTQ | ViT-EE | N/A | N/A | 10.26 | 3.58E11 | 76.56 |
| ViT-B (MP4/4) | MAQEE | 自适应 | 7.04 | 2.85E11 | 7.26 | 2.93E11 | 79.41 |
表 2:语义分割与目标检测跨任务性能(ViT-B,W4A4)
| 方法配置 | 分割受控 \(\hat{L}\downarrow\) | 分割受控 \(\widehat{\text{BOPs}}\downarrow\) | 分割标准 mIoU \(\uparrow\) | 检测受控 \(\hat{L}\downarrow\) | 检测受控 \(\widehat{\text{BOPs}}\downarrow\) | 检测标准 [email protected] \(\uparrow\) |
|---|---|---|---|---|---|---|
| FP32 全精度基准 | / | / | 34.14 | / | / | 68.23 |
| ViT-EE + ERQ (FP4/4) | N/A | N/A | 27.34 | N/A | N/A | 40.42 |
| ViT-EE + MPTQ (MP4/4) | 11.46 | 3.94E11 | 26.52 | N/A | N/A | 41.04 |
| MAQEE (本文) | 5.74 | 2.60E11 | 31.20 | 10.40 | 3.77E11 | 50.83 |
消融实验¶
表 3:各风险控制分量对 ViT-B ImageNet 分类的消融实验
| 模型消融配置 | 受控退出层 \(\hat{L}\downarrow\) | 受控 \(\widehat{\text{BOPs}}\downarrow\) | 标准退出层 \(\bar{L}\downarrow\) | 标准 \(\overline{\text{BOPs}}\downarrow\) | Top-1 准确率 (%) \(\uparrow\) |
|---|---|---|---|---|---|
| MAQEE 完整模型 | 7.04 | 2.85E11 | 7.26 | 2.93E11 | 79.41 |
| 去掉 PGR (无性能差距风险) | N/A | N/A | 4.23 | 2.12E11 | 61.18 |
| 去掉 BSR (无边界敏感风险) | 7.59 | 3.12E11 | 7.01 | 2.74E11 | 76.53 |
| 去掉 \(\text{SQNR}^{-1}\) (无信噪比约束) | 7.35 | 2.98E11 | 7.28 | 2.95E11 | 78.35 |
| 去掉 QID (无量化诱发漂移) | 7.26 | 3.01E11 | 7.18 | 2.95E11 | 79.06 |
关键发现¶
- PGR 是防止恶性过早退出的最关键刹车:去掉 PGR 后,外层阈值搜索缺乏对准确率下跌的惩罚,导致模型盲目追求低延迟而产生极其激进的早退决策(平均退出层降至 4.23 层),准确率断崖式跌落至 61.18%,直接无法达到受控性能目标。
- BSR 决定了量化早退的决策稳定性:去除 BSR 会使退出决策对数值扰动极度敏感。为了满足既定的准确率约束,模型被迫采取过于保守的退出策略,使得受控退出层从 7.04 层推迟至 7.59 层,BOPs 开销显著上升约 9.5%。
- 静态量化在动态网络中全面失效:传统 SOTA 混合精度量化 MPTQ 假设固定前向路径,在受控性能下于多个骨干网上均显示为“N/A”(无法在可接受延迟内达到基线精度),证明了考虑层利用率的 APQ 是动态网络量化的必经之路。
亮点与洞察¶
- 将“层利用率”作为资源分配权重的理论跃升:突破了传统混合精度量化仅依据静态二阶梯度或重建误差的固有思维,证明了在动态网络中静态敏感度排序存在逆序反例,巧妙地将静态敏感度与动态到达概率 \(u_\ell(\bm{\phi})\) 相结合。
- 双层解耦与热启动设计规避了 NP-hard 的联合求解困境:利用校准集前向缓存置信度的 coordinate search 使外层阈值更新几乎零额外计算开销,而基于边际效用 \(\Psi_\ell\) 的位宽重分配使得离散预算分配可在一维网格上快速收敛。
- 可平替拓展至广义输入自适应架构:APQ 的摊销精度理论不局限于基于分类头的早期退出,还可以自然推广至专家混合模型(MoE)中的路由激活、递归循环 Transformer 以及动态层跳跃(LayerSkip)等具有随机计算拓扑的现代模型。
局限与展望¶
- 密集预测任务(如目标检测)的下界仍有提升空间:在 MS COCO 上虽然领先 baseline 近 10 个点,但相较于 FP32 基准(68.23%)仍出现明显掉点(50.83%)。这主要是因为检测依赖 FPN/多尺度深浅层特征融合,早期退出截断深层特征导致特征金字塔不完整,叠加量化噪声后多目标定位严重退化。
- 硬件异构性与非规则位宽的落地瓶颈:当前硬件执行引擎对逐层 3、4、5、6 bit 混合位宽的加速支持有限,通常依赖特定的整数模拟或定制硬件加速器,需进一步探索通道级或分组级规则对齐的硬件友好型 APQ。
相关工作与启发¶
- vs MPTQ-ViT (Tai & Wu, 2024):MPTQ 采用基于二阶 Hessian 信息的静态混合精度搜索,但在早退网络中会盲目给极少执行的深层高位宽,导致浅层量化噪声诱发大量错误退出;MAQEE 引入利用率加权与决策风险,在同等平均 4-bit 下精度高出 2%~8%。
- vs QuEE (Regol et al., 2024) / McQueen (Saxena & Roy, 2023):早期的量化早退探索要么仅支持有限的位宽选择(如 8-bit),要么受限于固定的退出头结构且缺乏严密的风险分解;MAQEE 在极具挑战的 W4A4 超低精度下,建立了严谨的误差界定理与通用双层求解器。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从理论与算法上系统剖析并解决了量化噪声与动态随机深度执行路径之间的内在冲突。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 种视觉骨干、3 项核心视觉任务、4 种风险分量消融及双评估协议,对比坚实。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰(定理 1 至 5),符号定义严格,逻辑链条层层递进。
- 价值: ⭐⭐⭐⭐⭐ 为边缘端部署自适应深度 Vision Transformers 和大型多模态模型提供了兼顾延迟与精度的全新范式。