Efficient AlphaFold2 Training using Parallel Evoformer and Branch Parallelism

TL;DR

This paper introduces Parallel Evoformer and Branch Parallelism techniques to significantly accelerate AlphaFold2 training without sacrificing accuracy, enabling faster development in structural biology.

Contribution

The paper proposes novel parallelism methods that improve AlphaFold2 training efficiency while maintaining high prediction accuracy.

Findings

Branch Parallelism improves training performance by over 36%.

Parallel Evoformer maintains accuracy comparable to AlphaFold2.

Experiments conducted on UniFold and HelixFold demonstrate substantial speedups.

Abstract

The accuracy of AlphaFold2, a frontier end-to-end structure prediction system, is already close to that of the experimental determination techniques. Due to the complex model architecture and large memory consumption, it requires lots of computational resources and time to train AlphaFold2 from scratch. Efficient AlphaFold2 training could accelerate the development of life science. In this paper, we propose a Parallel Evoformer and Branch Parallelism to speed up the training of AlphaFold2. We conduct sufficient experiments on UniFold implemented in PyTorch and HelixFold implemented in PaddlePaddle, and Branch Parallelism can improve the training performance by 38.67% and 36.93%, respectively. We also demonstrate that the accuracy of Parallel Evoformer could be on par with AlphaFold2 on the CASP14 and CAMEO datasets. The source code is available on https://github.com/PaddlePaddle/PaddleFleetX