This is the official code repository for the paper "Circuit Tracing in Autoregressive Protein Language Models", by Darin Tsui, William Deinzer, Daniel Saeedi, and Amirali Aghazadeh, accepted into the Mechanistic Interpretability Workshop at ICML 2026. A link to the paper can be found here.
Additionally, one can explore protein circuits through our web-based visualizer!
The easiest way to get started with ProtoMech is through our interactive Google Colab notebook. No local installation is required.
- Models: ProGenMech currently supports ProGen3-Small!
- Circuit Discovery (optional): Identify circuits in either generation or zero-shot.
- Interactive Visualization: Generate files required for our website and visualize circuits!
circuit_utils/— Core CLT/PLT circuit utility code.function_circuit/— Function circuit discovery and circuit JSON outputs.generation_circuit/— Generation tasks and discovery workflows.training/— Training code for the CLT model.training_transcoder/— Training code for the PLT/transcoder model.visualization/— Analysis and visualization scripts for circuit outputs.steering/— Steering and replacement model experiments.data/— Input DMS and sequence data.external/— Vendored external ProGen3 model source.models/— Pretrained and checkpointed model directories.
- Run discovery:
sh circuit_discovery.sh - Run generation analysis:
sh generation_circuit/main.sh - Run steering:
sh steering/main_steering.sh - Use visualization scripts from
visualization/, e.g.python visualization/get_edge_weights.py
