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

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

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

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

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

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

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

作為一個開源模型， 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 開源的第四個模型，并且是首個 Mamba SSLM 架構(gòu)模型。

首個通用的大型純 Mamba 模型

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

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

但是如果不相應(yīng)地擴(kuò)展計算資源，模型推理速度就會變慢，超過一定長度的文本就沒法處理了。為了克服這些障礙，狀態(tài)空間語言模型 (SSLM) 架構(gòu)應(yīng)運而生，該架構(gòu)通過在處理單詞時不斷更新狀態(tà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ù)輸入動態(tài)調(diào)整其參數(shù)。這樣，模型可以關(guān)注或忽略特定輸入，類似于注意力機(jī)制在 Transformer 中的工作方式，同時提供處理長文本序列（例如整本書）的能力，而無需額外的內(nèi)存或計算資源。

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

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

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

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

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

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

Falcon Mamba 7B 的大部分訓(xùn)練是在 256 個 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)過了 AdamW 優(yōu)化器、WSD（預(yù)熱 - 穩(wěn)定 - 衰減）學(xué)習(xí)率計劃的訓(xùn)練， 并且在前 50 GT 的訓(xùn)練過程中，batch 大小從 b_min=128 增加到了 b_max=2048。

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

整個模型訓(xùn)練花費了大約兩個月時間。

模型評估

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

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

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

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

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

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

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

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