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
TL;DR
Existing protein language models are typically specialized for either understanding or generation, leaving unified handling of both objectives an open challenge. xTrimoPGLM addresses this gap with a unified, extremely large-scale pre-training framework, and the authors report strong downstream structure-prediction results while identifying practical limitations for deployment.
Problem
Existing protein language models are designed around specific pre-training frameworks, while integrating understanding and generation remains an open question.
Method
xTrimoPGLM unifies understanding and generation pre-training and scales the model to 100 billion parameters.
Results
The experiments suggest that scaling extends to downstream tasks, including key 3D structure prediction, and xTrimoPGLM-100B surpassed previous state-of-the-art methods in 15 of 18 tasks.
Takeaways & Limitations
The study supports unified, large-scale protein language modeling as a framework for protein understanding and generation.
Takeaways & Limitations
Effective real-world drug-design deployment still requires broader task adaptation, more accurate structure prediction, and reduced generative protein hallucinations.
Abstract
from arXiv · showhide
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
xTrimoPGLM unifies protein understanding and generation objectives in a foundation model scaled to 100 billion parameters. It improves benchmark performance, structure prediction, sequence generation, and scaling behavior across protein tasks.
- Method: The model combines autoencoding and autoregressive objectives to support protein understanding and generation within one unified framework.Its design uses a GLM backbone with bidirectional attention and an autoregressive objective.
- Results: xTrimoPGLM-100B surpassed previous state-of-the-art methods in 15 of 18 protein tasks spanning structure, interactions, functionality, and developability.The evaluation used linear probing and advanced fine-tuning techniques.
- Structure prediction: xTrimoPGLM-based xT-Fold improved CASP15 TM-score by 5 points over ESMFold while providing faster inference across scenarios.The model is presented as a high-performance 3D structural prediction system.
- Sequence generation: xTrimoPGLM generates diverse de novo sequences with 11.7% median sequence identity and supports property-directed generation after supervised fine-tuning.The reported programmable targets include structural and biophysical properties.
- Limitations: The authors identify adaptation to diverse protein tasks, improved structure accuracy, and reduced generative hallucinations as remaining requirements for real-world drug design.These limitations are described as necessary to bridge theoretical capability and practical application.
- Scaling: Scaling from smaller models to xTrimoPGLM-100B followed the anticipated power-law relationship between computational resources and training loss.Models from 1 million to 1 billion parameters were trained with 100 billion tokens, and their loss trajectory aligned with the 100-billion-parameter model.
Discussion
The discussion highlights xTrimoPGLM’s capabilities in downstream structure prediction and sequence design, alongside limitations in computational cost, out-of-distribution generalization, and generated-sequence reliability. It proposes expanded data modalities, efficient scaling, retrieval or MSA integration, and constrained generation as directions for improvement.
- Computational limitations: The 100B-parameter model has high computational cost, requiring at least four A100 80G GPUs, which may limit accessibility.The authors suggest quantization, kernel fusion, and related compression methods to reduce training and deployment resource requirements.
- Future directions: The authors propose richer sequence, structure, and functional data, multimodal modeling, compute-optimal training, MSA or retrieval integration, and constrained sampling to improve efficiency and OOD generalization.They also note that increasing model parameters can yield diminishing performance returns and that only a few efficiency methods have been validated so far.
- Structure prediction: xT-Fold shows strong TM-scores and inference speed, but its performance gap on out-of-distribution proteins persists as the pre-training model becomes more powerful.The discussion identifies CASP proteins as an example of challenging OOD data and notes that PLM-based methods still struggle against MSA-based or retrieval-augmented approaches.
- Sequence generation: Generated sequences can contain repetitive amino-acid fragments, and approximately 20% were not confidently predicted to fold by xT-Fold.After supervised fine-tuning, only 17.8% to 88% of generated sequences found similar structures in the Protein Data Bank.
- Capabilities: xTrimoPGLM extends unified understanding-and-generation pre-training to downstream protein tasks, including 3D structure prediction and supervised sequence design.The authors describe the model as a foundation for structure prediction and programmable protein sequence generation.
Competing interests
The paper declares no competing interests; the supplied passages otherwise consist of figure and metric descriptions rather than section-specific content.
- The authors declare no competing interests.
Methods
xTrimoPGLM combines bidirectional masked-token understanding with autoregressive span generation in a unified GLM-based training framework. Joint optimization and staged training are used to support both capabilities while addressing training stability.
- Pre-Training Objectives: GLM reconstructs masked consecutive spans autoregressively after span permutation, while MLM predicts masked tokens in place.The framework includes short-blank span prediction, long-text generation, and masked language modeling.
- Unified Pre-Training: xTrimoPGLM jointly optimizes MLM understanding and GLM generation objectives to support both protein-sequence representation and generation.The unified loss balances the two objectives, with α controlling their relative contribution.
- Objective Compatibility: The two objectives optimize simultaneously, while MLM-adapted GLM reaches comparable loss with a 2.2× speedup and GLM-adapted MLM achieves a 2× speedup.
Pre-Training Datasets
The pre-training corpus combines UniRef90 and ColabFoldDB, expanding coverage with metagenomic sequences and yielding approximately 939 million protein sequences and 200 billion tokens after filtering.
- Data Sources: The corpus combines UniRef90 with ColabFoldDB, whose metagenomic sequences span diverse environmental niches and support discovery of novel proteins.
- Training Set: The final pre-training dataset contains approximately 939M protein sequences with 200B tokens after deduplication, filtering, and composition-based refinement.The initial ColabFoldDB collection contained about 950M sequences, and roughly 780M unique sequences remained after early filtering and cross-deduplication.
- Dataset Composition: UniRef90 contributes about 156M sequences and 53B residue tokens, while ColabFoldDB comprises clustered and member collections totaling about 783M sequences.The reported ColabFoldDB components include approximately 208M clustered sequences and 575M member sequences.
- Sampling Strategy: UniRef90 is sampled with greater weight, producing an approximately 60% ColabFoldDB sampling ratio and increasing the contribution of higher-quality data.
- Validation Set: The validation set uses UniProt sequences released from January 1 through March 30, 2023, after similarity-based filtering against reference databases and training data.
Data availability
The study uses publicly available data sources for pre-training, downstream evaluation, and structure prediction, with dataset versions and release-date boundaries specified.
- Public Data Sources: The pre-training corpus is publicly sourced from UniRef90 and ColabFoldDB, with the UniRef90 version preceding December 2022 used.
- Downstream Evaluation: All 18 downstream task datasets are available online through the stated protein-language-model repository.
- Structure Prediction: Structure-prediction datasets come from the AlphaFold Database and the PDB, with the cited PDB release date restricted to before May 2020.
Code availability
The paper provides trained xTrimoPGLM weights and downstream datasets, alongside code and software used for modeling, analysis, visualization, and evaluation.
- Trained xTrimoPGLM weights and downstream datasets are made available.
- The implementation uses DeepSpeed v0.6.1 for model training and AlphaFold2 official implementations for protein 3D structure prediction.
- Analysis and evaluation rely on Python, NumPy, SciPy, seaborn, Matplotlib, pandas, and TM-align, with PyMOL used for structure visualization.
Supplementary materials
The supplementary materials document benchmark coverage, task definitions, datasets, evaluation settings, comparative results, and additional figures and tables.
- Additional supplementary figures and tables report structural prediction, sequence generation, scaling, benchmark tasks, compute comparisons, and model configurations.
- The supplementary materials cover 18 downstream tasks spanning protein structure, developability, interactions, and functions.
- xTrimoPGLM-100B significantly outperforms current state-of-the-art approaches in most protein-related tasks when combined with fine-tuning.
- Supplementary sections define evaluation tasks including contact maps, fold classification, secondary structure, and solubility, with their datasets and metrics.
- Contact Map: LoRA improves xTrimoPGLM-100B contact-map accuracy from 76.86 to 93.32.
- The model significantly outperforms ESM models on solubility and stability, while gains on temperature-related tasks remain marginal.
2 xT-Fold Acceleration Methods
The xT-Fold acceleration methods combine quantization and FlashAttention to reduce training and inference costs while supporting longer protein sequences.
- Quantizing the model to 4-bit reduces communication time and memory consumption, producing at least a 4x training-speed increase on the same devices.
- Longer-sequence inference initially caused memory overflow beyond length 700, motivating the incorporation of FlashAttention.
- Inference with FlashAttention supports sequences up to length 2000 in approximately 200 seconds, achieving at least a ~5x GPU inference-time reduction versus standalone AlphaFold2 excluding MSA processing.
3 xT-Fold Training Settings
xT-Fold is trained on distilled and experimentally determined structures using a two-stage procedure, a refined Evoformer-style architecture, and a frozen protein language model.
- xT-Fold is trained on approximately 380K experimentally determined structures from about 25K PDB clusters, with training and inference speed enhancements from architectural adjustments.
- The structure module adapts AF2’s Evoformer with query-only MSARowAttentionWithPairBias, no MSAColumnAttention, 48 layers, and up to three recycling stages.
- The protein language model remains frozen during training, while Evoformer outputs are initialized as a residual quasi-identity mapping.
- Training data combines AlphaFoldDB-derived predictions with real PDB structures, sampled in a 75:25 predicted-to-real ratio after filtering and clustering.
- Training uses two stages, beginning with 10 million samples at crop size 256 and then increasing the crop size to 384 with additional structural losses.
4 AlphaFold2 with FlashAttention
The evaluation uses OpenFold’s AlphaFold2 implementation with DeepMind weights and FlashAttention, then compares multiple AlphaFold2 runs against xT-Fold under differing template inputs.
- Implementation: FlashAttention significantly reduces original Self-Attention’s time and memory consumption without changing numerical results.It reduces reads and writes between high-bandwidth memory and GPU on-chip SRAM.
- Evaluation setup: Twenty-five AlphaFold2 models were generated from five checkpoints and five random seeds without AMBER relaxation.TM-score and GDT_TS score were calculated for all models.
- Comparison protocol: Model_4 and model_5 use PDB templates, whereas xT-Fold does not, so model_3_0 is selected for fairer comparison.The representative AlphaFold2 result is the first run of model_3.
5 The SFT & ReST Pipelines
The pipeline fine-tunes xTrimoPGLM-100B on curated task-specific sequences, then applies predictor-guided filtering for one-step reinforcement self-training.
- Datasets: Five benchmark datasets are used for supervised fine-tuning and reinforcement self-training.The curated tasks target fluorescence, fitness, temperature stability, fold, and localization properties.
- Datasets: Task-specific datasets range from 981 training sequences for immunoglobulins to 141,598 for low-thermal-stability proteins.Other datasets include nucleus localization and strong fluorescence intensity targets.
- Supervised fine-tuning: xTrimoPGLM-100B is trained with a GPT objective, while comparator models use their respective GPT or fine-tuning procedures.The xTrimoPGLM setup uses task-dependent gradient accumulation, whereas ProtGPT2 uses LoRA fine-tuning.
- ReST: The ReST pipeline uses task predictors as reward models to filter generated sequences and fine-tune the supervised checkpoint.Predictors score regression outputs or desired-class likelihoods before selected sequences are reused for training.
- Filtering: Generated sequences are filtered using task-specific thresholds or top-ranked predicted scores before iterative fine-tuning.Examples include fluorescence scores above 2.5, the top 300 fitness sequences, and the top 200 instability sequences.
6 Statistical Analysis of Generated Sequences
The study evaluates generated sequences through sampling diversity, disorder predictions, and structural analyses, finding natural-like disorder tendencies while identifying sampling-related failure modes.
- Sampling analysis: Higher temperature and nucleus-sampling factors broaden sequence diversity, but higher factors also increase <eos> selection and can shorten sequences.The study samples combinations of five temperatures and four nucleus probabilities.
- Sampling limitations: Temperatures below 0.8 and nucleus factors below 0.7 produce repetitive tokens and abnormal long sequences.The authors therefore recommend calibrating temperature and nucleus sampling together.
- Disorder statistics: Generated sequences closely match natural sequences across four disorder metrics, with the largest gap being 3.89% in LONG.The authors report comparable minimal, maximal, and structured predicted-disorder tendencies.
- Structural analysis: Removing the n-gram penalty tends to improve generated structural patterns relative to low-complexity sequences with long loop-disorder regions.The authors hypothesize that the penalty can impede grammatically correct sequence generation.
7 xTrimoPGLM-Ab: OAS Fine-tuning for Antibody
The antibody extension fine-tunes xTrimoPGLM on OAS data and develops antibody-specific models for naturalness, expression, structure prediction, and sequence generation.
- Antibody structure prediction: xTrimoPGLM-AbFold outperforms all compared models on every antibody-structure metric, including the existing xTrimoAbFold state-of-the-art model.It uses a single additional Evoformer block and operates without MSA or template information.
- Antibody structure prediction: xTrimoPGLM-AbFold achieves 6,300× and 103× speedups over original and faster MSA-searching AlphaFold-Multimer, respectively.With 16 Evoformer blocks, it retains best performance while achieving 2,400× and 40× speedups over those baselines.
- Antibody generation: CDR3 short masking generates antibodies closely resembling ground truth, while CSM and CWM generate varied lengths without mutations or deletions.Parallel CRM and CRR baselines display considerable disorder, and larger edit distances tend to degrade CDR3 organization.
8 Model FLOPs Comparisons
The paper compares computational resources across pretrained protein language models using parameter counts, training tokens, and estimated FLOPs, alongside downstream performance and structural benchmarks.
- Comparison methodology: The comparison derives model parameters from released implementations and hyperparameters, while training tokens are computed from steps, global batch size, and sequence length.When theoretical and practical configurations differ, the analysis favors metrics reported in the original paper.
- Compute estimation: Training compute is estimated from per-sample forward FLOPs, a 1F1B multiplier, the number of pre-training samples, and selected matrix multiplications.Embedding, element-wise addition, softmax, and layer-normalization operations are excluded from the FLOP count.
- Compute estimation: The accounting includes attention projections and products, feed-forward transformations, vocabulary projections, and language-model-head transformations.For hidden states of size (B, L, D) multiplied by a (D, 4D) weight matrix, the stated FLOPs expression is BLD4D2.
- Compute comparison: 128 A100 GPUs would achieve approximately 5.3e+21 FLOPs over four days under the stated ProtGPT2 throughput assumption.The estimate assumes each A100 GPU performs 120 TFLOPs per second.