[논문 리뷰] Knowledge-aware Contrastive Molecular Graph Learning

Yin Fang, Haihong Yang외 4명|arXiv (Cornell University)|2021. 03. 24.

Computational Drug Discovery Methods참고 문헌 19인용 수 10

한 줄 요약

이 논문은 대조적 지식 기반 GNN(CKGNN)을 제안하며, 화학 분야의 전문 지식(예: 지문과 기능기초)을 대조 학습을 통해 그래프 신경망에 명시적으로 통합하는 자기지도 학습 기반 분자 표현 학습 프레임워크이다. 이 방법은 유사한 구조를 가졌지만 기능적으로 다른 분자를 더 잘 구분할 수 있도록 하여, 여덟 개의 데이터셋에서 강력한 기준 모델 대비 평균 6%의 절대적 성능 향상을 달성한다.

ABSTRACT

Leveraging domain knowledge including fingerprints and functional groups in molecular representation learning is crucial for chemical property prediction and drug discovery. When modeling the relation between graph structure and molecular properties implicitly, existing works can hardly capture structural or property changes and complex structure, with much smaller atom vocabulary and highly frequent atoms. In this paper, we propose the Contrastive Knowledge-aware GNN (CKGNN) for self-supervised molecular representation learning to fuse domain knowledge into molecular graph representation. We explicitly encode domain knowledge via knowledge-aware molecular encoder under the contrastive learning framework, ensuring that the generated molecular embeddings equipped with chemical domain knowledge to distinguish molecules with similar chemical formula but dissimilar functions. Extensive experiments on 8 public datasets demonstrate the effectiveness of our model with a 6\% absolute improvement on average against strong competitors. Ablation study and further investigation also verify the best of both worlds: incorporation of chemical domain knowledge into self-supervised learning.

연구 동기 및 목표

• 분자의 기능적 변화(예: 기능기초 치환)와 같은 미세한 화학적 변화를 포착하지 못하는 GNN의 한계를 해결하기 위해.
• 지문과 기능기초와 같은 전문 지식을 그래프 신경망에 명시적으로 통합하여 분자 표현 학습을 향상시키기 위해.
• 분자 임베딩에서 원자 종류가 적고 빈도가 높은 원자(예: 탄소)로 인한 모호성을 줄이기 위해.
• 화학 지식을 활용한 자기지도 사전 훈련을 통해 최종 작업에서의 일반화 능력을 향상시키기 위해.
• 데이터나 구조에만 의존하는 모델과 비교해, 전문 지식과 대조 학습을 융합한 모델이 더 뛰어난 성능을 내는지 입증하기 위해.

제안 방법

• 기능기초와 지문과 같은 하위구조 수준의 전문 지식을 노드 및 그래프 임베딩에 명시적으로 인코딩하는 지식 기반 분자 인코더를 도입한다.
• 구조적 변형에 대해 불변인 표현을 학습하기 위해 지문 유사도 기반으로 분자 쌍을 비교하는 대조 학습 프레임워크를 활용한다.
• 유사한 지문을 가진 분자를 그룹화하는 분자 클러스터 전략을 사용하여 대조 학습을 위한 양성 쌍을 형성한다.
• 화학 전문 지식에 따라 유사한 분자의 임베딩 간 일치를 최대화하고, 비유사한 분자 간의 거리를 벌이는 대조 손실 함수를 설계한다.
• GNN 메시지 전달과 지식 기반 특징을 융합하여 잠재 공간에서 구조적 정보와 화학적 정보를 모두 유지한다.
• 240만 개 이상의 분자를 자기지도 방식으로 훈련하여, 최종적으로 성능 예측 작업으로의 전이를 가능하게 한다.

실험 결과

연구 질문

• RQ1 GNN에 화학 전문 지식을 명시적으로 통합하면, 단순히 구조적 정보만을 고려하는 것보다 분자 표현 학습 성능이 향상되는가?
• RQ2 지문 기반 양성 쌍을 활용한 대조 학습은 유사한 화학식을 가졌지만 기능이 다른 분자를 더 잘 구분할 수 있도록 모델의 능력을 향상시키는가?
• RQ3 지식 기반 특징 주입은 탄소와 같이 빈도가 매우 높은 원자에 대해 임베딩의 모호성을 어느 정도 줄이는가?
• RQ4 전문 지식과 대조 사전 훈련을 융합하면, 대규모 데이터나 구조 학습에만 의존하는 모델보다 성능이 뛰어나지는가?
• RQ5 개선된 임베딩 해석 해소 기능을 통해 모델은 동일한 화학식을 가졌지만 다른 기능기초를 가진 이성체 분자를 효과적으로 구분할 수 있는가?

주요 결과

• CKGNN는 강력한 경쟁 모델 대비 여덟 개의 벤치마크 데이터셋에서 평균 6%의 절대적 성능 향상을 달성한다.
• 제거 실험 결과, 지식 기반 특징이나 분자 클러스터 전략을 제거하면 성능이 크게 떨어지며, 이는 이 요소들이 핵심적인 역할을 한다는 것을 확인한다.
• CKGNN 임베딩의 유사도 분포는 분산이 0.067로 가장 크며, 이는 GCC와 사전 훈련된 GIN에 비해 더 풍부하고 다양한 표현을 가지는 것을 나타낸다.
• 사례 연구에서, 유사한 구조(예: 카복실기와 벤젠 고리)를 가진 분자를 유일하게 정확히 식별한 것은 CKGNN였으며, 단순히 구조에 의존하는 GNN은 실패했다.
• 기능 이성체의 경우, CKGNN는 GCC와 사전 훈련된 GIN(근처 1.0)에 비해 유사도 점수를 크게 낮춘다(예: 0.2661 및 0.1180), 이는 뛰어난 해석 해소 능력을 보여준다.
• 모델는 전문 지식 전용 또는 데이터 집약적 모델을 모두 능가하며, 지식과 자기지도 학습의 융합이 최적임을 입증한다.