🤖 AI

Multi-Granular Node Pruning for Causal Circuit Discovery

This paper proposes a scalable, multi-granular node pruning framework that uses learnable masks and sparsity penalties to efficiently discover smaller, more precise causal circuits in large language models by identifying relevant individual neurons while significantly reducing memory requirements compared to existing edge-pruning methods.

Original authors: Muhammad Umair Haider, Hammad Rizwan, Hassan Sajjad, A. B. Siddique

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Muhammad Umair Haider, Hammad Rizwan, Hassan Sajjad, A. B. Siddique

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Large Language Model (like the ones that write stories or answer questions) as a massive, bustling city. This city has millions of buildings (neurons), roads (connections), and districts (layers). When the model performs a specific task—like figuring out who "Kristi" gave a mango to in a story—it doesn't use the whole city. It only uses a tiny, specific neighborhood.

Finding that specific neighborhood is called Circuit Discovery. The goal is to map out exactly which parts of the city are doing the work and which parts are just standing there doing nothing.

The Problem with Old Maps

Previously, researchers tried to find these neighborhoods by looking at the roads (connections) between buildings. They would try to close off roads to see if the city still worked.

The Flaw: This is like trying to find a specific house in a city by blocking every single street. It takes forever, requires a massive amount of memory (like needing a map of every single street in the world), and it's too blunt. If you block a road, you might accidentally cut off a whole block of houses, even if only one house on that block was needed.
The Result: The old maps were "coarse." They could tell you that a whole district (like an "Attention Head") was important, but they couldn't tell you that only one specific room inside a building in that district was actually doing the work.

The New Solution: Multi-Granular Node Pruning

The authors of this paper propose a smarter way to map the city. Instead of looking at roads, they look directly at the buildings (nodes) themselves, and they do it at different levels of detail all at once.

Think of it like a set of Russian nesting dolls or a zoom lens:

The Big Picture: They can turn off entire districts (Transformer Blocks).
The Neighborhood: They can turn off specific streets within a district (Attention Heads).
The House: They can turn off specific rooms inside a house (Individual Neurons).

They use a special "learnable mask" (a digital switch) for every single building, from the biggest district down to the smallest room. They run a simulation where they mix a "clean" version of the city with a "corrupted" version (where the story makes no sense). By watching which buildings must stay on to keep the story correct, they figure out exactly what is essential.

What They Found (The Results)

When they tested this on different "cities" (models like GPT-2 and Llama), they found some surprising things:

Huge Savings: Their method found circuits that were much smaller than anyone else's. In the worst-case scenario, they removed 33% more building blocks and 60% more individual rooms than previous methods.
Memory Efficiency: Because they don't need to store every single street map (intermediate activations), their method uses 3 to 11 times less computer memory. It's like needing a small notebook instead of a library to draw the map.
Task-Specific Blueprints: They discovered that different tasks use different city layouts:
- Indirect Object Identification (Who did what to whom): This task relies heavily on the "MLP" buildings (the non-linear processing rooms) throughout the whole city, while the "Attention" streets are mostly empty.
- Gendered Pronouns (He vs. She): This task uses a very sparse, scattered network. Most of the city is turned off; only a few specific layers and rooms are active.
- Greater Than (Math/Numbers): This is the most extreme. It uses a "skip" mechanism, where the city ignores huge chunks of the middle layers and jumps straight to the end to do the calculation.

The Takeaway

The paper claims that by looking at the model at the level of individual "rooms" (neurons) rather than just whole "districts" (blocks), we can strip away a massive amount of unnecessary machinery.

They proved that many parts of these AI models that we thought were necessary are actually just dead weight. By pruning them down to the finest grain, they found the true, minimal "circuit" that drives the behavior, doing it faster and with much less computer power than before. They also showed that this works even on very large models (like Llama 3.1-8B) on a single computer card, which was previously too difficult for other methods to handle.

Technical Summary: Multi-Granular Node Pruning for Causal Circuit Discovery

1. Problem Statement

Causal circuit discovery aims to identify minimal subnetworks within Large Language Models (LLMs) that causally drive specific behaviors. While existing methods have advanced mechanistic interpretability, they face two primary limitations:

Computational Scalability: Current approaches predominantly rely on edge pruning (removing connections between components). Since the number of edges in a transformer grows quadratically with the number of nodes ( $O(|N|^2)$ ), edge-level pruning becomes computationally intractable as model size increases. Furthermore, these methods often require storing disentangled residual-stream representations and intermediate activations, leading to prohibitive memory usage.
Granularity Limitations: Existing methods typically operate at a coarse granularity, treating entire attention heads or MLP blocks as atomic units. This overlooks the possibility that specific neurons (nodes) within these components are responsible for target behaviors, resulting in recovered circuits that miss fine-grained structures and retain unnecessary components.

2. Methodology

The authors propose a multi-granular node-pruning framework that discovers minimal circuits across multiple levels of granularity within a single unified optimization step.

2.1 Hierarchical Granularity Design

Instead of pruning edges between coarse components, the framework learns masks over nodes at four distinct levels:

