[논문 리뷰] Accelerating Molecular Graph Neural Networks via Knowledge Distillation

Filip Ekström Kelvinius, Dimitar Georgiev외 2명|arXiv (Cornell University)|2023. 06. 26.

Machine Learning in Materials Science인용 수 4

한 줄 요약

이 논문은 구조적 변경 없이 분자 그래프 신경망(GNN)의 속도를 향상시키기 위해 노드-노드, 엣지-엣지, 엣지-노드, 벡터-벡터 등 지식 증류(KD) 전략을 제안한다. 대규모 정확도가 높은 교사 GNN에서 유사한 표현을 소규모 학생 모델로 증류함으로써, 에너지 예측에서는 최대 96.7%, 힘 예측에서는 최대 62.5%의 교사 성능을 달성하면서도 학생 모델의 추론 속도를 유지한다.

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의 효과를 평가하기 위해.
• 학생 모델 성능 향상을 위한 데이터 증강 및 학습 프로토콜 최적화를 탐색하기 위해.

제안 방법

• 방향성 및 등변성 GNN를 위한 특화된 네 가지 새로운 KD 전략(노드-노드(n2n), 엣지-엣지(e2e), 엣지-노드(e2n), 벡터-벡터(v2v))를 제안한다.
• 사전 훈련된 고정확도 교사 GNN을 활용해 에너지, 힘, 은닉 표현 등의 타겟 표현을 생성하여 지식 증류를 수행한다.
• 표준 회귀 손실과 은닉 표현에서 유도된 KD 손실을 조합한 다중 구성 요소 손실을 사용해 학생 GNN을 훈련한다.
• 스칼라 및 벡터 특징을 포함한 노드 수준 및 엣지 수준의 표현에 대해 증류를 적용하여 기하학적 및 구조적 정보를 유지한다.
• 데이터 혼합 비율 및 실수 데이터와 합성 데이터 간 손실 가중치 등 학습 하이퍼파ram터를 최적화한다.
• 랜덤 래터링 및 합성 데이터 생성(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 데이터 혼합 비율 및 손실 가중치와 같은 하이퍼파ram터가 증류 모델의 최종 성능에 어떤 영향을 미치는가?

주요 결과

• 제안된 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).