Source-linked AI summary

xTrimoPGLM: Unified 100B-Scale Pre-trained Transformer for Deciphering the Language of Protein

Bo Chen, Xingyi Cheng, Pan Li, Yangli-ao Geng, Jing Gong, Shen Li, Zhilei Bei, Xu Tan, Boyan Wang, Xin Zeng, Chiming Liu, Aohan Zeng, Yuxiao Dong, Jie Tang, Le Song

arXiv:2401.06199v2q-bio.QMcs.AIcs.LG

TL;DR

Existing protein language models are designed for specific tasks depending on their pre-training framework, leaving unified protein understanding and generation as a challenge. xTrimoPGLM uses a unified pre-training framework balancing two objectives, and its 100B model significantly surpassed state-of-the-art methods in 15 of 18 tasks.

Abstract

from arXiv · show

Protein language models have shown remarkable success in learning biological information from protein sequences. However, most existing models are limited by either autoencoding or autoregressive pre-training objectives, which makes them struggle to handle protein understanding and generation tasks concurrently. We propose a unified protein language model, xTrimoPGLM, to address these two types of tasks simultaneously through an innovative pre-training framework. Our key technical contribution is an exploration of the compatibility and the potential for joint optimization of the two types of objectives, which has led to a strategy for training xTrimoPGLM at an unprecedented scale of 100 billion parameters and 1 trillion training tokens. Our extensive experiments reveal that 1) xTrimoPGLM significantly outperforms other advanced baselines in 18 protein understanding benchmarks across four categories. The model also facilitates an atomic-resolution view of protein structures, leading to an advanced 3D structural prediction model that surpasses existing language model-based tools. 2) xTrimoPGLM not only can generate de novo protein sequences following the principles of natural ones, but also can perform programmable generation after supervised fine-tuning (SFT) on curated sequences. These results highlight the substantial capability and versatility of xTrimoPGLM in understanding and generating protein sequences, contributing to the evolving landscape of foundation models in protein science.

Introduction · Results

This section examines xTrimoPGLM’s unified approach to protein understanding and generation, including pre-training and fine-tuning strategies. It reports findings across protein-understanding tasks, structure prediction, and natural-like or targeted sequence generation, while noting practical and out-of-distribution limitations.

Introduction

  • xTrimoPGLM was introduced as a 100-billion-parameter unified model for protein understanding and generation, addressing the task-specific limits of existing models.The paper presents it as a unified pre-training framework and foundation model.Source: Introduction · p. 4
  • xTrimoPGLM-100B surpassed previous state-of-the-art methods on 15 of 18 tasks spanning protein structure, interactions, functionality, and developability.The reported benchmark experiments used linear probing and advanced fine-tuning techniques.Source: Introduction · p. 4
  • xTrimoPGLM showed lower perplexity than other models on two out-of-distribution protein sets, while scaling experiments showed that larger models commonly yielded better performance.Figure 1 describes the perplexity comparisons and the scaling analysis.Source: Introduction · p. 4, Figure 1 · p. 6
  • xT-Fold achieved a 5-point TM-score increase over ESMFold on CASP15 and faster inference across various scenarios.xT-Fold is the structure-prediction tool developed from xTrimoPGLM.Source: Introduction · p. 5
  • Generated sequences had a median sequence identity of 11.7% and could be tailored to structural and biophysical properties through supervised fine-tuning.The paper describes this fine-tuning capability as enabling programmable protein generation.Source: Introduction · p. 5
  • The paper identifies practical limitations requiring adaptation to diverse protein tasks, improved structure-prediction accuracy, and fewer generative hallucinations for real-world drug design.The authors say these enhancements are needed to bridge the gap between model capabilities and practical application.Source: Introduction · p. 5

