【论文阅读】MAMBA系列学习

Mamba

code：state-spaces/mamba: Mamba SSM architecture

paper：https://arxiv.org/abs/2312.00752

背景

1. 研究问题：如何在保持线性时间复杂度的同时，提升序列建模的性能，特别是在处理长序列和密集数据（如语言和基因组数据）时。
2. 研究难点：Transformer架构在长序列上的计算效率低，现有的结构化状态空间模型（SSMs）在处理离散数据时的局限性，以及如何在硬件资源有限的情况下实现高效的并行计算。
3. 相关工作：Transformer架构及其变体、线性注意力、门控卷积、循环模型和结构化状态空间模型（SSMs）。这些工作虽然在某些方面有所改进，但在处理长序列和密集数据时仍存在不足。

方法

这篇论文提出了一种新的选择性状态空间模型（Selective State Space Models, SSMs），用于解决序列建模中的效率和性能问题。具体来说，

1. 选择机制：首先，论文识别了现有SSMs的一个关键弱点：无法有效地根据输入进行数据选择。为了解决这个问题，论文提出了一种简单但强大的选择机制，通过将SSMs参数表示为输入的函数，使模型能够根据当前令牌选择性地传播或遗忘信息。

    

   

2. 硬件感知算法：其次，尽管这种变化阻止了高效卷积的使用，论文设计了一种硬件感知的并行算法，在递归模式下进行计算。该算法通过扫描而不是卷积来实现模型的计算，并且只在GPU内存层次结构中更有效的层次上实现扩展状态，从而避免了不同层次之间的IO访问。

3. 架构设计：最后，论文将选择性SSMs集成到一个简化的端到端神经网络架构中，该架构不包含注意力或MLP块，称为Mamba。Mamba架构结合了SSMs的设计和Transformers的MLP块，形成了一个简单且同质的架构。

   

实验

1. 合成任务：在选择性复制和诱导头任务上测试Mamba，这些任务需要模型具有内容感知能力。

    

   

2. 音频和基因组数据集：在YouTubeMix音频数据集上进行自回归预训练，并在SC09语音生成数据集上进行自动评估。

    

   

3. 语言建模：在Pile数据集上进行预训练，并在多个零样本评估任务上进行下游评估。

结果

1. 合成任务：在选择性复制任务中，Mamba模型轻松解决了任务，并且在组合长度超过1M令牌时表现出色。在诱导头任务中，Mamba模型能够完美地解决任务，并且在测试序列长度达到训练序列长度的4000倍时仍能保持高性能。

2. 音频和基因组数据集：在YouTubeMix音频预训练中，Mamba模型在长上下文下的性能优于现有的SaShiMi模型。在SC09语音生成任务中，Mamba模型在生成质量和生成速度上均优于现有的基线模型。

    

   

3. 语言建模：在Pile数据集上，Mamba模型在预训练困惑度和下游评估中均达到了Transformer模型的性能。Mamba-3B模型在预训练困惑度上比同大小的Transformer模型低4个点，并且在常识推理任务上超过了Pythia-7B模型。

结论

这篇论文提出了一种新的选择性状态空间模型（Mamba），通过引入选择机制和硬件感知算法，实现了线性时间复杂度的序列建模。Mamba在多个领域（如语言、音频和基因组）上达到了最先进的性能，并且在长序列数据处理上表现出色。Mamba的提出为构建不同领域的通用基础模型提供了一个强有力的候选者，特别是在需要长上下文的新兴模态（如基因组学、音频和视频）中。

创新与不足

1. 选择机制：提出了基于输入的选择机制，使模型能够根据当前令牌选择性传播或遗忘信息，从而克服了传统结构化状态空间模型（SSM）在处理离散模态时的弱点。
2. 硬件感知算法：设计了一种硬件感知的并行算法，通过扫描而不是卷积来递归计算模型，避免了在不同GPU内存层次结构之间进行IO访问，从而提高了计算效率。
3. 简化架构：将SSM架构与Transformer的MLP块结合，形成了一个简单的同质化架构（Mamba），无需注意力模块甚至MLP块，实现了线性时间复杂度的序列建模。
4. 快速推理：Mamba在推理速度上比Transformer快5倍，且序列长度线性扩展，实际数据上的性能提升可达百万序列长度。
5. 多模态性能：Mamba在语言、音频和基因组等多种模态上达到了最先进的性能，特别是在语言建模任务中，Mamba-3B模型在预训练和下游评估中均优于同大小的Transformer，甚至与两倍大小的Transformer相当。
6. 广泛的验证：在合成任务、音频和基因组建模、语言建模等多个领域进行了广泛的验证，展示了Mamba作为通用序列模型骨干的潜力。

