Sparse Weight Decomposition

Reference implementation for Sparse Weight Decomposition for Efficient Circuit Extraction. SWD rewrites a dense projection as two sparse factors whose bottleneck coordinates are directly intervenable circuit units.

This repository covers the experimental surfaces in the paper:

Experiment	Model / surface	Config
Single projection	GPT-2 Small layer 8 mlp.c_proj	configs/gpt2_single_projection.yaml
Model-size scaling	Qwen2.5 0.5B, 1.5B, 3B middle mlp.down_proj	configs/qwen25_*_single_projection.yaml
Full MLP	GPT-2 Small layer 8 mlp.c_fc + mlp.c_proj	configs/gpt2_full_mlp.yaml
Full transformer	all 48 GPT-2 Small body projections	configs/gpt2_full_model.yaml
Zero data	identity-Gram GPT-2 layer 8 mlp.c_proj	configs/zero_data.yaml
Dense controls	exact SVD and random orthogonal factorizations	configs/exact_dense_controls.yaml
Baselines	Transcoder, VPD-Recon-CI, Original VPD-KL, sparse pretraining	configs/baselines/

The runnable pipeline includes activation-Gram collection, DSF factorization, fixed-support recovery, CE/KL/reconstruction evaluation, task-margin attribution, train-only mean ablation, validation random controls, held-out sufficiency/necessity frontiers, exact dense controls, and unit inspection.

Install

Python 3.10+ and PyTorch 2.2+ are supported. A CUDA GPU is strongly recommended for paper-scale runs.

python -m venv .venv
source .venv/bin/activate
pip install -e '.[data,plot,dev]'

The upstream VPD implementation requires Python 3.13 and is intentionally kept in a separate environment; see docs/BASELINES.md.

Quick Start

Validate all release configs without downloading models or data:

swd validate --all

Prepare deterministic token blocks, then run one experiment stage at a time:

swd prepare-lm-data \
  --model gpt2 \
  --input '/path/to/fineweb-edu/*.parquet' \
  --output "$SWD_DATA_ROOT/gpt2/calibration.pt" \
  --seq-len 1024 \
  --num-sequences 24576

swd factorize --config configs/gpt2_single_projection.yaml
swd evaluate  --config configs/gpt2_single_projection.yaml
swd circuit   --config configs/gpt2_single_projection.yaml

swd run accepts a comma-separated stage list for an end-to-end run:

swd run --config configs/gpt2_single_projection.yaml \
  --stages factorize,evaluate,circuit

Use --smoke to cap data and task examples. It checks plumbing and artifact schemas; it is not a paper result.

Paper Results

The release contains compact, path-sanitized CSV snapshots for the paper plots. They let readers inspect and redraw results without the training checkpoints:

python scripts/plot_paper_results.py --output-dir figures/reproduced

See docs/EXPERIMENTS.md for the command and expected artifacts for every paper section, docs/DATA.md for exact splits and public task provenance, and docs/PAPER_RESULTS.md for the released table schemas.

Reproducibility Contract

• SWD token exposure means calibration tokens; trained baselines report optimizer-replay tokens.
• Circuit comparisons use one fixed checkpoint per method after matching replacement CE to within 0.001.
• circuit_train is used for attribution and ablation means, circuit_val selects the smallest statistically admissible prefix, and circuit_test is reporting-only.
• Circuit cost is reported both as selected bottleneck units and effective active read/write edges.
• Original VPD-KL is a replacement-quality reference, not a circuit baseline.

The current release deliberately contains no model checkpoints, FineWeb-Edu text, credentials, internal cluster launchers, or experiment-history dumps.

License

This project is licensed under the Apache License 2.0. The vendored Double Sparse Factorization solver is also Apache-2.0; its license and provenance are recorded under third_party/double_sparse/.