Official Implemetation of DPLM (ICML'24) - Diffusion Language Models Are Versatile Protein Learners

The repository is an official implementation of ICML24 paper Diffusion Language Models Are Versatile Protein Learners, which introduces diffusion protein language model (DPLM), a versatile protein language model that demonstrates strong generative and predictive capabilities for protein sequences. Specifically, DPLM exhibits impressive performance in protein sequence generation, motif scaffolding, inverse folding, and representation learning.

We develop DPLM based on the ByProt. This repository contains pretraining scripts of DPLM, and running scripts of various protein generation and understanding tasks as below:

• Unconditional protein sequence sampling: DPLM is capable of generating protein sequences with reasonable predicted structures unconditionally.
• Sequence-conditioned generation: motif scaffolding: DPLM can generate reasonable scaffold sequences given the specific functional motifs.
• Structure-conditioned generation: inverse folding: DPLM can produces sequences that can confidently fold into the given backbone structure.
• Representation learning: DPLM is a superior protein sequence representation learner, demonstrating impressive performance across a variety of protein predictive tasks.
• Controllable generation: DPLM enjoys plug-and-play programmability, generating samples satisfying provided secondary structure annotations.

If you have any questions, please feel free to raise an issue or contact us at wangxinyou@smail.nju.edu.cn. Thank you!

News

• 2024/10/18: check out our new work DPLM-2, a multimodal protein foundation model that extends DPLM to simultaneously model, understand, and generate both sequences and structures!

Installation

# clone project
git clone --recursive https://url/to/this/repo/dplm.git
cd DPLM

# create conda virtual environment
env_name=dplm

conda create -n ${env_name} python=3.9 pip
conda activate ${env_name}

# automatically install everything else
bash install.sh

Pretraining

Data

Download the preprocessed UniRef50 dataset

We pretrain DPLM on the UniRef50 dataset, which contains about 42 million protein sequences. We obtain the preprocessed UniRef50 dataset provided by EvoDiff (Alamdari et al, 2023).

bash scripts/download_uniref50.sh

Training

We train DPLM with approximate 1 million tokens per batch and 100,000 training steps. The training script is as follows:

• The following script is runned on one node with 8 A100 GPU. If you want to train on multi nodes, you can adjust the total number of tokens by ensuring that the max_tokens * accumulate_grad_batches*GPU_number is approximately 1 million.

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

max_tokens=8192
accumulate_grad_batches=16
# this means the effective batch size is GPU_number(8) * max_tokens(8192) * accumulate_grad_batcher(16), resulting in approximately 1 million.

exp=lm/dplm_650m
model_name=dplm_650m

python train.py \
    experiment=${exp} name=${model_name} \
    datamodule.max_tokens=${max_tokens} \
    trainer.accumulate_grad_batches=${accumulate_grad_batches}

You can adjust the other training configurations in the configs/experiment/lm/dplm_650m.yaml as needed.

Model checkpoints

Access pretrained models in varying sizes:

Users can load DPLM checkpoint by:

from byprot.models.lm.dplm import DiffusionProteinLanguageModel

model_name = "airkingbd/dplm_650m"
dplm = DiffusionProteinLanguageModel.from_pretrained(model_name)

Unconditional protein sequence generation

The results of unconditional protein sequence generation of DPLM of different scales (150M, 650M, 3B) are shown in the table below. For more details, please refer to our paper.

To generate new protein sequences using a pre-trained DPLM model:

model_name=dplm_650m # choose from dplm_150m, dplm_650m, dplm_3b
output_dir=generation-results/${model_name}

mkdir -p generation-results

python generate.py --model_name ${model_name} \
    --seq_lens 100 200 300 400 500 \
    --saveto ${output_dir}

# Evaluation
bash anylasis/plddt_compute.sh ${output_dir} # compute pLDDT using ESMFold

We also provide evaluation scripts in the analysis folder. Users can use the analysis/uncond_analysis.ipynb to obtain average pLDDT score of each length and draw the line chart of the pLDDT score.

Sequence-conditioned generation: motif scaffolding

We examine 17 motif-scaffolding problems, and for each problem, we sample 100 sequences and then calculate the success rate according to the motif-RMSD < 1$\AA$ and pLDDT > 70. Success rate of each problem is shown below.

We provide the following script to sample sequences for each motif-scaffolding problem. Note that before generation, you should download all the motif pdb files which are provided by EvoDiff, and place them in the data-bin/scaffolding-pdbs folder.

export CUDA_VISIBLE_DEVICES=0

model_name=dplm_650m
output_dir=./generation-results/${model_name}_scaffold

mkdir -p generation-results

# Generate scaffold 
python scaffold_generate.py \
    --model_name airkingbd/${model_name} \
    --num_seqs 100 \
    --saveto $output_dir

# Predict structure by ESMFold
max_tokens=1024
pdb_path=$output_dir/scaffold_fasta/esmfold_pdb

# folding
mkdir -p $pdb_path

echo 'folding by ESMFold'
output_filename_list=$(ls ${output_dir}/scaffold_fasta)
echo $output_filename_list

python analysis/cal_plddt_dir.py -i ${output_dir}/scaffold_fasta -o ${pdb_path} --max-tokens-per-batch ${max_tokens}

For evaluation, users can use the analysis/motif_analysis.ipynb to obtain success rate of each problem.

Structure-conditioned generation: inverse folding

The partial results on the CATH 4.3 dataset are shown in the table below. For more details, please refer to our paper.

Data

Download the preproceesd CATH datasets

• CATH 4.2 dataset provided by Generative Models for Graph-Based Protein Design (Ingraham et al, NeurIPS'19)
• CATH 4.3 dataset provided by Learning inverse folding from millions of predicted structures (Hsu et al, ICML'22)

bash scripts/download_cath.sh

Training

We train structure-conditional DPLM based on the LM-Design framework, designating the pre-trained protein language model as DPLM. The training script is as below.

exp=lm/cond_dplm_650m
dataset=cath_4.3
name=${dataset}/cond_dplm_650m

python train.py \
    experiment=${exp} datamodule=${dataset} name=${name} \
    logger=tensorboard trainer=ddp_fp16 

Evaluation on valid/test datasets

Users can set the eval_sc to true to calculate the self-consistency TMscore and pLDDT, which will result in a significant evaluation time overhead.

dataset=cath_4.3
exp_path=${dataset}/cond_dplm_650m
eval_sc=false 
# if set ${eval_sc} to true, the program will calculate the self-consistency 
# TMscore and pLDDT during generation, 
# thus siginificantly increase the evaluation time.

python test.py \                                                                 
    experiment_path=${exp_path} \
    data_split=test ckpt_path=best.ckpt mode=predict \
    task.generator.max_iter=100 task.generator.eval_sc=${eval_sc}

Representation Learning

DPLM excels in various protein prediction tasks, demonstrating its versatility and effectiveness. The following table summarizes its performance:

The training and evaluation pipeline is based on the Saprot repository, and we slightly modify the code to support DPLM. Users can select the "representationlearning" branch for the evaluation of protein predictive tasks.

Acknowledgements

DPLM extends its gratitude to the following projects and individuals.

We draw inspiration and leverages/modifies implementations from:

• microsoft/evodiff for the preprocessed UniRef50 dataset, sequence sampling evaluation implementation and data pipeline.

• westlake-repl/SaProt for the representation learning evaluation pipeline.

• jingraham/neurips19-graph-protein-design for the preprocessed CATH dataset.

• facebook/esm for their ESM implementations and pretrained model weights.

• jasonkyuyim/se3_diffusion for their self-consistency structural evaluation implementation.

We express our sincere appreciation to the authors of these repositories for their invaluable contributions to the development of DPLM.

Citation

@inproceedings{wang2024diffusion,
  title={Diffusion Language Models Are Versatile Protein Learners},
  author={Wang, Xinyou and Zheng, Zaixiang and Ye, Fei and Xue, Dongyu and Huang, Shujian and Gu, Quanquan},
  booktitle={International Conference on Machine Learning},
  year={2024}
}