不足

1. 模型规模限制：目前的实验主要集中在较小的模型规模上，尚未在大于7B参数的模型上进行充分评估，未来需要进一步研究Mamba在更大规模模型上的表现。
2. 长序列处理的挑战：尽管Mamba在处理长序列时表现出色，但在极长序列（如1M序列）上的表现仍需进一步优化，特别是在计算和内存消耗方面。
3. 硬件依赖性：当前的硬件感知算法主要针对NVIDIA的GPU进行了优化，未来需要考虑在其他硬件平台（如TPU）上的实现和优化。
4. 下游任务的适应性：虽然Mamba在多个模态上表现出色，但在特定领域的下游任务中可能还需要进一步的调整和优化，以充分发挥其潜力。

关键问题

问题1：Mamba模型的选择机制是如何设计的？它如何提高模型的效率和性能？

Mamba模型的选择机制通过将结构化状态空间模型（SSMs）参数表示为输入的函数，使模型能够根据当前令牌选择性地传播或遗忘信息。具体来说，Mamba模型的选择机制通过以下步骤实现：

1. 参数化SSM参数：将SSMs的参数（如A、B、C）表示为输入x的函数，即sB(x)、sC(x)和sA(x)。例如，sA(x)可以通过Broadcast操作将输入x投影到维度D。
2. 动态调整状态：通过引入一个可学习的参数Δ，控制模型在不同输入下的状态变化。Δ的值取决于输入x，从而实现选择性传播或遗忘信息。
3. 硬件感知算法：为了提高计算效率，Mamba模型采用了一种硬件感知的并行算法。该算法通过扫描而不是卷积来实现模型的计算，并且只在GPU内存层次结构中更有效的层次上实现扩展状态，从而避免了不同层次之间的IO访问。

这种选择机制使得Mamba模型在处理长序列和密集数据（如语言和基因组数据）时表现出更高的效率和性能。模型能够根据当前令牌选择性地传播或遗忘信息，从而在保持线性时间复杂度的同时，提升序列建模的性能。

问题2：Mamba模型在音频和基因组数据集上的表现如何？与现有的基线模型相比有哪些优势？

1. YouTubeMix音频预训练：在YouTubeMix音频数据集上，Mamba模型在长上下文下的性能优于现有的SaShiMi模型。具体来说，Mamba模型在训练序列长度达到60秒（960000个样本）时，性能显著提升，且在更长序列下继续保持优势。
2. SC09语音生成：在SC09语音生成数据集上，Mamba模型在生成质量和生成速度上均优于现有的基线模型。例如，Mamba-6.1M模型在生成质量上显著优于WaveNet和SampleRNN等基线模型，并且在生成速度上也表现出色。

与现有的基线模型相比，Mamba模型的优势主要体现在以下几个方面：

1. 长上下文处理能力：Mamba模型通过选择机制和硬件感知算法，能够在长序列数据处理上表现出色，尤其是在需要长上下文的新兴模态（如基因组学、音频和视频）中。
2. 高效的并行计算：Mamba模型采用了一种硬件感知的并行算法，通过扫描而不是卷积来实现模型的计算，从而在GPU内存层次结构中更有效地利用内存，提高了计算效率。
3. 综合性能提升：在音频和基因组数据集上，Mamba模型在预训练和下游评估中均达到了最先进的性能，显著提升了生成质量和生成速度。

问题3：Mamba模型在语言建模任务中的表现如何？与现有的Transformer模型相比有哪些优势？

