Mamba 架构模型这次终于要「站」起来了?自 2023 年 12 月首次推出以来,Mamba 便成为了 Transformer 的强有力竞争对手。
此后,采用 Mamba 架构的模型不断出现,比如 Mistral 发布的首个基于 Mamba 架构的开源大模型 Codestral 7B。
今天,阿布扎比技术创新研究所(TII)发布了一个新的开源 Mamba 模型 ——Falcon Mamba 7B。
先来总结一波 Falcon Mamba 7B 的亮点:无需增加内存存储,就可以处理任意长度的序列,并且能够在单个 24GB A10 GPU 上运行。
目前可以在 Hugging Face 上查看并使用 Falcon Mamba 7B,这个仅用因果解码器的模型采用了新颖的 Mamba 状态空间语言模型(State Space Language Model, SSLM)架构来处理各种文本生成任务。
从结果来看,Falcon Mamba 7B 在一些基准上超越同尺寸级别的领先模型,包括 Meta 的 Llama 3 8B、Llama 3.1 8B 和 Mistral 7B。
Falcon Mamba 7B 分为四个变体模型,分别是基础版本、指令微调版本、4bit 版本和指令微调 4bit 版本。
作为一个开源模型, Falcon Mamba 7B 采用了基于 Apache 2.0 的许可证「Falcon License 2.0」,支持研究和应用目的。
Hugging Face 地址:https://huggingface.co/tiiuae/falcon-mamba-7b
Falcon Mamba 7B 也成为了继 Falcon 180B、Falcon 40B 和 Falcon 2 之后,TII 开源的第四个模型,并且是首个 Mamba SSLM 架构模型。
首个通用的大型纯 Mamba 模型
一直以来,基于 Transformer 的模型一直占据着生成式 AI 的统治地位,然而,研究人员注意到,Transformer 架构在处理较长的文本信息时可能会遇到困难。
本质上,Transformer 中的注意力机制通过将每个单词(或 token)与文本中的每个单词进行比较来理解上下文,它需要更多的计算能力和内存需求来处理不断增长的上下文窗口。
但是如果不相应地扩展计算资源,模型推理速度就会变慢,超过一定长度的文本就没法处理了。为了克服这些障碍,状态空间语言模型 (SSLM) 架构应运而生,该架构通过在处理单词时不断更新状态来工作,已成为一种有前途的替代方案,包括 TII 在内的很多机构都在部署这种架构。
Falcon Mamba 7B 采用了卡内基梅隆大学和普林斯顿大学研究人员最初在 2023 年 12 月的一篇论文中提出的 Mamba SSM 架构。
该架构使用一种选择机制,允许模型根据输入动态调整其参数。这样,模型可以关注或忽略特定输入,类似于注意力机制在 Transformer 中的工作方式,同时提供处理长文本序列(例如整本书)的能力,而无需额外的内存或计算资源。
TII 指出,该方法使模型适用于企业级机器翻译、文本摘要、计算机视觉和音频处理任务以及估计和预测等任务。
训练数据
Falcon Mamba 7B 训练数据高达 5500GT ,主要由 RefinedWeb 数据集组成,并添加了来自公共源的高质量技术数据、代码数据和数学数据。所有数据通过 Falcon-7B/11B 标记器进行 tokenized 操作。
与其他 Falcon 系列模型类似,Falcon Mamba 7B 采用多阶段训练策略进行训练,上下文长度从 2048 增加到了 8192。此外,受到课程学习概念的启发,TII 在整个训练阶段精心选择了混合数据,充分考虑了数据的多样性和复杂性。
在最后的训练阶段,TII 使用了一小部分高质量精选数据(即来自 Fineweb-edu 的样本),以进一步提升性能。
训练过程、超参数
Falcon Mamba 7B 的大部分训练是在 256 个 H100 80GB GPU 上完成的,采用了 3D 并行(TP=1、PP=1、DP=256)与 ZeRO 相结合的策略。下图为模型超参数细节,包括精度、优化器、最大学习率、权重衰减和 batch 大小。
具体而言,Falcon Mamba 7B 经过了 AdamW 优化器、WSD(预热 - 稳定 - 衰减)学习率计划的训练, 并且在前 50 GT 的训练过程中,batch 大小从 b_min=128 增加到了 b_max=2048。
在稳定阶段,TII 使用了最大学习率 η_max=6.4×10^−4,然后使用超过 500GT 的指数计划将其衰减到最小值。同时,TII 在加速阶段采用了 BatchScaling 以重新调整学习率 η,使得 Adam 噪声温度保持恒定。
整个模型训练花费了大约两个月时间。
模型评估
为了了解 Falcon Mamba 7B 与同尺寸级别领先的 Transformer 模型相比如何,该研究进行了一项测试,以确定使用单个 24GB A10GPU 时模型可以处理的最大上下文长度。
结果显示,Falcon Mamba 能够比当前的 Transformer 模型适应更大的序列,同时理论上能够适应无限的上下文长度。
接下来,研究者使用批处理大小为 1 ,硬件采用 H100 GPU 的设置中测量模型生成吞吐量。结果如下图所示,Falcon Mamba 以恒定的吞吐量生成所有 token,并且 CUDA 峰值内存没有任何增加。对于 Transformer 模型,峰值内存会增加,生成速度会随着生成的 token 数量的增加而减慢。
即使在标准的行业基准测试中,新模型的性能也优于或接近于流行的 transformer 模型以及纯状态空间模型和混合状态空间模型。
例如,在 Arc、TruthfulQA 和 GSM8K 基准测试中,Falcon Mamba 7B 的得分分别为 62.03%,53.42% 和 52.54%,超过了 Llama 3 8 B, Llama 3.1 8B, Gemma 7B 和 Mistral 7B。然而,在 MMLU 和 Hellaswag 基准测试中,Falcon Mamba 7B 远远落后于这些模型。
TII 首席研究员 Hakim Hacid 在一份声明中表示:Falcon Mamba 7B 的发布代表着该机构向前迈出的重大一步,它激发了新的观点,并进一步推动了对智能系统的探索。在 TII,他们正在突破 SSLM 和 transformer 模型的界限,以激发生成式 AI 的进一步创新。
目前,TII 的 Falcon 系列语言模型下载量已超过 4500 万次 —— 成为阿联酋最成功的 LLM 版本之一。
Falcon Mamba 7B 论文即将放出,大家可以等一等。