PaperlayerView exact revision on arXiv

Source-linked AI summary

Pre-training Co-evolutionary Protein Representation via A Pairwise Masked Language Model

Liang He, Shizhuo Zhang, Lijun Wu, Huanhuan Xia, Fusong Ju, He Zhang, Siyuan Liu, Yingce Xia, Jianwei Zhu, Pan Deng, Bin Shao, Tao Qin, Tie-Yan Liu

arXiv:2110.15527v1cs.CLcs.AI

TL;DR

Protein representation learning seeks to capture co-evolutionary information in residue sequences, often extracted from MSAs. This paper pre-trains a Pairwise Masked Language Model that predicts masked residue pairs jointly. It reports improved contact prediction and performance beyond the MSA baseline on TAPE under the stated sequence-only setting.

  • Problem

    Protein sequence representation learning seeks to capture inter-residue co-variation, which is relevant to protein structure and function, without relying solely on MSA extraction.

    Source: §1 · p. 2

  • Method

    PMLM pre-trains a protein sequence encoder to predict masked residue pairs jointly alongside masked tokens.

    Source: §3.1 · p. 5, §3.2 · p. 5

  • Results

    PMLM significantly improves contact prediction over baselines and surpasses the MSA baseline on the TAPE contact benchmark with sequence-only input.

    Source: §1 · p. 3, §5 · p. 10

  • Takeaways & Limitations

    The results support sequence pre-training as a way to capture co-evolutionary information for contact prediction without MSA input at evaluation.

    Source: §1 · p. 3, §5 · p. 10

    PMLM’s potential is not fully developed, and extending pairwise supervision to triples would entail cubic storage and computation costs in the number of masked tokens.

    Source: §5 · p. 11

Abstract

from arXiv · show

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.

1 INTRODUCTION

Protein co-evolutionary information is reflected in inter-residue co-variation and matters for both structure and function. The section motivates modeling masked residue pairs directly and reports contact-prediction gains over MLM and MSA baselines.

1 INTRODUCTION

  • Inter-residue co-variation is important to protein structure and function, while labels are costly and unlabeled sequences are growing rapidly.The text connects residue co-evolution with contact inference and mutational effects.Source: §1 · p. 2
  • The section asks whether sequence pre-training can capture co-evolutionary information directly instead of extracting it from multiple sequence alignments.The target information is pairwise statistics such as Q(x_i, x_j|X).Source: §1 · p. 2
  • Conventional masked language models predict masked residues independently, whereas PMLM models the probability of masked residue pairs jointly.The paper argues P(x_i, x_j|X/M) need not equal P(x_i|X/M) · P(x_j|X/M).Source: §1 · p. 3
  • PMLM improves contact prediction over the MLM baseline and surpasses the MSA baseline on TAPE when pre-trained on UR50 with sequence-only input.Contact prediction is used because of its established relationship to co-evolutionary information.Source: §1 · p. 3

2 RELATED WORK

Prior protein representation work spans MSA-based coupling analysis, sequence pre-training, and applications of learned representations. The section also describes scaling, auxiliary objectives, and generative approaches.

2 RELATED WORK

  • MSA-based methods model residue couplings, but their performance drops when a target’s MSA has few homologous sequences.Neural methods may use DCA-derived features or raw MSAs; raw-MSA methods avoid hand-crafted feature information loss.Source: §2 · p. 3
  • Protein sequence pre-training improves over one-hot representations, although TAPE reports that sequence-only models lag alignment-based methods on downstream tasks.RNN approaches include Bi-LSTM, UniRep, and AWD-LSTM models.Source: §2 · pp. 3–4
  • Transformer protein models include ESM, PRoBERTa, contrastive-learning methods, and MSA Transformer, which depends on database searches to generate alignments.ESM also reports effects of sequence diversity and model size on encoder performance.Source: §2 · p. 4
  • Large-scale sequence models have captured evolutionary information, and ProtTrans embeddings outperform alignment-free baselines for secondary structure prediction.The section describes models trained with thousands of GPUs or up to 3 billion parameters.Source: §2 · p. 4
  • Other approaches add auxiliary labels or generate sequences, while learned-model outputs have been related to contacts, evolutionary order, mutational effects, and virus escape.Examples include same-family or profile prediction, and ProGen’s taxonomically conditioned sequence generation.Source: §2 · p. 4

3 METHOD

The method adds a pairwise masked-token objective to protein sequence pre-training. A Transformer encoder and prediction heads support pair prediction, with fine-tuning for contact prediction.

3 METHOD

  • The method builds language models from protein sequence data, using masked-token prediction as the vanilla MLM objective.The loss predicts each masked token conditioned on the masked sequence; D is the sequence set and θ denotes encoder parameters.Source: §3.1 · p. 4, §3 · p. 4

