recAMPLIFY:递归蛋白质语言模型的大规模预训练
recAMPLIFY: Large-Scale Pretraining of Recursive Protein Language Models
精选理由
以递归预训练和自适应深度提升蛋白质预测性能,并降低模型计算成本。
AI 摘要
蛋白质语言模型(pLM)是蛋白质注释和突变效应预测的强大工具,但其发展在很大程度上依赖于扩大模型规模和数据量,这使得模型成本高昂且容易出现记忆化。测试时扩展提供了一条互补路径,但为自回归 LLM 开发的方法无法直接迁移:掩码编码器缺少生成循环,而氨基酸也无法用推理语言进行 verbalize。为此,我们将递归预训练应用于掩码式 pLM:由权重共享的 Transformer 核心在可调次数的循环中迭代优化与残基对齐的隐藏状态。我们将这一方案实例化为 recAMPLIFY:该 pLM 具有 155M 个独立参数,并使用与 AMPLIFY-350M 相同的 2T 个 token 进行训练。尽管参数量减少了 56%,它仍提升了 CASP 接触恢复率和四项监督任务中的三项表现,其中折叠准确率显著提高了 22.5%。由于 recAMPLIFY 的最优深度会因任务和测定而异,我们进一步提出门控混合递归(Gated Mixture Recurrence,GMR),这是一种通过递归过程中访问到的隐藏状态的概率加权混合来解码每个蛋白质的框架。GMR 在标量、低秩或逐维粒度上采用学习得到的残差门,以控制每个循环提出的更新有多少被纳入状态。这一方法将提前停止与门控统一为单一的自适应计算目标,同时在无门控时将 PonderNet 式提前停止作为特殊情况恢复出来。在 recAMPLIFY 之上叠加 GMR 后,预期层深度进一步降低了最多 37%,不过准确率上的权衡因任务而异。机制分析表明,残差门控能够在超出训练深度后控制递归稳定性,而停止深度则与生物学属性相对应。