力压Transformer,详解Mamba和状态空间模型 精华

发布于 2024-7-23 11:03
浏览
0收藏

大型语言模型(LLMs)之所以能够在语言理解与生成上取得巨大成功,Transformer架构是其强大的支撑。从开源的Mistral,到OpenAI开发的闭源模型ChatGPT,都采用了这一架构。

然而,技术的探索从未止步。为进一步提升LLMs的性能,学界正在研发能够超越Transformer的新架构。其中,Mamba模型以其创新的状态空间模型(State Space Model)成为研究的焦点。

本文介绍Mamba模型及其在语言建模领域的应用,逐步解析状态空间模型的基本概念,并通过丰富的可视化内容,让读者直观地理解这一技术如何有望挑战现有的Transformer架构。

1 Transformer架构的挑战

我们先对Transformer架构做一个快速回顾,并指出其存在的一个主要缺陷。

Transformer架构将文本输入视为由一系列token构成的序列。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

其核心优势在于,无论面对何种输入,都能追溯到序列中的早期token,以此来推导出其深层的语义表示。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

1.1 Transformer的核心组件

Transformer由架构两部分组成:编码器和解码器。编码器负责解析文本,而解码器则负责生成文本。这种结构的结合,使之能够胜任从文本翻译到内容创作的多种任务。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

进一步地,我们可以仅利用解码器部分来创建生成式模型。这种基于Transformer的模型,即生成预训练Transformer(GPT),通过解码器来续写或补全输入的文本,展现出其在文本生成方面的强大能力。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

1.2 自注意力的高效训练

单个解码器块由两个主要部分组成,即“掩蔽自注意力机制(Masked Self-attention)”和“前馈神经网络(Feedforward Neural Network)”。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

自注意力是这些模型运行如此良好的主要原因。它提供了整个序列的未压缩视图,并加快了训练速度。

具体来说,自注意力机制的工作原理是通过创建一个矩阵,该矩阵对序列中的每个token与之前所有token进行比较,并通过计算它们之间的相关性来确定权重。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

在训练过程中,自注意力矩阵是一次性整体构建的,这表示不需要依次等待每个token的注意力计算完成,而是可以同时进行整个序列的注意力计算。例如,在处理“我”和“名字”的关联之前,无需先完成“名字”和“是”的关联计算。

这种设计实现了训练过程的并行化,极大地提升了训练速度,使Transformer架构在处理大规模数据集时更加高效。

1.3 训练与推理的矛盾

然而,Transformer架构也有其局限性。每当生成新的token时,必须对整个序列的注意力权重重新进行计算,哪怕此前已经生成了若干token。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

生成长度为L的序列需要大约L²次计算,随着序列的延长,计算成本会急剧上升。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

这种对序列全面重新计算的需求,是Transformer架构的一个主要瓶颈。

接下来,让我们看看传统的递归神经网络(RNN)是如何克服这一推理过程中的效率问题。

1.4 RNN的潜力

递归神经网络(RNN)是一种处理序列数据的网络结构。在序列的每个时间点,RNN接收两个输入:当前时间点t的输入数据和上一个时间点t-1的隐藏状态,以此来计算下一个隐藏状态并预测输出结果。

RNN具有循环机制,能够将历史信息传递至下一步,类似于将每一步的信息“串联”起来。这种机制可以通过可视化的方式“展开”,以便更清晰地理解其工作原理。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

在生成输出时,RNN仅依赖于前一步骤的隐藏状态和当前的输入数据,避免了像Transformer那样需要重新计算整个序列的历史隐藏状态。

正因如此,RNN在进行推理时速度较快,因为它的计算量与序列长度呈线性关系,理论上能够处理无限长的上下文。

举例来说,当我们将RNN应用于之前的输入文本。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

每个隐藏状态都是对之前所有状态的压缩汇总。

但这里存在一个问题:随着时间的推移,比如在生成名字"Maarten"时,最后一个隐藏状态可能已经丢失了对"Hello"的记忆,因为RNN在每一步只考虑了前一个状态的信息。

此外,RNN的这种顺序依赖性也导致了另一个问题:它的训练过程无法并行化,必须按顺序逐步进行。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

与Transformer相比,RNN在推理速度上具有优势,但在训练并行化方面却存在不足。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

这就引出了一个问题:能否找到一种结合了Transformer训练并行化优势和RNN线性推理优势的架构呢?

答案是肯定的,Mamba模型就是。在深入了解Mamba架构之前,先来了解状态空间模型的世界。

2 状态空间模型(SSM)

状态空间模型(SSM),像Transformer和RNN一样,处理信息序列,如文本和信号。在这一部分中,我们将了解SSM的基础知识以及其与文本数据的关系。

2.1 什么是状态空间

状态空间模型是一种通过数学方法全面描述系统状态的方式,包含了描述系统所必需的全部最小变量。

简单来说,就像我们在迷宫中寻找路径,状态空间就是那张展示所有可能位置(即状态)的地图。在这张地图上,每个点都代表一个独特的位置,并且携带了如距离出口远近等具体信息。

进一步简化这个概念,可以将“状态空间表示”理解为这张地图的提炼,它不仅告诉我们当前所在的位置(即当前状态),还展示了可能的目的地(未来状态),以及如何通过特定的行动(比如右转或左转)达到下一个状态。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

虽然状态空间模型依赖方程和矩阵来捕捉系统的行为,但其核心目标是追踪系统的位置、可能的移动方向及其变化路径。

在这个模型中,用以描述状态的变量,如示例中的X和Y坐标或者到出口的距离,统称为“状态向量”。

力压Transformer,详解Mamba和状态空间模型-AI.x社区

这听起来有点熟悉,因为在语言模型中,类似的嵌入或向量经常用来描述输入序列的“状态”。例如,你当前的位置状态就可以通过一个向量来表示:

力压Transformer,详解Mamba和状态空间模型-AI.x社区

在神经网络的语境下,系统的“状态”通常指的是其隐藏状态,这在生成新token的过程中,尤其是在大型语言模型的背景下,扮演着至关重要的角色。

本文转载自 AI科技论谈​,作者: AI科技论谈

收藏
回复
举报
回复
相关推荐