MT-RevSNN：一种面向分层脉冲变换器的低内存低时间开销可逆框架

《Neural Networks》：MT-RevSNN: Memory and Time Efficient Reversible Framework for Hierarchical Spiking Transformer

　　李思睿|宋泽阳|张学毅|张少辰|林哲源|蔡思奇|李海洲

摘要：脉冲变换器能够实现高效能的类脑计算，但在时空反向传播过程中需要O(L×T)的内存消耗，其中L为网络深度，T为模拟时间步数。可逆架构可通过在反向传播时重构激活值来降低内存需求，但在应用于分层变换器时面临两大限制：（1）不支持带有维度重缩放的分层脉冲变换器；（2）与整数型脉冲神经元不兼容，会导致梯度不稳定。为克服这些限制，我们提出了MT-RevSNN，这是首个可实现分层脉冲变换器O(1)内存消耗的可逆框架。具体而言，我们引入信息融合块和双因子缩放机制，以实现可逆的维度重缩放并稳定整数型神经元的动态行为。为降低实现成本，我们还提出了低秩压缩脉冲自注意力机制。在分层QKFormer架构上实现的MT-RevSNN，在ImageNet-1K数据集上的训练内存消耗降低了12.4倍，训练速度提升了1.9倍，同时保持了与QKFormer模型相当的准确率。这些结果展示了可逆架构在实现大规模脉冲变换器的高效训练方面的潜力。相关代码可在补充材料中获取。

