Causal Circuit-Guided Pruning (CC-Prune)

This repository contains the official implementation and reproduction scripts for "Causal Circuit-Guided Pruning: A Principled Framework for Functionality-Preserving Compression of Transformers".

We introduce CC-Prune, a method that identifies and protects causally important computational circuits via activation patching, allowing for high-sparsity model compression without catastrophic performance degradation. This framework is compared against state-of-the-art baselines such as Wanda.

📂 Directory Structure

.
├── Dockerfile_CWP               # Docker configuration for the experimental environment
├── casual_pruning_execution.sh  # Main execution script (Builds Docker & runs experiments)
├── casual_active_scoring.py     # [Step 1] Calculates causal importance via Activation Patching
├── casual_active_pruning.py     # [Step 2] Performs CC-Prune using scores from Step 1
├── casual_active_visualize.py   # [Step 3] Generates heatmaps and score plots
├── casual_wanda_scoring.py      # [Baseline Step 1] Calculates Wanda importance scores
├── casual_wanda_pruning.py      # [Baseline Step 2] Performs Wanda pruning using scores
├── results/                     # Directory for output CSVs and PNGs (Mounted to container)
└── README.md                    # This file

🛠️ Prerequisites

To replicate the experiments, you need a machine with NVIDIA GPUs and the following software installed:

• Docker Engine
• NVIDIA Container Toolkit (required for --gpus all)

🚀 Quick Start

The entire pipeline (environment build, scoring, pruning, and visualization) is automated via the shell script.

1. Clone the repository:

   git clone <your-repo-url>
   cd <repo-name>
   

2. Run the reproduction script:

   bash casual_pruning_execution.sh
   

What the script does:

1. Builds the Docker image causal-pruning using Dockerfile_CWP.
2. Phase 1 (CC-Prune):
   • Runs casual_active_scoring.py to generate causal importance scores.
   • Runs casual_active_pruning.py to prune the model while protecting the circuit.
   • Runs casual_active_visualize.py to generate figures.
3. Phase 2 (Wanda Baseline):
   • Runs casual_wanda_scoring.py to calculate Wanda metrics.
   • Runs casual_wanda_pruning.py for comparison.
4. Automatically cleans up containers and images after execution.

Note: By default, the script uses GPU 0 (CUDA_VISIBLE_DEVICES="0"). Modify casual_pruning_execution.sh if you need to use different GPUs.

🔬 Experiment Pipeline

The experiments consist of sequential stages. Scoring must be executed before Pruning for both methods.

1. Causal Circuit-Guided Pruning (CC-Prune)

• Step 1: Scoring
  • Script: casual_active_scoring.py
  • Description: Calculates the causal importance of each neuron/head using Activation Patching on the JSTS dataset. Saves scores to causal_scores_*.csv.
• Step 2: Pruning
  • Script: casual_active_pruning.py
  • Description: Loads the CSV from Step 1, identifies the "Causal Circuit" (top $k$% important components), and prunes the rest using magnitude pruning.
• Step 3: Visualization
  • Script: casual_active_visualize.py
  • Description: Generates heatmaps (e.g., protection_L3_H5.png) proving that specific weights are protected.

2. Baseline Comparison (Wanda)

• Step 1: Scoring
  • Script: casual_wanda_scoring.py
  • Description: Calculates importance metrics based on Weights $\times$ Activation Norms.
• Step 2: Pruning
  • Script: casual_wanda_pruning.py
  • Description: Loads the Wanda scores and performs pruning to provide a baseline for performance comparison.

📊 Results

All experimental results, including CSV logs and generated figures, will be saved in the ./results directory.

Key outputs include:

• causal_scores_*.csv: Raw importance scores for all components (Generated by scoring step).
• hybrid_pruning_results_*.csv: Performance metrics of CC-Prune at various sparsity levels.
• wanda_pruning_results_*.csv: Performance metrics of Wanda pruning.
• protection_L*_H*.png: Proof of circuit protection mechanism.

---

Note: The filenames in this repository use "casual" for historical consistency, but the methodology refers to "Causal" (Cause-and-Effect).