Research map / Large language models / Language, safety and interpretability
Extracting Interpretable Task-Specific Circuits from Large Language Models for Faster Inference
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes an automatic, training-free method that prunes GPT-2 Small down to a task-specific circuit (attention heads and optionally MLPs) using mean/zero ablation and KL-divergence thresholds, yielding standalone smaller models for three circuit-studied tasks.
Problem
Circuits found by mechanistic interpretability are localized via activation patching, which only adds hooks and does not shrink or speed up the model; there is no way to extract a task-specific standalone submodel from an LLM.
Approach
Traverse attention heads (and optionally MLPs) from last layer to first, temporarily ablate each node on a small patching set, and permanently remove it if the KL-divergence increase on a validation set is below a threshold alpha, so the pruned model is physically smaller and faster rather than hook-patched.
Key results
- Mean-ablation pruning of GPT-2 Small on greater-than without MLP pruning-free setting keeps 100% accuracy with 32.65% parameter reduction; with MLP pruning it reaches 82.77% reduction at 99.84% accuracy.
- On IOI, pruning yields 28.66% parameter reduction at 100% accuracy (MLPs kept) and 72.31% reduction at 96.53% accuracy (MLPs pruned).
- Distillation baseline at matched size collapses to ~2-13% accuracy on acronym and IOI, versus 78.64-100% for the proposed method.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · Acronyms | accuracy: 99.92 \pm 0.17 -> 78.64 \pm 5.16 | Hurt |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · IOI | accuracy: 100.00 \pm 0.00 -> 96.53 \pm 1.52 | Hurt |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · Greater-than | accuracy: 100.00 \pm 0.00 -> 99.84 \pm 0.36 | No clear effect |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · Acronyms | parameter reduction: 32.88 \pm 0.00 -> 63.70 \pm 3.10 | Helped |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · IOI | parameter reduction: 28.66 \pm 0.50 -> 72.31 \pm 2.98 | Helped |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · Greater-than | parameter reduction: 32.65 \pm 0.00 -> 82.77 \pm 0.13 | Helped |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · Acronyms | inference time reduction: 83.56 \pm 0.53 -> 86.20 \pm 1.02 | Helped |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · IOI | inference time reduction: 57.58 \pm 8.76 -> 82.78 \pm 0.72 | Helped |
| Include MLPs in pruning (ablation) | Circuit extraction on GPT-2 Small · Greater-than | inference time reduction: 77.78 \pm 0.15 -> 88.33 \pm 1.00 | Helped |
| Use mean ablation instead of zero ablation (ablation) | Circuit extraction on GPT-2 Small · Acronyms | accuracy at comparable model size | Helped |
| Use mean ablation instead of zero ablation (ablation) | Circuit extraction on GPT-2 Small · IOI and Greater-than | resulting model size | Helped |
| Increase the pruning threshold (ablation) | Circuit extraction on GPT-2 Small · Acronyms | resulting model size | Hurt |
Surprising findings
- Mean-ablation pruning often preserves or slightly improves accuracy relative to the original GPT-2 Small.
- Zero ablation is less consistent and requires larger submodels than mean ablation for the same threshold.
Limitations
- Evaluated only on GPT-2 Small and three single-token-style tasks with manually identified circuits.
- Node-level pruning yields less precise circuits than edge-level methods; TPR against manual circuits is low (20-57%).
- Zero ablation gives noisy, less consistent results and larger submodels.
Experimental setup
Datasets: Acronym prediction prompts (Lee et al. style template) IOI templated sentences (one of 15 templates) Greater-than year prompts
Benchmarks: Acronym prediction Indirect Object Identification Greater-than
Metrics: accuracy parameter reduction (%) inference time (ms) TPR and FPR of recovered attention heads KL divergence
Baselines: knowledge distillation (logit matching) zero ablation
Builds on: activation patching automatic circuit discovery knowledge distillation
Closest papers by meaning
- Towards Global-level Mechanistic Interpretability: A Perspective of Modular Circuits of Large Language Models (ICML 2025 · similarity 0.92)
- Sparsity-Guided Holistic Explanation for LLMs with Interpretable Inference-Time Intervention (AAAI 2024 · similarity 0.88)
- Transcoders find interpretable LLM feature circuits (NeurIPS 2024 · similarity 0.88)
- Concept Bottleneck Large Language Models (ICLR 2025 · similarity 0.88)
- Towards Understanding Fine-Tuning Mechanisms of LLMs via Circuit Analysis (ICML 2025 · similarity 0.88)
- MIB: A Mechanistic Interpretability Benchmark (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.