Protenix-Mini+: efficient structure prediction model with scalable pairformer

TL;DR

This work targets the bottlenecks of AF3-style structure prediction—namely high inference latency and cubic scaling with token count—by introducing Protenix-Mini+, a lightweight variant that combines pruning of redundant blocks, scalable Pairformer components with triangular linear attention and chunked updates, and a few-step ODE diffusion sampler. The approach achieves large efficiency gains (over 90% faster in some settings) with modest accuracy losses (roughly a 3% drop in intra-protein LDDT for challenging cases), enabling scalable predictions for long biomolecular sequences and complexes. Extensive experiments across RecentPDB, PoseBusters, and AF3AB demonstrate competitive performance relative to larger models and highlight practical gains in inference time, especially when longer sequences or multi-seed sampling are involved. Overall, Protenix-Mini+ advances deployment-ready, diffusion-based structure prediction by delivering a favorable accuracy-efficiency balance and outlining pathways for hardware-aware acceleration and broader applicability to antibody–antigen and multi-sequence contexts.

Abstract

Lightweight inference is critical for biomolecular structure prediction and downstream tasks, enabling efficient real-world deployment and inference-time scaling for large-scale applications. While AF3 and its variants (e.g., Protenix, Chai-1) have advanced structure prediction results, they suffer from critical limitations: high inference latency and cubic time complexity with respect to token count, both of which restrict scalability for large biomolecular complexes. To address the core challenge of balancing model efficiency and prediction accuracy, we introduce three key innovations: (1) compressing non-scalable operations to mitigate cubic time complexity, (2) removing redundant blocks across modules to reduce unnecessary overhead, and (3) adopting a few-step sampler for the atom diffusion module to accelerate inference. Building on these design principles, we develop Protenix-Mini+, a highly lightweight and scalable variant of the Protenix model. Within an acceptable range of performance degradation, it substantially improves computational efficiency. For example, in the case of low-homology single-chain proteins, Protenix-Mini+ experiences an intra-protein LDDT drop of approximately 3% relative to the full Protenix model -- an acceptable performance trade-off given its substantially 90%+ improved computational efficiency.

Figures (4)

• Figure 1: An overview of our proposed architecture. (a) Overview of the AF3-style model architecture. The model consists of three stages: ❶ a feature aggregation stage that collects atomic representations; ❷ a token-level refinement stage that iteratively enhances representations with the deep pairformer and MSA to create a rich representation for diffusion conditioning; and ❸ a diffusion stage that transforms randomly initialized diffusion noise into the final atomic structure. (b) By eliminating stacked blocks, we reduce the model size. (c) Triangle linear attention replaces the cubic triangle attention with a kernelized formulation that projects queries and keys into feature maps, enabling efficient updates. (d) Chain-aware chunking reduces the cubic cost of triangular updates by partitioning sequences into blocks and mean-pooling within each block, yielding compressed representations that enable efficient chunk triangle updates.
• Figure 2: Comparison of computational cost (TeraFLOPs per Pairformer block) versus sequence length for different pairformer variants.
• Figure 3: Runtime and efficiency comparison across Protenix model variants. (a) Pairformer module runtime on GPUs with compute capability comparable to NVIDIA H100. Protenix scales cubically with sequence length, while lighter variants achieve reduced inference time. (b) Pairformer computational cost in PFLOPs. (c) Full model sampling time (minutes), showing substantial speedups for Mini, Mini+, and Tiny models from fewer blocks, scalable operations, and fewer diffusion steps.
• Figure 4: Normalized Pade approximation across 16 pairformer blocks.