Active Circuit Discovery

Active Inference-guided circuit discovery for mechanistic interpretability in Large Language Models.

Combines attribution graph analysis with a POMDP agent (powered by pymdp) to efficiently identify causally important transcoder features in LLMs, using Anthropic's circuit-tracer library. Evaluated on Gemma-2-2B and Llama-3.2-1B.

Key Results

Benchmark	Model	POMDP vs Random	Oracle Efficiency
IOI (5 prompts)	Gemma-2-2B	+16.3%	58.3%
IOI (5 prompts)	Llama-3.2-1B	-25.4%	37.5%
Multi-step (3 prompts)	Gemma-2-2B	+30.4%	73.3%
Multi-step (3 prompts)	Llama-3.2-1B	-88.8%	6.5%
Feature Steering	Gemma-2-2B	8/50 prediction changes at 10x	---
Feature Steering	Llama-3.2-1B	9/50 prediction changes at 10x	---

All results from real feature_intervention API calls with GemmaScope / Llama transcoders. No synthetic or fabricated data.

Architecture

Prompt → Gemma-2-2B / Llama-3.2-1B
           ↓
   circuit-tracer (EAP + Transcoders)
           ↓
   Attribution Graph (active features, adjacency matrix)
           ↓
   Active Inference POMDP Agent (pymdp)
   ├─ Variational state inference
   ├─ Expected Free Energy scoring
   └─ Dirichlet observation-model learning
           ↓
   feature_intervention(layer, pos, fidx, value)
           ↓
   KL Divergence Measurement → Update Beliefs

Hypotheses and Outcomes

ID	Hypothesis	Criterion	Outcome
H1	POMDP agent more efficient than random	Paired t-test, α=0.05	Positive trend on Gemma; not significant (p=0.27)
H2	Features causally control predictions	Binomial test > 1% baseline	Accepted (17/100, p < 10⁻¹⁵)
H3	Oracle efficiency ≥ 50%	Point estimate	Accepted for Gemma (58.3%); not for Llama (37.5%)
H4	Findings transfer across architectures	Qualitative replication	Accepted (all experiments replicated on Llama)

Installation

Local (DGX Spark / GPU machine)

git clone https://github.com/SharathSPhD/ActiveCIrcuitDiscovery.git
cd ActiveCIrcuitDiscovery
python -m venv .venv
source .venv/bin/activate

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt
pip install git+https://github.com/safety-research/circuit-tracer.git

Docker (DGX Spark / aarch64)

docker compose up active-circuit-discovery

Docker (standard x86_64 / amd64, T4-compatible)

A portable image is provided for any x86_64 host with an NVIDIA GPU of compute capability >= 7.5 (including a single 16 GB T4):

# build
docker build -f Dockerfile.amd64 -t acd:amd64 .

# run one experiment (results land in ./results)
docker compose -f docker-compose.amd64.yml run --rm acd \
    python -m src.experiments.run_real_experiments --model gemma --experiment ioi

The model loads in float32 and uses ~5 GB VRAM, so the full pipeline fits on a T4. Absolute KL values may shift slightly across GPU generations and driver/cuDNN versions, but the relative method orderings are hardware-independent.

Google Colab (free T4)

Click any Colab badge above -- dependencies install automatically and the notebooks reproduce the main experiments on a free T4 GPU.

Quick Start

import torch
from circuit_tracer import ReplacementModel, attribute
from circuit_tracer.graph import prune_graph

model = ReplacementModel.from_pretrained(
    model_name="google/gemma-2-2b",   # or "unsloth/Llama-3.2-1B"
    transcoder_set="gemma",            # or "llama"
    backend="transformerlens",
    device=torch.device("cuda"),
    dtype=torch.float32,
)

prompt = "When John and Mary went to the store, John gave the bag to"
raw_graph = attribute(prompt=prompt, model=model, max_n_logits=5,
                      desired_logit_prob=0.9, batch_size=256)

# Ablate a feature and measure causal effect
clean_logits, _ = model.feature_intervention(prompt, [])
clean_probs = torch.softmax(clean_logits[0, -1, :], -1)

iv_logits, _ = model.feature_intervention(
    prompt, [(25, 14, 4717, 0)]  # ablate L25 feature 4717
)
iv_probs = torch.softmax(iv_logits[0, -1, :], -1)
kl = torch.nn.functional.kl_div(
    torch.log(iv_probs + 1e-10), clean_probs, reduction='sum'
).item()
print(f"KL divergence from ablation: {kl:.6f}")

Running Experiments

# Full experiment suite on both models
python -m src.experiments.run_real_experiments --model both --experiment all

# Single model / single benchmark
python -m src.experiments.run_real_experiments --model gemma --experiment ioi
python -m src.experiments.run_real_experiments --model llama --experiment steering

# Results saved to results/ as JSON
ls results/*.json

Project Structure

ActiveCIrcuitDiscovery/
├── src/
│   ├── experiments/
│   │   └── run_real_experiments.py    # Main experiment runner
│   ├── circuit_analysis/
│   │   └── circuit_tracer_backend.py  # circuit-tracer wrapper
│   ├── active_inference/
│   │   └── pomdp_agent.py            # POMDP agent (pymdp)
│   ├── core/
│   │   ├── data_structures.py        # Core data types
│   │   └── metrics.py                # Statistical metrics
│   └── visualization/
│       ├── visualizer.py             # Result plotting
│       └── research_dashboard.py     # Interactive dashboard
├── scripts/
│   └── generate_figure_data.py       # LaTeX figure generation
├── notebooks/
│   ├── 01_circuit_discovery_gemma.ipynb   # Circuit discovery demo
│   ├── 02_active_inference_demo.ipynb     # POMDP agent comparison
│   └── 03_reproduce_biology_paper.ipynb   # Feature steering
├── paper/                                 # LaTeX paper source
├── results/                               # Experiment JSON outputs
├── Dockerfile.dgx-spark                   # aarch64 / DGX Spark image
├── Dockerfile.amd64                       # x86_64 / T4-compatible image
├── docker-compose.yml                     # aarch64 compose target
├── docker-compose.amd64.yml               # x86_64 compose target
└── requirements.txt

Hardware Requirements

Component	Minimum	Recommended
GPU VRAM	8 GB (T4)	128 GB (DGX Spark)
RAM	16 GB	64 GB
Storage	20 GB	50 GB

License

MIT License