The Family of Diffusion Protein Language Models (DPLM)

Overview 🌟

This repository contains the official implementation of training and inference as well as the pre-trained weights for the Family of Diffusion Protein Language Models (DPLM), including: - DPLM from ICML'24 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. - DPLM-2 from ICLR'25 paper "DPLM-2: A Multimodal Diffusion Protein Language Model", a multimodal protein foundation model that extends discrete diffusion protein language model to accommodate both sequences and structures. - ICML'25 spotlight paper "Elucidating the Design Space of Multimodal Protein Language Models", where we elucidate the challenges of structure modeling of multimodal protein language models (e.g., DPLM-2 and ESM3) and propose advanced designs for better structure modeling. We have released the finer-grained bit-based generative modeling (DPLM-2 Bit). The full implementation of the paper will be released soon.

Key Features 🔑

Specifically, the DPLM family exhibits impressive performance in protein (structure and sequence) co-generation, any-to-any conditional generation (e.g., folding, inverse folding, and motif scaffolding), and representation learning. We develop DPLM based on the ByProt. This repository contains pretraining scripts for DPLM and running scripts for various protein generation and understanding tasks, as detailed below: - Unconditional protein generation: DPLM is capable of unconditionally generating protein sequences with reasonable predicted structures. DPLM-2 can generate diverse and highly plausible proteins through simultaneous structure-sequence co-generation. - Sequence-conditioned generation (forward folding): DPLM-2 can generate reasonable protein structure given the input protein sequence, achieving close performance with the strong folding model (e.g., ESMFold). - Structure-conditioned generation (inverse folding): DPLM and DPLM-2 can produce sequences that can confidently fold into the given backbone structure. - Motif scaffolding: DPLM can generate reasonable scaffold sequences given specific functional motifs. DPLM-2 achieves more successful motif scaffolding through multimodal motif conditioning. - Representation learning: DPLM is a superior protein sequence representation learner, while DPLM-2 offers structure-aware protein represenrations, 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.

TODOs

• [ ] Controllable/guided generation with discrete diffusion classifier guidance.
• [ ] Representation learning of DPLM-2

DPLM

"Diffusion Language Models Are Versatile Protein Learners." Wang et al., In ICML 2024

DPLM-2

"DPLM-2: A Multimodal Diffusion Protein Language Model." Wang et al., In ICLR 2025

Updates 📢

• ​​[2025-07]​ We update the default sampling strategy of DPLM-2 to annealing@2.0:0.1.
• ​​[2025-04]​​ Our latest work DPLM-2.1, which focuses on analysis and better protein structure modeling of multimodal protein language models, is accepted to ICML'25 Spotlight! Check Elucidating the Design Space of Multimodal Protein Language Models. We have release the implementation of finer-grained and better structure modeling (DPLM-2 Bit). The full implementation will be released soon.
• ​​[2024-10]​​ 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!
• ​​[2024-03]​​ We release DPLM, a versatile protein language model that demonstrates strong generative and predictive capabilities for protein sequences!

Table of Contents 📚

• Quick Start
• Installation
• Load Pretrained Models
• Generation Examples
• Model Checkpoints
• Advanced Usage
• Training
• Unconditional protein (co-)generation
  • Protein sequence generation (DPLM)
  • Protein sequence-structure co-generation (DPLM-2 & DPLM-2-Bit)
• Sequence-conditioned Generation: Forward Folding
• Structure-conditioned generation: inverse folding
• Motif scaffolding
• Representation Learning
• Acknowledgements
• Citation

Quick Start

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 scripts/install.sh

Load Pretrained Models

Users can load DPLM/DPLM-2 checkpoint by:

from byprot.models.dplm import DiffusionProteinLanguageModel as DPLM
from byprot.models.dplm2 import MultimodalDiffusionProteinLanguageModel as DPLM2
from byprot.models.dplm2 import DPLM2Bit

dplm = DPLM.from_pretrained("airkingbd/dplm_650m").cuda()
dplm2 = DPLM2.from_pretrained("airkingbd/dplm2_650m").cuda()
dplm2_bit = DPLM2Bit.from_pretrained("airkingbd/dplm2_bit_650m").cuda()

Generation Examples

Protein sequence generation

from generate_dplm import initialize_generation

