《Nature Communications》：Model-agnostic linear-memory online learning in spiking neural networks
编辑推荐：
《自然·通讯》编辑推荐：本研究针对脉冲神经网络（SNN）在线学习系统在长时程训练中内存消耗巨大、模型通用性受限的难题，开发了名为BrainTrace的模型无关、线性内存的自动化在线学习系统。该系统通过AlignPre/AlignPost建模框架标准化模型构建，提出线性内存的pp-prop学习算法，并提供了可自动生成优化代码的编译器。研究证实BrainTrace在多种任务中实现了高性能、低内存消耗和高计算吞吐量，并成功应用于全脑尺度的果蝇SNN建模，为大规模脉冲神经网络建模奠定了基础。
大脑以其惊人的能效和毫秒级的时间精度处理着复杂信息，而脉冲神经网络（SNN）被认为是模拟这种生物神经计算的理想范式。然而，在SNN中进行有效学习一直面临一个根本性挑战：长时程时间信用分配问题。这意味着，突触的更新需要依赖于跨越数千甚至数百万个离散时间步长的状态轨迹。目前的主流解决方案——基于替代梯度的通过时间反向传播（BPTT）算法，虽然能计算精确的序列级梯度，但其离线学习特性要求存储所有中间状态，导致内存成本随序列长度T线性增长。对于需要模拟生物时间尺度的任务（例如，一秒可能对应数千个求解器步骤），这种内存占用很快变得难以承受，严重限制了可训练序列的时长以及SNN模型的规模和复杂性。
在线学习算法作为有吸引力的替代方案应运而生。它们通过仅向前传播的方式更新参数，将内存使用限制在一小组每时间步变量上，从而避免存储整个状态轨迹。然而，尽管取得了令人鼓舞的进展，现有方法未能同时实现模型无关的通用性和线性内存缩放。那些接近BPTT精度水平的先进方法，如e-prop和OSTL，需要实例化大量的资格变量，其内存和计算成本随网络规模呈二次方增长，使得它们对于大规模脉冲网络不切实际。而那些实现线性内存复杂度的方法（如OTTT、OTPE等）通常依赖于简化的神经元抽象（例如简化的泄漏积分发放（LIF）动力学），无法推广到生物保真度所需的丰富神经元动力学。最关键的是，目前缺乏类似于现代自动微分（AD）生态系统的SNN在线学习工具：用户无法编写任意的脉冲动力学模型并期望系统自动合成数值稳定、高效且准确的在线更新规则。这种自动化、模型无关框架的缺失，阻碍了其在专业团体之外的广泛应用，并限制了其现实影响力。
为了填补这一空白，王超明、董行思等研究人员在《自然·通讯》（Nature Communications）上发表了他们的研究成果，引入了BrainTrace——一个模型无关、线性内存且自动化的脉冲神经网络在线学习系统。该系统提供了一个理论保证的、用于训练通用脉冲网络的线性内存在线梯度算法，以及一个能够从用户定义的SNN模型自动生成高效在线学习代码的编译器。在效率方面，BrainTrace利用脉冲神经元稀疏、事件驱动的特性，重新表述了实时循环学习（RTRL），在保持有效训练所需梯度保真度的同时，实现了相对于神经元数量的线性内存缩放。在通用性方面，这种效率扩展到了一大批脉冲网络动力学模型，其突触状态可以沿着突触前或突触后轴进行组织。在易用性方面，BrainTrace的编译器能为任何用户定义的SNN自动生成高效的在线学习代码，便于在中央处理器（CPU）、图形处理器（GPU）和张量处理器（TPU）等各种硬件平台上实现。
为了验证BrainTrace，研究团队在多种SNN架构和任务上进行了评估。首先，他们评估了BrainTrace编译算法在梯度计算近似方面的有效性。结果表明，无论是在单步隐藏雅可比矩阵还是长程隐藏雅可比矩阵上，BrainTrace的D-RTRL和pp-prop算法与BPTT计算的结果均保持了较高的余弦相似度，验证了其梯度近似的有效性。特别是在延迟匹配样本（DMTS）任务中，研究发现D-RTRL和pp-prop即使不使用替代梯度也能成功学习，显示出其捕捉长时程依赖关系的能力。在内存消耗和计算速度方面，与BPTT相比，pp-prop算法实现了恒定的内存使用（不随序列长度增长），并且比D-RTRL内存效率高35倍，速度提升10倍。
进一步地，研究团队在神经形态计算基准任务（如N-MNIST、SHD、DVS Gesture数据集）上评估了pp-prop的性能。结果显示，pp-prop在大多数设置下匹配甚至超过了BPTT的性能，同时使用了在线、内存高效的更新规则。与包括OSTL、e-prop在内的先进在线学习算法相比，pp-prop consistently实现了强劲的性能，例如在SHD数据集上超过了最强基线10多个百分点。
为了展示BrainTrace在处理生物真实电路方面的能力，研究人员构建了一个包含兴奋-抑制（E/I）群体、传导性突触和稀疏连接性的生物启发脉冲网络，并在证据积累任务上对其进行训练。结果表明，pp-prop在保持高性能的同时，内存占用远低于BPTT和D-RTRL。训练后的网络动力学显示出高度选择性、时间精确的发放模式，与小鼠在执行类似任务时后顶叶皮层观察到的神经活动高度吻合。
最具突破性的展示是BrainTrace在线性内存复杂度方面的优势，使其能够训练大规模脉冲网络。研究人员利用果蝇全脑连接组（FlyWire connectome）数据，构建了一个包含超过12.5万个神经元和5000万突触连接的全脑脉冲模型，并利用pp-prop算法训练其突触权重，以重现实验记录的全脑静息态钙成像信号及其区域间功能连接。训练后，模型不仅在单个神经毡（neuropil）水平上再现了实验观察到的自发波动，在全脑水平上生成的功能连接模式也与实验数据高度一致。这项工作首次实现了在单GPU上对如此大规模的全脑脉冲网络进行长时程训练，为利用日益丰富的大规模神经科学数据和连接组数据集构建功能模型提供了可行的计算工具。
本研究主要采用了以下几种关键技术方法：1) 提出了AlignPre和AlignPost两种规范化的脉冲网络建模框架，用于结构化描述多样的神经元和突触动力学；2) 开发了基于实时循环学习（RTRL）理论、但通过利用脉冲动力学的内在特性（如稀疏性、符号一致性）进行近似优化的在线学习算法，包括对角线近似的RTRL（D-RTRL）和预-后资格迹传播（pp-prop）算法，实现了线性内存复杂度；3) 构建了一个自动化在线学习编译器，能够分析用户定义的模型（基于Python/JAX），自动生成高效的在线学习代码（包括资格迹更新和梯度计算逻辑），并利用即时（JIT）编译技术部署到不同硬件平台（CPU/GPU/TPU）；4) 在算法验证中，使用了包括果蝇全脑静息态钙成像数据（来自公开数据集）在内的多种神经形态数据集（如IBM DVS Gesture, SHD, N-MNIST）和认知任务（如延迟匹配样本、证据积累任务）进行性能评估。
BrainTrace支持广泛的脉冲网络
BrainTrace支持具有特定离散化动力学的多种SNN。它将SNN分解为两部分：一是支配神经元状态时间演化的“内在动力学”部分；二是通过相互作用算符将突触前输入转换为突触后电流的“神经元相互作用”部分。这种分解通过AlignPre和AlignPost两种抽象实现，分别将突触变量与突触前或突触后神经元维度对齐，从而显著降低了计算复杂性和内存消耗。
BrainTrace实现线性内存在线学习算法
SNN训练的目标是找到一组突触权重θ以最小化损失函数L。研究人员通过利用SNN的固有特性，对标准的RTRL算法进行了大幅简化。首先，他们发现由内在动力学引起的对角线雅可比矩阵D
t主导了完整的隐藏雅可比矩阵。其次，他们证明了隐藏状态对权重的雅可比矩阵?h
t/?θ
t可以分解为对角矩阵D
ft和输入x
t的克罗内克积。基于这些观察，他们将资格迹近似为对角线形式（D-RTRL），将内存复杂度从O(H
3)降至O(H
2)。进一步地，他们利用脉冲神经元输出符号一致的特性，提出了一个秩一近似，将资格迹最终简化为分别追踪突触前和突触后活动的两个变量（ε
xt和ε
ft）的外积，这就是pp-prop算法，最终实现了O(H)的线性内存复杂度。
t/?θ
texhibit cubic memory complexity, while the hidden Jacobian ?h
t/?h
t-1demonstrates quadratic memory complexity. B Cosine similarity between the diagonal Jacobian D
tinduced by intrinsic dynamics in Eq. 1 and the full hidden Jacobian ?h
t/?h
t-1, over different network dynamics and fraction of spiking neurons. C In D-RTRL computation, the eligibility trace ? is reduced to exhibit quadratic memory complexity D Cosine similarity between the rank-one approximation(Eq. 5) and D-RTRL eligibility trace(Eq 4) for LIF neurons with exponential synapse(Supplementary Eq. 72). Similarity is computed for membrane potential(v) and synaptic conductance(g). E Cosine similarity between the rank-one approximation(Eq. 5) and D-RTRL eligibility trace(Eq. 4) for ALIF neurons with exponential synapse(Supplementary Eq. 73). Similarity is computed for membrane potential(v), synapse(g),and adaptation(a) variables. F In pp-prop computation, the eligibility trace ? is approximated using two variables with linear memory complexity that separately track pre-synaptic and post-synaptic neuronal activities.'>
BrainTrace实现自动化在线学习编译
受现代深度学习框架（如PyTorch, TensorFlow）自动化训练成功的启发，BrainTrace引入了一个编译器，用于将SNN自动转换为高效的在线学习实现。该编译器在Primitive层定义了三种基本元素来捕获通用SNN的基本组成部分。在ETraceGraph阶段，将用户定义的基于Python的模型转换为符号中间表示（IR），分析模型参数和神经元状态之间的数据依赖性。在ETraceAlgorithm阶段，根据用户选择的算法（D-RTRL或pp-prop）自动生成优化的在线学习实现代码，并进行高级优化。最后，利用JAX和XLA的即时（JIT）编译能力，将用户定义的模型和编译器生成的算法自动部署到目标硬件平台。
算法在近似性、长时程依赖、内存和速度上的评估
研究人员首先评估了BrainTrace编译算法在近似梯度计算方面的有效性。在不同脉冲数据集、神经元动力学和突触模型上的测试表明，D-RTRL和pp-prop均能实现与BPTT计算梯度的高余弦相似度。在DMTS任务中，只有D-RTRL和pp-prop能够在不使用替代梯度的情况下成功学习，显示了其捕捉长时程依赖关系的能力。分析表明，pp-prop通过分别追踪突触前活动和突触后雅可比动力学信息的资格迹来维持时间信息。在内存和速度方面，pp-prop表现出恒定且低的内存消耗，以及优于D-RTRL的计算速度。
在神经形态分类任务上的性能评估
在标准基准测试上的广泛实验表明，pp-prop在多种SNN架构（RLIF, RadLIF, EGRU）上，在大多数设置下匹配了BPTT的性能，同时使用了在线、内存高效的更新规则。与最先进的在线学习算法（如OSTL, e-prop, ETLP等）相比，pp-prop consistently实现了强劲的性能，特别是在SHD数据集上表现突出。
为认知任务训练兴奋-抑制脉冲网络
研究人员构建了一个包含四变量广义泄漏积分发放（GIF）神经元、不同E/I群体、传导性突触和稀疏连接性的生物启发网络，并训练其完成证据积累任务。结果表明，pp-prop在保持高性能的同时，内存占用显著低于BPTT和D-RTRL。训练后的网络动力学显示出高度选择性、时间精确的发放模式，与生物实验中观察到的神经活动相似。
通过全脑脉冲网络训练模拟果蝇静息态功能连接
BrainTrace的线性内存复杂度使其能够训练大规模脉冲网络。研究人员利用FlyWire连接组数据构建了果蝇全脑脉冲模型，并使用pp-prop算法训练其突触权重以重现实验记录的静息态神经活动。训练后，模型不仅在神经毡水平上再现了实验观察到的自发波动，在全脑水平上生成的功能连接模式也与实验数据高度一致，证明了模型泛化到未训练神经动力学的能力。
综上所述，BrainTrace通过联合优化通用性、效率和可用性这三个很少能同时实现的维度，推动了SNN在线学习的发展。在算法层面，pp-prop揭示了脉冲神经元的事件驱动动力学天然支持仅具线性复杂度的准确在线梯度近似。在系统层面，BrainTrace将这一见解转化为实用的工具链，扩展了通用脑模型平台BrainPy的在线学习支持能力。在应用层面，BrainTrace在大规模脑模拟模型的长时程训练中展现了独特优势，其成功拟合果蝇全脑神经活动的案例表明，它为解决神经科学领域如何利用丰富数据集构建功能模型这一紧迫挑战提供了强大的计算基础。尽管存在一些局限性（如近似质量对发放统计和网络深度的依赖性，以及需要扩展更多可塑性机制等），但BrainTrace无疑为大规模、基于生物学的脉冲网络建模铺平了道路，有望在神经科学、神经形态计算和人工智能领域产生深远影响。