The xTrimoPGLM Framework

  • The framework simultaneously optimizes GLM blank infilling with bidirectional input processing and an MLM objective to enhance understanding.The two objectives use distinct indicator tokens.Source: The xTrimoPGLM Framework · p. 7
  • Pre-training first uses MLM for approximately 400 billion tokens, then combines MLM and GLM at a 20%/80% ratio for 600 billion tokens to develop understanding and generation.The second phase is intended to refine both representational and generative abilities.Source: The xTrimoPGLM Framework · pp. 7–8
  • The loss-scaling curve from models with 1 million to 1 billion parameters trained on 100 billion tokens closely aligns with the observed training loss of xTrimoPGLM-100B.This supports the anticipated scaling law for the 100-billion-parameter model.Source: The xTrimoPGLM Framework · p. 8
  • Perplexity was 10.81 and 6.70 for xTrimoPGLM-100B versus 10.98 and 7.25 for ESM2-15B, and 13.35 and 8.78 versus PROGEN2-xlarge’s 14.30 and 9.75.The paired values correspond to datasets filtered at 0.5 and 0.9 sequence identity, respectively.Source: The xTrimoPGLM Framework · p. 9
  • The models reached reported perplexity levels with less training data: under 480 billion MLM tokens matched ESM2 trained on 1 trillion, while under 120 billion GLM tokens paralleled PROGEN2-xlarge trained on 350 billion.The comparisons used 400 billion phase-1 and 80 of 100 billion phase-2 MLM tokens, and 120 of 150 billion phase-2 GLM tokens; the PROGEN2 learning-rate schedule had not ended, potentially overestimating its loss.Source: The xTrimoPGLM Framework · p. 9

Evaluation on Protein Understanding benchmarks

  • LoRA consistently improves task performance over static probing while adding less than 1% in trainable parameters at rank 8.The reported fine-tuning results also show a scaling trend in task performance with supervised fine-tuning.Source: Evaluation on Protein Understanding… · p. 11
  • Across protein-related downstream tasks, exponentially greater pre-training compute correlated with linear performance improvements, and most tasks improved with training FLOPs and model size.The authors report this scaling behavior across four categories of protein-related tasks.Source: Evaluation on Protein Understanding… · pp. 11–13
  • Across tasks, the reported performance patterns were: 44% showed a gradual increase, 28% exhibited a “Grokking” phenomenon, and 16% did not follow a strictly linear pattern.The Grokking phenomenon was described as a sudden increase in test-data performance as the model’s knowledge base expanded to out-of-distribution scenarios.Source: Evaluation on Protein Understanding… · p. 12

Evaluation on Protein Structure Prediction

  • xT-Fold scored 0.86 on CAMEO and 0.70 on CASP15, versus ESMFold’s 0.85 and 0.65 and OmegaFold’s 0.80 and 0.60.On CAMEO, PLM-based models closely matched RosettaFold; on the more OOD CASP15 set, PLM-based methods lagged behind MSA-augmented approaches.Source: Evaluation on Protein Structure… · pp. 14–15
  • MSA/template-based models were approximately 10x to 50x slower than PLM-based models, while xT-Fold had a marginal speed advantage over ESMFold and OmegaFold, especially on longer sequences.The speed gap persisted even against AlphaFold2 accelerated with FlashAttention.Source: Evaluation on Protein Structure… · p. 15
  • xTrimoPGLM perplexity correlated with TM-score, and xT-Fold performed better than ESMFold and OmegaFold on samples with intermediate perplexity.At moderately high perplexity, xT-Fold was less effective than AlphaFold2 and RosettaFold.Source: Evaluation on Protein Structure… · p. 15
  • PPL–TM-score pairs were (-0.415, 1.7e-9) on CAMEO, (-0.579, 1e-5) on CASP14, and (-0.239, 8e-2) on CASP15; correlations were clearly negative only on CAMEO and CASP14.Lower perplexities were associated with more accurate predicted structures; scaling single-sequence models improved performance but struggled with OOD data, which MSA augmentation effectively addressed.Source: Evaluation on Protein Structure… · p. 16

