[论文解读] AugLiChem: Data Augmentation Library of Chemical Structures for Machine Learning

Rishikesh Magar, Yuyang Wang+5 more|arXiv (Cornell University)|Nov 29, 2021

Machine Learning in Materials Science被引用 5

一句话总结

本文介绍了 AugLiChem，一个基于 Python 的分子和晶体材料数据增强库，通过生成多样化且化学上有效的结构变体来提升机器学习模型的性能。通过在分子图上应用图级别变换（如节点掩码和键删除），以及在晶体上应用结构扰动，作者在多个基准数据集上展示了显著的性能提升——尤其在图神经网络（GNN）上，ROC-AUC 和归一化 RMSE 均有改善。

ABSTRACT

Machine learning (ML) has demonstrated the promise for accurate and efficient property prediction of molecules and crystalline materials. To develop highly accurate ML models for chemical structure property prediction, datasets with sufficient samples are required. However, obtaining clean and sufficient data of chemical properties can be expensive and time-consuming, which greatly limits the performance of ML models. Inspired by the success of data augmentations in computer vision and natural language processing, we developed AugLiChem: the data augmentation library for chemical structures. Augmentation methods for both crystalline systems and molecules are introduced, which can be utilized for fingerprint-based ML models and Graph Neural Networks(GNNs). We show that using our augmentation strategies significantly improves the performance of ML models, especially when using GNNs. In addition, the augmentations that we developed can be used as a direct plug-in module during training and have demonstrated the effectiveness when implemented with different GNN models through the AugliChem library. The Python-based package for our implementation of Auglichem: Data augmentation library for chemical structures, is publicly available at: https://github.com/BaratiLab/AugLiChem.

研究动机与目标

• 解决机器学习中分子和晶体材料数据集有限且噪声较多的挑战。
• 通过开发针对性的数据增强策略，克服化学性质预测中的数据稀缺问题。
• 设计一个模块化、开源的 Python 软件包，可无缝集成到现有的分子和晶体机器学习工作流中。
• 在包括分类和回归任务在内的多种机器学习基准上，评估基于图的增强策略的有效性。
• 证明数据增强可提升图神经网络和基于指纹的模型的泛化能力和鲁棒性，尤其是在数据量较少的场景下。

提出的方法

• 为晶体材料开发了五种不同的数据增强策略，包括原子取代、晶格畸变和超胞扩展。
• 提出了三种用于分子的图级别增强技术：节点掩码、键删除和子结构移除，同时保持化学有效性。
• 将所有增强方法整合到一个模块化的 Python 软件包（AugLiChem）中，支持基于指纹的模型和图神经网络（GNNs）。
• 仅在训练集上应用增强操作，以保持评估的完整性并防止数据泄露。
• 在多个数据集上使用标准的机器学习评估协议，包括 k 折交叉验证和随机划分。
• 使用标准指标评估性能：分类任务使用 ROC-AUC，回归任务使用归一化 RMSE。

实验结果

研究问题

• RQ1 图级别的数据增强策略是否能提升图神经网络在分子和晶体材料性质预测中的性能？
• RQ2 哪些特定的增强技术（如节点掩码、键删除）在多种化学数据集上能带来最一致的性能提升？
• RQ3 在材料科学中常见的低数据场景下，数据增强如何影响机器学习模型的泛化能力和鲁棒性？
• RQ4 一个统一的、即插即用的库是否能简化将数据增强集成到现有分子和晶体机器学习工作流中的过程？
• RQ5 组合多种增强策略会产生怎样的影响？是否存在收益递减现象，还是产生协同增益？

主要发现

• 节点掩码和键删除单独使用时，在大多数基准测试中均提升了模型性能，其中节点掩码在 Lipo 回归数据集上表现尤为突出。
• 在五项分类任务中的四项（BBBP、HIV、Tox21、SIDER）和五项回归任务中的三项中，节点掩码与键删除的组合优于单一策略。
• 子结构移除整体提升有限，但在 BACE 和 MUV 数据集上取得了具有竞争力的结果，表明其具有任务特定的适用性。
• 将三种分子图增强方法（NM+BD+SR）全部组合使用时，性能反而低于更简单的组合，可能是因为结构偏离原始图的程度过高。
• 在五个晶体材料数据集上，使用 AugLiChem 显著提升了图神经网络的性能，展示了在四种不同 GNN 架构上的稳定增益。
• 开源的 AugLiChem 软件包可无缝集成到现有机器学习流程中，支持自动数据处理、缺失数据插补，以及单目标和多目标学习。