[论文解读] Knowledge-Design: Pushing the Limit of Protein Design via Knowledge Refinement

Zhangyang Gao, Cheng Tan+1 more|arXiv (Cornell University)|May 20, 2023

Machine Learning in Bioinformatics被引用 6

一句话总结

Knowledge-Design 引入了一种自信感知的优化框架，利用预训练模型（ESM、ESM-IF、GearNet）的多模态知识，迭代修正蛋白质序列设计中低置信度的残基。通过集成门控注意力融合与记忆检索机制，其在 CATH4.2 上实现了 60.77% 的恢复率——比 PiFold 高出 9.11%，并成为首个在 CATH、TS50 和 TS500 所有基准测试中均突破 60% 恢复率的方法。

ABSTRACT

Recent studies have shown competitive performance in protein design that aims to find the amino acid sequence folding into the desired structure. However, most of them disregard the importance of predictive confidence, fail to cover the vast protein space, and do not incorporate common protein knowledge. After witnessing the great success of pretrained models on diverse protein-related tasks and the fact that recovery is highly correlated with confidence, we wonder whether this knowledge can push the limits of protein design further. As a solution, we propose a knowledge-aware module that refines low-quality residues. We also introduce a memory-retrieval mechanism to save more than 50\% of the training time. We extensively evaluate our proposed method on the CATH, TS50, and TS500 datasets and our results show that our Knowledge-Design method outperforms the previous PiFold method by approximately 9\% on the CATH dataset. Specifically, Knowledge-Design is the first method that achieves 60+\% recovery on CATH, TS50 and TS500 benchmarks. We also provide additional analysis to demonstrate the effectiveness of our proposed method. The code will be publicly available.

研究动机与目标

• 解决现有蛋白质设计模型在处理预测置信度方面存在的局限性，并缺乏对通用蛋白质知识的整合。
• 通过从预训练模型中蒸馏知识，突破迭代优化中的性能饱和，实现对局部极小值的逃离。
• 通过整合来自 ESM、ESM-IF 和 GearNet 的结构与序列知识，提升设计蛋白质序列的泛化能力与结构准确性。
• 通过引入记忆检索机制，缓存中间模块输出，从而缩短大规模蛋白质设计的训练时间。
• 通过结合自信感知优化与多模态知识融合，在多个基准测试（包括 CATH、TS50 和 TS500）上实现最先进性能。

提出的方法

• 提出一种自信感知模块，利用预测氨基酸的最大概率识别低置信度残基。
• 实现一种门控注意力融合机制，根据预测置信度自适应地结合结构（GearNet）、序列（ESM）和历史预测嵌入。
• 引入虚拟 MSA 和循环迭代技术，以提升优化性能与模型收敛性。
• 设计一种记忆检索机制，缓存中间模块输出，实现更快的推理速度，并将训练时间减少 50% 以上。
• 采用多阶段优化流水线训练 Knowledge-Design，基于多模态知识与置信度分数，迭代提升残基预测。
• 利用 ESM 和 ESM-IF 作为教师模型，在训练过程中注入归纳偏置，帮助逃离局部最优。

实验结果

研究问题

• RQ1 预测置信度能否被有效利用，以识别并优化蛋白质序列设计中的低质量残基？
• RQ2 整合来自预训练模型（ESM、ESM-IF、GearNet）的多模态知识，如何提升蛋白质设计的准确性和泛化能力？
• RQ3 与单一模型相比，结合多个预训练模型（如 ESM + ESM-IF）对恢复性能的影响如何？
• RQ4 记忆检索机制是否能显著缩短大规模蛋白质设计中的训练时间，且不损害模型性能？
• RQ5 与 PiFold 等基线模型相比，Knowledge-Design 在结构相似性方面提升了多少？

主要发现

• Knowledge-Design 在 CATH4.2 基准测试中实现了 60.77% 的恢复率，成为首个在 CATH、TS50 和 TS500 上均突破 60% 恢复率的方法。
• 在 CATH4.2 数据集上，该方法比 PiFold 提升了 9.11% 的恢复率，且当结合 ESM 与 ESM-IF 知识时达到最佳性能。
• ESM 与 ESM-IF 知识的结合带来了 9.11% 的性能提升，超过两者单独贡献之和（5.72% + 2.86%），表明存在协同效应。
• Knowledge-Design 提升了正确残基的置信度，同时减少了低置信度预测，表现为正确残基的置信度分布更趋近于 1.0。
• 通过 ESMFold 进行结构比较显示，与 PiFold 相比，Knowledge-Design 在 1a73、1a81 和 1ac1 上分别将 RMSD 降低了 15.9%、35.3% 和 60%。
• 记忆检索机制通过缓存中间模块结果，避免冗余前向传播，将训练时间减少了 50% 以上。