ProteinAE: Protein Diffusion Autoencoders for Structure Encoding
This work addresses the problem of protein structure representation for researchers in computational biology, offering a streamlined approach that simplifies optimization and improves generative modeling, though it is incremental in advancing existing autoencoder and diffusion methods.
The authors tackled the challenge of representing protein structures for generative modeling by introducing ProteinAE, a diffusion autoencoder that maps protein backbone coordinates into a continuous latent space, achieving state-of-the-art reconstruction quality and enabling efficient, high-quality structure generation that outperforms prior latent-based methods.
Developing effective representations of protein structures is essential for advancing protein science, particularly for protein generative modeling. Current approaches often grapple with the complexities of the SE(3) manifold, rely on discrete tokenization, or the need for multiple training objectives, all of which can hinder the model optimization and generalization. We introduce ProteinAE, a novel and streamlined protein diffusion autoencoder designed to overcome these challenges by directly mapping protein backbone coordinates from E(3) into a continuous, compact latent space. ProteinAE employs a non-equivariant Diffusion Transformer with a bottleneck design for efficient compression and is trained end-to-end with a single flow matching objective, substantially simplifying the optimization pipeline. We demonstrate that ProteinAE achieves state-of-the-art reconstruction quality, outperforming existing autoencoders. The resulting latent space serves as a powerful foundation for a latent diffusion model that bypasses the need for explicit equivariance. This enables efficient, high-quality structure generation that is competitive with leading structure-based approaches and significantly outperforms prior latent-based methods. Code is available at https://github.com/OnlyLoveKFC/ProteinAE_v1.