基于因果归因的剪枝方法 CAP 可保持大语言模型推理性能

View PDF HTML (experimental)

Abstract:Large language models (LLMs) excel at multi-step reasoning but incur substantial inference cost. We introduce Causal Attribution Pruning (CAP), a training-free method that identifies critical attention heads by measuring their causal impact on reasoning tasks and uses these head-level scores to guide fine-grained weight pruning. For each attention head, CAP estimates the expected performance degradation when the head is masked during forward passes on a small calibration set of reasoning problems. These causal scores are then converted into weight-level importance values for the corresponding projection matrices. Unlike magnitude-only or activation-based criteria, CAP's interventional measurement directly captures each head's functional contribution, yielding relative accuracy gains of up to 61% over Wanda on ARC-Challenge at 20% sparsity. We evaluate CAP on GSM8K, StrategyQA, and ARC-Challenge using Llama-3-8B-Instruct and Mistral-7B-Instruct at 10%, 20%, and 50% sparsity. At moderate sparsity (10-20%), CAP improves over Wanda in most model-benchmark configurations. with especially large gains on ARC-Challenge for Llama-3. Our results suggest that attention-head-level causal attribution can better preserve reasoning performance on downstream benchmarks than correlational pruning criteria at equivalent sparsity, while remaining limited by coarse MLP attribution at 50% sparsity.

Comments:	Accepted at the ICLR 2026 Workshop on LLM Reasoning. 13 pages, 2 figures
Subjects:	Computation and Language (cs.CL)
Cite as:	arXiv:2606.19350 [cs.CL]
	(or arXiv:2606.19350v2 [cs.CL] for this version)
	https://doi.org/10.48550/arXiv.2606.19350 arXiv-issued DOI via DataCite

Submission history

From: Yuhao Ge [view email]
[v1] Mon, 27 Apr 2026 01:44:10 UTC (89 KB)
[v2] Tue, 27 Oct 2026 05:55:55 UTC (89 KB)