近年來，大語言模型(LLM)在各個(gè)領(lǐng)域取得了顯著成效。但現(xiàn)有的Transformer架構(gòu)存在計(jì)算復(fù)雜度高、內(nèi)存消耗大等問題。而狀態(tài)空間模型(SSM)如Mamba雖然具有常數(shù)復(fù)雜度和優(yōu)化的硬件性能，但在記憶回溯任務(wù)上表現(xiàn)較弱。針對(duì)這一問題，NVIDIA提出了Hymba架構(gòu)，通過在同一層中結(jié)合注意力頭和SSM頭，以實(shí)現(xiàn)兩種架構(gòu)優(yōu)勢的互補(bǔ)。

核心創(chuàng)新

Hymba的核心創(chuàng)新主要包括三個(gè)方面：

1.并行混合頭設(shè)計(jì)

• 在同一層內(nèi)并行集成注意力頭和SSM頭
• 注意力機(jī)制提供高分辨率記憶回溯能力
• SSM提供高效的上下文總結(jié)能力
• 這種設(shè)計(jì)相比Zamba和Jamba等只在不同層使用兩種機(jī)制的方法更加靈活

2.可學(xué)習(xí)的元令牌(Meta Tokens)

• 在輸入序列前添加可學(xué)習(xí)的元令牌
• 這些令牌與所有后續(xù)令牌交互
• 充當(dāng)知識(shí)的壓縮表示
• 提高了回溯和通用任務(wù)性能

3.KV緩存優(yōu)化

• 在層間共享KV緩存
• 大多數(shù)層使用滑動(dòng)窗口注意力機(jī)制
• 顯著減少了內(nèi)存和計(jì)算成本

架構(gòu)設(shè)計(jì)

如論文圖1所示，Hymba的混合頭模塊包含：

1.輸入處理

• 輸入序列前添加Meta Tokens
• 通過投影層將輸入轉(zhuǎn)換為查詢、鍵、值以及SSM特征

2.并行處理

• 注意力頭處理高精度記憶回溯
• SSM頭進(jìn)行高效的上下文總結(jié)
• 兩種頭并行處理相同的輸入信息

3.輸出融合

• 對(duì)注意力頭和SSM頭的輸出進(jìn)行歸一化
• 通過可學(xué)習(xí)的向量進(jìn)行重新縮放
• 最后取平均得到最終輸出

性能優(yōu)勢

相比現(xiàn)有模型,Hymba-1.5B在多個(gè)方面都展現(xiàn)出顯著優(yōu)勢:

1.與Llama 3.2 3B相比

• 準(zhǔn)確率提高1.32%
• 緩存大小減少11.67倍
• 吞吐量提高3.49倍

2.與同等規(guī)模(2B以下)模型相比

• 在常識(shí)推理任務(wù)上取得最好性能
• 需要的緩存大小顯著減小
• 具有更高的處理速度

3.指令微調(diào)后的變體Hymba-1.5B-Instruct

• 在GSM8K和GPQA等基準(zhǔn)測試上表現(xiàn)優(yōu)異
• 經(jīng)常超越更大規(guī)模的模型

Hymba架構(gòu)實(shí)現(xiàn)與實(shí)驗(yàn)評(píng)估

1. 融合混合頭模塊設(shè)計(jì)

Hymba提出了一個(gè)統(tǒng)一且對(duì)稱的模塊設(shè)計(jì)公式。對(duì)于輸入序列 X?（原始輸入序列 X 加上元令牌），主要包括：

輸入投影:

• 使用 Win_proj = [WQ, WK, WV, WSSM, WG] 進(jìn)行投影
• 生成注意力頭的查詢、鍵、值
• 生成SSM頭的輸入特征和門控信號(hào)

注意力頭輸出：

SSM頭輸出：

輸出融合：

其中β1和β2是可學(xué)習(xí)的向量，用于重新縮放各通道的輸出。

2. KV緩存優(yōu)化策略

全局與局部注意力結(jié)合：