Evaluation on Protein Sequence Generation

  • 84.0 vs 74.3 mean pLDDT and 0.695 vs 0.522 mean TM-score favor xTrimoPGLM over PROGEN2-xlarge, alongside mean sequence identity of 0.224 vs 0.165.The comparisons describe generated sequences; higher structural resemblance to PDB entries was reported across perplexity ranges.Source: Evaluation on Protein Sequence… · p. 17
  • Among 11,048 generated sequences with pLDDT>80, mean sequence similarity to natural proteins was 0.699 and mean TM-score was 0.864.Generated sequences span a broad spectrum of folds while maintaining low sequence identity, suggesting access to a larger sequence space in the protein manifold.Source: Evaluation on Protein Sequence… · p. 18
  • A one-step ReST process uses task predictors as reward models to evaluate SFT-generated sequences and select higher-quality sequences for self-training.The same filtering and fine-tuning settings are applied during the ReST stages.Source: Evaluation on Protein Sequence… · p. 19
  • Under identical SFT protocols, xTrimoPGLM sequences score higher on targeted properties than its untuned baseline and surpass PROGEN2 and ProtGPT2 on most tasks.The reported results also indicate alignment with minimal data or fewer tuning steps; the authors implement one-step ReST using task predictors as reward models to guide self-training.Source: Evaluation on Protein Sequence… · p. 19
  • ReST produces sequences that closely resemble natural datasets through an iterative process that selects sequences for further fine-tuning.The iterative procedure develops the Reinforcement Self-Training model.Source: Evaluation on Protein Sequence… · pp. 19–20
  • Evaluation examined relationships among xT-Fold pLDDT structural-prediction confidence, sequence identity, and PDB counterparts’ TM-scores.The analyses relate the scores and identities to structural counterparts in the Protein Data Bank.Source: Evaluation on Protein Sequence… · p. 20
  • After SFT alignment, xTrimoPGLM generates sequences embodying specific structural characteristics and aligning closely with established PDB entries.The model was fine-tuned on four SCOP fold-type sequences.Source: Evaluation on Protein Sequence… · p. 20

Discussion

The discussion presents xTrimoPGLM as a large-scale unified model with applications to structure prediction and sequence design, while emphasizing limits in computational accessibility, out-of-distribution performance, and generation reliability. It outlines efficiency, data-scaling, and retrieval-based directions for future work.

Discussion

  • Fine-tuning xTrimoPGLM-100B without quantization requires at least four A100 80G GPUs, limiting practical access.The discussion identifies high computational cost as a barrier to application.Source: Discussion · p. 20
  • Quantization (32), kernel fusion (28), and acceleration methods (33–36) could reduce resource needs, but few methods have been validated and real-world effectiveness remains uncertain.The authors suggest such techniques could enable training and deploying larger models with fewer computational resources.Source: Discussion · p. 20
  • Performance gains diminished as model parameters increased, and xT-Fold retained an OOD gap on CASP proteins even as pre-training models became more powerful.The authors speculate that protein sequences are less smooth in semantic space and more factual in nature.Source: Discussion · pp. 20–21
  • PLM-based methods still struggle against MSA-based or retrieval-augmented approaches on OOD data; the authors propose richer data and modalities, compute-optimal pre-training, and proportional data-model scaling.Suggested data sources include sequences, structures, and functional descriptions (37), with proteins, DNA, and RNA as additional modalities (2); compute-optimal pre-training is cited as (38).Source: Discussion · p. 21
  • Integrating MSA modules with neural network retrievers (39–41) could improve OOD generalization while accelerating processing compared with traditional MSA methods.The authors present this as a potential route to better end-to-end optimization and faster inference.Source: Discussion · p. 21
  • At low sampling temperatures, xTrimoPGLM can generate repetitive amino-acid fragments, including repeated alanine, that may receive high-confidence structures despite not occurring in nature.The discussion identifies hallucination as a generation limitation.Source: Discussion · p. 21
  • Approximately 20% of generated sequences had xT-Fold pLDDT < 70; after SFT, 17.8% to 88% found similar structures in the PDB, leaving fold stability uncertain.The authors suggest adding constraints during training and sampling (39) to reduce hallucinations.Source: Discussion · pp. 21–22
  • The authors conclude that unified, large-scale pre-training extends to downstream tasks including 3D structure prediction and enables protein sequence design through supervised fine-tuning.They describe the work as a stepping stone for future protein foundation-model research.Source: Discussion · p. 22

Methods

