PLM Reverse Distillation

Protein language models (PLMs) scale poorly: for many tasks, mid-sized models often outperform the largest in the same family. Reverse Distillation addresses this by decomposing large PLM representations into orthogonal subspaces guided by smaller models of the same family. The resulting embeddings have a Matryoshka-style nested structure — the first k dimensions of a larger model's embedding exactly match the smaller model's representation — ensuring larger reverse-distilled models consistently outperform smaller ones.

On ProteinGym benchmarks, reverse-distilled ESM-2 variants outperform their respective baselines at the same embedding dimensionality, with the reverse-distilled 15B model achieving the strongest performance.

Installation

Requires Python ≥ 3.12 and uv.

git clone https://github.com/rohitsinghlab/plm_reverse_distillation.git
cd plm_reverse_distillation
uv lock && uv sync
uv pip install -e '.[dev]'
source .venv/bin/activate

Quick Start

See inference_tutorial.ipynb for a step-by-step walkthrough of loading pretrained models and extracting embeddings.

Pretrained scalers for all ESM-2 model pairs (8M → 35M → 150M → 650M → 3B → 15B) are available on HuggingFace and loaded automatically via the model registry:

singhlab/plm_reverse_distillation

Available Models

All models use PCR regression and PCA for dimensionality reduction. Each model applies the full chain of scalers from ESM-2 8M up to the target size.

Scripts

Embedding extraction

Extract embeddings from a FASTA file using a pretrained RD model:

python scripts/extract.py \
    --fasta_file proteins.fasta \
    --output_dir embeddings/ \
    --repr_type mean \
    --batch_size 32

Key arguments: --repr_type (per_tok / mean / bos), --repr_layers, --batch_size, --truncation_seq_length.

Training scalers

Train new scalers on your own data:

python scripts/train.py \
    --dataset_path proteins.fasta \
    --scalar_path scalers/ \
    --regressor_type pcr \
    --scaler_type rd \
    --n_pretrained_seqs 5000

Key arguments: --regressor_type (linear / ridge / pcr), --scaler_type (rd / naive), --pca_type (incremental / fbpca), --n_pretrained_seqs.

Citation

If you use reverse distillation, please cite:

@inproceedings{catrina2026reverse,
  title   = {Reverse Distillation: Consistently Scaling Protein Language Model Representations},
  author  = {Catrina, Darius and Bepler, Christian and Sledzieski, Samuel and Singh, Rohit},
  booktitle = {International Conference on Learning Representations},
  year    = {2026}
}

License

This project is licensed under the MIT License — see LICENSE for details.