论文解读（GSAT）《Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism》

论文信息

论文标题：Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism
论文作者：Siqi Miao, Mia Liu, Pan Li
论文来源：2022,ICML
论文地址：download 
论文代码：download

1 Introduction

1.1 引入

　　GNN的可解释性问题：通常旨在从原始的输入图中提取一个子图：人们希望提取的子图中仅包含最能帮助标签预测的信息。

　　例子：如下图，我们知道-OH官能团能够使得一个分子具有水溶性。因此对于一个用来预测分子水溶性的GNN来说，人们希望给定下图的分子后，模型能够告诉我们对预测最重要的部分是-OH官能团所代表的子图。这样一来，人们就能从模型中获取更多的关于数据关键特征的理解，从而指导进一步的研究。​

　　

　　GNN可解释性主要有两大类方法：

1. 固有的可解释模型（Inherently Interpretable Models）；
2. 事后解释方法（Post-hoc Interpretation Methods）；

　　第一类方法主要旨在设计自身即可提供解释性的GNN模型。这些模型的良好可解释性往往以其预测精度为代价，如注意力机制，多篇研究显示其无法为GNN带来值得信任的可解释结果；

　　第二类方法，这些工作通常假设人们会提供一个预先训练好的GNN。随后它们会将该GNN的参数固定，然后训练一个新的模型，即解释器（Explainer），来从输入图中找出一个子图。它们希望这些子图能够：1）尽可能小；2）尽可能保持原有预测分数。最后这些子图即被认为是GNN捕捉到的数据的关键特征。

　　最近，基于不变因果特征学习（Invariant Learning）的工作也逐渐被提出。这些工作认为训练数据中可能会存在数据偏见（Data Bias），使得模型最终学习到一些和标签具有伪相关性（Spurious Correlations）的特征，并因此遭受严重的泛化问题。下图展示了伪相关特征的一个例子。这些特征可能是收集或生成训练数据时的偏见造成的，它们实质上并不是真正决定样本类别的特征。而当测试集不存在这些伪相关的特征时，模型的效果将大打折扣。因此，这些工作引入了因果分析理论（Causality Analysis），希望迫使模型学习数据中不变的、与标签具有因果关系的特征（Invariant Causal Patterns），来解决上述 OOD 泛化问题（Out-of-distribution Generalization）。这类方法在寻找那些不变的因果特征时，也能提供一定程度的自身可解释性。不过也由于这些方法引入了因果分析，它们的架构往往十分复杂且需要大量的计算。

　　

　　而在这篇工作中，作者们指出了事后解释方法的诸多问题，并同样专注于设计自身可解释模型。这篇工作提出了一种全新的随机注意力机制（Stochastic Attention Mechanism），该机制显示出了强大的可解释能力和泛化能力。对比过去的可解释工作，该机制在6个数据集上提升了至多20%、平均12%的可解释性能；在11个数据集上提升了平均3%的模型准确率，并且在OGBG-MolHiv榜单上达到SOTA（在不使用手工设计的专家特征的模型中）。

　　除此之外，该机制对可解释能力和泛化能力的提升同样具有理论保障。在一定假设下，该机制天然的不受伪相关特征的影响，从而能够抓取出真正重要的数据特征。在去除伪相关特征的能力上，该机制以远远更小的复杂度，对比基于因果分析的方法提升了平均12%的OOD泛化能力。

1.2 事后解释方法的问题

　　作者们在文中首先指出了事后可解释方法的四个问题，并认为这些事后解释方法擅长于检查预先训练好的模型对一些特征的敏感程度，但它们并不能提取出对预测真正重要的数据特征，而这才应该是可解释方法需要解决的最有趣的问题。具体来说，作者们指出的四个问题是：

1. 数据分布偏移（Data Distribution Shifts）

　　首先，事后解释方法将不可避免的遭受数据分布偏移的影响。直觉上，这是因为给定的预先训练好的模型（记作），总是在原始输入图上进行训练的：它从来没有在任何子图上进行过训练。因而极有可能在上是欠拟合的，故而导致并不能真正反应各个子图的重要性。

2. 与标签伪相关的数据特征（Spuriously Correlated Patterns）

　　其次，预先训练得到的模型可能会过拟合训练数据中与标签信息伪相关，甚至是无关的特征。这是由于大多数模型本身是基于最大化互信息法则（Maximum Mutual Information Principle）来进行训练的，因此在训练中自然会捕捉尽可能多的输入特征，而这也是不变因果特征学习这个方向产生的主要动机。

　　在这种情况下，事后解释方法很可能会将这些伪相关或者无关的特征提取出来，当作数据中的关键特征，而这可能会将人们引入到一个错误的方向。【自编码器预训练 接 线性分类器.....】

3. 初始化问题（Initialization Issuses）

　　随后，作者们从优化和信息瓶颈理论切入，指出事后解释方法对不同的的初始化是敏感的。在同一个数据集上，基于不同的随机种子训练得到的，事后解释方法可能会得出差异较大解释结果。而过去的事后解释方法，在评估时往往会忽略这一点，只基于某一个固定的，仅在不同的随机种子上训练解释器。这可能会得到过于乐观的结果，而使得事后解释方法的性能没有得到全面的评估。【使用不同随机种子训练的，我就看到一个论文这样用......】

