[论文解读] Accelerating Molecular Graph Neural Networks via Knowledge Distillation

Filip Ekström Kelvinius, Dimitar Georgiev+2 more|arXiv (Cornell University)|Jun 26, 2023

Machine Learning in Materials Science被引用 4

一句话总结

该论文提出了一种知识蒸馏（KD）策略——节点对节点（n2n）、边对边（e2e）、边对节点（e2n）和向量对向量（v2v）——以在不改变网络架构的前提下加速分子图神经网络（GNN）。通过将大型高精度教师GNN的隐藏表征蒸馏到小型学生模型中，该方法在能量和力预测任务上分别实现了教师模型62.5%和96.7%的性能，同时保持了学生模型的推理速度。

ABSTRACT

Recent advances in graph neural networks (GNNs) have enabled more comprehensive modeling of molecules and molecular systems, thereby enhancing the precision of molecular property prediction and molecular simulations. Nonetheless, as the field has been progressing to bigger and more complex architectures, state-of-the-art GNNs have become largely prohibitive for many large-scale applications. In this paper, we explore the utility of knowledge distillation (KD) for accelerating molecular GNNs. To this end, we devise KD strategies that facilitate the distillation of hidden representations in directional and equivariant GNNs, and evaluate their performance on the regression task of energy and force prediction. We validate our protocols across different teacher-student configurations and datasets, and demonstrate that they can consistently boost the predictive accuracy of student models without any modifications to their architecture. Moreover, we conduct comprehensive optimization of various components of our framework, and investigate the potential of data augmentation to further enhance performance. All in all, we manage to close the gap in predictive accuracy between teacher and student models by as much as 96.7% and 62.5% for energy and force prediction respectively, while fully preserving the inference throughput of the more lightweight models.

研究动机与目标

• 解决最先进分子GNN日益增长的计算成本问题，该问题限制了其在大规模模拟中的应用。
• 通过在不修改架构的前提下加速轻量级GNN，改善分子性质预测中的速度-精度权衡。
• 为分子模拟中使用的定向和等变GNN开发专门的KD策略。
• 在包括OC20-2M和COLL在内的多种教师-学生配置和数据集上评估KD的有效性。
• 研究数据增强和训练协议优化，以进一步提升学生模型的性能。

提出的方法

• 提出四种新颖的KD策略：面向等变GNN的节点对节点（n2n）、边对边（e2e）、边对节点（e2n）和向量对向量（v2v）蒸馏。
• 使用预训练的高精度教师GNN生成目标表征（能量、力、隐藏状态）用于知识蒸馏。
• 通过结合标准回归损失与来自隐藏表征的KD损失，训练学生GNN。
• 对节点级和边级表征（包括标量和向量特征）应用蒸馏，以保留几何和结构信息。
• 优化训练超参数，包括真实数据与合成数据的混合比例以及真实与合成数据之间的损失权重。
• 通过随机扰动和合成数据生成（d1M）探索数据增强，以提升学生模型的泛化能力。

Figure 1: Using knowledge distillation, we manage to significantly boost the predictive accuracy of different student models on the OC20-2M [ 17 ] and COLL [ 6 ] datasets while fully preserving their inference throughput.

实验结果

研究问题

• RQ1 知识蒸馏是否能有效加速分子GNN，同时保持推理速度？
• RQ2 不同基于特征的蒸馏策略（n2n、e2e、e2n、v2v）对能量和力预测任务中学生模型性能的影响如何？
• RQ3 在分子模拟中，蒸馏能在多大程度上缩小大教师模型与小学生模型之间的准确率差距？
• RQ4 数据增强（如随机扰动或合成数据）是否能进一步提升学生模型的性能？
• RQ5 学习率、数据混合比例和损失权重等超参数如何影响蒸馏模型的最终性能？

主要发现

• 所提出的KD策略在不改变架构的前提下，一致地提升了多种数据集和教师-学生配置下学生模型的性能。
• 该方法在能量预测任务上实现了教师模型96.7%的性能，在力预测任务上实现了62.5%的性能，显著缩小了准确率差距。
• 推理吞吐量完全得以保留，因为学生模型轻量化且推理高效。
• 通过随机扰动进行的数据增强未带来显著改进，基于梯度的扰动效果甚至不如随机噪声。
• 与其它领域预期相反，合成数据生成（d1M）在性能上并未显著优于仅使用真实数据的联合训练基线。
• 全面的消融实验验证了各KD组件的有效性，其中向量对向量（v2v）和节点对节点（n2n）策略表现出显著增益。

Figure 2: Similarity analysis between the node features of SchNet, PaiNN and GemNet-OC using CKA (averaged over $n=987$ nodes).