Mamba 架構(gòu)模型這次終于要「站」起來(lái)了？自 2023 年 12 月首次推出以來(lái)，Mamba 便成為了 Transformer 的強(qiáng)有力競(jìng)爭(zhēng)對(duì)手。

此后，采用 Mamba 架構(gòu)的模型不斷出現(xiàn)，比如 Mistral 發(fā)布的首個(gè)基于 Mamba 架構(gòu)的開源大模型 Codestral 7B。

今天，阿布扎比技術(shù)創(chuàng)新研究所（TII）發(fā)布了一個(gè)新的開源 Mamba 模型 ——Falcon Mamba 7B。

先來(lái)總結(jié)一波 Falcon Mamba 7B 的亮點(diǎn)：無(wú)需增加內(nèi)存存儲(chǔ)，就可以處理任意長(zhǎng)度的序列，并且能夠在單個(gè) 24GB A10 GPU 上運(yùn)行。

目前可以在 Hugging Face 上查看并使用 Falcon Mamba 7B，這個(gè)僅用因果解碼器的模型采用了新穎的 Mamba 狀態(tài)空間語(yǔ)言模型（State Space Language Model, SSLM）架構(gòu)來(lái)處理各種文本生成任務(wù)。

從結(jié)果來(lái)看，F(xiàn)alcon Mamba 7B 在一些基準(zhǔn)上超越同尺寸級(jí)別的領(lǐng)先模型，包括 Meta 的 Llama 3 8B、Llama 3.1 8B 和 Mistral 7B。

Falcon Mamba 7B 分為四個(gè)變體模型，分別是基礎(chǔ)版本、指令微調(diào)版本、4bit 版本和指令微調(diào) 4bit 版本。

作為一個(gè)開源模型， Falcon Mamba 7B 采用了基于 Apache 2.0 的許可證「Falcon License 2.0」，支持研究和應(yīng)用目的。

Hugging Face 地址：https://huggingface.co/tiiuae/falcon-mamba-7b

Falcon Mamba 7B 也成為了繼 Falcon 180B、Falcon 40B 和 Falcon 2 之后，TII 開源的第四個(gè)模型，并且是首個(gè) Mamba SSLM 架構(gòu)模型。

首個(gè)通用的大型純 Mamba 模型

一直以來(lái)，基于 Transformer 的模型一直占據(jù)著生成式 AI 的統(tǒng)治地位，然而，研究人員注意到，Transformer 架構(gòu)在處理較長(zhǎng)的文本信息時(shí)可能會(huì)遇到困難。

本質(zhì)上，Transformer 中的注意力機(jī)制通過(guò)將每個(gè)單詞（或 token）與文本中的每個(gè)單詞進(jìn)行比較來(lái)理解上下文，它需要更多的計(jì)算能力和內(nèi)存需求來(lái)處理不斷增長(zhǎng)的上下文窗口。

但是如果不相應(yīng)地?cái)U(kuò)展計(jì)算資源，模型推理速度就會(huì)變慢，超過(guò)一定長(zhǎng)度的文本就沒(méi)法處理了。為了克服這些障礙，狀態(tài)空間語(yǔ)言模型 (SSLM) 架構(gòu)應(yīng)運(yùn)而生，該架構(gòu)通過(guò)在處理單詞時(shí)不斷更新?tīng)顟B(tài)來(lái)工作，已成為一種有前途的替代方案，包括 TII 在內(nèi)的很多機(jī)構(gòu)都在部署這種架構(gòu)。

Falcon Mamba 7B 采用了卡內(nèi)基梅隆大學(xué)和普林斯頓大學(xué)研究人員最初在 2023 年 12 月的一篇論文中提出的 Mamba SSM 架構(gòu)。

該架構(gòu)使用一種選擇機(jī)制，允許模型根據(jù)輸入動(dòng)態(tài)調(diào)整其參數(shù)。這樣，模型可以關(guān)注或忽略特定輸入，類似于注意力機(jī)制在 Transformer 中的工作方式，同時(shí)提供處理長(zhǎng)文本序列（例如整本書）的能力，而無(wú)需額外的內(nèi)存或計(jì)算資源。

TII 指出，該方法使模型適用于企業(yè)級(jí)機(jī)器翻譯、文本摘要、計(jì)算機(jī)視覺(jué)和音頻處理任務(wù)以及估計(jì)和預(yù)測(cè)等任務(wù)。

訓(xùn)練數(shù)據(jù)

Falcon Mamba 7B 訓(xùn)練數(shù)據(jù)高達(dá) 5500GT ，主要由 RefinedWeb 數(shù)據(jù)集組成，并添加了來(lái)自公共源的高質(zhì)量技術(shù)數(shù)據(jù)、代碼數(shù)據(jù)和數(shù)學(xué)數(shù)據(jù)。所有數(shù)據(jù)通過(guò) Falcon-7B/11B 標(biāo)記器進(jìn)行 tokenized 操作。

