Randomized Forward Mode Gradient for Spiking Neural Networks in Scientific Machine Learning

TL;DR

RFG以随机权重扰动替代反向传播,PDE回归误差最低达0.0347并降低约66%训练计算。

cs.NE 🔴 高级 2024-11-11 19 次浏览
Ruyin Wan Qian Zhang George Em Karniadakis
脉冲神经网络 随机前向梯度 科学机器学习 神经算子 神经形态硬件

核心发现

方法论

论文提出Randomized Forward-Mode Gradient(RFG):对参数施加Bernoulli随机扰动δθ,仅通过前向传播计算方向导数,再以扰动向量重构梯度。脉冲神经元采用IF模型,并用Gaussian surrogate gradient或基于Stein引理的Weak Surrogate Gradient近似Heaviside函数导数;同时比较全局扰动与逐层扰动。

关键结果

  • 在二维Mexican hat wavelet回归中,SG-RFG-G、SG-RFG-L和SG-BP的relative L2误差分别为0.0426、0.0342和0.0298,说明逐层RFG明显优于全局RFG,但仍略逊于反向传播。
  • 在1D Poisson方程DeepONet任务中,训练和测试各使用800个GRF样本;SG-BP与WSG-BP误差为0.0246和0.0241,而SG-RFG-G、WSG-RFG-G、SG-RFG-L、WSG-RFG-L分别为0.0422、0.0412、0.0347和0.0359。
  • 反向传播需要前向加反向计算,作者估计RFG训练成本降低约66%;逐层RFG在reaction-diffusion PDE上获得与BP相近的预测,局部损失LL反而未改善Poisson结果。

研究意义

该工作回应了SNN长期依赖反向传播的矛盾:反向传播需要链式法则、跨层误差信息和存储中间激活,与生物突触可塑性及Loihi 2等神经形态芯片的局部、并行计算不完全匹配。RFG把学习转化为扰动—观察—更新过程,在科学机器学习回归和PDE求解中保持可用精度,为低功耗物理建模提供了新的训练路线。

技术贡献

核心贡献是将随机前向模式自动微分引入脉冲DeepONet与SepONet。全局形式估计g(θ)≈(∇L·δθ)δθ,逐层形式则分别计算每层方向导数;JVP负责前向传播中的内积,surrogate gradient处理脉冲不可导性。该设计消除完整反向传播,并允许不同层并行扰动,形成更适合局部突触更新和专用采样硬件的工程接口。

新颖性

相较ANN-to-SNN转换和使用surrogate gradient的反向训练,本文直接在SNN中采用随机权重扰动与forward-mode gradient联合优化,并系统比较SG/WSG、BP/RFG及global/layer-wise组合。创新重点不是提出新的神经元,而是把RFG机制具体化为科学神经算子训练方案。

局限性

  • 实验规模较小,网络仅含1至2个隐藏层、16或32个神经元,尚未证明RFG在深层、大参数SNN中的稳定性与扩展性。
  • 当前每次迭代只使用一个扰动方向;随机估计方差、学习率和扰动分布对收敛的影响未被系统量化。

未来方向

作者计划研究多方向扰动,并在Intel Loihi-2上实现验证。进一步需要分析扰动方差与样本复杂度、设计更低方差的分布或控制变量,并测试更深网络、更多维PDE及真实硬件能耗。

AI 总览摘要

脉冲神经网络以离散尖峰传递信息,具有稀疏计算和低能耗潜力,尤其适合Loihi 2等神经形态硬件。但主流SNN训练仍依赖反向传播:它需要链式法则、跨层误差传递和中间状态保存,既缺乏生物合理性,也可能削弱硬件优势。

Wan、Zhang和Karniadakis提出Randomized Forward-Mode Gradient(RFG),通过随机扰动权重并观察损失变化估计方向梯度。IF神经元负责脉冲动力学,Gaussian surrogate gradient或Stein引理构造的WSG近似Heaviside导数;训练时采用全局或逐层扰动,不执行完整反向传播。RFG随后被用于脉冲MLP、DeepONet和SepONet。

结果显示,该方法虽通常不及BP精确,却达到竞争性表现。Mexican hat回归的relative L2误差为0.0426(全局RFG)、0.0342(逐层RFG)和0.0298(BP);1D Poisson任务中逐层SG-RFG误差为0.0347,而SG-BP为0.0246。作者估计训练计算量下降约66%,并观察到逐层RFG在reaction-diffusion PDE上接近BP。局限是实验规模小、单方向扰动方差未充分研究;多方向采样和Loihi-2部署将决定其实际价值。

