LoopQ：递归Transformer的量化

# LoopQ：递归变换器的量化
来源：https://arxiv.org/html/2605.16343
Rui Fang 台湾大学 rfang@arbor\.ee\.ntu\.edu\.tw  
Hsi\-Wen Chen 台湾大学 hwchen@arbor\.ee\.ntu\.edu\.tw  
Ming\-Syan Chen 台湾大学 mschen@ntu\.edu\.tw

###### 摘要

循环语言模型（LoopLM）通过递归重用Transformer块提高了参数效率，从而在固定模型大小下实现更深的计算。然而，这种重用使得LoopLM在训练后量化（PTQ）下更加脆弱。我们首次对LoopLM中的量化问题进行了系统研究，并识别出三个挑战：角色间的分布偏移、循环转换间的状态重用以及递归误差累积。为解决这些挑战，我们提出了LoopQ，一种面向循环的PTQ框架，它保留共享的量化主干，同时引入轻量级的自适应调整。LoopQ结合了激活缩放、选择性变换、跨循环状态对齐以及轨迹感知优化，以减少循环内的分布不匹配和跨循环的误差累积。在七个基准上的实验表明，在W4A4量化下，与最强的静态PTQ基线相比，LoopQ将平均下游准确率提高了68.8%，并将平均困惑度降低了87.7%。

## 1 引言