與其他 Falcon 系列模型類似，F(xiàn)alcon Mamba 7B 采用多階段訓(xùn)練策略進(jìn)行訓(xùn)練，上下文長(zhǎng)度從 2048 增加到了 8192。此外，受到課程學(xué)習(xí)概念的啟發(fā)，TII 在整個(gè)訓(xùn)練階段精心選擇了混合數(shù)據(jù)，充分考慮了數(shù)據(jù)的多樣性和復(fù)雜性。

在最后的訓(xùn)練階段，TII 使用了一小部分高質(zhì)量精選數(shù)據(jù)（即來(lái)自 Fineweb-edu 的樣本），以進(jìn)一步提升性能。

訓(xùn)練過(guò)程、超參數(shù)

Falcon Mamba 7B 的大部分訓(xùn)練是在 256 個(gè) H100 80GB GPU 上完成的，采用了 3D 并行（TP=1、PP=1、DP=256）與 ZeRO 相結(jié)合的策略。下圖為模型超參數(shù)細(xì)節(jié)，包括精度、優(yōu)化器、最大學(xué)習(xí)率、權(quán)重衰減和 batch 大小。

具體而言，F(xiàn)alcon Mamba 7B 經(jīng)過(guò)了 AdamW 優(yōu)化器、WSD（預(yù)熱 - 穩(wěn)定 - 衰減）學(xué)習(xí)率計(jì)劃的訓(xùn)練， 并且在前 50 GT 的訓(xùn)練過(guò)程中，batch 大小從 b_min=128 增加到了 b_max=2048。

在穩(wěn)定階段，TII 使用了最大學(xué)習(xí)率 η_max=6.4×10^?4，然后使用超過(guò) 500GT 的指數(shù)計(jì)劃將其衰減到最小值。同時(shí)，TII 在加速階段采用了 BatchScaling 以重新調(diào)整學(xué)習(xí)率 η，使得 Adam 噪聲溫度保持恒定。

整個(gè)模型訓(xùn)練花費(fèi)了大約兩個(gè)月時(shí)間。

模型評(píng)估

為了了解 Falcon Mamba 7B 與同尺寸級(jí)別領(lǐng)先的 Transformer 模型相比如何，該研究進(jìn)行了一項(xiàng)測(cè)試，以確定使用單個(gè) 24GB A10GPU 時(shí)模型可以處理的最大上下文長(zhǎng)度。

結(jié)果顯示，F(xiàn)alcon Mamba 能夠比當(dāng)前的 Transformer 模型適應(yīng)更大的序列，同時(shí)理論上能夠適應(yīng)無(wú)限的上下文長(zhǎng)度。 

接下來(lái)，研究者使用批處理大小為 1 ，硬件采用 H100 GPU 的設(shè)置中測(cè)量模型生成吞吐量。結(jié)果如下圖所示，F(xiàn)alcon Mamba 以恒定的吞吐量生成所有 token，并且 CUDA 峰值內(nèi)存沒(méi)有任何增加。對(duì)于 Transformer 模型，峰值內(nèi)存會(huì)增加，生成速度會(huì)隨著生成的 token 數(shù)量的增加而減慢。

即使在標(biāo)準(zhǔn)的行業(yè)基準(zhǔn)測(cè)試中，新模型的性能也優(yōu)于或接近于流行的 transformer 模型以及純狀態(tài)空間模型和混合狀態(tài)空間模型。

例如，在 Arc、TruthfulQA 和 GSM8K 基準(zhǔn)測(cè)試中，F(xiàn)alcon Mamba 7B 的得分分別為 62.03%，53.42% 和 52.54%，超過(guò)了 Llama 3 8 B, Llama 3.1 8B, Gemma 7B 和 Mistral 7B。然而，在 MMLU 和 Hellaswag 基準(zhǔn)測(cè)試中，F(xiàn)alcon Mamba 7B 遠(yuǎn)遠(yuǎn)落后于這些模型。

TII 首席研究員 Hakim Hacid 在一份聲明中表示：Falcon Mamba 7B 的發(fā)布代表著該機(jī)構(gòu)向前邁出的重大一步，它激發(fā)了新的觀點(diǎn)，并進(jìn)一步推動(dòng)了對(duì)智能系統(tǒng)的探索。在 TII，他們正在突破 SSLM 和 transformer 模型的界限，以激發(fā)生成式 AI 的進(jìn)一步創(chuàng)新。

目前，TII 的 Falcon 系列語(yǔ)言模型下載量已超過(guò) 4500 萬(wàn)次 —— 成為阿聯(lián)酋最成功的 LLM 版本之一。

Falcon Mamba 7B 論文即將放出，大家可以等一等。