The Methods describe xTrimoPGLM’s unified bidirectional-input and autoregressive-generation approach, its staged combination of masked-token and span-generation objectives, and the evaluation and training procedures used to examine objective adaptation and stabilize 100B-scale pre-training.

Backbone Framework: General Language Model (GLM)

  • GLM combines bidirectional input processing with autoregressive blank infilling, providing a backbone for both understanding and generation.The framework reconstructs randomly blanked continuous spans sequentially and differs from causal decoder-only models, which use unidirectional attention.Source: Backbone Framework: General Language… · p. 29

Pre-Training Objectives

  • The unified objective jointly optimizes in-place MLM predictions for understanding and autoregressive GLM span reconstruction for generation, balanced by weighting factor α.The GLM objective randomly permutes masked spans, while α balances the two objectives in the unified loss.Source: Pre-Training Objectives · pp. 30–31

The Pre-Training Strategy

  • Training begins with MLM alone over about 400 billion tokens, masking 15% of input tokens and setting GLM loss to 0; a later 600-billion-token stage combines MLM and GLM, with MLM at 20%.Source: The Pre-Training Strategy · p. 31

Empirical Analysis of Unified Training

  • Validation losses show that MLM and GLM optimize simultaneously in alternating 50%-of-batch-time schedules, with objective shifts at 100B and 200B consumed tokens.The compatibility experiment uses xTrimoPGLM-150m, configured to align with xTrimoPGLM-100B hyperparameters.Source: Empirical Analysis of Unified Training · p. 33
  • MLM-adapted GLM matches the loss of the GLM from-scratch model with a 2.2× speedup (110B tokens); GLM-adapted MLM achieves a 2× speedup (100B tokens).Source: Empirical Analysis of Unified Training · p. 34

The Training Stability of Unified Training

  • A fixed GLM-loss ratio can trigger catastrophic collapse even at 1%, so the authors propose gradually increasing GLM loss before restoring the standard learning rate.In practice, the increase uses K = 1000 steps and a learning rate of 1e-7; the 100B run diverged only during the transition stage.Source: The Training Stability of Unified… · pp. 34–35

Pre-Training Configurations

  • xTrimoPGLM-100B was trained on 1 trillion UniRef90 and ColabFoldDB tokens using 96 DGX-A100 servers and 3D parallelism.The model has 72 transformer layers, 80 attention heads, 10,240 embedding dimensions, and 31,744 feed-forward embedding dimensions; training used FP16 for forward and backward passes.Source: Pre-Training Configurations · p. 35

Pre-Training Datasets

xTrimoPGLM’s pre-training data combines UniRef90 and ColabFoldDB, with filtering and weighted sampling yielding a final corpus of about 939M sequences and 200B tokens.

  • The dataset combines UniRef90, with approximately 153M initial entries, and ColabFoldDB, with approximately 950M entries and metagenomic sequences from diverse environmental niches.Source: Pre-Training Datasets · pp. 36–37

Training Set

  • Deduplication, short-sequence filtering, cross-referencing with UniRef90, and removal of sequences with over 80% single-amino-acid composition reduced ColabFoldDB to around 200M sequences.The initial 950M sequences were narrowed to 780M unique sequences before composition filtering.Source: Pre-Training Datasets · pp. 37–38
  • Approximately 939M protein sequences totaling 200B tokens comprise the final pre-training dataset.Source: Pre-Training Datasets · p. 38
  • ColabFoldDB is sampled at approximately 60%, weighting UniRef90 three- to four-fold more heavily to leverage its high-quality data.The paper states this weighting boosts the model’s fine-tuning capability with high-quality data.Source: Pre-Training Datasets · p. 38

Validation Set

  • Validation uses UniProt sequences released January 1–March 30, 2023, after filtering out sequences exceeding the stated 90% or 0.5% similarity thresholds against UniRef50 or training data.The candidate increment contained 18M sequences before filtering.Source: Pre-Training Datasets · p. 38

Data availability

All data used in the study are publicly available, with sources for pre-training, downstream-task, and structure-prediction datasets identified.

