AAAI 2023 | 利用脉冲神经网络扩展动态图表示学习

2023-02-25 14:02

©PaperWeekly 原创 · 作者 | 李金膛

单位 | 中山大学博士生

研究方向 | 可信图学习

2020 年国家双碳战略的确立与实施，绿色低碳已经成为全社会的重要议题，也是科技从业者的重要使命和责任。有文献指出，从 2012 年到 2018 年，用于深度学习的计算资源消耗增长了 30 万倍，高能耗的特点使得深度学习广受诟病。在深度学习算法在工业届的应用日益增多的今天，算法效率和可扩展性成为制约其更加广泛应用的一个关键因素，如何构建鲁棒、有效、低能耗的预测模型，是一个重要的研究方向。

本文基于第三代人工神经网络——脉冲神经网络，聚焦于动态图表示学习场景，提出一种绿色高效的动态图表示学习方法。

论文题目：

Scaling Up Dynamic Graph Representation Learning via Spiking Neural Networks

论文链接：

https://arxiv.org/abs/2208.10364

代码链接：

https://github.com/EdisonLeeeee/SpikeNet

脉冲神经网络技术背景

随着深度学习技术的不断发展，模型算力得到了很大提升，但与此同时也面临着一个关键问题：神经网络涉及的计算成本和能耗究竟有多大？

如图 1 所示，根据马萨诸塞大学研究人员公布的最新研究 [1]，训练一个常见的大型深度学习模型，该过程可排放超过 626,155 磅二氧化碳，相当于一辆普通汽车寿命周期碳排放量的 5 倍，一个普通人一生碳排放量的 57 倍。该研究结果揭示了目前 AI 领域面临的一个严峻问题：深度学习模型训练过程成本高昂，并且随着算力的提升，这一问题会日趋严重。

尽管传统的人工神经网络在多项任务上取得了突破性的进展，但是能耗问题却限制了其更加广泛的部署应用，为解决这一问题，脉冲神经网络（Spiking Neural Network，SNN）被提了出来。基于类脑计算框架下的 SNN 以脉冲神经元为基本计算单元，通过稀疏脉冲序列传递信息，被称为能耗更低的新一代绿色人工智能技术。

IF（Integrate-and-Fire）[2] 是一类经典的脉冲神经元模型，其假设脉冲神经元是一个理想的电容器，随着时间流逝逐渐累积电荷并且不存在任何泄露。因此，当施加在电容器上的电压越大，进入电容器的电荷就越多。这一关系可由公式表示，其中为电荷量，为膜电压，为电容常数。为描述神经元状态随时间变化情况，通过对时间进行微分，可以改写成如下的微分方程：

其中为当前时刻的输入电流。当膜电压积累的电荷超过阈值时发射脉冲，并消耗神经元之前积累的电荷，电压下降至静息电位，神经元处于静息状态。然而，实际情况中，神经元并不是完美的电容器。相反，随着时间的推移，它们会慢慢地泄漏电流，将膜电压拉回到静息电位。

为了更好地模拟这一现象，研究人员提出了 LIF（Leaky Integrate-and-Fire）[3] 模型。LIF 通过额外引入一个衰减项，并将神经元的状态变化活动描述为电压输入和环境损耗的共同作用控制。

如图 2 所示，基于 LIF 模型的神经元可以看出是额外与一个电阻并联的电容器。因此，可以将上述的微分方程改写为：

其中，为当前膜电压输入。是一个与神经元有关的超参数，用于控制膜电位衰减的速度。将上式微分方程作一阶展开，即可得到 LIF 神经元的膜电压更新公式：

方法介绍

现有的动态图表示学习方法，以 TGAT [4] 和 EvolveGCN [5] 为例，主要通过将静态图场景上的图表示学习方法（GNN）与时序模型结合（如 RNN 类方法 LSTM、GRU 等）。尽管在建模动态图数据上取得了一定的效果，但是仍然不可避免地受到了计算效率的制约而无法扩展至更大规模数据场景下。

为解决动态图表示学习在大规模图数据场景下的效率问题，本文题出了 SpikeNet，一个基于 SNN 的可扩展动态图表示学习框架。

SpikeNet 整体框架如图 4 所示。SpikeNet 聚焦于离散时间图序列，对于每个时间戳的图快照，SpikeNet 首先采样节点的邻居子集并聚合邻居信息，然后使用 LIF 模型将邻域聚合得到的信号作为输入，通过 Integrate-and-Fire 机制来捕捉图结构的动态变化，并输出脉冲序列信号。最后，通过对每个时间步长得到的脉冲信号使用池化操作得到最终的节点表征，用于后续的下游任务（如节点分类）学习。