1. 预训练困惑度：在Pile数据集上，Mamba模型在预训练困惑度上达到了Transformer模型的性能。具体来说，Mamba-3B模型在预训练困惑度上比同大小的Transformer模型低4个点。
2. 下游评估：在多个零样本评估任务上，Mamba模型的表现优于现有的基线模型。例如，在常识推理任务上，Mamba-3B模型超过了Pythia-7B模型，显示出更强的推理能力。
3. 生成吞吐量：Mamba模型的生成吞吐量比同大小的Transformer模型高5倍，表明其在实际应用中具有更高的生成效率。

与现有的Transformer模型相比，Mamba模型的优势主要体现在以下几个方面：

1. 线性时间复杂度：Mamba模型通过选择机制和硬件感知算法，实现了线性时间复杂度的序列建模，能够在长序列数据处理上保持高效。
2. 高效的并行计算：Mamba模型采用了一种硬件感知的并行算法，通过扫描而不是卷积来实现模型的计算，从而在GPU内存层次结构中更有效地利用内存，提高了计算效率。
3. 更强的生成能力：Mamba模型在生成质量和生成速度上均优于现有的基线模型，尤其是在大规模参数和长序列生成任务中表现出色。

Mamba-2

paper：[2405.21060] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

背景

1. 研究问题：这篇文章要解决的问题是如何将结构化状态空间模型（SSMs）与变体注意力机制联系起来，以便在保持语言建模性能的同时提高模型的效率。
2. 研究难点：该问题的研究难点包括：SSMs在训练和推理过程中需要线性时间复杂度，但在硬件上的实现效率较低；变体注意力机制虽然计算效率高，但难以在小到中等规模上匹配或超越Transformers的性能。
3. 相关工作：该问题的研究相关工作包括：Transformers在语言建模中的成功应用，SSMs如Mamba在小到中等规模上匹配或超越Transformers的性能，以及Linear Attention（LA）框架通过结构化矩阵对自回归注意力和线性RNN之间建立联系。

方法

这篇论文提出了结构化状态空间对偶（SSD）框架，用于解决SSMs与变体注意力机制之间的联系问题。具体来说，

1. 结构化状态空间模型（SSM）：首先，论文定义了SSM作为参数化的序列变换，表示为矩阵乘法形式：

y=SSM(A,B,C)(x)=Mx

其中，M 是一个结构化矩阵，A 是状态转移矩阵，B 是输入投影矩阵，C 是输出投影矩阵，x 是输入序列。
2. 半可分离矩阵：论文证明了SSM可以表示为一类称为半可分离矩阵的结构化矩阵。半可分离矩阵具有子二次参数和快速矩阵乘法算法。
3. 线性注意力（LA）：论文通过张量收缩提供了LA的简洁证明，并将其推广为结构化掩码注意力（SMA）。LA的核心思想是将softmax折叠到核特征图中，并通过矩阵乘法的结合律重写注意力计算。
4. 结构化状态空间对偶（SSD）：论文展示了SSM和SMA的对偶性，证明了任何具有快速递归形式的核注意力方法必须是SSM。SSD的对偶形式与注意力密切相关，定义为：

(L∘QK⊤)⋅V

其中，L 是一个输入依赖的掩码矩阵。

实验

1. 数据集：实验使用了Pile数据集，这是一个包含800GB多样化文本的大型数据集，适用于语言建模任务。
2. 模型规模：实验中使用了从125M到2.7B参数的模型，涵盖了多种模型规模。
3. 训练设置：所有模型使用AdamW优化器，梯度裁剪值为1.0，权重衰减为0.1，不使用dropout，学习率采用线性预热和余弦衰减。
4. 对比模型：实验中对比了Mamba-2、Mamba-1、Transformer++等模型，评估了它们在标准语言建模预训练和下游评估中的表现。

结果

1. 合成任务：在多查询联想记忆（MQAR）任务中，Mamba-2在所有设置下表现良好，显著优于Mamba-1，甚至在状态大小受控的情况下也优于Mamba-1。

    

   

2. 语言建模：在标准自回归语言建模预训练和下游评估中，Mamba-2在多个基准测试中匹配或超越了Mamba和开源Transformers。例如，Mamba-2在Pile数据集上训练300B令牌时，性能优于Mamba-2.8B、Pythia-2.8B和Pythia-6.9B。

    

   

