SaDiT: Efficient Protein Backbone Design via Latent Structural Tokenization and Diffusion Transformers

TL;DR

SaDiT is a novel protein backbone design framework that combines structural tokenization and diffusion transformers to achieve faster generation while maintaining structural quality, outperforming existing models in speed and viability.

Contribution

SaDiT introduces a discrete latent space for protein geometry and an IPA Token Cache to accelerate diffusion-based backbone design, a novel approach in the field.

Findings

SaDiT significantly outperforms RFDiffusion and Proteina in speed and structural quality.

SaDiT effectively captures complex topological features in protein structures.

The model demonstrates high designability in fold-class conditional generation tasks.

Abstract

Generative models for de novo protein backbone design have achieved remarkable success in creating novel protein structures. However, these diffusion-based approaches remain computationally intensive and slower than desired for large-scale structural exploration. While recent efforts like Proteina have introduced flow-matching to improve sampling efficiency, the potential of tokenization for structural compression and acceleration remains largely unexplored in the protein domain. In this work, we present SaDiT, a novel framework that accelerates protein backbone generation by integrating SaProt Tokenization with a Diffusion Transformer (DiT) architecture. SaDiT leverages a discrete latent space to represent protein geometry, significantly reducing the complexity of the generation process while maintaining theoretical SE(3) equivalence. To further enhance efficiency, we introduce an IPA Token Cache mechanism that optimizes the Invariant Point Attention (IPA) layers by reusing computed token states during iterative sampling. Experimental results demonstrate that SaDiT outperforms state-of-the-art models, including RFDiffusion and Proteina, in both computational speed and structural viability. We evaluate our model across unconditional backbone generation and fold-class conditional generation tasks, where SaDiT shows superior ability to capture complex topological features with high designability.