《Allen-Cahn Message Passing for Graph Neural Networks with Particle Phase Transition》阅读笔记

一.概述

本文从多粒子系统的角度引入排斥力到消息传递框架中，提出了ACMP模型。现有的大多数消息传递神经网络是由与Dirichlet能量相关的吸引力驱动的，相应的图神经网络存在oversmoothing问题，无法进行异配数据集预测。但倘若相互排斥的粒子始终保持排斥，随着时间推进，粒子之间会不断远离。为避免粒子被推到无穷远处，作者引入了Allen-Cahn项。

二.背景

消息传递图神经网络

给定无向图$\mathcal{G}(\mathcal{V},\mathcal{E})$，$x_i^{(k-1)}\in {\mathbb{R}}^d$表示节点$i$在第$(k-1)$层的特征，$a_{j,i}\in {\mathbb{R}}^D$表示节点$j$到$i$的边特征。
$${\mathbf{x}}_i^{(k)}={\gamma }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\square }_{j\in {\mathcal{N}}_i}{\varphi }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},a_{j,i}\right)\right)$$
其中$\square$表示可微、转置不变函数（如sum、mean等），$\gamma$和$\varphi$表示可微函数如（MLPs），${\mathcal{N}}_i$表示节点$i$的一跳邻居集。

消息传递网络通过聚合邻居的特征和自身上一轮的特征来生成本轮新的特征表示。许多代表性GNNs都follow这一范式，如GCN的聚合规则为 x i ′ = Θ ⊤ ∑ j ∈ N i ∪ { i } a j , i d j d i x j \mathbf{x}_i^{\prime}=\Theta^{\top} \sum_{j \in \mathcal{N}_i \cup\{i\}} \frac{a_{j, i}}{\sqrt{d_j d_i}} \mathbf{x}_j xi′​=Θ⊤∑j∈Ni​∪{i}​dj​di​ ​aj,i​​xj​，其中${\hat{d}}_i=1+{\sum }_{j\in \mathcal{N}(i)}{a}_{j,i}$。

图神经扩散

图上神经扩散方程（GRAND）为一些消息传递提供了统一的数学框架：
$$\frac{\partial }{\partial t}\mathbf{x}(t)=\mathrm{div}[\mathbf{G}(\mathbf{x}(t),t)\nabla \mathbf{x}(t)]$$
其中$\mathbf{G}=\mathrm{diag}\left(a\left(x_i(t),x_j(t),t\right)\right)$，$a$为反映节点$i$和节点$j$间相似性的函数，$x_i$为节点$i$的特征。

三.动机

引力和斥力

本文研究了与图神经扩展相似的相互作用的粒子系统的神经方程，将改式子改写为按分量计算的形式，得到粒子系统：
$$\frac{\partial }{\partial t}{x}_i(t)=\sum _{j\in {\mathcal{N}}_i}a\left(x_i,x_j\right)\left(x_j-x_i\right)$$
从粒子系统的公式可以看出，当$a\left(x_i,x_j\right)>0$大于时，$x_i$的速度方向向着$x_j$，即$x_i$被$x_j$吸引。相反地，若$a\left(x_i,x_j\right)<0$，则$x_i$有趋势远离$x_j$。因此，$a\left(x_i,x_j\right)$ 表示$x_i$与$x_j$之间的吸引力或排斥力。在上述的扩展模型中，所有的$a\left(x_i,x_j\right)$都是正的，因此连通分量中的所有节点特征都是相互吸引的。

然而，在某些问题中negative消息可能更重要，例如在二部图上进行节点分类，邻居消息为负，因为因为连接的节点属于不同的类。在粒子系统中可以加入偏置${\beta }_{ij}$来模拟正、负信息机制。
$$\frac{\partial }{\partial t}{x}_i(t)=\sum _{j\in {\mathcal{N}}_i}\left(a\left(x_i,x_j\right)-{\beta }_{i,j}\right)\left(x_j-x_i\right)$$
系数项$a\left(x_i,x_j\right)-{\beta }_{i,j}$对应相互作用力，可能为吸引力也可能为排斥力，因此通过这种方式消息传递既能处理同配图又能处理异配图。

Allen-Cahn项

图的Dirichled能量定义为：
$$\mathbf{E}(\mathbf{x})=\frac{1}{N}\sum _{i\in \mathcal{V}}\sum _{j\in {\mathcal{N}}_i}{a}_{i,j}{\Vert {\mathbf{x}}_i-{\mathbf{x}}_j\Vert }^2$$
然而，加入排斥力可能会导致粒子被推到无穷远，这样Dirichlet能量就变得无界了。为避免该情况，作者添加了一个Allen-Cahn项$\delta x_i\left(1-x_i^2\right)$，即：
$$\frac{\partial }{\partial t}{x}_i(t)=\alpha \sum _{j\in {\mathcal{N}}_i}\left(a\left(x_i,x_j\right)-{\beta }_{i,j}\right)\left(x_j-x_i\right)+\delta x_i\left(1-x_i^2\right)$$

四.Allen-Cahn消息传递

基于上述内容，作者提出了Allen-Cahn Message Passing (ACMP)神经网络。

令${\mathrm{x}}^{\text{in }}$表示输入特征矩阵，在作者的解决方案中，首先对嵌入过一个MLP，即 $\mathbf{x}(0)=\mathrm{MLP}\left({\mathbf{x}}^{\mathrm{in}}\right)$ 。

ACMP的更新方程为：
$$\frac{\partial }{\partial t}{\mathbf{x}}_i(t)=\mathbf{\alpha }\odot \sum _{j\in {\mathcal{N}}_i}\left(a\left({\mathbf{x}}_i(t),{\mathbf{x}}_j(t)\right)-\beta \right)\left({\mathbf{x}}_j(t)-{\mathbf{x}}_i(t)\right)+\mathbf{\delta }\odot {\mathbf{x}}_i(t)\odot \left(1-{\mathbf{x}}_i(t)\odot {\mathbf{x}}_i(t)\right)$$
ACMP的模型架构如下：