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

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.

💡 Research Summary

The paper “SaDiT: Efficient Protein Backbone Design via Latent Structural Tokenization and Diffusion Transformers” introduces a novel framework that significantly accelerates the process of de novo protein backbone generation while maintaining high structural fidelity. The core challenge addressed is the computational intensity of existing diffusion-based protein design models (e.g., RFDiffusion, FrameDiff), which operate directly on high-dimensional coordinate or frame spaces, requiring hundreds of iterative denoising steps.

SaDiT’s innovation lies in shifting the generative process from this continuous coordinate space to a compressed, discrete latent space. It achieves this through three key components:

1. Structural Tokenization via SaProt: The framework utilizes a pre-trained SaProt encoder to map a protein’s continuous backbone geometry into a sequence of discrete “structural tokens.” This acts as a bottleneck, regularizing the generative search space to physically plausible local geometries and providing a compressed representation.
2. Diffusion Transformer (DiT) Backbone: Instead of conventional U-Nets or graph networks, SaDiT employs a Diffusion Transformer as its generative model. The Transformer architecture is well-suited for capturing the long-range dependencies inherent in protein folding. The diffusion process is applied directly to the embeddings of these structural tokens, and Adaptive Layer Normalization (adaLN) conditions the network on the diffusion timestep.
3. IPA Token Cache Mechanism: To address the computational cost of the crucial SE(3)-equivariant Invariant Point Attention (IPA) layers, SaDiT introduces a novel caching mechanism. During the iterative reverse diffusion sampling, as the structure converges and token states become stable, the model caches and reuses the computed key/value states from IPA. This drastically reduces redundant computations in later sampling steps, leading to near-linear scaling of inference time.

The paper provides a theoretical proof that the entire SaDiT pipeline (Encoder -> DiT -> Decoder) is SE(3)-equivariant. This is essential because it ensures the generated protein structures are invariant to global rotations and translations, a fundamental requirement for biological validity. The encoder uses only relative distances and orientations, making the latent tokens SE(3)-invariant. The decoder predicts local relative frame transformations, which are then assembled into a global structure, guaranteeing equivariant output.

Extensive experiments demonstrate SaDiT’s superiority over state-of-the-art baselines like RFDiffusion and Proteína. Evaluations on unconditional backbone generation and fold-class conditional generation tasks show that SaDiT achieves a substantial speedup (approximately 2-5x faster sampling) on identical hardware. Crucially, it matches or exceeds the performance of these baselines in terms of “designability”—the likelihood that a generated backbone can be filled with a stable amino acid sequence. The model particularly excels at capturing complex topological features when conditioned on specific fold classes.

In summary, SaDiT represents a paradigm shift in efficient protein design by successfully combining latent structural compression with a powerful transformer-based diffusion model and an inference-time optimization strategy. It opens the door for large-scale structural exploration and real-time protein design applications that were previously hindered by computational bottlenecks.