[Paper Review] UniIF: Unified Molecule Inverse Folding

Zhangyang Gao, Jue Wang+6 more|arXiv (Cornell University)|May 29, 2024

Synthesis and properties of polymersMaterials Science3 citations

TL;DR

UniIF proposes the first unified deep learning framework for inverse folding of all molecules—small molecules, proteins, RNA, and materials—by introducing a unified block graph representation and a geometric block attention network with virtual long-term dependency modules. It achieves state-of-the-art performance across protein, RNA, and material design tasks, outperforming prior models by significant margins.

ABSTRACT

Molecule inverse folding has been a long-standing challenge in chemistry and biology, with the potential to revolutionize drug discovery and material science. Despite specified models have been proposed for different small- or macro-molecules, few have attempted to unify the learning process, resulting in redundant efforts. Complementary to recent advancements in molecular structure prediction, such as RoseTTAFold All-Atom and AlphaFold3, we propose the unified model UniIF for the inverse folding of all molecules. We do such unification in two levels: 1) Data-Level: We propose a unified block graph data form for all molecules, including the local frame building and geometric feature initialization. 2) Model-Level: We introduce a geometric block attention network, comprising a geometric interaction, interactive attention and virtual long-term dependency modules, to capture the 3D interactions of all molecules. Through comprehensive evaluations across various tasks such as protein design, RNA design, and material design, we demonstrate that our proposed method surpasses state-of-the-art methods on all tasks. UniIF offers a versatile and effective solution for general molecule inverse folding.

Motivation & Objective

• Address the lack of unified models for inverse folding across small molecules, proteins, and materials.
• Overcome unit discrepancy by treating atoms, amino acids, and nucleotides as unified blocks with fixed-size representations.
• Develop a consistent geometric featurization strategy that captures 3D interactions across diverse molecular types.
• Enable long-range dependency modeling in large molecular systems without quadratic computational cost.
• Create a single, general-purpose model that performs well across diverse inverse folding tasks

Proposed method

• Introduce a unified block graph representation where each block (atom, amino acid, nucleotide) is encoded with decoupled equivariant basis and invariant features.
• Construct geometric features using local coordinate frames and dot products between virtual atoms to model 3D interactions.
• Design a geometric block attention network with three core modules: geometric interaction, interactive attention, and virtual long-term dependency.
• Use sparse GNNs with global virtual blocks to maintain long-range dependencies while reducing computational cost.
• Learn local frame orientations via a differentiable, GNN-based mechanism to improve geometric feature extraction.
• Reuse the same geometric featurizer across all model layers to enable iterative refinement of structural representations.

Experimental results

Research questions

• RQ1 Can a single deep learning model effectively perform inverse folding across small molecules, proteins, RNA, and materials?
• RQ2 How can a unified data representation be designed to handle the unit discrepancy between atoms and predefined microstructures?
• RQ3 What geometric featurization strategy enables consistent 3D interaction modeling across diverse molecular types?
• RQ4 Can long-range dependencies be effectively captured in large molecular systems without incurring quadratic computational costs?
• RQ5 Does the integration of virtual long-term dependency modules improve performance on inverse folding tasks?

Key findings

• UniIF achieves state-of-the-art performance on protein design, with a recovery rate of 48.94% ± 0.37, surpassing the previous best (RDesign at 41.53% ± 0.38).
• On RNA design, UniIF achieves 48.94% ± 0.37 recovery rate, outperforming PiFold (24.48% ± 1.13) and GraphTrans (24.73% ± 0.93).
• In material design on the CHILI-3K dataset, UniIF achieves 75.3% ± 1.2 recovery rate, significantly outperforming GIN (58.7% ± 0.2) and EdgeCNN (63.2% ± 0.9).
• Ablation studies show that learning local frames via a GNN improves performance, with recovery dropping to 65.2% ± 3.9 when using Householder orthogonalization.
• The virtual long-term dependency module enhances performance, with UniIF (drop 0.2) achieving 47.19% ± 0.45 recovery on protein design, compared to 44.29% ± 0.29 without it.
• UniIF successfully designs sequences that fold into near-native structures, as confirmed by AlphaFold3 re-folding, with improved RMSD and recovery over baselines.