Protenix

A Reproduction of AlphaFold3

This blog post summarizes the research paper "Protenix: A Reproduction of AlphaFold3" by ByteDance's AML AI4Science team. The paper introduces Protenix, a fully open-source implementation of AlphaFold3, aiming to enhance accessibility and improve upon the original model's performance and methodology.

Introduction

AlphaFold3 (AF3) has significantly advanced protein structure prediction. However, limited accessibility (lack of code and ambiguities in the original paper) hinders broader adoption. Protenix directly addresses this by providing fully open-sourced code, model weights, and pre-processed data. The key contributions of Protenix are:

• Improved Model Performance: Protenix achieves comparable or better performance than AF3, AlphaFold-Multimer 2.3 (AF2.3), and RoseTTAFold2NA (RF2NA) across various benchmarks, including protein-ligand, protein-protein, protein-nucleic acid, and RNA structure prediction. The authors benchmark against these models but cannot compare against HelixFold3 and Chai-1 due to commercial restrictions.

• Refined Methodology: Protenix meticulously implements the AF3 methodology, clarifying ambiguous steps and correcting typographical errors. This detailed reproduction effort enhances understanding and aids further advancements in the field.

• Enhanced Accessibility: Protenix offers open-sourced code, model weights, inference code, and trainable code. This significantly improves accessibility for researchers, promoting broader adoption and collaborative development.

The following table compares the accessibility features of Protenix with related projects:

Results

Protenix demonstrates strong performance across diverse benchmarks:

2.1 Overview: Protenix uses a training dataset curated from the PDB (cutoff date: September 30, 2021) and predicted structures from AlphaFold2 and OpenFold. The evaluation datasets are diverse.

2.2 Ligands (PoseBusters V2): Protenix outperforms AF3-2019 in terms of RMSD success rate, achieving higher rates for both RMSD and PB-Valid metrics. Analysis shows Protenix's superior performance stems from utilizing a more recent training dataset (cutoff 2021-09-30 vs. 2019-09-30 for AF3) and not solely data leakage.

2.3 Proteins (AF2.3 Comparison): Protenix achieves comparable or better performance than AF2.3 in protein-protein interaction prediction, using DockQ scores, across different interface types.

2.4 Nucleic Acids (CASP15 and other datasets): Protenix demonstrates performance comparable to AF3 and superior to RF2NA in predicting RNA and DNA structures.

2.5 Discussion: The authors discuss confidence score modifications, improvements to the interface prediction and the limitations of the memorization effect, highlighting areas for future work. Analysis of correlation between predicted chain pair ipTM and DockQ score shows a positive correlation (p > 0.5).

Figure 1 in the paper displays visual representations of Protenix's performance across several benchmarks. Figures 2-4 present in-depth analyses of PoseBusters results and case studies that show Protenix's ability to predict novel and challenging protein-ligand interactions. Figures 6 and 7 highlight Protenix's capabilities in RNA and DNA structure prediction. Figure 8 displays the relationship between the predicted scores and interface DockQ scores for several interface types. Finally, Figure 9 illustrates how similar structures in the training set can influence the model’s accuracy.

Methods

The paper details modifications to the data pipeline, model, and training process compared to the AF3 implementation.

3.1 Data Pipeline: Key differences include modifications to the parser, MSA approach (no MSA for nucleic chains), featurization, and cropping method. Table 2 in the paper details specific corrections to the algorithms presented in the original AF3 paper.

3.2 Model and Training: Protenix uses a similar model architecture to AF3 but introduces several optimizations, including parameter initialization strategies, scalability improvements (BF16 mixed precision training, custom CUDA kernels, and DS4Sci_EvoformerAttention), and training modifications. These optimizations significantly reduce training time and memory usage. Table 3 summarizes the training stage configurations.

A code snippet showcasing a correction introduced to Algorithm 18 of the AF3 supplementary materials:

Algorithm 18
Original Line 9:  x₁ = (xnoisy – denoised)/î
Corrected Line 9: x₁ = (xnoisy – denoised)/î

The correction ensures mathematical consistency in solving Ordinary Differential Equations (ODEs).

Future Plan

Future development focuses on improving model performance, adding features (predicting structures with incorporated priors or without MSA), and enhancing the evaluation process. The authors plan to create a general-purpose evaluation platform to support fair model comparisons.

Acknowledgements

The authors express gratitude to the creators of AlphaFold, OpenFold, and ColabFold, whose contributions provided a strong foundation for this work. The paper also acknowledges various open-source tools and datasets.

Conclusion

Protenix represents a valuable contribution to the field of protein structure prediction. Its open-source nature and improved performance, alongside the detailed documentation of its development process, make it a significant resource for the research community, facilitating further development and exploration. The authors encourage readers to consult the repository for detailed information and usage examples.