ProteinAE

ProteinAE is a protein structure autoencoder that compresses backbone coordinates into a compact, continuous latent space, providing a foundation for latent-space generative protein design. Diffusion and flow-matching models that operate directly on 3D coordinates have driven recent progress in structure generation, but working in raw coordinate space is expensive and awkward. Learning a good latent representation of structure lets a generator operate in a smaller, smoother space, mirroring the latent-diffusion strategy that transformed image generation.

The core of ProteinAE is a non-equivariant Diffusion Transformer trained end-to-end with a single flow-matching objective. It maps protein backbone coordinates from E(3) into a continuous latent representation with an aggressive bottleneck, then reconstructs the structure with high fidelity, avoiding the quantization losses of discrete tokenizers such as VQ-VAE-based approaches. Building on this encoder, the authors train a protein latent diffusion model (PLDM) that generates new structures by sampling in the learned latent space.

ProteinAE was introduced in October 2025 by Shaoning Li, Le Zhuo, Yusong Wang, Mingyu Li, Xinheng He, Fandi Wu, Hongsheng Li, and Pheng-Ann Heng, a collaboration led from The Chinese University of Hong Kong with Tencent, Shanghai Jiao Tong University, and additional partners. It extends NVIDIA's Proteina codebase from Cα-only to full-backbone-atom modeling, is described in an arXiv preprint accepted to ICLR 2026, and is released under the MIT license.

#Key Features

• Continuous latent encoding: ProteinAE maps backbone coordinates into a continuous latent space, sidestepping the quantization error of discrete structure tokenizers and enabling high-fidelity reconstruction.
• Non-equivariant Diffusion Transformer: The autoencoder is a scalable, non-equivariant Diffusion Transformer trained end-to-end with a single flow-matching objective rather than a multi-term loss.
• Aggressive bottleneck: Backbone geometry is compressed from a 256-dimension representation down to an 8-dimension latent, yielding a compact code for downstream generation.
• Latent diffusion generator (PLDM): A 200M-parameter protein latent diffusion model generates new structures by sampling in the learned space, applying the latent-diffusion recipe to proteins.
• Strong reconstruction and generation: The model reports substantially lower reconstruction error than prior latent methods and high designability for unconditional generation.

#Technical Details

ProteinAE comes in Base (~20M parameters) and Large (~100M parameters) variants, with the PLDM generator at 200M parameters. The autoencoder downsamples the per-residue representation from dimension 256 to a latent dimension of 8 while keeping sequence length fixed (downsampling ratio r=1). It is trained on AFDB-FS, a filtered subset of the AlphaFold Protein Structure Database comprising 588,318 structures of length 32-256 residues after sequence and structure clustering. On CASP15 TS-domains, ProteinAE reconstructs backbones to a Cα RMSD of 0.28 +/- 0.20 Angstrom, compared with 1.23 Angstrom for an ESM3 VQ-VAE and 1.15 Angstrom for ProToken. For unconditional generation, the ProteinAE-PLDM system reaches 93% designability with high structural diversity and competitive novelty, matching leading coordinate-space methods while substantially outperforming prior latent-based generators.

#Applications

ProteinAE is aimed at researchers building generative protein-design pipelines who want to work in a compact latent space rather than directly on 3D coordinates. The autoencoder provides a reusable structural representation for downstream tasks, while the PLDM demonstrates de novo backbone generation. The released code, a public checkpoint on HuggingFace, and a from-pretrained workflow let groups reconstruct structures, sample new backbones, or plug the latent space into their own conditional generators and property predictors.

#Impact

ProteinAE ports the latent-diffusion paradigm that reshaped image generation into protein structure modeling, showing that a continuous, heavily compressed latent can preserve near-atomic reconstruction fidelity while supporting high-quality generation. By avoiding discrete tokenization and its quantization error, it offers an alternative to VQ-based structure tokenizers and a compact substrate for future conditional design models. Accepted to ICLR 2026, its evaluations are computational; experimental validation of generated designs remains future work.