神经形态计算是一种类脑计算范式,一般是指在神经形态芯片上运行脉冲神经网络(Spiking Neural Network,SNN)。
本质来讲,神经形态计算,是一种由算法驱动硬件的设计范式。凭借低功耗的优点,神经形态计算也被认为是替换传统 AI 的“潜力股”。
对神经形态计算的理解应当从系统层面展开,不应该孤立地只看算法、或者只看硬件。
而神经形态计算中的“类脑”,指的是脉冲神经元能够模拟生物神经元的结构和功能。
通过这种模拟:一方面神经形态计算具有仿生复杂时空动态,另一方面神经形态计算可以利用脉冲信号来传递信息。
前者,让脉冲神经元模型的表达能力在理论上高于基于传统人工神经网络(Artificial Neural Network,ANN)的人工神经元模型;后者,让脉冲神经元具有脉冲驱动的计算特征。
当脉冲神经网络运行在神经形态芯片上时,只有在输入脉冲信号的时候,才会触发稀疏计算。否则,神经元就会处于静息状态。因此,要想实现低功耗的神经形态,脉冲驱动是一个必不可少的要素。
目前,神经形态计算领域面临这样一个严峻的现实问题:相比传统的人工神经网络算法,脉冲神经网络算法在任务性能上远远不及前者,很难满足各种复杂场景的需求。
对于边缘计算场景来说,往往要满足低功耗、低时延的要求。而一旦解决脉冲神经网络在算法层面的性能问题,再结合神经形态芯片的优势,神经形态计算的优势将能得到极大凸显。
中国科学院自动化所李国齐研究员和团队认为,神经形态计算的性能潜力远远未被发掘。
例如,在神经网络架构方面,目前绝大多数神经形态计算的应用,都围绕脉冲卷积神经网络(Convolutional Neural Network,CNN)展开,当前的神经形态芯片也只能支持脉冲 CNN。
相比之下,传统深度学习早已借助 Transformer 架构,在各种任务上取得了巨大突破。而直到该团队提出一系列 Spike-driven Transformer 模型之后,神经形态计算领域才将脉冲驱动范式融入 Transformer 架构。
图 | 李国齐(来源:李国齐)
脉冲神经网络和 Transformer,到底该怎样结合?
对于李国齐来说,围绕脉冲神经网络的工作,最早可以追溯到发表于 2018 年的工作。那时,他还在清华大学类脑计算中心施路平教授团队工作。
他表示:“施老师团队提出了替代梯度时空的反向传播算法,解决了脉冲神经网络领域基本的训练问题。”
但是,由于基础编程框架缺乏、二值脉冲不可导、深度网络脉冲退化等问题的存在,导致直到 2021 年之前脉冲神经网络最多也只有十几层。
如此之小的规模导致相比传统深度学习的性能而言,脉冲神经网络的性能远远落后于前者。后来,脉冲神经网络也开始朝着越来越深的方向发展。
比如说,2021 年李国齐团队在国际先进人工智能协会会议(Association for the Advancement of Artificial Intelligence,AAAI)上发表的论文,解决了脉冲神经网络的深层训练问题。
加入中国科学院自动化所之后,李国齐和北京大学田永鸿教授在 Science Advance 合作发表了关于脉冲神经网络开源训练框架 SpikingJelly 的论文。
这篇论文解决了领域内训练框架缺失的问题,极大降低了脉冲神经网络的学习门槛。
与此同时,李国齐团队和田永鸿团队还分别提出两种不同残差深度的脉冲网络,目前已经成为领域内的通用残差架构。
这两种架构能让脉冲神经网络做到几百层的深度,而且能够避免脉冲退化的问题,解决了大规模脉冲神经网络训练在深度和规模上的技术瓶颈。
虽然脉冲神经网络与人工神经网络之间的性能间隙已经被极大缩小,但是这还远远不够。Transformer 架构是深度学习的里程碑,也引起了脉冲神经网络领域内学者们的兴趣。
大约从 2022 年开始,脉冲 Transformer 的相关工作陆续面世。这些研究基本都是将 Transformer 架构中的一部分人工神经元,替换成脉冲神经元。
一些关键的操作比如自注意力算子等都被保留,从而让任务性能得到保障。
这些早期工作为李国齐团队的工作带来了启发。但是,他们觉得这更像是一种人工神经网络/脉冲神经网络的异构。
于是,课题组提出这样一个问题:“脉冲神经网络和 Transformer,究竟该以怎样的方式结合,才能同时汲取二者的优势?”
经过反复思考和不断讨论,课题组最终选择“脉冲驱动自注意力算子”作为问题的突破口。
原因在于:目前脉冲神经网络领域内的脉冲驱动算子,还只有卷积和全连接这两种。
而自注意力机制是 Transformer 成功的关键,那么能否把自注意力机制改为脉冲驱动?
确定这个思路之后,他们进行了反复实验,最终设计出一些可以正常工作的脉冲驱动自注意力算子。
结果发现:脉冲驱动自注意力算子拥有许多优良特性,比如天然就是线性算子以及稀疏运算等。
当脉冲驱动 Transformer 可以正常工作以后,他们试图通过改进架构来进一步提升性能。
但是,Transformer 架构的变体太多,甚至令人眼花缭乱。
于是,他们开始思考:能否设计一种脉冲神经网络的 meta 架构?这样一来,脉冲神经网络与人工神经网络在架构上的差距,立马就能被大大缩减。
后来,该团队将这一系列工作主要分为两步:
第一步:提出脉冲驱动自注意力算子。这也是脉冲神经网络领域的第三类算子,借此能让整个 Spike-driven Transformer 中只有稀疏加法。
第二步,探索脉冲神经网络的 meta 架构。通过此,可以缩小脉冲神经网络和传统人工神经网络在架构设计上的间隙。
完成上述步骤之后,他们为脉冲神经网络领域成功引入了新算子和新架构,让神经形态计算在拥有低能耗优势的同时,任务性能也能迈上一个台阶。
课题组相信预计在两年之内,继续沿着这个方向进展的话,脉冲神经网络的性能将完全可以和人工神经网络比肩,并且前者的能效优势将更加突出。
在当前比较主流的视觉任务、自然语言处理任务、生成式任务上,假如神经形态计算可以在算法层面上解决性能瓶颈,一定会启发基于新脉冲算子和新脉冲网络架构的神经形态芯片的设计。同时,对于低功耗人工智能的实现也具有重要意义。
不久之前,关于上述研究的相关论文以《脉冲驱动 Transformer V2:元脉冲神经网络架构启发下一代神经形态芯片设计》(Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips)为题收录于 2024 国际表征学习大会(ICLR 2024,International Conference on Learning Representations 2024)上。
中国科学院自动化所助理研究员姚满为论文第一作者,李国齐研究员为通讯作者。
(来源:ICLR2024)
一方面,本次成果可被用于边缘神经形态计算场景,例如采取“神经形态视觉+神经形态计算”的组合。
这里的神经形态视觉,指的是通过动态视觉传感器(Dynamic Vision Sensor,DVS)来感知视觉场景中的亮度变化,进而只输出异步稀疏事件流的类眼仿生感知范式。
对于神经形态计算来说,它天然具有事件驱动计算的特性,所以非常适合处理这种稀疏事件流。
最近,该团队还与一家类脑初创公司合作,将脉冲神经网络部署到异步感算一体芯片。
芯片处理器部分的静息功耗只有 0.42mW,在典型神经形态视觉任务场景中的功耗也在 10mW 以下。
这让该款芯片具有“永远在线(always-on)”的特性,在一些边缘低功耗计算场景中具有显著优势。
如果能将 Spike-driven Transformer 架构融入到异步类脑芯片中,它不仅能继续保持低功耗的特性。同时,随着模型表达能力的提升,还能被用于更多场景之中。
另一方面,本次成果为基于神经形态的超大规模网络设计提供了技术支撑。
目前,大多数基于人工神经网络的大模型,都是基于 Transformer 架构设计而来。而本次工作将脉冲驱动范式融入 Transformer 架构,带来了纯加法的 Transformer。
同时,本次所设计的算子与输入 token 个数以及特征维度都是线性的。因此模型规模越大,模型的能耗优势也就越明显。
众所周知,目前人工智能已经迎来大模型时代,大模型也有望成为未来人类社会的基础服务设施。
但是,随着用户量和使用频次的增长,AI 的高能耗问题将成为不容忽视的问题。
在这种背景之下,融合类脑时空动态的新一代线性脉冲神经网络架构的探索,就显得尤为重要。这也意味着本次成果可以为低功耗类脑脉冲大模型提供技术支撑。
(来源:ICLR2024)
预计神经形态计算领域将迎来大发展
一路走来,李国齐深感不易。他说:“无论对于行外人、还是对于行内人,脉冲神经网络领域其实一直都饱受质疑。甚至我们组的部分同学也不是很有信心,因为他们经常会看到网友对于这一方向的质疑。”
对此,他也表示理解。尽管脉冲神经网络拥有类脑、低功耗等优势,但是这些优势只能在系统层面有所体现。
如前所述,相比已经成熟的人工神经网络,脉冲神经网络在各个方面仍然存在一定差距,因此脉冲神经网络领域到底能够走向何方,一直不够明朗。
好在这几年脉冲神经网络领域已经迎来了长足发展,他对于这一领域的信心也越来越足。
其表示:“我个人对神经形态计算领域的发展呈乐观态度,预计最近几年神经形态计算领域将迎来大发展。”
尤其是随着大模型时代的到来,AI 想要成为人类社会运行的底层设施,就不能忽视巨大的能耗问题。
因此,他和团队非常希望本次成果能够推进脉冲神经网络走向实际应用,为设计下一代神经形态芯片带来启发。
总体来说,神经形态计算领域内还存着诸多需要攻克的难题,这需要整个领域共同的努力。
而基于本次成果,他们将继续围绕如下几个方面开展工作:
一是实现更大规模的脉冲神经网络模型。由于脉冲神经网络的时空动态复杂,因此比人工神经网络更难以训练,这就需要设计新的训练方法,来实现高效的训练。
二是把脉冲神经网络推向更多的任务类型。本次工作主要围绕计算机视觉任务展开,未来他们还想尝试将所设计的结构,用于更多的任务上,比如长时序任务等。
三是提出基于脉冲神经网络的类脑脉冲大模型架构。可以预见的是,这将是一项艰巨的任务,需要在训练速度、架构设计、模型规模、任务性能、长距依赖关系建模等方面,让目前的脉冲神经网络取得系统性的突破。
四是设计适配类脑脉冲大模型的硬件计算架构。眼下,围绕本次工作的硬件实现,该团队已经展开了一些探索。
如果在硬件上可以实现高效的脉冲驱动自注意力算子,再结合大规模脉冲神经网络的稀疏计算特性,必将实现更多的功能。
参考资料:
1. Yao, Man, Jiakui Hu, Zhaokun Zhou, Li Yuan, Yonghong Tian, Bo Xu, and Guoqi Li.“Spike-driven transformer.”Advances in Neural Information Processing Systems 36 (NeurIPS2023).
2. Yao, Man, JiaKui Hu, Tianxiang Hu, Yifan Xu, Zhaokun Zhou, Yonghong Tian, Bo Xu, and Guoqi Li.“Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips.”In The Twelfth International Conference on Learning Representations. (ICLR2024).
排版:朵克斯