量化鲁棒性[1(https://arxiv.org/html/2605.16343#bib.bib17),34(https://arxiv.org/html/2605.16343#bib.bib18),38(https://arxiv.org/html/2605.16343#bib.bib19)]对于在资源受限设备上高效部署LLM至关重要[39(https://arxiv.org/html/2605.16343#bib.bib15)]。循环语言模型（LoopLM）[5(https://arxiv.org/html/2605.16343#bib.bib36),21(https://arxiv.org/html/2605.16343#bib.bib40),45(https://arxiv.org/html/2605.16343#bib.bib42),17(https://arxiv.org/html/2605.16343#bib.bib43),3(https://arxiv.org/html/2605.16343#bib.bib41),8(https://arxiv.org/html/2605.16343#bib.bib5)]通过递归重用共享的Transformer块，在固定参数预算下提高了表达能力，使其对内存和带宽受限的场景具有吸引力[43(https://arxiv.org/html/2605.16343#bib.bib6),20(https://arxiv.org/html/2605.16343#bib.bib1),19(https://arxiv.org/html/2605.16343#bib.bib4)]。然而，现有的LoopLM主要针对全精度推理，而大多数训练后量化（PTQ）方法优化层级的量化误差[37(https://arxiv.org/html/2605.16343#bib.bib32),2(https://arxiv.org/html/2605.16343#bib.bib33),22(https://arxiv.org/html/2605.16343#bib.bib31),18(https://arxiv.org/html/2605.16343#bib.bib14),14(https://arxiv.org/html/2605.16343#bib.bib34),35(https://arxiv.org/html/2605.16343#bib.bib10),32(https://arxiv.org/html/2605.16343#bib.bib35)]，没有考虑通过递归重用状态传播的误差。

这种限制对于LoopLM尤其成问题，因为相比类似的非循环架构，它们对量化敏感得多[40(https://arxiv.org/html/2605.16343#bib.bib2)]。在我们的初步实验中，直接对LoopLM应用PTQ会导致性能下降超过50%。在相同参数预算下，量化的非循环LLM甚至可以胜过更高精度的LoopLM。这表明，全精度设置下参数重用的优势并不会自动转移到量化场景中，在量化场景中，递归重用会放大量化引起的误差。

参照图注图1：左图：LoopLM表现出与循环相关的激活漂移（挑战1）、循环转换处的误差尖峰（挑战2）以及跨循环的误差累积（挑战3）。右图：LoopQ保留共享的量化主干，同时添加轻量级的循环感知缩放、选择性变换和跨循环状态对齐，以减少递归量化误差。图1(https://arxiv.org/html/2605.16343#S1.F1)通过量化递归计算所诱发的两种耦合现象，展示了这种脆弱性。首先，左上图显示，相同的物理层在不同循环遍历中表现出不同的归一化P99激活范围，这表明单一的静态量化配置无法适应多种与循环相关的状态。其次，左下图显示，量化误差并非局限于局部。相反，它们随着递归的进行而增长，并在循环转换处（一个循环产生的隐藏状态传递给下一个循环时）出现明显的尖峰。这些观察表明，LoopLM中的量化既受到每层分布不匹配的支配，也受到递归误差累积的支配。

根本原因在于，随着递归的进行，相同的共享块在不同的计算角色下被重用。因此，单个物理层不再表现为固定的前馈位置，而是跨循环在依赖于阶段的角色下运行。这种角色变化导致了三个关键挑战，我们通过理论分析将其形式化：(1) 角色间分布偏移：与循环相关的激活统计量导致在共享量化下的尺度和几何不匹配；(2) 转换间状态重用：一个循环产生的隐藏状态被重用为下一个循环的输入，引入了转换不一致性；(3) 递归误差累积：量化误差沿着递归轨迹传播，并被后续的循环相关变换放大。

一个简单的解决方案是为所有循环分配循环特定的量化配置，从而完全适应循环相关的分布。然而，这会破坏递归计算的共享参数结构，并引入大量的内存和部署开销，损害LoopLM的效率优势。因此，我们提出LoopQ，一种面向循环的训练后量化框架，它保留共享的量化主干，同时引入影响不到5%变换参数（对应于不到0.12%的总模型参数）的轻量级循环相关自适应调整。

如图1(https://arxiv.org/html/2605.16343#S1.F1)右图所示，循环感知激活缩放（棕色块）使用循环相关尺度来追踪幅度漂移，而选择性循环感知变换（蓝色块）则通过共享间隙分析识别出的共享敏感模块，选择性引入循环相关变换，纠正残差几何不匹配。为了稳定循环转换，跨循环转换适配器（绿色块）使用轻量级适配器对齐跨循环隐藏状态。最后，在定点收敛性[17(https://arxiv.org/html/2605.16343#bib.bib43)]的指导下，轨迹感知校准在真实的递归动态下执行渐进式的轨迹感知校准，以抑制轨迹级别的误差累积。

我们的贡献总结如下：(i) 我们首次对LoopLM中的量化进行了系统研究，识别出递归计算中的三个基本挑战：角色间分布偏移、转换间状态重用和递归误差累积。(ii) 我们提出了LoopQ，一种面向循环的训练后量化框架，它保留共享的量化主干，同时引入轻量级循环相关自适应调整，包括激活缩放、选择性变换、跨循环状态对齐和轨迹感知优化。(iii) 在七个基准上的实验表明，在W4A4量化下，LoopQ优于五种PTQ基线，将平均下游准确率提高了68.8%，并将平均困惑度降低了87.7%。

## 2 问题形式化

我们首先介绍LoopLM架构，然后是训练后量化设置。

**LoopLM。** LoopLM是一种递归Transformer，由共享的L层堆栈组成，迭代应用T次循环。令\(\mathbf{X}\)表示输入，\(f_\ell(\cdot;\theta_\ell)\)表示第\(\ell\)层。我们将循环\(t\)中第\(\ell\)层的隐藏状态表示为\(\mathbf{H}_{t,\ell}\in\mathbb{R}^{n\times d}\)，其中每一行对应一个token表示。初始状态为\(\mathbf{H}_{0,0}=\mathrm{Embed}(\mathbf{X})\)。

循环计算定义为

\[
\mathbf{H}_{t,\ell}=f_\ell(\mathbf{H}_{t,\ell-1};\theta_\ell),\qquad\mathbf{H}_{t+1,0}=\mathbf{H}_{t,L},
\tag{1}
\]

其中\(t\in\{0,\ldots,T-1\}\)，\(\ell\in\{1,\ldots,L\}\)。令\(\mathbf{H}_t:=\mathbf{H}_{t,0}\)。每个循环可以概括为\(F:=f_L\circ\cdots\circ f_1\)，得到\(\mathbf{H}_{t+1}=F(\mathbf{H}_t)\)，最终输出为\(\mathbf{Z}=\mathrm{Proj}(\mathbf{H}_T)\)。

为了分析清晰，我们使用单一权重矩阵\(\mathbf{W}_\ell\)来表示\(f_\ell(\cdot;\theta_\ell)\)，即\(\mathbf{H}_{t,\ell}=\mathbf{H}_{t,\ell-1}\mathbf{W}_\ell^\top\)。这种抽象自然可以推广到具有多个权重矩阵的通用Transformer块。

**对LoopLM的训练后量化。** 我们考虑标准的训练后量化（PTQ）设置，其中每层\(f_\ell\)由量化后的对应层\(\widetilde{f}_\ell\)（使用量化权重\(\hat{\theta}_\ell\)）近似。由于模型权重和激活通常表现出不平衡的分布，例如重尾或通道方差差异，直接量化是次优的，因为量化级别在动态范围上的分配效率低下。

为了解决这个问题，我们采用可逆变换\(\mathbf{P}_\ell\)来重塑表示，同时保留原始函数。对于每一层\(\ell\)，全精度映射可以等价地写为

\[
\mathbf{H}_{t,\ell}=(\mathbf{H}_{t,\ell-1}\mathbf{P}_\ell)(\mathbf{W}_\ell\mathbf{P}_\ell^{-\top})^\top.
\tag{2}
\]

相应地，量化推理执行为

\[
\widetilde{\mathbf{H}}_{t,\ell}=Q_a(\widetilde{\mathbf{H}}_{t,\ell-1}\mathbf{P}_\ell;c_\ell)\,Q_w(\mathbf{W}_\ell\mathbf{P}_\ell^{-\top})^\top,\qquad\widetilde{\mathbf{H}}_{t+1,0}=\widetilde{\mathbf{H}}_{t,L}.
\tag{3}
\]

这里，\(Q_w\)和\(Q_a\)分别表示权重和激活量化。对于第\(\ell\)层的激活元素\(x\)，对称均匀量化定义为

\[
Q_a(x;c_\ell)=c_\ell\cdot\mathrm{clip}\!\left(\left\lfloor x/c_\ell\right\rceil,q_{\min},q_{\max}\right),
\tag{4}
\]

其中\(c_\ell\)决定了量化级别之间的间距，\(q_{\min}\)和\(q_{\max}\)分别表示最小和最大可表示的整数级别。

与常规量化相比，量化LoopLM引入了两种不同的不匹配来源：(i) 每个共享层内激活分布的循环间变化，其中量化参数\((\mathbf{P}_\ell,c_\ell)\)跨循环固定，而激活跨循环变化；(ii) 隐藏状态的递归演化，其中\(\mathbf{H}_{t,\ell}\)跨循环变化，导致复合分布偏移，单一的静态量化配置无法捕获。

## 3 LoopLM中量化误差产生的原因

我们分析了静态量化与LoopLM中动态隐藏状态演化之间的不匹配如何导致两种耦合效应。所有证明见附录A(https://arxiv.org/html/2605.16343#A1)。

**循环内的分布偏移。** 一个关键挑战在于，相同的量化层\(\widetilde{f}_\ell\)必须处理由递归计算产生的与循环相关的隐藏状态。

这种不匹配表现为两种形式。首先，当激活幅度变化时，共享的激活尺度\(c_\ell\)无法适应所有循环：大的\(c_\ell\)会降低小激活的分辨率，而小的\(c_\ell\)会导致大激活被裁剪。其次，共享的变换\(\mathbf{P}_\ell\)只能对齐循环轨迹中有限部分的激活几何结构。随着隐藏状态\(\mathbf{H}_{t,\ell}\)的演化，它们的主要方向和通道能量分布发生变化，产生与循环相关的异常模式。因此，固定的\(\mathbf{P}_\ell\)可能会在某些循环中抑制对量化敏感的维度，但在其他循环中放大它们，使量化噪声与关键特征错位。

###### 命题1。

定义第\(\ell\)层的变换后激活为\(\hat{\mathbf{H}}_{t,\ell-1}:=\mathbf{H}_{t,\ell-1}\mathbf{P}_\ell\)。假设对于每个循环\(t\)，\(\hat{\mathbf{H}}_{t,\ell-1}\sim s_t\mathbf{R}_t\)，其中\(s_t>0\)是一个循环相关的尺度，\(\mathbf{R}_t\)是一个协方差为\(\boldsymbol{\Sigma}_t\)的零均值随机矩阵。令\(Q_a(\cdot;c_\ell)\)是一个具有裁剪参数\(c_\ell\)的对称量化器。期望激活量化误差为

\[
\varepsilon_{t,\ell}:=\mathbb{E}\!\left[\left\|Q_a(\hat{\mathbf{H}}_{t,\ell-1};c_\ell)-\hat{\mathbf{H}}_{t,\ell-1}\right\|_F^2\right].
\]

如果存在两个循环\(t\neq t'\)，使得要么\(s_t\neq s_{t'}\)（尺度漂移），要么\(\boldsymbol{\Sigma}_t\neq\boldsymbol{\Sigma}_{t'}\)（协方差漂移），那么没有任何共享对\((\mathbf{P}_\ell,c_\ell)\)可以同时对\(\varepsilon_{t,\ell}\)和\(\varepsilon_{t',\ell}\)最优。因此，共享校准在至少一个循环上不可避免地产生严格为正的过量误差。

**跨循环的误差传播。** 除了分布偏移之外，PTQ在LoopLM中尤其脆弱，因为量化扰动的是整个递归隐藏状态轨迹，而不是孤立的层计算。回顾一下，\(\mathbf{H}_{t+1}=F(\mathbf{H}_t)\)和\(\widetilde{\mathbf{H}}_{t+1}=\widetilde{F}(\widetilde{\mathbf{H}}_t)\)分别表示跨越L层的全精度和量化递归。它们的不匹配源于两个耦合来源：(i) 通过量化层映射在每个循环内引入的近似误差；(ii) 由于将量化隐藏状态\(\widetilde{\mathbf{H}}_t\)（而非其全精度对应物\(\mathbf{H}_t\)）馈入后续计算而导致的状态不匹配。结果，量化误差并非在每个层或循环独立产生，而是递归地改变未来的隐藏状态，并沿着递归轨迹复合，特别是在循环转换\(t\rightarrow t+1\)处（见图1(https://arxiv.org/html/2605.16343#S1.F1)）。

###### 命题2。

定义循环误差为\(\varepsilon_t:=\|\widetilde{\mathbf{H}}_t-\mathbf{H}_t\|_F\)。假设两种递归从相同的初始状态开始，并且对于每个循环\(t\)，全精度函数\(F\)在\(\widetilde{\mathbf{H}}_t\)和\(\mathbf{H}_t\)之间的线段上是\(\gamma_t\)-Lipschitz的。那么

\[
\varepsilon_{t+1}\leq\varepsilon_t^{\mathrm{quant}}+\gamma_t\varepsilon_t
\]