4. 潜在的有偏见的约束（Potentially Biased Constraints）

　　最后，由于上述各种问题，事后解释方法有时很难得出符合人们直觉的解释子图。故这些方法中往往嵌入稀疏化约束（Sparsity Constraint），或连接性约束（Connectivity Constraint）等，来得到人们更能理解的数据特征。这些约束极大的要求人们对数据集和任务自身具有一定的先验知识，否则这些约束很可能极大的影响模型的解释结果。一个优秀的可解释模型应当自身即能够抓取适当的数据关键特征而不用附加其它约束。本文提出的随机注意力机制能够在没有上述约束的情况下，取得远远更好的可解释性能。【图结构学习吧，很重要的一点就是加图正则化.......】

1.3 Preliminaries

Label-relevant Subgraph

　　指：子图和图标签相关。如：为确定分子的溶解度，羟基-OH 是一个正相关的子图，就好像它存在，分子通常溶于水。寻找与标签相关的子图是可解释图学习的一个共同目标。

Attention Mechanism

　　具有注意力的 GNN 经常会产生低保真度的注意力权重，因为它为每条边学习多个权值，因此将这些权值与不规则的图结构结合起来，进行与图标签相关的特征选择是有问题的。

　　两种类型的注意力模型：

• • 注意力权重归一化，求和为 1 ，如 GAT；
  • 在没有归一化的情况下学习 $[0,1]$ 之间的权重，如 GGNN；

Learning to Explain (L2X)

　　L2X [ 2018 ] 研究了正则特征空间中的特征选择问题，并提出了一种互信息(MI)最大化规则来选择固定数量的特征。

　　具体来说，让 $I(a ; b) \triangleq \sum\limits _{a, b} \mathbb{P}(a, b) \log \frac{\mathbb{P}(a, b)}{\mathbb{P}(a) \mathbb{P}(b)}$  表示两个随机变量 $a$ 和 $b$ 之间的 MI。较大的 MI 表明两个随机变量之间存在一定的高相关性。因此，根据输入特征 $X \in \mathbb{R}^{F}$，L2X 是搜索一组大小为 $k$  的索引 $S \subseteq\{1,2, \ldots, F\}$，其中 $k=|S|<F$。使由索引 $S$ (用 $X_{S}$ 表示) 的子空间中的特征与标签 $Y$ 的互信息最大化，即，

　　　　$\underset{S \subseteq\{1,2, \ldots, F\}}{\text{max}}I\left(X_{S} ; Y\right), \quad \text { s.t. }|S| \leq k .\quad\quad\quad(1)$

　　本文模型的灵感来自于 L2X ，由于 图特征 及 可解释的特征 并没有固定的维数，所以不能直接应用 L2X ，因此，本文建议使用 Sec 3.1 中使用信息约束。

　　熵定义为 ：$H(a)   \triangleq-\sum\limits _{a} \mathbb{P}(a) \log \mathbb{P}(a)$

　　KL 散度定义为： $\mathrm{KL}(\mathbb{P}(a) \| \mathbb{Q}(a)) \triangleq \sum\limits _{a} \mathbb{P}(a) \log \frac{\mathbb{P}(a)}{\mathbb{Q}(a)}$

　　信息瓶颈 IB ：$\underset{Z}{\text{max }} \mathcal{L}_{I B}=I(Z ; Y)-\beta I(X ; Z)$

3 Graph Learning Interpretation via GIB

3.1 GIB-based Objective for Interpretation

　　关键子图可能具有不用的大小，以及高度不规则，所以本文建议使用信息约束来选择与标签相关的子图：

　　　　$\begin{array}{l}\underset{G_{S}}{\text{max}} &  \quad I\left(G_{S} ; Y\right) \\\text {s.t. }&  \quad I\left(G_{S} ; G\right) \leq \gamma,\quad G_{S} \in \mathbb{G}_{s u b}(G)\end{array}\quad\quad\quad(2)$

　　其中 $\mathbb{G}_{s u b}(G)$ 表示 $G$ 的子图的集合。

　　注意，GIB 不施加任何潜在的偏差约束，如所选子图的大小或连通性。相反，GIB 使用信息约束 $I\left(G_{S} ; G\right) \leq \gamma$ 来选择只从 $G$ 中继承了最具指示性的信息的 $G_{S}$，通过最大化 $I\left(G_{S} ; Y\right)$ 来预测标签 $Y$。因此，GS提供了模型解释。

3.2 Issues of Post-hoc GNN Interpretation Methods

　　几乎所有之前的方法都是事后方法，如 GNNExplainer [2019 ]，PGExplainer [2020] 和 GraphMask [2021 ]。给定一个预先训练的预测器 $f_{\theta}(\cdot): \mathcal{G} \rightarrow \mathcal{Y}$，他们试图找出影响模型预测最大的子图 $G_{S}$，同时保持预先训练的模型不变。这个过程首先最大化了 $f_{\theta}(G)$ 和 $Y$ 之间的 MI，并得到模型参数：

