PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Language Models Are Versatile Protein Learners

Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, Quanquan Gu

arXiv:2402.18567v2cs.LGq-bio.BM

TL;DR

Protein language models have faced a trade-off between sequence understanding and generation, while autoregressive models also have limited access to global sequence context. DPLM uses discrete diffusion pretraining and conditioning strategies for protein sequences. It supports both sequence generation and representation learning, with inverse-folding performance that can match or exceed strong baselines.

  • Problem

    Masked protein language models support sequence understanding but lack a well-defined formulation for generation, while autoregressive models often fall short in protein sequence understanding and use only one-sided context.

    Source: §1 · p. 1

  • Method

    DPLM learns to denoise protein sequences across noise levels and uses classifier guidance to steer discrete diffusion generation toward desired properties.

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

  • Results

    DPLM supports both protein sequence generation and representation learning, and inverse-folding sequences can outperform or match strong baselines including LM-DESIGN.

    Source: §5 · p. 9, §4.3.2 · p. 9

  • Takeaways & Limitations

    Plug-and-play guidance lets DPLM adapt generation to user preferences, including desired secondary structures.

    Source: §3.2 · p. 5, §4.3.3 · p. 9

    Scaffold pass rate and average success rate decline slightly when evaluated with ESMFold rather than OmegaFold.

    Source: §D.3 · p. 21

Abstract

from arXiv · show

