APACE: AlphaFold2 and Advanced Computing as a Service for Accelerated Discovery in Biophysics

The 3D structure of a protein is entirely dependent on its amino acid sequence, and this in turn directly determines its function and the role it plays in various processes ranging from catalyzing chemical reactions to transporting molecules across cell membranes. The prediction of this 3D structure from an amino acid sequence represents a computational grand challenge in biophysics. Models developed from AI algorithms, such as AlphaFold, have had a transformative impact on biology and drug discovery, but can often require considerable computational expense, and the sheer size of the AI models can prohibit full exploitation of their capabilities. A team of researchers led by Argonne National Laboratory demonstrated how to combine large AI models with high-performance computing (HPC).

Challenge

Predicting protein structures is computationally challenging as there are 20 known amino acids, resulting in an immense number of possible conformations for a protein chain to adopt, and the conformations are themselves subject to further changes from the surrounding environment.

Approach

AlphaFold2 was trained on protein structures from the Protein Data Bank to predict the structure of a protein given an amino acid sequence. AlphaFold2 lacks parallelization and is only capable of leveraging a single GPU during inference, limiting it to the prediction of just one protein structure at a time. To overcome its limitations, the researchers developed APACE, a computational tool to enhance use of AlphaFold2 on HPC systems. In addition to several optimizations, including the use of node-local storage and checkpointing, APACE leverages multiple instances of AlphaFold2 to greatly accelerate the prediction of a protein structure. The ALCF Catalyst team provided access to an umbrella project for testing chemistry and material science codes and helped carry out jobs on the ALCF’s Polaris system.

Results

Supercomputers, including ALCF’s Polaris system, were used to demonstrate the capabilities of APACE. CPU speedups around 2x and GPU speedups of 4x and 40x (depending on protein) were observed using the Ray compute engine for the distribution of AI frameworks. Ensemble sizes spanned 100 to 300 members, reducing time-to-solution from several days to approximately an hour. The protein structures obtained generally agreed well with ground truth structures for which AlphaFold2 was expected to provide reliable results. Minor misfoldings were addressed with standard methods like molecular dynamics simulations. AlphaFold2 generated erroneous structures for the proteins for which the application was expected to have difficulty (e.g., GABA_A). The authors provided instructions to reproduce the AlphaFold2 installation, and APACE is available in GitHub.

Impact

The developed APACE framework can leverage large supercomputers to accelerate the capabilities of AI models, such as AlphaFold2. APACE was found to be up to two orders of magnitude faster than off-the-shelf AlphaFold2 implementations, reducing time-to-solution from weeks to minutes and potentially enabling accelerated innovation and discovery. This presents promising implications for pharmaceutical research, where employing the framework can dramatically shorten timelines for the screening of candidate compounds and the subsequent development of specialized drugs tailored for specific viral strains. Successful integration of a multi-terabyte AI model and database in modern supercomputer environments paves the way to iteratively improved research cycles. This computational approach could be readily linked with self-driving laboratories to enable automation and further accelerate scientific discovery.