Graph Mamba Networks

Graph Mamba Networks (GMNs),一种新的基于状态空间模型图神经网络。被认为是消息传递神经网络Graph Transformer的改进。

img-2024-04-14 09-38-56.png

GMN 需要四个必须步骤和一个可选步骤

  1. Tokenization:将图映射到一个 token 序列。
  2. PE/SE:可选,通过关于节点位置和图结构的信息传递误差
  3. 局部编码:通过子图向量化机制,将每个节点周围的局部结构编码
  4. Token 排序:通过上下文对 token 进行拍去
  5. 双向 Mamba:扫描并选择相关的节点或子图来输入隐藏状态

挑战

序列和二维数据

Mamba 期望的输入为因果数据,因此对于二维的图结构数据处理起来比较棘手。需要某种对 token 的排序机制。

长序列建模

领域中,token序列,如节点,边和子图,都可以作为认为是序列。但 Transformer 无法处理长序列。但 Mamba 可以有效过滤不相关信息,任何时刻都可以重制其状态。

此外,随着序列增长,其表现也得到了提升。

可扩展性

由于图结构的复杂性,Transformer 和 Mamba 这类的序列编码器都要求恰当的位置和就够编码,这一过程通常是二次计算复杂度,导致训练起来非常困难

子图尺度

Graph Mamba Networks

设图 ,其中 为节点集合,且 为边的集合。我们假定节点 维的特征向量为 ,其中 为特征矩阵。

对于节点 ,假定 表示它的邻节点。对于节点子集 ,用 表示由 组成的对应子图, 表示对应的子图邻矩阵

Tokenization

将图映射为 token 序列。这里引入了一种新的 Tokenization 方式。首先对每个节点采样一些能够表征此节点的邻居关系,局部和全局位置的子图,然后将这些子图向量化来得到其节点表示

邻节点采样

对于给定的节点 和正整数 ,从 出发进行 步随机游走。记 表示了第 步后所有被访问到的节点集合,就可以定义这 步带来的 token 表示为

得到的子图 可以认为是对节点 步采样。最终对于每个节点都可以得到这样的 token。

由于之间提到,SSSM 的性能随着序列长度而提升,因此对应的,我们可以重复此采样过程 次。对于每个节点 ,得到了序列

可选的结构/位置编码

PE 是为了提供节点在图中的位置信息,在图中相近的两个节点的 PE 和 SE 应该是相近的。

其中 为节点 对应的位置信息。

邻节点编码

给定节点 及其子图,我们可以通过编码器对子图进行编码,为此,构建一系列

其中 。这里的编码器可以是一个消息传递神经网络。此公式把节点特征,游走时获取的边特征以及局部结构信息编码到了一个特征向量

Token 排序

由于状态空间模型这类序列模型需要输入的 token 是序列的,因此需要构建合适的顺序。尤其是 Mamba 的 都是上一次输入的函数,因此选取合适的序列至关重要。

令当