[Paper Review] Probing Graph Representations

Mohammad Sadegh Akhondzadeh, Vijay Lingam+1 more|arXiv (Cornell University)|Mar 7, 2023

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper introduces a probing framework to investigate what chemically relevant information is encoded in graph representations learned by GNNs and Graph Transformers. Using linear probing, Bayesian mutual information, and intervention-based paired probing, the authors find that transformer-based models capture significantly more chemically meaningful information than message-passing GNNs, and even randomly initialized models encode useful features, highlighting probing as a diagnostic tool for model interpretability and development.

ABSTRACT

Today we have a good theoretical understanding of the representational power of Graph Neural Networks (GNNs). For example, their limitations have been characterized in relation to a hierarchy of Weisfeiler-Lehman (WL) isomorphism tests. However, we do not know what is encoded in the learned representations. This is our main question. We answer it using a probing framework to quantify the amount of meaningful information captured in graph representations. Our findings on molecular datasets show the potential of probing for understanding the inductive biases of graph-based models. We compare different families of models and show that transformer-based models capture more chemically relevant information compared to models based on message passing. We also study the effect of different design choices such as skip connections and virtual nodes. We advocate for probing as a useful diagnostic tool for evaluating graph-based models.

Motivation & Objective

• To investigate what types of chemically relevant information are encoded in learned graph representations.
• To compare the representational capacity of message-passing GNNs versus Graph Transformers using empirical probing.
• To evaluate the impact of architectural choices—such as skip connections and virtual nodes—on information encoding in representations.
• To assess whether pretraining is necessary for models to learn useful structural and chemical features.
• To advocate for probing as a diagnostic and explainability tool in graph representation learning.

Proposed method

• Employing a multi-strategy probing framework: linear probing, Bayesian mutual information (BMI) estimation, and paired probing via input intervention.
• Using RDKit to compute chemical properties (e.g., functional groups, atom counts, 3D structure features) without supervision.
• Applying incremental probing to reduce dataset bias and enable fair comparison across models.
• Using paired probing to isolate the effect of specific structural changes (e.g., removing a functional group) on representation shifts.
• Measuring transferability of representations by evaluating probing performance on downstream tasks.
• Leveraging pre-trained and randomly initialized models to assess the role of training in feature learning.

Experimental results

Research questions

• RQ1 Which types of chemical and structural properties are reliably encoded in graph representations of molecules?
• RQ2 How do Graph Transformers compare to traditional GNNs in capturing chemically relevant information?
• RQ3 To what extent do architectural components like skip connections or virtual nodes enhance information encoding in representations?
• RQ4 Can untrained (randomly initialized) models encode useful chemical features, suggesting that pretraining may not be the sole driver of representation quality?
• RQ5 How does the richness of encoded information correlate with downstream transferability and model generalization?

Key findings

• Transformer-based models encode significantly more chemically relevant information than message-passing GNNs, particularly for functional groups and 3D structural properties.
• Even randomly initialized models capture a substantial amount of useful chemical information, suggesting that architectural inductive biases alone can yield informative representations.
• Skip connections lead to representations that encode more information, indicating their role in preserving structural and chemical details.
• The probing framework reveals that models can learn properties not strictly required for the main prediction task, indicating potential for spurious inductive biases.
• There is a strong correlation between the richness of encoded information (as measured by probing) and the transferability of representations to downstream tasks.
• Paired probing confirms that specific structural changes (e.g., deletion of a functional group) lead to measurable and interpretable shifts in representations, validating the method’s sensitivity.