Mamba再次挑戰(zhàn)霸主Transformer！首個通用Mamba開源大模型一鳴驚人

---

  新智元報道  

今天，阿布扎比支持的技術(shù)創(chuàng)新研究所(TII) 開源了全球第一個通用的大型Mamba架構(gòu)模型——Falcon Mamba 7B。

雖然之前Mistral已經(jīng)發(fā)過Mamba架構(gòu)的Codestral Mamba模型，但僅針對編碼；Falcon Mamba則是通用模型，能夠處理各種文本生成任務。

它是繼Falcon 180B、Falcon 40B和Falcon 2之后TII的第四個開放模型，與Falcon系列之前的型號不同，F(xiàn)alcon Mamba 7B完全采用SSLM架構(gòu)而不是傳統(tǒng)的Transformer架構(gòu)。

Mamba架構(gòu)橫空出世后，體現(xiàn)出了內(nèi)存效率方面的顯著優(yōu)勢，無需額外的內(nèi)存需求即可生成大量文本。

如今，SSLM正在逐漸蠶食Transformer架構(gòu)原本「大一統(tǒng)」的地位。

測評數(shù)據(jù)顯示，F(xiàn)alcon Mamba 7B性能已經(jīng)超越同尺寸級別的領(lǐng)先模型，例如Meta最新開源模型Llama 3.1 8B和Mistral 7B。

Falcon Mamba 7B將根據(jù)TII Falcon License 2.0發(fā)布，這是一個基于Apache 2.0的許可證，其中包括促進負責任地使用人工智能的使用政策。

雖然Transformer模型仍然主導著AI大模型領(lǐng)域，但研究人員指出，該架構(gòu)在處理較長文本時可能會遇到困難。

Transformer的自注意力機制（Self-Attention）讓模型可以關(guān)注輸入序列中的所有位置，并為每個位置分配不同的注意力權(quán)重。

這使得模型能夠更好地處理長距離的依賴關(guān)系，也就是說，對于句子中距離較遠的單詞，模型也能有效地捕獲其關(guān)系。

這種通過比較文本中每個token來理解上下文的方式，需要更多的計算能力和內(nèi)存來處理不斷增長的上下文窗口。

如果資源沒有相應擴展，推理速度會變慢，最終無法處理超過某個固定長度的文本。

為了解決這些難題，狀態(tài)空間語言模型（SSLM）架構(gòu)應運而生，該架構(gòu)通過在處理單詞時持續(xù)更新「狀態(tài)」，已成為一種有前途的替代方案。它已經(jīng)被一些組織部署，TII是最新的采用者。

這個全新的Falcon模型正是使用了CMU和普林斯頓大學的研究人員在2023年12月的一篇論文中最初提出的Mamba SSM架構(gòu)。

該架構(gòu)使用選擇機制，使模型能夠根據(jù)輸入動態(tài)調(diào)整其參數(shù)。

通過這種方式，模型可以專注于或忽略特定輸入，類似于Transformer中的注意力機制，但同時具備處理長文本序列（如整本書）的能力，而無需額外的內(nèi)存或計算資源。

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

首個通用大規(guī)模Mamba模型

上面提到，基于注意力機制的Transformer是當今所有最強大語言模型中占主導地位的架構(gòu)。然而，由于計算和內(nèi)存成本隨著序列長度的增加而增加，注意力機制在處理長序列時存在根本限制。

各種替代架構(gòu)，特別是SSLM，試圖解決序列擴展限制，但性能不及最先進的Transformer。

Falcon Mamba模型在不損失性能的前提下，可以突破序列擴展限制。

Falcon Mamba基于去年12月提出的第一版Mamba架構(gòu)，增加了RMS歸一化層以確保在大規(guī)模訓練中保持穩(wěn)定性。

這種架構(gòu)選擇確保了Falcon Mamba模型：

- 可以在不增加任何內(nèi)存存儲的情況下處理任意長度的序列，特別是可以在單張A10 24GB GPU上運行；

- 無論上下文大小，生成新token所需的時間恒定。

模型訓練

Falcon Mamba使用約5500GT（相當于5.5B token）的數(shù)據(jù)進行訓練，主要由RefinedWeb數(shù)據(jù)組成，并添加了公共來源的高質(zhì)量技術(shù)數(shù)據(jù)和代碼數(shù)據(jù)。

在大部分訓練中使用了恒定的學習率，隨后進行了一個較短的學習率衰減階段。

在最后階段，還加入了一小部分高質(zhì)量的精選數(shù)據(jù)，以進一步提升模型性能。