　　　　$\tilde{\theta} \triangleq \arg \underset{\theta}{\text{max }} I\left(f_{\theta}(G) ; Y\right)\quad\quad\quad(3)$

　　然后优化子图提取器 $g_{\phi}$ ：

　　　　$\begin{array}{l}\tilde{\phi} \triangleq \arg   \underset {\phi}{\text{max }}  I\left(f_{\tilde{\theta}}\left(G_{S}\right) ; Y\right) \\\text {s.t. } G_{S}=g_{\phi}(G) \in \Omega \text {. }\end{array}\quad\quad\quad(4)$

　　其中，$\Omega$ 表示满足一些约束的子图子集$\mathbb{G}_{s u b}(G)$。例如，GNNExcrener 和 PGExplainer 采用的基数约束。让我们暂时忽略不同约束之间的差异，只关注优化目标。

　　事后方法的目标函数 $\text{Eq.4}$ 和 GIB（$\text{Eq.2}$）有些类似。然而，事后方法可能不能给出，甚至不能近似于 $\text{Eq.2}$ 的最优解，因为 $f_{\theta} \circ g_{\phi}$ 并不是经过联合训练的。从优化的角度来看，事后方法只从无约束空间的模型 $f_{\tilde{\theta}}$ 到信息约束空间 $\Omega$ 的单步投影到 $f_{\tilde{\theta}} \circ g_{\tilde{\phi}} $（见 Fig. 4），投影规则遵循诱导的 MI 降低最小化 $I\left(f_{\tilde{\theta}}(G) ; Y\right)-I\left(f_{\tilde{\theta}}\left(g_{\tilde{\phi}}(G)\right) ; Y\right)$ 的 。

　　

　　后果：

• 首先，$f_{\tilde{\theta}}$ 可能不能完全从 $G_{S}=g_{\phi}(G)$ 中提取信息来预测 $Y$，因为 $f_{\tilde{\theta}}$ 最初被训练使 $I\left(f_{\tilde{\theta}}(G) ; Y\right)$ 近似 $I(G, Y)$ ，而 $\left(G_{S}, Y\right)=\left(g_{\phi}(G), Y\right)$ 遵循不同的 $(G, Y)$ 的分布。因此，$I\left(f_{\tilde{\theta}}\left(G_{S}\right) ; Y\right)$ 可能不能很好地近似 $I\left(G_{S} ; Y\right)$，从而可能误导 $g_{\phi}$ 的优化，并使 $g_{\phi}$ 不能选择真实表示 $Y$ 的 $G_{S}$。 　　
• 其次，对 $g_{\phi}$ 的积极优化可能会给出一个大的 MI ：$\hat{I}\left(f_{\tilde{\theta}}\left(g_{\phi}(G)\right) ; Y\right)$ （或者 相当小的训练损失）通过选择有助于区分训练标签的特征，但本质上与标签无关或与总体水平上的标签虚假相关。即：无用的特征有时也能很好的区分图标签。根据经验，我们确实观察到 Mutag 上的过拟合问题，如 Fig.5 所示，特别是 PGExplainer 和 GraphMask。在前 5 到 10 个epoch，这两种模型成功地选择了良好的解释，同时有很大的训练损失。进一步的训练成功地减少了损失（在 10 个 epoch 之后），但却大大降低了解释性能。这也可能是为什么在这些事后方法的原始文献中，只建议在少数epoch上进行训练的原因。然而，在实际任务中，很难有地面真实解释标签来验证结果，并确定一个可靠的停止标准。

　　

　　事后方法存在初始化问题。它们的可解释性对预先训练好的模型$f_{\tilde{\theta}}$ 高度敏感，Fig.5 中的大方差证明了这一点。只有当预训练的 $f_{\tilde{\theta}}$ 近似于最优 $f_{\theta^{*}}$ 时，才能大致保证其性能。因此，根据 GIB 原理等式对 $f_{\theta} \circ g_{\phi}$ 进行联合训练通常需要根据 $\text{Eq.}$ 的 GIB 准则。

4 Stochastic Attention Mechanism for GIB

　　在本节中，首先给出 GIB 目标的变分界（$\text{Eq.2}$），然后引入带有随机注意机制的模型GSAT。

4.1 A Tractable Objective for GIB

　　GSAT 是学习一个参数 $\phi$ 的子图提取器来提取 $G_{S} \in \mathbb{G}_{\text {sub }}(G)$。$g_{\phi}$ 通过注入的随机性阻断数据 $G$ 中与标签无关的信息，同时允许 $G_{S}$ 中保存的与标签相关的信息做出预测。

　　在 GSAT中，$g_{\phi}(G)$ 本质是 $\mathbb{G}_{\text {sub }}(G)$ 上的分布，将该分布表示为 $\mathbb{P}_{\phi}\left(G_{S} \mid G\right)$。$g_{\phi}(G)$ 即为 $\mathbb{P}_{\phi}\left(G_{S} \mid G\right)$ 。

　　对于一些 $\beta>0$，通过 GIB 得到了 $g_{\phi}$ 的优化：

