【蛋白设计】LM-DESIGN: 改造蛋白质语言模型 进行基于结构的蛋白序列设计

——背景——

基于结构的蛋白序列设计是蛋白质理性设计领域的重要问题，随着深度学习的发展，GVP-GNN [1]、ESM-IF1 [2]、ProteinMPNN [3]、PiFold [4]等方法基于Encoder-Decoder架构将蛋白序列设计建模为从图到序列的端到端学习问题，取得了显著的成效。但是，目前已有的蛋白质序列设计方法仍面临着以下挑战：第一，实验所确定的蛋白质结构数据非常有限，例如CATH数据集中已知的蛋白质结构比UniRef序列数据库中的序列数据小几个数量级（<0.1%），这使得在这些数据上训练得到的模型无法充分探索蛋白质序列空间。第二，蛋白质表面、loop区等结构灵活可变，其残基类型更多取决于序列进化信息而非结构信息，因而现有的完全基于结构信息的模型在对这些区域进行序列设计时常常生成不具有期望生物功能的序列，例如ProteinMPNN在对酪氨酸激酶（Tyrosine kinase）进化上高度保守的activation loop进行设计时生成了无效序列（图1）。

图1 酪氨酸激酶activation loop序列设计

以ESM-1b [5]、ESM-2 [6] 为代表的蛋白质语言模型（Protein Language Models, PLMs）通过学习大量蛋白质序列数据获得有关序列进化的先验知识，通过改造蛋白质语言模型进行序列设计可以更充分地利用这些从蛋白序列中获得的信息，从而提升已有模型的表现。近日，字节跳动提出了基于结构的蛋白序列设计方法LM-DESIGN [7]，该方法将预训练的蛋白质语言模型与轻量级的结构适配模块（structural adapter）相结合，在蛋白质序列设计任务上取得了新SOTA，并提供了一种高度模块化、可扩展的模型框架，该工作已被ICML 2023接收并入选Oral。

——方法——‍

图2 LM-DESIGN的模型架构

LM-DESIGN的模型架构如图2所示。模型的主体部分是预训练的蛋白质语言模型，作者使用预训练的蛋白质结构编码器（structure encoder）对蛋白质骨架的三维结构进行编码，通过结构适配模块将蛋白质结构信息与来自蛋白质语言模型的序列信息相融合，通过条件掩码语言建模（Conditional Masked Language Modeling, CMLM）的方式，基于初始的蛋白序列和蛋白质主链三维结构信息，多次迭代生成最终的蛋白序列。

图3 结构适配模块的结构

结构适配模块如图3所示。该模块由多头注意力（Multihead ATTN）层与Bottleneck FFN层组成，多头注意力层把蛋白序列信息作为query，在蛋白序列信息和结构信息之间做cross attention，并使用RoPE（Rotary Position Embedding）作为位置编码。Bottleneck FFN层对多头注意力层的输出依次进行降维、非线性变换和升维，提高模型的表现力。LM-DESIGN模型中的结构适配模块类似于Adapter-BERT [8] 的adapter层，类似的模型架构在自然语言处理（Natural Language Processing, NLP）领域被广泛应用，能够实现对预训练模型的高效微调。考虑到训练速度等因素，作者仅在蛋白质语言模型最后一层Encoder后增加了结构适配模块。

LM-DESIGN模型在训练时输入蛋白质主链结构X和序列S，随机掩盖部分序列S_marked，模型根据可见的序列S_obs预测概率

通过最大化上述概率进行模型训练。在进行推断时，LM-DESIGN模型根据结构信息X在初始序列S⁽⁰⁾基础上迭代采样

‍

经过T步迭代（T为超参数，默认T=5）后生成最终的蛋白序列。在实际的序列生成任务中，作者使用ProteinMPNN作为默认的蛋白质结构编码器、使用ESM-1b作为默认的蛋白质语言模型，默认的序列采样方式为确定性贪心解码（greedy deterministic decoding）即直接最大化概率p(S|·)，初始序列S⁽⁰⁾为原始的ProteinMPNN生成的序列。为控制生成序列的多样性和收敛速度，作者额外引入了温度τ作为超参数对log p(S|·)进行缩放。

——结果——‍

作者首先在常用的CATH 4.2和CATH 4.3数据集上评估了LM-DESIGN模型的表现。如图4所示，在前述默认设定下，LM-DESIGN在CATH 4.2和CATH 4.3上均达到了新SOTA，在生成序列困惑度（perplexity）、序列恢复率（recovery）等指标上表现提升显著，且训练过程不需要额外的蛋白质结构数据。由于使用了预训练的蛋白质语言模型和结构编码器，LM-DESIGN可训练的参数量占总参数量的比例较小，训练成本较低。此外，LM-DESIGN模型能够对不同的预训练模型加以适配，具有高度的可扩展性。

图4 LM-DESIGN在CATH 4.2和CATH 4.3上的表现

