[Paper Review] PiFold: Toward effective and efficient protein inverse folding

Zhangyang Gao, Cheng Tan+2 more|arXiv (Cornell University)|Sep 22, 2022

Protein Structure and DynamicsBiochemistry, Genetics and Molecular Biology49 citations

TL;DR

PiFold introduces a one-shot graph neural network framework with a novel residue featurizer and PiGNN layers, achieving state-of-the-art recovery while dramatically speeding up inference compared to autoregressive models.

ABSTRACT

How can we design protein sequences folding into the desired structures effectively and efficiently? AI methods for structure-based protein design have attracted increasing attention in recent years; however, few methods can simultaneously improve the accuracy and efficiency due to the lack of expressive features and autoregressive sequence decoder. To address these issues, we propose PiFold, which contains a novel residue featurizer and PiGNN layers to generate protein sequences in a one-shot way with improved recovery. Experiments show that PiFold could achieve 51.66\% recovery on CATH 4.2, while the inference speed is 70 times faster than the autoregressive competitors. In addition, PiFold achieves 58.72\% and 60.42\% recovery scores on TS50 and TS500, respectively. We conduct comprehensive ablation studies to reveal the role of different types of protein features and model designs, inspiring further simplification and improvement. The PyTorch code is available at \href{https://github.com/A4Bio/PiFold}{GitHub}.

Motivation & Objective

• Motivate improved structure-based protein design by addressing both accuracy and inference speed.
• Propose a comprehensive residue featurizer that includes learnable virtual atoms to enrich representations.
• Develop PiGNN layers that capture multi-scale residue interactions at node, edge, and global levels.
• Eliminate autoregressive decoding to enable one-shot sequence generation without sacrificing accuracy.
• Empirically validate PiFold on CATH, TS50, and TS500 with ablations to understand feature and module contributions.

Proposed method

• Construct a k-NN protein graph (k=30) with enhanced node and edge features, including rotation- and translation-invariant distance, angle, and direction features.
• Introduce learnable virtual atoms V_i^k to capture complementary information to real atoms for distance features.
• Define PiGNN layers that combine local node updates, edge updates, and a global context gate for efficient multi-scale learning.
• Use AttMLP-based node message passing, with optional edge updating and global context attention to improve expressivity.
• Adopt a one-shot graph-based generation without autoregressive decoding, enabling linear-time (O(1) with respect to length) inference for long proteins.

Experimental results

Research questions

• RQ1 Can a one-shot graph-based model surpass autoregressive decoders in accuracy while dramatically increasing inference speed?
• RQ2 What residue features and multi-scale interactions most improve recovery in structure-based protein design?
• RQ3 How well does PiFold generalize to multiple benchmarks (CATH, TS50, TS500) compared to existing graph-based methods?
• RQ4 What is the impact of removing autoregressive decoding on performance and efficiency?

Key findings

• PiFold achieves 51.66% recovery on CATH 4.2, outperforming several baselines.
• PiFold attains 58.72% recovery on TS50 and 60.42% on TS500, the first graph model to exceed 55% on TS50 and 60% on TS500.
• Inference speed is ~70x faster than autoregressive competitors for long proteins.
• Ablation shows edge features and virtual atoms significantly contribute to performance, and removing autoregressive decoding with deeper encoders maintains or improves recovery.
• 20 training epochs suffice to reach state-of-the-art performance on the CATH benchmark.