Mamba再次挑战霸主Transformer！首个通用Mamba开源大模型一鸣惊人

直接说重点：阿布扎比的技术创新研究所（TII）今天开源了全球第一个通用型的大型Mamba架构模型——Falcon Mamba 7B。这可不是个小动静。

之前Mistral也发过一款Mamba架构的模型（Codestral Mamba），但那个只针对代码生成。而Falcon Mamba是通用的，能处理各种文本生成任务。它是TII继Falcon 180B、Falcon 40B和Falcon 2之后推出的第四个开放模型，但和前面三个不同，这回完全抛弃了Transformer，改用SSLM架构（也就是状态空间语言模型）。

Mamba架构自诞生以来，就在内存效率上展现出了碾压式的优势——生成大量文本时几乎不额外占用内存。如今，SSLM正在一点点蚕食Transformer原本“大一统”的地位。从测评数据看，Falcon Mamba 7B的表现已经超过了同尺寸的顶级模型，比如Meta刚开源的Llama 3.1 8B和Mistral 7B。模型将根据TII Falcon License 2.0发布，这是个基于Apache 2.0的许可证，附带负责任使用AI的使用政策。

Falcon Mamba 7B有什么特别之处？

虽然Transformer依然主导着大模型领域，但研究人员早就指出：它在处理长文本时会遇到瓶颈。Transformer的自注意力机制需要比较文本中每个token来理解上下文，随着上下文窗口增长，计算和内存需求也跟着疯长。要是资源跟不上，推理速度就会变慢，最终超出某个长度就彻底没法干了。

为了解决这个问题，状态空间语言模型（SSLM）架构应运而生。它通过在处理单词时持续更新“状态”，成了一条很有前景的出路。TII是最近一批吃螃蟹的组织之一。他们使用的Mamba SSM架构，最早由CMU和普林斯顿大学的研究人员在2023年12月的论文中提出（论文地址：https://arxiv.org/pdf/2312.00752）。

这个架构用了一种选择机制，让模型能根据输入动态调整参数。这样一来，模型可以像Transformer那样关注或忽略特定输入，但处理长文本序列（比如整本书）时却不需要额外内存或计算资源。TII指出，这种方法特别适合企业级的机器翻译、文本摘要、计算机视觉、音频处理以及估计预测等任务。

首个通用大规模Mamba模型

上面提到，基于注意力机制的Transformer是当今最强语言模型的主导架构。但注意力机制在处理长序列时有个根本限制：计算和内存成本随序列长度线性增长。各种替代架构（尤其是SSLM）试图解决这个问题，但性能一直追不上最先进的Transformer。

Falcon Mamba的出现打破了这一局面：它在不损失性能的前提下，突破了序列扩展的限制。它基于去年12月提出的第一版Mamba架构，额外增加了RMS归一化层来保证大规模训练时的稳定性。这种架构设计确保了：第一，可以在不增加任何内存存储的情况下处理任意长度的序列——单张A10 24GB GPU就能运行；第二，无论上下文多大，生成新token所需的时间是恒定的。

模型训练

Falcon Mamba使用了约5500GT（相当于5.5B token）的数据进行训练，主体是RefinedWeb数据，外加来自公开渠道的高质量技术数据和代码数据。大部分训练中用了恒定的学习率，最后接了一个较短的学习率衰减阶段。在最后阶段还加入了一小部分高质量的精选数据，来进一步提升模型性能。

性能评估

使用lm-evaluation-harness包对新排行榜版本的所有基准测试进行模型评估，然后用HuggingFace分数归一化处理结果。从下图可以看到，Falcon Mamba 7B获得了15.04的均分，超过了Llama 3.1 8B的13.41分和Mistral 7B的14.50分。

此外，还使用了lighteval对大语言模型排行榜第一版的基准测试进行评估。可以看到，Falcon Mamba 7B仅次于Transformer架构的Falcon 2 11B，分数仍然超过了Gemma、Llama、Mistral等同等规模的知名模型。

处理大规模序列

理论上，SSM模型在处理大规模序列时就有效率优势。为了验证这一点，团队用optimum-benchmark库，在内存使用和生成吞吐量方面将Falcon Mamba和流行的Transformer模型做了比较。为了公平，将所有Transformer模型的词汇大小调整为与Falcon Mamba一致，因为词汇大小对内存需求影响很大。

在查看结果之前，先区分一下序列中的提示词（预填充）和生成（解码）部分。预填充的细节对SSM来说比对Transformer更重要：当Transformer生成下一个token时，需要关注上下文里所有先前token的键和值，内存需求和生成时间都随上下文长度线性增长。而SSM只关注并存储它的递归状态，因此在生成大规模序列时不需要额外内存或时间。

虽然这解释了SSM在解码阶段的优势，但预填充阶段要用新方法来充分利用SSM架构。标准方法是并行处理整个提示词以充分利用GPU（称为并行预填充），这需要把提示词每个token的隐藏状态存在内存里。对于Transformer，额外内存主要由KV缓存占据；对于SSM，不需要缓存，存储隐藏状态的内存是唯一与提示词长度成比例的部分。因此，如果采用并行预填充，内存需求随提示词增长，SSM也会失去处理任意长序列的能力。

替代方案是逐个处理token（顺序预填充），类似于序列并行处理，可以大规模处理提示词，更好地利用GPU。顺序预填充对Transformer意义不大，但给SSM模型带来了处理任意长提示词的可能性。

实验首先测试了单张24GB A10 GPU上能适应的最大序列长度（批大小固定为1，使用float32精度）。即使在并行预填充中，Falcon Mamba也能适应比Transformer更大的序列；而在顺序预填充中，它发挥出全部潜力，可以处理任意长度的提示词。

接下来，在提示词长度为1、生成token数量最多为130k的情况下测量生成吞吐量（批大小1，H100 GPU）。结果如图：Falcon Mamba生成所有token时保持恒定吞吐量，GPU峰值内存没有增加；而Transformer模型随着生成token数量增加，峰值内存上升，生成速度变慢。

如何使用？

Falcon Mamba架构将在HuggingFace transformers库的下一个版本（4.45.0以上）中提供。使用模型需要安装最新版本或从源代码安装。它与HuggingFace提供的大部分API兼容，例如：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "tiiuae/falcon-mamba-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
inputs = tokenizer("Hello world, today", return_tensors="pt").to(0)
output = model.generate(**inputs, max_new_tokens=100, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=True))

它还支持像bitsandbytes库量化这样的功能，方便在GPU内存较小的情况下运行模型，例如：

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "tiiuae/falcon-mamba-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
quantization_config = BitsAndBytesConfig(load_in_4bit=True)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config)
inputs = tokenizer("Hello world, today", return_tensors="pt").to(0)
output = model.generate(**inputs, max_new_tokens=100, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=True))

此外，还推出了Falcon Mamba的指令微调版本，经过额外50亿个token的监督微调（SFT），提高了执行指令任务时的精确性和有效性。用户可以通过演示体验该指令模型的功能，聊天模板格式如下：

<|im_start|>user
prompt<|im_end|>
<|im_start|>assistant

用户也可以直接使用基础模型和指令模型的4-bit转换版本，但要保证GPU与bitsandbytes库兼容才能运行量化模型。还可以通过torch.compile获得更快的推理速度：加载模型后，只需调用model = torch.compile(model)。

参考资料：
https://huggingface.co/blog/falconmamba
https://venturebeat.com/ai/falcon-mamba-7bs-powerful-new-ai-architecture-offers-alternative-to-transformer-models/
https://medium.com/@puneetthegde22/mamba-architecture-a-leap-forward-in-sequence-modeling-370dfcbfe44a