Antibody Crosscoders — model-diffing dictionaries for ProtBERT → IgBERT

Sparse crosscoders (shared encoder, per-model decoders) trained to diff the activations of three frozen protein/antibody language models across two fine-tuning stages:

ProtBERTIgBERT-unpairedIgBERT (paired) (all 30-layer / 1024-hidden BERT encoders).

These are the trained crosscoder checkpoints from the study. Activation stores (.h5, ~5.8 TB) are intentionally not included — they are regenerable from the companion dataset repo plus the extraction scripts; see REPRODUCE.md in the code repo (https://github.com/WionaGlaenzer/ab_interp).

Checkpoint format

Each .pt is a PyTorch state_dict for a BatchTopKCrossCoder from the jkminder/dictionary_learning fork (commit d1fd6f3). Load with:

import sys; sys.path.insert(0, "dictionary_learning")  # jkminder fork
from dictionary_learning import BatchTopKCrossCoder
ae = BatchTopKCrossCoder.from_pretrained("cc_paired_small_s0_L20.pt",
                                         code_normalization="CROSSCODER")

Main config: dict 32768 (expansion 32), k 32, num_layers 3, layer 20, target_rms 1.0. Model axis order: [prot_bert, igbert_unpaired, igbert_paired] (2-model / delta checkpoints use [unpaired, paired]).

What's here (naming → role)

pattern role paper
cc_{unpaired,paired,mixed}_small_s{0,1,2}_L20.pt main 3-model, size-matched, 3 seeds Fig 2, Fig 3, Tables S2/S3
cc_{paired,single,mixed}_L20.pt full-data 3-model Fig S5/S6/S7 (V/J-gene, conserved residues)
cc_{unpaired,paired,mixed}_velayer_L{4,12,28}.pt layer sweep Fig 1c, Fig S1, Fig S3
cc2_paired_alllayer_L{1..30}.pt 2-model, all 30 layers Fig S4
cc2_paired_rebuild95_L20.pt 2-model (IgBERT vs unpaired) Fig 2d
delta_{paired_rebuild95,deltaCC_unpVpair}_L20.pt Delta-Crosscoder Fig 2d, failure analysis
cc_{unpaired,paired,mixed}_scale_L20.pt, cc_paired_scale{500k,FULL}_L20.pt 10× scale-up Table S1
cc_single{4k,8k,16k}_L20.pt expansion-factor sweep Fig S2a/b
cc_singlek{8,16,64,128}_L20.pt BatchTopK-K sweep Fig S2c/d
sae_{prot,unp,pair}_L20*.pt per-model single-SAE baselines

Exploratory runs (e.g. cc_mixed64x, cc_mixed96k, cc_paired_ep*) are included for completeness but are not reported figures.

Full figure→checkpoint→dataset provenance: REPRODUCE.md in the code repo.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support