Parallel EvoFormer in Protein Prediction

• Parallel Evoformer is a decoupled design of the Evoformer block that enables concurrent processing of MSA and pair representations.
• It leverages branch parallelism by assigning MSA and pair-stack computations to separate devices, achieving 38-40% speed improvements.
• Experimental evaluations show it maintains predictive accuracy while halving activation memory, boosting scalable protein structure prediction.

The Parallel Evoformer is a restructuring of the Evoformer block in AlphaFold2, devised to enable efficient parallel computation of its primary branches during end-to-end protein structure prediction model training. By separating the multiple sequence alignment (MSA) track and pair interaction track into independently executable computation graphs, and introducing a synchronized cross-communication at the block's close, the Parallel Evoformer—combined with Branch Parallelism—achieves significant wall-clock speedups and activation memory reduction, while maintaining the original model’s predictive accuracy (Wang et al., 2022).

1. Standard Evoformer Block in AlphaFold2

The classical Evoformer block operates on two high-dimensional data structures: the MSA representation $M\in\mathbb{R}^{b\times s\times r\times c_m}$ and the pair representation $Z\in\mathbb{R}^{b\times r\times r\times c_z}$, where $b$ is batch size, $s$ is MSA depth, $r$ is sequence length, and $c_m$, $c_z$ are channel dimensions. The standard computation proceeds sequentially as follows:

• Row-wise gated self-attention on MSA, with pair bias derived from $Z$, followed by column-wise gated self-attention across sequences.
• Transition (feed-forward) applied to MSA features.
• Outer-Product Mean (OPM) computes an MSA-to-pair cross-term, updating $Z$.
• Pair-stack operations: triangle multiplicative updates and self-attention on $Z$, followed by a feed-forward network.

These components are tightly coupled, as cross-branch OPM occurs mid-block. Each sub-block is $Z\in\mathbb{R}^{b\times r\times r\times c_z}$0 or $Z\in\mathbb{R}^{b\times r\times r\times c_z}$1 in computational cost and comparable in activation storage.

2. Structural Modification: The Parallel Evoformer Concept

The Parallel Evoformer relocates the cross-branch OPM operation to the end of each block, effecting a decoupled architecture where:

• The MSA-stack (steps a–c) and pair-stack (steps e–g) are computed independently of each other.
• Both branches operate in parallel rather than sequentially, requiring only synchronization at block boundaries.
• Finalization per block involves $Z\in\mathbb{R}^{b\times r\times r\times c_z}$2, the independently updated pair representation, receiving the cross-term via $Z\in\mathbb{R}^{b\times r\times r\times c_z}$3.

Theoretically, this reduces serial runtime from $Z\in\mathbb{R}^{b\times r\times r\times c_z}$4 to $Z\in\mathbb{R}^{b\times r\times r\times c_z}$5, where $Z\in\mathbb{R}^{b\times r\times r\times c_z}$6 is the communication time for the OPM and synchronization. Activation memory usage per device decreases by roughly half when the branches run on distinct accelerators.

3. Branch Parallelism Execution Scheme

Branch Parallelism concretely maps the MSA- and pair-stack branches to separate devices (“ranks”):

• Rank 0: computes MSA operations, produces OPM output, and broadcasts it to rank 1.
• Rank 1: receives OPM, computes pair-stack (without OPM during forward), incorporates the OPM, and broadcasts the new $Z\in\mathbb{R}^{b\times r\times r\times c_z}$7 back to rank 0.

The backward pass mirrors this exchange with gradient broadcast for OPM and all-reduce for the shared $Z\in\mathbb{R}^{b\times r\times r\times c_z}$8 gradients. Only two small tensors are communicated per block per step, so communication overhead remains low compared to per-branch computation. The splitting is efficient because Evoformer block timing is naturally load-balanced, with MSA-stack and pair-stack consuming approximately 60% and 40% of runtime, respectively.

4. Performance Evaluation and Experimental Setup

Experiments were conducted using NVIDIA A100 40 GB GPUs (up to 256 GPUs over two nodes), employing one protein per GPU (batch size 128 for Data Parallel (DP), 256 for BP+DP), with frameworks UniFold (PyTorch) and HelixFold (PaddlePaddle). Mixed precisions were used: BF16 activations and FP32 weights.

Benchmark results demonstrate the following speedups with BP+DP versus DP alone:

Framework	Initial train (DP)	Initial train (BP+DP)	Throughput gain	Fine-tune (DP)	Fine-tune (BP+DP)	Gain
UniFold	4.16 s/step, 30.76 prot/s	3.02 s/step, 42.38 prot/s	+37.74%	15.02 s/step, 8.52 prot/s	10.70 s/step, 11.96 prot/s	+40.37%
HelixFold	4.92 s/step, 26.01 prot/s	3.55 s/step, 36.05 prot/s	+38.59%	16.45 s/step, 7.78 prot/s	12.30 s/step, 10.41 prot/s	+33.84%

Complete end-to-end training time on UniFold fell from 5.80 days (128 A100 DP) to 4.18 days (256 A100 BP+DP), a 38.7% improvement. HelixFold showed analogous gains.

5. Accuracy on Structure Prediction Benchmarks

The Parallel Evoformer architecture yields accuracy indistinguishable from the original sequential block arrangement. On CASP14 (87 domains) and CAMEO (371 targets):

Benchmark	AlphaFold2 Evoformer	Parallel Evoformer	$Z\in\mathbb{R}^{b\times r\times r\times c_z}$9TM	$b$0lDDT-C$b$1
CASP14	TM = 0.891, lDDT = 0.783	TM = 0.893, lDDT = 0.785	<0.005	<0.005
CAMEO	TM = 0.842, lDDT = 0.756	TM = 0.845, lDDT = 0.759	<0.005	<0.005

$b$2TM and $b$3lDDT-C$b$4 are both $b$5, falling within the established run-to-run stochasticity of the model (Wang et al., 2022).

6. Communication, Scalability, and Memory Considerations

The added communication per block consists of two small broadcasts/all-reduces linked to branch cross-terms, with negligible impact on overall step time due to moderate tensor sizes ($b$6, $b$7). BP’s device scaling is constrained by the number of heavy branches, here limited to two; scaling beyond requires combining BP with Data Parallel or Activation-split schemes. Splitting the tracks allows each device to store only the relevant activations, maximizing per-GPU memory efficiency until additional activation partitioning yields diminishing returns.

7. Extensions and Limitations

Parallel Evoformer and BP provide a template for extending efficient branch parallelism in models with more than two heavy branches, such as Evoformer variants for multimeric complexes. The scheme has plausible applicability to other multi-branch deep learning architectures in areas such as vision and NLP, especially when per-branch loads are balanced. Further potential enhancements include pipeline-fused inter-block communication to hide latency and integrating tensor-parallel strategies for scaling to extremely large models.

In summary, Parallel Evoformer with Branch Parallelism recasts each Evoformer block as two concurrently executable branches with a lightweight synchronization mechanism, producing substantial end-to-end speedup (~38–40%) and halved activation memory, all while preserving state-of-the-art accuracy for protein structure prediction (Wang et al., 2022).