Circuit Survival Under Quantization

Final research package for Quantization Can Degrade the IOI Circuit’s Behavior Without Reorganizing It: A Case Study (Ramon Vazquez, University of Illinois Chicago, August 31, 2026).

The project studies Indirect Object Identification (IOI) under fp16, int8, and 4-bit NF4 quantization in GPT-2 small, GPT-2 XL, and Mistral-7B. The final mechanistic analysis compares clean GPT-2 models with NF4-simulated copies because the bitsandbytes quantized layer classes are not directly hookable for the activation-patching sweep.

Final findings

• int8: no statistically detectable IOI degradation in any model.
• int4 / NF4: GPT-2 small loses 17.5% of its IOI logit-difference margin (n=500); GPT-2 XL loses 4.8% (n=500); Mistral-7B is statistically flat (n=200).
• Mechanism: the measured head-level causal structure is preserved under the NF4 simulation: Spearman ρ = 0.81 for GPT-2 small and 0.93 for GPT-2 XL; top-9 overlap is 8/9 and 9/9 respectively.
• Positive control: coarse per-channel RTN at 4 bits inverts the IOI behavior and reorganizes the head-level ranking (ρ = −0.074), showing that the sweep can detect reorganization.
• Evaluation pitfall: BOS handling shifts the GPT-2 IOI metric by about 18%, so all headline behavioral comparisons use the BOS frame.

The manuscript is included under paper/.

Reproducing the final pipeline

The released CSVs and figures are the authoritative final artifacts. The commands below describe how they were produced.

1. Generate the prompt data

python src/data_generation.py --task ioi --n 500 --out data/ioi_prompts.csv
python src/data_generation.py --task greater_than --n 500 --out data/gt_prompts.csv

2. Run the sanity checks

python src/sanity_checks.py

The pipeline uses self-patch, noise-patch, and known-circuit reproduction checks before trusting the mechanistic measurements.

3. Reproduce the final behavioral GPT-2 reruns

python src/bos_rerun_n500.py

This is the final headline behavioral path for GPT-2 small and GPT-2 XL: HuggingFace forward passes, a prepended BOS token, 500 prompts, and fp16/int8/NF4 int4 conditions. It writes results/bos_frame_results_n500.csv.

Mistral-7B uses its tokenizer's BOS behavior by default; its released final measurements are in results/mistral_final_results.csv.

4. Reproduce the mechanistic NF4 simulations

python src/mech_circuit_comparison.py
python src/mech_circuit_comparison_xl.py
python src/granular_analysis.py
python src/topk_sensitivity.py

These steps use blocksize-64 NF4 quantization/dequantization on TransformerLens model weights, validate the simulation against the measured behavioral effect, and then run exhaustive single-head activation patching over 25 ABC-swap prompt pairs.

• GPT-2 small: 144 heads; results/mech_per_head_comparison.csv
• GPT-2 XL: 1,200 heads; results/mech_xl_per_head_comparison.csv
• Granular preservation metrics: results/granular_summary.csv
• Top-k sensitivity: Appendix D; src/topk_sensitivity.py

5. Reproduce the positive control

python src/rtn_circuit_ablation.py

This applies coarse per-channel symmetric RTN at 4 bits to GPT-2 small and repeats the same head-patching analysis. Outputs are results/rtn_per_head_comparison.csv, results/rtn_scatter.pdf, and results/rtn_summary.txt.

Final result files

results/ contains the final, paper-facing artifacts:

• final_behavioral_results.csv — consolidated Table 1 data
• gpt2_final_results.csv
• gpt2-xl_final_results.csv
• mistral_final_results.csv
• bos_frame_results_n500.csv — original final GPT-2 n=500 table
• mech_per_head_comparison.csv
• mech_xl_per_head_comparison.csv
• granular_summary.csv
• rtn_per_head_comparison.csv
• final behavioral and mechanistic figures
• mechanistic summary text files

Historical exploratory behavioral CSVs and pre-final behavioral plots are preserved byte-for-byte in results/archive/ rather than mixed with the final outputs.

Important methodological note

The final mechanistic claim is not based on hooking bitsandbytes quantized layers directly. Instead, the project uses an NF4 simulation with blocksize 64 and a validation gate requiring the simulated model to reproduce the observed behavioral degradation closely enough before drawing mechanistic conclusions.

This distinction is also stated in the manuscript's Methods and Limitations. The simulation omits bitsandbytes' double quantization of scales, so the mechanistic conclusion is appropriately limited to an NF4-faithful approximation rather than the exact deployed kernels.

Legacy / exploratory code

src/run_experiment.py and the original results/archive/*_results.csv files are retained for provenance and historical reproducibility. They are not the source of the final n=500 headline numbers. In particular, the old generic runner uses a 200-example default and predates the final BOS-framed rerun.

Environment

The final experiments used:

• NVIDIA RTX 5080, 16 GB VRAM (local runs)
• NVIDIA T4, 15 GB VRAM (initial GPT-2 small run)
• PyTorch 2.11
• bitsandbytes
• transformers
• TransformerLens
• pinned environments in requirements-local.txt and requirements-colab.txt

No model weights or large activation caches are committed.

Repository structure

src/
  data_generation.py
  metrics.py
  circuit_analysis.py
  quantization.py
  sanity_checks.py
  run_experiment.py             # legacy/original runner
  bos_rerun_n500.py             # final GPT-2 behavioral rerun
  mech_circuit_comparison.py    # final GPT-2 small NF4 mechanism
  mech_circuit_comparison_xl.py # final GPT-2 XL NF4 mechanism
  granular_analysis.py
  topk_sensitivity.py
  rtn_circuit_ablation.py        # positive control
  make_plots.py

data/
  ioi_prompts.csv
  gt_prompts.csv

results/
  final_behavioral_results.csv
  *_final_results.csv
  mechanistic CSVs / PDFs / summaries
  archive/                       # historical exploratory artifacts

paper/
  Vazquez_Quantization_Degrades_IOI_Circuit.pdf

provenance/
  execution traces, environment snapshots, timestamps, hashes, and notebook

Reproducibility and provenance

PROVENANCE.md documents the human execution trail and the captured environment. The historical SHA-256 manifest is preserved as a provenance snapshot; archived exploratory artifacts retain their original bytes.

The manuscript source (.tex) is not included because no LaTeX source file was present in the supplied research package. The final manuscript PDF is included for a self-contained release.