Falcon Mamba 第一个具有竞争力的无注意7B语言模型 ！

点击下方卡片，关注 「AI视界引擎」 公众号

( 添加时备注：方向+学校/公司+昵称/姓名 )

在本技术报告中，作者提出了Falcon Mamba 7B，一种基于新颖Mamba架构的新一代大型语言模型。Falcon Mamba 7B在5.8万亿个 Token 上进行训练，这些 Token 是通过精心选择的数据混合物得到的。

作为纯Mamba架构的模型，Falcon Mamba 7B在基于 Transformer 的开源权重模型中超越了 Mistral 7B，Llama3.1 8B，和Falcon2 11B。

它的性能与Gemma 7B相当，并在不同的架构设计模型（如RecurrentGemma 9B和RWKV-v6 Finch 7B/14B）中表现出色。

目前，Falcon Mamba 7B是文献中在这个规模上表现最好的Mamba模型，超越了现有的Mamba和混合Mamba-Transformer模型，根据Open LLM Leaderboard（Fourrier等，2024）。由于其架构，Falcon Mamba 7B在推理方面显著 faster，且需要更少的内存进行长序列生成。

尽管最近的研究表明，混合Mamba-Transformer模型优于纯架构设计，但作者证明即使纯Mamba设计也可以与Transformer和混合设计相比实现类似，甚至更好的结果。

作者在https://huggingface.co/tiiuea/falcon-mamba-7b上公开了作者实现的Falcon Mamba 7B的权重，并采用了一种许可宽松的许可证。

1 Introduction

现代的foundation模型主要基于Transformer及其核心关注层 。由于其与窗口长度的平方复杂度相关，最近的研究试图提出纯注意力的更高效替代方案，例如FlashAttention（Dao等人，2022年；Dao，2024年）、滑动窗口关注 。最近提出了超越Transformer的新架构，例如 Griffin（De等人，2024年）、RWKV（Peng等人，2023年）和Mamba ，并已显示出与Transformer相当的性能。然而，它们中的大多数要么在小型规模上证明了自己的性能，要么与最近的基于Transformer的性能良好的LLM仍存在性能差距。

社区已经做出了努力，将Mamba LLMs扩展到原始测试目的2.8B的Mamba模型之外。值得注意的是，这些模型采用了混合Mamba-Transformer设计，与纯Transformer模型相比，表现出了优越的性能。

然而，在大数据和模型规模方面，纯无注意力模型是否能与高度优化的Transformer相媲美，仍然不清楚。

作者介绍了Falcon Mamba 7B —— 一个基于（预训练）的纯Mamba架构设计的基础模型，同时也是FalconLLM系列中的首个状态空间语言模型（SSLM）。作者认为Falcon Mamba 7B正面回答了上述问题，并且，据作者所知，这是首个做到如此的模型。通过使用Open LLM Leaderboard（Fourrier等人，2024年）的基准测试集进行衡量，Falcon Mamba 7B与强大的基于 Transformer 的预训练LLM，Mistral 7B 和Falcon2 11B 相当或超越。

此外，它还超过了其他架构设计的模型，如基于Griffin和RWKV-v6 Finch 7B和14B（Peng等人，2024年）的RecurrentGemma 9B（Botev等人，2024年）。更重要的是，由于采用了纯Mamba设计，Falcon Mamba 7B在上下文长度不同时，维持了恒定的内存成本，同时在大规模上下文数据生成方面提供了极端高效推理。

在这份技术报告中，作者详细介绍了Falcon Mamba 7B模型架构、训练配方和预训练数据准备。之后，作者将对在流行基准测试上采用不同架构的LLM进行详细比较。

最后，作者将展示Falcon Mamba 7B的更广泛含义、局限性和优势，并得出结论。

2 Model Architecture

《猛禽响尾蛇》7B模型架构基于Mamba（郭 & 道，2023）。该架构的核心参数总结在表1中。

作者在整个训练过程中解耦了输入嵌入层与输出权重，以提高模型的灵活性。根据作者的实验结果，这种方法在7B规模上提高了模型性能。

