Graph-Mamba

论文：Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces

对于Graph-Mamba，需要了解一些相关知识。

Related Work

Mamba结构

Mamba建立在状态空间模型的概念之上，但引入了一个新的变化。它利用选择性状态空间 (SSM: selective state space model)，支持跨长序列更高效和有效地捕获相关信息。

Mamba结构最初是在Mamba: Linear-Time Sequence Modeling with Selective State Spaces中提出的，

为了解决 Transformer 在长序列上的计算效率低下问题，人们开发了许多亚二次时间架构，如线性注意力、门控卷积和递归模型以及结构化状态空间模型（SSM）。此类模型的一个关键弱点是无法进行基于内容的推理。

作者提出了一类新的选择性状态空间模型，该模型在多个维度上改进了先前的工作，从而在序列长度线性缩放的同时，实现了Transformer的建模能力。主要的改进有以下几点：

1. 选择机制 先前的模型，特别是那些用于处理序列数据的模型，可能在有效选择数据方面存在局限，它们可能不够有效地关注重要的输入信息或忽略不相关的输入信息。选择性复制和归纳头等合成任务在理解和改进模型的选择机制方面提供了重要的直觉。于是作者设计了一种简单的选择机制，根据输入对 SSM 参数进行参数化（让 SSM 参数成为输入的函数）。这样，模型就能过滤掉无关信息，并无限期地记住相关信息。
2. 硬件感知算法 先前的SSM模型为了计算效率，必须是时间不变和输入不变的，这意味着它们的计算方式不随时间或输入数据的变化而改变。而上一点提到的改进，虽然提高了模型的灵活性和效果，但也带来了技术上的挑战。 于是作者开发了一种硬件感知算法。这种算法考虑了硬件的特性和限制，特别是在使用GPU进行计算时。算法使用递归的方式通过扫描来计算模型，而不是使用传统的卷积。 算法设计考虑到了GPU内存层次结构，为了避免不同级别之间的IO访问，它不会实体化扩展的状态，这一步骤有助于减少内存访问和相关的时间延迟。 在理论上，这种实现方式使得处理速度随序列长度线性增长，与基于卷积的所有SSM模型的伪线性增长相比有显著优势。
3. 模型架构 作者将先前的 SSM 架构设计（Dao、Fu、Saab 等人，2023 ）与 Transformers 的 MLP 模块合并为一个模块，集成到一个简化的端到端神经网络架构中，该架构没有注意力，甚至没有 MLP 块，简化了先前的深度序列模型架构，形成了一种包含选择性状态空间的简单、同质的架构设计（Mamba）。

其算法的核心内容可以用下图表示：

公式以及结构的推导可见下链接：

一文读懂Mamba：具有选择状态空间的线性时间序列建模 - 知乎 (zhihu.com)

代码实现：

挑战Transformer的新架构Mamba解析以及Pytorch复现 - 知乎 (zhihu.com)

GNN and GCN

Graph Neural Networks (GNN): Graphical-Based Learning Environments for Pattern Recognition | SpringerLink

Graph Convolutional Networks (GCN):[1609.02907] Semi-Supervised Classification with Graph Convolutional Networks (arxiv.org)

Graph Neural Networks (GNN)

原论文过于枯燥，下面这个专题的文章讲解很详细只粗略浏览了一些

GNN综述——从入门到入门 - 知乎 (zhihu.com)

Graph Convolutional Networks (GCN)

核心原理图如下：

Graph Transformers

[2012.09699v2] A Generalization of Transformer Networks to Graphs (arxiv.org)

• Transformer and GNN

• Graph Transformer结构

GraphGPS

[2205.12454] Recipe for a General, Powerful, Scalable Graph Transformer (arxiv.org)

the replacement of fully-connected Transformer attention with its sparse alternatives

State Space Models(SSM)

Utilizing a recurrent scan, SSM stores context in its hidden states, and updates the output by combining these hidden states with input.

Graph-Mamba

Graph-Mamba采用选择性SSM来实现依赖于输入的图形稀疏化，设计了一个Graph-Mamba块(GMB)，并将其合并到流行的GraphGPS框架中，以便与其他图形注意实现进行公平的比较。GMB利用循环扫描序列建模和选择机制来实现两个级别的图形稀疏化。第一级涉及Mamba模块中的选择机制，它有效地过滤远程上下文中的相关信息。第二级是通过提出的节点优先级方法实现的，允许图中的重要节点访问更多上下文。因此，这些序列建模特征为图稀疏化提供了一种结合数据依赖和启发式信息选择的有前途的途径。

Structured state space models for sequence

概述Graph-Mamba架构，通过合并GMB来取代GraphGPS框架中的注意力模块。

• A) GMB层，是对GPS层的改进，结合了基于边缘的MPNN和以节点为中心的GMB来输出更新的节点和边缘嵌入。
• B) Graph-Mamba将GatedGCN模型作为MPNN的默认模型。
• C) GMB使用节点优先级和排列技术的专门训练配方，执行图的稀疏化。
• D)带有Mamba模块的选择机制，方便了依赖于输入的上下文过滤

Graph-dependent selection mechanism

在以节点为输入序列的图学习中，选择机制允许隐藏状态基于先前序列中的相关节点进行更新，由当前输入节点进行门控，并随后影响当前节点的输出嵌入。Gt的范围在0到1之间，允许模型在需要时完全过滤掉不相关的上下文。选择和重置的能力使Mamba能够在给定长期上下文的情况下提取相关的依赖关系，同时最大限度地减少复发的每个步骤中不重要节点的影响。因此，它提供了一种上下文感知的替代方案，通过仅在长输入序列中保留相关依赖关系来稀疏图的注意力。