　　　　$\underset{\phi}{\text{min}} -I\left(G_{S} ; Y\right)+\beta I\left(G_{S} ; G\right) \text {, s.t. } G_{S} \sim g_{\phi}(G) \quad\quad\quad(5)$

　　接下来，我们推导出 $\text{Eq.5}$ 中这两项的一个可处理的变分上界。详细的推导在附录 B 中给出。

　　对于项 $I\left(G_{S} ; Y\right)$，引入针对 $\mathbb{P}\left(Y \mid G_{S}\right)$ 的参数化变分近似 $\mathbb{P}_{\theta}\left(Y \mid G_{S}\right)$，得到一个下界：

　　　　$I\left(G_{S} ; Y\right) \geq \mathbb{E}_{G_{S}, Y}\left[\log \mathbb{P}_{\theta}\left(Y \mid G_{S}\right)\right]+H(Y)\quad\quad\quad(6)$

　　在本文模型，$\mathbb{P}_{\theta}\left(Y \mid G_{S}\right)$ 本质上可以作为参数 $\theta$ 的预测器 $f_{\theta}: \mathcal{G} \rightarrow \mathcal{Y}$。

　　对于项 $I\left(G_{S} ; G\right)$，引入边际分布 $\mathbb{P}\left(G_{S}\right)=   \sum_{G} \mathbb{P}_{\phi}\left(G_{S} \mid G\right) \mathbb{P}_{\mathcal{G}}(G)$ 的变分近似 $\mathbb{Q}(G_{S})$。并得到上界：

　　　　$I\left(G_{s} ; G\right) \leq \mathbb{E}_{G}\left[\operatorname{KL}\left(\mathbb{P}_{\phi}\left(G_{S} \mid G\right) \| \mathbb{Q}\left(G_{S}\right)\right)\right]\quad\quad\quad(7)$

　　插入上述两个不等式，得到了 $\text{Eq.5}$ 的一个变分上界作为 GSAT 的目标：

　　　　$\begin{array}{l}  \underset{\theta, \phi}{\text{min}}   -\mathbb{E}\left[\log \mathbb{P}_{\theta}\left(Y \mid G_{S}\right)\right]+\beta \mathbb{E}\left[\operatorname{KL}\left(\mathbb{P}_{\phi}\left(G_{S} \mid G\right) \| \mathbb{Q}\left(G_{S}\right)\right)\right] \\\text { s.t. } \quad G_{S} \sim \mathbb{P}_{\phi}\left(G_{S} \mid G\right)\end{array}\quad\quad\quad(8)$

　　接下来，在 GSAT 中指定 $\mathbb{P}_{\theta}$（又名 $f_{\theta}$）、$\mathbb{P}_{\phi}$（又名 $g_{\phi}$）和 $Q$。

4.2 GSAT and Stochastic Attention Mechanism

Stochastic Attention via $\mathbb{P}_{\phi}$

　　提取器 $g_{\phi}$ 首先通过GNN将输入图 $G$ 编码为一组节点表示形式 $\left\{h_{v} \mid v \in V\right\}$。对于每条边 $(u, v) \in E$ ，$g_{\phi}$ 包含一个 MLP 层后接 $sigmoid$，它将边连接 $\left(h_{u}, h_{v}\right)$ 映射到 $p_{u v} \in[0,1]$。然后，对于训练的每一次向前传递，从伯努利分布 $\alpha_{u v} \sim \operatorname{Bern}\left(p_{u v}\right)$ 中抽取随机注意。为确保梯度 $p_{u v}$ 是可计算的，应用 gumbel-softmax 重新参数化技巧。提取的图 $G_{S}$ 将有一个被注意选择的子图为 $A_{S}=\alpha \odot A$。这里 $\alpha$ 是 $(u, v) \in E$ 项为 $\alpha_{u v}$ 的矩阵，或者非边项为零。上述程序给出 $G$ 的 $G_{S}$ 的分布，描述了 $\mathbb{P}_{\phi}\left(G_{S} \mid G\right) $，所以 $\mathbb{P}_{\phi}\left(G_{S} \mid G\right)=\prod\limits _{u, v \in E} \mathbb{P}\left(\alpha_{u v} \mid p_{u v}\right)$，其中 $p_{u v}$ 是 $G$ 的一个函数，这本质上使得注意 $\alpha_{u v}$ 在给定输入图 $G$ 的不同边上是有条件地独立的。

Prediction via $\mathbb{P}_{\theta}$

　　预测器 $f_{\theta}$ 采用相同的 GNN 将提取的图 $G_{S}$ 编码为图表示，最后将该表示通过 MLP 层加 softmax 对 $Y$ 的分布进行建模。该过程给出了变分分布 $\mathbb{P}_{\theta}\left(Y \mid G_{S}\right)$。

Marginal Distribution Control via $\mathbb{Q}$

　　$\text{Eq.7}$ 对于任何 $\mathbb{Q}\left(G_{S}\right)$ 都总是正确的。我们对 $\mathbb{Q}\left(G_{S}\right)$ 的定义如下。

