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)

Cite this repo

@inproceedings{
  miller2024transformer,
  title={Transformer Circuit Evaluation Metrics Are Not Robust},
  author={Joseph Miller and Bilal Chughtai and William Saunders},
  booktitle={First Conference on Language Modeling},
  year={2024},
  url={https://openreview.net/forum?id=zSf8PJyQb2}
}