2.1 脉冲神经网络：更加绿色高效的RNN替代方法

本文探索了一个新的方向，即我们可以使用 SNNs（如 LIF）而不是 RNNs（如 LSTM 和 GRU）作为时序架构来捕捉图序列中的动态模式。相比于 RNNs 类模型，SNNs 参数量少（几乎无训练参数）、计算效率高（可以完全不涉及矩阵乘法），具有构建鲁棒、有效、低能耗的预测模型的巨大潜力。

回顾前文，LIF 的神经元行为是由一系列的事件来表征的，即 Integrate，Fire 和 Reset。对于每个时间戳，LIF 模型中的每个神经元根据其记忆的状态和当前输入更新膜电位，然后在膜电位达到阈值时发射脉冲。发射脉冲之后，膜电位就会被重置到一个特定的值，然后重新开始下一轮的输入-发射-重置过程。每一层的神经元根据从上一层收到的输入信号经历这一过程。相应地，每个神经元都会为后续层产生一个脉冲序列作为输入。

然而，先前的工作通常对不同层的每个 LIF 神经元采用固定的发射阈值，从而潜在限制了其灵活性和表达能力。为了解决这个问题，本文首先提出了一个神经元阈值的自适应更新策略，能够根据当前阈值和接收的脉冲输入计算神经元的动态阈值变化：

其中和是在用于调整发射阈值的衰减因子。是输入 0-1 脉冲序列，代表一个神经元在时间是否发射脉冲。根据上式，发射阈值就可以根据神经元的动态变化进行适应性调整。最后，结合了自适应阈值调整策略的 LIF 模型可由如下的式子表示：

其中表示当神经元接收来自先前神经元的输入时模型的更新行为，本文采用了 LIF 模型的衰减式更新方式。在神经元输出中发射和产生脉冲的决定是根据Heaviside 阶跃函数进行离散化的，其定义为：如果，否则为。

结合了神经元阈值动态更新策略后，LIF 模型的整个更新过程如图 5 所示。通过不断充电、放电以及动态重置调整，LIF 模型可以捕获动态时间信息并在每个时间段输出一定数量的脉冲。

为方便起见，本文将上述过程简化为，为接收输入时的 LIF 模型。对于每个时间步，输入是图神经网络聚合得到的邻域信息，而输出是中间节点表征的脉冲序列。

2.2 动态邻居采样与邻居聚合

接下来，本文将介绍 SpikeNet 如何在每一层和时间步中聚合邻域信息为作为 LIF 模型的输入。对于每个时间戳输入的图结构快照，我们可以直接利用 GraphSAGE 论文 [6] 中提出的邻居采样方法，采样输入节点的跳子图作为图神经网络模型的输入。然而，GraphSAGE 主要应用于静态图数据，如果直接将其采样方法应用于动态图数据场景容易忽略动态时间变化信息。因此，本文提出使用如下的采样方式：

其中，为图采样器，可以为 GraphSAGE 论文所提的均匀采样器或者随机游走采样器。表示在时刻新增的子图结构，且。上述采样步骤结合了每个时间步的宏观演化信息（macro-dynamics）和微观动态信息（micro-dynamics）来学习节点表征。直观地说，微观动态信息捕捉到了节点表征的细粒度结构和时间属性，而宏观演化信息则从全局角度探索动态图的演化模式。

接着，可以根据采样的邻居子图进行节点表征更新：

每一层得到的脉冲序列都会输入 LIF 模型，进行神经元充放电更新，从而得到节点的隐藏层表征（脉冲序列）。最后，采用对每个时间步得到的脉冲序列进行池化操作得到最终的输出：

这里的池化函数可以有很多选择，例如 Mean, Sum, 或者采用更加高级的如 Attention 注意力机制。本文仅简单采用了一层全连接网络得到最终的输出：

即为节点的表征，可用于后续的节点分类或链路预测任务。

2.3 Masked Summation运算机制

脉冲神经网络的一大优势是使用离散脉冲而不是实值进行消息传递，这一特性使得我们可以更进一步地优化其计算方式。如图 8 所示，在消息聚合阶段，我们可以将矩阵乘法简化为 Masked Summation 运算，一种简单的索引求和的计算方式，从而避免了复杂的矩阵运算。

2.4 替代梯度训练

相较于人工神经网络的广泛使用，脉冲神经网络受限于难以训练的问题无法大规模进行应用。这是由于生成二值脉冲序列的阶跃函数不可导性，导致无法采用梯度下降方法直接进行训练。梯度下降方法的广泛使用造就了人工神经网络的大流行，相反，至今鲜有能够有效训练脉冲神经网络的方法。

