doi.org/10.48550/arXiv.2303.03951

本文介绍于 2023 年 德国亥姆霍兹信息安全中心研究团队发表在 AISTATS 2023 上的一篇文章，文章原标题为 Probing Graph Representations，文章设计了多种分子表示的探测模型，并通过探测模型研究了图模型在预训练后所编码分子信息。

随着基于图的深度学习模型不断出现，亟需回答的一个问题是「图模型将什么信息编码进了表示中？」为了研究这一问题，文章构建了探测模型测试预训练图模型得到的分子表示。

探测图表示的思路很简单，如果能从图模型输出的分子表示中提取出分子性质，那么就可以认为该性质被编码进分子表示中，所以文章的工作流程是「预训练-预测」（略不同于「预训练-微调」）。通过该流程，文章测试了传统 GNN 与基于 Transformer 的图模型等不同架构、不同数据集、不同优化算法等因素对于模型编码得到的潜变量的影响。

方法

在分子性质预测中，对于分子 \(\boldsymbol{x}\) 与其性质 \(y\)，完成该任务的模型就是映射 \(f:\boldsymbol{x}\mapsto y\)。取出 GNN 或图 Transformer 模型中 \(d\) 维的 \(l\) 层输出 \(f_l(\boldsymbol{x})=\boldsymbol{z}\)，该潜变量 \(\boldsymbol{z}\) 可以作为输入 \(\boldsymbol{x}\) 的一种表示，进一步得到 \(y\)。

文章使用不同的图模型得到分子表示，再通过另一模型测试分子表示 \(\boldsymbol{z}\) 预测分子性质 \(y\) 的性能，从而对比不同图模型提取特征信息的能力。

所构建的预测分子性质任务包括较为基础的判断是否具有某些官能团、更高层次的毒性、血脑屏障渗透性等。

探测策略

1. 线性探测（Linear Probing）：使用最简单的线性层，将分子表示映射为分子性质。
2. 贝叶斯探测（Bayesian Probing）：互信息可以用于 \(Z\) 与 \(P\) 两个随机变量之间的依赖程度，文中通过计算潜变量与分子性质间的贝叶斯互信息进行评估。
3. 成对探测（Pairwise Probing）：将结构相近而性质差异大的分子构成一对 \((\boldsymbol{x}_i,\boldsymbol{x}'_i)\)，通过主成分分析等方法分子潜变量与分子性质之间的关系。

实验

首先使用线性模型用 \(\boldsymbol{z}\) 预测了分子中是否具有某种子结构，结果如上图所示，基于 Transformer 的一类图模型显然具有比 GCN 和 GIN 具有更好的表现，同时 GCN 模型得到的表示又比以 Morgan 指纹作为分子表示更好。

在更高层次的分子性质数据集上测试各种分子表示，结果如上图所示，以 Morgan 指纹作为分子表示的任务效果比部分图模型更好，Morgan 指纹作为一种可以简单获得的分子表示，仍然适合用于许多机器学习模型中完成预测任务。

基于 Transformer 的图模型在更高层次的分子性质数据集上同样具有更好的表现，是具有潜力的新一代分子表示方式。这一点也可以从下图中看出，在左图中，基于 Transformer 图模型的结果都位于右上角，既能表示低层次的子结构信息，也能有效编码高层次的分子性质信息，而其他分子表示则位于左下角。右图使用贝叶斯互信息评估了样本数量与 \(Z\) 和 \(Y\) 之间的依赖程度的关系，就整体趋势而言，仍然是基于 Transformer 图模型效果更好。

最后，文章通过主成分分析评估了相似分子间不同的分子表示，两个相似分子仅在官能团上有所不同，文中选择的官能团为硝基。结果如下图左侧一列所示，with FG 表示含硝基分子，w/o FG 表示去除该官能团的分子，可以明显看出，相比于 GCN，GraphGPS 这一基于 Transformer 的图模型所产生的特征中，两种结构相似的分子也具有较大的区分子，是更好的分子表示。

结论

文章设计探测模型研究了图模型在预训练后编码的分子信息，最终发现相比于使用消息传递聚合信息的传统 GNN 模型，基于 Transformer 的图模型能够学习到更多与化学相关的化学信息，得到更好的分子表示。文章中提出的分析方法为预训练模型的测试以及分子表示的评估提供了指导。