深度分析

研究背景

SNN以LIF/IF神经元模拟膜电位和尖峰,利用稀疏事件降低能耗。ANN-to-SNN转换可继承ANN精度,但转换成本高且可能损失稀疏优势;直接训练通常使用surrogate gradient加BP。科学机器学习中的DeepONet学习函数到函数的算子,SepONet则为各空间维度设置独立trunk网络,为SNN提供了PDE建模平台。

核心问题

Heaviside尖峰函数不可导,使标准梯度无法直接传播。BP还需要跨时间、跨层传递误差信号并保存活动,难以对应局部突触学习,也不适合部分神经形态硬件。问题是:能否仅依靠前向计算和随机扰动,在保持PDE回归精度的同时降低训练成本?

核心创新

本文把RFG用于SNN科学算子学习。第一,使用随机参数扰动估计∇L方向信息,而非完整链式反传;第二,比较Gaussian SG与Stein-based WSG;第三,提出global和layer-wise两种扰动粒度。逐层方案可并行计算各层JVP,更接近局部更新,并在实验中通常优于全局扰动。

方法详解

  • �� IF动态:U_l[t]=U_l[t−1]+W_lS_{l−1}[t]+B_l,并在越过U_threshold后产生S_l=H(U_l−U_threshold),随后重置膜电位。
  • �� 代理导数:SG使用f(x)=exp(−x²/2σ²)/(σ√(2π));WSG用δ_k∼N(0,σ²I)近似Stein弱形式。
  • �� RFG:采样Bernoulli扰动δθ,前向计算∇L·δθ,再令g≈(∇L·δθ)δθ;layer-wise版本逐层估计。
  • �� 更新:θ_l←θ_l−ηg(θ_l)。网络使用脉冲MLP构成DeepONet/SepONet。

实验设计

任务包括二维Mexican hat wavelet回归、1D Poisson方程和非线性reaction-diffusion PDE。前者使用10,000个[−1,1]²样本、2层16单元和32时间步;Poisson使用GRF生成的800训练与800测试样本、branch/trunk各1层32单元;PDE使用200个GRF函数,100个训练、100个测试,网格100×100,SepONet宽度16。比较SG/WSG、BP/RFG及global/layer-wise扰动。

结果分析

Mexican hat任务误差依次为0.0426、0.0342和0.0298,逐层RFG优于全局RFG。Poisson中BP约0.024,而RFG为0.0347–0.0422;WSG未稳定优于SG。加入局部损失LL后BP和RFG均未改善。reaction-diffusion可用仅1层16单元建模,逐层RFG预测与BP相近。

应用场景

适合需要低功耗、事件驱动推理与物理建模的场景,例如边缘PDE代理模型、实时动力系统预测和神经形态科学计算。实际应用需将随机扰动、脉冲编码和训练流程映射到Loihi 2等芯片,并验证真实能耗,而不仅是算法操作数。

局限与展望

目前证据来自小型网络和有限PDE任务,未报告大规模数据集、长时间序列或真实芯片实验。单扰动方向可能带来较高估计方差,且RFG精度仍落后BP。未来应进行多方向并行采样、方差控制、深层网络测试和硬件能效基准。

通俗解读 非专业人士也能看懂

把训练想成调试一座会发信号的工厂。工厂里不是连续流水线,而是工人偶尔亮灯;灯亮过阈值就代表一个“尖峰”。传统方法像总经理从最终产品一路倒查每个工位,告诉所有工人自己该改多少。这很精确,却需要保存整个生产过程,而且不符合工人只根据邻近信息学习的方式。

RFG采用另一种办法:每轮随机把一些旋钮轻轻拧动一次,然后只看最终产品变好还是变坏。如果变好,就记录这次旋钮组合;如果变坏,就反向调整。多次重复后,这些“试一下”的结果就能指示改进方向。逐层扰动相当于一次只检查一个车间,信息更清楚,所以通常比全厂同时乱动更有效。

论文把这套办法用于函数回归和PDE求解。它的准确度略低于传统倒查法,但Mexican hat误差达到0.0342,Poisson误差达到0.0347,并估计训练计算减少约66%。因此,它像一种更适合低功耗智能工厂的学习方式,但还需要更大规模测试和真实芯片验证。