　　对于 $G$ 中的每个图 $G \sim \mathbb{P}_{\mathcal{G}}$ 和 $G$ 中的每两个有向节点对 $(u, v)$，我们采样 $\alpha_{u v}^{\prime} \sim \operatorname{Bern}(r)$，其中 $r \in[0,1]$ 是一个超参数。如果 $\alpha_{u v}^{\prime}=1$，我们删除 $G$ 中的所有边，并添加所有边 $(u，v)$。假设得到的图为 $G_{S}$。此过程定义了分布 $\mathbb{Q}\left(G_{S}\right)= \sum_{G} \mathbb{P}\left(\alpha^{\prime} \mid G\right) \mathbb{P}_{\mathcal{G}}(G)$ 。由于 $\alpha^{\prime}$ 独立于图 $G$，假设其大小为 $n$，$\mathbb{Q}\left(G_{S}\right)=\sum_{n} \mathbb{P}\left(\alpha^{\prime} \mid n\right) \mathbb{P}_{\mathcal{G}}(G=n)= \mathbb{P}(n) \prod_{u, v=1}^{n} \mathbb{P}\left(\alpha_{u v}^{\prime}\right)$。一个大小为 $n$ 的图的概率 $\mathbb{P}(n)$ 是一个常数，因此不会影响模型。请注意，我们对 $\mathbb{Q}\left(G_{S}\right)$ 的选择具有使用标准高斯分布作为变分自动编码器的潜在分布的相似精神。

　　使用上面的 $\mathbb{P}_{\theta}$，$\text{Eq.8}$ 中的第一项减少为一个标准的交叉熵损失。使用 $\mathbb{P}_{\phi}$ 和 $\mathbb{Q}$ ，对于每一个 $G \sim \mathbb{P}_{\mathcal{G}}$，$n$ 为 $G$ 的大小，kl-散度项变为

　　　　$\begin{array}{l}\mathrm{KL}\left(\mathbb{P}_{\phi}\left(G_{S} \mid G\right) \| \mathbb{Q}\left(G_{S}\right)\right)= \\\sum\limits_{(u, v) \in E} p_{u v} \log \frac{p_{u v}}{r}+\left(1-p_{u v}\right) \log \frac{1-p_{u v}}{1-r}+c(n, r)\end{array}\quad\quad\quad(9)$

　　其中 $c(n，r)$ 是一个没有任何可训练参数的常数。

4.3 The Interpretation Mechanism of GSAT

　　GSAT 的可解释性本质上来自于信息控制：GSAT通过注意向 $G_{S}$ 中注入随机性来减少输入图中的信息。在训练中，$\text{Eq.9}$ 中的正则项将尝试为所有边缘分配较大的随机性，但在分类损失 $\min -I\left(G_{S} ; Y\right)$（相当于交叉熵损失）的驱动下，GSAT 可以学习减少在任务相关子图上的注意力的随机性。因此，提供模型解释的不是整个 $G_{S}$，而是具有随机减少注意力的部分，又名 $p_{u v} \rightarrow 1$。因此，当 GSAT 提供解释时，在实践中，我们可以根据 $p_{u v}$ 对所有的边进行排序，并使用那些排名靠前的边（如果需要，给定一定的预算）作为检测到的子图进行解释。正如实验中（Table 5）所示的那样，注入随机性对性能的贡献非常显著，而我们的正则化项的贡献( $\text{Eq.9}$ )，当我们将其与稀疏性驱动的 $\ell_{1}-\text { norm }$ 进行比较时（Fig. 7）。

　　GSAT与以前的方法有本质的不同，因为我们没有使用任何稀疏性约束，如 $\ell_{1}-\text { norm }$ ，$\ell_{0}-\text { norm }$或 $\ell_{2}-\text { norm-regression}$ 到 {0,1}来选择大小约束（或连通性约束）子图。我们实际上观察到，在边缘正则化( $\text{Eq.9}$ )中，设置 $r$ 远离 $0$ ，也就是说，使 $G_{S}$ 远离稀疏，通常会提供更稳健的解释。这与我们的直觉相吻合，即GIB根据定义并不对选定的子图做出任何假设，而只是限制了来自原始图的信息。我们的实验表明，即使在优化过程中没有利用这些假设，GSAT也显著优于基线，即使与标签相关的子图满足这些假设。如果与标签相关的子图确实断开或大小变化，GSAT的改进预计会更大。

4.4. Further Comparison on Interpretation Mechanism

　　PGExcwaner 和 GraphMask 在他们的模型中也具有随机性。然而，他们的主要目标是在一个离散的子图选择空间上实现基于梯度的搜索，而不是像 GSAT 那样控制信息。因此，他们在原则上并没有像我们的那样得出信息正则化( $\text{Eq.9}$ )，但采用稀疏性约束来提取一个直接用于解释的小子图 $G_{S}$。

