MP-NeRF: A Massively Parallel Method for Accelerating Protein Structure Reconstruction from Internal Coordinates

Abstract

The conversion of proteins between internal and cartesian coordinates is a limiting step in many pipelines, such as molecular dynamics simulations and machine learning models. This conversion is typically carried out by sequential or parallel applications of the Natural extension of Reference Frame (NeRF) algorithm. This work proposes a massively parallel NeRF implementation which, depending on the polymer length, achieves speedups between 400-1200x over the previous state-of-the-art. It accomplishes this by dividing the conversion into three main phases: parallel composition of the monomer backbone, assembly of backbone subunits, and parallel elongation of sidechains; and by batching these computations into a minimal number of efficient matrix operations. Special emphasis is placed on reusability and ease of use. We open source the code (available at https://github.com/EleutherAI/mp_nerf) and provide a corresponding python package.

NeRFax: An efficient and scalable conversion from the internal representation to Cartesian space

“…A comparative assessment of performance gain through NeRFax method. The PyTorch mp-NeRF implementation is used as a performance benchmark; A) relative performance on CPU and GPU with respect to the state-of-the-art NeRF code (Alcaide et al (2022)) as a function of protein length, B) a 1,000-molecule model of an IDP aggregate.…”

“…Unlike a number of previous implementations of NeRF that leveraged C++ - like performance and accessibility of the high-level Python libraries Tensorflow and Pytorch (Alcaide et al . (2022); AlQuraishi (2019)), our algorithm was programmed using JAX.…”

“…An arbitrary reconstruction task was taking 2 seconds for a 100 residue protein. Alcaide et al . (2022) proposed a generalised implementation of the pNeRF algorithm, denoted mp-NeRF.…”

“…An arbitrary reconstruction task was taking 2 seconds for a 100 residue protein. Alcaide et al (2022) proposed a generalised implementation of the pNeRF algorithm, denoted mp-NeRF. Their approach utilized the branched chain strategy of Bayati et al (2020) bringing down the reconstruction time to 10 ms on average.…”

“…Unlike a number of previous implementations of NeRF that leveraged C++ -like performance and accessibility of the high-level Python libraries Tensorflow and Pytorch(Alcaide et al (2022); AlQuraishi (2019)), our algorithm was programmed using JAX. Superior numerical performance of NeRFax transpires from optimisations applied during compilation, including fusion and constant folding.…”