类脑计算开启大模型计算新范式？——挑战获得诺贝尔奖的ANN

wisemodel开源社区

2024-10-24 21:28始智AI wisemodel官方账号

BrainTransformers是首个使用脉冲神经网络实现SOTA性能的大型语言模型，它将类脑计算中的脉冲神经网络（SNN）与大语言模型相结合，开创了大模型研究的新方向。

目前BrainTransformers已经开源，并发布到始智AI wisemodel开源社区，欢迎大家前往使用。

模型地址：

人工神经网络（ANN）是当前深度学习的主流模型，2024年的诺贝尔物理学奖就颁发给了ANN理论的奠基人。随着模型规模的不断扩大，计算资源消耗、能源效率等问题也日益凸显。在这个背景下，一种新的计算范式——类脑计算，正在悄然兴起，有望为AI的未来发展开辟新的道路。

从人工神经网络到脉冲神经网络：追寻更像大脑的AI

要理解BrainTransformers的创新之处，我们首先需要了解传统人工神经网络（ANN）和脉冲神经网络（SNN）的区别。通过模拟人脑神经元的连接方式来处理信息，在图像识别、自然语言处理等领域取得了巨大成功。然而，ANN的工作方式是连续的、同步的，每个神经元在每次计算中都会输出一个实数值。这种方式虽然在数学上便于处理，但与真实大脑神经元的工作方式存在较大差异。

想象一下，如果我们的大脑像ANN那样工作，那么我们的每个脑细胞都需要在每一秒内不停地进行计算和输出，这显然是非常耗能的。事实上，我们的大脑要高效得多。

脉冲神经网络（SNN）更接近生物神经系统的工作原理。在SNN中，神经元只在接收到足够强度的输入时才会发放脉冲，这种事件驱动的计算方式更符合大脑的工作机制。就像我们的大脑细胞，只有在必要时才会产生电信号，大部分时间都处于"待机"状态，这样就大大降低了能量消耗。

生物神经元产生动作电位的过程：

· 静息状态：
在静息状态下，神经元内部相对于外部环境呈现负电位（约-70mV），这称为静息膜电位。这种电位差主要由钾离子（K+）的浓度梯度维持，因为钾离子通道在静息状态下略微开放。

· 去极化：
当神经元接收到刺激时（比如神经递质结合到受体上），钠离子通道开始打开。钠离子（Na+）快速流入细胞内，导致膜电位变得更加正向，这个过程称为去极化。

· 阈值和启动：
如果去极化达到某个阈值（约-55mV），就会触发更多的电压门控钠离子通道打开，导致大量钠离子涌入。这个过程是自我强化的，膜电位迅速上升到约+40mV。

· 动作电位：
这个快速的电位变化就是动作电位，也就是神经脉冲。它沿着轴突传播，最终达到突触末梢。

· 复极化：
动作电位达到峰值后，钠离子通道关闭，而钾离子通道开始大量开放。钾离子（K+）快速外流，使膜电位迅速回落，甚至短暂地比静息电位更负（过极化）。

· 不应期：
在动作电位之后，神经元会经历一个短暂的不应期。在此期间，无论刺激多强，神经元都不能产生新的动作电位。这是因为钠离子通道需要时间从失活状态恢复。

· 恢复静息状态：最后，钠-钾泵（Na+-K+ ATPase）开始工作，将钠离子泵出细胞，将钾离子泵入细胞，逐渐恢复离子的浓度梯度和静息膜电位。

脉冲神经网络(SNN)是一种模仿生物神经系统工作方式的人工神经网络。与传统的人工神经网络(ANN)不同,SNN的信息传递和处理是通过离散的脉冲(或称为"尖峰")来完成的,这更接近真实神经元的工作方式。让我们深入了解SNN的核心工作原理:

神经元模型

在SNN中,最常用的神经元模型是"积分发放"(Integrate-and-Fire, IF)模型或其变体。这种模型的工作方式如下:

· 膜电位累积: 神经元接收来自其他神经元的输入脉冲,这些输入会改变神经元的内部状态,即膜电位。

