批量自适应剪枝：面向语言推理模型的周期性神经元激活感知权重剪枝

# 批量自适应剪枝：面向语言推理模型的周期性神经元激活感知权重剪枝
来源：https://arxiv.org/html/2608.14003  
作者：Shota Takashiro<sup>11</sup>、Yusuke Iwasawa、Takeshi Kojima、Yutaka Matsuo  
所属单位：东京大学

###### 摘要
大型推理模型（LRMs）通过扩展链式思维生成在复杂任务上取得了优异性能，但推理过程产生了巨大的计算成本。在生产环境中，批处理推理对于实现高吞吐量至关重要，而我们评估的现有免训练自适应剪枝方法在此场景下性能严重下降。由于一个批次必须共享单个剪枝掩码，这些方法需要跨样本聚合激活值，然后应用基于阈值的筛选；该阈值是在非聚合激活值上离线校准的，无法匹配聚合后的分布，导致实际稀疏率发生漂移，并在批处理推理中使推理任务的准确率崩溃。本文提出一种专门针对LRMs批处理推理的免训练自适应剪枝方法，该方法基于两个核心组件。首先，我们用**周期性top-k选择**替代基于阈值的筛选，该选择基于聚合后的重要性分数，不受聚合引起的激活分布偏移影响，且每个更新周期只执行一次选择（而非每个token都执行），从而保持加速效果。其次，基于重要神经元在长推理生成过程中会周期性重新激活的观察，我们引入**激活记忆**机制，在更新阶段间累积重要性，从而保留重复激活的神经元。在多个推理基准上的实验表明，我们的方法在DeepSeek-R1-Distill-Qwen-7B上，以4的批处理大小和50%目标稀疏率进行评估时，平均准确率比现有最先进的自适应剪枝方法提高了39.7个百分点，并在50%实际稀疏率下实现了相比密集推理1.40倍的加速。<sup>11</sup>
<sup>11</sup> 我们的代码已开源：https://github.com/matsuolab/batch-wise-prune

参见图1说明  
图1：（左上）现有自适应剪枝方法<cit.>使用离线校准的阈值选择神经元，因此在聚合的激活分布下实际稀疏率发生漂移，性能下降。（左中）top-k选择不受该偏移影响，但在每个解码步骤都执行会产生计算开销，抵消了加速效果。（左下）我们的方法**周期性**通过top-k选择更新剪枝掩码，并通过激活记忆跨更新周期传递重要性，在保持计算效率的同时维持了推理性能。（右）DeepSeek-R1-Distill-Qwen-7B某前馈层中\(\mathbf{Z}\)的激活值可视化，蓝色水平线表示输入边界。具有大激活幅度的神经元（红色高亮显示）表现出高重复激活率和短重复激活周期，这启发了我们的周期性掩码更新策略。