引言：人脑通过数十亿神经元之间的稀疏、异步电信号传递来处理信息，这种机制具有无与伦比的能效优势。脉冲神经网络旨在通过将信息编码为离散的脉冲事件而非连续值激活函数来模拟这种生物计算方式（Maass，1997；Tavanaei等人，2019）。这种以事件驱动、稀疏激活为特征的架构为脉冲神经网络带来了两大优势。首先，在硬件层面，诸如Intel Loihi（Davies等人，2018）和IBM TrueNorth之类的类脑处理器可以利用脉冲的二元特性，用简单的累加操作替代能耗较高的乘累加操作，从而大幅降低推理时的能耗。其次，在算法层面，脉冲神经网络天然地将时间维度融入其动态行为中，使其具备天然的时空处理能力（Wu等人，2018），非常适合处理序列数据和基于事件的数据。这些特性推动了脉冲神经网络在静态图像识别、点云处理以及类脑感知领域的应用（Deng等人，2020；Tavanaei等人，2019）。该研究领域的一个重要发现是，脉冲神经网络并非仅仅是传统人工神经网络的近似替代品：只要采用合适的训练方法，深度脉冲神经网络即可在保持显著能效优势的同时，在大规模基准测试中取得与传统神经网络相当的准确率。然而，要将脉冲神经网络的理论能效优势转化为实际的大规模应用方案仍面临巨大挑战，因为必须同时考虑架构设计、学习算法以及硬件兼容性等问题。脉冲神经网络的架构发展大致跟随了人工神经网络研究的轨迹，尽管存在一定时间差。早期的脉冲神经网络模型大多采用较浅的拓扑结构（Kheradpisheh等人，2018），通常只有两到三层，且采用赫布规则或基于奖励的塑性规则进行训练。随着替代梯度方法的出现（Neftci等人，2019；Wang等人，2020），这一状况发生了重大转变。这类方法通过在反向传播过程中使用平滑的替代函数来规避阶跃脉冲函数的非可微性，从而使得深度脉冲神经网络架构可以通过时空反向传播方法进行训练（Wu等人，2018），进而获得处理复杂感知任务所需的表征能力。借助时空反向传播方法，该领域出现了深度脉冲神经网络版本的经典人工神经网络架构。例如，ResNet的脉冲版本（Fang等人，2021）证明了残差连接机制能够保持不同层级间的脉冲动态特性，而脉冲变换器则将自注意力机制改造成完全基于二元脉冲的运算方式（Li等人，2024；Yao等人，2024；Zhou等人，2022）。最近，该领域又出现了向分层脉冲变换器发展的趋势（Yao等人，2025；Zhou等人，2024），这类架构通过金字塔结构提取多尺度视觉特征。通过分阶段进行维度重缩放（如像素块合并），这类设计克服了单分辨率柱状模型的固有局限，在ImageNet-1K等大规模基准测试中取得了当前最优的准确率。尽管有了这些架构上的进步，但训练效率问题——尤其是时空反向传播过程中的GPU内存消耗问题——仍然是一个亟待解决的瓶颈。在正向传播过程中，时空反向传播方法会将网络动态在所有T个时间步中逐层展开，因此需要在每一层都存储中间阶段的神经元膜电位和激活值，以便后续进行梯度计算。这就导致了O(L×T)的训练内存需求，其中L为网络深度。对于那些包含数十层及多个时间步的大型分层脉冲变换器来说，如此高的内存成本显得极为高昂，严重限制了在普通硬件上可使用的批量大小、模型深度以及训练效率。为解决这一内存瓶颈，人们已经探索了多种策略。梯度检查点技术（Chen等人，2016）通过在后向传播过程中重新计算中间激活值来提供解决方案，但为此需要额外的正向计算，虽然有效，却会带来额外的计算开销，而脉冲神经网络训练本身就具有时间展开的特性，进一步加剧了这一问题。将人工神经网络转换为脉冲神经网络的方法（Li等人，2022；Wu等人，2021）则通过使用预训练的人工神经网络权重来规避这一问题。尽管近期的转换方法在减少所需模拟时间步数方面取得了显著进展，有些方法甚至只需T=1步即可完成训练（Chen等人，2025a），但与直接训练方法相比，它们在相同时间步数约束下的准确率仍有明显差距。另一种思路是采用受生物学机制启发的在线学习规则进行直接训练（Deng等人，2020），但目前这类方法的准确率在处理复杂视觉任务时仍落后于时空反向传播方法。越来越受到重视的另一种替代方案是可逆架构理念，这类架构的设计使得在反向传播过程中可以从后续层的输出中精确重建中间激活值，从而无需缓存这些激活值。这一理念最初在人工神经网络领域由RevNets（Gomez等人，2017）提出，随后也被应用于视觉变换器领域（Mangalam等人，2022），这类方法能够在不牺牲模型表达能力的前提下实现恒定内存消耗的训练。然而，将可逆性原理应用到脉冲神经网络中并非易事，因为脉冲神经元具有复杂的时空动态特性。近期的一些研究将可逆性原理引入了类脑计算领域。空间可逆方法（S-RevSNN）（Zhang和Zhang，2024）通过构建空间（深度）维度上的双射映射，将内存消耗降至O(T)。时间可逆方法（T-RevSNN）（Hu等人，2024a）则利用脉冲神经元在时间维度上的动态特性，将内存复杂度降低到O(L)。此外还有二元可逆脉冲神经网络（Cao等人，2025）以及并行可逆框架（Xu等人，2025）等进一步的发展。尽管取得了这些进展，目前的可逆脉冲神经网络实现方案仍存在两个相互关联的致命缺陷，阻碍了其在先进分层架构中的应用：（1）与维度重缩放不兼容。可逆网络的核心要求是各模块的输入与输出之间必须保持严格的维度一致性，而分层脉冲变换器则在层次转换时通过像素块合并和通道扩展等方式逐步聚合多尺度特征，这类维度改变操作与现有的可逆脉冲神经网络架构完全不相容，因此无法直接应用于分层架构中。（2）整数型脉冲神经元会导致梯度不稳定。为将内存复杂度从O(T)进一步降低到完全恒定的O(1)，整数型脉冲神经元（如I-LIF（Luo等人，2024）、IBRA-LIF（Ye等人，2025））具有很大应用潜力，因为它们的时间动态特性可以被压缩，从而避免需要按时间存储数据。然而，将这些神经元与深度可逆路径结合使用时，会导致严重的梯度不稳定以及脉冲消失现象。整数类型转换带来的量化误差会在深度的可逆残差连接中不断累积，最终导致在大规模数据集上的优化过程失败。这两个缺陷使得可逆脉冲神经网络在理论上虽有望为分层脉冲变换器实现完全恒定的O(1)训练内存消耗，但实际应用仍存在巨大差距。为克服这些限制，如图1所示，我们提出了MT-RevSNN，这是一种专为分层架构设计的、兼具内存效率和时间效率的可逆脉冲神经网络框架。MT-RevSNN通过三项协同创新的手段，首次实现了大规模分层脉冲变换器的O(1)训练内存复杂度。我们的主要贡献如下：• 可逆分层设计。我们提出了首个用于分层脉冲神经网络的可逆框架MT-RevSNN。通过将灵活的重连技术与信息融合块相结合，我们解决了可逆性要求与维度重缩放需求之间的矛盾。信息融合块负责处理所有层次间的维度转换，将训练内存复杂度从O(L×T)降低到O(T)，同时还能完整保留多尺度特征提取功能。• 通过双因子缩放稳定整数型神经元动态。我们引入了双因子缩放机制，该机制包含膜电位灵敏度因子和脉冲幅度因子，我们将该机制应用于I-LIF和IBRA-LIF两种神经元，从而得到了SI-LIF和SIBRA-LIF两种新型神经元。双因子缩放机制能够解决深度可逆模型中的梯度不稳定问题，使整数型脉冲神经元的动态特性得以稳定，进而将训练内存复杂度降低到O(1)。• 低秩压缩脉冲自注意力机制。受多头潜在注意力的启发（Liu等人，2024），我们提出了一种轻量级的脉冲注意力模块，该模块可将键矩阵和值矩阵投影到共享的低维潜在空间中，从而大幅减少参数数量和计算量，且不会影响模型准确率。• 在保持相当准确率的同时实现更高效率。在分层QKFormer架构上实现的MT-RevSNN，在ImageNet-1K数据集上的准确率可与当前最先进的脉冲神经网络模型相媲美，同时训练内存消耗降低了12.4倍，训练速度提升了1.9倍。

