[论文解读] Orientation-Aware Graph Neural Networks for Protein Structure Representation Learning

Jiahan Li, Shitong Luo+6 more|arXiv (Cornell University)|Jan 28, 2022

Machine Learning in Bioinformatics被引用 4

一句话总结

该论文提出了定向权重（\overrightarrow{{\mathbf{W}}}）图神经网络，将标量权重扩展为三维定向向量，以更好地捕捉蛋白质三维结构中的细粒度几何关系，例如主链取向和侧链二面角。通过将 \overrightarrow{{\mathbf{W}}}-感知器（\overrightarrow{{\mathbf{W}}}-perceptrons）集成到蛋白质图上的新型 SO(3)-等变消息传递框架中，该模型在多个蛋白质结构预测任务中实现了最先进性能，同时保持了几何表达能力和生物学合理性。

ABSTRACT

By folding into particular 3D structures, proteins play a key role in living beings. To learn meaningful representation from a protein structure for downstream tasks, not only the global backbone topology but the local fine-grained orientational relations between amino acids should also be considered. In this work, we propose the Orientation-Aware Graph Neural Networks (OAGNNs) to better sense the geometric characteristics in protein structure (e.g. inner-residue torsion angles, inter-residue orientations). Extending a single weight from a scalar to a 3D vector, we construct a rich set of geometric-meaningful operations to process both the classical and SO(3) representations of a given structure. To plug our designed perceptron unit into existing Graph Neural Networks, we further introduce an equivariant message passing paradigm, showing superior versatility in maintaining SO(3)-equivariance at the global scale. Experiments have shown that our OAGNNs have a remarkable ability to sense geometric orientational features compared to classical networks. OAGNNs have also achieved state-of-the-art performance on various computational biology applications related to protein 3D structures. The code is available at https://github.com/Ced3-han/OAGNN/tree/main.

研究动机与目标

• 为解决经典及现有等变图神经网络在捕捉蛋白质三维结构中氨基酸取向和二面角等细微几何特征方面的局限性。
• 开发一种几何表达能力强的神经网络框架，支持标量和向量特征上的丰富几何运算。
• 设计一种在蛋白质图上保持全局几何不变性且不约束模型架构的 SO(3)-等变消息传递范式。
• 实现所提框架与现有图神经网络架构的无缝集成，仅需极少修改。
• 与经典及等变基线模型相比，证明其在蛋白质结构表征任务中具有更优的表达能力和性能。

提出的方法

• 提出 \overrightarrow{{\mathbf{W}}}-感知器，将标量权重推广为三维定向向量，实现对标量和向量特征的几何运算。
• 引入一种新型消息传递范式，利用局部主链取向（{\bm{O}}_{u}）将特征转换到局部坐标系中进行等变处理。
• 采用多层 \overrightarrow{{\mathbf{W}}}-感知器（\mathcal{F} 和 \mathcal{H}）处理聚合后的节点和边特征，并将结果转换回全局坐标系。
• 将该框架应用于多种图神经网络变体：DW-GCN、DW-GIN 和 DW-GAT，其中注意力机制针对标量和向量表征进行了适配。
• 使用弗罗贝尼乌斯内积和标准内积分别计算向量和标量注意力得分，实现对邻居的有选择性聚合。
• 通过利用每个氨基酸残基的刚体变换，形式化证明了全局 SO(3)-等变性。

实验结果

研究问题

• RQ1 将标量权重扩展为三维定向向量，是否能显著提升对蛋白质三维结构中细粒度几何特征的建模能力？
• RQ2 如何在不施加架构约束的前提下，有效保持通用图神经网络框架中的 SO(3)-等变性？
• RQ3 \overrightarrow{{\mathbf{W}}}-感知器在捕捉蛋白质特异性几何关系（如侧链二面角）方面，相较于经典及现有等变图神经网络的性能提升程度如何？
• RQ4 所提框架能否灵活集成到多种图神经网络架构中，同时保持几何表达能力和性能？
• RQ5 该等变消息传递范式是否能在多个蛋白质结构学习任务中提升整体性能？

主要发现

• 所提出的 \overrightarrow{{\mathbf{W}}}-GNN 框架在建模几何关系方面表现出优于经典标量加权网络及现有 SO(3)-等变模型的优越表达能力。
• 在合成数据集上，模型成功捕捉了角度关系（如侧链二面角），而基于 VNN 的方法对此类关系建模效果较差。
• 该框架在多个蛋白质结构学习任务中实现了最先进性能，包括功能预测、错误构象排序、蛋白质对接和致病突变识别。
• DW-GAT 变体通过引入对标量和向量特征的可学习注意力机制，实现了选择性邻居聚合，提升了对异质性蛋白质相互作用的表征学习能力。
• SO(3)-等变消息传递范式确保了全局几何一致性，同时不限制 \mathcal{F} 和 \mathcal{H} 的设计，从而支持灵活且强大的网络架构。
• 所有模型和代码将在论文被接受后公开发布，以支持可复现性，并为未来在其他三维学习任务中的扩展提供基础。