[论文解读] SparseChem: Fast and accurate machine learning model for small molecules

Ádám Arany, Jaak Simm+2 more|arXiv (Cornell University)|Mar 9, 2022

Computational Drug Discovery Methods被引用 6

一句话总结

SparseChem 是一个为高维稀疏分子数据（如 ECFP 指纹）设计的快速、可扩展的机器学习框架，支持分类、回归和右删失回归的高效多任务学习。在单张 V100 GPU 上，仅用 2 分钟即可训练包含 42 万种化合物和 3,500 项任务的数据集，在基准数据集上达到最先进性能，AUC-ROC 为 0.794，R² 为 0.358。

ABSTRACT

SparseChem provides fast and accurate machine learning models for biochemical applications. Especially, the package supports very high-dimensional sparse inputs, e.g., millions of features and millions of compounds. It is possible to train classification, regression and censored regression models, or combination of them from command line. Additionally, the library can be accessed directly from Python. Source code and documentation is freely available under MIT License on GitHub.

研究动机与目标

• 为解决在药物发现中常见的高维稀疏分子数据上训练高精度机器学习模型的挑战。
• 通过高效的稀疏神经网络，实现支持数百万特征和数千项任务的可扩展多任务学习。
• 在一个统一框架内支持多样化的学习设置，包括分类、回归和删失回归。
• 在 CPU 和 GPU 硬件上实现高速训练和推理，适用于现实世界的药物发现流程。
• 通过命令行界面和 Python API 提供可扩展、用户友好的访问，促进广泛采用。

提出的方法

• SparseChem 使用稀疏神经网络层作为输入层，可高效处理如 ECFP 指纹等百万维稀疏特征向量。
• 采用稀疏训练损失，仅对观测到的输出元素计算梯度，从而实现对数万项任务的高效训练。
• 内部小批量处理技术可实现大于 GPU 内存容量的有效批量大小，提升训练效率。
• 该框架支持混合模型，可在同一网络中联合处理分类和（删失）回归任务，提升泛化能力和训练速度。
• 实现任务特定的损失函数：分类使用二元交叉熵，回归使用均方误差，上下删失回归使用单边平方损失。
• 系统支持每项任务的加权，并为每项任务计算标准指标（AUC-ROC、AUC-PR、R²、RMSE），实现全面评估。

实验结果

研究问题

• RQ1 机器学习框架能否在保持高预测精度的同时，高效处理百万维稀疏分子指纹？
• RQ2 使用稀疏神经网络的多任务学习在药物发现的 QSAR 建模中如何提升性能？
• RQ3 删失回归损失函数能否有效建模具有上下限的生物活性数据？
• RQ4 在包含 42 万种化合物和 3,500 项任务的工业规模数据集上，深度学习框架的训练速度和可扩展性如何？
• RQ5 与单任务学习相比，分类和回归任务的混合建模在多大程度上提升了模型性能？

主要发现

• SparseChem 在单张 NVIDIA V100 GPU 上，仅用 2 分钟即完成对 42 万种化合物和 3,500 项任务的训练，展现出卓越的训练速度。
• 在 CheMBL 29 数据集上，分类模型在每分钟 12 个周期的训练速度下，AUC-ROC 达到 0.794，AUC-PR 达到 0.717。
• 回归模型在相同硬件上实现 R² 为 0.358，皮尔逊相关系数为 0.620，训练速度达每分钟 15 个周期。
• 删失回归模型实现 R² 为 0.356，相关系数为 0.621，与标准回归模型性能差异极小。
• 该框架支持最多 30,000 项任务，并支持每项任务加权，可针对辅助任务或优先任务进行精细调优。
• 系统支持在 CPU 和 GPU 上进行完整训练和推理，通过命令行和 Python API 实现灵活集成。