Reverse Distillation: Consistently Scaling Protein Language Model Representations
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 12 of its 27 references that are in the map. First posted on arXiv 8 Mar 2026.
Builds on
12 of 27 references.
References · 12
- 2015Distilling the Knowledge in a Neural Networkcited by 10,076
- 2020Scaling Laws for Neural Language Modelscited by 6,411
- 2022Training Compute-Optimal Large Language Modelscited by 2,927
- 2018UMAP: Uniform Manifold Approximation and Projection for Dimension Reductioncited by 2,841
- 2022Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference timecited by 1,366
- 2019What Does BERT Look At? An Analysis of BERT's Attentioncited by 994
- 2020A Primer in BERTology: What we know about how BERT workscited by 807
- 2023Orthogonal Subspace Learning for Language Model Continual Learningcited by 296
- 2022Matryoshka Representation Learningcited by 268
- 2020BERTology Meets Biology: Interpreting Attention in Protein Language Modelscited by 85
- 2024xTrimoPGLM: Unified 100B-Scale Pre-trained Transformer for Deciphering the Language of Proteincited by 49
- 2025Sculpting Subspaces: Constrained Full Fine-Tuning in LLMs for Continual Learningcited by 15
Built on by
Not cited yet.