3. 效率基准：SSD算法在训练和推理效率上比Mamba的选择性扫描实现快2-8倍，并且在序列长度达到2K时比FlashAttention-2快。

    

   

结论

这篇论文提出了一个基于结构化矩阵的理论框架，连接了SSMs和变体注意力机制。通过SSD框架，论文设计了一个新的架构Mamba-2，其核心层是对Mamba选择性SSM的改进，速度提高了2-8倍，同时在语言建模任务中保持了竞争力。论文的贡献包括：

1. 建立了SSMs与半可分离矩阵之间的理论联系。
2. 提出了结构化掩码注意力的广义抽象。
3. 设计了高效的SSD算法，并在多个基准测试中展示了其优越性。
4. 通过SSD框架，展示了如何利用Transformers的系统优化技术来改进SSMs。

创新与不足

1. 理论框架：论文提出了一个基于结构化矩阵的理论框架，揭示了结构化状态空间模型（SSMs）和注意力变体之间的紧密联系。
2. 高效算法：通过结构化矩阵的分块分解，设计了一种新的SSD算法，显著提高了SSMs的计算效率，特别是在序列长度为2K及以上时比FlashAttention-2快6倍。
3. 新架构：提出了Mamba-2架构，其核心层是对Mamba选择性SSM的改进，速度提升了2-8倍，同时在语言建模任务上仍与Transformers竞争。
4. 多视角连接：通过张量收缩和矩阵变换，展示了SSMs和注意力之间的多种视角连接，提供了新的理解和优化方法。
5. 系统优化：将SSMs与Transformers的系统优化技术相结合，提出了张量并行和序列并行的大规模训练方法，以及处理变长序列的有效方法。
6. 开源模型和代码：提供了模型代码和预训练检查点，促进了社区的贡献和进一步的研究。

不足

1. 局限性：尽管SSD算法在大多数情况下表现出色，但在某些特定场景下可能不如标准SSM或注意力模型高效。需要进一步研究如何改进这些特定情况下的性能。
2. 下一步工作：未来的研究方向包括进一步优化SSD算法，特别是针对更一般的对角SSM情况；探索更多的结构化矩阵表示和分解方法；以及将SSMs与其他模型和技术（如混合模型、上下文学习等）结合，以进一步提升性能。

关键问题

问题1：论文中提出的结构化状态空间对偶（SSD）框架是如何连接SSM和注意力机制的？

结构化状态空间对偶（SSD）框架通过将SSM和注意力机制都表示为结构化矩阵的矩阵乘法来连接它们。具体来说，SSM被表示为一个半分离矩阵的矩阵乘法形式：

y=SSM(A,B,C)(x)=Mx

其中，M 是一个半分离矩阵，A 是状态转移矩阵，B 是输入投影矩阵，C 是输出投影矩阵，x 是输入序列。注意力机制则被表示为：

Y=MX

其中，M 是一个半分离矩阵，X 是输入序列，Y 是输出序列。SSD的两个主要形式是线性形式和二次形式。线性形式对应于SSM的递归计算，而二次形式对应于注意力机制的计算。通过这种表示，SSD框架揭示了SSM和注意力机制之间的深层联系，并为优化这两种模型提供了新的思路。

问题2：论文中提出的SSD算法是如何提高SSM训练和推理效率的？

SSD算法通过块分解半分离矩阵来提高SSM的训练和推理效率。具体步骤如下：

1. 块分解：将半分离矩阵M分解为多个子矩阵，每个子矩阵可以在更小的计算复杂度内处理。
2. 线性计算：对于每个子矩阵，使用线性计算方法（如矩阵乘法）进行处理。
3. 二次计算：对于需要二次计算的子矩阵，使用专门的快速算法（如稀疏矩阵乘法）进行处理。
4. 并行化：利用现代硬件（如GPU）的并行计算能力，加速矩阵乘法和加法操作。

通过这种块分解方法，SSD算法在保持线性时间复杂度的同时，能够有效地利用硬件资源，显著提高SSM的训练和推理效率。实验结果表明，SSD算法在序列长度为2K及以上时比FlashAttention-2更快，并且在训练和推理过程中显著优于Mamba的融合扫描实现。

问题3：论文中提出的Mamba-2架构是如何利用SSD框架进行优化的？

