引言

图神经网络 (GNN) 彻底改变了我们处理结构化数据的方式,从预测分子性质到分析社交网络,其应用无处不在。然而,标准的 GNN——特别是消息传递神经网络 (MPNN) ——有一个众所周知的弱点: 过挤压 (Oversquashing)

在标准的 MPNN 中,信息是从邻居聚合而来的。要了解 10 步之外的节点,你需要 10 层消息传递。随着感受野的增长,需要压缩到单个固定大小向量中的信息量呈指数级增长。信号因此被“挤压”,导致模型无法捕捉长程依赖关系。

为了解决这个问题,研究人员将目光投向了 状态空间模型 (SSMs) , 如 Mamba 和 S4。由于其线性复杂度和处理海量上下文窗口的能力,它们目前正席卷序列建模领域。但这里有一个问题: SSM 是为序列 (有序数据) 设计的。图不是序列;它们是具有置换不变性的结构。强行将图转换为序列 (例如,通过对节点进行排序) 往往会破坏图的基本对称性。

在这篇文章中,我们将深入探讨一篇新论文,该论文提出了一个优雅的解决方案: GRAMA (Graph Adaptive Autoregressive Moving Average,图自适应自回归移动平均)

GRAMA 并没有将图变成节点序列,而是将特征变成了图的序列。通过将经典的自回归移动平均 (ARMA) 系统理论与现代深度学习相结合,作者创建了一个模型,该模型在理论上等同于 SSM,具有完全的置换等变性,并且在长程任务上非常有效。

背景: GNN 与动力系统的融合

要理解 GRAMA,我们需要快速回顾两个概念: 当前 GNN 的局限性和 ARMA 模型的基础知识。

传播的问题

标准的 GNN,如 GCN 或 GIN,依赖于局部聚合。虽然效率很高,但它们很难在不丢失信号保真度的情况下在整个图中传播信息。

  • 图 Transformer (GTs) 通过使用连接每个节点到所有其他节点的注意力机制来解决这个问题。虽然功能强大,但它们通常带有二次计算成本 (\(O(N^2)\)) ,这使得它们难以扩展到大型图。
  • 图 SSMs 尝试在图上使用高效的序列模型 (如 Mamba) 。然而,大多数现有方法使用启发式方法对节点进行排序 (例如,按度数) 。这意味着如果你重新排列输入文件中的节点顺序,模型可能会给出不同的预测——这违反了 置换等变性 (permutation equivariance)

什么是 ARMA?

自回归移动平均 (ARMA) 模型是时间序列分析 (例如股票价格预测) 中的主要工具。它基于两件事来建模时间 \(t\) 的值:

  1. 自回归 (AutoRegressive, AR) : 系统的先前值 (历史) 。
  2. 移动平均 (Moving Average, MA) : 先前的“误差”或残差 (冲击) 。

在数学上,经典的 ARMA 模型如下所示:

经典的 ARMA 方程

在这里,\(\phi\) 代表 AR 系数 (历史有多重要) ,\(\theta\) 代表 MA 系数 (过去的残差有多重要) 。

GRAMA 论文的关键洞察在于,我们可以将这种动力过程应用于神经网络的深度而非时间,从而允许节点特征以稳定、可控的方式演变。

GRAMA 方法

GRAMA 的核心创新在于它如何构建图学习问题。GRAMA 不会将图视为单个静态快照,而是将输入嵌入到 图序列 中,并使用神经 ARMA 块对其进行处理。

1. 从静态图到图序列

我们如何在不对节点施加任意顺序的情况下,在静态图上使用序列模型?

作者提出了“特征扩展”。给定一个输入特征矩阵,他们将其投影到一个长度为 \(L\) 的序列中。想象一下,每个节点不再是一个特征向量,而是拥有一小段特征历史轨迹

特征序列初始化的方程

同时,他们初始化了一个 残差 (\(\Delta\)) 序列,用于跟踪系统中的差异或“冲击”。

残差序列初始化的方程

这种转换至关重要。它创建了一个 ARMA 模型可以操作的“时间”维度 (序列长度 \(L\)) ,而不需要触及空间维度 (图拓扑结构) 。这保留了置换等变性。

2. 架构: 深入解析

一旦图被转换为序列,它就会通过 GRAMA 块 。 该块使用可学习的 ARMA 过程更新节点特征。

让我们可视化完整的框架:

图 1: GRAMA 框架概览

如上图 1 所示,该模型由初始化和随后的堆叠 GRAMA 块组成。在每个块内部,更新规则结合了三个组件:

A. 自回归 (AR) 更新

节点的新状态取决于其在序列中的自身历史。这使得模型能够保持特征演变的长期记忆。

AR 更新方程

B. 移动平均 (MA) 更新

新状态还取决于残差的历史。在控制理论中,这允许系统根据过去的“噪声”或输入进行自我修正。

MA 更新方程

C. 图骨干 (“空间”组件)

到目前为止,AR 和 MA 部分只关注节点自身的历史序列。我们仍然需要与邻居交换信息!

这就是 骨干 GNN 发挥作用的地方。“当前”残差 \(\delta\) 是通过在前一个状态上运行标准 GNN (如 GCN、GAT 或 GPS) 生成的。这将空间混合 (GNN) 与时间演变 (ARMA) 结合在一起。

递归步骤的最终更新方程结合了这三者:

完整的递归方程

这里,\(\delta^{(\ell+L)}\) 代表 GNN 骨干网络的输出。通过这种方式构建层,GRAMA 允许 GNN 专注于局部聚合,而 ARMA 动力学则处理深度 (时间) 上的信息保留。

3. 深度 GRAMA 与非线性

