Reverse Distillation: Consistently Scaling Protein Language Model Representations

Abstract: Unlike the predictable scaling laws in natural language processing and computer vision, protein LLMs (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.

• The paper introduces Reverse Distillation, a method that decomposes large protein language model representations into smaller validated subspaces to ensure consistent, monotonic scaling.
• It employs principal component regression and singular value decomposition to isolate orthogonal features, preserving essential biological signals while enhancing interpretability.
• Empirical evaluations demonstrate that reverse-distilled models outperform baselines on variant effect prediction and structural inference tasks, providing more granular and reliable insights.

Reverse Distillation: Consistently Scaling Protein LLM Representations

Introduction

Scaling protein LLMs (PLMs) stands in sharp contrast with NLP and vision models, where performance improvements are generally monotonically linked to increasing model and representation size. In PLMs, scaling fails to deliver consistent improvement, and larger models sometimes underperform compared to smaller ones, particularly on functional and variant effect prediction benchmarks. This paper introduces the Reverse Distillation framework, which decomposes large PLM representations into orthogonal subspaces constructed from smaller models, reconstructing model families into Matryoshka-style embeddings. This enables consistent and interpretable scaling, ensuring that larger reverse-distilled models strictly subsume and surpass the representational power of smaller models.

The theoretical underpinning is that smaller models, being capacity-bounded, encode broadly-shared biological features optimally for many tasks. Directly concatenating the full expressivity of larger models dilutes these signals via task-irrelevant variance. Reverse Distillation isolates and preserves features from each scale, maximizing residual information gain and preventing destructive feature interference.

Figure 1: Overview of the Reverse Distillation procedure: decomposing the representational space of large PLMs into smaller-scale features and orthogonal residuals via SVD, yielding Matryoshka-structured embeddings across model scales.

Methodology

The framework operates over a hierarchy of PLMs with increasing parameter counts and embedding dimensions (e.g., ESM-2: 8M–15B parameters). For models $M_r$ ($k_r$-dimensional) and $M_p$ ($k_p$-dimensional, $k_p>k_r$), the representation is decomposed as:

$$H_p \approx [H_r, H_{res}], \quad H_{res} \perp H_r$$

Here, $H_r$ comes from the smaller model and $H_{res}$ contains features unique to the larger model, orthogonalized using principal component regression (PCR) followed by SVD. The decomposition procedure is recursively chained across all models in a family, producing embeddings where, for any prefix dimension $d$, the first $d$ features correspond to a valid, smaller reverse-distilled model.

This construction ensures:

• Matryoshka property: Prefixes of embeddings are themselves valid representations at smaller scales, enabling dynamic inference-time truncation and adaptive computational cost.
• Monotonicity: Larger models' reverse-distilled representations strictly outperform smaller ones at the same embedding dimensionality.
• Interpretability: Each contribution from a model scale is isolated and directly attributable, opening new avenues for interpretability analyses.

Empirical Evaluation

Scaling on Protein Fitness Prediction

On the ProteinGym benchmark for deep mutational scanning (DMS), the reverse-distilled ESM-2 models (rd.650M, rd.3B, rd.15B) consistently restore monotonic scaling and outperform their respective baselines at equivalent dimensionality. The rd.15B achieves the strongest variant effect prediction performance (single-mutation Spearman $\rho=0.904\pm0.037$, two-mutation $\rho=0.720\pm0.120$). Critically, reverse-distilled models exhibit strictly superior performance in the vast majority of datasets compared to both baseline and lower-scale models, reversing the previous trend where medium-sized models often outperformed the largest.

Downstream Protein Property Prediction

Reverse-distilled models consistently outperform their base counterparts on diverse structural and functional prediction tasks, including secondary structure, metal ion binding, and localization prediction—all evaluated with AUPR metrics, with rd.15B achieving highest or near-highest scores in almost all tasks.

Functional Attribution in Embedding Spaces

Probing with sparse autoencoders (SAE), rd.35M’s features are associated with a larger number of distinct Gene Ontology (GO) terms (40 per feature vs. 32 for ESM-2 35M), and these terms are significantly less general, demonstrating that reverse-distilled embeddings capture more specific and distinct biological functions.

Figure 2: Reverse-distilled embeddings support a higher number of enriched and less general GO terms in SAE analysis, enhancing functional interpretability and specificity of learned representations.

Algorithmic and Theoretical Analysis

This linear framework is MSE-optimal within the class of Matryoshka-constrained embeddings. PCR (with noise component removal via random matrix theory) is preferred over standard OLS for the initial mapping, further improving downstream performance, as demonstrated in ablation studies.

Figure 3: PCR-based mapping outperforms OLS consistently in cross-scale linear decomposition, emphasizing the need to denoise small-scale representations for effective hierarchical scaling.

Additionally, while simple PCA+concatenation baselines can approximate similar overall performance, they lack the Matryoshka property, making them unsuitable for scalable and interpretable model chaining.

Figure 4: Ablation demonstrates that while naive PCA+concatenation offers strong performance, only reverse distillation provides cross-scale structural consistency via the Matryoshka property.

Practical Considerations

Inference with reverse-distilled models is efficient: the required chained forward passes through several smaller models introduce only a moderate overhead (e.g., rd.650M takes 1.69x the time of baseline 650M), and the prefix structure allows reusing computed embeddings at lower scales. Code and pre-trained weights are openly available for integration and extension.

Implications and Future Directions

This work establishes that the apparent “scaling wall” for PLMs is largely a function of representational entanglement, not a limitation of upstream model capacity. By realigning how information is hierarchically partitioned and reused, the authors demonstrate that monotonic scaling and consistent improvement are achievable without additional parameter growth, nonlinear transforms, or re-training.

From a practical perspective, this enables robust, interpretable, and storage-efficient application of large PLMs across diverse biological prediction tasks, with reduced compute and infrastructure cost. Theoretically, it paves the way for further extensions along multiple axes:

• Nonlinear decompositions: The utility of nonlinear mappings is evidenced by further improved reconstruction $R^2$; future work can systematize these advances while maintaining interpretability and scalability.
• Single-model reverse-distillation via efficient fine-tuning (e.g., LoRA): Embedding the full hierarchy at inference time from a single large model forward pass will enhance efficiency and generative task applicability.
• Transfer to other foundation models: While biology provides a stringent testbed, the approach is applicable wherever scaling laws degrade due to entanglement and redundancy (e.g., genomics, materials, chemistry, non-biological sequence domains).

Conclusion

Reverse Distillation introduces a theoretically grounded, efficient, and empirically validated solution to the PLM scaling problem. Linear Matryoshka-style decomposition of PLM representations recovers monotonic scaling behavior, boosts cross-task performance, and yields more granular, biologically specific representations. This approach invites further exploration of hierarchical, interpretable, and parameter-efficient model scaling strategies in computational biology and beyond.