下载掌阅APP，畅读海量书库
立即打开
Transformer架构无疑是大型语言模型取得显著成功的基石。从开源模型（如Mistral）到闭源（如ChatGPT），Transformer架构几乎无处不在，其影响力深远。然而，为了突破大型语言模型（LLM）的性能边界，科研人员一直在不懈地探索更为先进的架构。
在这些探索中，一种名为Mamba的新型架构引起了业界的广泛关注。Mamba架构不仅继承了Transformer架构的核心优势（如自注意力机制和强大的序列处理能力），而且在多个方面进行了创新和改进。
首先，Mamba架构通过引入更高效的注意力机制，显著降低了计算复杂度。在Transformer中，自注意力机制的计算成本随着序列长度的增加而迅速增长，这限制了模型在处理长文本时的性能。而Mamba则采用了更为精细的注意力计算方法，既能更准确地捕捉文本中的关键信息，又能保持较低的计算开销。
其次，Mamba架构在模型结构设计上进行了优化。它采用了更为灵活的解码器结构，使得模型在处理不同任务时能够更加灵活和高效。
更重要的是，Mamba架构在保持高性能的同时，还具备出色的可扩展性。这意味着随着模型规模的增大，Mamba的性能提升将更加显著。这对于构建更大规模、更强大的LLM具有重要意义。
总的来说，作为一种新兴的大型语言模型架构，Mamba架构展现了巨大的潜力和优势。随着研究的不断深入和技术的持续进步，我们有理由相信，Mamba将成为未来LLM发展的重要方向之一。
在讲解为什么需要使用Mamba模型之前，我们首先需要回顾一下前面介绍的注意力模型Transformer，并深入探讨它的一个显著缺点。
Transformer在处理自然语言处理（Natural Language Processing，NLP）任务时，将任何文本输入都视为由一系列Token（词元，俗称标记）组成的序列。这种方法在处理文本数据时确实有其优势，但也存在一个不容忽视的问题。Transformer处理流程如图2-1所示。
图2-1 Transformer的处理流程
Transformer在处理文本数据时，其核心理念是将任何文本输入视为由一系列Token（词元）组成的序列。这种处理方式赋予了Transformer强大的能力，即无论接收到何种输入，它都能够回顾序列中任一先前的Token，并据此推导出它们的表示形式。这一特性使得Transformer在理解长文本和捕捉文本内部依赖关系方面具有显著优势。Transformers的编码方法如图2-2所示。
图2-2 Transformers的编码方法
Transformer主要由两部分组成：编码器（Encoder）和解码器（Decoder），如图2-3所示。编码器负责将输入的文本序列转换为便于模型处理的内部表示，而解码器则依据这种内部表示生成文本。这两部分结构协同工作，使Transformer能够胜任机器翻译、文本摘要等多种NLP任务。
图2-3 编码器和解码器
另外，在训练过程中，单个解码器块的核心组件包括屏蔽自注意力机制和前馈神经网络。屏蔽自注意力机制使模型在训练时仅关注当前位置之前的Token，从而避免了未来信息的泄露。该机制通过创建权重矩阵，将每个Token与之前的Token进行比较（见图2-4），并根据它们之间的相关程度确定权重，从而帮助模型快速获取整个序列的上下文信息，进而提升文本生成质量。
图2-4 Token与之前的Token进行比较
然而，尽管Transformer在训练过程中表现出色，但在推理阶段却面临“推理的诅咒”问题。具体来说，即使我们已经根据先前的Token生成了一些文本，在生成下一个Token时仍需重新计算整个序列的注意力。由于这种计算量随着序列长度的增加而呈平方增长，导致在生成长文本时的计算成本极高。全序列的注意力计算如图2-5所示。
图2-5 全序列的注意力计算
例如，对于长度为 L 的序列，生成每个Token大约需要 L 2 的计算量。随着长度 L 的增大，这一计算成本迅速上升，使Transformer在处理长文本时变得力不从心。这不仅成为Transformer架构的主要瓶颈，也限制了其在需要生成长文本的应用场景中的使用。图2-6展示了使用注意力机制的训练与推理的对比。
图2-6 使用注意力机制的训练与推理比较
为了解决这一问题，研究人员开始探索新的架构和方法。Mamba架构正是在这种背景下应运而生的。Mamba架构通过引入新的技术和机制，旨在降低生成文本时的计算成本，同时保持甚至提升文本生成质量。
循环神经网络（Recurrent Neural Network，RNN）是一种专门设计用于处理序列数据的神经网络架构。其核心思想是，RNN在每个时间步（ t ）内依赖于两个输入：当前时间步 t 的输入和前一个时间步 t -1的隐藏状态。基于这两个输入，RNN会生成下一个隐藏状态，并据此预测输出，如图2-7所示。
图2-7 循环神经网络的输出
RNN的特点在于其内置的循环机制，这种机制使得网络能够将信息从当前时间步传递到下一个时间步。为了更直观地理解这一过程，可以将RNN在时间维度上展开，观察其随时间变化的动态行为，如图2-8所示。
图2-8 循环神经网络的展开
在生成输出时，RNN仅依赖当前的输入和上一个时间步的隐藏状态。这种机制使得RNN在推理时非常高效，因为它不需要重新计算所有先前的隐藏状态。相比之下，Transformer这样的架构在推理时需要计算整个序列的隐藏状态，导致计算成本大幅度增加。因此，RNN的推理速度是随序列长度线性增长的，理论上甚至可以处理无限长的上下文。
为了具体说明这一点，我们可以将RNN应用于一个实际的文本输入示例。在这个例子中，随着RNN读取文本中的每个单词，它的隐藏状态会不断更新，包含当前及之前所有单词的信息。
然而，由于RNN只考虑前一个时间步的隐藏状态，随着时间的推移，一些早期信息可能会被逐渐遗忘。例如，在处理包含Hello和Maarten这两个单词的文本时，当RNN生成Maarten这个名称时，它可能不再包含关于Hello这个单词的信息，如图2-9所示。
图2-9 推理过程只受前一个状态影响
此外，RNN的顺序性质也带来了另一个挑战：训练过程无法并行进行。这是因为RNN在每个时间步都需要等待前一个时间步的计算结果，才能继续计算下一个时间步。这种串行计算方式限制了RNN的训练速度，尤其在处理长序列时。
因此，尽管RNN在推理速度上具有优势，但其无法并行训练的问题却限制了其在实际应用中的性能。那么，是否存在一种既能像Transformer那样并行训练，又能保持RNN随序列长度线性扩展推理速度的架构呢？答案是肯定的，但这需要我们深入了解一种被称为状态空间模型（State Space Model，SSM）的新概念。训练与推理速度的对比如图2-10所示。
图2-10 训练与推理速度的比较
为了融合Transformer模型的并行训练优势与RNN在处理序列时聚焦于前一个Token的特点，同时实现对前文信息的选择性整合，Mamba团队创新性地采用了SSM（状态空间模型）这一新架构来处理信息的续写。这一架构不仅能够有效结合Transformer的高效并行计算能力，还能模拟RNN的逐步处理机制，从而在保持全局信息处理能力的同时，也能聚焦于序列中的局部依赖关系。使用状态空间模型的Mamba如图2-11所示。
图2-11 使用状态空间模型的Mamba
SSM的核心思想在于其独特的状态空间设计，使得模型能够在处理每个新的Token时，动态地更新和维护一个内部状态。这一状态不仅编码了之前所有Token的信息，还能根据当前的处理需求进行选择性的信息整合。通过这种方式，SSM能够在处理长序列时保持高效的计算性能，同时又能捕捉到序列中的重要信息。
在实际应用中，SSM首先通过Transformer的并行训练机制快速处理输入序列，提取全局特征。随后，在推理阶段，它借鉴了RNN的工作方式，每次只关注前一个Token，并根据空间状态进行信息的更新和传递。这种机制使得SSM在处理长文本时能够更有效地利用历史信息，从而提高了模型的续写能力和文本生成的连贯性。
可以看到，Mamba通过引入SSM，成功地将Transformer的高效并行性与RNN的序列依赖性结合起来，为自然语言处理和文本生成任务提供了一种全新的解决方案。