通过知识蒸馏加速分子图神经网络
机器学习
2023-10-31 v2 化学物理
摘要
图神经网络（GNNs）的最新进展使得对分子及分子系统的建模更加全面，从而提升了分子性质预测与分子模拟的精度。然而，随着该领域向更大、更复杂的架构发展，最先进的 GNNs 对于许多大规模应用而言已变得难以负担。本文探讨了知识蒸馏（KD）在加速分子 GNNs 方面的效用。为此，我们设计了能够促进方向性与等变 GNNs 中隐藏表示蒸馏的 KD 策略，并在能量与力预测的回归任务上评估其性能。我们在不同的教师-学生配置与数据集上验证了我们的方案，并证明它们能在不修改学生模型架构的情况下持续提升其预测精度。此外，我们对框架中各个组件进行了全面优化，并探究了数据增强进一步提升性能的潜力。总体而言，我们在完全保持更轻量模型推理吞吐量的同时，将教师与学生模型在能量与力预测上的预测精度差距分别缩小了多达 96.7% 和 62.5%。
引用
@article{arxiv.2306.14818,
title = {Accelerating Molecular Graph Neural Networks via Knowledge Distillation},
author = {Filip Ekström Kelvinius and Dimitar Georgiev and Artur Petrov Toshev and Johannes Gasteiger},
journal= {arXiv preprint arXiv:2306.14818},
year = {2023}
}
备注
Accepted as a conference paper at NeurIPS 2023