[论文解读] Pre-training Co-evolutionary Protein Representation via A Pairwise Masked Language Model

Liang He, Shizhuo Zhang+11 more|arXiv (Cornell University)|Oct 29, 2021

Machine Learning in Bioinformatics参考文献 53被引用 11

一句话总结

本文提出了一种成对掩码语言模型（PMLM），在预训练过程中捕捉蛋白质序列中的残基间共进化的依赖关系，通过建模被掩码残基对的联合概率而非独立预测，从而提升表征学习效果。该方法在标准MLM基础上，接触预测性能相对提升高达9%，在TAPE基准测试中超越基于MSA的基线模型超过7%，证明了仅使用序列的预训练方法在蛋白质建模中可超越传统的依赖MSA的方法。

ABSTRACT

Understanding protein sequences is vital and urgent for biology, healthcare, and medicine. Labeling approaches are expensive yet time-consuming, while the amount of unlabeled data is increasing quite faster than that of the labeled data due to low-cost, high-throughput sequencing methods. In order to extract knowledge from these unlabeled data, representation learning is of significant value for protein-related tasks and has great potential for helping us learn more about protein functions and structures. The key problem in the protein sequence representation learning is to capture the co-evolutionary information reflected by the inter-residue co-variation in the sequences. Instead of leveraging multiple sequence alignment as is usually done, we propose a novel method to capture this information directly by pre-training via a dedicated language model, i.e., Pairwise Masked Language Model (PMLM). In a conventional masked language model, the masked tokens are modeled by conditioning on the unmasked tokens only, but processed independently to each other. However, our proposed PMLM takes the dependency among masked tokens into consideration, i.e., the probability of a token pair is not equal to the product of the probability of the two tokens. By applying this model, the pre-trained encoder is able to generate a better representation for protein sequences. Our result shows that the proposed method can effectively capture the inter-residue correlations and improves the performance of contact prediction by up to 9% compared to the MLM baseline under the same setting. The proposed model also significantly outperforms the MSA baseline by more than 7% on the TAPE contact prediction benchmark when pre-trained on a subset of the sequence database which the MSA is generated from, revealing the potential of the sequence pre-training method to surpass MSA based methods in general.

研究动机与目标

• 为解决从大量未标注数据中学习有效蛋白质序列表征的挑战，且不依赖计算密集型多序列比对（MSA）。
• 通过一种新型预训练目标，直接从序列中捕捉残基间的共进化信息，该目标建模被掩码残基对之间的依赖关系。
• 通过联合建模残基对，学习更丰富、更具上下文感知的表征，从而提升下游蛋白质结构与功能预测任务的性能。
• 评估仅使用序列的预训练方法是否能在蛋白质建模的表征质量上超越基于MSA的方法。

提出的方法

• 提出一种成对掩码语言模型（PMLM），通过建模两个被掩码残基的联合概率，扩展标准掩码语言建模方法，而非独立建模其概率。
• PMLM目标通过学习残基对的概率不等于其各自概率的乘积，显式捕捉共进化信号，反映残基间的依赖关系。
• 模型在大规模蛋白质序列数据库（如Pfam、UR50）上进行预训练，采用多任务设置，同时使用标准MLM和新型PMLM目标。
• 预训练编码器通过微调下游头用于接触预测、二级结构预测和远距离同源检测等任务。
• 采用残差网络（ResNet）头以增强下游任务的表征学习能力，性能优于简单的线性头。
• 在TAPE基准上评估该方法，并与强基线模型（包括基于MSA的模型和标准Transformer）进行比较。

实验结果

研究问题

• RQ1 是否能从原始蛋白质序列中有效捕捉共进化信息，而无需依赖多序列比对（MSA）？
• RQ2 与标准MLM中独立掩码相比，建模被掩码残基对的联合概率是否能带来更优的蛋白质表征学习？
• RQ3 仅使用序列的预训练方法是否能在下游蛋白质结构预测任务中超越基于MSA的方法？
• RQ4 模型规模和下游架构如何影响所提出的PMLM方法的性能？

主要发现

• 在相同训练设置下，PMLM模型在接触预测性能上相比标准MLM基线最高提升9%。
• 在TAPE接触预测基准上，PMLM相比基于MSA的基线模型性能提升超过7%，证明了仅使用序列预训练的优势。
• 模型规模从base增至large，最终到XL（715M参数）时，接触预测的F1-score从57.7%提升至71.7%，性能提升达14%。
• 与线性头相比，使用ResNet头将UR50数据集上的精确率从39.2%提升至41.6%，表明更复杂的下游模块具有优势。
• 在二级结构预测任务中，PMLM-base与TAPE Transformer基线表现相当，CB513数据集上的得分分别为0.728和0.730。
• 在远距离同源预测任务中，PMLM-base在家族级别和超家族级别基准上分别优于TAPE Transformer达10%和10.6%，表明其能有效捕捉进化信号。