Multi-Granular Node Pruning for Causal Circuit Discovery

Causal circuit discovery aims to identify minimal subnetworks that causally drive specific behaviors in large language models (LLMs). Existing approaches focus on edge pruning or unstructured weight pruning. These methods are computationally expensive and typically operate on coarse-grained components, such as attention heads or MLP blocks, thereby missing finer-grained structure. 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 more compact circuits than prior methods, 33.34% more MLPs, and 59.8% more neurons in the least favorable setting, with larger gains overall. We further demonstrate that many neurons deemed important by coarse methods are, in fact, irrelevant and can be removed with negligible impact on task performance. Our method is also memory-efficient, requiring at least 3× less memory as it avoids storing intermediate activations in memory to work.