Reverse Distillation: Consistently Scaling Protein Language Model Representations

2026-03-08

AI Summary

This work addresses the performance saturation or degradation often observed when scaling up protein language models. To overcome this limitation, the authors propose a reverse distillation framework that decomposes embeddings from a large model into orthogonal subspaces guided by smaller models from the same family, thereby constructing a nested Matryoshka embedding structure. In this architecture, representations from smaller models serve as prefix subspaces of the larger model’s embeddings, effectively disentangling general-purpose features from incremental information. This approach achieves, for the first time, consistent performance gains with increasing model scale. Evaluated on the ProteinGym benchmark, reverse-distilled variants of ESM-2 consistently outperform baseline models at equivalent embedding dimensions, with the 15-billion-parameter model achieving state-of-the-art performance.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: (Large) Language Models
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Search and Retrieval-Augmented AI: Large language models for search

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.

Problem

• protein language models
• scaling laws
• model performance
• representation scaling
• embedding dimensionality

Innovation

• Reverse Distillation
• Protein Language Models
• Matryoshka Embeddings
• Orthogonal Subspaces
• Model Scaling