超越Llama 3.1！TII发布首个纯无注意力大模型，非Transformer架构站起来了

最新推荐文章于 2024-09-14 09:59:20 发布

PaperWeekly

最新推荐文章于 2024-09-14 09:59:20 发布

阅读量97

点赞数

文章标签： llama transformer 架构深度学习人工智能

原文链接：https://mp.weixin.qq.com/s?__biz=MzIwMTc4ODE0Mw==&mid=2247677995&idx=3&sn=1a29c9ab07448a45db794a2d6512c69b&chksm=97a8da2c5e42960cc47d9d0c02eaeb0ca271275957a034ed1e0d28eaa8ee8ada239ea1709171&scene=126&sessionid=0

版权

©作者 | 杜伟、陈陈

来源 | 机器之心

Mamba 架构的大模型又一次向 Transformer 发起了挑战。

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 远远落后于这些模型。