专业知识

论文阅读汇报：SPATIO-TEMPORAL APPROXIMATION: A TRAINING- FREE SNN CONVERSION FOR TRANSFORMERS

1.研究背景

ANN在处理复杂任务方面表现出色，但能耗较高。相反，SNN由于其生物学上的可解释性和低能耗特性，被视为下一代神经网络的有力候选。然而，SNN在处理某些复杂任务时存在性能限制。

因此,将预训练的ANN转换为SNN成为一种吸引人的方法。这种方法无需额外训练数据和资源，即可保持强大的性能。尽管现有的转换方法在卷积网络上表现良好，但新兴的Transformer模型引入了自注意力和测试时规范化等独特机制，导致非因果非线性交互，现有的SNN难以实现。

为解决这一问题，本研究提出了一种时空近似（STA）转换流程，首次实现了Transformer到SNN的转换。通过在时空两个维度上逼近这些操作，我们提出了全局群体操作器来空间近似非线性操作，并设计了一种时间校正自注意力层，通过估计-校正方法在推理时近似脉冲乘法。


2. 预备知识

2.1 ANN vs SNN

这部分内容将介绍传统人工神经网络和脉冲神经网络的不同,以及脉冲神经网络的一些优点。

Artificial neural networks

上图描绘了典型的人工神经元的模型。

计算过程由：y=ψ\psi(b + ∑j\sum_jxjx_jwjw_j)

  • x：输入激活
  • y：输出激活
  • w：突触权重
  • b：偏差
  • j：输入神经元的索引
  • ψ\psi：非线性激活函数

Spiking neural networks

上图展示了一个典型的脉冲神经元，与ANN神经元相比，它具有相似的结构，但行为不同。相比之下，脉冲神经元通过二进制事件编码的脉冲序列进行通信，而不是人工神经网络中的连续激活。树突整合输入脉冲，soma进而进行非线性变换，产生输出脉冲序列。

行为通常由流行的 LIF 模型建模，描述为:

  • (t): 时间步长，时间离散化单位，模型中的微分方程会根据t进行更新
  • τ: 时间常数
  • u: 膜电位
  • s: 输出脉冲
  • ur1：静息电位，指神经元在没有接收到任何突触输入时的电位水平。它通常被设置为一个负值,比如-70mV。这个负值意味着当神经元不被激活时,内部是趋向稳定的状态。
  • ur2：复位电位，当神经元放电产生脉冲后,为了准备下一次积分放电,需要将电位重新设置到一个复位值。这个值通常接近甚至等于静息电位,比如也设为-70mV。它使得神经元在放电后快速返回静息,准备下一次信息编码。
  • wjw_j: 第 j 个输入神经元的突触权重
  • tjkt^k_j:第 j 个输入神经元的第 k 个脉冲在TWT_W的积分时间窗口内激发的时间（总共 SjTwS^{Tw}_j脉冲序列），发放脉冲的时间
  • K():描述时间衰减效应的核函数
  • uthu_{th}:决定是否发射脉冲的发射阈值

ANN vs SNN

下面表格是两者的对比：

特性 人工神经网络 (ANN) 脉冲神经网络 (SNN)
通信方式 使用高精度和连续值编码的激活 通过二进制事件编码的脉冲序列通信
表达信息 连续激活值 脉冲模式
信息传播 主要通过空间传播 空间传播和时间多功能性
动态行为 通常较为简单 丰富的动态行为，当前状态受到过去历史的影响
信号特性 稠密的连续信号 稀疏的二进制脉冲，事件驱动的计算

2.2 ANN-TO-SNN CONVERSION

将ANN转换为SNN的两种主要方法，下面我将详细解释这两种方法：：

  • Training-free conversion
  • Training-dependent conversion

Training-free conversion

这种转换方法直接应用于预训练的ANN,转换过程包括阈值平衡、参数校准以及功能性脉冲发射。这些技术允许将ANN直接转换为SNN，且仅需少量样本进行校准，无需重新训练或微调。优势在于可以应用于高性能的开源ANN模型，但大多限于CNN，缺乏适用性于Transformers模型。

Training-dependent conversion

在转换之前，首先定制ANN以增强其与SNN的兼容性,或者在转换后对SNN进行微调。这种方法可以减少转换损失和延迟，但依赖于特定数据集，意味着更高的训练成本和较弱的泛化能力，并且通常也会维持类似CNN的结构限制。


本文研究属于上文之一的Training-free conversion方法，但将转换扩展到Transformer模型。

3. 面临的挑战

这部分内容将介绍为什么传统转化方法不适用于Transformer模型，也就是CNN和Transformer的关键性区别。

下面图分别是Transformer的一个模块(Transformer模型确实太过于庞大了，想展示详细结构一张图片放不下)和卷积神经网络结构。


3.1 非线性运算符

