[论文解读] ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs

Limei Wang, Yi Liu+3 more|arXiv (Cornell University)|Jun 17, 2022

Machine Learning in Materials Science被引用 27

一句话总结

ComENet 引入了一种完整且高效的单跳邻域消息传递方案用于三维分子图，达到 SE(3) 完整表示，复杂度为 O(nk)，在大型数据集上具备强劲的经验表现。

ABSTRACT

Many real-world data can be modeled as 3D graphs, but learning representations that incorporates 3D information completely and efficiently is challenging. Existing methods either use partial 3D information, or suffer from excessive computational cost. To incorporate 3D information completely and efficiently, we propose a novel message passing scheme that operates within 1-hop neighborhood. Our method guarantees full completeness of 3D information on 3D graphs by achieving global and local completeness. Notably, we propose the important rotation angles to fulfill global completeness. Additionally, we show that our method is orders of magnitude faster than prior methods. We provide rigorous proof of completeness and analysis of time complexity for our methods. As molecules are in essence quantum systems, we build the \underline{com}plete and \underline{e}fficient graph neural network (ComENet) by combing quantum inspired basis functions and the proposed message passing scheme. Experimental results demonstrate the capability and efficiency of ComENet, especially on real-world datasets that are large in both numbers and sizes of graphs. Our code is publicly available as part of the DIG library (\url{https://github.com/divelab/DIG}).

研究动机与目标

• 动机：在分子图中充分整合三维信息以实现准确学习的必要性。
• 定义并保证三维图表示的几何完整性（全局和局部）。
• 开发具有旋转角度的高效单跳消息传递方案，以实现全局完整性。
• 在 ComENet 中整合量子启发特征，以提升三维分子学习。
• 在大规模数据集（OC20、Molecule3D、QM9）上展示可扩展性和具有竞争力的准确性。

提出的方法

• 在 SE(3) 下形式化定义三维图的完整性，并通过边的旋转角度证明全局完整性。
• 对每条边计算单跳几何特征（d、theta、phi、tau）以形成 SE(3) 不变表示。
• 引入在 1 跳邻域内工作的局部和全局卷积层，以实现 O(nk) 复杂度。
• 使用基于最近邻的旋转角度 tau 来唯一确定 3D 构型并实现构型体识别。
• 在 ComENet 架构中整合量子启发的基底特征和基于神经网络的消息传递框架（嵌入层、交互层、自身原子层、池化）。
• 给出完整性证明并分析时间复杂度；在 OC20、Molecule3D 和 QM9 上进行评估。

实验结果

研究问题

• RQ1 是否可以仅使用单跳邻域信息就实现三维分子图的全局完整性（信息不丢失）？
• RQ2 如何利用旋转角度来实现全局完整性并实现构型水平的辨别？
• RQ3 与现有的两跳方法相比，完整的三维 GNN 的实际时间复杂度是多少？
• RQ4 将量子启发特征与所提出的消息传递结合起来，是否在大规模分子数据集上取得有竞争力的性能？

主要发现

• 证明了单跳完整几何变换（d、theta、phi、tau）对强连通的三维图可产生全局完整性。
• 通过使用单跳邻域和旋转角度实现 O(nk) 时间复杂度，明显快于 O(nk^2) 或 O(nk^3) 方法。
• 在 OC20、Molecule3D 和 QM9 上证明了具有竞争力或优于最先进基线的性能，并在某些数据集上实现了训练/推理速度的大幅提升（6-10 倍）。
• 显示旋转角度在 OC20 中实现构型级辨别并识别重要的构象差异。
• 提出将完整消息传递与量子启发特征整合的 ComENet 架构，用于可扩展的 3D 分子学习。