· 阈值比较: 当膜电位达到一个预定的阈值时,神经元就会"发放"一个脉冲。

· 重置: 发放后,神经元的膜电位会重置到一个基准值,通常是静息电位。

· 不应期: 发放后,神经元会进入一个短暂的不应期,在此期间它不会对输入做出响应。

通过这种对比，我们可以看到SNN模型如何简化和抽象了生物神经元的复杂过程，同时保留了其关键的功能特性。这种简化使得我们能够在计算机中高效地模拟大规模神经网络，同时仍然保持了与生物神经系统的相似性。这也是为什么SNN被认为比传统的人工神经网络更接近生物脑的工作方式，有潜力在某些任务上展现出独特的优势。

信息编码

在SNN中,信息是通过脉冲的时间和频率来编码的。主要有两种编码方式:

· 速率编码: 信息由一段时间内脉冲的平均频率来表示。这类似于生物神经元中的发放率编码。

· 时间编码: 信息由精确的脉冲时间来表示。这种编码方式可以传递更精细的时间信息。

突触连接

神经元之间通过突触连接。在SNN中,突触的强度(权重)决定了预突触脉冲对后突触神经元的影响程度。突触可以是兴奋性的(正权重)或抑制性的(负权重)。

学习机制

SNN的学习通常基于生物学启发的规则,最著名的是"尖峰时间依赖可塑性"(Spike-Timing-Dependent Plasticity, STDP)。STDP规则如下:

· 如果预突触神经元在后突触神经元之前不久发放,则增强这个突触的权重。

· 如果预突触神经元在后突触神经元之后不久发放,则减弱这个突触的权重。

这种学习机制使得网络能够对输入模式的时间特征特别敏感。

网络动态

SNN的整体行为是由所有神经元的集体动态决定的。信息在网络中的传播是异步的,只有当神经元发放时才会影响其他神经元。这种事件驱动的特性使得SNN在处理时间序列数据时特别有优势。

计算效率

由于SNN的事件驱动特性,只有在有脉冲发生时才需要更新神经元状态。这意味着在硬件实现时,SNN可能比传统ANN更加节能,特别是在处理稀疏数据时。

SNN相比ANN具有以下主要优势：

1. 生物学合理性：SNN的工作方式更接近真实神经元，有助于我们更好地理解和模拟人脑功能。

2. 能效优势：由于SNN采用事件驱动的计算方式，只有在必要时才进行信息处理，理论上可以大大降低能耗。

3. 时序信息处理能力：SNN天然适合处理时间序列数据，这在语音识别、动作识别等领域具有潜在优势。

SNN的困局

那么，这么好用的算法理论为什么没有被广泛使用，反而让看起来既不仿生又能耗极大的ANN成为主流甚至获得了跨领域的诺贝尔奖呢？

首先，SNN的不连续性和非差分性是一个主要障碍。神经元通过离散的脉冲进行通信，这导致网络行为呈现不连续性。脉冲函数（通常是阶跃函数）在数学上是不可微的，这使得无法直接应用基于梯度的优化方法，我们需要寻找近似函数进行替代或者使用和反向传播完全不同的训练方式。与之相比，ANN中使用的激活函数如ReLU或sigmoid是连续可微的，能够直接应用反向传播等成熟的优化技术。

其次，SNN的时间依赖性增加了训练的复杂度。信息被编码在脉冲的精确时间中，这意味着我们需要考虑神经元在不同时间点的状态。这种时序特性使得传统的基于批处理的训练方法难以直接应用。而ANN则不需要处理这种复杂的时间动态，可以更容易地并行化训练过程。

另一个棘手的问题是信用分配。在SNN中，确定哪个神经元的脉冲对最终输出贡献最大是非常困难的。这使得难以有效地更新网络权重以改善性能。ANN中的反向传播算法可以相对直接地解决这个问题，但在SNN中，我们需要更复杂的机制来追踪和分配贡献。

SNN中神经元活动的稀疏性也是一个挑战。许多神经元在给定时间内可能不会发放脉冲，这种稀疏性可能导致梯度消失问题，使得网络难以学习。相比之下，ANN中的激活通常是密集的，即使在使用ReLU等会导致稀疏性的激活函数时，也不会达到SNN那样的程度。