input_tokens = initialize_generation(
  length=200,
  num_seqs=5,
  tokenizer=dplm.tokenizer,
  device=next(dplm.parameters()).device
)
samples = dplm.generate(
  input_tokens=input_tokens,
  max_iter=500,
)
print([''.join(seq.split(' ')) for seq in dplm.tokenizer.batch_decode(samples, skip_special_tokens=True)])

Protein sequence-structure co-generation

User can check the generated sequence and structure in the ./generation-results folder.

from generate_dplm2 import initialize_generation, save_results

input_tokens = initialize_generation(
  task="co_generation",
  length=200,
  num_seqs=5,
  tokenizer=dplm2.tokenizer,
  device=next(dplm2.parameters()).device
)[0]

samples = dplm2.generate(
  input_tokens=input_tokens,
  max_iter=500,
)
save_results(
    outputs=samples,
    task="co_generation",
    save_dir="./generation-results/dplm2_generation",
    tokenizer=dplm2.tokenizer,
    struct_tokenizer=dplm2.struct_tokenizer, save_pdb=True
)

samples = dplm2_bit.generate(
  input_tokens=input_tokens,
  max_iter=500,
)
save_results(
    outputs=samples,
    task="co_generation",
    save_dir="./generation-results/dplm2_bit_generation",
    tokenizer=dplm2_bit.tokenizer,
    struct_tokenizer=dplm2_bit.struct_tokenizer
)

Model Checkpoints

Access pretrained models in varying sizes:

Advanced Usage

Training

DPLM

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), which can be downloaded from this link. After downloading, please place the dataset in the ./data-bin/uniref50 folder.

We also provide the preprocessed dataset in HuggingFace datasets format, which we recommend to use. User can download the HF dataset locally in advance for faster loading by:

bash scripts/download_uniref50_hf.sh

Example of training

We train DPLM with approximately 1 million tokens per batch for 100,000 training steps.

The following command is run on one node with 8 A100 GPUs. If you want to train on multiple nodes, you can adjust the total number of tokens by ensuring that max_tokens * accumulate_grad_batches*#GPUs 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 #GPUs(8) * max_tokens(8192) * accumulate_grad_batches(16), resulting in approximately 1 million.

exp=dplm/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/dplm/dplm_650m.yaml as needed.

DPLM-2

Dataset

We use the experimental structures from PDB and AF2-predicted structures from SwissProt dataset as training data for DPLM-2. We provide a preprocessed HuggingFace dataset of PDB and SwissProt. User can download the HF dataset locally in advance for faster loading by:

bash scripts/download_pdb_swissprot.sh

Example of training

As noted in section 3.2 in DPLM-2 paper, we propose an efficient warm-up training strategy to mitigate the scarcity of structure training data. During training, we initialize the DPLM-2 model with pretrained DPLM checkpoint, to leverage the evolutionary knowledge captured by sequence-based pLM during large-scale sequence pretraining, which is beneficial for structure modeling.

We train DPLM-2 with approximately 64,000 tokens per batch for 100,000 training steps. To preserve the evolutionary knowledge captured by DPLM, we use the LoRA to prevent large parameter shifts. The training command is as follows:

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

max_tokens=8192
accumulate_grad_batches=1
# this means the effective batch size is #GPUs(8) * max_tokens(8192) * accumulate_grad_batches(1), resulting in approximately 64 thousand.

exp=dplm2/dplm2_650m
model_name=dplm2_650m

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

DPLM-2 Bit-based Modeling

In our latest work DPLM-2.1, we show that the index-based structure token is challenging for the model to predict. A finer-grained, bit-based modeling approach in the latent space (i.e., predicting each bit of the quantized structure feature instead of the index) leads to better structural modeling and generation performance.

The training dataset is same to DPLM-2, and the training command is as below:

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

max_tokens=8192
accumulate_grad_batches=1
# this means the effective batch size is #GPU(8) * max_tokens(8192) * accumulate_grad_batches(1), resulting in approximately 64 thousand.

exp=dplm2/dplm2_bit_650m
model_name=dplm2_bit_650m

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

Unconditional protein (co-)generation

Protein sequence generation (DPLM)

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.

| Length | 100 | 200 | 300 | 400 | 500 | 600 | 700 | 800 | 900 | 1000 | | ------ | ------------- | ------------- | ------------- | ------------- | ------------- | --------