Mamba: Linear-Time Sequence Modeling with Selective State Spaces——Introduction 2

一、

架构解释

新的架构被命名为Mamba，通过结合以前结构化状态空间模型（SSMs）的设计和Transformer中的多层感知机（MLP）块，简化了以前的深度序列模型架构。这种设计旨在创建一个更简单、更同质的架构，同时结合选择性状态空间的优点。以下是详细解释：

架构的关键组成部分

1. 选择性状态空间：

   • 状态空间在这里指的是能够有效表示和处理序列的模型。
   • 选择性状态空间机制允许模型动态决定输入序列的哪些部分是重要的，应该保留或传播，以及哪些部分可以忽略或遗忘。

2. Transformer中的MLP块：

   • 在Transformer模型中，MLP块用于自注意力机制之后进一步处理信息。
   • MLP块通常由两个线性变换和中间的非线性激活函数（通常是ReLU）组成。

简化和同质的架构（Mamba）

Mamba架构结合了SSMs的优势和MLP块的简洁和高效。具体工作方式如下：

步骤1：输入嵌入

• 与其他序列模型类似，输入序列首先被嵌入到一个高维空间。这个嵌入将输入序列中的每个元素（例如，句子中的单词、时间序列中的数据点）转换为向量。

步骤2：选择性状态空间块

• 这个块负责选择性地传播信息。它使用根据当前输入生成的动态参数来决定哪些信息应该保留或遗忘。
• 选择性状态空间块通过更新隐藏状态来处理嵌入的输入序列，以选择性过滤和保留重要信息。

步骤3：MLP块

• 在选择性状态空间块之后，处理过的序列被输入到MLP块。
• MLP块包含：
  • 一个线性变换，将输入投影到一个高维空间。
  • 一个非线性激活函数（例如ReLU），引入非线性。
  • 另一个线性变换，将向量投影回原始或期望的维度。

步骤4：输出处理

• MLP块的输出可以被输入到后续的层中，或者直接用于最终任务，具体取决于模型设计。
• 如果使用多层架构，相同的架构（选择性状态空间块和MLP块）会重复。

详细例子

让我们通过一个文本处理的例子来说明Mamba架构的工作原理：

1. 输入句子：“The quick brown fox jumps over the lazy dog.”

2. 步骤1：输入嵌入：

   • 句子中的每个单词被嵌入到一个高维向量中。
   • 例如：“The” -> [0.1, 0.3, 0.5, …], “quick” -> [0.2, 0.6, -0.1, …]

3. 步骤2：选择性状态空间块：

   • 每个嵌入的单词被选择性状态空间块处理。
   • 这个块为每个单词生成动态参数，决定其重要性。
   • 例如：对于单词"quick"，动态参数可能指示其重要性较高，应该在隐藏状态中保留。

4. 步骤3：MLP块：

   • 选择性状态空间块的输出被输入到MLP块。
   • MLP块应用线性变换和非线性激活来进一步处理信息。
   • 例如："quick"的向量可能被转换以强调上下文相关的重要特征。

5. 步骤4：输出处理：

   • 处理后的序列可以用于最终任务，例如文本分类、语言建模等。
   • 如果架构有多层，MLP块的输出可以被输入到另一个选择性状态空间块和MLP块中。

Mamba架构的优势

• 简洁性：通过结合SSMs和MLP块，架构更加直接，易于实现。
• 效率：选择性状态空间机制允许高效处理长序列，保留重要信息。
• 同质性：设计在整个模型中是一致的，便于扩展和优化。

总结

Mamba架构利用SSMs和MLP块的优势，创建了一个简单、高效且强大的序列模型。通过结合选择性状态空间，模型能够动态调整输入数据的重要性，从而在各种任务中表现更好。

二、

选择性SSMs以及Mamba架构是全递归模型，具有使其适合作为通用基础模型在序列上运行的关键特性。以下是详细解释：

关键特性

1. 高质量（High Quality）

• 选择性带来的强大性能：
  • 选择性SSMs通过动态调整模型参数，根据输入数据的不同选择性地传播或遗忘信息。这种机制使模型在处理密集模态（如语言和基因组学）时表现出色。
  • 例如，在自然语言处理任务中，选择性SSMs能够高效地捕捉句子的语义和结构，从而提高文本分类、翻译和生成等任务的准确性和表现。

2. 快速训练和推理（Fast Training and Inference）

• 线性扩展的计算和内存需求：

  • 选择性SSMs在训练过程中，其计算和内存需求随序列长度线性扩展。这意味着即使处理长序列，训练时间和内存使用量也能保持在可控范围内。
  • 例如，如果一个模型处理的序列长度从1000增加到10000，训练时间和内存需求仅增加10倍，而不是平方增长。

• 推理时的常数时间展开：

  • 在推理过程中，选择性SSMs以自回归方式展开序列，这种展开只需常数时间，而无需缓存前面的元素。
  • 例如，在生成文本时，模型只需根据当前单词和前一个隐藏状态生成下一个单词，而不需要存储整个上下文，极大地提高了推理效率。

3. 长上下文（Long Context）

• 高质量和高效性带来的性能提升：
  • 选择性SSMs的高质量和高效性结合，使其在处理长序列数据时表现出色，能够在实际数据中处理长度高达100万的序列。
  • 例如，在基因组学研究中，选择性SSMs能够处理长达数百万碱基对的DNA序列，捕捉其中的长程依赖关系，从而在基因注释和变异检测等任务中取得更好的结果。

示例解释

语言处理示例

假设我们有一个长文本序列，例如一本小说，我们希望通过选择性SSMs模型对文本进行理解和生成。

1. 高质量性能：

   • 模型能够动态调整参数，根据每个单词的重要性选择性地保留或遗忘信息。
   • 例如，对于描述场景的句子，模型可以保留关键的描述性词汇，忽略冗余信息，提高对文本的理解。

2. 快速训练和推理：

   • 在训练过程中，模型能够高效处理长文本，其计算和内存需求线性扩展。
   • 在推理过程中，生成每个单词只需常数时间，无需存储和访问整个文本上下文。

3. 长上下文处理：

   • 模型能够处理长达数百万字符的文本，捕捉长程依赖关系。
   • 例如，在小说中，模型能够记住前面章节的重要情节，帮助在后续生成时保持情节一致性。

基因组学示例

假设我们要分析一个长达数百万碱基对的DNA序列，选择性SSMs模型可以应用于基因注释和变异检测。

1. 高质量性能：

   • 模型能够根据每个碱基对的特性动态调整参数，选择性地保留重要的序列特征。
   • 例如，在识别基因编码区时，模型能够保留关键的序列片段，过滤掉非编码区。

2. 快速训练和推理：

   • 在处理长DNA序列时，模型的计算和内存需求线性扩展，保持高效的训练速度。
   • 在推理过程中，每一步仅需常数时间，能够快速分析和注释长DNA序列。

3. 长上下文处理：

   • 模型能够处理长达数百万碱基对的序列，捕捉基因组中的长程依赖关系。
   • 例如，模型能够识别跨越多个基因的调控区域，提高基因组分析的准确性。

总结

选择性SSMs和Mamba架构通过结合高质量性能、快速训练和推理以及处理长上下文的能力，成为通用基础模型在序列上运行的理想选择。这些特性使其在语言处理、基因组学等多个领域表现出色，处理长序列数据时具有显著优势。