[논문 리뷰] Fragment-based Pretraining and Finetuning on Molecular Graphs

Kha-Dinh Luong, Ambuj K. Singh|arXiv (Cornell University)|2023. 10. 05.

Computational Drug Discovery Methods인용 수 4

한 줄 요약

이 논문은 분자 그래프를 위한 프ragment 기반 사전 훈련 프레임워크인 GraphFP를 제안한다. 이는 구조적 정보를 다중 해상도에서 포괄적으로 포착하기 위해 대비 학습을 통해 분자 및 프ragment GNN을 동시에 훈련시킨다. 화학적으로 의미 있는 프ragment의 압축된 어휘를 추출하고, 프ragment 임베딩과 그 구성 원자들의 집계된 노드 임베딩 간의 일관성을 강제함으로써, 8개의 분자 벤치마크 중 5개에서 성능 향상을 이끌어내며, 장거리 생물학적 벤치마크 성능을 최소 11.5% 향상시킨다.

ABSTRACT

Property prediction on molecular graphs is an important application of Graph Neural Networks. Recently, unlabeled molecular data has become abundant, which facilitates the rapid development of self-supervised learning for GNNs in the chemical domain. In this work, we propose pretraining GNNs at the fragment level, a promising middle ground to overcome the limitations of node-level and graph-level pretraining. Borrowing techniques from recent work on principal subgraph mining, we obtain a compact vocabulary of prevalent fragments from a large pretraining dataset. From the extracted vocabulary, we introduce several fragment-based contrastive and predictive pretraining tasks. The contrastive learning task jointly pretrains two different GNNs: one on molecular graphs and the other on fragment graphs, which represents higher-order connectivity within molecules. By enforcing consistency between the fragment embedding and the aggregated embedding of the corresponding atoms from the molecular graphs, we ensure that the embeddings capture structural information at multiple resolutions. The structural information of fragment graphs is further exploited to extract auxiliary labels for graph-level predictive pretraining. We employ both the pretrained molecular-based and fragment-based GNNs for downstream prediction, thus utilizing the fragment information during finetuning. Our graph fragment-based pretraining (GraphFP) advances the performances on 5 out of 8 common molecular benchmarks and improves the performances on long-range biological benchmarks by at least 11.5%. Code is available at: https://github.com/lvkd84/GraphFP.

연구 동기 및 목표

• 분자 성질 예측을 위한 GNN에서 노드 수준 및 그래프 수준 사전 훈련의 한계를 해결한다.
• 기존의 프ragment 기반 방법에서의 과도한 스무딩과 구조적 정밀도 손실 문제를 해결하기 위해 이중 인코더 프레임워크를 도입한다.
• 데이터 증강이나 3D 좌표와 같은 사전 정보 없이도 화학적으로 충실한 자기지도 학습 전략을 개발한다.
• 프ragment 수준의 표현 학습을 통해 다양한 분자 및 장거리 생물학적 벤치마크에서의 최종 성능 향상을 도모한다.
• 주요 부분그래프 추출을 통한 원칙적인 프ragment 어휘를 활용해 계층적인 구조 표현을 학습함으로써 더 나은 일반화를 가능하게 한다.

제안 방법

• 대규모 사전 훈련 데이터셋에서 흔한 프ragment의 압축적이고 다양한, 화학적으로 의미 있는 어휘를 추출하기 위해 주요 부분그래프 추출 기법을 사용한다.
• 노드 수준의 분자 그래프와 고차원 연결성의 프ragment 그래프에 각각 별도의 GNN을 훈련시키며, 대비 학습 목표를 적용한다.
• 분자 그래프 내에서 프ragment 임베딩과 그 구성 원자들의 집계된 임베딩 간의 일관성을 강제하여 다중 척도 표현을 정렬한다.
• 프ragment 그래프에서 유도된 레이블(예: 구조적 특성)을 사용하는 예측 사전 훈련 과제를 도입하여 분자 GNN의 이해도를 향상시킨다.
• 微조정 단계에서 사전 훈련된 분자 및 프ragment GNN을 모두 활용하여 프ragment 수준의 구조 지식을 유용하게 한다.
• 프ragment 그래프를 대비 학습 외에도 그래프 수준의 예측 사전 훈련을 위한 보조 레이블 생성에 활용한다.

Figure 1: Fragment-based contrastive pretraining framework. $\textsc{GNN}_{M}$ processes molecular graphs while $\textsc{GNN}_{F}$ processes fragment graphs. The fragment-based pooling function FragPool aggregates node embeddings into a combined embedding that forms a positive contrastive pair with

실험 결과

연구 질문

• RQ1 노드 수준 또는 그래프 수준 사전 훈련과 비교했을 때, 프ragment 수준 사전 훈련이 분자 그래프 표현 학습에 어떤 영향을 미치는가?
• RQ2 대비 학습을 통해 분자 및 프ragment GNN을 공동으로 사전 훈련할 경우, 다양한 분자 벤치마크에서 최종 성능에 어떤 영향을 미치는가?
• RQ3 프ragment 어휘의 크기와 품질이 자기지도 사전 훈련의 효과성에 어떤 영향을 미치는가?
• RQ4 프ragment 기반 표현이 장거리 생물학적 성질 예측 과제에서 성능 향상에 기여하는가?
• RQ5 데이터 증강이나 3D 좌표와 같은 사전 정보가 없는 상황에서, 학습된 표현의 타당성과 일반화 능력에 어떤 영향을 미치는가?

주요 결과

• GraphFP는 8개의 표준 분자 성질 예측 벤치마크 중 5개에서 최신 기준 또는 경쟁 가능한 성능을 달성한다.
• 기존의 기준 방법 대비 장거리 생물학적 벤치마크 성능을 최소 11.5% 향상시킨다.
• t-SNE 시각화 결과, 프ragment 임베딩이 고차원 구조 패턴을 잘 포착하고 있으며, 분자 기초 구조에 따라 명확한 군집이 형성됨을 확인할 수 있다.
• 동일한 프ragment에 속한 원자들의 임베딩이 잠재 공간에서 서로 가까이 모여 있음을 확인하여, 풀링된 노드 표현이 고차원 연결성을 효과적으로 캡처하고 있음을 시사한다.
• 최적 점(예: 3,200)을 초과해 프ragment 어휘 크기를 늘일 경우 성능 저하가 발생하며, 이는 파rameter 수 증가와 구조적 정밀도 손실로 인한 것으로 보인다.
• GNN 아키텍처의 선택이 성능에 상당한 영향을 미치며, GIN 아키텍처가 다양한 벤치마크에서 평균적으로 가장 우수한 성능을 보였다.

Figure 2: Fragment-based predictive pretraining. $\textsc{GNN}_{M}$ processes molecular graphs and produces graph-level embeddings used for prediction. The right upper box shows unique fragments that exist in the input molecule while the right lower box shows the ground-truth structural backbone. A