提出多粒度节点剪枝法，更精细高效地发现大模型行为电路。
Multi-Granular Node Pruning for Causal Circuit Discovery
- 在块到单个神经元的多粒度上统一优化，用可学习掩码实现剪枝。
- 剪出的电路节点数减少30%以上，且性能不受影响。
- 内存占用降低5-10倍，无需保存中间激活值，适合大规模模型。
电路发现旨在识别大型语言模型中负责特定行为的最小子网络。现有方法主要依赖迭代边剪枝，计算成本高，且仅限于注意力头或MLP块等粗粒度单元，忽略了如单个神经元等更细粒度结构。本文提出一种面向电路发现的节点级剪枝框架，同时解决可扩展性与粒度限制问题。该方法在统一优化目标下引入跨多粒度层级（从完整模块到单个神经元）的可学习掩码，并通过粒度特异性稀疏惩罚引导剪枝过程，可在一次微调运行中实现全面压缩。实验证明，本方法发现的电路在节点数量上比以往方法减少30%以上；此外，许多被粗粒度方法认为重要的神经元实际无关紧要，但任务性能得以保持。更重要的是，该方法内存占用降低5-10倍，因无需存储中间激活值即可运行。
原文摘要 · Abstract (English)
Circuit discovery aims to identify minimal subnetworks that are responsible for specific behaviors in large language models (LLMs). Existing approaches primarily rely on iterative edge pruning, which is computationally expensive and limited to coarse-grained units such as attention heads or MLP blocks, overlooking finer structures like individual neurons. We propose a node-level pruning framework for circuit discovery that addresses both scalability and granularity limitations. Our method introduces learnable masks across multiple levels of granularity, from entire blocks to individual neurons, within a unified optimization objective. Granularity-specific sparsity penalties guide the pruning process, allowing a comprehensive compression in a single fine-tuning run. Empirically, our approach identifies circuits that are smaller in nodes than those discovered by prior methods; moreover, we demonstrate that many neurons deemed important by coarse methods are actually irrelevant, while still maintaining task performance. Furthermore, our method has a significantly lower memory footprint, 5-10x, as it does not require keeping intermediate activations in the memory to work.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。