This paper introduces diffusion protein language model (DPLM), a versatile protein language model that demonstrates strong generative and predictive capabilities for protein sequences. We first pre-train scalable DPLMs from evolutionary-scale protein sequences within a generative self-supervised discrete diffusion probabilistic framework, which generalizes language modeling for proteins in a principled way. After pre-training, DPLM exhibits the ability to generate structurally plausible, novel, and diverse protein sequences for unconditional generation. We further demonstrate the proposed diffusion generative pre-training makes DPLM possess a better understanding of proteins, making it a superior representation learner, which can be fine-tuned for various predictive tasks, comparing favorably to ESM2 (Lin et al., 2022). Moreover, DPLM can be tailored for various needs, which showcases its prowess of conditional generation in several ways: (1) conditioning on partial peptide sequences, e.g., generating scaffolds for functional motifs with high success rate; (2) incorporating other modalities as conditioner, e.g., structure-conditioned generation for inverse folding; and (3) steering sequence generation towards desired properties, e.g., satisfying specified secondary structures, through a plug-and-play classifier guidance. Code is released at \url{https://github.com/bytedance/dplm}.

1 Introduction

Protein language models have typically separated strong sequence prediction from sequence generation, while autoregressive models also face limits in modeling proteins’ global interactions. DPLM applies discrete diffusion pretraining to combine sequence generation and representation learning, with several forms of conditional generation.

1 Introduction

  • Masked language models support protein sequence understanding but lack a well-defined generative formulation, while autoregressive models often fall short in understanding protein sequences.Autoregressive models are also limited by unidirectional context for capturing global amino-acid interactions.Source: §1 · p. 1

B learned representation

  • DPLM uses discrete diffusion pretraining on evolutionary-scale protein sequences to provide a generative framework for protein language modeling.It denoises sequences across noise levels, from fully noisy to clean.Source: §1 · p. 1, B learned representation · p. 2
  • DPLM is pretrained at model sizes up to 3B and is designed for both protein sequence generation and representation learning.The paper reports unconditional generation, downstream prediction, and conditional generation capabilities.Source: B learned representation · pp. 2–3
  • DPLM supports conditional generation from partial sequences, other modalities such as structures, and desired properties using plug-and-play classifier guidance.Examples include motif scaffolding, inverse folding, and secondary-structure steering.Source: B learned representation · p. 3

2 Preliminaries

Protein language modeling uses different probabilistic approaches for sequence understanding and generation, with different strengths and constraints. Diffusion models instead perturb data through a forward process and learn to reverse that process through denoising.

2.1 Language Modeling for Protein

  • Masked language models use bidirectional context to predict masked residues, whereas autoregressive models factorize sequence probabilities and generate residues left to right.Masked models excel at sequence understanding but are difficult to apply to generation.Source: §2.1 · p. 3

2.2 Diffusion Probabilistic Models

  • Diffusion models perturb data through a forward Markov process and learn a backward denoising process that moves samples toward the data distribution.The forward process progresses toward a noise distribution over increasingly noisy steps.Source: §2.2 · p. 3
  • Generation begins with a sample from the noise distribution and proceeds through iterative denoising.Each step samples a less noisy state conditioned on the current state and the model’s prediction.Source: §2.2 · p. 3

3 DPLM: A Versatile Protein LM

DPLM uses discrete diffusion to model protein sequences, combining sequence generation with representations for prediction. It supports partial-sequence, cross-modal, and classifier-guided generation, and is presented as a more versatile alternative to EvoDiff.

3.1 Protein Language Modeling w/ Discrete Diffusion

  • Discrete diffusion operates directly on discrete sequence states, avoiding the mismatch of applying continuous Gaussian diffusion to discrete protein sequences.With absorbing diffusion, residues are either preserved or masked; the generalized learning objective includes masked and autoregressive language models as special cases.Source: §3 · p. 4, §3.1 · p. 4
  • Pre-trained on UniRef50’s around 45 million sequences (about 14 billion tokens), DPLM generates by iterative denoising and provides noise-free embeddings for predictive tasks.The embeddings are obtained by inputting the sequence at noise level 0.Source: §3.1 · p. 4

3.2 Conditioning

  • For partial-sequence conditioning, DPLM preserves observed residues; for cross-modal conditions, a cross-attention adapter is tuned while DPLM and the modality encoder remain frozen.The described cross-modal cases include inverse folding from a backbone and conditioning on small-molecule ligands.Source: §3.2 · p. 5
  • Discrete classifier guidance steers generation toward desired properties by approximating conditional reverse steps with a learned classifier or regressor.The approach uses a first-order Taylor approximation because the guidance probability does not factorize across sequence positions.Source: §3.2 · p. 5

3.3 Comparisons with The Most Related Work

  • Compared with EvoDiff, DPLM reports stronger representation learning and broader conditioning, using discrete diffusion rather than order-agnostic autoregressive diffusion.EvoDiff’s better performance relies on multiple-sequence-alignment-based MSA-Transformer parameterization.Source: §3.3 · pp. 5–6

4 Experiments

Across unconditional generation, downstream prediction, and conditional-generation tasks, DPLM performs strongly, with larger models generally attaining better results than smaller ones. The experiments assess sequence generation, representation learning, motif scaffolding, inverse folding, and secondary-structure-guided sampling.

4.1 Evaluation of Unconditional Generation

  • DPLM generates structurally plausible, novel, and diverse sequences, with the highest pLDDT across tested lengths and lower similarity to known PDB structures as sequences lengthen.Generated sequences maintain pLDDT > 75 as pdb-TM decreases beyond length 300; DPLM also achieves low average inner-TM scores.Source: §4.1 · pp. 6–7

4.2 Evaluation of Protein Representation Learning on Downstream Predictive Tasks

  • DPLM outperforms ESM2 across all evaluated predictive tasks and closely approaches structure-aware SaProt.The tasks include function, localization, annotation, protein-protein interaction, and secondary-structure prediction.Source: §4.2 · p. 7, Table 1 · p. 7

4.3.1 SEQUENCE-COND.: MOTIF-SCAFFOLDING

  • DPLM outperforms EvoDiff in solved motif-scaffolding problems and average success rate, while structure conditioning improves four problems but lowers performance on six others.Across shared solvable problems, DPLM has a higher success rate except for 3ixt; the reported scaffolding criterion requires motif-RMSD < 1 Å and pLDDT > 70.Source: §4.3.1 · p. 8, Figure 3 · p. 8

4.3.2 STRUCTURE-CONDITIONED: INVERSE FOLDING

  • On CATH 4.2 and 4.3, DPLM matches or outperforms strong baselines, including LM-DESIGN, with scTM = 0.85 and pLDDT > 76.DPLM uses a structural adapter and GVP-Transformer Encoder; its reported results use argmax decoding rather than sampling.Source: Table 2 · p. 8, §4.3.2 · p. 9

4.3.3 CONTROLLABLE GENERATION: SECONDARY STRUCTURE GUIDED PROTEIN SAMPLING

  • Discrete classifier guidance steers pre-trained DPLM toward sequences satisfying supplied secondary-structure annotations without retraining for each condition.The annotations are extracted from template natural proteins.Source: §4.3.3 · p. 9, Figure 4 · p. 9

5 Discussions

DPLM combines protein sequence generation and representation learning with several conditioning strategies. The discussion identifies broader conditioning, adaptation of language-model techniques, and integrating structure as future directions.

5 Discussions

  • DPLM supports both protein sequence generation and representation learning, with sequence, cross-modal, and plug-and-play classifier-guided conditioning strategies.Source: §5 · p. 9
  • Broader conditional-generation applications could use additional modalities, including MSA, small molecules, or antigens, and guidance for more user preferences.The paper names structural symmetry, superfamily, binding affinity, thermostability, and fluorescence as examples.Source: §5 · p. 9
  • The authors propose adapting long-context methods to handle proteins beyond DPLM’s training length limit, potentially extending modeling to DNA and RNA.Source: §5 · p. 9
  • Integrating protein structure with sequence in a next-generation DPLM is identified as future work.The paper points to structure prediction, design, and full-atom modeling advances as context for this direction.Source: §5 · pp. 9–10

Impact Statement

The impact statement describes potential applications in therapeutic and other computer-aided design, while the RDM section explains its transition reparameterization, training loss, and sampling rule.

Impact Statement

  • Protein generation and representation learning may support therapeutic macromolecule development and accelerate drug-discovery research.The authors also suggest possible adaptation to small-molecule, material, and chip design, while calling for responsible use.Source: Impact Statement · p. 10

A Reparameterizaed Discrete Diffusion Models (RDM)

B.1 Pre-training of DPLM

The section contrasts training DPLM from scratch with a two-stage masked-language-model initialization, describing instability in scratch training and measures that can improve it.

B.1 Pre-training of DPLM

  • Scratch training often showed loss spikes that hurt performance, while absorbing diffusion exposes the model to masking ratios ranging from 0% to 100%.Source: §B.1 · pp. 17–18
  • The proposed two-stage approach initializes DPLM from a pretrained masked language model and then continues training with the discrete-diffusion objective.This approach also avoids gradual masking warmup.Source: §B.1 · pp. 17–18
  • The authors suggest that gradual masking warmup could ease early training by limiting masking before increasing it toward the full diffusion objective.They describe this as a proposed principle and leave direct exploration of the strategy for future work.Source: §B.1 · p. 18
  • With gradient-norm clipping and 300k training steps, scratch-pretrained DPLM can achieve performance comparable to the two-stage approach.Source: §B.1 · p. 18

B.2 Pre-training of AR-LM baseline

The appendix describes an autoregressive baseline and motivates absorbing discrete diffusion over alternative diffusion choices. It reports better unconditional sampling for absorbing diffusion than multinomial diffusion across lengths, and says the discrete approach excels over a latent-diffusion comparison.

B.2 Pre-training of AR-LM baseline

  • The autoregressive baseline uses DPLM’s architecture with causal attention and next-token prediction, but forcing a predefined sample length significantly reduced sequence quality.Source: §B.2 · p. 18

C Reasons for choosing absorbing discrete diffusion

  • Absorbing diffusion extends masked-language modeling, whereas multinomial diffusion resembles denoising autoencoders with less established evidence for scaling as a sequence-learning objective.Source: §C.1 · p. 19

C.1 Regarding discrete diffusion (DD) definitions (multinomial vs absorbing as proposed in D3PM)

  • Absorbing-DD generally performs better than multinomial/uniform-DD in unconditional sampling across different sequence lengths.Source: Figure 5 · p. 19, §C.1 · p. 19

C.2 Regarding latent diffusion

  • Latent diffusion on protein embeddings requires a lossy continuous relaxation, while the paper reports that its discrete-diffusion approach still excels on unconditional-sampling pLDDT.Source: §C.2 · p. 19

D Additional Experimental Details

The top-k denoising sampler can collapse to repetitive sequences, so the authors perturb prediction scores with the Gumbel-Max trick. This reduces trivial samples and improves diversity while retaining the categorical sampling distribution.

D.1 A modified unconditional sampling strategy

  • Without conditioning, top-k unmasking can favor frequent amino acids and spread them through the sequence, producing trivial single-amino-acid repetitions.The authors attribute this collapse to high-confidence frequent tokens causing other amino acids to be dropped at later steps.Source: §D.1 · p. 19
  • The Gumbel-Max trick perturbs prediction scores with Gumbel noise while sampling from the original categorical distribution.The maximum of gi + log pi selects the sampled category, but differs from the original log-probability score.Source: §D.1 · p. 19
  • The modified scores reduce trivial repetitive cases and improve sequence diversity.Source: §D.1 · p. 19

D.2 Delve deeply into the pLDDT score of unconditional sampling

The authors examine why unconditional DPLM samples have high pLDDT and whether that reflects mode collapse. They report similar pseudo-perplexity for proteins with high and low pLDDT, consistent with learning diverse structural patterns.

D.2 Delve deeply into the pLDDT score of unconditional sampling

  • DPLM unconditional samples have higher pLDDT than UniRef50, whose average is lowered by less-plausible repetitive sequences; PDB pLDDT is similar to DPLM.The passage gives ADADAD...ADADAD as an example of a UniRef50 sequence with pLDDT 35.14.Source: §D.2 · p. 20
  • Pseudo-perplexity is similar for less-structural proteins (pLDDT < 50) and structural proteins (pLDDT > 70), suggesting DPLM learns diverse structural patterns.Source: §D.2 · p. 20
  • The authors hypothesize that learning first captures more-structured proteins before long-tailed, less-structural patterns, leaving this learning dynamic for future study.Source: §D.2 · p. 20

D.3 Sequence-conditional generation: motif-scaffolding

These sections report motif-scaffolding results, describe adapter-based inverse folding and classifier-free guidance, and review related language-modeling and protein-design approaches.

D.3 Sequence-conditional generation: motif-scaffolding

  • 12 of 17 motif-scaffolding problems were solved; overall success rate was 0.27 for sequence-only and 0.31 for structure-conditioned sampling.ESMFold evaluation slightly lowers overall pass rate and average success rate compared with OmegaFold.Source: §D.3 · pp. 20–21, Table 6 · p. 21

D.4 Structure-conditional generation: inverse folding

  • DPLM's structural adapter adds encoded 3D structure to its sequence decoder; draft-sequence noising mitigates exposure bias and boosts inverse-folding performance.Inference uses greedy deterministic decoding; the unconditional and adapter-conditioned models also enable classifier-free guidance.Source: §D.4 · p. 21

D.5 Classifier-free guidance

  • Classifier-free guidance combines conditional and unconditional models, and the paper reports improved structure-conditioned generation with guidance.For η > 1, guidance prioritizes the conditional model while discouraging samples from moving away from the unconditional distribution.Source: §D.5 · p. 22, Figure 6 · p. 22

E.1 Language Models

  • Language-model research spans autoregressive and non-autoregressive generation, scaling, instruction tuning, and emergent abilities.Source: §E.1 · p. 22

E.2 Protein Language Models

  • Protein language models trained on sequences have shown downstream generalization and can capture structural information and support protein-design tasks.Source: §E.2 · pp. 22–23

E.3 Diffusion Language Models

  • Discrete diffusion models operate on discrete tokens, whereas continuous approaches use continuous surrogates; the former are presented as suited to language data.Source: §E.3 · p. 23

E.4 Protein Structure Diffusion Models

  • Protein-structure diffusion methods generate structures using representations such as coordinates, inter-residue angles, residue frames, or distance-and-angle matrices.The reviewed approaches include designs for functional structures, novel proteins, and large proteins or complexes.Source: §E.4 · p. 23

E.5 Protein Inverse Folding

  • Inverse folding generates sequences conditioned on protein structures, commonly encoded as residue graphs, and limited structure data is a key challenge.The review notes data augmentation and pretrained protein language models as approaches used in this area.Source: §E.5 · pp. 23–24

F Visualization of Unconditional Samples

The figures visualize unconditional protein samples across two length ranges, from 50 to 1000 residues.

F Visualization of Unconditional Samples

  • Figure 7 visualizes unconditional samples from 50 to 500 in length.Source: Figure 7 · p. 25
  • Figure 8 visualizes unconditional samples from 600 to 1000 in length.Source: Figure 8 · p. 25
Loading 2402.18567v2…