ARCHead: 面向大语言模型输出头的激活度量残差校正

# 大语言模型输出头的激活度量残差校正

来源：https://arxiv.org/html/2608.02703
Şuayp Talha Kocabay 独立研究者 kocabaysuayptalha08@gmail\.com &Talha Rüzgar Akkuş 独立研究者 talharuzgarakkus@gmail\.com

###### 摘要

仅权重量化大幅降低了大型语言模型（LLM）Transformer 块的存储占用，但实际后端通常仍将最终语言建模头（LM-head）保留为 BF16 或 FP16。对这种投影进行朴素量化会强烈扰动词汇 logit 分布。我们提出 ARCHead，一种打包式 LM-head 压缩器，它结合了量化低秩核心、分组 INT4 残差，以及在激活派生度量中拟合的低秩校正。ARCHead 不存储稠密 BF16 头，并将持久性 LM-head 存储减少 3.7–3.9×。在 Qwen3-8B-Base 上，它仅使用 BF16 头存储的 25.6%，同时达到 1.007 的相对困惑度；存储匹配的朴素 INT4 则得到 1.14–1.16。替换 AWQ 或 bitsandbytes 遗留的 BF16 头仅增加 0.006–0.007 的交叉熵，在我们的测量中吞吐变化小于 2%。因此 ARCHead 通过压缩块量化器可能保持不变的大型输出投影来补充块量化器。代码见 https://github.com/suayptalha/archead。

ARCHead：大语言模型输出头的激活度量残差校正

Şuayp Talha Kocabay
独立研究者
kocabaysuayptalha08@gmail\.com

Talha Rüzgar Akkuş
独立研究者
talharuzgarakkus@gmail\.com

Kamer Ali Yüksel
aiXplain, Inc\.
kamer@aixplain\.com

## 1 引言

