[Paper Review] Token-Level Guided Discrete Diffusion for Membrane Protein Design

Shrey Goel, Schray, Peregrine M.+4 more|arXiv (Cornell University)|Oct 22, 2024

RNA and protein synthesis mechanismsBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

MeMDLM introduces the first masked diffusion protein language model (MDLM) for de novo membrane protein design, leveraging the ESM-2 pLM to generate sequences with high transmembrane (TM) character. By applying discrete noise via mask tokens and reverse diffusion, MeMDLM outperforms autoregressive methods in generating biologically realistic membrane proteins while preserving key physicochemical properties.

ABSTRACT

Reparameterized diffusion models (RDMs) have recently matched autoregressive methods in protein generation, motivating their use for challenging tasks such as designing membrane proteins, which possess interleaved soluble and transmembrane (TM) regions. We introduce the Membrane Diffusion Language Model (MemDLM), a fine-tuned RDM-based protein language model that enables controllable membrane protein sequence design. MemDLM-generated sequences recapitulate the TM residue density and structural features of natural membrane proteins, achieving comparable biological plausibility and outperforming state-of-the-art diffusion baselines in motif scaffolding tasks by producing lower perplexity, higher BLOSUM-62 scores, and improved pLDDT confidence. To enhance controllability, we develop Per-Token Guidance (PET), a novel classifier-guided sampling strategy that selectively solubilizes residues while preserving conserved TM domains, yielding sequences with reduced TM density but intact functional cores. Importantly, MemDLM designs validated in TOXCAT beta-lactamase growth assays demonstrate successful TM insertion, distinguishing high-quality generated sequences from poor ones. Together, our framework establishes the first experimentally-validated diffusion-based model for rational membrane protein generation, integrating de novo design, motif scaffolding, and targeted property optimization.

Motivation & Objective

• To develop a generative framework for de novo membrane protein design that overcomes limitations of autoregressive models.
• To adapt masked diffusion language modeling (MDLM) to protein language models (pLMs) for improved sequence generation.
• To preserve critical physicochemical and structural features of membrane proteins during generation.
• To enable scaffolding of transmembrane and soluble motifs with high biological fidelity.
• To provide a sequence-first platform for experimental applications in drug discovery and synthetic biology.

Proposed method

• Pre-train the ESM-2-150M pLM on a comprehensive protein sequence space to establish a rich embedding space.
• Fine-tune the pre-trained ESM-2 on membrane protein sequences to specialize in membrane-specific features.
• Implement a discrete noise schedule that replaces amino acid tokens with <mask> tokens over time steps, following an absorbing-state diffusion process.
• Use a time-conditioned noise schedule αₜ to control the probability of masking tokens at each timestep t.
• Reverse the diffusion process by denoising from a fully masked sequence to generate complete, realistic protein sequences.
• Apply the model to both unconditional generation and conditional scaffolding of transmembrane and soluble domains.

Experimental results

Research questions

• RQ1 Can masked diffusion language modeling (MDLM) be effectively adapted to protein language models (pLMs) for membrane protein sequence generation?
• RQ2 Does MDLM-based generation outperform autoregressive methods in generating membrane proteins with high transmembrane (TM) character?
• RQ3 Can MeMDLM accurately reconstruct functional motifs (TM and soluble domains) with higher fidelity than state-of-the-art inpainting methods?
• RQ4 To what extent does MeMDLM preserve key physicochemical properties such as solubility and membrane localization?
• RQ5 Can MeMDLM generate biologically plausible membrane protein sequences suitable for downstream experimental validation?

Key findings

• MeMDLM-generated sequences exhibit significantly greater transmembrane (TM) character compared to autoregressive methods, indicating improved structural realism.
• MeMDLM achieves a pseudo-perplexity of 3.819 on transmembrane domains and 7.029 on soluble domains, outperforming EvoDiff (20.554 and 16.991, respectively).
• Cosine similarity between ESM-2-650M embeddings of true and reconstructed sequences reaches 0.768 for TM domains and 0.778 for soluble domains, demonstrating high biological similarity.
• MeMDLM embeddings achieve AUROC scores of 0.949 for solubility and 0.541 for membrane localization, closely matching SOTA pLMs.
• The model successfully generates diverse, novel membrane protein sequences with realistic structural and functional features suitable for experimental applications.
• MeMDLM demonstrates that MDLM can be effectively applied to encoder-only pLMs to unlock powerful generative capabilities for protein design.