　　IB-subgraph 考虑使用 GIB 作为目标，但不注入任何随机性来生成 $G_{S}$，因此其选择的子图 $G_{S}$ 是 $G$ 的确定性函数。具体来说，IB-subgraph 采样图 $G$ 来估计 $I(GS；G)$ 并优化确定性函数 $G_{S}=g_{\phi}(G)$ 以最小化这种MI估计。在这种情况下，$I\left(G_{S} ; G\right)\left(=H\left(G_{S}\right)-H\left(G_{S} \mid G\right)\right)$ 简化为熵 $H\left(G_{S}\right)$，这倾向于给出一个小尺寸的 $G_{S}$，因为小图的空间很小，并且有一个熵的下上界。相比之下，$G_{S} \sim g_{\phi}(G)$ 在GSAT中是随机的，而 GSAT 主要通过注入随机性来增加 $H\left(G_{S} \mid G\right)$ 来实现GIB。

4.5 Guaranteed Spurious Correlation Removal

　　GSAT可以消除训练数据中的虚假相关性，并保证了可解释性。我们可以证明，如果子图模式 $G_{S}^{*}$ 与标签 $Y$ 之间存在对应关系，则模式 $G_{S}^{*}$ 是 GIB 目标的最优解( $Eq.2$)。

　　Theorem 4.1. Suppose each $G$ contains a subgraph $G_{S}^{*}$ such that $Y$ is determined by $G_{S}^{*}$ in the sense that $Y= f\left(G_{S}^{*}\right)+\epsilon$ for some deterministic invertible function $f$ with randomness $\epsilon$ that is independent from $G$ . Then, for any $\beta \in[0,1]$, $G_{S}=G_{S}^{*}$ maximizes the GIB $I\left(G_{S} ; Y\right)- \beta I\left(G_{S} ; G\right)$ , where $G_{S} \in \mathbb{G}_{\text {sub }}(G)$ .

　　Proof. Consider the following derivation:

　　　　$\begin{aligned}& I\left(G_{S} ; Y\right)-\beta I\left(G_{S} ; G\right) \\=& I\left(Y ; G, G_{S}\right)-I\left(G ; Y \mid G_{S}\right)-\beta I\left(G_{S} ; G\right) \\=& I\left(Y ; G, G_{S}\right)-(1-\beta) I\left(G ; Y \mid G_{S}\right)-\beta I\left(G ; G_{S}, Y\right) \\=& I(Y ; G)-(1-\beta) I\left(G ; Y \mid G_{S}\right)-\beta I\left(G ; G_{S}, Y\right) \\=&(1-\beta) I(Y ; G)-(1-\beta) I\left(G ; Y \mid G_{S}\right)-\beta I\left(G ; G_{S} \mid Y\right)\end{aligned}$

　　其中第三个相等式是由于 $G_{S} \in \mathbb{G}_{s u b}(G)$，那么 $\left(G_{S}, G\right)$ 没有比 $G$ 持有更多的信息。

　　如果 $\beta \in[0,1]$，$G_{S}$ 的最大化 $I\left(G_{S}, Y\right)-\beta I\left(G_{S} ; G\right)$ 也可以最小化 $(1-\beta) I\left(G ; Y \mid G_{S}\right)+\beta I\left(G ; G_{S} \mid Y\right)$，$(1-\beta) I\left(G ; Y \mid G_{S}\right)+\beta I\left(G ; G_{S} \mid Y\right)$ 的下界为 0。

　　$G_{S}^{*}$ 是生成 $(1-\beta) I\left(G ; Y \mid G_{S}^{*}\right)+   \beta I\left(G ; G_{S}^{*} \mid Y\right)=0$ 的子图。这是因为(a) $Y=f\left(G_{S}^{*}\right)+\epsilon$，其中 $\epsilon$ 独立于 $G$，所以 $I\left(G ; Y \mid G_{S}^{*}\right)=0$ 和 (b) $G_{S}^{*}=f^{-1}(Y-\epsilon)$，其中 $\epsilon$ 独立于 $G$，所以 $I\left(G ; G_{S}^{*} \mid Y\right)=0$。因此，$G_{S}=G_{S}^{*}$ 最大限度地提高了GIB $I\left(G_{S} ; Y\right)-\beta I\left(G_{S} ; G\right)$，其中 $G_{S} \in \mathbb{G}_{\mathrm{sub}}(G)$。

　　虽然 $G_{S}^{*}$ 决定了 $Y$，但在训练数据集中，数据 $G$ 和 $Y$ 可能存在一些由环境造成的虚假相关性。也就是说，$G\text{\ }G∗S$ 可能与标签有一定的相关性，但这种相关性是虚假的，并不是决定其标签的真正原因（如 $Fig.6$ 所示）。一个经过训练的 $G$，通过MI最大化来预测 $Y$ 的模型可以捕捉到这种虚假的相关性。如果这种相关性在测试阶段发生变化，模型就会出现性能下降。

　　然而，Theorem 4.1 表明，GSAT通过优化GIB目标，能够仅通过提取 $G_{S}^{*}$ 来解决上述问题，从而消除了虚假的相关性，也提供了有保证的可解释性。

4.6. Fine-tuning and Interpreting a Pre-trained Model