Block Level: Entire Attention Blocks and MLP Blocks.
Head Level: Individual Attention Heads.
Neuron Level: Individual neurons within Attention layers and the two linear layers of MLP blocks (MLP1/Keys and MLP2/Values).

2.2 Two-Stream Causal Pruning Setup

To identify causal importance without introducing distribution shifts (a risk with simple zero-ablation), the method employs a two-stream forward pass:

Clean Stream: The model processes an input designed to elicit a specific behavior (e.g., predicting a year > 50).
Corrupted Stream: The model processes a minimally perturbed version of the input (e.g., inverting gender pronouns or changing the starting year) to elicit a different prediction.
Interpolation: For each node $n_i$ , a learnable mask $m_i \in [0, 1]$ interpolates between the clean activation ( $h^{clean}$ ) and the corrupted activation ( $h^{corrupted}$ ):
$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$
The optimization seeks to find masks that allow the model to produce the correct output (matching the clean behavior) while maximizing sparsity.

2.3 Optimization and Mask Parameterization

Hard Concrete Distribution: Masks are parameterized using the Hard Concrete distribution to enable gradient-based optimization while inducing approximately binary behavior (0 or 1).
Unified Objective: The training objective combines three components:
1. Faithfulness Loss ( $L_{faith}$ ): A KL divergence ensuring the pruned circuit's output distribution matches the base model.
2. Task Loss ( $L_{task}$ ): A margin-based ranking loss ensuring the circuit correctly distinguishes between ground-truth and distractor options.
3. Sparsity Regularizer ( $L_0$ ): Encourages the masks to be zero, minimizing the circuit size.
Hierarchical Consistency: Post-training, masks are binarized with a constraint ensuring that if a parent node (e.g., an MLP block) is pruned, all its child nodes (e.g., neurons within that block) are also deactivated.

3. Key Contributions

Multi-Granular Framework: The introduction of a unified optimization framework that learns masks simultaneously across transformer blocks, attention heads, MLP blocks, and individual neurons.
Scalability and Efficiency: The method avoids the quadratic scaling of edge pruning and the need to store intermediate activations for all edges. It reduces GPU memory usage by 3× to 11× compared to prior methods, enabling circuit discovery on Llama 3.1-8B using a single 40GB GPU.
Fine-Grained Discovery: The approach identifies significantly more compact circuits than baselines, specifically achieving high sparsity at the MLP block and neuron levels, which previous methods largely failed to achieve.

4. Experimental Results

The method was evaluated on three standard circuit discovery tasks (Indirect Object Identification, Gendered Pronouns, and Greater Than) using GPT-2, GPT-2-XL, and Llama models.

4.1 Sparsity and Compactness

MLP and Neuron Pruning: The proposed method outperforms edge-pruning baselines (EAP, EP) and unstructured weight pruning (DiscoGP) significantly. In the least favorable setting, it pruned 33.34% more MLP blocks and 59.8% more neurons than prior methods.
Task-Dependent Structures: The discovered circuits exhibit distinct topologies based on the task:
- IOI: Dominated by distributed MLP computation; attention is largely pruned except in final layers.
- Gendered Pronouns (GP): Highly localized and discontinuous; MLP computation is restricted to a few layers, while attention is sparse.
- Greater Than (GT): Exhibits extreme sparsity with contiguous blocks of inactive layers and "skip-like" execution patterns.

4.2 Fidelity and Consistency

Performance: The pruned circuits maintain high task accuracy (often >95%) and low KL divergence relative to the full model.
Consistency: Experiments across five random seeds show that the discovered circuits (fidelity metrics and sparsity profiles) are highly stable.
Qualitative Analysis: In Llama 3.2-1B, the surviving attention heads in the IOI task were found to perform specific functional roles (IO-Movers, S-Inhibitors) mirroring manually discovered circuits in GPT-2, validating the method's ability to recover mechanistic insights.

4.3 Compute Efficiency

Memory: The method requires at least 3× less memory than edge-pruning baselines (e.g., 6.27 GB vs. 72.79 GB for EAP on GPT-2).
Time: Training completes in under 6 minutes, significantly faster than edge-pruning methods (approx. 50 minutes), though slightly slower than the single-pass EAP method (which trades time for massive memory usage).

5. Significance and Claims

The paper claims that its primary significance lies in overcoming the scalability and granularity bottlenecks of current circuit discovery. By shifting from edge pruning to multi-granular node pruning, the authors demonstrate that:

Many neurons retained by coarse-grained methods are unnecessary and can be removed with negligible impact on performance.
Larger models (e.g., Llama 8B) support equally strong or stronger pruning than smaller models, suggesting that increased scale introduces redundancy rather than requiring denser task-specific circuits.
The method provides a computationally feasible path to discovering fine-grained causal circuits in large-scale models, revealing that different behaviors are encoded by distinct patterns of attention, MLP, and neuron activity.

The authors acknowledge a limitation: the framework identifies which nodes are necessary but does not explicitly recover the interaction structure (edges) between the surviving nodes. They suggest a future direction of combining node selection with edge recovery to infer interaction structures.