让大蛋白模型始终比小模型强,靠的是把大模型拆成小模型的叠加版。
Reverse Distillation: Consistently Scaling Protein Language Model Representations
- 用小模型引导大模型,将大模型嵌入分解为正交子空间,形成嵌套结构。
- 在ProteinGym上,反向蒸馏的ESM-2模型性能超越同维基线,150亿参数最强。
- 适合研究蛋白语言模型扩展性、需高效多尺度表示的科研人员。
与自然语言处理和计算机视觉中可预测的缩放规律不同,蛋白质语言模型(PLMs)的缩放效果差:许多任务中,同一家族的模型性能会停滞甚至下降,中等规模模型常优于最大模型。本文提出反向蒸馏(Reverse Distillation),一种基于同家族小模型引导的大模型表示分解框架,将大模型嵌入分解为正交子空间。所得嵌入具有嵌套的马特罗什卡式结构:大模型嵌入的前k维恰好等于小模型的表示。这确保了更大反向蒸馏模型始终优于更小模型。核心思想是:小模型因容量受限,优先编码广泛共享的蛋白质特征;反向蒸馏分离这些共性特征,并正交提取大模型的额外贡献,避免干扰。在ProteinGym基准测试中,反向蒸馏的ESM-2变体在相同嵌入维度下优于各自基线,其中150亿参数的反向蒸馏模型表现最优。该框架可推广至所有存在缩放挑战的模型家族。代码与训练模型见https://github.com/rohitsinghlab/plm_reverse_distillation。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。