ktalreja/ProtoMechModels

This repository contains models from the paper "Protein Circuit Tracing via Cross-Layer Transcoders" by Darin Tsui, Kunal Talreja, Danial Saeedi, and Amirali Aghazadeh,to be presented at ICML 2026. You can find the preprint here.

ProtoMech Models

1. CLT_L{num_layers}_D{clt_dim}: CLT model with checkpoints. Can be loaded with CLT_L{num_layers}_D{clt_dim}/checkpoints/last.ckpt
2. PLT_L{num_layers}_D{plt_dim}: PLT model with checkpoints. Can be loaded with PLT_L{num_layers}_D{clt_dim}/checkpoints/last.ckpt
3. esm2_t6_8M_UR50D.pt: Base ESM model used for experiments (same for esm2_t12_35M_UR50D.pt)
4. eval_models: CNN evaluation models used for probe steering, one per DMS assay (same for eval_models_35M)
5. probe: probe evaluation models for function circuit, one per DMS assay and split configuration