面对这些挑战，研究人员提出了几种SNN训练方法，如SpikeProp、近似梯度方法、ANN到SNN的转换方法、强化学习方法以及基于STDP的学习。然而，这些方法都有其局限性。SpikeProp仅适用于单次脉冲编码，难以处理更复杂的时间编码模式。近似梯度方法可能导致精度损失，且不能完全捕捉SNN离散特性。转换方法在将ANN转换为SNN时可能导致性能下降，且无法充分利用SNN独特特性。强化学习方法的训练过程可能很慢，难以扩展到大型网络。基于STDP的学习则难以直接优化特定的目标函数。

BrainTransformers：SNN遇上大语言模型

面对这些挑战，来自LumenScopeAI的研究团队提出了BrainTransformers项目。该项目基于BrainTransformers架构，实现了一个使用脉冲神经网络的大语言模型——BrainGPTForCausalLM。

这个创新性的项目将SNN应用于Transformer架构，这在大语言模型领域是一次全新的尝试。研究团队通过三个训练阶段，包括SNN特定的神经元突触可塑性训练，成功构建了这个基于SNN的语言模型，并取得了同尺寸模型的SOTA性能，目前已开源3B小模型。

SNN兼容的Transformer组件

研究团队重新设计了Transformer的关键组件，使其能够在SNN框架下工作。这包括：

1. SNNMatmul：一种基于脉冲的矩阵乘法操作。

2. SNNSoftmax：使用累积脉冲和归一化来近似Softmax函数。

3. SNNSiLU：SNN版本的SiLU激活函数。

Excitatory-Inhibitory Integrate-and-Fire神经元模型

设计了一种新的神经元模型，称为EI-IF（Excitatory-Inhibitory Integrate-and-Fire）。这种神经元可以产生正负脉冲，更好地模拟生物神经元的行为。其动态可以表示为：

其中V(t)是膜电位，I(t)是输入电流，θ(t)是发放阈值，S(t)是神经元输出。Synapsis模块：

这个模块模拟了生物神经系统中的突触连接。它将预突触神经元、线性层和后突触神经元结合，实现了SNN中的高效线性变换。其数学表达式为：

这种设计允许模型在保持SNN特性的同时，执行类似于传统神经网络的线性变换操作。

基于SNN的数学函数近似

为了在SNN中实现复杂的数学运算，研究团队设计了一系列创新的基于脉冲的函数近似器。这些近似器不仅让SNN能够执行复杂的数学运算，还保持了SNN的核心特性。让我们深入了解这个突破性的设计：

CustomNeuron：函数近似的基础

研究团队首先设计了一种称为CustomNeuron的特殊神经元。这种神经元可以通过调整其参数来近似特定的数学函数。CustomNeuron的工作原理如下：

对于输入x，CustomNeuron的动态可以描述为：

其中V(t)是膜电位，S(t)是输出脉冲，a是脉冲幅度，θ是发放阈值。通过调整a、θ和时间步长T，CustomNeuron的累积输出可以近似特定函数。这就像是设计了一个可编程的神经元，我们可以通过调整其"程序"（即参数）来让它模仿不同的数学函数。

SquareApproximator：平方函数的SNN实现

基于CustomNeuron，研究团队设计了SquareApproximator来近似平方函数。这个近似器的工作原理是：

1) 将输入区间[0, b]分成N个子区间。

2) 每个子区间由一个CustomNeuron负责，该神经元被训练来近似该区间内的平方函数。

3) 对于输入x，SquareApproximator的输出是：

其中Si(x)是第i个CustomNeuron的输出。这种方法允许我们用一系列简单的、基于脉冲的操作来近似复杂的平方函数。想象一下，我们用一群特殊训练的神经元，每个只对一定范围内的脉冲输入产生相应，拟合一小段曲线，共同"拼凑"出一个完整的平方函数。

SqrtApproximator：平方根函数的SNN实现

