Multi-Granular Node Pruning for Causal Circuit Discovery

AI summary

Overview

Research area: Mechanistic interpretability of large language models — specifically automated causal circuit discovery, the task of isolating the minimal subnetwork inside a trained LLM that is responsible for a particular behavior.

Technical level: Intermediate. The paper assumes familiarity with transformer internals (attention heads, MLP blocks, residual streams, up/down projections) and with the circuit-discovery literature (ACDC, EAP, edge pruning, activation patching), but its core idea — learning masks over nodes at several levels of granularity — is conceptually straightforward.

Scope: The paper proposes a node-level, multi-granularity pruning framework that learns which transformer blocks, attention heads, MLP blocks, and individual neurons to keep in a single optimization run, and evaluates it on three circuit-discovery tasks in GPT-2 plus transfer to Llama 3.2-1B and Llama 3.1-8B.

What This Paper Is About

Existing circuit-discovery methods mostly prune edges between coarse components (attention heads, MLP blocks), which is computationally expensive and cannot see structure below the component level. The authors instead learn trainable masks directly over nodes at several granularities at once — from whole transformer blocks down to single neurons — so the model can be compressed at multiple levels in one training run. The goal is to find smaller, more faithful circuits while using far less GPU memory.

Key Contributions

1. A multi-granular node-pruning framework for causal circuit discovery that learns masks at multiple levels of granularity — transformer blocks, attention heads, MLP blocks, attention neurons, and MLP neurons (first and second linear layers) — inside a single unified optimization objective, with granularity-specific sparsity penalties.

2. Demonstrated compactness gains with lower memory cost. The method prunes 33.34% more MLPs and 59.8% more neurons than prior methods in the least favorable setting, with larger gains elsewhere, while requiring at least 3× less memory and scaling to Llama 3.1-8B on a single 40GB GPU.

3. Evidence that coarse-grained methods retain unnecessary neurons. The authors show that many neurons kept by coarse-grained baselines can be removed with negligible impact on task performance.

4. A task-dependent topological analysis of discovered circuits, showing that different behaviors are encoded by distinct patterns of attention, MLP, and neuron retention.

Main Findings

• MLP and neuron pruning is where the gap is largest. EAP and EP retain all MLP blocks on every task, and DiscoGP achieves only limited compression on a single task. The authors' method uniquely achieves block-level MLP sparsity and prunes the vast majority of MLP neurons, while all baselines — including DiscoGP's unstructured weight pruning — induce effectively zero neuron-level sparsity.

• Attention-head sparsity is broadly comparable but highest for their method. On GPT-2 the method retains as few as 9 heads on the GT task, with reported attention-head pruning of 76.40% (IOI), 77.8% (GP), and 93.80% (GT).

• Fidelity is preserved or improved. Reported accuracy on GPT-2 is 0.958 (IOI), 0.984 (GP), and 0.992 (GT). On GT the KL divergence to the full model is 0.048 and logit difference is 0.371; on IOI KL is 0.466 with logit difference 5.012; on GP KL is 0.337 with logit difference 4.735.

• Circuits are task-dependent, not uniform. IOI is dominated by MLP-based computation, with several layers (e.g., 1, 6, 8, 11) retaining substantial MLPs while attention is entirely pruned. GP is highly localized and discontinuous: layers 1 and 2 are fully inactive and MLP computation is restricted to just three layers (0, 3, and 7). GT shows the most extreme sparsity, with a contiguous block of inactive layers (2–5), a large MLP in Layer 0, and computation resuming in Layers 8–11 — described as consistent with a skip-like mechanism.

• Different MLP sublayers compress differently. MLP1 (Keys) outputs can be sparsified to a greater extent than MLP2 (Values) outputs, consistent with Geva et al. (2020). This pattern holds in the reported GPT-2 numbers (M1 pruning exceeds M2 pruning across all three tasks).