3.1 PAIRWISE MASKED LANGUAGE MODEL

  • PMLM adds a loss for masked residue pairs, whose joint probability can differ from the product of individual token probabilities.The combined MLM and PMLM losses are balanced by λ, set to 1; labels come from the sequences and remain self-supervised.Source: §3.1 · p. 5

3.2 MODEL ARCHITECTURE

  • A stacked Transformer encoder maps raw protein sequences to vectors, with separate two-layer MLP heads predicting 20 residue classes and 400 residue-pair classes.The pre-training uses both masked-token prediction and masked-pair prediction.Source: §3.2 · p. 5

3.3 MASKED PAIR PREDICTION

  • The pair head concatenates the dot product and difference of residue-pair vectors, then predicts labels for distinct masked-position pairs.Pairs with identical positions are excluded because they are already represented in the MLM loss; the number of pair labels is |M|^2−|M|.Source: §3.3 · p. 5, Figure 1 · p. 6

3.4 FINE-TUNING FOR CONTACT PREDICTION

  • For contact prediction, the fine-tuning network is built on the first fully connected layer of the pair prediction head.Other downstream tasks are fine-tuned on encoder outputs.Source: §3.4 · p. 6

4 EXPERIMENT EVALUATION

Across validation and contact-prediction experiments, PMLM’s joint residue predictions differ from independent token predictions and outperform the reported MLM and MSA baselines. Ablations also show performance varies with pre-training data, model size, and downstream module.

4.1 EXPERIMENT SETUP

  • Experiments used Pfam and UR50, with 12-layer base, 34-layer large, and 36-layer, 1280-hidden XL encoders; TAPE and RaptorX used shallow and ResNet decoders, respectively.The XL model was pre-trained on UR50.Source: §4.1 · pp. 6–7

4.2 PRE-TRAINING VALIDATION

  • On Pfam and UR50 validation, non-negligible ΔACC and KL histograms indicate PMLM joint residue probabilities differ from products of individual probabilities.Figure 3 shows KL divergence for PMLM-large on the two validation datasets.Source: §4.2 · pp. 7–8, Figure 3 · p. 8, Table 2 · p. 8

4.3 COMPARISON ON TAPE-CONTACT

  • PMLM-base gains over 9% over MLM-base on TAPE-Contact; UR50-pretrained PMLM-xl beats the MSA baseline by about 8%.Table 3 reports medium-/long-range Precision@L/5.Source: §4.3 · p. 8, Table 3 · p. 9

4.4 COMPARISON ON RAPTORX-CONTACT

  • PMLM-base beats ESM Transformer-12 by about 4%; PMLM-xl exceeds ESM Transformer-34 by more than 9% and ESM-1b by 3%.Table 4 reports long-range Precision@L.Source: Table 4 · p. 9, §4.4 · p. 9

4.5 ABLATION STUDY

  • On UR50, diverse-data PMLM-base rises from 47.2 to 57.7; PMLM-xl scores 71.7 versus 57.7, and XL (+ ResNet) improves by more than 18% over base (+ ResNet).The data and encoder-size comparisons use Table 3; the ResNet comparison uses Table 5.Source: §4.5 · pp. 9–10, Table 3 · p. 9, Table 5 · p. 10
  • With PMLM-base fixed, the eight-layer ResNet variant scores 41.6 versus 39.2 for the linear layer.The comparison concerns downstream contact-prediction modules.Source: Table 4 · p. 9, §4.5 · p. 10
  • Secondary-structure accuracy is close to TAPE; remote-homology performance is slightly lower at fold level but up to 10% better at superfamily and family levels.The reported remote-homology comparisons are on the TAPE benchmark.Source: §4.5 · p. 10, Table 6 · p. 10, Table 7 · p. 10

5 DISCUSSION

The paper proposes PMLM to capture co-evolutionary information during protein pre-training and reports improved contact prediction, including surpassing the MSA baseline on TAPE. The authors note that PMLM’s potential remains underdeveloped and that more efficient pre-training is needed for proposed triple-residue supervision.

5 DISCUSSION

  • PMLM captures co-evolutionary information during pre-training and improves contact prediction over baselines, surpassing the MSA baseline on the TAPE contact benchmark.The authors state that its pre-trained model generates a better representation for global structure.Source: §5 · p. 10
  • The authors say PMLM’s potential is not fully developed, and increasing token mask probability might improve its representation.Source: §5 · pp. 10–11
  • Triple Masked Language Model supervision from multiple residues might help, but its storage and computation costs grow cubically with the number of masked tokens.Efficient pre-training is left for future work.Source: §5 · p. 11

Related papers

Loading 2110.15527v1…