Skip to content

Latest commit

 

History

3 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Marking is not using

A causal test for sparse-autoencoder features in a protein language model.

Code, results and manuscript for a study of what sparse-autoencoder (SAE) features in ESM-2 650M actually do, as opposed to what they correlate with.

An SAE feature is normally validated by showing it fires on the residues carrying a known concept — a signal peptide, a binding site — and reporting a detection score. That answers "does this feature mark the concept?" It does not answer "does the model use it?" This repository contains the pipeline that separates the two questions, and the evidence that they come apart.

Headline results

Question Answer
Are SAE features causally used by ESM-2? Yes, for signal peptide, transit peptide, propeptide and DNA-binding region — at two depths and under two SAE sparsity rules (95% CIs exclude zero).
Does the detection score predict causal importance? No. The propeptide feature has the lowest out-of-sample F1 yet among the largest causal effects; a high-scoring decoy is causally inert.
Do SAEs beat raw activations at decoding concepts? No. Under a matched linear probe the SAE significantly out-decodes raw neurons on 0 of 19 concepts at any layer.
Is best-single-feature scoring trustworthy? Not by default. Selecting the reported feature on the data it is scored on inflates it — about 4× more for a sparse dictionary than a dense baseline, and not because the dictionary is larger.

Detection score (x) versus causal effect (y) for the four ablated layer-16 features. Blue = causally used (interval excludes zero); red = the null control. The score does not order the causal effect.

What is here

src/            Libraries:  dataset (Swiss-Prot parsing), plm_adapter (ESM-2 activation
                extraction), topk_sae (TopK/BatchTopK SAE + training), evaluate (scoring)
run/            Pipeline stages, numbered in execution order (see docs/PIPELINE.md)
data/           Sampled corpus + residue label matrix (activation cache is regenerated)
results/        Every result JSON behind every number in the paper, + SHA-256 manifest
figures/        Paper figures (PNG + SVG) and the code that makes them
paper/          Manuscript source (LaTeX, elsarticle) and compiled PDF
docs/           Pipeline reference, reproduction guide, full results record

Quick start

git clone https://github.com/Biodyn-AI/marking-is-not-using.git
cd marking-is-not-using
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

Inspect the published results without recomputing anything (seconds — the result JSONs are in the repo):

make figures     # regenerates all three paper figures from results/
make summary     # prints the headline numbers straight from the result files

Reproduce from scratch (needs ~90 GB disk for the activation cache and a GPU or a patient CPU — see docs/REPRODUCING.md):

make all

The pipeline in one paragraph

Stage 01 runs 20,000 Swiss-Prot proteins through ESM-2 and caches the residual stream at six layers. Stage 02 turns the UniProt feature table into a per-residue label matrix. Stage 03 trains one TopK SAE per layer. Stages 04–05 score features against concepts and compare them with a raw-neuron baseline; stage 04c repeats that scoring with honest (out-of-sample) feature selection, which is what exposes the winner's curse. Stage 06 fits a whole-dictionary linear probe against the same neuron baseline. Stage 07 rules out a positional confound. Stage 08 is the causal test: it ablates one feature inside a forward hook and measures the KL shift in ESM-2's own masked-token predictions, on-span versus off-span versus matched random directions, bootstrapped over proteins. Stages 09–12 add confidence intervals, figures and the candidate-matched control. Full detail: docs/PIPELINE.md.

Data and model provenance

  • Model: facebook/esm2_t33_650M_UR50D, pinned to revision 08e4846e537177426273712802403f7ba8261b6c. The pin is in the code; do not unpin it if you want to reproduce the numbers.
  • Annotations: UniProt Swiss-Prot (CC BY 4.0). data/corpus/swissprot_20k.jsonl is the exact sampled subset used, and data/labels/residue_labels.npz the derived label matrix, both redistributed under that licence.
  • Trained SAEs: the seven SAE checkpoints (~350 MB) are attached to the repository's latest release rather than tracked in git.
  • Integrity: results/MANIFEST.sha256 records a SHA-256 for every ground-truth artefact.

A note on numerical reliability

All causal results are computed in CPU float32. GPU (Apple MPS) execution of this workload is not reliable under memory pressure — a command-buffer fault can silently perturb the result, and in our runs produced a false negative for one concept. On an unloaded machine CPU and GPU agree to the last digit. If you adapt this code for intervention experiments, keep the CPU reference.

Citation

@article{kendiukhov2026marking,
  title  = {Marking is not using: a causal test for sparse-autoencoder features
            in a protein language model},
  author = {Kendiukhov, Ihor},
  year   = {2026}
}

Licence

MIT for the code (see LICENSE). Swiss-Prot-derived data files remain under CC BY 4.0.

About

Causal test for sparse-autoencoder features in a protein language model (ESM-2 650M): feature ablation read out through the model's own masked-LM predictions, winner's-curse quantification, and an honest raw-neuron baseline.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages