Probe Pruning: Accelerating LLMs through Dynamic Pruning via Model-Probing
作者: Qi Le, Enmao Diao, Ziyan Wang, Xinran Wang, Jie Ding, Li Yang, Ali Anwar
分类: cs.CL, cs.AI, cs.LG
发布日期: 2025-02-21
备注: ICLR 2025
🔗 代码/项目: GITHUB
💡 一句话要点
Probe Pruning：通过模型探测实现LLM的动态剪枝加速
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 模型剪枝 动态剪枝 模型加速 在线推理 模型探测 结构化剪枝 LLM
📋 核心要点
- 现有LLM剪枝方法缺乏动态性和样本感知能力，无法针对不同输入进行优化。
- Probe Pruning通过少量前向计算探测关键权重，并结合历史信息进行动态剪枝。
- 实验表明，仅用少量FLOPs探测即可显著提升LLM结构化剪枝效率，性能优于现有方法。
📝 摘要（中文）
本文提出了一种名为Probe Pruning (PP) 的新型框架，用于在线、动态、结构化地剪枝大型语言模型 (LLM)，并以批处理方式应用。PP 的核心思想是并非所有样本和 token 对模型的输出贡献相同。通过探测每个 batch 的一小部分，可以有效地识别关键权重，从而为不同的 batch 定制动态剪枝。PP 包括三个主要阶段：探测、历史信息指导的剪枝和完整推理。在探测阶段，PP 基于残差重要性选择一小部分关键的隐藏状态，并提前运行几个模型层。在历史信息指导的剪枝阶段，PP 将探测状态与历史状态进行策略性整合。随后，它基于整合后的状态和 PP 重要性得分（一种专门用于评估每个权重通道在维持性能方面的重要性的指标）进行结构化权重剪枝。在最后阶段，对剩余权重进行完整推理。PP 的一个主要优点是它与现有模型兼容，因为它无需额外的神经网络模块或微调。对 LLaMA-2/3 和 OPT 模型进行的全面评估表明，即使是最小的探测（仅使用 1.5% 的 FLOPs）也可以显着提高 LLM 结构化剪枝的效率。例如，在 WikiText2 上评估 LLaMA-2-7B 时，在 40% 的剪枝率下，PP 实现的每次运行时减少的性能下降率比最先进的方法低 2.56 倍。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLM）的计算效率问题，具体而言，通过动态剪枝减少推理过程中的计算量。现有静态剪枝方法无法适应不同输入样本的需求，导致性能下降或剪枝率受限。动态剪枝虽然可以根据输入调整模型结构，但通常需要额外的训练或复杂的控制机制，增加了部署成本。
核心思路：论文的核心思路是利用少量的前向计算（称为“探测”）来估计模型中不同权重的重要性，并根据这些重要性动态地剪枝模型。这种方法的核心在于，并非所有权重对所有输入都同等重要，通过选择性地保留关键权重，可以在显著减少计算量的同时，尽可能地保持模型性能。
技术框架：Probe Pruning (PP) 框架包含三个主要阶段：1) 探测 (Probing)：选择一小部分关键的隐藏状态，基于残差的重要性，并运行几个模型层。2) 历史信息指导的剪枝 (History-informed Pruning)：将探测阶段的状态与历史状态进行整合，以获得更全面的权重重要性评估。3) 完整推理 (Full Inference)：基于整合后的状态和 PP 重要性得分进行结构化权重剪枝，然后对剩余的权重进行完整推理。
关键创新：PP 的关键创新在于其动态性和高效性。它不需要额外的训练或微调，可以直接应用于现有的 LLM。通过少量的前向计算，PP 能够准确地估计权重的重要性，并根据这些重要性进行动态剪枝。此外，PP 框架还引入了 PP 重要性得分，这是一种专门用于评估每个权重通道在维持性能方面的重要性的指标。
关键设计：在探测阶段，选择残差较大的隐藏状态作为关键状态，因为这些状态通常包含更多的信息。历史信息的整合采用加权平均的方式，权重可以根据经验或通过学习得到。PP 重要性得分的设计需要平衡计算复杂度和准确性，论文中可能采用了某种启发式方法或可学习的函数来计算该得分。具体的剪枝策略（例如，剪枝比例、剪枝粒度）也需要根据实际情况进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Probe Pruning 在 LLaMA-2/3 和 OPT 模型上实现了显著的加速效果。例如，在 LLaMA-2-7B 模型上，仅使用 1.5% 的 FLOPs 进行探测，即可在 40% 的剪枝率下，实现比现有最佳方法低 2.56 倍的性能下降率。这表明 Probe Pruning 能够在保持模型性能的同时，有效地减少计算量。
🎯 应用场景
Probe Pruning 可应用于各种需要加速LLM推理的场景，如移动设备、边缘计算和高并发服务器。通过动态剪枝，可以在资源受限的环境中部署更大规模的LLM，提升用户体验。该技术还有助于降低云计算成本，提高AI服务的可扩展性，并促进LLM在更多领域的应用。
📄 摘要（原文）
We introduce Probe Pruning (PP), a novel framework for online, dynamic, structured pruning of Large Language Models (LLMs) applied in a batch-wise manner. PP leverages the insight that not all samples and tokens contribute equally to the model's output, and probing a small portion of each batch effectively identifies crucial weights, enabling tailored dynamic pruning for different batches. It comprises three main stages: probing, history-informed pruning, and full inference. In the probing stage, PP selects a small yet crucial set of hidden states, based on residual importance, to run a few model layers ahead. During the history-informed pruning stage, PP strategically integrates the probing states with historical states. Subsequently, it structurally prunes weights based on the integrated states and the PP importance score, a metric developed specifically to assess the importance of each weight channel in maintaining performance. In the final stage, full inference is conducted on the remaining weights. A major advantage of PP is its compatibility with existing models, as it operates without requiring additional neural network modules or fine-tuning. Comprehensive evaluations of PP on LLaMA-2/3 and OPT models reveal that even minimal probing-using just 1.5% of FLOPs-can substantially enhance the efficiency of structured pruning of LLMs. For instance, when evaluated on LLaMA-2-7B with WikiText2, PP achieves a 2.56 times lower ratio of performance degradation per unit of runtime reduction compared to the state-of-the-art method at a 40% pruning ratio. Our code is available at https://github.com/Qi-Le1/Probe_Pruning.