[Paper Review] A Latent Diffusion Model for Protein Structure Generation

Cong Fu, Keqiang Yan+8 more|arXiv (Cornell University)|May 6, 2023

Protein Structure and Dynamics9 citations

TL;DR

The paper introduces LatentDiff, a latent diffusion framework with an equivariant protein autoencoder to generate novel protein backbone structures in a reduced latent space, improving designability and sampling efficiency.

ABSTRACT

Proteins are complex biomolecules that perform a variety of crucial functions within living organisms. Designing and generating novel proteins can pave the way for many future synthetic biology applications, including drug discovery. However, it remains a challenging computational task due to the large modeling space of protein structures. In this study, we propose a latent diffusion model that can reduce the complexity of protein modeling while flexibly capturing the distribution of natural protein structures in a condensed latent space. Specifically, we propose an equivariant protein autoencoder that embeds proteins into a latent space and then uses an equivariant diffusion model to learn the distribution of the latent protein representations. Experimental results demonstrate that our method can effectively generate novel protein backbone structures with high designability and efficiency. The code will be made publicly available at https://github.com/divelab/AIRS/tree/main/OpenProt/LatentDiff

Motivation & Objective

• Motivate reducing the diffusion modeling space for 3D protein backbones to improve efficiency.
• Propose an equivariant autoencoder that preserves sequence and 3D geometry to produce latent representations.
• Develop a rotation-invariant latent diffusion process operating on latent protein representations.
• Demonstrate reconstruction quality, designability, and parallel sampling efficiency against baselines.

Proposed method

• Design an SE(3) equivariant protein autoencoder that down-/upsamples protein backbones while preserving C-alpha order.
• Construct a graph-based latent space with structure padding, edge building, graph expansion, and SE(3) equivariant message passing.
• Train a latent diffusion model in the autoencoder latent space with rotation-equivariant reverse diffusion.
• Enforce SE(3) invariance through zero-mean latent representations and rotation-equivariant diffusion (EDM-based) in latent space.
• Use a two-stage generation: encode to latent, perform diffusion to sample latent, then decode to 3D backbone.

Experimental results

Research questions

• RQ1 Can LatentDiff generate novel protein backbone structures with good designability?
• RQ2 Does an SE(3) equivariant autoencoder effectively compress protein backbones for diffusion modeling?
• RQ3 How does latent-space diffusion compare to full-space diffusion in terms of designability and sampling efficiency?
• RQ4 Is the latent diffusion process rotation-invariant and suitable for 3D protein geometry?

Key findings

• Autoencoder with downsampling factor 4 achieves balanced reconstruction (RMSD 1.2755 Å, Augment Acc 100%, Residue Acc 98%, Edge Stable 70.99%, Torsion MAE 0.8951 rad).
• LatentDiff designability: 66.9% of generated backbones have scTM > 0.5, outperforming FoldingDiff (14.2%) and ProtDiff (11.8%).
• LatentDiff achieves competitive designability while enabling faster sampling compared to some full-atom diffusion baselines (e.g., LatentDiff-P is more efficient).
• For autoencoder variants, smaller downsampling (factor 2) yields better reconstruction but larger latent space; factor 4 provides a balance between reconstruction and efficiency.
• The method supports parallel sampling efficiency, reducing diffusion steps relative to some baselines while maintaining designability.