　　GSAT还可以微调和解释一个预先训练好的GNN。给定一个由 $\max _{\theta} I\left(f_{\theta}(G) ; Y\right)$ 预先训练的 $f_{\tilde{\theta}}$，GSAT 可以通过 $\max _{\theta, \phi} I\left(f_{\theta}\left(G_{S}\right) ; Y\right)-\beta I\left(G_{S} ; G\right) $，$G_{S} \sim g_{\phi}(G)$ 进行微调，通过初始化 $g_{\phi}$ 和 $f_{\theta}$ 中使用的GNN作为预先训练的模型 $f_{\tilde{\theta}}$。

　　我们观察到，这个框架几乎永远不会损害原始的预测性能（有时甚至会提高它）。此外，与从头开始训练GNN相比，该框架往往能获得更好的解释结果。

5 Experiments

　　由上一节可见，GSAT架构简单直接，但同时其性能又具有理论保障。这一章节将通过实验结果具体展示GSAT的可解释能力、泛化能力和各模块的消融实验结果。

可解释性

　　作者们在真实数据集和合成数据集上都对GSAT的可解释性进行了评估。作者们基于这些数据集中已知的解释标注对每个方法的解释结果评估了ROC AUC。如下图，GSAT对比过去的可解释工作，在6个数据集上提升了至多20%、平均12%的可解释性能。

　　

泛化性能

　　由于GSAT能够帮助去除伪相关性，它同时也能帮助提升模型的分类泛化能力。如下图，GSAT在11个数据集上提升了平均3%的模型准确率，并且在OGBG-MolHiv榜单上达到SOTA（在不使用手工设计的专家特征的模型中）。

　　

OOD泛化性能（伪相关性移除）

　　为了对比GSAT的移除伪相关特征的能力，作者们同时提供了和不变因果特征学习的方法的直接对比。如下图，可见GSAT能够在不利用因果分析框架的情况下，以更为简单的架构提升平均12%的OOD泛化能力。

　　

消融实验

　　作者们提供了GSAT中各个模块的消融实验结果，如下表，可见当不注入随机性（NoStoch），或者不添加正则项（$\beta=0$）时，模型效果均会大幅下降。而当不注入随机性时，模型效果将遭受最大的下降。这一消融实验展示了注入的随机性在GSAT中扮演着极其重要的角色。

　　

　　作者们同样实验了将从信息瓶颈中推导得来的KL散度正则项替换成过去的方法常用的 $\ell_{1}$ 正则。下图对各正则项的系数进行了网格搜索，可见文中提出的信息正则项显著优于$\ell_{1}$正则。

　　

6 Conclusion

　　图随机注意(GSAT)是一种建立可解释图学习模型的新型注意机制。GSAT注入了随机性来阻断与标签无关的信息，并利用减少随机性来选择与标签相关的子图。这种基本原理是基于信息瓶颈原理建立的。GSAT具有许多具有变革性的特征。例如，它消除了图学习解释中的稀疏性、连续性或其他潜在的偏差假设，而没有性能衰减。它还可以消除伪相关，更好地提高模型泛化。作为一个副产品，我们还从信息瓶颈的优化角度揭示了事后解释方法背后的一个潜在的严重问题。

---

7 GSAT 简单版

　　在很长的一段时间里，人们认为注意力机制无法提供较好的可解释性，尤其是在图学习领域。而该论文的作者们提出了一种随机注意力机制，并特别的在图学习领域进行了推导和评估，作者们称该机制为GSAT，即图随机注意力（Graph Stochastic Attention）。后续实验表明该机制能够同时提供强大的可解释能力和泛化能力。

机制原理

　　随机注意力机制，顾名思义，即是在学习注意力时注入随机性。下图提供了其在图学习领域的一个例子。该任务目标是预测图中是否存在五节点环（由图中粉色节点包围），这些环中的边是自然则是对预测结果重要的边。该机制原理的直觉如下：

1. 首先，每一条边将会习得一个 $[0,1]$ 之间的注意力权重，该权重将指代每一条边在训练中的抽样概率。一个正则项需要被引入来鼓励每一条边习得较小的抽样概率，即维持较大的随机性。如下图中间样本。
2. 随后，倘若对预测结果重要的边存在较大的随机性，那么它们在训练中将会被过于频繁的丢弃，而这将极大的影响分类损失（交叉熵）。因此被分类损失推动，重要的边最后则会维持较小的随机性，即习得较大的抽样概率（理想情况下接近于 $1$）。如下图右侧样本。
3. 最后，每条边的随机程度指代其对预测性能的重要程度，而越重要的边应有越大的抽样概率。如下图右侧样本虚线框中的子图即为低随机性的子图，它代表着对预测最为重要的子图。

　　

训练目标

　　现在的问题即是，上述的正则项应当如何选取呢？事实上这也非常直觉。因为作者们的目标是控制训练图中的随机性，而从信息论的角度来说，作者们即是希望控制图中的信息量。那么一个显而易见的选择就是信息瓶颈理论（Information Bottleneck Principle）。通过注入信息瓶颈，GSAT能够天然的控制图中的信息量，从而达到预期的效果。

　　具体而言，图信息瓶颈损失可以写作：

