English

Reverse Distillation: Consistently Scaling Protein Language Model Representations

Machine Learning 2026-03-10 v1 Biomolecules

Abstract

Unlike the predictable scaling laws in natural language processing and computer vision, protein language models (PLMs) scale poorly: for many tasks, models within the same family plateau or even decrease in performance, with mid-sized models often outperforming the largest in the family. We introduce Reverse Distillation, a principled framework that decomposes large PLM representations into orthogonal subspaces guided by smaller models of the same family. The resulting embeddings have a nested, Matryoshka-style structure: the first k dimensions of a larger model's embedding are exactly the representation from the smaller model. This ensures that larger reverse-distilled models consistently outperform smaller ones. A motivating intuition is that smaller models, constrained by capacity, preferentially encode broadly-shared protein features. Reverse distillation isolates these shared features and orthogonally extracts additional contributions from larger models, preventing interference between the two. On ProteinGym benchmarks, reverse-distilled ESM-2 variants outperform their respective baselines at the same embedding dimensionality, with the reverse-distilled 15 billion parameter model achieving the strongest performance. Our framework is generalizable to any model family where scaling challenges persist. Code and trained models are available at https://github.com/rohitsinghlab/plm_reverse_distillation.

Keywords

Cite

@article{arxiv.2603.07710,
  title  = {Reverse Distillation: Consistently Scaling Protein Language Model Representations},
  author = {Darius Catrina and Christian Bepler and Samuel Sledzieski and Rohit Singh},
  journal= {arXiv preprint arXiv:2603.07710},
  year   = {2026}
}

Comments

Proceedings of ICLR 2026

R2 v1 2026-07-01T11:09:16.493Z