Data availability

  • All data used in the study are publicly available, and their uses are described in the methods.Source: Data availability · p. 39
  • The xTrimoPGLM-100B pre-training dataset uses UniRef90 downloaded from the version preceding December 2022 and ColabFoldDB.Source: Data availability · p. 39
  • All 18 downstream-task datasets are available online at proteinglm on Hugging Face.Source: Data availability · p. 39
  • Structure-prediction datasets come from the AlphaFold Database and PDB, with PDB release dates before May 2020.Source: Data availability · p. 39

Code availability

The xTrimoPGLM trained weights and downstream datasets are available from the Hugging Face repository. The paper also lists software used for training, data analysis, and structural work.

Code availability

  • Trained xTrimoPGLM weights and downstream datasets are available at https://huggingface.co/proteinglm.Source: Code availability · p. 39
  • Model training used DeepSpeed v0.6.1.Source: Code availability · p. 39
  • Data analysis used Python v3.8, NumPy v1.16.4, SciPy v1.2.1, seaborn v0.11.1, Matplotlib v3.3.4, and pandas v1.1.5.Source: Code availability · p. 39
  • TM-align v20190822 was used to compute TM-scores, structure visualizations were created in PyMOL v2.3.0, and protein structures were predicted using official AlphaFold2 implementations.Source: Code availability · p. 39

Supplementary materials … 1 Protein Understanding Benchmarks and Evaluations

This supplementary section evaluates xTrimoPGLM-100B on protein-understanding tasks spanning four domains, describing downstream evaluation protocols and comparisons with ESM2 in probing and LoRA modes. The results report performance trends and gains associated with model scale and/or LoRA, with task-dependent margins.

Additional Figures

  • Additional figures are listed on benchmark performance, xT-Fold structure prediction, structural diversification, sequence-generation alignment and control, and model scaling.The listed figures also include generated-sequence cases, perplexity comparisons, sequence identity and diversity, the pre-training dataset, and human protein sequence mapping.Source: Additional Figures · p. 60

1 Protein Understanding Benchmarks and Evaluations

  • The evaluation covers 18 downstream tasks across protein structure, developability, interactions, and functions.For datasets without established benchmarks, the authors use ESM2-15B with LoRA fine-tuning results.Source: §1 · p. 61
  • xTrimoPGLM-100B, combined with fine-tuning, significantly outperforms current SOTA approaches on most protein-related tasks.Figure 2 is identified as presenting protein-understanding benchmark performance.Source: Additional Figures · p. 60, §1 · p. 61
  • The benchmark suite spans classification and regression tasks, assessed with AUC, accuracy, and correlation metrics.Examples include cloning failure, metal-ion binding, enzyme catalytic efficiency, peptide-HLA/MHC affinity, and fitness prediction.Source: §1 · pp. 65–66
  • The zero-shot fitness evaluation used DMS_ProteinGym_substitutions, filtering sequences exceeding 1,024 residues and retaining 201 assays; performance was measured by Spearman correlation, averaged within each assay and then across assays.Source: §1 · pp. 67–68
  • Zero-shot mutation scores used masked-position log-odds for understanding models, an additive assumption for multiple mutations, and whole-sequence log-likelihood for generative models.The scoring procedures draw on ESM-1v and PROGEN2, respectively.Source: §1 · p. 68
  • Zero-shot fitness prediction showed no apparent scaling benefit, while SFT xT100B performed better than smaller models on supervised fitness prediction.The authors hypothesize that stronger differentiation of original and mutated sequences may not translate into accurate mutation scores without supervision.Source: Table 7
  • The benchmarks compared xTrimoPGLM-100B with ESM2 across downstream tasks using probing and LoRA fine-tuning modes.Metric values are reported for both modes, with underlining marking the best probing performance.Source: §1 · pp. 68–69
  • LoRA raises xTrimoPGLM-100B accuracy from 76.86 to 93.32 in the reported structure evaluation, alongside substantial gains from large pretrained models.The evaluated structure tasks span secondary structure, contact maps, and fold classification.Source: §1 · p. 69
  • Larger models consistently outperform smaller counterparts, with notably higher Spearman correlation on fitness and 10-class classification accuracy on localization.The comparison includes xTrimoPGLM-100B and ESM2-15B versus ESM2-150M.Source: §1 · p. 69
  • LoRA fine-tuning consistently outperforms probing, including on unseen sequence-pair patterns; its advantage is small on peptide-MHC and TCR-pMHC tasks, where baseline AUC exceeds 0.9.The small margin is attributed to high baseline performance on those interaction tasks.Source: §1 · p. 70
  • xTrimoPGLM-100B significantly outperforms ESM on solubility and stability, but temperature-task gains are marginal; both models exceed MCC 0.93 on Temperature Stability.xTrimoPGLM-100B is slightly better than ESM on the temperature-related datasets.Source: §1 · p. 70
  • Optimal Temperature has the smallest benchmark training set, approximately 1.7k samples, potentially constraining achievable model performance.Source: §1 · p. 70

