[论文解读] UniIF: Unified Molecule Inverse Folding

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

Synthesis and properties of polymersMaterials Science被引用 3

一句话总结

UniIF 提出首个统一的深度学习框架，用于所有分子（小分子、蛋白质、RNA 和材料）的逆折叠问题，通过引入统一的块图表示和带有虚拟长期依赖模块的几何块注意力网络，实现了在蛋白质、RNA 和材料设计任务中的最先进性能，显著优于以往模型。

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.

研究动机与目标

• 解决小分子、蛋白质和材料之间逆折叠统一模型的缺乏问题。
• 通过将原子、氨基酸和核苷酸统一视为具有固定大小表示的块，克服单位差异问题。
• 设计一致的几何特征化策略，以捕捉不同类型分子之间的三维相互作用。
• 在大型分子系统中实现长程依赖建模，同时避免二次方计算成本。
• 构建一个单一、通用的模型，可在多种逆折叠任务中表现出色。

提出的方法

• 引入统一的块图表示，其中每个块（原子、氨基酸、核苷酸）通过解耦的等变基和不变特征进行编码。
• 利用局部坐标系和虚拟原子之间的点积构建几何特征，以建模三维相互作用。
• 设计几何块注意力网络，包含三个核心模块：几何交互、交互注意力和虚拟长期依赖模块。
• 使用带有全局虚拟块的稀疏图神经网络，以在降低计算成本的同时保持长程依赖。
• 通过可微的、基于图神经网络的机制学习局部框架方向，以提升几何特征提取效果。
• 在所有模型层中复用相同的几何特征提取器，以实现结构表示的迭代优化。

实验结果

研究问题

• RQ1 单一深度学习模型能否有效实现小分子、蛋白质、RNA 和材料的逆折叠？
• RQ2 如何设计统一的数据表示，以处理原子与预定义微观结构之间的单位差异？
• RQ3 何种几何特征化策略可实现在多种分子类型中的一致三维相互作用建模？
• RQ4 能否在不产生二次方计算成本的前提下，有效捕捉大型分子系统中的长程依赖？
• RQ5 虚拟长期依赖模块的集成是否能提升逆折叠任务的性能？

主要发现

• 在蛋白质设计任务中，UniIF 实现了最先进性能，恢复率达到 48.94% ± 0.37，显著优于此前最佳模型（RDesign 为 41.53% ± 0.38）。
• 在 RNA 设计任务中，UniIF 达到 48.94% ± 0.37 的恢复率，优于 PiFold（24.48% ± 1.13）和 GraphTrans（24.73% ± 0.93）。
• 在 CHILI-3K 数据集上的材料设计任务中，UniIF 达到 75.3% ± 1.2 的恢复率，显著优于 GIN（58.7% ± 0.2）和 EdgeCNN（63.2% ± 0.9）。
• 消融实验表明，通过图神经网络学习局部框架可提升性能，若改用 Householder 正交化，恢复率下降至 65.2% ± 3.9。
• 虚拟长期依赖模块显著提升性能，UniIF（drop 0.2）在蛋白质设计任务中恢复率达到 47.19% ± 0.45，而无该模块时仅为 44.29% ± 0.29。
• UniIF 能成功设计出可折叠为接近天然结构的序列，经 AlphaFold3 重新折叠验证，其 RMSD 和恢复率均优于基线模型。