由仅解码器的转换器模型构建的大型语言模型 (LLM) 的开发在自然语言处理 (NLP) 领域的转型以及推进各种深度学习应用程序方面发挥了至关重要的作用，包括 强化学习、时间序列分析、图像处理等等。然而，尽管具有可扩展性和强大的性能，由仅解码器变压器模型构建的法学硕士仍然面临着重大缺陷。虽然富有表现力，但 Transformer 衍生的 LLM 中的注意力机制在推理和训练过程中都需要大量计算资源，因此需要大量内存来存储序列长度和二次 FLOP。这种高计算要求限制了变压器模型的上下文长度，使得自回归生成任务的规模成比例地昂贵，并且阻碍了从连续数据流中学习和真正无限的序列处理的能力。

在最近一个时期， 状态空间模型 （SSM）已经展示了卓越的功能和性能，在大规模建模基准中与 Transformer 架构模型竞争，同时实现了作为序列长度和线性时间函数的内存复杂性。此外，最近发布的状态空间模型Mamba在一系列语言建模和长序列处理任务中表现出了出色的性能。同时，专家混合 (MoE) 模型也表现出了令人印象深刻的性能，同时显着降低了推理的延迟和计算成本，尽管是以更大的内存占用为代价。本文将在 Mamba 和 MoE 模型的基础上讨论 BlackMamba，这是一种新颖的架构，它将 Mamba 状态空间模型与 MoE 模型相结合，以利用这两个框架提供的优势。 BlackMamba 上的实验证明了它在训练 FLOP 和推理方面都优于现有 Mamba 框架和 Transformer 基线的能力。 BlackMamba 框架的卓越性能表明，它可以有效地结合 Mamba 和 MoE 框架的功能，通过 MoE 和 Mamba 的线性复杂度生成提供快速且经济高效的推理。

本文旨在深入介绍 BlackMamba 框架。我们探讨了该框架的机制、方法和架构，以及它与最先进的图像和视频生成框架的比较。让我们开始吧。

BlackMamba：状态空间模型 MoE 简介

大型语言模型（LLM）的进展，特别是基于仅解码器变压器架构的模型，显着影响了 自然语言处理（NLP） 领域并扩展到各种深度学习应用，包括强化学习、时间序列分析、图像处理等。然而，尽管它们具有可扩展性和强大的性能，这些仅解码器的基于变压器的 LLM 遇到了显着的挑战。注意力机制是基于 Transformer 的一个关键特征 法学硕士s，需要大量的计算资源来进行推理和训练。这涉及到对内存的需求随着序列长度和计算操作 (FLOP) 的二次增加而增长。如此密集的计算需求限制了模型的上下文长度，随着模型的扩展提高了自回归生成任务的成本，并阻碍了模型有效地从连续数据流中学习或处理无限长度序列的能力。 

在过去的几年里，为了克服这些限制，人们做出了巨大的努力，并且注意力已经转向设计规范的密集注意力变压器模型的架构替代方案，其中 SSM 和 MoE 模型是最有前途的候选架构。相对于 Transformer 架构模型，采用状态空间模型所获得的主要好处是 SSM 提供的输入序列长度的线性计算复杂性，而不是 Transformer 提供的二次复杂性。理论上，对于给定的 FLOPS 或每秒浮点运算预算，相对于输入序列长度的线性计算复杂性使状态空间模型能够处理比 Transformer 架构模型更大的序列，并且在没有 KV 缓存的情况下在计算中呈现自回归生成常数。最近开发的状态空间模型（包括 Mamba、RetNet 和其他一些模型）已经展示了高效的长序列推理和训练，以及具有类似缩放属性的 Transformer 的有竞争力的语言建模任务性能。另一方面，混合专家模型架构作为密集变压器的替代方案越来越受欢迎，因为它有助于显着减少推理和训练 FLOP，这对于实现与密集模型相当的质量至关重要。 MoE（专家混合）模型通过在单次前向传递期间仅激活总参数的稀疏选择来运行。他们利用路由功能来确定根据给定的上下文调用哪些“专家”采取行动。这种方法在推理的计算成本和参数总数之间建立了分离，从而允许在固定的推理预算内增强性能，尽管参数数量增加并且内存需求更大。

与传统变压器相比，这种架构上的进步提供了显着的优势，并代表了进一步发展的令人兴奋的方向。我们认为，将这些增强功能集成到组合的 Mamba-MoE 模型中可以显着加速语言建模能力和效率，超越标准 Transformer 模型。与传统的密集变压器模型相比，Mamba-MoE 架构的预期优势包括：

曼巴舞： 在训练和推理阶段实现相对于输入序列长度的线性计算复杂性。它使自回归生成能够在恒定的时间范围内和恒定的内存使用情况下发生。

教育部： 提供与较小的密集基线模型相当的推理速度和训练计算效率，同时保持与具有与密集版本相同参数数量的模型相媲美的模型质量水平。

话虽如此，有必要指出 Transformer 架构模型仍然是最先进的，并且在语言建模任务和序列处理任务上表现出了一致且卓越的性能。 Transformer 架构的核心是采用自注意力机制，对序列中不同标记的嵌入之间的点积相似性执行二次全面比较，并执行到输出向量的线性映射。 Transformer 模型由堆叠在 MLP 或多层感知器块之间的自注意力块组成，多层感知器块进一步由具有给定激活函数的两层 MLP 组成。 

BlackMamba：架构和方法论

状态空间模型