平方根函数的近似更具挑战性，因为它在接近零时变化很快。研究团队采用了类似但更复杂的方法：

使用对数尺度划分输入区间[x_start, x_end]:

每个子区间仍由一个CustomNeuron负责。

SqrtApproximator的输出为:

这种对数尺度的划分允许在接近零的区域有更密集的采样点，从而更准确地近似平方根函数的快速变化部分。

SNNRMSNorm：基于SNN的RMS归一化

有了这些基本的函数近似器，研究团队就能够实现更复杂的操作，如RMS（Root Mean Square）归一化。SNNRMSNorm的实现过程如下：

这个过程展示了如何将基本的函数近似器组合起来，实现更复杂的数学操作，同时保持SNN的特性。

受神经元突触可塑性启发的三阶段训练方法

BrainTransformers采用了一种独特的三阶段训练方法，巧妙地结合了ANN的训练效率和SNN的生物合理性，其中具体的等价转化数学证明与突触可塑性训练函数可以参考LumenScopeAI提供的Paper（链接置于文末）。

ANN训练阶段：首先，团队训练了一个量化的ANN模型。这个模型的行为近似于SNN，但可以使用标准的反向传播算法进行高效训练。

ANN到SNN转换阶段：接下来，将训练好的ANN模型转换为等效的SNN。这个过程涉及到将量化的ANN组件替换为它们的SNN对应物，如将量化器替换为EI-IF神经元。

STDP启发的SNN训练阶段：最后，使用受生物启发的学习规则对转换后的SNN进行微调。这个阶段融合了局部的STDP（Spike-Timing-Dependent Plasticity）学习规则和全局的任务性能反馈。

SOTA性能！

我们的开源版本BrainTransformers-3B尽管只有30亿参数，相比其他大型模型如Llama3.1-70B（700亿参数）要小得多，但它在多项基准测试中展现出了令人惊喜的性能。以下是BrainTransformers-3B在一些关键任务上的表现，详细指详见本项目GitHub仓库或Huggingface仓库（链接置于文末）：

这些成绩虽然还不及最顶尖的大型模型，但考虑到BrainTransformers-3B的规模，其表现已经相当出色。特别是在GSM8K数学推理任务上的76.3%的成绩，展示了模型强大的逻辑推理能力。值得注意的是，在某些任务上，如Truthfulqa，BrainTransformers-3B甚至超过了一些更大的模型，这可能暗示了SNN在处理某些类型的任务时可能具有独特优势。

未来展望：开启AI新纪元？

BrainTransformers项目为我们展示了类脑计算在AI领域的巨大潜力。通过结合生物启发的计算模型与先进的深度学习技术，我们有望实现更高效、更智能的AI系统。

研究团队表示，他们正在进行多项优化工作，包括：

1. 开发专门针对SNN的GPU算子，充分利用现有硬件架构的并行计算能力。

2. 与半导体领域的合作伙伴共同设计开发专门适配其算法的神经形态芯片，旨在从硬件层面实现SNN的高效计算。

这些努力有望进一步提升BrainTransformers的性能和效率，为AI的未来发展开辟新的道路。

BrainTransformers项目向我们展示了类脑计算在AI领域的巨大潜力。它不仅推动了技术的进步，更为我们思考AI的未来发展方向提供了新的视角。这种结合生物启发计算和深度学习的方法，可能会为我们带来更高效、更智能、更像人脑的AI系统。

虽然目前BrainTransformers在某些任务上的表现还不及最顶尖的大型模型，但考虑到它的创新性和潜力，我们有理由相信，随着进一步的研究和优化，类脑计算可能会在AI领域掀起一场革命。

无论您是AI研究者、开发者，还是对这一领域感兴趣的普通读者，BrainTransformers项目都值得您的关注。因为它可能代表了AI的一个重要发展方向，也可能是我们通往更智能、更高效的AI系统的一个重要里程碑。同时LumenScopeAI团队正在寻找有意向共建BrainTransformers开原生态的朋友！

让我们一起见证AI向着更像人脑、更高效、更智能的方向迈进！

编辑：成蕴年

查看原图 258K