动态重连延迟消息传递图神经网络
延迟消息传递
消息传递图神经网络(MPNN)倾向于出现过度压缩现象,从而导致依赖长距离交互的任务性能下降。这可以主要归因于消息传递仅在节点的直接邻居之间本地发生。传统的静态图重连技术通常尝试通过允许远程节点即时通信(在Transformer的极端情况下,通过使所有节点在每个层次上都可访问)来抵消这种效应。然而,这会产生计算价格,并以破坏输入图结构提供的归纳偏差为代价。在本文中,我们描述了两种克服过度压缩并抵消静态重连方法的副作用的新机制:动态重连和延迟消息传递。这些技术可以被合并到任何MPNN中,并在长距离任务上比图形Transformer具有更好的性能。
本文由Francesco Di Giovanni和Ben Gutteridge共同撰写,基于B. Gutteridge等人的论文,DRew:具有延迟的动态重连消息传递(2023),ICML。
经典的消息传递图神经网络(MPNN)通过聚合每个节点的1-hop邻居的信息来操作。因此,需要具有多个消息传递层的深度MPNN来学习需要长距离交互的任务(即存在节点v,其表示需要考虑最短路径(测地线)距离d(u,v)= r > 1的某个节点u中包含的信息)。如果图形结构使得接收场随着跳跃距离指数级扩展[1],则可能需要将太多的消息“挤”进固定的节点特征向量中——这种现象称为过度压缩[2]。
过度压缩
在我们之前的工作中[3-4],我们将过度压缩形式化为MPNN输出在节点u处对距离为r的节点的输入的缺乏灵敏度。这可以通过以下偏导数(Jacobian)的上限来量化
- 速度就是你所需要的:通过 GPU 感知优化在设备上加速大型扩散模型
- Meta AI通过Voicebox打破了障碍:一个前所未有的生成式人工智能模型——革命性地改变了语音合成领域
- 从GPT-3到未来的语言模型
|∂ x ᵤ ⁽ ʳ ⁾/∂ x ᵥ ⁽⁰⁾| < c ( A ʳ ) ᵤᵥ。
这里c是一个与MPNN体系结构相关的常数(例如激活函数的Lipschitz正则性,深度等),A是图的归一化邻接矩阵。当A ʳ的条目按指数方式随着距离r衰减时,就会发生过度压缩。事实上,现在已知过度压缩更一般地可以与图的局部结构(例如负曲率[3])或其最短路径距离以外的全局结构(例如交通时间或有效电阻[4,5])有关。
上述表达式中的幂次Ar反映了在MPNN中,距离为r的节点u和v之间的通信是由不同路径连接的相邻节点的交互序列组成的。因此,节点u和v仅从第r层开始交换信息,并且延迟等于它们之间的距离r。过度压缩是由于这些路径上的中间节点通过重复的消息传递“稀释”了这些信息。
图重连
可以通过部分解耦输入图结构与用于计算消息的结构来解决过度压缩的问题,这个过程称为图重连[6]。通常,重连是作为预处理步骤执行的,其中输入图G被替换为一些其他图G’,根据某些空间或频谱连通度度量,这些图对于消息传递更加“友好”。
实现这一点的最简单方法是连接一定距离内的所有节点,从而使它们能够直接交换信息。这就是多跳消息传递方案的想法[7]。图变压器[8]将其发挥到了极致,通过注意力加权边连接所有节点对。
这样,信息不再与沿途其他节点的信息“混合”,从而避免了过度压缩。然而,这样的重连使得图在第一层就变得更加密集,增加了计算占用空间,并且部分损害了输入图所提供的归纳偏差,因为每层本地和全局节点都以相同的方式和瞬时互动。
动态图重连
查看我们之前的两个节点u和v之间距离为r>1的示例,在经典的MPNN中,必须等待r层,然后u和v才能交互,并且这种交互永远不是直接的。相反,我们认为一旦到达第r层,这两个节点已经等待“足够长的时间”,因此可以允许它们直接交互(通过插入额外的边,而不是通过中间邻居进行)。
因此,在第一层,我们仅在输入图的边上传播消息(与经典MPNN相同),但在每个后续层中,节点u的接受域会扩展一个跳[9]。这使得远程节点可以在不经过中间步骤的情况下交换信息,同时保留了输入图拓扑所提供的归纳偏差:随着深层逐渐密集化,图根据距离进行扩展。
我们称这种机制为动态图重连或DRew [10]。DRew-MPNN可以看作是经典MPNN在输入图上的本地操作和同时考虑所有成对交互的图变压器之间的“中间地带”。
延迟消息传递
在经典的MPNN中,距离r的两个节点u和v始终以恒定的延迟r层进行交互,这是信息从一个节点到另一个节点所需的最短时间。因此,节点v从r层之前“看到”节点u的状态(与其他节点的特征混合)。在DRew-MPNN中,当两个节点交互时,它们通过插入的边瞬间使用它们的当前状态进行交互。
延迟消息传递是这两种极端情况的一个折衷方案:我们为节点之间发送的消息添加全局延迟(一个超参数 𝜆)。
为了简单起见,我们在这里考虑两种简单情况:没有延迟(像DRew一样),或者最大延迟的情况,其中距离为r的两个节点u和v从第r层开始直接相互作用,但具有恒定的r延迟(就像经典的MPNNs一样):在第r层,节点u可以与节点v的状态交换信息,就好像它是r层之前[11]的节点状态。
延迟控制信息在图上的流动速度。没有延迟意味着消息传递更快,远程节点一旦添加边缘即可立即交互;相反,延迟越多,信息流就越慢,远程节点访问添加边缘时的过去状态。
𝜆DRew框架
我们将结合动态重连和延迟消息传递的架构称为𝜆DRew(发音为“Andrew”)。
一种观看𝜆DRew的方法是作为具有稀疏跳过连接的架构,允许消息不仅“水平”地传播(在同一层的图节点之间传递,如经典MPNN中),而且“垂直”地传播(跨不同层的图节点)。依赖于GNN中垂直边缘的思想并不新颖,实际上可以将残差连接视为连接每个节点到前一层的相同节点的垂直链接。
延迟机制通过创建连接节点u和不同节点v的垂直边缘来扩展此方法,具体取决于u和v之间的图距离。这样,我们可以利用跨越不同层的跳过连接固有的好处,同时将它们的条件限制在我们拥有的额外几何信息上,以图形距离的形式。
𝜆DRew缓解了过度压缩,因为远程节点现在可以访问多个(较短的)路径来交换信息,绕过重复的本地消息传递中的“信息稀释”。与静态重连不同,𝜆DRew通过减缓图的密集化并使其依赖于层,从而降低内存占用。
𝜆DRew适合以不同的速度探索图形,处理长距离交互,并普遍增强非常深的GNNs的能力。由于𝜆DRew确定消息何时何地被交换,但不确定如何被交换,因此它可以被视为可以增强现有MPNNs的元架构。
实验结果
在我们的论文[10]中,我们使用固定的参数预算,对𝜆DRew与经典MPNN基线、静态重连和Transformer类型的架构进行了广泛比较,在最近由Vijay Dwivedi和共同作者[11]引入的长距离基准测试(LRGB)中,𝜆DRew在大多数情况下表现优异。
对 LRGB 任务中 𝝼DRew 的消融研究揭示了我们框架的另一个重要贡献:调整 𝝼 以适应任务。我们观察到,使用更多的延迟(较低的 𝝼 值)对于大量层 L 的性能更好,而使用较少的延迟(高 𝝼)可以确保在较少的层后填充计算图并增加连接的密度。因此,在浅层架构(小 L)中,完全去除延迟(𝝼=∞)表现更好。相反,在深层架构(大 L)中,更多的延迟(小 𝝼)“减缓”了信息传递图的稠密化,从而提高了性能。
传统的 MPNN 类型架构在消息交换方面不同 [12]。图重连技术增加了对它们在图上发送的位置的额外控制。我们的动态图重连和延迟消息传递的新方法允许进一步控制消息何时交换。
这种方法似乎非常强大,我们的工作只是利用了访问图神经网络中过去状态的想法的第一次尝试,这取决于底层数据的“几何属性”。我们希望这种新范式可以引导出更具理论基础的框架,并挑战 MPNN 无法解决长距离任务的观念,除非使用二次注意力层进行增强。
[1] 这通常在“小世界”图中(例如社交网络)。
[2] U. Alon 和 E. Yahav,关于图神经网络的瓶颈及其实际影响(2021),ICLR。
[3] J. Topping 等人,通过曲率理解图上过度压缩和瓶颈(2022),ICLR。
[4] 通勤时间是随机游走从节点 v 到节点 u 再返回所需的期望时间。参见 F. Di Giovanni et al. ,关于消息传递神经网络中的过度压缩:宽度、深度和拓扑的影响(2023),ICML。
[5] 请参见 F. Di Giovanni 等人的定理 4.3,了解过度压缩如何影响 GNN 的能力?(2023)arXiv:2306.03589。
[6] 图重连在 GNN 社区中有些争议,因为有些人认为输入图是神圣不可侵犯的,不应该被触及。事实上,大多数现代 GNN 都采用某种形式的图重连,无论是显式地(作为预处理步骤)还是隐式地(例如通过邻居采样或使用虚拟节点)。
[7] R. Abboud、R. Dimitrov 和 I. Ceylan,用于图属性预测的最短路径网络(2022),arXiv:2206.01003。
[8] 请参见 V. P. Dwivedi 和 X. Bresson 的文章,将 Transformer 网络推广到图形(2021),arXiv:2012.09699 和 C. Ying 等人的文章,Transformer 真的在图形表示方面表现不佳吗?(2021),NeurIPS。
[9] 动态重连导致以下消息传递公式:
m ᵤ ⁽ ˡ ᵏ ⁾=AGG({ x ᵥ ⁽ ˡ ⁾ : v∈𝒩ₖ ( u )}),其中 1 ≤ k ≤ l +1
x ᵤ ⁽ ˡ ⁺¹⁾=UP( x ᵤ ⁽ ˡ ⁾, m ᵤ ⁽ ˡ ¹⁾,…, m ᵤ ⁽ ˡ ˡ ⁺¹⁾)
其中 AGG 是一个排列不变的聚合运算符,𝒩ₖ ( u ) 是节点 u 的 k -hop 邻域,UP 是一个从每个 k -hop 单独接收信息的更新操作。见 [10] 中的方程 5。
[10] B. Gutteridge 等人,DRew: Dynamically rewired message passing with delay (2023),ICML。
[11] 延迟消息传递的形式为
m ᵤ ⁽ ˡ ᵏ ⁾=AGG({ x ᵥ ⁽ ˡ ᐨ ˢ ⁽ ᵏ ⁾ ⁾ : v∈𝒩ₖ ( u )}) ,其中 1 ≤ k ≤ l +1。
x ᵤ ⁽ ˡ ⁺¹⁾=UP( x ᵤ ⁽ ˡ ⁾, m ᵤ ⁽ ˡ ¹⁾,…, m ᵤ ⁽ ˡ ˡ ⁺¹⁾)
其中 s ( k )=max{0, k ﹣𝝼},见 [10] 中的方程 6。选择 𝝼=∞ 对应于没有延迟(如 DRew),𝝼=1 对应于经典的 MPNN(两个距离为 r 的节点 u 和 v 从第 r 层开始直接交互,但具有恒定的延迟 r 。
[11] V. P. Dwivedi 等人,Long range graph benchmark (2022),arXiv:2206.08164。
[12] 在我们的原型书《Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges》(2021)中,我们区分了三种 MPNN 的“风味”:卷积、注意和通用消息传递。
我们感谢 Federico Barbero、Fabrizio Frasca 和 Emanuele Rossi 对本文进行校对并提供有见地的评论。有关图上深度学习的其他文章,请参阅 Michael 在 Towards Data Science 上的其他文章,订阅他的文章和 YouTube 频道,获得小猪AI会员,或关注他的 Twitter。