参考上图Transformer模型在处理数据时，使用的非线性函数比CNN更复杂。CNN通常使用ReLU函数，这是一种简单的非线性函数。相比之下，Transformer使用如GELU、平方根、指数这样的复杂非线性函数。这些函数的复杂性在于它们的计算过程不是简单的线性或分段线性（如ReLU那样），但是SNN中常用的神经元主要适用于处理线性或分段线性的响应。

因此，要在Transformer模型中使用这些复杂的非线性函数，需要采用特殊的方法在空间域内近似这些函数的行为。


3.2 变量乘标量与矩阵乘积

在CNN中，推理通常通过将变量特征（如图像的像素值）乘以恒定的权重矩阵来完成，对于下图中，使用不同的输入数据进行卷积操作时，卷积核是不会改变的，所以变量只有输入数据。

相比之下，Transformer模型涉及更多的变量-变量乘法操作。对于下图中自注意力机制中的query-key 乘积，Q和K的值是根据输入数据分别乘以WQW_Q和WkW_k得到的，所以Q和K进行相乘的时候，输入数据不同，两者也会不同。


在脉冲神经网络（SNN）中处理这种变量-变量的乘法操作时面临着额外的挑战。SNN通常使用时间序列来编码信息。在推理过程中，输入数据被编码成不完整的时间序列，意味着在某个特定时间点上，整个输入序列可能还没有完全呈现。如果在这样的时间点上直接计算两个特征矩阵的乘积，就可能产生非因果的结果——即，计算可能依赖于未来尚未出现的数据。这违反了因果关系，即原因必须先于结果。


当两个输入都在动态变化时（如转换器模型中的查询-键乘积），不仅需要实时处理当前的数据，还需考虑每个数据点的时间关联性。这增加了计算过程的复杂性，因为每个时间步的输出都依赖于当前以及之前时间步的输入。这种依赖性和实时处理需求在变量与固定权重的乘法中不那么明显。


综上所述，为了解决这个问题，需要采用特殊的方法来确保乘法操作在时间上是一致和稳定的，即在时间维度进行近似，从而克服非因果性的问题。
































4. 研究方法

基于目前存在的问题，作者提出了一种名为"Spatio-Temporal Approximation"（STA）的转换流程，包括两个新的脉冲模块作为通用逼近器。空间上，使用局部神经元群体模拟精确的非线性操作；时间上，采用估算-校正过程模拟自注意力中的乘法操作。


4.1 非线性的空间近似

根据上文，Transformer的浮点非线性对SNN转换提出了挑战，作者受通用近似定理启发，开发脉冲对应体来对非线性操作进行空间近似。

1
通用近似定理：一个包含足够多隐含层神经元的多层前馈网络，能以任意精度逼近任意预定的连续函数

所提出的近似器有以下特性:

  • 1)仅由神经元组成
  • 2)普遍适用于所有操作、模型和数据

这个提出的近似器被称为通用群运算符（UGO），它包含两个全连接层和一个隐藏的IF神经元层。

这个空间近似过程分为三个阶段：

  • 数据合成:在这个阶段,目的是生成覆盖函数 ff (这个函数可以使任意非线性函数，比如指数函数，gelu函数或者倒数函数等) 输入范围的合成数据集。由于Transformer的LayerNorm限制了函数的输入范围在一个小的连续区间 DD(例如GELU的 D=[−10,10]D = [-10, 10]),因此合成数据集 D^\hat{D} 会被创建为覆盖这个区间的均匀/正态分布混合体。然后,从这个合成数据集中采样 MM 个点以涵盖所有可能的输入。这些浮点数据对 {xi,f(xi)}\{x_i, f(x_i)\} 将作为合成训练数据,使得UGO能够在没有真实训练数据的情况下近似函数 ff。

  • ANN构建:在此步骤中,基于函数 ff 的复杂性手动选择一个合适的超参数大小 NN 来定义一个人工神经网络(ANN)f^\hat{f} 的规模。这个ANN f^\hat{f} 然后在合成数据上使用ReLU或其他定制激活函数进行训练,以近似函数 ff。

  • SNN转换:在预训练的ANN f^\hat{f} 被成功训练之后,它会被转换为一个脉冲神经网络(SNN)f^\hat{f},该网络由IF(Integrate-and-Fire)神经元组成,并在 TT 时间步内进行操作。


下面是相关代码实现细节：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class X2X(nn.Module):
    def __init__(self):
        super().__init__()
        self.approximator = nn.Sequential(
          """模拟UGO基本构成"""
            nn.Linear(1, 2, bias=False),
            nn.ReLU(),
            nn.Linear(2, 1, bias=False)
        )
        self.approximator.requires_grad = False
        self.approximator[0].weight.data[0, 0] = 1.0
        self.approximator[0].weight.data[1, 0] = -1.0
        self.approximator[2].weight.data[0, 0] = 1.0
        self.approximator[2].weight.data[0, 1] = -1.0

    def forward(self, x):
        dim = x.dim()
        if dim == 0:
            return self.approximator(x)
        else:
            return torch.squeeze(self.approximator(torch.unsqueeze(x, -1)))
