Repurposing Protein Folding Models for Generation with Latent Diffusion

The 2024 Nobel Prize in Chemistry was awarded to AlphaFold2, marking broad recognition of the breakthrough in AI for protein structure prediction. But after predicting structures, the next key question is how to design new proteins. Existing generative models have several limitations: many models generate only backbone atoms and cannot directly produce all-atom structures; protein drugs need to be humanized to avoid immune rejection; and drug design also needs to satisfy complex constraints such as solubility. PLAID addresses these pain points by proposing a multimodal generative model that can simultaneously generate discrete protein sequences and continuous all-atom structural coordinates.

The core idea of PLAID is to learn the latent space of a protein folding model and train a diffusion model on this space. During training, only protein sequences are needed to obtain embeddings; during inference, after sampling from the latent space, the frozen folding model weights are used to decode sequences and structures. The paper uses ESMFold as the folding model, which is a successor to AlphaFold2 that replaces the retrieval step with a protein language model. This approach is similar to vision-language-action models in robotics, using prior knowledge from pretrained large models to assist downstream tasks.

Directly applying the above method encounters a challenge: the latent space of Transformer models such as ESMFold is very large and requires substantial regularization, making it as difficult to learn embeddings as high-resolution image synthesis. To this end, the authors propose CHEAP (Compressed Hourglass Embedding Adaptations of Proteins), which learns a compressed model of the joint embedding of sequences and structures. Through mechanistic interpretability analysis, they found that a "massive activation" phenomenon exists in the latent space, and that this space is highly compressible. CHEAP makes training an all-atom protein generative model feasible.

An important advantage of PLAID is that training requires only sequence data, not structural data. This is because sequence databases (such as UniProt) are 2–4 orders of magnitude larger than experimental structure databases, and obtaining sequences costs far less than experimentally resolving structures. By learning from a larger database, PLAID can capture a broader distribution of proteins. This property gives PLAID significant scalability in structure generation tasks where data is scarce.

PLAID supports controlling generation through compositional text prompts, and the current proof of concept covers two dimensions: function and organism. For example, through function keyword prompts, PLAID can generate metalloproteins with a tetrahedral cysteine-Fe²⁺/Fe³⁺ coordination pattern while maintaining high sequence diversity; through transmembrane protein keyword prompts, the generated proteins have hydrophobic residues in the core region, consistent with transmembrane characteristics. This control capability offers the possibility of designing proteins with specific functions.

Currently, PLAID's prompt control is limited to a single function or single organism dimension, and more complex constraints (such as solubility) have not yet been achieved. The generated proteins have not yet undergone large-scale wet-lab validation, and their activity and stability in real biological environments remain unknown. In addition, whether the low-dimensional manifold assumption of the latent space fully holds, and whether the semantic granularity of text prompts is sufficient to meet the needs of actual drug design, both await further study. The authors point out that this method can be extended to other domains where a predictor exists from an abundant modality to a scarce modality, such as protein-nucleic acid complexes, and they welcome collaboration for wet-lab testing.