Context-aware Graph Causality Inference for Few-Shot Molecular Property Prediction

TL;DR

CaMol tackles few-shot molecular property prediction by linking chemical priors to causal inference through a context graph that connects functional groups, molecules, and properties. It introduces a backdoor-adjusted causality framework with a learnable atom masking mechanism to extract a causal substructure C while treating remaining atoms as confounders S, and it employs a distribution intervention to robustly estimate the effect of C on the target Y within a meta-learning setting. The approach combines context-graph learning, a causal substructure extractor, and semantic interventions within a MAML-style training loop, achieving superior few-shot accuracy and interpretable substructures that align with chemical knowledge. Experiments across six MoleculeNet benchmarks demonstrate strong performance gains, improved interpretability, and favorable efficiency, suggesting practical impact for rapid, data-scarce molecular property prediction.

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.

Figures (9)

• 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.
• Figure 2: Causal relationships between variables in MPP.
• Figure 3: Overview of the CaMol architecture, consisting of a causal substructure extractor and a distribution intervener.
• Figure 4: A sample efficiency comparison on Few-shot MPP in terms of ROC-AUC according to the number of labeled samples.
• Figure 5: Heatmaps of property–property similarities obtained from (left) representations of causal substructure discovered by CaMol and (right) co-occurrence frequencies of functional groups. The two similarities are highly correlated (Pearson $r = 0.67$, Spearman $\rho = 0.74$; both $p < 0.01$).