请注意，与 Transformer 不同，序列长度并不是Mamba架构的一部分。在推理过程中，可以使用任何序列长度，而模型处理长序列的实际能力取决于用于训练的序列长度。

设计决策：最近的研究表明，一种混合架构，具有交错注意力和支持向量机（SSM）层，可以超越纯Transformer或SSM模型。这种改进被认为源于两种模型的互补特性：SSM的一般序列到序列映射能力以及注意力层的快速检索属性。

最近基于Mamba的语言模型遵循这种直觉，将混合设计扩展到超过2.8亿模型，例如Samba 3.8B（Ren等人，2024年），Zamba 7B，Jamba 12B/52B。

然而，引入注意力层破坏了Mamba架构的线性可扩展性，引发了这样一个问题：在规模上，一个纯粹基于Mamba的架构能否实现与最先进的（state-of-the-art，SoTA）开源语言模型竞争，同时保持其线性可扩展性？最近的无注意力模型，如RWKV-v6（Peng等人，2024年），在小型尺度或在某些学术基准上展示了其性能。然而，在各种基准上的更全面的比较中，它们远远落后于流行的语言模型。

在预训练过程中，作者观察到一致的损失峰值随机且不可预测地出现。值得注意的是，当作者应用更高的学习率时，模型表现出更明显的损失峰值，并更容易出现偏差。这一现象也在Falcon2（Malartic等人，2024年）的训练中观察到，最近的一些论文，如Jamba（Lieber等人，2024年）和Mamba2（Dao和Gu，2024年）也报告了类似的问题。

特别的是，作者发现Mamba架构对学习率比Transformer更敏感。

仔细的模型初始化和减少模型的学习率敏感性对于解决这个问题至关重要。与（Dehghani等人，2023年）一致，应用预范数和后范数与RMSNorm已经成为一种常见的做法，以稳定预训练。

同样，作者在B、C和Δ之后添加了RMSNorm层。从作者的实验中，似乎比其他设置（如在每个块的最终输出投影之前添加RMSNorm层）带来更稳定的训练损失。这与Jamba模型的设计（Lieber等人，2024年）一致。

3 Pre-training

Training stategy

falcon-mamba-7B 的大部分训练时间是在256个H100 80GB GPU上进行的，使用仅有的数据并行（DP=256）。这还结合了ZeRO优化，以高效管理内存和训练过程。

该模型使用AdamW优化器进行训练，其中$\beta\_{1}=0.9$，$\beta\_{2}=0.95$，$\epsilon=10^{-8}$，权重衰减值$0.1$。尽管在Falcon-Mamba-7B预训练期间未在输出对数概率上应用Z-损失，但在后续实验中作者观察到，它有助于稳定训练，与（Wortsman等人，2024年）一致。

作者应用了恒定-衰减（WSD）学习率计划（Hu等人，2024年），其中固定暖身持续时间为1GT，稳定阶段的学习率为$\eta\_{\text{max}}=6.4\times 10^{-4}$。这样，作者的模型在大多数相关阶段使用相对较高的学习率进行训练，导致在不同训练阶段之间和衰减阶段开始时引入的数据分布变化快速适应（见3.2.2节）。

在衰减阶段，作者使用指数计划将学习率降低到最小值$\eta\_{\text{min}}=\frac{\eta\_{\text{max}}}{256}$，其中$t\_{\text{decay}}$是衰减阶段的持续时间。与大多数技术报告不同，作者发现较长的LR衰减阶段在评估方面提供了更好的结果。作者将总训练 Token 的约10%保留给衰减阶段，以实现最佳性能，这与最近miniCPM的结论一致（Hu等人，2024年）。

在训练开始时，作者使用了批量大小逐渐增加的方式。具体来说，作者将批量大小初始值 $b\_{\text{min}}=128$ 线性增加到最大值 $b\_{\text{max}}=2048$，在整个前 50 个 GPU 周期内。在作者的实验中，作者发现批量大小逐渐增加会影响损失曲线和最终模型性能。这种影响最方便地用梯度噪声温度 $T\_{\text{噪声}}$ 解释，对于 Adam 优化器，它被定义为（Malladi 等人，2022）。