章节节选：脉冲神经网络训练：目前脉冲神经网络的训练方法主要分为两类：人工神经网络到脉冲神经网络的转换，以及使用替代梯度进行直接训练。人工神经网络到脉冲神经网络的转换方法是将ReLU激活函数替换为脉冲神经元，将连续输出转换为离散的脉冲频率（Li等人，2022；Wu等人，2021）。虽然这种方法可以利用预训练的人工神经网络权重，且通常只需少量训练轮次，但由此得到的脉冲神经网络仍需要较长的模拟时间。相比之下，直接训练方法……

方法：我们提出了MT-RevSNN，这是一个旨在解决大规模分层脉冲神经网络训练瓶颈的综合性框架。MT-RevSNN通过三个关键组件的协同作用实现了O(1)的训练内存消耗：（1）采用带有信息融合块的可逆设计，该设计能够实现维度重缩放，将内存复杂度从O(L×T)降低到O(T）；（2）将整数型脉冲神经元与我们提出的双因子缩放机制相结合，进一步将内存复杂度降低到O(1)；（3）低秩压缩脉冲自注意力机制，这是一种低秩结构的脉冲注意力模块……

实验：我们将MT-RevSNN集成到了分层脉冲神经网络变换器QKFormer（Zhou等人，2024）中：我们利用信息融合块配合重连技术将原有的非可逆架构转变为可逆架构，用低秩压缩脉冲自注意力机制替代传统的自注意力机制，同时用SI-LIF和SIBRA-LIF这两种新型神经元替代普通的LIF神经元。这样便得到了两种变体：MT-RevSNN（SI-LIF）和MT-RevSNN（SIBRA-LIF）。我们在静态数据集（ImageNet-1K、CIFAR-10/100）以及类脑计算数据集DVS128上对这两种模型进行了测试。

结论：我们提出了MT-RevSNN，这是一种基于QKFormer构建的可逆脉冲神经网络框架，它将三种关键创新技术——通过信息融合块实现灵活维度重缩放的重连技术、用于稳定基于整数型神经元动态的双因子缩放机制，以及用于降低计算成本的低秩压缩注意力机制——整合到了一个模块化设计中。这种组合不仅提升了内存效率，还提高了训练速度，使得MT-RevSNN在准确率方面能够与当前最先进的模型相媲美，甚至更优。

CRediT作者贡献说明：李思睿：撰写——初稿撰写、软件实现、方法设计、形式分析、数据整理、概念构思。宋泽阳：撰写——审阅与编辑、指导工作。张学毅：撰写——审阅与编辑。张少辰：撰写——审阅与编辑。林哲源：撰写——审阅与编辑。蔡思奇：撰写——审阅与编辑、指导工作、资金筹集。李海洲：撰写——审阅与编辑、指导工作、资金筹集。

利益冲突声明：作者声明不存在任何可能影响本文研究结果的已知财务利益或个人关系。