• 僅在關(guān)鍵層(第一層、中間層和最后一層)使用全局注意力
• 其他層使用滑動(dòng)窗口注意力(SWA)
• 該策略在維持性能的同時(shí)顯著提升效率

跨層KV共享：

• 相鄰層間共享鍵值緩存
• 減少參數(shù)冗余
• 節(jié)省的參數(shù)可以重新分配給其他模型組件

3. 元令牌的創(chuàng)新應(yīng)用

主要功能：

• 防止令牌重寫：為模型提供獨(dú)立于輸入的令牌
• 處理"強(qiáng)制關(guān)注"問題：通過修改softmax的分母來優(yōu)化注意力分布
• KV緩存和SSM狀態(tài)的初始化：可以看作是一種學(xué)習(xí)到的提示調(diào)優(yōu)

實(shí)現(xiàn)效果：

• 降低了注意力圖的熵
• 幫助模型更好地聚焦于重要信息
• 提升了回溯能力和常識(shí)推理性能

實(shí)驗(yàn)評(píng)估

1.基準(zhǔn)測試性能

如論文表2所示，在1.5T預(yù)訓(xùn)練數(shù)據(jù)條件下，Hymba-1.5B相比同規(guī)模模型具有明顯優(yōu)勢：

(1)與SmolLM2-1.7B比較

• 平均準(zhǔn)確率提升1.02%
• 緩存大小減少19.91倍
• 吞吐量提高2.79倍

(2)與其他2T以下訓(xùn)練數(shù)據(jù)的模型比較

• 相比Phi-1.5提升平均準(zhǔn)確率5.21%
• 相比h2o-danube2-1.8B提升5.41%

2、指令微調(diào)效果

(1)基礎(chǔ)指令微調(diào)

• 采用兩階段策略：全量微調(diào)(FFT)和直接偏好優(yōu)化(DPO)
• 在GSM8K、GPQA等任務(wù)上達(dá)到同類最佳性能

(2)DoRA參數(shù)高效微調(diào)

• 在RoleBench上超越了Llama-3.1-8B-Instruct約2.4%
• 展示了模型在參數(shù)高效微調(diào)場景的潛力

3、消融實(shí)驗(yàn)結(jié)果

(1)架構(gòu)組件分析

• 混合頭結(jié)構(gòu)比順序疊加提升顯著
• KV緩存優(yōu)化在保持性能的同時(shí)大幅提升效率
• 元令牌的引入進(jìn)一步提升了模型表現(xiàn)

(2)頭部重要性分析

• SSM頭在第一層對(duì)語言建模至關(guān)重要
• 移除單個(gè)注意力頭平均導(dǎo)致0.24%性能下降
• 移除單個(gè)SSM頭平均導(dǎo)致1.1%性能下降

這些實(shí)驗(yàn)結(jié)果充分證明了Hymba架構(gòu)的有效性和優(yōu)勢。

Hymba模型訓(xùn)練實(shí)現(xiàn)細(xì)節(jié)

1.預(yù)訓(xùn)練策略

如論文圖8所示，Hymba采用了多階段的訓(xùn)練流程：

基礎(chǔ)預(yù)訓(xùn)練階段:

• 使用較大學(xué)習(xí)率(3e-3)
• 采用DataCompLM數(shù)據(jù)集
• 訓(xùn)練1T個(gè)token

學(xué)習(xí)率退火階段:

• 逐漸將學(xué)習(xí)率降至1e-5
• 使用高質(zhì)量數(shù)據(jù)集
• 總共處理約500B個(gè)token

上下文擴(kuò)展:

• 將序列長度從2K擴(kuò)展到8K
• 調(diào)整ROPE基礎(chǔ)參數(shù)
• 進(jìn)一步提升長序列處理能力

2.模型系列規(guī)格

根據(jù)論文表11的描述，Hymba提供了三種不同規(guī)格的模型：

（1）Hymba-125M

• 24個(gè)模塊
• 隱藏層大小512
• 8個(gè)注意力頭
• 總參數(shù)量約125M