• Scale does not require denser circuits. On Llama 3.2-1B and Llama 3.1-8B, GP is consistently more compressible than IOI, and the larger model generally supports equally strong or stronger structural pruning while preserving close agreement with the original: on IOI, MLP pruning rises from 25.0% (1B) to 28.1% (8B); on GP, from 37.5% (1B) to 40.6% (8B). Accuracy is 0.96 on IOI for both models and 0.99 on GP for both.

• Compute efficiency. On GPT-2 the method completes training in 350 seconds using 6,270 MB of peak memory, versus EAP at 21 seconds and 72,794 MB, edge pruning at 2,756 seconds and 33,354 MB, and DiscoGP at 2,926 seconds and 17,633 MB. The method adds only 55,465 trainable parameters to GPT-2 small, which has 124.5M parameters.

• Circuits are stable across seeds. Evaluating five seeds for the mask parameters on Llama 3.2-1B, fidelity metrics, per-layer sparsity profiles, and attention head selections remain highly consistent.

• Post-hoc edge pruning adds further compression. Applying edge sparsification on top of the discovered node circuits yields an additional 3.7 and 3.6 percentage points of edge compression on IOI and GP respectively, at a small additional memory cost.

• Qualitative analysis mirrors manual findings. Attention-score and Direct Logit Attribution analysis of the Llama 3.2-1B IOI circuit shows surviving late-layer heads implement distinct roles — IO-Movers, S-Inhibitors, and S-Promoters — mirroring roles identified in the manually discovered GPT-2 IOI circuit (Wang et al., 2022), while early layers contribute primarily through MLPs.

Methodology in Plain English

The core idea is to attach a trainable "volume knob" (a mask between 0 and 1) to every candidate node in the model — whole attention blocks, whole MLP blocks, individual attention heads, individual attention neurons, and individual MLP neurons in both the up and down projections. Turning a knob to 1 keeps the node's real activation; turning it to 0 replaces it with a corrupted activation. Because Knobs at coarser levels override finer ones (pruning an MLP block also disables its neurons), the search happens hierarchically.

Training runs the model twice per batch: once on a clean prompt that elicits the target behavior, and once on a minimally corrupted version of that prompt that changes the correct answer (for instance, swapping a starting year of 1950 for 1901 so the model no longer prefers numbers greater than 50). Corrupted patching is used rather than zero or mean ablation because the authors argue ablation creates distribution shift and can let the pruned model learn a new, compensatory solution instead of revealing the original mechanism.

The loss combines three terms: a faithfulness term (KL divergence between the full model's output distribution and the circuit's), a margin-based ranking task loss that keeps the circuit's preference for the ground-truth answer over the distractor, and a sparsity penalty on the masks. The sparsity penalty uses the L0 objective from Louizos et al. (2017), and masks are parameterized with the Hard Concrete distribution so they are differentiable during training but approximately binary, with standard settings ζ = 1.1, γ = −0.1, β = 0.66.

After training, masks are binarized with a threshold on the learned log α parameter, and hierarchical consistency is enforced so that deactivating a parent also zeroes its children. The framework introduces no extra forward passes beyond the two streams, and because it operates on nodes rather than edges — where the number of edges grows quadratically with the number of nodes — it avoids both the quadratic search space and the need to store intermediate activations.

Experiments cover three tasks on GPT-2: Indirect Object Identification (30 syntactic templates, 200 train and 200 validation examples, 1000 test instances), Gendered Pronouns (top 1,000 most popular male and female baby names from 2000; 150 train/validation examples each and 378 test examples), and Greater Than (the Hanna et al. 2023 dataset of 5 templates, 120 noun choices, years spanning 1100–2199, and 12,540 total samples). Training uses a batch size of 32, 500 epochs, max sequence length 64, and the Adam optimizer at learning rate 3e-2. Baselines are EAP, EP (edge pruning), and DiscoGP. Because edge-pruning methods do not natively report node-level results and this method does not natively report edge-level results, the authors apply strict proxy checks — counting a node as pruned only if it is an island with no incoming or outgoing edges, and reporting a conservative edge estimate that assumes all remaining nodes are fully connected.

