[论文解读] Probing Graph Representations

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

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

本文提出一种探针框架，用于研究图神经网络（GNNs）和图注意力网络（Graph Transformers）所学习的图表示中编码了哪些与化学相关的资讯。通过线性探针、贝叶斯互信息估计以及基于干预的成对探针方法，作者发现基于Transformer的模型所捕捉的化学意义信息显著多于消息传递型GNNs，甚至随机初始化的模型也编码了有用的特征，凸显了探针作为模型可解释性与开发诊断工具的重要性。

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.

研究动机与目标

• 探究在图表示中编码了哪些类型的化学相关资讯。
• 通过实证探针比较消息传递型GNNs与图注意力网络的表征能力。
• 评估架构选择（如跳跃连接和虚拟节点）对表征中信息编码的影响。
• 评估预训练是否为模型学习到有用结构与化学特征的必要条件。
• 倡导将探针作为图表示学习中诊断与可解释性的工具。

提出的方法

• 采用多策略探针框架：线性探针、贝叶斯互信息（BMI）估计，以及通过输入干预的成对探针。
• 使用RDKit在无监督条件下计算化学性质（例如官能团、原子计数、三维结构特征）。
• 应用增量探针以减少数据集偏差，并实现不同模型间的公平比较。
• 使用成对探针隔离特定结构变化（例如移除一个官能团）对表征偏移的影响。
• 通过在下游任务中评估探针性能来衡量表征的可迁移性。
• 利用预训练模型与随机初始化模型，评估训练在特征学习中的作用。

实验结果

研究问题

• RQ1 哪些类型的化学与结构性质在分子图表示中被可靠地编码？
• RQ2 图注意力网络与传统GNNs在捕捉化学相关资讯方面有何差异？
• RQ3 跳跃连接或虚拟节点等架构组件在多大程度上增强了表征中的信息编码？
• RQ4 未经训练（随机初始化）的模型是否能编码有用的化学特征，表明预训练并非表征质量的唯一驱动因素？
• RQ5 编码信息的丰富程度与下游可迁移性及模型泛化能力之间有何关联？

主要发现

• 基于Transformer的模型所编码的化学相关资讯显著多于消息传递型GNNs，尤其在官能团与三维结构特性方面。
• 即使随机初始化的模型也捕捉到了大量有用的化学信息，表明架构的归纳偏置本身即可产生信息丰富的表征。
• 跳跃连接使表征编码了更多资讯，表明其在保留结构与化学细节方面具有重要作用。
• 探针框架揭示，模型可学习到并非主预测任务严格所需之性质，提示可能存在虚假的归纳偏置。
• 编码信息的丰富程度（通过探针测量）与表征向下游任务的可迁移性之间存在强烈相关性。
• 成对探针证实，特定结构变化（例如官能团的删除）会导致表征出现可测量且可解释的偏移，验证了该方法的敏感性。