作者随后在ProteinMPNN所使用的多链蛋白复合物（multi-chain complex）数据集上测试了LM-DESIGN进行序列设计的表现。如图5所示，LM-DESIGN的表现相比ProteinMPNN有显著提升，且仅把ProteinMPNN的自回归生成任务替换成CMLM任务即可实现序列恢复率4.39%的提升。

图5 LM-DESIGN在多链蛋白复合物序列设计任务上的表现

作者随后对于LM-DESIGN模型表现进行了详细的探究，探讨了迭代次数T（图6A）、温度τ（图6B）、使用额外训练数据（AlphaFold 2预测的结构数据，图6C）、预训练模型参数量（图6D）等因素对LM-DESIGN模型表现的影响。可以看出，迭代的序列生成过程提升了LM-DESIGN模型的表现，通过控制温度$\tau$可以平衡生成序列的质量与多样性，在一定范围内进行数据增强（data augmentation）有助于提升模型表现，而使用参数量更大的预训练模型也有助于模型表现的提升。‍

图6 LM-DESIGN模型表现的影响因素

作为案例研究，作者测试了LM-DESIGN对CATH数据集里结构高度相似、但特定区域具有不同的序列行使不同的生物功能的四个蛋白进行序列设计的表现，结果如图7所示，可见LM-DESIGN对于蛋白结构较小的区别高度敏感，能够很好地设计出具有相应生物功能的蛋白序列，而ProteinMPNN的设计成功率相对较低。

作者最后测试了LM-DESIGN模型的zero-shot生成能力。抗体的互补决定区（Complementary Determining Regions, CDRs）序列设计是蛋白序列设计领域最困难的任务之一，作者测试了LM-DESIGN在RAbD数据集上进行zero-shot CDR-H3序列生成的表现，在基于真实结构（图8A）的条件下，LM-DESIGN在多项指标上超越了先前抗体结构-序列co-design的SOTA模型MEAN [9]，而基于MEAN预测的结构进行序列设计（图8B）时，向初始结构加入一定扰动也能提升LM-DESIGN的模型表现。上述结果证明LM-DESIGN模型能够很好地泛化到未训练过的序列生成任务上。

图8 LM-DESIGN zero-shot抗体设计表现

——总结——

本文介绍了一种通过向预训练的蛋白质语言模型中引入结构信息来进行蛋白序列设计的方法LM-DESIGN，该方法通过轻量级的结构适配模块融合了蛋白序列信息和主链结构信息，基于蛋白结构迭代生成序列，在多项序列设计任务上取得SOTA，该模型框架具有高度的可扩展性和泛化性能。不过，该模型尚不能对生成序列的性质进行约束或引导，并且设计出的蛋白序列并未经湿实验合成与表征，其生成序列的合理性仍有待进一步检验。

LM-DESIGN方法证明经大规模预训练的蛋白语言模型所学到的蛋白序列进化信息对于蛋白质序列设计任务具有重要的意义，结合这些信息与基于物理的先验知识能够更好地提升序列设计模型的表现，节约模型训练的成本。此外，LM-DESIGN模型架构在蛋白质序列-结构co-design方面也有一定的潜力，对于后续工作具有很好的参考价值。

参考文献

[1] Jing, B., Eismann, S., Suriana, P., Townshend, R. J., & Dror, R. (2020). Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411.

[2] Hsu, C., Verkuil, R., Liu, J., Lin, Z., Hie, B., Sercu, T., ... & Rives, A. (2022, June). Learning inverse folding from millions of predicted structures. In International Conference on Machine Learning (pp. 8946-8970). PMLR.

[3] Dauparas, J., Anishchenko, I., Bennett, N., Bai, H., Ragotte, R. J., Milles, L. F., ... & Baker, D. (2022). Robust deep learning–based protein sequence design using ProteinMPNN. Science, 378(6615), 49-56.

[4] Gao, Z., Tan, C., & Li, S. Z. (2022). PiFold: Toward effective and efficient protein inverse folding. arXiv preprint arXiv:2209.12643.

[5] Rives, A., Meier, J., Sercu, T., Goyal, S., Lin, Z., Liu, J., ... & Fergus, R. (2021). Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15), e2016239118.

[6] Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., ... & Rives, A. (2022). Language models of protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022, 500902.

[7] Zheng, Z., Deng, Y., Xue, D., Zhou, Y., Ye, F., & Gu, Q. (2023). Structure-informed language models are protein designers. bioRxiv, 2023-02.

[8] Houlsby, N., Giurgiu, A., Jastrzebski, S., Morrone, B., De Laroussilhe, Q., Gesmundo, A., ... & Gelly, S. (2019, May). Parameter-efficient transfer learning for NLP. In International Conference on Machine Learning (pp. 2790-2799). PMLR.

[9] Kong, X., Huang, W., & Liu, Y. (2022). Conditional antibody design as 3d equivariant graph translation. arXiv preprint arXiv:2208.06073.

---

作者：王宇哲

审稿：黄志贤

编辑：黄志贤

GoDesign

ID：Molecular_Design_Lab