Mamba-2架构通过以下方式利用SSD框架进行优化：

1. 并行参数投影：在Mamba-2中，状态转移矩阵A、输入投影矩阵B、输出投影矩阵C和输入序列X在块开始时并行生成，而不是作为SSM输入X的函数。这减少了参数数量，并使得更大规模的模型更容易进行张量并行。
2. 额外的归一化层：在块的最后添加一个额外的归一化层（如LayerNorm、GroupNorm或RMSNorm），以提高稳定性。
3. 多头结构：引入多头结构，将SSM的每个头独立处理，类似于多头注意力机制。Mamba-2采用了多值注意力（MVA）模式，其中B和C矩阵在所有输入通道上共享。
4. 块分解的张量并行：通过块分解的方法实现张量并行，使得更大规模的模型能够在多个GPU上进行并行计算。

这些优化措施使得Mamba-2在保持语言建模性能的同时，显著提高了训练和推理效率。实验结果表明，Mamba-2在多个任务和规模上均表现出色，验证了SSD框架的有效性和潜力。

Jamba

code：kyegomez/Jamba: PyTorch Implementation of Jamba: "Jamba: A Hybrid Transformer-Mamba Language Model"

paper：Jamba: A Hybrid Transformer-Mamba Language Model

背景

1. 研究问题：本文提出了一种新的大型语言模型Jamba，该模型基于Transformer和Mamba层的混合架构，并结合了专家混合（MoE）模块。Jamba旨在解决Transformer在高内存和计算要求方面的局限性，同时提高长上下文处理能力。
2. 研究难点：Transformer模型在处理长上下文时，键值（KV）缓存大小成为限制因素，且由于其缺乏单一总结状态，推理速度慢，吞吐量低。相比之下，RNN模型虽然能处理长上下文，但训练成本高且难以捕捉长距离关系。
3. 相关工作：现有的混合模型尝试结合注意力机制和状态空间模型（SSM），但大多在模型规模或性能上存在不足。例如，H3和StripedHyena在长上下文处理上表现不佳。

方法

Jamba模型用于解决大语言模型在处理长上下文时的效率和性能问题。具体来说，

1. 混合架构：Jamba模型结合了Transformer层和Mamba层，并在某些层中加入了MoE模块。Transformer层负责处理短序列和高计算任务，而Mamba层则更适用于长序列和低计算任务。

   

2. MoE模块：MoE模块允许在不增加计算需求的情况下增加模型容量。在Jamba中，MoE应用于某些MLP层，每个层有16个专家，每次选择前2个专家进行计算。

3. 灵活配置：Jamba架构提供了多种配置选项，包括层数、注意力层与Mamba层的比例、何时使用MoE以及每层的专家数量等。这些配置可以根据硬件和性能要求进行优化。

实验

1. 数据集：Jamba模型在一个包含网络文本、书籍和代码的自定义数据集上进行训练，数据集最后更新于2024年3月。数据处理流程包括质量过滤和去重。
2. 训练基础设施：模型在NVIDIA H100 GPU上进行训练，使用了高效的分布式训练框架，包括FSDP、张量并行性、序列并行性和专家并行性。
3. 模型配置：为了适应单个80GB GPU，Jamba模型的配置如下：
   • 4个Jamba块，每个块包含8层，注意力层与Mamba层的比例为1:7。
   • 每隔一层使用MoE模块，每层有16个专家，每次选择前2个专家。

结果

1. 学术基准测试：Jamba在多个标准学术基准测试中表现优异，与类似大小的领先模型（如Llama-2 70B和Mixtral）相比，性能相当但吞吐量更高。

    

   

2. 长上下文评估：Jamba在合成和自然主义长上下文评估中表现出色，支持长达256K的上下文长度，并且在大多数长上下文问答任务中优于Mixtral。

    

   

3. 吞吐量分析：在不同批处理和上下文长度设置下，Jamba的吞吐量比Mixtral高出3倍，特别是在长上下文情况下表现尤为显著。

    

   

结论

本文提出的Jamba模型通过结合Transformer和Mamba层，并加入MoE模块，实现了高性能和长上下文处理能力。Jamba模型在多个基准测试中表现优异，且具有更高的吞吐量，适用于单个80GB GPU。未来的研究可以进一步探索混合注意力-状态空间模型的有效性和应用。