2 xT-Fold Acceleration Methods

The xT-Fold model based on xTrimoPGLM-100B was accelerated through quantization and FlashAttention, improving training speed and enabling inference on longer sequences.

  • At least a 4x training speed increase on the same devices followed W8A16 and further 4-bit quantization, which reduced model communication time and memory consumption.Source: §2 · p. 71
  • FlashAttention (28) was added after inference overflowed on sequences exceeding 700 residues; operator fusion and softmax reduction enabled inference on longer sequences.Source: §2 · p. 71
  • At least ~5x GPU inference time reduction versus standalone AlphaFold2 was reported, with sequences up to 2,000 residues processed in ~200 seconds, excluding MSA feature processing.Source: §2 · p. 71

3 xT-Fold Training Settings

xT-Fold adapts an Evoformer-style structure module for protein structure prediction and trains it in two stages on experimental and distilled structures, with the pretrained protein language model frozen.

Training Settings

  • Training uses ~380K experimental structures from ~25k PDB clusters and ~4M high-confidence AlphaFoldDB structures, sampling predicted and real structures 75% and 25% of the time, respectively.The PDB release date is before May 2020, and the distilled structures have pLDDT greater than 90%.Source: §3 · p. 71

Training Procedure

  • The two-stage procedure uses crop sizes 256 then 384; stage 2 adds structural violation loss with coefficient 0.01 and cosine learning-rate decay from 5e-5.The training description specifies 78K stage-1 steps and 70K stage-2 steps; stage 2 warms up for 3000 steps and decays to 1e-5.Source: §3 · pp. 71–73

Model Architecture

  • Restricting pair-biased row attention to the query, removing column attention, and using a 48-layer block with up to three recycling stages significantly enhances training and inference speeds.Inference uses recyclenum = 3.Source: §3 · p. 72

Initialization

  • Initialization zeroes Evoformer output parameters to create a quasi-identity mapping from the PLM, dynamically selects Evoformer modules based on gradients, and freezes the PLM.The PLM is not fine-tuned during training.Source: §3 · p. 72

Data Distillation

  • Data distillation filters about 60 million AFDB samples using AF2 pLDDT 0.9 criteria, yielding about 4 million after clustering at sequence maximum identity 0.5.The procedure also filters samples similar to the test set by 90% and integrates 38K PDB samples in a 75:25 distill-to-PDB blend.Source: §3 · p. 72

4 AlphaFold2 with FlashAttention

This section adapts OpenFold’s AlphaFold2 inference to use FlashAttention, reporting lower runtime and memory use without changing numerical results, alongside similar scores across model runs.

4 AlphaFold2 with FlashAttention

  • The implementation loads DeepMind’s AlphaFold2 weights into OpenFold, whose predictions were confirmed to match DeepMind’s official inference results.Source: §4 · p. 73
  • FlashAttention significantly reduces AlphaFold2’s runtime and memory consumption by reducing HBM–GPU SRAM reads and writes, without changing numerical results.Source: §4 · p. 73
  • The evaluation used five AlphaFold2 checkpoints with five random seeds each, producing 25 models without AMBER relaxation for TM-score and GDT_TS assessment.Source: §4 · p. 73
  • CAMEO TM-score was 0.87 ~0.88 and CASP15 TM-score was 0.74 ~0.77 across the predicted models.The authors report little difference among the models.Source: §4 · p. 73
  • Because AlphaFold2 model_4 and model_5 use PDB templates while xT-Fold does not, model_3_0 was chosen as AlphaFold2’s representative result for a fairer comparison.Source: §4 · pp. 73–74