简单解释 像给14岁少年讲一样

想象你在玩一个会自动调节角色属性的游戏。角色不是连续说话,而是只有在能量条超过某个值时才发出“砰”的信号,这就是脉冲神经网络。问题是,游戏结束后,如果分数不好,传统训练要从结局倒着检查每一步,告诉每个按钮该往哪边按。这叫反向传播,效果不错,但过程复杂,也不像大脑那样只靠附近神经元的活动学习。

这篇论文的方法更像试装备:把攻击力、防御力等参数轻轻随机改一下,打一局,再看分数变高还是变低。如果变高,就朝这个方向继续;如果变低,就换方向。它不需要从结局一路倒查,而是只做前向尝试。作者还比较了“全队一起换装备”和“每个小队分别换装备”,后者通常更好。

实验包括波浪形函数、Poisson方程和反应扩散方程。Mexican hat任务中,逐层方法误差为0.0342,传统方法为0.0298;Poisson任务中逐层方法为0.0347,传统方法约为0.0246。也就是说,准确度稍差,但训练计算量大约减少66%。

为什么重要?因为未来芯片可能像大脑一样用很少的电力处理信息。现在它还只在小模型上测试,而且每次只尝试一个随机方向,像只试一套装备。下一步要同时尝试多个方向,并把算法放进Loihi 2芯片,看看现实中是否真的更省电!

术语表

Spiking Neural Network (脉冲神经网络)

用离散尖峰而非连续激活传递信息的网络。尖峰稀疏性可降低计算和能耗。

本文所有MLP、DeepONet和SepONet均采用脉冲网络。

Randomized Forward-Mode Gradient (随机前向模式梯度)

通过随机参数方向和一次前向方向导数估计梯度的方法。其核心估计为(∇L·δθ)δθ。

本文用RFG替代SNN的完整BP。

Surrogate Gradient (代理梯度)

用平滑函数近似不可导尖峰函数导数的训练技术。它使Heaviside脉冲能够参与梯度计算。

论文比较SG和Stein-based WSG。

DeepONet (深度算子网络)

由branch和trunk网络组成,用于学习函数空间之间的映射。输出通常是两者特征内积。

用于1D Poisson算子学习。

SepONet (可分离算子网络)

为不同空间维度设置独立trunk网络的算子架构。各维特征外积后与branch表示组合。

用于reaction-diffusion PDE。

开放问题 这项研究留下的未解疑问

  • 1 RFG在深层、大参数SNN中的收敛性和样本复杂度尚不清楚;需要理论分析扰动方差、学习率与网络深度的关系。
  • 2 单方向Bernoulli扰动可能造成高方差。多方向采样是否能在硬件并行下同时提升精度与效率,仍缺乏定量验证。
  • 3 论文未提供Loihi-2实测能耗。算法操作数减少是否能转化为芯片级优势,还取决于随机采样、存储和通信开销。

应用场景

近期应用

边缘PDE代理模型

工程团队可用脉冲DeepONet学习参数化PDE的输入函数到解的映射,再部署到低功耗边缘设备。前提是训练任务规模较小,并需先验证RFG与BP的精度差距。

神经形态动力学预测

在实时传感器系统中,事件驱动SNN可预测扩散、反应或结构响应。RFG减少反向训练依赖,适合探索局部更新和并行硬件实现。

远期愿景

Loihi-2上的科学计算

若多方向RFG能降低估计方差并在Loihi-2上高效执行,未来可形成面向物理仿真的低能耗训练—推理平台,服务机器人、能源和数字孪生。

原文摘要

Spiking neural networks (SNNs) represent a promising approach in machine learning, combining the hierarchical learning capabilities of deep neural networks with the energy efficiency of spike-based computations. Traditional end-to-end training of SNNs is often based on back-propagation, where weight updates are derived from gradients computed through the chain rule. However, this method encounters challenges due to its limited biological plausibility and inefficiencies on neuromorphic hardware. In this study, we introduce an alternative training approach for SNNs. Instead of using back-propagation, we leverage weight perturbation methods within a forward-mode gradient framework. Specifically, we perturb the weight matrix with a small noise term and estimate gradients by observing the changes in the network output. Experimental results on regression tasks, including solving various PDEs, show that our approach achieves competitive accuracy, suggesting its suitability for neuromorphic systems and potential hardware compatibility.

cs.NE math.NA math.OC