（2）Hymba-350M

• 32個(gè)模塊
• 隱藏層大小768
• 12個(gè)注意力頭
• 總參數(shù)量約350M

（3）Hymba-1.5B

• 32個(gè)模塊
• 隱藏層大小1600
• 25個(gè)注意力頭
• 總參數(shù)量約1.52B

3.指令微調(diào)實(shí)現(xiàn)

（1）監(jiān)督微調(diào)(SFT)

• 第一階段：使用900K樣本/3B tokens
• 第二階段：使用6.5M樣本/10B tokens
• 涵蓋代碼、數(shù)學(xué)、MMLU等多個(gè)領(lǐng)域

（2）DPO優(yōu)化

• 使用200K樣本/0.7B tokens
• 進(jìn)一步改進(jìn)指令遵循能力
• 采用余弦學(xué)習(xí)率調(diào)度

實(shí)際應(yīng)用與局限性分析

Hymba模型在實(shí)際應(yīng)用中展現(xiàn)出獨(dú)特的優(yōu)勢，特別是在處理長序列文本時(shí)表現(xiàn)突出。通過SSM實(shí)現(xiàn)的高效上下文編碼和滑動(dòng)窗口注意力機(jī)制，顯著降低了內(nèi)存消耗，使其非常適合在資源受限的環(huán)境中部署。在特定任務(wù)上，如數(shù)學(xué)推理、函數(shù)調(diào)用和角色扮演等場景，Hymba表現(xiàn)出與大型模型相媲美的性能，這使其成為一個(gè)極具實(shí)用價(jià)值的輕量級(jí)選擇。

但是作為一個(gè)相對(duì)小型的語言模型，Hymba也存在一些固有的局限性。由于參數(shù)量的限制，在處理某些需要深度推理或廣泛知識(shí)儲(chǔ)備的復(fù)雜任務(wù)時(shí)，其表現(xiàn)可能不如參數(shù)量更大的模型。此外混合架構(gòu)的設(shè)計(jì)雖然創(chuàng)新，但也帶來了實(shí)現(xiàn)和優(yōu)化方面的挑戰(zhàn)。模型訓(xùn)練過程需要更復(fù)雜的調(diào)參策略，這增加了模型開發(fā)和部署的技術(shù)門檻。

未來展望

從技術(shù)發(fā)展的角度來看，Hymba的創(chuàng)新架構(gòu)為語言模型的發(fā)展開辟了新的方向。未來的研究可能會(huì)進(jìn)一步探索注意力機(jī)制和SSM的最優(yōu)配比，以及更高效的融合策略。隨著計(jì)算資源的提升和算法的優(yōu)化，研究者們可能會(huì)嘗試擴(kuò)展模型規(guī)模，同時(shí)保持其高效處理的特性。特別值得關(guān)注的是，如何在保持計(jì)算效率的同時(shí)進(jìn)一步提升模型性能，這個(gè)平衡點(diǎn)的探索將是未來研究的重要方向。

在應(yīng)用拓展方面，Hymba展現(xiàn)出的混合架構(gòu)思路可能會(huì)被引入到更多領(lǐng)域。例如，將這種架構(gòu)應(yīng)用到多模態(tài)任務(wù)中，探索在視覺-語言交互等場景下的效果。同時(shí)，針對(duì)特定垂直領(lǐng)域的優(yōu)化也是一個(gè)重要方向，通過專門的微調(diào)策略，可能會(huì)在特定場景下取得更好的表現(xiàn)。

Hymba的出現(xiàn)為解決語言模型在效率和性能之間的權(quán)衡提供了新的思路。雖然目前仍存在一些局限性，但其創(chuàng)新的架構(gòu)設(shè)計(jì)和實(shí)驗(yàn)結(jié)果表明，這種混合架構(gòu)很可能成為未來語言模型發(fā)展的一個(gè)重要方向。隨著技術(shù)的不斷進(jìn)步和應(yīng)用場景的拓展，我們有理由期待基于這種架構(gòu)的更多突破性進(jìn)展。