单个 GRAMA 块执行 \(R\) 个递归步骤。为了构建深度网络,可以堆叠多个块。至关重要的是,ARMA 递归本身是一个 线性系统 。 为了赋予神经网络表达能力,非线性激活函数 (\(\sigma\)) 被应用在块之间

块之间非线性的方程

4. 自适应与选择性学习

僵化的 ARMA 模型具有固定的系数 \(\phi\) 和 \(\theta\)。然而,不同的图 (甚至特征序列的不同部分) 可能需要不同的动力学。有些可能需要长记忆 (高 AR) ,而另一些则依赖于即时输入。

GRAMA 引入了 选择性 (Selectivity) 。 该模型不学习静态系数,而是使用 注意力机制 根据输入特征动态预测系数 \(\phi\) 和 \(\theta\)。

注意力机制对序列进行池化并计算归一化分数:

系数的注意力机制

这导致系数会自适应地变化。作者在论文中可视化了这些学习到的系数:

学习系数的热力图

在上图 3(b) 中,请注意训练后的系数 (Post-training) 不是均匀的。模型学会了关注特定的过去步骤,有效地“选择”历史的哪一部分对当前更新很重要。这种适应性是 GRAMA 性能的一个主要因素。

理论基础: SSM 连接

这篇论文最强有力的贡献之一是理论方面的。作者证明了他们特定的神经 ARMA 公式在数学上等同于 线性状态空间模型 (SSM)

SSM 通常由隐藏状态 \(x_t\) 和输出 \(f_t\) 的转换定义:

标准 SSM 方程

作者推导出了一个显式映射,表明 GRAMA 的 AR/MA 系数构成了 SSM 中状态矩阵 \(\mathbf{A}\) 的特定结构。

状态矩阵 A 的构造

这为什么重要?

  1. 稳定性: 我们可以通过观察该矩阵的特征值来分析神经网络的稳定性。如果它们位于单位圆内,则梯度传播是稳定的。
  2. 长程交互: 理论表明,如果特征值接近单位圆,状态矩阵的幂 (\(A^k\)) 衰减缓慢,允许信息在非常长的距离上传播而不会消失。

这一理论基础为 GRAMA 为何能缓解过挤压提供了正式的解释。

实验与结果

研究人员在 26 个不同的数据集上评估了 GRAMA,从合成压力测试到现实世界的分子基准测试。

1. “图转移”压力测试

为了明确测试过挤压,他们使用了“图转移 (Graph Transfer) ”任务。目标是将信号从源节点传递到特定距离 (跳数) 的目标节点。

随着距离的增加,标准 GNN 通常会失败,因为信号被稀释了。

图转移结果

在图 2 中,请看 紫色线 (GRAMA)

  • GCN (蓝色)GAT (橙色) 的错误率随着距离 (x 轴) 的增加而飙升。
  • GRAMA 即使在 50 跳时也能保持接近零的错误率,在稳定性方面显著优于标准 MPNN 甚至像 GPS 这样具有竞争力的图 Transformer。

2. 长程图基准 (LRGB)

LRGB 是测试现实世界场景 (肽) 中长程能力的标准。

表 2: LRGB 结果

在表 2 中,我们看到了 GRAMA 在 Peptide-func 和 Peptide-struct 上的表现。

  • 搭配简单 GCN 骨干的 GRAMA (GRAMA_GCN) 实现了 70.93 的 AP,比标准 GCN (59.30) 高出 11 个点以上。
  • 它的表现与专业的 Graph Mamba 模型和 Graph Transformer 相当甚至更好,而且通常参数更少或复杂度更低。

3. 效率与速度

SSM 方法的主要承诺之一是效率。具有全注意力的 Graph Transformer 是 \(O(N^2)\)。GRAMA 由于依赖于骨干 GNN (通常是线性 \(O(E)\)) 和通道级递归,因此扩展性更好。

表 9: 运行时间

表 9 比较了 Roman-Empire 数据集上的运行时间。

  • GPS (图 Transformer) : 在深度 32 时内存溢出 (OOM) ,在深度 4 时训练耗时约 1139 毫秒。
  • GRAMA: 即使在深度 32 时也不会 OOM。在同等深度下,它比 Transformer 快得多,同时实现了更高的准确率 (86-88% vs 81-82%) 。

4. 标准基准测试

作者还在标准的异配数据集 (Roman-empire, Amazon-ratings 等) 上进行了测试。

表 17: 结果汇总

表 17 中的总结令人信服。几乎在每项任务中——无论是节点分类、回归还是图分类——将骨干网络 (如 GCN) 包装在 GRAMA 框架中都能持续提升性能,而且通常提升幅度很大。

结论

GRAMA 代表了动力系统与深度图学习融合迈出的成熟一步。通过承认图需要置换等变性,同时也认识到序列状态空间建模的威力,作者创造了一种兼具两者优点的混合模型。

关键要点:

  1. 图序列: 将静态输入转换为时间序列,使得能够使用强大的 ARMA/SSM 动力学而不破坏图的对称性。
  2. 即插即用的增强: GRAMA 充当一个包装器。你可以将 任何 GNN 骨干 (GCN, GAT, GatedGCN) 放入其中,以增强其长程能力。
  3. 自适应动力学: 基于注意力的系数学习允许模型 针对每个图 决定是否需要长期记忆或局部关注。
  4. 理论稳定性: 与 SSM 的联系保证了模型能有效地对抗梯度消失和过挤压问题。

对于处理大型图 (如生物网络或供应链) ,其中远距离节点相互影响的学生和研究人员来说,GRAMA 提供了一个理论基础扎实且实际上高效的工具来捕捉那些难以捉摸的长程信号。