状态空间模型属于相对于输入序列的长度具有线性复杂度的序列模型组。状态空间模型的架构更多地与循环神经网络和卷积神经网络而不是基于注意力的架构保持一致，并且受到连续动态系统的启发，该系统通过隐式潜在空间映射一维函数。线性动力系统使用关联扫描或卷积扫描使并行计算变得高效。在实际场景中，状态空间模型的循环性质是其在 GPU 等高度并行人工智能硬件上仍需采用的原因。然而，像 RWKV 和 曼巴 使用并行扫描内核将循环操作有效地映射到 GPU，从而促进新颖架构的训练，其效率与 Transformer 模型所达到的效率相当。 

Transformer 中与序列长度相关的固有二次复杂度是一个众所周知的限制，它阻碍了对很长上下文的推理和理解。最近的创新引入了扩展上下文长度的想法，使变压器能够在推理过程中应用于更长的上下文之前在可行的规模上进行训练。尽管取得了这些进步，推理过程仍然需要大量的计算资源和内存，尤其是维护键值 (KV) 缓存，使其成为资源密集型工作。最近的研究工作重点是通过结合输入相关的门控机制来增强状态空间模型的表达能力，类似于注意机制中的查询、键、值（QKV）矩阵。 

这些努力的目的是保持状态空间递归固有的线性进展，从而允许通过卷积或选择性扫描过程高效执行。这种方法显着缩小了实际应用中与变压器的性能差距。在这些进步中，Mamba 作为状态空间模型脱颖而出，它反映了先前研究的目标，在高达 2.8 亿个参数的规模上显示出与 Transformer 相当的令人印象深刻的性能水平。它通过将依赖于输入的门控应用于状态空间模型（SSM）递归的输入来实现这一点，同时通过使用定制的选择性扫描内核来确保高效计算。

专家模型的混合

专家混合 (MoE) 模型通过在前向传递期间选择性地激活参数来实现推理成本和总参数计数之间的分离。这些模型没有使用所有参数，而是将令牌定向给特定的多层感知器 (MLP) 专家。理想情况下，每个专家都经过定制，可以处理特定类型的输入，并通过路由机制（本质上是一个紧凑的神经网络）确定每个令牌最合适的专家。这种方法的目的是在更密集的配置中保留具有相同数量参数的模型的综合表达能力，但大大减少了计算需求。通常，路由器是从令牌到专家索引的线性层的映射，每个专家只是一个标准转换器多层感知器。然而，开发人员尚未找出路由器的最佳训练方法，因为专家分配问题是不可微分的，并且专家模型的混合通常会在不同专家之间的负载平衡和训练稳定性方面遇到困难，以提高硬件效率。 

建筑

BlackMamba 的核心采用了标准 Transformer 模型，该模型由交错的 MLP 块和沿残差流按顺序添加的注意力块组成。现在，大多数专家混合模型只是用路由专家层替换多层感知器块。另一方面，BlackMamba 框架不仅用路由专家层替换了 Transformer 中的多层感知器块，而且还用 Mamba 状态空间模型层替换了注意力层。 BlackMamba框架的架构如下图所示。 

训练和数据集

BlackMamba 模型在自定义数据集上对超过 300 亿个标记进行训练，并使用专家多层感知器的 SwiGLU 激活函数。该框架由 8 名专家进行训练，开发人员发现这个数字是模型的内存占用和推理成本之间的正确平衡和权衡。用于训练 BlackMamba 框架的自定义数据集由现有的开源数据集组成，包括 Starcoder、SlimPajama、Pile 等。下表展示了用于训练 BlackMamba 框架的每个数据集的权重。总的来说，数据集中有 1.8 万亿个代币。 

黑曼巴：结果

为了确保 Mamba 和 BlackMamba 之间的公平比较，开发人员在相同的训练数据上使用相同的训练参数来训练这两个模型。对于相同的前向传递模型大小，BlackMamba 框架能够在推理时间以及每秒训练浮点运算方面优于 Mamba 和 Transformer 模型。下图演示了从初始单标记提示以自回归方式生成给定长度的序列所需的时间，作为序列长度的函数。 

此外，BlackMamba 框架中结合了 Expert 模型和 Mamba 模型混合的延迟优势，与 Transformer 模型、纯 Mamba 模型和 MoE 模型相比，推理时间显着加快。此外，BlackMamba 框架的推理优势与序列长度成正比，使得 BlackMamba 在长序列生成方面极其有效。接下来，下图显示了分配给分别具有 340 亿和 640 亿参数的 BlackMamba 模型的代币数量。可以看出，由于 BlackMamba 模型实施了改进的 Sinkhorn 算法，大多数层都表现出高水平的专家平衡。 

下表涵盖了 BlackMamba 框架与一系列开源预训练语言模型相比的评估分数。可以看出，BlackMamba 框架能够在所有基线上与大多数框架竞争并超越。此外，值得注意的是，性能优于 BlackMamba 的模型具有相当多的参数数量，并且性能差距很小，这表明 BlackMamba 框架具有较少参数的能力。 

最后的思考

在本文中，我们讨论了 BlackMamba，这是一种新颖的架构，它将 Mamba 状态空间模型与混合专家模型相结合，以获得这两个框架提供的好处。 BlackMamba 上的实验表明，它在训练 FLOP 和推理方面都优于现有的 Mamba 框架和 Transformer 基线。 BlackMamba 框架的卓越性能表明，它能够很好地继承和结合 Mamba 和 MoE 框架的能力，因为它将 MoE 的廉价且快速的推理与 Mamba 的线性复杂度生成相结合。我们讨论了 BlackMamba 框架的架构如何在训练 FLOP 和推理成本方面超越经过严格训练的大型语言模型、现有 Mamba 框架和混合专家模型。此外，BlackMamba 框架还同时继承了 Mixture of Expert 模型和 Mamba 框架的生成 FLOP 并减少了训练。