Source-linked AI summary
PiFold: Toward effective and efficient protein inverse folding
Zhangyang Gao, Cheng Tan, Pablo Chacón, Stan Z. Li
TL;DR
Structure-based protein design needs sequences that fold into desired structures, but existing methods struggle to combine expressive features with efficient decoding. PiFold introduces a richer residue featurizer and PiGNN layers for one-shot sequence generation. It reports strong recovery across CATH 4.2, TS50, and TS500 while achieving a 70-times inference-speed improvement over autoregressive competitors for long proteins.
Problem
Existing structure-based protein design methods have limited residue-feature expressiveness and rely on autoregressive decoding, making it difficult to improve accuracy and efficiency together.
Method
PiFold combines a comprehensive residue featurizer, learnable virtual atoms, multi-scale PiGNN layers, and one-shot sequence generation without an autoregressive decoder.
Results
PiFold achieves 51.66% recovery on CATH 4.2, 58.72% on TS50, and 60.42% on TS500, while running 70 times faster than autoregressive competitors for long proteins.
Takeaways & Limitations
PiFold demonstrates that protein inverse folding can improve recovery and inference efficiency simultaneously through richer features and one-shot generation.
Abstract
from arXiv · showhide
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}.
1 INTRODUCTION
Structure-based protein design seeks sequences that fold into desired 3D structures, but existing graph models face limited feature expressiveness and slow autoregressive decoding. PiFold addresses both issues with richer residue features, PiGNN layers, and one-shot sequence generation.
- Motivation: Structure-based protein design aims to create amino acid sequences that fold into desired structures for biological applications.The paper frames this as a crucial problem for understanding living systems and fighting disease.
- Challenges: Existing graph models improve representations through feature engineering, architectural modules, or larger datasets, yet expressive residue representations remain challenging.Examples include angle and distance features, geometric message passing, equivariant vectors, global attention, and additional training data.
- Challenges: Autoregressive decoding dramatically slows inference in most graph-based protein design models.The paper notes that efficiency improvements may be perceived as requiring an accuracy sacrifice.
- PiFold: PiFold combines a novel residue featurizer with stacked PiGNNs to learn from richer, multi-scale residue interactions.The featurizer uses comprehensive features and learnable virtual atoms, while PiGNN models dependencies at node, edge, and global levels.
2 RELATED WORKS
Structure-based protein design is formulated as finding amino acid sequences that fold into desired structures, an underdetermined inverse problem. Prior methods use MLP, CNN, or graph-based structural encoders, often combining structural encoding with autoregressive sequence decoding.
- Problem formulation: Protein folding has an inverse design problem: finding amino acid sequences that fold into a desired 3D structure.The paper represents structures with residue coordinates and sequences with amino acid identities.
- Problem formulation: The inverse folding problem is underdetermined because homologous proteins can share similar structures while valid amino acid sequences are not unique.Recent approaches therefore consider both sequential and structural dependencies.
- Prior methods: Existing structure-based design methods are grouped into MLP-based, CNN-based, and GNN-based models according to their structural encoders.These families differ in how they extract residue or structural representations.
- Prior methods: MLP-based methods primarily focus on feature construction, whereas CNN-based methods extract protein features from 2D or 3D representations.SPIN uses torsion angles, sequence profiles, and energy profiles; SPROF uses a distance matrix with a 2D CNN.
- Prior methods: Graph-based methods represent 3D structures as k-NN graphs with residue nodes and pairwise interactions encoded as node and edge features.Graph models incorporate structural constraints through graph neural networks and avoid rotating each residue separately as CNN methods do.
3 METHOD
PiFold combines comprehensive, invariant residue features with PiGNN layers that model node-, edge-, and global-level interactions. It uses these representations to generate protein sequences in one shot rather than through autoregressive or iterative decoding.
- PiFold takes protein structures as input and produces sequences expected to fold into those structures.
- Decoding: PiFold generates protein sequences in a one-shot manner without autoregressive or iterative decoding schemas.The framework stacks PiGNN layers to remove the autoregressive decoder.
- Featurizer: The featurizer constructs rotation- and translation-invariant distance, angle, and direction features for individual or paired residues.It uses local residue coordinate systems and features derived from backbone atoms.
- Featurizer: Learnable virtual atoms provide additional distance features intended to capture information complementary to real atoms.
- PiGNN: PiGNN learns geometric residue representations from multi-scale node, edge, and global context interactions.The architecture includes local node updates, iterative edge updates, and global context attention.
4 EXPERIMENTS
PiFold is evaluated for accuracy, efficiency, ablations, and generalization across CATH, TS50, and TS500. Results assess whether one-shot generation preserves recovery while reducing inference cost.
- 4.1 PERFORMANCE (Q1): PiFold is evaluated against recent graph models on CATH using a standardized split of 18,024 training, 608 validation, and 1,120 test proteins.Baselines include StructGNN, StructTrans, GCA, GVP, GVP-large, AlphaDesign, ESM-IF, and ProteinMPNN.
- 4.1 PERFORMANCE (Q1): 20 epochs are enough for PiFold to achieve state-of-the-art performance during training.The efficiency evaluation reports state-of-the-art perplexity and recovery with fewer training epochs.
- 4.1 PERFORMANCE (Q1): 70 times faster inference is achieved by PiFold than autoregressive competitors on 100 long chains averaging 1,632 residues.PiFold uses one-shot generation with O(1) complexity, whereas autoregressive decoding has O(L) complexity in protein length L.
- 4.2 ABLATION (Q2): Edge features rank above node features, with edge distances identified as essential and Cα-O and C-N distances especially significant.Virtual atoms also improve recovery, with more virtual atoms producing higher recovery in the reported ablation.
- 4.2 ABLATION (Q2): AttMLP messaging improves recovery by at least 0.9%, while global context attention and edge updating add 0.44% and 1.22%, respectively.These results come from systematic ablations of PiGNN modules.
- 4.2 ABLATION (Q2): Replacing autoregressive decoding with one-shot PiGNN improves recovery while reducing inference time, and deeper encoders further improve recovery.The ablation compares progressively fewer autoregressive layers with more encoder layers.
- 4.3 GENERALIZATION (Q3): PiFold achieves consistent improvements across TS50 and TS500, exceeding 55% recovery on TS50 and 60% on TS500.The evaluation reports perplexity, median recovery, and worst recovery, and includes all-alpha, all-beta, and mixed-protein examples.
5 CONCLUSION
PiFold simultaneously improves recovery and inference efficiency for structure-based protein design, achieving strong recovery across three test sets while greatly accelerating long-protein design.
- PiFold achieves 51.66% recovery on CATH 4.2, 58.72% on TS50, and 60.42% on TS500.
- PiFold is 70 times faster than autoregressive competitors for designing long proteins.
- The method combines a protein featurizer with PiGNN to improve recovery and efficiency simultaneously.
- Comprehensive ablation studies examine the roles of protein features and model designs for future simplification and improvement.
A APPENDIX
The appendix examines edge-distance features and clarifies a baseline comparison setting for structure-based sequence design. Distance ablations identify influential geometric features, while baseline handling addresses residues with missing coordinates.
- Distance ablation: Cα-O and C-N edge distances are the most significant features in the distance ablation.The analysis also finds that distances between learnable virtual atoms significantly improve recovery.
- Distance ablation: Distances between learnable virtual atoms significantly improve recovery, suggesting that the model can discover critical residual points automatically.
- Baseline details: ProteinMPNN supports partial-residue design, whereas this study focuses on designing sequences from structures without known residues.
- Baseline details: 51% recovery on CATH4.2 is obtained by ProteinMPNN when residues with missing coordinates are directly assigned sequence labels.The authors remove this operation to compare all methods under the same setting.