5 The SFT & ReST Pipelines

The section describes curated datasets and supervised fine-tuning for xTrimoPGLM-100B, then introduces one-step reinforcement self-training that uses task-predictor feedback to iteratively improve the model.

5.1 Dataset

  • The authors curated five datasets from the 18 benchmarks for supervised fine-tuning and reinforcement self-training.The datasets target elevated fluorescence, high-fitness mutations, low thermal stability, immunoglobulin folds, and nuclear localization.Source: §5.1 · p. 74

5.2 Supervised Fine-tuning.

  • xTrimoPGLM-100B was fine-tuned with the GPT objective using LoRA with rank 16 and α=32.For comparison, ProGen2-xlarge was trained at full scale with the GPT objective, while ProtGPT2 used its official Hugging Face fine-tuning script.Source: §5.2 · p. 75

5.3 Reinforcement Self-Training

  • One-step reinforcement self-training starts from the supervised fine-tuned model and uses task predictors as reward models to guide iterative improvement.Predictors score regression tasks by numeric features and classification tasks by the likelihood of the desired label; selected high-quality sequences are used for further fine-tuning until convergence.Source: §5.3 · p. 75
  • For reinforcement self-training, task-specific filtering selects generated sequences using predicted scores or label likelihoods before further fine-tuning.The selected sequences are used to fine-tune the checkpoint saved after supervised fine-tuning until convergence.Source: §5.3 · p. 76

6 Statistical Analysis of Generated Sequences

The section analyzes generated protein sequences statistically, comparing their diversity, disorder predictions, and predicted structures. It finds disorder predictions close to those for natural sequences, while sampling settings and n-gram constraints affect sequence and structural properties.

Statistical properties of the sampled artificial sequences

  • Higher temperature and nucleus sampling broadened sequence diversity but increased <eos> selection and could shorten sequences; low settings produced repetitive, abnormally long sequences.Repetitive tokens were observed below temperature 0.8 and nucleus sampling probability 0.7, motivating careful calibration of the two parameters.Source: §6 · pp. 76–77

Intrinsically unstructured/disordered proteins/domains

  • Generated sequences had disorder predictions close to the natural dataset across all four metrics, with the largest gap of 3.89% in LONG.The generated and natural datasets contained 6,523 and 10,000 sequences, respectively; the natural dataset was sampled uniformly from the original training data.Source: §6 · pp. 77–78

Trial & Error of Generated Structures with N-gram penalty

  • In structure-generation trials, low-complexity sequences with local repeats had predicted long loop-disorder regions, and the authors hypothesize that n-gram penalties impede grammatically correct sequence generation.The trials compared generation with an n-gram penalty of 3 against generation without the constraint.Source: §6 · p. 78

7 xTrimoPGLM-Ab: OAS Fine-tuning for Antibody

This section examines adapting xTrimoPGLM to antibodies through OAS fine-tuning, covering naturalness, antibody structure prediction, and sequence generation, including CDR3 redesign. It reports structure-prediction evaluations and analyzes sequence organization across generation strategies and edit conditions.

7 xTrimoPGLM-Ab: OAS Fine-tuning for Antibody

  • The antibody model starts from xTrimoPGLM-1B pretrained on general proteins for 500B tokens, then uses OAS antibody sequences with CDR masking, random-span masking, and MLM.For fine-tuning, 40% of samples mask one or two whole CDRs, 40% use random-span masking, and 20% use MLM; the OAS dataset comprises 1 billion antibody sequences.Source: §7 · p. 79