## 1 引言
大型语言模型（LLMs）在众多任务上展现出卓越的能力<cit.>，但它们伴随着巨大的计算成本。特别是，近期推理能力的进步催生了大型推理模型（LRMs）<cit.>，这些模型采用扩展的链式思维过程来解决复杂问题。虽然有效，但这类模型生成的输出长度显著更长，进一步加剧了推理过程的计算成本。在生产环境中，通过批处理推理（将多个输入样本作为一个批次一起处理）同时服务多个请求对于实现高服务效率至关重要<cit.>，这些计算瓶颈变得尤为关键。为了进一步降低请求延迟，我们需要能在批处理设置中高效运行的更轻量级模型。一种减少推理成本的有前景的方法是剪枝。现有的剪枝方法可分为**静态**或**自适应**两类。静态方法<cit.>在解码前确定剪枝模式（基于校准数据或提示），并在整个生成过程中重复使用；然而，我们观察到，在推理任务中，当激活模式在链式思维生成过程中演变时，这种固定方法会导致性能下降。自适应方法<cit.>通过根据激活幅度在每个步骤动态选择神经元来解决此问题，在复杂推理任务上取得了强劲性能。然而，现有的自适应剪枝方法（如CATS<cit.>和TEAL<cit.>）在批处理推理场景中面临根本性限制。在GPU上，批处理推理要求批次中所有样本共享单个剪枝掩码，因此必须先将每个样本的激活值聚合成一个分数向量。失败之处发生在聚合之后：这些方法使用在外部语料库上离线校准的阈值选择神经元。该阈值是针对单个样本的激活分布调整的，而跨批次的聚合改变了其应用的分布，因此在生成过程中实际稀疏率偏离目标。在我们方法的其他部分保持不变的情况下，用top-k选择替代该阈值选择，是在批处理下保持准确率的关键（第4.4节），这确定了基于阈值的选择是TEAL在批处理下崩溃的主要原因。第二个困难特定于推理工作负载：激活模式在数千个生成的token上不断演变，因此在生成结束前很久选择的掩码就已过时，而重新在每个解码步骤选择掩码则代价高昂，无法留下任何加速效果。本文，我们提出一种针对LRMs的免训练自适应剪枝方法，能够在保持推理任务强劲性能的同时实现高效的批处理推理。我们的贡献有两方面。（i）据我们所知，我们首次通过实验证明，我们评估的现有免训练剪枝方法在批处理推理下会严重降低推理任务性能。这两个系列方法失败的原因不同：静态方法（Wanda、Griffin；在两个DeepSeek-R1-Distill模型上评估）在任何批处理大小下都会在推理任务上性能下降，因为解码前固定的剪枝模式（Wanda基于校准数据，Griffin基于提示）无法跟踪长链思维中演变的激活；而自适应方法在批处理大小为1时表现良好，但一旦必须共享掩码就会崩溃：TEAL在所有评估的四个模型上，在50%目标稀疏率下，相比密集模型丢失了58.4到67.7个点的平均准确率。（ii）我们提出一种方法，通过两个组件在此场景下保留了显著更高的准确率：**周期性top-k选择**，不受聚合引起的激活分布偏移影响，且每个更新周期只运行一次选择（而非每个token）；以及**激活记忆**，在更新阶段间累积重要性，从而保留随时间重复激活的神经元，这基于重要神经元在长推理生成过程中会周期性重新激活的观察（图1；扩展可视化见附录Q）。跨样本聚合是第三个设计轴，但我们并未声称其是我们性能提升的来源：我们使用逐元素最大值而非均值进行聚合（详见第3.2节中的批量聚合），且第4.4节的实验表明准确率对此选择基本不敏感。仅依赖运行时激活也使我们的方法无需校准，而Wanda和TEAL在部署前都必须在外部语料库上进行校准。通过在多个推理任务上的大量实验，我们证明我们的方法在批处理设置中保持了良好性能，而现有的自适应剪枝方法则出现显著性能下降。具体而言，在DeepSeek-R1-Distill-Qwen-7B（DS-R1-Qwen-7B）上，以4的批处理大小和50%目标稀疏率评估时，我们的方法在平均准确率上比先前最先进的自适应剪枝方法高出39.7个百分点。此外，在50%实际稀疏率下，我们的方法实现了相比密集推理1.40倍的加速。该方法适用及不适用的场景在附录A中说明。

表1：各种剪枝方法在不同推理设置下的实验总结。我们的方法（*）在所有设置中均能保持性能一致，而现有方法在推理任务或批处理推理上性能下降。
| 任务类型 | 批处理大小 | 静态剪枝 | 现有自适应剪枝 | 我们的方法 |
| :--- | :---: | :---: | :---: | :---: |
| 非推理 | 1 | ✓ | ✓ | ✓ |
| 非推理 | ≥2 | ✓ | ✓ | ✓ |
| 推理 | 1 | ✗ | ✓ | ✓ |
| 推理 | ≥2 | ✗ | ✗ | ✓* |

## 2 相关工作
针对大型语言模型的剪枝方法可沿多个维度分类：粒度、输入依赖性和自适应性。表1总结了现有方法与我们的方法之间的关键差异。

