AutoCircuit

A library for efficient patching and automatic circuit discovery.

Read the paper

Transformer Circuit Metrics are not Robust (Oral spotlight, COLM 2024)

Getting Started

pip install auto-circuit

Easy and Efficient Edge Patching

auto-circuit/experiments/demos/zero_ablate_an_edge.py

Lines 20 to 26 in 03ce552

	patch_edges = [
	"Resid Start->MLP 2",
	"MLP 2->A2.4.Q",
	"A2.4->Resid End",
	]
	with patch_mode(model, ablations, patch_edges):
	patched_out = model(tokens)

Different Ablation Methods

auto-circuit/experiments/demos/patch_an_edge.py

Line 37 in 03ce552

ablations = src_ablations(model, test_loader, AblationType.TOKENWISE_MEAN_CORRUPT)

Automatic Circuit Discovery

auto-circuit/experiments/demos/patch_an_edge.py

Lines 49 to 55 in 03ce552

	attrution_patching_scores: PruneScores = mask_gradient_prune_scores(
	model=model,
	dataloader=test_loader,
	official_edges=None,
	grad_function="logit",
	answer_function="avg_diff",
	)

Visualization

auto-circuit/experiments/demos/patch_some_edges.py

Line 49 in 5777057

fig = draw_seq_graph(model, prune_scores)