[论文解读] Graph Neural Network Architecture Search for Molecular Property Prediction

Shengli Jiang, Prasanna Balaprakash|arXiv (Cornell University)|Aug 27, 2020

Computational Drug Discovery Methods参考文献 35被引用 5

一句话总结

本文提出一种基于正则化进化（regularized evolution）的神经架构搜索（NAS）框架，用于自动发现适用于分子性质预测的最优消息传递神经网络（MPNN）架构。通过引入边特征、跳跃连接以及全面的搜索空间，该方法发现了性能优越的MPNN模型，在MoleculeNet基准数据集上的表现优于人工设计的模型，凸显了架构定制在分子机器学习中的关键作用。

ABSTRACT

Predicting the properties of a molecule from its structure is a challenging task. Recently, deep learning methods have improved the state of the art for this task because of their ability to learn useful features from the given data. By treating molecule structure as graphs, where atoms and bonds are modeled as nodes and edges, graph neural networks (GNNs) have been widely used to predict molecular properties. However, the design and development of GNNs for a given data set rely on labor-intensive design and tuning of the network architectures. Neural architecture search (NAS) is a promising approach to discover high-performing neural network architectures automatically. To that end, we develop an NAS approach to automate the design and development of GNNs for molecular property prediction. Specifically, we focus on automated development of message-passing neural networks (MPNNs) to predict the molecular properties of small molecules in quantum mechanics and physical chemistry data sets from the MoleculeNet benchmark. We demonstrate the superiority of the automatically discovered MPNNs by comparing them with manually designed GNNs from the MoleculeNet benchmark. We study the relative importance of the choices in the MPNN search space, demonstrating that customizing the architecture is critical to enhancing performance in molecular property prediction and that the proposed approach can perform customization automatically with minimal manual effort.

研究动机与目标

• 自动化设计用于分子性质预测的消息传递神经网络（MPNN）以减少人工调参并提升性能。
• 通过开发一种针对数据特性的自动化架构搜索方法，解决分子数据集中的低数据效率和非欧几里得结构问题。
• 构建一种新型搜索空间，包含边特征、跳跃连接以及可配置的MPNN组件（如注意力机制、聚合函数、激活函数），以实现更好的泛化能力。
• 评估神经架构搜索（NAS）在发现任务特定MPNN架构方面的有效性，使其在MoleculeNet数据集上超越人工设计的基线模型。
• 通过敏感性分析探究MPNN中不同架构选择的相对重要性，以指导未来模型开发。

提出的方法

• 定义一个全面的MPNN搜索空间，包含状态维度、注意力机制、注意力头数、聚合函数、激活函数、更新函数、聚合操作以及跳跃连接。
• 在DeepHyper AutoML框架内实现正则化进化（RE）搜索策略，以高效且异步地探索搜索空间。
• 使用基于验证损失的奖励信号在NAS过程中评估候选架构，通过延长训练周期重新训练模型，以确保性能比较的稳健性。
• 通过可学习的神经网络将边特征集成到消息传递过程中，生成注意力权重，从而增强模型捕捉化学键特性的能力。
• 采用基于随机森林的非参数敏感性分析，量化不同数据集上各架构超参数的重要性。
• 在15、30和60台计算节点上，比较RE与随机搜索（RS）在收敛速度、高性能架构发现数量以及可扩展性方面的表现。

实验结果

研究问题

• RQ1 神经架构搜索（NAS）能否自动发现超越人工设计模型的高性能MPNN架构，用于分子性质预测？
• RQ2 在搜索空间中引入边特征和跳跃连接，对MPNN在分子数据集上的性能有何影响？
• RQ3 在不同分子性质预测任务中，MPNN架构组件（如注意力机制、聚合函数、激活函数）的相对重要性如何？
• RQ4 正则化进化（RE）搜索策略在发现高质量MPNN架构方面，是否在效率和性能上优于随机搜索（RS）？
• RQ5 最优MPNN架构在不同分子数据集上的差异程度如何，是否表明需要进行数据特定的定制？

主要发现

• 使用正则化进化（RE）自动发现的MPNN架构在所有数据集的验证损失和测试损失上均优于MoleculeNet基准中所有人工设计的MPNN模型。
• 正则化进化（RE）发现的高性能架构数量多于随机搜索（RS），且在更短的时钟时间（wall-clock time）内达到更优性能，证明其具有更高的搜索效率。
• 包含边特征、注意力机制和跳跃连接的搜索空间显著提升了模型性能，相比不包含这些组件的标准MPNN有明显改进。
• 敏感性分析表明，注意力机制、聚合函数和激活函数等架构选择的重要性因数据集而异，凸显了数据特定架构设计的必要性。
• 基于RE的NAS方法在15、30和60台计算节点上均表现出良好的可扩展性，性能增益保持一致，未出现搜索质量下降。
• 结果表明，架构定制对分子性质预测至关重要，且NAS可实现该过程的自动化，仅需极少的人工干预。