First commit.

Author: danfriedman0

README.md

@@ -0,0 +1,57 @@
+# Extracting rule-based descriptions of attention features in transformers
+
+This repository contains the code for our paper, "Extracting rule-based descriptions of attention features in transformers".
+Please see our paper for more details.
+
+## Quick links
+* [Setup](#Setup)
+* [Attention output SAEs](#Attention-output-SAEs)
+* [Data](#Data)
+* [Rule extraction](#Rule-extraction)
+* [Questions?](#Questions)
+* [Citation](#Citation)
+
+## Setup
+
+Install [PyTorch](https://pytorch.org/get-started/locally/) and then install the remaining requirements: `pip install -r requirements.txt`.
+This code was tested using Python 3.12 and PyTorch version 2.3.1.
+
+## Attention output SAEs
+
+We train [attention output SAEs](https://arxiv.org/abs/2406.17759) for every attention head in GPT-2 small, using a fork of https://github.com/ckkissane/attention-output-saes.
+These SAEs can be downloaded from: https://huggingface.co/danf0/attention-head-saes/.
+
+## Data
+
+Code for generating datasets of feature activations can be found in [src/get_exemplars.py](src/get_exemplars.py).
+See [scripts/generate_data.sh](scripts/generate_data.sh) for the command to generate the datasets used in our paper, which are based on [OpenWebText](https://huggingface.co/datasets/Skylion007/openwebtext).
+
+## Rule extraction
+
+Code for extracting and evaluating skip-gram rules can be found in [src/run_rules.py](src/run_rules.py).
+For example, the following command will extract rules for 10 features from head 0 in layer 0.
+```bash
+python src/run_rules.py \
+    --layer 0 \
+    --head 0 \
+    --num_features 10 \
+    --rule_type "v1" \
+    --output_dir "output/skipgrams/L0H0";
+```
+Code for finding and generating rules containing "distractor" features is in [src/find_distractors.py](src/find_distractors.py) and [src/generate_distractors.py](src/generate_distractors.py)
+The [scripts](scripts) directory contains example contains example commands for running these scripts.
+
+# Questions?
+
+If you have any questions about the code or paper, please email Dan ([email protected]) or open an issue.
+
+# Citation
+
+```bibtex
+@article{friedman2025extracting,
+    title={Extracting rule-based descriptions of attention features in transformers},
+    author={Friedman, Dan and Wettig, Alexander and Bhaskar, Adithya and Chen, Danqi},
+    journal={arXiv preprint},
+    year={2025}
+}
+```

requirements.txt

@@ -0,0 +1,9 @@
+einops
+matplotlib
+numpy
+pandas
+sae-lens
+scikit-learn
+seaborn
+transformers
+tqdm

scripts/extract_rules.sh

@@ -0,0 +1,14 @@
+#!/bin/bash 
+
+LAYER=0
+HEAD=0
+METHOD="magnitude"
+# METHOD="importance"
+
+srun python src/run_rules.py \
+  --layer "${LAYER}" \
+  --head "${HEAD}" \
+  --method "${METHOD}" \
+  --num_features 100 \
+  --example_dir "data/openwebtext_n50000_bins2x50/" \
+  --output_dir "output/skipgrams/${METHOD}/L${LAYER}H${HEAD}/";

scripts/find_distractor_examples.sh

@@ -0,0 +1,11 @@
+#!/bin/bash 
+
+LAYER=0
+HEAD=0
+
+python src/find_distractors.py \
+  --cmd "get_distractor_examples" \
+  --layer "${LAYER}" \
+  --head "${HEAD}" \
+  --example_dir "data/openwebtext_n50000_bins2x50/" \
+  --output_dir "output/distractor_examples/L${LAYER}H${SLURM_ARRAY_TASK_ID}/";

scripts/find_distractor_features.sh

@@ -0,0 +1,11 @@
+#!/bin/bash 
+
+LAYER=0
+HEAD=0
+
+srun python src/find_distractors.py \
+  --cmd "get_distractors" \
+  --layer "${LAYER}" \
+  --head "${HEAD}" \
+  --example_dir "data/openwebtext_n50000_bins2x50/" \
+  --output_dir "output/distractor_features/L${LAYER}H${HEAD}/";

scripts/generate_data.sh

@@ -0,0 +1,22 @@
+#!/bin/bash 
+
+LAYER=0
+HEAD=0
+DATASET_SIZE=50000
+NUM_BINS=2
+EXAMPLES_PER_BIN=50
+SEED=0
+
+python src/get_exemplars.py \
+  --layer "${LAYER}" \
+  --head "${HEAD}" \
+  --dataset_size "${DATASET_SIZE}" \
+  --num_bins "${NUM_BINS}" \
+  --examples_per_bin "${EXAMPLES_PER_BIN}" \
+  --max_length 64 \
+  --num_features 100 \
+  --min_count 150 \
+  --max_count 49850 \
+  --dataset_path "Skylion007/openwebtext" \
+  --seed "${SEED}" \
+  --output_dir "data/openwebtext_n${DATASET_SIZE}_bins${NUM_BINS}x${EXAMPLES_PER_BIN}/L${LAYER}H${HEAD}/";

scripts/generate_distractor_examples.sh

@@ -0,0 +1,9 @@
+#!/bin/bash 
+
+HEAD=0
+
+python src/generate_distractors.py \
+  --layer 0 \
+  --head "${HEAD}" \
+  --example_dir "output/distractor_examples/" \
+  --output_dir "output/generated_distractor_examples/L0H${HEAD}/";

setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup, find_packages
+
+setup(name="src", version="0.1", packages=find_packages())