使用lm-evaluation-harness包對新排行榜版本的所有基準測試進行模型評估，然后使用HuggingFace分數(shù)歸一化處理評估結(jié)果。

如下圖所示，F(xiàn)alcon Mamba 7B獲得15.04的均分，超過Llama 3.1 8B 13.41分和Mistral 7B 14.50分。

此外，還使用了lighteval對大語言模型排行榜第一版的基準測試進行評估。

可以看到，F(xiàn)alcon Mamba 7B僅次于Transformer架構(gòu)的Falcon 2 11B，分數(shù)仍然超過Gemma、Llama、Mistral等同等規(guī)模的知名模型。

處理大規(guī)模序列

理論上來說，SSM模型在處理大規(guī)模序列時具有效率優(yōu)勢。

為了驗證模型的大規(guī)模序列處理能力，使用optimum-benchmark庫，對Falcon Mamba和流行的Transformer模型在內(nèi)存使用和生成吞吐量方面進行了比較。

為了公平比較，將所有Transformer模型的詞匯大小調(diào)整為與Falcon Mamba一致，因為這對模型的內(nèi)存需求有很大影響。

在查看結(jié)果之前，先討論序列中提示詞（預填充）和生成（解碼）部分的區(qū)別。

預填充的細節(jié)對于SSM而言，比對于Transformer模型更為重要。

當Transformer生成下一個token時，它需要關(guān)注上下文中所有先前token的鍵和值，這意味著內(nèi)存需求和生成時間都會隨著上下文長度線性增長。

而SSM僅關(guān)注并存儲其遞歸狀態(tài)，因此在生成大規(guī)模序列時不需要額外的內(nèi)存或時間。

雖然這解釋了SSM在解碼階段相對于Transformer的優(yōu)勢，但在預填充階段需要使用新方法來充分利用SSM架構(gòu)。

預填充的標準方法是并行處理整個提示詞以充分利用GPU。這種方法在optimum-benchmark庫中使用，我們稱之為并行預填充。

并行預填充需要將提示詞每個token的隱藏狀態(tài)存儲在內(nèi)存中。對于Transformer，這額外的內(nèi)存主要由存儲的KV緩存占據(jù)。

對于SSM模型，不需要緩存，存儲隱藏狀態(tài)的內(nèi)存成為唯一與提示詞長度成比例的部分。

因此，內(nèi)存需求將隨提示詞長度增長，SSM模型將失去處理任意長序列的能力，類似于Transformer。

并行預填充的替代方法是逐個處理token提示詞，我們稱之為順序預填充。

類似于序列并行處理，它也可以大規(guī)模地處理提示詞，而不是單個token，以更好地利用GPU。

雖然順序預填充對Transformer意義不大，但它為SSM模型帶來了處理任意長提示詞的可能性。

考慮到這些觀點，實驗首先測試了可以在單個24GB A10 GPU上適應的最大序列長度。

其中，批大小固定為1，使用float32精度。

即使在并行預填充中，F(xiàn)alcon Mamba也能適應比Transformer更大的序列，而在順序預填充中發(fā)揮了全部潛力，可以處理任意長度的提示詞。

接下來，在提示詞長度為1，生成token數(shù)量最多為130k的情況下測量生成吞吐量，使用批大小為1，并在H100 GPU上進行。

結(jié)果如圖所示。可以觀察到，F(xiàn)alcon Mamba在生成所有token時保持恒定的吞吐量，且GPU峰值內(nèi)存沒有增加。

而對于Transformer模型，隨著生成token數(shù)量的增加，峰值內(nèi)存增加，生成速度變慢。

Falcon Mamba架構(gòu)將在HuggingFace transformers庫的下一個版本（4.45.0以上）中提供。

使用Falcon Mamba 7B模型，需要安裝最新版本的HuggingFace transformers，或從源代碼安裝庫。

Falcon Mamba與HuggingFace提供的大多數(shù)API兼容，這些API已經(jīng)比較熟悉，例如：

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內(nèi)存較小的情況下運行模型，例如：

此外，還推出了Falcon Mamba的指令微調(diào)版本，該版本經(jīng)過額外50億個token的監(jiān)督微調(diào)（SFT），這種擴展訓練提高了模型在執(zhí)行指令任務時的精確性和有效性。

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))

可以通過演示體驗該指令模型的功能，對于聊天模板，可以使用以下格式：

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

用戶還可以直接使用基礎(chǔ)模型和指令模型的4-bit轉(zhuǎn)換版本，但要保證GPU與bitsandbytes庫兼容，才能運行量化模型。

用戶還可以通過torch.compile獲得更快的推理速度，加載模型后，只需調(diào)用model = torch.compile(model)。