1
2
3
4
5
6
7
8
9
10
11
def x2x_to_spike_module(ann_module, T, belong_to_ln=False):
    """将 X2X 类的模型转换为Spike模块"""
    snn_module = copy.deepcopy(ann_module)
    snn_module.approximator[0]
    snn_module.approximator[0] = SpikeLinear_ReLU(T=T, module=ann_module.approximator[0])
    snn_module.approximator[0].relu = nn.ReLU()
    snn_module.approximator[0].belong_to_x2x = True
    snn_module.approximator[0].belong_to_ln = belong_to_ln
    snn_module.approximator[1] = StraightThrough()
    snn_module.approximator[2] = SpikeLinear_ReLU(T=T, module=ann_module.approximator[2])
    return snn_module

4.2 乘法的时间近似

在传统网络中，乘法通常是固定权重矩阵和变量特征矩阵之间的操作。然而，在Transformer的自注意力机制中，涉及的是变量特征矩阵之间的乘法。由于在推理过程中，这些矩阵由不完整的时间序列编码，直接计算它们的乘积是非因果的。为了解决这个问题，提出了时间校正自注意力层（TCSA），使用估计-修正机制。

考虑一个简单的情况,其中 AA 和 BB 是一维矩阵(向量),且长度为 TT。在时间 tt( t≤Tt \leq T),我们只收到了 AA 和 BB 的前 tt 个元素的尖峰,A(i)A(i) 是 AA 在时间 ii 的脉冲值,B(i)B(i) 是 BB 在时间 ii 的脉冲值。

估计-修正机制步骤

初始化: 初始化一个累积值 Φ(t)\Phi(t) 用于存储估计的乘积,初始时 Φ(0)=0\Phi(0) = 0。

对于每个时间步 t=1,2,…,Tt = 1, 2, \ldots, T: 收到 AA 和 BB 在时间 tt 的脉冲值 A(t)A(t) 和 B(t)B(t)。

估计和修正:

- 估计 ABAB 的值:A^(t)=1t2∑i=1tA(i)\hat{A}(t) = \frac{1}{t^2} \sum_{i=1}^t A(i),B^(t)=1t2∑i=1tB(i)\hat{B}(t) = \frac{1}{t^2} \sum_{i=1}^t B(i)。

- 计算估计的乘积:Ψ(t)=A^(t)⋅B^(t)\Psi(t) = \hat{A}(t) \cdot \hat{B}(t)。

- 修正增量 Q(t)=t⋅Ψ(t)−(t−1)⋅Ψ(t−1)Q(t) = t \cdot \Psi(t) - (t - 1) \cdot \Psi(t - 1)。

- 更新累积值 Φ(t)\Phi(t):Φ(t)=Φ(t−1)+Q(t)\Phi(t) = \Phi(t - 1) + Q(t)。

最终乘积: 在时间 TT 时,Φ(T)\Phi(T) 是 ABAB 的最终估计值。












5. 研究结果

5.1 零样本分类

该部分涉及使用CLIP模型进行零样本分类(CLIP是一个多模态人工神经网络，训练于图像-文本对，使用了不同的图像编码器背骨架，包括ResNet和视觉Transformer)。实验中，使用预训练的ResNet-50作为基线，评估了CIFAR-10/100、ImageNet-200以及分布偏移的CIFAR-10.1/10.2数据集。结果显示，转换后的ViT模型在所有数据集和时间设置上显著超越了ResNet。















5.2 标准分类实验

在这部分内容中，研究人员对视觉Transformer (ViT) 模型进行了微调，并将其在传统图像分类任务上的性能与resnet-20和CLIP中的预训练ResNet-50基线进行了比较。结果表明，与其他转换方法相比，STA算法在较少的步骤（T=32或64）内达到接近峰值的准确度，而大多数基线需要超过128步才能实现最佳准确度。

5.3 能耗估计

SNN的高能效来源于两个方面：

  • 稀疏和事件驱动的计算
  • 低功耗的突触操作

研究发现，UGO在计算成本上相比GELU节省了约41%，而脉冲乘法在注意力模块上平均节省了33%，在最稀疏的情况下可达75%。

6. 研究意义

这项工作首次将预训练的主流Transformer模型成功转换为纯事件驱动的SNN，这对于神经形态硬件的部署具有重要意义。同时，这种转换方法在保持原始模型泛化性能的同时，避免了性能下降，促进了SNN在各种下游应用中的实际使用。

7. 研究局限性

尽管这项工作展示了将Transformer模型转换为SNN的潜力，但其能源节省优势不如传统的脉冲CNN。此外，转换的ViT模型相较于传统脉冲CNN，在保持较高性能和鲁棒性的同时，计算量略高，模拟步骤较少。

8. 有关感想

这项研究展现了ANN到SNN转换领域的显著进步，尤其是在处理复杂的Transformer模型方面。通过在空间和时间两个维度上的创新近似方法，这项工作在低能耗神经形态硬件上部署大规模预训练模型提供了新的可能性。