大型語言模型(LLMs),例如 GPT-4 和 Llama-2,雖然非常強大,但需要大量的計算資源,這使得它們不太適合在小型設備上使用。特別是基於注意力的變壓器模型,對記憶體的需求很高,計算複雜度是平方級的,這限制了它們的效率。狀態空間模型(SSMs),例如 Mamba,提供了一個替代方案,具有較低的複雜度,但它們的記憶回憶能力有限,對於複雜任務的表現受到影響。現有的混合模型通常是將變壓器和 SSM 層按順序結合,但往往缺乏最佳性能所需的協同效應。
NVIDIA 發布 Hymba:混合頭並行架構
NVIDIA 推出了 Hymba,一個新的小型語言模型系列,具有混合架構,結合了 Mamba 和並行運行的注意力頭。這個模型擁有 15 億個參數,旨在解決小型自然語言處理(NLP)模型面臨的效率和性能挑戰,並在 1.5 兆個標記上進行訓練。
NVIDIA 的 Hymba 模型具有混合頭並行架構,將變壓器的注意力機制與 SSM 結合,以提高效率。這種架構允許注意力頭和 SSM 頭並行處理輸入數據,結合了兩種方法的優勢。注意力頭提供高解析度的記憶回憶,而 SSM 頭則能有效地總結上下文。
Hymba 還引入了可學習的元標記,這些標記會添加到每個輸入提示之前,以幫助存儲關鍵信息並減輕注意力機制的負擔。該模型的架構進一步優化,使用跨層的鍵值(KV)共享和部分滑動窗口注意力,以保持緊湊的快取大小,有效解決記憶體限制。
技術細節
Hymba-1.5B 模型結合了並行運行的 Mamba 和注意力頭,並使用元標記來提高效率。這種設置減少了變壓器的計算負擔,並不影響記憶回憶。Hymba 包含 16 個 SSM 狀態和 3 個完整的注意力層,其餘則使用滑動窗口注意力,以平衡效率和記憶解析度。它還具有來自 PyTorch 2.5 的 FlexAttention,為模型的訓練和推理增添了靈活性。
Hymba 的一個關鍵特徵是能夠在多個層之間以及同一層內的頭之間共享 KV 快取,顯著減少了記憶體使用。滑動窗口注意力和共享 KV 快取的結合最小化了計算複雜度,使 Hymba 相比其他同類型模型更為高效。
效率、性能和多功能性
Hymba 展示了小型語言模型可以在計算上高效的同時達到競爭性的性能。在基準測試中,Hymba-1.5B-Base 模型超越了所有 20 億以下的公共模型,並以 1.32% 的平均準確率超過 Llama-3.2-3B,快取大小減少了 11.67 倍,吞吐量提高了 3.49 倍。這使得 Hymba 適合在小型、性能較低的硬體上部署。
Hymba 的混合注意力和 SSM 設置在各種任務中提高了性能,包括一般基準測試和需要大量回憶的任務。它的吞吐量約為每秒 664 個標記,顯著高於其他模型,如 SmolLM2 或 Llama-3.2-3B,後者在類似測試場景中遇到了記憶體不足的問題。這些指標突顯了 Hymba 在需要速度和記憶體效率的實際部署場景中的適用性。
結論
NVIDIA 的 Hymba 小型語言模型系列代表了 NLP 技術在效率和多功能性方面的重要進展。通過其混合頭並行架構結合變壓器注意力和狀態空間模型,Hymba 為在資源有限的設備上部署有效的 NLP 能力提供了一條途徑。該模型減少的記憶體需求、提高的吞吐量以及創新的元標記和跨層 KV 共享的使用,使其成為未來語言模型應用中效率和準確性都至關重要的有希望的選擇。
查看論文。對於有興趣進一步探索 Hymba 模型的人,NVIDIA 已經在 Hugging Face 上提供了 Hymba-1.5B-Base 和 Hymba-1.5B-Instruct。所有這項研究的功勞都歸於這個項目的研究人員。此外,別忘了在 Twitter 上關注我們,加入我們的 Telegram 頻道和 LinkedIn 群組。如果你喜歡我們的工作,你會喜歡我們的電子報。別忘了加入我們的 55k+ ML SubReddit。
[免費 AI 虛擬會議] SmallCon:免費虛擬 GenAI 會議,參加者包括 Meta、Mistral、Salesforce、Harvey AI 等。請於 12 月 11 日加入我們,了解如何從 AI 先驅者如 Meta、Mistral AI、Salesforce、Harvey AI、Upstage、Nubank、Nvidia、Hugging Face 等建立大型模型。
本文由 AI 台灣 運用 AI 技術編撰,內容僅供參考,請自行核實相關資訊。
歡迎加入我們的 AI TAIWAN 台灣人工智慧中心 FB 社團,
隨時掌握最新 AI 動態與實用資訊!