7.1 Zero-shot Naturalness

  • xTrimoPGLM-Ab-1B surpasses other baselines in two datasets; its GLM and MLM variants have similar Dataset 1 results, with little AUC difference on pair test, and additional training consistently improves Dataset 2 AUC by 0.02.Source: §7.1 · p. 81
  • Combining CDR-span and random-span tasks outperforms either task alone on both naturalness datasets, supporting their joint use in the GLM objective.The combined model used a 50% CDR-span and 50% random-span task mix.Source: §7.1 · pp. 81–82

7.2 xTrimoPGLM-AbFold: Antibody structure prediction

  • xTrimoPGLM-AbFold, based on xTrimoPGLM-Ab-1B, predicts antibody structures directly from amino acid sequences for single chains and VH-VL complexes.The approach uses the model as an encoder for structure prediction, avoiding the substantial computational resources and time required by MSA.Source: §7.2 · p. 82
  • xTrimoPGLM-AbFold significantly outperforms all other models, including state-of-the-art xTrimoAbFold, on every antibody structure-prediction metric.The result is achieved with a pre-trained antibody model fine-tuned with a single additional Evoformer block, without MSA or template information.Source: §7.2 · pp. 84–85
  • For VH-VL complexes, xTrimoPGLM-AbFold performs comparably to AlphaFold-Multimer without MSA or template information and is 6,300× faster than the original model.It is also 103× faster than the faster MSA-searching AlphaFold-Multimer.Source: §7.2 · p. 85
  • Increasing xTrimoPGLM-AbFold from 1 to 16 Evoformer blocks yields only marginal improvement, despite the 16-block model achieving the best performance on all metrics.The 16-block model retains a 2,400× speedup over original AlphaFold-Multimer and a 40× speedup over its accelerated version.Source: §7.2 · p. 85

7.3 Antibody-specific Generation using xTrimoPGLM-Ab.

  • The CDR3 redesign study used four strategies—partial masking, whole-region masking, random mutation, and random retrieval—on a SARS-CoV-2 wild-type-binding heavy-chain antibody.The model generated redesigns with xTrimoPGLM-Ab-1B, and xTrimoPGLM-AbFold was used to predict their structures.Source: §7.3 · pp. 85–86
  • Both CSM and CWM generate antibody sequences of varying lengths; CSM tends to yield sequences close to ground truth, so outputs with edit distance 2 or less were filtered.The experiments generated 6,000 antibodies using xTrimoPGLM-Ab-1B.Source: §7.3 · p. 87
  • CRM and CRR sequences display considerable disorder, whether edits involve few mutations or deletions or complete CDR3 replacement.The passage contrasts these parallel baselines with the preceding method, but the supplied fragments do not identify that method's result.Source: §7.3 · p. 87
  • As edit distance grows, the generated antibody CDR3 region tends to lose structural organization, indicating limitations in large generative models.The analysis identifies a relationship between edit distance and CDR3 structure.Source: §7.3 · p. 87

8 Model FLOPs Comparisons

This section compares computational resources used by pre-trained protein language models, deriving model and training-compute estimates from source-paper configurations and released resources. It describes how training tokens and FLOPs are estimated and what the FLOP counts include.

8 Model FLOPs Comparisons

  • Total training tokens are calculated as training steps multiplied by global batch size and sequence length, while parameters follow released implementations and hyperparameters, excluding tied weights and buffers.Source: §8 · pp. 87–88
  • Training compute is estimated as three times single-sample forward-pass FLOPs, multiplied by the number of training samples, calculated as total training tokens divided by sequence length.The factor of three represents one forward and one backward propagation without activation recomputation.Source: §8 · p. 88
  • FLOP counts include matrix multiplications in attention, feed-forward, vocabulary projection, and any language-model head, but exclude embeddings, element-wise addition, softmax, and layer normalization.Source: §8 · p. 88
  • Approximately 5.3e+21 FLOPs are estimated for ProtGPT2 over four days on 128 A100 GPUs.The estimate assumes each A100 performs 120 TFLOPs per second.Source: §8 · p. 88
  • Table 8 compares training compute across pre-trained protein language models using implementations and configurations from their source papers and accompanying resources.When a paper’s theoretical account differs from its practical implementation, the paper’s reported metrics are favored.Source: §8 · p. 87, Table 8 · p. 97
Loading 2401.06199v2…