[논문 리뷰] Context-aware Graph Causality Inference for Few-Shot Molecular Property Prediction

Van Thuy Hoang, O. Lee|arXiv (Cornell University)|2026. 01. 16.

Advanced Graph Neural Networks인용 수 0

한 줄 요약

CaMol은 컨텍스트 그래프, 원자 마스킹, 백도어 보정을 사용하여 Few-shot 분자 특성 예측을 위한 컨텍스트 인식 인과 프레임워크를 도입하고 정확도와 해석가능성을 향상시킨다.

ABSTRACT

Molecular property prediction is becoming one of the major applications of graph learning in Web-based services, e.g., online protein structure prediction and drug discovery. A key challenge arises in few-shot scenarios, where only a few labeled molecules are available for predicting unseen properties. Recently, several studies have used in-context learning to capture relationships among molecules and properties, but they face two limitations in: (1) exploiting prior knowledge of functional groups that are causally linked to properties and (2) identifying key substructures directly correlated with properties. We propose CaMol, a context-aware graph causality inference framework, to address these challenges by using a causal inference perspective, assuming that each molecule consists of a latent causal structure that determines a specific property. First, we introduce a context graph that encodes chemical knowledge by linking functional groups, molecules, and properties to guide the discovery of causal substructures. Second, we propose a learnable atom masking strategy to disentangle causal substructures from confounding ones. Third, we introduce a distribution intervener that applies backdoor adjustment by combining causal substructures with chemically grounded confounders, disentangling causal effects from real-world chemical variations. Experiments on diverse molecular datasets showed that CaMol achieved superior accuracy and sample efficiency in few-shot tasks, showing its generalizability to unseen properties. Also, the discovered causal substructures were strongly aligned with chemical knowledge about functional groups, supporting the model interpretability.

연구 동기 및 목표

• Few-shot 분자 특성 예측(MPP)을 동기화하고 기능 그룹 인과관계 활용의 필요성을 제시한다.
• 컨텍스트 그래프를 통해 화학적 사전 정보를 통합하여 인과 서브구조를 발견하는 CaMol을 제안한다.
• 학습 가능한 원자 마스킹과 분포 기반 백도어 개입을 사용하여 인과 서브구조를 혼란변수로부터 구분한다.
• 발견된 서브구조를 화학 지식과 정렬시켜 해석가능성과 전이성을 향상시킨다.

제안 방법

• 각 에피소드 내에서 기능군, 분자 및 특성을 인코딩하는 컨텍스트 그래프를 구성한다.
• 분자를 BRICS 기반 기능군으로 분해하고 GNN 인코더를 통해 맥락 표현을 학습한다.
• 인과 서브구조 C를 교란 변수 S에서 분리하기 위한 학습 가능한 원자 마스킹 메커니즘을 도입한다.
• 화학적으로 근거 있는 교란을 사용하여 S에 대해 주변화를 통해 P(Y|do(C))를 추정하기 위한 백도어 보정을 포함한 분포 개입을 적용한다.
• 인과 예측 손실, S에 대한 균등 사전으로의 KL 발산, 개입 서브그래프 간의 분산/불변성 항을 결합한 총 손실을 최적화한다.
• 내부 루프의 인과 업데이트와 외부 루프 평가를 가진 MAML 스타일의 메타 학습을 사용하여 few-shot 일반화를 촉진한다.

Figure 1: (a) The seen properties are relevant to the unseen property prediction. (b) The causal substructures vary and depend on molecular property prediction tasks.

실험 결과

연구 질문

• RQ1 기능군, 분자 및 특성을 연결하는 컨텍스트 그래프가 어떻게 Few-shot 분자 특성 예측을 개선할 수 있는가?
• RQ2 학습 가능한 원자 마스킹이 분자 그래프에서 인과 서브구조를 혼란 변수로부터 효과적으로 구분해 낼 수 있는가?
• RQ3 백도어 보정 분포 개입이 분자와 특성 전반에 걸친 교란 변수에 대한 강건성을 향상시키는가?
• RQ4 발견된 인과 서브구조가 화학 지식과 일치하고 모델 해석 가능성을 높이는가?

주요 결과

• CaMol은 강력한 베이스라인에 비해 여섯 가지 MoleculeNet 데이터셋에서 few-shot 설정에서 우수한 정확도를 달성한다.
• 발견된 인과 서브구조는 알려진 기능군과의 강한 정합성을 보이며 해석가능성을 뒷받침한다.
• 이 프레임워크는 특히 다양성이 높은 데이터셋과 불균형 데이터에서 강한 샘플 효율성을 보이다 (예: MUV, PCBA).
• 맥락 지도를 가진 백도어 보정 인과 추론은 단독으로 분자-특성 관계에 의존하는 모델보다 더 강건한 예측을 산출한다.
• 이 접근법은 예측된 특성에 대해 충실하고 모델 일관된 설명을 제공한다.

Figure 2: Causal relationships between variables in MPP.