[논문 리뷰] Motif-aware Attribute Masking for Molecular Graph Pre-training

Eric Inae, Gang Liu외 1명|arXiv (Cornell University)|2023. 09. 08.

Machine Learning in Materials ScienceMaterials Science인용 수 3

한 줄 요약

이 논문은 화학적으로 의미 있는 모티프 내의 모든 노드 특징을 마스킹하여 상호 모티프 지식 전이를 향상시키는 분자 그래프 사전학습을 위한 모티프 인식 특성 마스킹 전략 MoAMa를 제안한다. 모델이 교차 모티프 정보를 사용하여 특징을 재구성하도록 유도함으로써, MoAMa는 8개의 분자 성질 예측 데이터셋에서 최신 기준(SOTA) 기반보다 평균 1.3% 향상된 예측 성능을 달성한다.

ABSTRACT

Attribute reconstruction is used to predict node or edge features in the pre-training of graph neural networks. Given a large number of molecules, they learn to capture structural knowledge, which is transferable for various downstream property prediction tasks and vital in chemistry, biomedicine, and material science. Previous strategies that randomly select nodes to do attribute masking leverage the information of local neighbors However, the over-reliance of these neighbors inhibits the model's ability to learn from higher-level substructures. For example, the model would learn little from predicting three carbon atoms in a benzene ring based on the other three but could learn more from the inter-connections between the functional groups, or called chemical motifs. In this work, we propose and investigate motif-aware attribute masking strategies to capture inter-motif structures by leveraging the information of atoms in neighboring motifs. Once each graph is decomposed into disjoint motifs, the features for every node within a sample motif are masked. The graph decoder then predicts the masked features of each node within the motif for reconstruction. We evaluate our approach on eight molecular property prediction datasets and demonstrate its advantages.

연구 동기 및 목표

• 그래프 신경망 사전학습에서 무작위 특성 마스킹의 한계를 해결하기 위해, 지역 이웃 특징에 과도하게 의존하고 상호 모티프 구조 지식을 포착하지 못하는 문제를 해결한다.
• 사전학습 중에 화학 모티프 간의 지식 전달을 명시적으로 개선하여 분자 성질 예측 성능을 향상시킨다.
• 모티프 내 특징 전파에 대한 의존도를 줄이고 GNN의 메시지 전달 병목 현상을 완화한다.
• 모티프 내 구조 패tern과 상호 모티프 의존성 양측을 활용하는 사전학습 전략을 설계한다.
• 다양한 분자 성질 예측 벤치마크에서 모티프 인식 마스킹의 효과를 평가한다.

제안 방법

• 모든 분자 그래프를 사전 정의된 모티프 어휘를 사용하여 상호 배타적인 화학 모티프로 분해한다.
• 각 모티프 내 모든 노드 특징을 학습 가능한 마스크 토큰으로 마스킹하여, 모델이 교차 모티프 정보를 사용해 재구성하도록 유도한다.
• 그래프 디코더는 이웃 모티프의 맥락과 내부 모티프 구조를 기반으로 마스킹된 노드 특징을 예측한다.
• 모델은 모티프 간 마스킹된 속성을 재구성하는 목적을 가진 마스킹 오토에인코더 목적함수를 사용해 사전학습한다.
• 상호 모티프 지식 전이 정도는 MRR 기반 지표로 측정되며, 낮은 점수일수록 강한 상호 모티프 영향을 의미한다.
• 주로 AUC를 주요 지표로 사용하여 8개의 분자 성질 예측 데이터셋에서 방법을 평가한다.

Figure 1: Our MoAMa masks every node in a motif to pre-train GNNs. The full masking of a motif forces the GNNs to learn to (1) pass feature information across motifs and (2) pass local structural information within the motif. Compared to the traditional random attribute masking strategies, the motif

실험 결과

연구 질문

• RQ1 무작위 마스킹과 비교해 모티프 인식 특성 마스킹이 분자 그래프 사전학습에서 상호 모티프 지식 전이를 향상시키는가?
• RQ2 모티프 내 모든 특징을 명시적으로 마스킹함으로써 지역 이웃 특징에 대한 의존도가 감소하고 메시지 전달 병목 현상이 완화되는가?
• RQ3 MoAMa의 사전학습 전략이 최종 분자 성질 예측 성능 향상에 어느 정도 기여하는가?
• RQ4 다양한 모델 간에 상호 모티프 지식 전이와 예측 성능 간의 상관관계는 어떻게 되는가?
• RQ5 Tanimoto 유사도나 도메인 전용 히우리스틱에 의존하지 않고도 최신 기준 사전학습 방법과 경쟁 가능한 성능을 달성할 수 있는가?

주요 결과

• MoAMa는 8개의 분자 성질 예측 데이터셋에서 평균 테스트 AUC 75.3을 기록하여 최고의 베이스라인 대비 1.3% 향상된 성능을 달성했다.
• 모델은 높은 상호 모티프 영향 비율(노드 수준 0.80, 그래프 수준 0.51)을 보이며, 모티프 간 지식 전이가 강화됨을 나타냈다.
• MoAMa의 MRR 점수는 가장 낮았으며, 각각 노드 수준 0.59, 그래프 수준 0.55, 모티프 수준 0.41로, 모티프 내 특징에 대한 의존도 감소와 함께 강화된 상호 모티프 지식 전이를 확인했다.
• 예측 AUC와 상호 모티프 영향 비율 간에 중간 정도의 양의 상관관계(r = 0.52–0.55)를 관찰하여, 상호 모티프 학습과 성능 향상 간의 연관성을 뒷받침했다.
• AUC와 MRR 점수 간에 중간 정도의 음의 상관관계(r = -0.39 ~ -0.44)를 관찰하여, 모티프 내 특징 의존도 감소가 최종 정확도 향상에 기여함을 추가로 확인했다.
• 모델는 명시적인 모티프 내 구조와 상호 모티프 의존성 인코딩을 통해 학습 붕괴를 효과적으로 완화하고 모델 표현력을 향상시켰다.

Figure 2: Inter-motif knowledge transfer score by motif count. A higher $\text{MRR}_{\text{inter}}^{(n)}$ score denotes greater inter-motif knowledge transfer.