Peeling Context from Cause for Molecular Property Prediction

TL;DR

This paper introduces CLaP, a layerwise framework that isolates causal signals from spurious context in molecular models, enhancing interpretability and predictive accuracy across multiple benchmarks.

Contribution

CLaP is a novel method that progressively peels away context to focus on causal structure in molecular property prediction, improving both interpretability and performance.

Findings

CLaP improves predictive metrics like MAE, MSE, and R^2 across four benchmarks.

It generates atom-level causal saliency maps aligned with chemical intuition.

The approach enhances model interpretability for molecular design.

Abstract

Deep models are used for molecular property prediction, yet they are often difficult to interpret and may rely on spurious context rather than causal structure, which reduces reliability under distribution shift and harms predictive performance. We introduce CLaP (Causal Layerwise Peeling), a framework that separates causal signal from context in a layerwise manner and integrates diverse graph representations of molecules. At each layer, a causal block performs a soft split into causal and non-causal branches, fuses causal evidence across modalities, and progressively removes batch-coupled context to focus on label-relevant structure, thereby limiting shortcut signals and stabilizing layerwise refinement. Across four molecular benchmarks, CLaP consistently improves MAE, MSE, and $R^2$ over competitive baselines. The model also produces atom-level causal saliency maps that highlight substructures responsible for predictions, providing actionable guidance for targeted molecular edits. Case studies confirm the accuracy of these maps and their alignment with chemical intuition. By peeling context from cause at every layer, the model yields predictors that are both accurate and interpretable for molecular design.