GPTQ（Frantar 等人，2023 (https://arxiv.org/html/2608.02703#bib.bib1)）、AWQ（Lin 等人，2025 (https://arxiv.org/html/2608.02703#bib.bib2)）和 bitsandbytes NF4（Dettmers 等人，2023a (https://arxiv.org/html/2608.02703#bib.bib5)）等训练后量化（PTQ）方法使大型语言模型（LLM）的部署变得容易得多。然而，它们的实际实现通常聚焦于 Transformer 块权重，并将最终语言建模头（LM-head）保留为 BF16 或 FP16。该头将最终隐藏状态直接映射为每个词汇项的 logit。与 Transformer 块内部的误差不同，LM-head 误差之后没有归一化、残差路径或另一个学习变换；它直接到达 softmax。

这对大词表模型越来越重要。Qwen3-8B-Base（Yang 等人，2025 (https://arxiv.org/html/2608.02703#bib.bib11)）和 Gemma-4-E4B（Team，2026 (https://arxiv.org/html/2608.02703#bib.bib12)）的 BF16 输出投影分别约占 1.18 和 1.28 GB。一旦 Transformer 块被压缩到四位，这样的投影可能成为最大的剩余稠密张量之一。在我们对量化 Qwen3-8B-Base 检查点的检查中，AWQ 和 bitsandbytes NF4 都将 `lm_head` 保留为形状为 `151,936×4,096` 的 BF16 `Linear` 层。

朴素低比特头量化不是足够的补救措施。它最小化的是权重空间重建误差，该误差平等对待每个隐藏方向，尽管最终隐藏状态呈现各向异性分布。沿频繁激活方向的小误差可能改变许多词汇 logit，而沿几乎不激活方向的较大误差可能影响很小。因此，相关目标是由隐藏状态分布引起的期望输出误差。

我们提出 ARCHead，一种专门的、打包式 LM-head 压缩器。ARCHead 首先用量化低秩因子和分组 INT4 残差表示头。然后在激活派生度量变换之后用低秩分支逼近剩余误差。所得模块是可直接替换的输出头，且不存储原始矩阵的稠密 BF16 副本。

我们的贡献是：

- • 我们在激活派生度量中形式化输出头压缩，并表明在固定量化核心下，ARCHead 校正在因子量化之前是该度量中的最佳秩-\(r\) 逼近。
- • 我们开发了一种打包表示，其测量的状态字典占用为五个模型族对应 BF16 头的 25–27%。
- • 我们表明 ARCHead 避免了存储匹配的朴素 INT4 的较大质量损失，可泛化到三个输出头，并能以很小的额外损失压缩 AWQ 和 bitsandbytes 遗留的 BF16 头。
- • 我们评估了 logit 保真度、下游准确率、校准敏感性、构建时间和生成吞吐量，并将持久参数存储与后端相关的运行时峰值内存分开。

## 2 相关工作

#### LLM 量化。

GPTQ 在顺序权重量化期间应用二阶信息（Frantar 等人，2023 (https://arxiv.org/html/2608.02703#bib.bib1)）；AWQ 保护激活显著通道（Lin 等人，2025 (https://arxiv.org/html/2608.02703#bib.bib2)）；SmoothQuant 在激活和权重之间转移量化难度（Xiao 等人，2023 (https://arxiv.org/html/2608.02703#bib.bib3)）。其他方法优化量化参数（Shao 等人，2024 (https://arxiv.org/html/2608.02703#bib.bib4)）、使用逐层蒸馏（Yao 等人，2022 (https://arxiv.org/html/2608.02703#bib.bib21)）、分离离群值（Dettmers 等人，2023b (https://arxiv.org/html/2608.02703#bib.bib8)；Kim 等人，2024 (https://arxiv.org/html/2608.02703#bib.bib19)），或引入不相干变换和旋转（Chee 等人，2023 (https://arxiv.org/html/2608.02703#bib.bib6)；Tseng 等人，2024 (https://arxiv.org/html/2608.02703#bib.bib7)；Ashkboos 等人，2024 (https://arxiv.org/html/2608.02703#bib.bib9)；Liu 等人，2025 (https://arxiv.org/html/2608.02703#bib.bib10)）。AQLM 使用加法码进行极限压缩（Egiazarian 等人，2024 (https://arxiv.org/html/2608.02703#bib.bib18)），而 LLM\.int8\(\) 和 NF4 提供了广泛使用的混合精度表示（Dettmers 等人，2022 (https://arxiv.org/html/2608.02703#bib.bib16)，2023a (https://arxiv.org/html/2608.02703#bib.bib5)）。这些方法主要面向 Transformer 块线性层；ARCHead 则针对实际流水线可能保持稠密的大型输出投影。

#### 低秩与残差校正。

低秩语言模型压缩（Hsu 等人，2022 (https://arxiv.org/html/2608.02703#bib.bib24)）和参数高效低秩更新（Hu 等人，2022 (https://arxiv.org/html/2608.02703#bib.bib20)）表明结构化矩阵分量可以紧凑表示。ARCHead 在目标和目标函数上都有所不同：它计算*量化残差*的低秩逼近，而该残差先经过由 LM-head 输入估计的度量变换。这直接将校正容量分配给影响观测 logit 的方向。

#### 实际存储。

一个预期的比特宽度本身并不能确立可部署的压缩：实现可能保留反量化张量，或在序列化状态中物化一个张量。因此，我们根据打包头实际注册的张量来测量字节数，并将持久存储与临时前向传播内存分开报告。

## 3 方法

### 3\.1 激活度量目标

设稠密输出头权重为 \(\mathbf{W}\in\mathbb{R}^{V\times D}\)，并设 \(\mathbf{H}\in\mathbb{R}^{N\times D}\) 包含来自 \(N\) 个校准 token 的最终隐藏状态。对于近似 \(\widehat{\mathbf{W}}\)，记 \(\boldsymbol{\Delta}=\mathbf{W}-\widehat{\mathbf{W}}\)。经验平方 logit 误差为

\[
\mathcal{E}=\frac{1}{N}\lVert\mathbf{H}\boldsymbol{\Delta}^{\top}\rVert_{F}^{2}=\operatorname{Tr}\!\left(\boldsymbol{\Delta}\mathbf{C}\boldsymbol{\Delta}^{\top}\right),
\tag{1}
\]
其中 \(\mathbf{C}=\mathbf{H}^{\top}\mathbf{H}/N\)。Frobenius 权重目标是 \(\mathbf{C}=\mathbf{I}\) 的特例，并忽略激活几何。

ARCHead 使用阻尼度量变换。设 \(\mathbf{C}_{\lambda}=\mathbf{C}+\lambda\bar{c}\mathbf{I}=\mathbf{Q}\boldsymbol{\Lambda}\mathbf{Q}^{\top}\)，其中 \(\bar{c}\) 是 \(\mathbf{C}\) 的对角线均值，定义

\[
\mathbf{T}_{p}=\mathbf{Q}\boldsymbol{\Lambda}^{p}\mathbf{Q}^{\top},\qquad
\mathbf{T}_{p}^{-1}=\mathbf{Q}\boldsymbol{\Lambda}^{-p}\mathbf{Q}^{\top}.
\tag{2}
\]
当 \(p=\tfrac{1}{2}\) 时，诱导目标是公式 1 (https://arxiv.org/html/2608.02703#S3.E1) 的阻尼形式；其他值调整主导激活方向被强调的强度。我们在校准数据上选择 \(p\)，并对 Qwen3-8B-Base 使用 \(p=0.75\)。

### 3\.2 核心与残差校正

参见图 1 (https://arxiv.org/html/2608.02703#S3.F1) 标题：ARCHead 将稠密输出头分解为打包量化核心 \(\mathbf{W}_{d}\) 和激活度量低秩校正。原始 BF16 矩阵在构建后被丢弃。

量化核心为

\[
\mathbf{W}_{d}=Q_{5}(\mathbf{A}_{c})Q_{8}(\mathbf{B}_{c})+Q_{\mathrm{SC4}}\!\left(\mathbf{W}-\mathbf{A}_{c}\mathbf{B}_{c}\right),
\tag{3}
\]
其中 \(\mathbf{A}_{c}\mathbf{B}_{c}\) 是秩 \(r_{c}\) 近似，\(Q_{b}\) 表示 \(b\) 位分组量化，SC4 表示带量化分组尺度的打包有符号 INT4 残差。该核心留下的残差为 \(\mathbf{E}=\mathbf{W}-\mathbf{W}_{d}\)。

ARCHead 计算秩 \(r_{r}\) 的随机截断 SVD

\[
\mathbf{E}\mathbf{T}_{p}\approx\mathbf{U}_{r_{r}}\boldsymbol{\Sigma}_{r_{r}}\mathbf{V}_{r_{r}}^{\top}
\tag{4}
\]
并将右因子映射回来：

\[
\mathbf{A}_{w}=\mathbf{U}_{r_{r}}\boldsymbol{\Sigma}_{r_{r}},\qquad
\mathbf{B}_{w}=\mathbf{V}_{r_{r}}^{\top}\mathbf{T}_{p}^{-1}.
\tag{5}
\]
最终近似和 logit 为

\[
\widehat{\mathbf{W}}=\mathbf{W}_{d}+\mathbf{A}_{w}\mathbf{B}_{w},
\tag{6}
\]
\[
\widehat{\mathbf{Y}}=\mathbf{H}\mathbf{W}_{d}^{\top}+(\mathbf{H}\mathbf{B}_{w}^{\top})\mathbf{A}_{w}^{\top}.
\tag{7}
\]
对于打包的 ARCHead 变体，两个校正因子都以行/分组 INT8 量化存储。

#### 条件最优性。

对于固定的 \(\mathbf{W}_{d}\) 和可逆的 \(\mathbf{T}_{p}\)，上述未量化因子在所有秩 \(r_{r}\) 乘积 \(\mathbf{A}\mathbf{B}\) 上最小化

\[
\left\lVert(\mathbf{E}-\mathbf{A}\mathbf{B})\mathbf{T}_{p}\right\rVert_{F}^{2}
\tag{8}
\]
这可以直接从 Eckart–Young–Mirsky 定理得出，只要将变量改为 \(\widetilde{\mathbf{E}}=\mathbf{E}\mathbf{T}_{p}\)。该保证以固定核心为条件，并在校正因子量化之前成立。

#### 证明。

设 \(\widetilde{\mathbf{B}}=\mathbf{B}\mathbf{T}_{p}\)。因为 \(\mathbf{T}_{p}\) 可逆，\(\operatorname{rank}(\mathbf{A}\widetilde{\mathbf{B}})=\operatorname{rank}(\mathbf{A}\mathbf{B})\)，所以问题等价于

\[
\min_{\operatorname{rank}(\mathbf{A}\widetilde{\mathbf{B}})\leq r_{r}}\left\lVert\widetilde{\mathbf{E}}-\mathbf{A}\widetilde{\mathbf{B}}\right\rVert_{F}^{2}.
\tag{9}
\]
根据 Eckart–Young–Mirsky 定理，最小化器是 \(\widetilde{\mathbf{E}}\) 的秩 \(r_{r}\) 截断 SVD \(\mathbf{U}_{r_{r}}\boldsymbol{\Sigma}_{r_{r}}\mathbf{V}_{r_{r}}^{\top}\)。选择 \(\mathbf{A}=\mathbf{U}_{r_{r}}\boldsymbol{\Sigma}_{r_{r}}\) 和 \(\mathbf{B}=\mathbf{V}_{r_{r}}^{\top}\mathbf{T}_{p}^{-1}\) 得到所述因子。对于 \(p=\tfrac{1}{2}\)，目标函数是阻尼经验 logit MSE。我们的随机截断 SVD 和随后的因子量化近似这一理想校正，因此结果并不声称量化端到端模块的全局最优性。

### 3\.3 打包表示与集成

打包的 ARCHead 仅注册 SC4 残差及尺度、量化核心因子、量化校正因子和形状元数据。其持久大小直接测量为

\[
\operatorname{bytes}(\mathrm{ARCHead})=\sum_{i}\operatorname{numel}(\theta_{i})
\tag{10}
\]
\[
\quad{}\times\operatorname{element\_size}(\theta_{i}),
\]
涵盖序列化模块中的所有参数和缓冲区 \(\theta_{i}\)。原始 \(V\times D\) BF16 矩阵不被注册。在模型转换期间，稠密头用于构建 ARCHead，然后被丢弃。该模块可以在块量化之后替换输出嵌入，因此 ARCHead 与 AWQ、bitsandbytes 或其他块后端互补。算法 1 (https://arxiv.org/html/2608.02703#alg1) 给出了完整构建过程。

#### 状态字典核算。

打包实现不仅仅是理论上的比特计数方案：与在加载时保留反量化稠密张量的原型不同，它不存储稠密 BF16 \(V\times D\) 头。稠密存储为 \(VD\times 2\) 字节，而打包存储是上面在注册 `state_dict` 上测量的总和。存储的分量包括：打包 INT4 残差及其 FP16/BF16 尺度；5 位左核心因子和 8 位右核心因子及尺度；行/分组 INT8 校正因子 \(\mathbf{A}_{w}\) 和 \(\mathbf{B}_{w}\)；量化尺度和零点元数据；以及组大小和形状元数据。表 1 (https://arxiv.org/html/2608.02703#S3.T1) 报告了从实例化 PyTorch 缓冲区测量的值，而非理论比特宽度。

表 1：详细的打包 ARCHead 存储验证，根据实例化 PyTorch `state_dict` 缓冲区计算。

#### 加载时与峰值内存。

持久 LM-head 内存在将模型加载到 GPU 时衡量参数分配，是评估存储输出头占用空间的适当指标。表 2 (https://arxiv.org/html/2608.02703#S3.T2) 报告了相应的节省。前向传播峰值内存是单独的：它包含激活、输出 logit、临时缓冲区、CUDA 分配器行为、后端工作区和 KV 缓存。ARCHead 减少持久输出头参数；我们不将其与后端和工作负载相关的运行时峰值内存混为一谈。

表 2：加载时持久 LM-head 内存节省。

#### 峰值内存测量协议。

测量之前，我们执行 `torch.cuda.empty_cache()`，重置 PyTorch 峰值内存统计，并在执行前后同步 CUDA 设备。我们分别测量和报告加载时持久参数内存与前向传播峰值，后者固有地包含动态 logit、临时缓冲区、分配器状态和后端工作区；因此，我们避免关于后端相关峰值内存减少的无根据声称。

### 3\.4 构建算法

输入为稠密 LM-head 权重 \(\mathbf{W}\in\mathbb{R}^{V\times D}\)、隐藏校准激活 \(\mathbf{H}\in\mathbb{R}^{N\times D}\)、核心秩和校正秩 \(r_{c}\) 与 \(r_{r}\)、组大小 \(g\)、度量幂 \(p\) 和阻尼 \(\lambda\)。该过程首先估计阻尼激活协方差 \(\mathbf{C}=