▲ 图9 替代梯度训练过程前向传播与反向传播示意图

尽管无法直接使用梯度下降优化脉冲神经网络，研究人员也提出了各种有效的训练方式，其中使用较为广泛的是替代梯度训练方法。本文即采用替代梯度训练方法对 SpikeNet 进行优化训练。具体优化方式如图 9 所示：

前向传播时，模型仍然采用阶跃函数得到输出的脉冲序列；

反向传播时，采用阶跃函数的一个平滑替代函数进行梯度计算求导。的平滑替代函数有很多，本文简单起见采用了如平滑化的 sigmoid 函数，定义为，其导数为：

其中，为平滑因子，用于控制 sigmoid 函数的平滑程度。越大，越接近于函数，但容易在函数两端出现梯度消失或梯度爆炸现象。

实验结果

实验在三个大规模动态图数据集上进行，任务为动态图节点分类，评价指标为Macro-F1 & Micro-F1。

从表 1 可以看出，尽管 SpikeNet 使用了离散脉冲序列作为层与层之间的消息传递载体，其分类准确率仍然优于大部分基于实值输入的对比方法。

▲ 图10 不同方法（a）参数量和（b）训练时间对比

图 10 关于参数量与训练时间的比较也说明了 SpikeNet 的优越性。借助于脉冲神经网络使用离散脉冲传递的特征，SpikeNet 展示了构建鲁棒、有效、低能耗的预测模型的巨大潜力。

▲ 图11 动态阈值调整策越相关参数实验

图11关于参数和说明了本文提出的动态阈值调整策的有效性。当和，SpikeNet 取得了最优的表现。

▲ 图12：SpikeNet在不同时间段的平均脉冲发射频率

此外，受益于内在的神经元动力模式及基于脉冲的消息通信范式，脉冲神经网络可以很容易地应用于一些专门的神经形态硬件加速训练和推理，如 Intel Loihi [7]。脉冲发射频率是衡量脉冲神经网络效率一个重要属性，因为能量消耗与脉冲的数量成正比。

图 12 可视化了不同时间段中 SpikeNet 的平均脉冲发射频率，可以看到，在初始的时间步长内，SpikeNet 的脉冲发射频率较低，而随着时间的推移，频率之间增大，说明了不同数据集中，图结构信息随着图演化逐渐增大。从整体来看，SpikeNet 在所有时间步的平均脉冲发射频率为 20%-30% 左右，这说明 SpikeNet 的中间节点表示是非常稀疏的（只有 20%-30% 的元素是非零的）。

换句话说，我们只需要大约20％-30％的内存来存储中间节点的表示。此外，由 SpikeNet 学习的嵌入向量的每个维度仅由 1 位编码（二进制脉冲），与实值嵌入向量不同，后者至少由 32 位编码。因此，二值化的节点表征也可以大大减少内存和时间成本。

总结

本文先介绍了脉冲神经网络的技术背景，并提出使用脉冲神经网络替代 RNN 作为动态图表征学习中的时序模型。通过算法创新提出 SpikeNet，以 LIF 模型结合动态阈值策略作为时序模型，使用动态邻居采样与邻居聚合捕获动态图结构信息，最后利用替代梯度训练方法解决脉冲神经网络难以训练的问题。实验结果表明，SpikeNet 以较低的时间与空间开销获得了更高的预测准确率，展示了构建鲁棒、有效、低能耗的预测模型的巨大潜力。

参考文献

[1] Strubell E, Ganesh A, McCallum A. Energy and policy considerations for deep learning in NLP[J]. arXiv preprint arXiv:1906.02243, 2019.

[2] Salinas E, Sejnowski T J. Integrate-and-fire neurons driven by correlated stochastic input[J]. Neural computation, 2002, 14(9): 2111-2155.

[3] Gerstner W, Kistler W M, Naud R, et al. Neuronal dynamics: From single neurons to networks and models of cognition[M]. Cambridge University Press, 2014.

[4] Xu D, Ruan C, Korpeoglu E, et al. Inductive representation learning on temporal graphs[J]. arXiv preprint arXiv:2002.07962, 2020.

[5] Pareja A, Domeniconi G, Chen J, et al. Evolvegcn: Evolving graph convolutional networks for dynamic graphs[C]. Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(04): 5363-5370.

[6] Hamilton W, Ying Z, Leskovec J. Inductive representation learning on large graphs[J]. Advances in neural information processing systems, 2017, 30.

[7] Davies M, Srinivasa N, Lin T H, et al. Loihi: A neuromorphic manycore processor with on-chip learning[J]. Ieee Micro, 2018, 38(1): 82-99.

更多阅读