Sparse Attention Decomposition Applied to Circuit Tracing

Code to reproduce the paper "Sparse Attention Decomposition Applied to Circuit Tracing".

Requirements

conda env create -f environment.yml
conda activate sparse-attn-decomp

Reproduce paper figures

See paper.ipynb.

Tracing

Run the tracing and saving its data

python3 full_tracing_data_collection.py

This script produces the file results.nms-p256-f1.0-folded-expandedO-scaled.pkl, which is used for the next experiments and plots. It also produces results.nms-p256-f1.0-folded-expandedO-scaled-all-svs.pkl, which is the tracing using all SVs instead of a subset of them. This is used to produce the Figure 15 in the paper.

Build the graphs

python3 full_tracing_build_graph.py

This script produces the files nms-p256-f1.0-folded-expandedO-scaled.graphml, which is used for the next experiments and plots.

Interventions

Single-edge

python3 interventions_single_edge.py

This script produces the file interventions_single-edge.parquet, which is used for the plots.

Multi-edge

python3 interventions_multi_edge.py

This script produces the file interventions_multi-edge.parquet, which is used for the plots.