Why This Matters

Impact on research. Circuit discovery has been bottlenecked by cost and coarseness: edge-based methods scale quadratically in the number of nodes and require storing internal representations. This work shows that a node-level, multi-granularity formulation can produce more compact circuits at a fraction of the memory, which changes what is practically feasible — it enables circuit analysis on 8B-parameter models on a single 40GB GPU. The finding that coarse-grained methods retain many neurons that can be removed with negligible effect is also a caution for interpreting component-level circuits as complete explanations. The authors note a limitation, however: the framework identifies which nodes are necessary but does not recover the interaction structure, directionality, or multiplicity of connections between them.

Real-world applications

• Model debugging and auditing. Identifying which subnetworks drive a behavior gives engineers a target for diagnosing failures such as factual errors, prompt sensitivity, or unwanted associations encoded from training data.

• Safety and guardrail analysis. Locating the internal pathways behind refusal or moderation behavior supports targeted oversight — though the authors explicitly flag the inverse risk, that targeted node removal may inadvertently disable moderation, refusal, or safety-related pathways and enable circumvention of guardrails.

• Efficient model compression research. Although the authors distinguish circuit discovery from deployment-oriented pruning, multi-granularity node masks that survive fine-tuning suggest a general recipe for structure-aware sparsification.

• Interpretability tooling for deployed LLMs. Reproducible, seed-stable circuits (verified across five seeds) that transfer to Llama models and mirror roles found by manual analysis make automated circuit discovery more usable as an analysis primitive in industry pipelines.

Industry relevance. The efficiency profile is the headline for practitioners: 350 seconds and 6,270 MB on GPT-2 for this method versus 2,756 seconds / 33,354 MB for edge pruning and 2,926 seconds / 17,633 MB for DiscoGP, plus scalability to Llama 3.1-8B on a single 40GB GPU. That combination of low memory and competitive wall-clock time is what makes routine, iterative circuit analysis inside a normal engineering workflow plausible rather than a research-only exercise. Code is released at https://github.com/MuhammadUmairHaider/Node-Pruning-for-Circuit-Discovery under a CC BY 4.0 license.

Future Directions

• Recovering interaction structure, not just node selection. The authors state as a limitation that binary masks show which nodes are necessary but not which active nodes exchange information, in which direction, or how many separate connections exist. They propose a hybrid node-and-edge procedure that infers interaction structure on top of node selections.

• Extending multi-granularity analysis to more models and tasks. The paper reports results on GPT-2, GPT-XL (Appendix P), Llama 3.2-1B, and Llama 3.1-8B, but larger models are evaluated only for the authors' own method because baselines are too expensive — leaving open whether the compactness advantage holds at greater scale.

• Understanding and tuning the sparsity–performance threshold. Task performance stays stable up to a model-dependent λ (0.975 for GPT-2, approximately 0.85 for Llama-3.2-1B) before degrading sharply, and the optimal value is fairly stable across tasks for a given model. What determines that threshold and whether it can be predicted rather than swept is not resolved.

• Safety implications of targeted node removal. The risk section raises an open question about which safety-relevant pathways are collocated with task circuits and whether node-level surgical edits could be used adversarially — an issue the paper flags but does not investigate.

Target Audience

This paper is most useful to mechanistic interpretability researchers working on circuit discovery and activation patching; NLP engineers who need to audit or debug deployed LLMs and are constrained by GPU memory; and graduate students entering the interpretability literature who want a concrete, reproducible method with released code. It is less suited to readers without a working understanding of transformer internals, since the granularity taxonomy (attention heads versus attention neurons versus MLP1/MLP2 neurons) is assumed rather than introduced from scratch.

Authors’ abstract

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.