　　　　$\begin{array}{l}\underset{\theta, \phi}{\text{min}}-I\left(f_{\theta}\left(G_{S}\right), Y\right)+\beta I\left(G_{S} ; G\right) \\\text { s.t. } G_{S} \sim g_{\phi}(G)\end{array}$

　　其中  $I(. ; .)$  代表两个随机变量之间的互信息量（Mutual Information）,  $\beta$ 是一个正则系数, 代表 信息瓶颈注入的强度。  $g_{\phi}$  是一个负责从原图  $G$  中提取子图  $G_{S}$  的模型, 而  $f_{\theta}$  则是负责对提取出的子 图  $G_{S}$  进行下游任务  $Y$  的预测的模型。

　　互信息量自身不易优化，作者们为上述目标中的两项分别推导出了变分上界（Variational  Upper Bound）来优化该目标。

1. 对于第一项 $-I\left(f_{\theta}\left(G_{S}\right), Y\right)$ , 易得其变分上界即为 $-\mathbb{E}\left[\log \mathbb{P}_{\theta}\left(Y \mid G_{S}\right)\right]$ , 而这事实上就是 $f_{\theta}$ 基于 $G_{S}$ 进行预测后产生的交叉熵损失。
2. 对于第二项中的 $I\left(G_{S} ; G\right)$ , 易得其变分上界为 $\mathbb{E}\left[\operatorname{KL}\left(\mathbb{P}_{\phi}\left(G_{S} \mid G\right) \| \mathbb{Q}\left(G_{S}\right)\right)\right]$ , 其中 $\mathbb{P}_{\phi}\left(G_{S} \mid G\right)$ 即 为 $g_{\phi}$ 基于 $G$ 得到的每条边的采样概率; 而 $\mathbb{Q}\left(G_{S}\right)$ 即是对各边采样概率分布的一个正则, 因为该 $\mathrm{KL}$ 散度本质上将鼓励习得的每条边的采样分布逼近 $\mathbb{Q}\left(G_{S}\right)$ 的分布。举例来说, 倘若 $\mathbb{Q}\left(G_{S}\right)$ 是 一个参数为 $0.7$ 的伯努利分布, 那么这一项则将鼓励每条边的采样概率接近 $0.7$ , 而这正好符合 作者们对随机注意力机制中正则项工作原理的期待。

有保障的可解释性和OOD泛化（伪相关性去除）能力

　　由上文可知，最终GSAT的训练目标即是一个分类损失（鼓励高分类性能），加上一个KL散度的正则项（鼓励高随机性）。理想情况下，我们期待当模型仅将重要的边维持较小的随机性时，该训练目标应该被最小化，因为在这种情况下我们可能可以在达到最高分类性能的同时，取得最高的整体随机性。而作者们则在论文的定理4.1中证明了这一点，使得GSAT的性能具有理论保障。

　　具体来说，论文中定理4.1表明: 给定一个任务，如果我们假设输入图 $G$ 中包含一个子图 $G_{S}^{*}$ , 并 且其标签 $Y$ 将由下式决定: $Y=f\left(G_{S}^{*}\right)+\epsilon$ , 其中 $f$ 是一个可逆的且无随机性的函数，$\epsilon$ 是与 $G$ 无关的随机噪声。那么对于任何的 $\beta \in[0,1]$，$G_{S}=G_{S}^{*}$ 能够最小化上文提出的信息瓶颈损失。

　　这意味着GSAT能够在不利用因果分析工具的情况下，天然的找出真正重要的子图 $G_{S}^{*}$ ，并且移除可能存在的伪相关特征，从而提供有保障的可解释性和OOD泛化能力。

GSAT 模型架构

　　有了文中提出的两个变分上界, 那么GSAT的模型架构问题则变得一目了然。现在只需要对  $\mathbb{P}_{\theta}\left(Y \mid G_{S}\right)$  和  $\mathbb{P}_{\phi}\left(G_{S} \mid G\right)$  进行适当的参数化。

　　直觉来说，如下图：

1. $\mathbb{P}_{\phi}\left(G_{S} \mid G\right)$ 的输入是原始图 $G$ ，其对每一条边输出一个注意力权重，那么 $g_{\phi}$ 显然可以是一个这 样工作的GNN。
2. $\mathbb{P}_{\theta}\left(Y \mid G_{S}\right)$ 的输入是子图 $G_{S}$ ，其输出对该样本的标签预测，那么 $f_{\theta}$ 显然可以是一个这样工作的 GNN。
3. 因此，在GSAT中， $g_{\phi}$ 将会接受原图 $G$ 作为输入，然后输出每一条边的随机注意力的值。紧接 着，基于随机注意力的取值，一个子图 $G_{S}$ 将会被采样出来。最后 $G_{S}$ 将会被喂给 $f_{\theta} $ 进行最后的 标签预测。
4. 尽管 $f_{\theta}$ 和 $g_{\phi}$ 可以是两个不同的 $\mathrm{GNN}$ ，但作者们发现这里用同一个 GNN 效果就足够好。另外，架构中 $g_{\phi}$ 最后的采样操作本身是不可导的，因此作者们提出利用Gumbel-softmax Trick来重 参数化这一步骤，使其可导。