因果归因剪枝：在保持大语言模型推理性能的同时降低计算成本

背景

大型语言模型（LLMs）在涉及多步逻辑推理的任务中表现出色，例如数学解题、复杂问答及科学推理。然而，随着模型参数规模的指数级增长，其推理成本（Inference Cost）也变得极其高昂。为了在资源受限的环境中部署这些模型，模型剪枝（Pruning）——即移除模型中冗余的参数或连接——成为了一种关键的压缩技术。

传统的剪枝方法通常基于“幅度”（Magnitude，即权重绝对值的大小）或“激活值”（Activation，即神经元在特定输入下的活跃程度）来评估参数的重要性。然而，这些相关性指标往往无法准确反映某个参数或模块在复杂推理链条中的实际功能贡献。特别是在注意力机制（Attention Mechanism）中，某些权重虽然在数值上较小，但在特定的因果推理路径中可能起着决定性作用。因此，如何在不显著牺牲推理精度的前提下，更精准地识别并剪枝冗余参数，是当前大模型优化领域的一大挑战。

核心内容

本文提出了一种名为**因果归因剪枝（Causal Attribution Pruning, CAP）**的训练无关（Training-free）方法。该方法旨在通过测量注意力头（Attention Heads）对推理任务的因果影响，来识别关键组件，并据此指导细粒度的权重剪枝。

1. 核心机制：从注意力头到权重的因果评估

CAP 的核心创新在于引入了“干预性测量”（Interventional Measurement）。其工作流程如下：

• 校准集构建：选取一小部分具有代表性的推理问题作为校准集（Calibration Set）。
• 因果得分估计：对于每一个注意力头，CAP 会在前向传播过程中将其屏蔽（Masked），即阻断该头的输出。通过比较屏蔽前后模型在推理任务上的性能变化，估算该头被移除时导致的预期性能下降。这种性能下降即为该注意力头的“因果得分”。
• 权重级重要性转化：将注意力头级别的因果得分转化为对应投影矩阵（Projection Matrices）中具体权重的重要性值。这意味着，如果一个注意力头对推理至关重要，那么构成该头的权重将被赋予较高的重要性评分，从而在剪枝时被保留；反之，则被剪枝。

2. 与现有方法的对比

与仅依赖幅度（Magnitude-only）或激活值（Activation-based）的剪枝标准不同，CAP 直接捕捉每个注意力头的功能贡献。

• Wanda：作为当前主流的训练无关剪枝基线，Wanda 主要基于激活值和权重的乘积幅度来评估重要性。
• CAP 的优势：实验数据显示，在 20% 的稀疏度（Sparsity，即移除 20% 的参数）下，CAP 在 ARC-Challenge 基准测试中相比 Wanda 实现了高达 61% 的相对准确率提升。这表明，因果归因能更准确地保留那些对推理逻辑至关重要的“小权重”或“低激活”参数。

3. 实验评估

研究团队在以下模型和基准上对 CAP 进行了全面评估：

• 模型：Llama-3-8B-Instruct 和 Mistral-7B-Instruct。
• 基准数据集：
  • GSM8K：小学级数学推理。
  • StrategyQA：需要多步策略推理的问答。
  • ARC-Challenge：科学推理挑战集。
• 稀疏度设置：10%、20% 和 50%。

4. 主要发现

• 中等稀疏度下的优势：在 10%-20% 的中等稀疏度范围内，CAP 在大多数模型-基准配置中均优于 Wanda。特别是在 Llama-3 模型上，ARC-Challenge 的性能提升尤为显著。
• 高稀疏度的局限性：当稀疏度达到 50% 时，CAP 的性能提升受到限制。作者指出，这主要是因为 MLP（多层感知机）部分的归因粒度较粗（Coarse MLP Attribution），导致在极高压缩率下难以精准保留所有关键信息。

关键要点

• 方法创新：提出了 CAP（Causal Attribution Pruning），一种无需重新训练的剪枝方法，通过干预性测量评估注意力头对推理任务的因果影响。
• 评估指标：不再单纯依赖权重幅度，而是通过屏蔽注意力头并观察性能下降幅度来量化其重要性，从而将头级重要性转化为权重级重要性。
• 性能提升：在 20% 稀疏度下，CAP 在 ARC-Challenge 上相比基线方法 Wanda 实现了高达 61% 的相对准确率增益。
• 适用场景：在 10%-20% 的中等稀疏度下表现最佳，能有效保留 LLM 的多步推理能力。
• 当前局限：在 50% 的高稀疏度下，由于 MLP 层归因粒度较粗，性能提升受限，表明该方法在极端压缩场景下仍有优化空间。
• 模型兼容性：已在 Llama-3-8B-Instruct 和 Mistral-7B-Instruct 等主流开源模型上验证有效性。

意义与影响

这项研究对大语言模型的轻量化部署具有重要的理论和实践意义：

1. 重新定义“重要性”：CAP 证明了在推理任务中，参数的“因果贡献”比其“数值幅度”或“激活频率”更能反映其真实价值。这为未来的模型压缩算法提供了新的评估视角，即从相关性分析转向因果干预分析。
2. 降低推理门槛：通过更高效的剪枝，可以在几乎不损失推理精度的情况下大幅减少模型的计算量和内存占用。这对于在边缘设备、移动端或资源受限的服务器上部署强大的推理模型至关重要。
3. 优化资源分配：CAP 方法无需额外的训练步骤，仅通过一次前向传播的干预测量即可完成剪枝指导，极大地降低了模型压缩的时间成本和计算开销。
4. 指明未来方向：研究指出的 MLP 归因粒度问题，为后续工作提供了明确的技术突破口。未来的研究可能需要结合更细粒度的 MLP 因果归因方法，以突破高稀疏度下的性能瓶颈，实现更极致的模型压缩。

总之，CAP 提供了一种在“效率”与“智能”之间取得更好平衡的新途径，使得保留大模型复杂推理能力的同时，能够更大幅度地削减冗余计算，推动了大模型向更普惠、更高效的方向发展。