优点与不足

1. 混合架构：Jamba提出了一种新颖的混合Transformer-Mamba架构，结合了Transformer和Mamba层的优势，提供了更高的吞吐量和更小的内存占用。
2. MoE模块：在部分层中引入MoE模块，增加了模型容量而不增加计算需求，使得模型能够在保持高效的同时扩展规模。
3. 长上下文支持：Jamba模型支持长达256K的上下文长度，这在生产级公开模型中是前所未有的。
4. 灵活性：通过调整Transformer和Mamba层的比例以及MoE的应用频率，Jamba提供了在性能、内存使用和吞吐量之间的灵活平衡。
5. 高性能：在多个标准语言模型基准测试和长上下文评估中，Jamba表现出色，与同等参数规模的Mixtral和Llama-2模型相当甚至更优。
6. 开源实现：Jamba的实现代码和权重在Apache 2.0许可证下公开，鼓励社区进一步研究和优化。

不足

1. 未对齐和微调：发布的Jamba模型是预训练的基础模型，没有经过对齐或指令微调，也没有 moderation 机制，不应用于生产环境或与终端用户互动，除非进行额外的适应。
2. 长上下文处理的优化：尽管Jamba在长上下文处理上表现出色，但作者提到，随着社区对纯Transformer模型的优化也在不断发展，Jamba在这些方面的吞吐量差距可能会进一步扩大。
3. 未来工作：未来的研究可以进一步探索大规模混合注意力-状态空间模型中的上下文学习能力的出现，作者计划发布较小规模训练运行的模型检查点，以方便此类研究。

关键问题

问题1：Jamba模型在处理长上下文时有哪些具体的优势？

1. 高效的内存使用：通过将注意力层与Mamba层结合，Jamba模型减少了键值（KV）缓存的大小。在256K上下文中，Jamba的KV缓存仅为4GB，而类似的纯Transformer模型则需要32GB。
2. 高吞吐量：Jamba模型在长上下文处理方面表现出色，特别是在128K上下文中，其吞吐量是Mixtral的3倍。
3. 强大的长上下文处理能力：Jamba模型支持长达256K的上下文长度，并且在多个长上下文评估任务中表现优异，例如在Needle-in-a-haystack任务中成功检索到放置在长上下文窗口中间的简单语句。

问题2：Jamba模型中的MoE模块是如何设计的？其在模型中起到了什么作用？

1. 设计：在Jamba模型中，MoE模块应用于某些多层感知器（MLP）层。每个MLP层有16个专家，每次选择前2个专家进行计算。这种设计允许在不增加计算需求的情况下增加模型容量。
2. 作用：MoE模块的主要作用是提高模型的容量，使其能够处理更多的参数和复杂任务，同时保持计算效率。通过在每个MLP层中使用MoE，Jamba模型能够在保持较低的计算复杂度的同时，增加模型的表达能力和性能。

问题3：Jamba模型在学术基准测试中的表现如何？与其他模型相比有何优势？

1. 表现：Jamba在多个标准学术基准测试中表现优异。例如，在HellaSwag任务中，Jamba的得分为87.1%，与Llama-2 70B的85.3%相当，但吞吐量提高了3倍。
2. 优势：Jamba模型在保持与类似大小模型（如Llama-2 70B和Mixtral）相当性能的同时，具有更高的吞吐量。此外，Jamba模型的内存占用更少，特别是在处理长上下文时，其KV缓存大小显著低于纯Transformer模型。

参考：Transformer和Mamba强强结合！最新混合架构全面开源，推理速度狂飙8倍_an empirical study of mamba-based language models-CSDN博客

（2024，Attention-Mamba，MoE 替换 MLP）Jamba：混合 Transformer-Mamba 语言模型 - 知乎

Mamba和Transformer成功合体！“两全其美”混合大模型Jamba诞生：单张GPU即可处理140k上下文 - 知乎 关于Mamba和Mamba-2经典论文的学习笔记总结，以及Mamba和Transformer的对比总结，欢迎交流_mamba2论文-CSDN博客