#### 剪枝粒度。
剪枝通过移除冗余参数来减小模型大小。非结构化剪枝移除单个权重<cit.>，但需要专用硬件进行加速。半结构化剪枝（如N:M稀疏性）也依赖特定的GPU架构<cit.>。结构化剪枝<cit.>通常需要训练后处理。我们的方法通过标准矩阵运算执行结构化剪枝，可在任何硬件上实现加速，无需专用内核或训练后处理。

#### 输入无关 vs. 输入依赖剪枝。
输入无关方法在推理前使用校准数据确定固定的剪枝模式<cit.>，受益于批处理兼容性，但无法利用特定于输入的激活模式。输入依赖方法根据运行时激活动态选择神经元<cit.>，但为批处理推理带来挑战，因为不同样本需要不同的剪枝模式。

#### 静态 vs. 自适应剪枝。
静态剪枝确定一个固定的稀疏模式并在整个生成过程中重复使用<cit.>，而自适应剪枝在解码过程中动态更新掩码<cit.>。静态剪枝无法适应演变的激活模式，导致在输出较长的推理模型上性能下降。现有自适应方法依赖于基于阈值的筛选，这在跨批次聚合时会遭受分布偏移问题。此外，DejaVu<cit.>和ProSparse<cit.>在单个样本级别运行，使批处理推理具有挑战性。我们的方法通过周期性top-k选择结合激活记忆解决了这些局限性，该机制在掩码更新间传递重要性，详见第3节。

图2：我们的批量自适应剪枝方法概览。（a）为每个样本计算重要性分数并聚合成一个分数向量，因为一个批次必须共享一个剪枝掩码；我们默认使用逐元素最大值作为聚合算子。（b）该方法分三个阶段运行：输入处理、密集探索和周期性自适应剪枝，激活记忆在连续的掩码更新间传递重要性分数。

## 3 方法
本节介绍我们用于高效LLM推理的批量自适应剪枝方法。我们的方法是免训练的，无需微调；剪枝决策实时做出，仅依赖推理过程中观察到的模型激活。我们首先介绍门控MLP结构的基础知识（第3.1节），该结构在现代基于Transformer的LLM中占绝大多数参数，因此是我们方法的剪枝目标。接下来，描述我们的批量重要性评分机制（第3.2节），该机制定义了确定在门控MLP中应剪枝哪些神经元的度量。最后，详述我们的三阶段剪枝算法（第3.3节），该算法使用定义的批量重要性分数进行周期性剪枝。

### 3.1 基础知识
现代Transformer架构在其前馈模块中采用门控MLP结构<cit.>。给定输入隐藏状态序列\(\mathbf{X} \in \mathbb{R}^{T \times D}\)，其中\(T\)是token数量，每个隐藏状态\(\mathbf{x} \in \mathbb{R}^{D}\)由门控MLP处理：
\[
\mathbf{z} = \sigma(\mathbf{W}_g \mathbf{x}) \odot (\mathbf{W}_1 \mathbf{x}) \tag{1}
\]
其中\(\sigma\)表示SiLU激活函数，\(\odot\)表示逐元素乘法。对于所有权重矩阵，\(\mathbf{W}_1, \mathbf{W}_g \in \mathbb{R}^{D_{\text{FF}} \times D}\)和\(\mathbf{W}_2 \in \mathbb{R}^{D \times D_{\text{FF}}}\)，通常\(D_{\text{FF}} \gg D\)。我们称\(\mathbf{z} = \text{FF}_1(\mathbf{x}) \in \mathbb{R}^{D_{\text{FF}}}\)为FF激活值，输出计算为\(\text{FF}_2(\mathbf{z}) = \mathbf{W}_2 \mathbf{z}\)。

### 3.2 批量重要性神经元评分
我们定义一个**重要性分数**\(\mathbf{s} \in \mathbb{R}^{D_{\text{FF}}}\)...