混合式 AI 模型如何平衡記憶體與效率

連結表格

摘要和1. 引言

1. 方法論

2. 實驗和結果

   3.1 在vQuality數據上的語言建模

   3.2 關於注意力和線性遞迴的探索

   3.3 高效長度外推

   3.4 長上下文理解

3. 分析

4. 結論、致謝和參考文獻

A. 實現細節

B. 額外實驗結果

C. 熵測量詳情

D. 局限性

\

A 實現細節

\ 對於滑動GLA架構中的GLA層，我們使用頭數dm/384、鍵擴展比率0.5和值擴展比率1。對於RetNet層，我們使用的頭數是注意力查詢頭數的一半，鍵擴展比率為1，值擴展比率為2。GLA和RetNet的實現來自Flash Linear Attention存儲庫[3] [YZ24]。我們使用基於FlashAttention的實現進行Self-Extend外推[4]。Mamba 432M模型的模型寬度為1024，而Mamba 1.3B模型的模型寬度為2048。除非另有說明，所有在SlimPajama上訓練的模型都具有相同的訓練配置和與Samba相同的MLP中間大小。SlimPajama上的訓練基礎設施基於TinyLlama代碼庫的修改版本[5]。

\

\ 在下游任務的生成配置中，我們對GSM8K使用貪婪解碼，對HumanEval使用溫度τ = 0.2和top-p = 0.95的核採樣[HBD+19]。對於MBPP和SQuAD，我們設置τ = 0.01和top-p = 0.95。

B 額外實驗結果

\

\

\

\

C 熵測量詳情

\

\

D 局限性

雖然Samba通過指令微調展示了有希望的記憶檢索性能，但其預訓練基礎模型的檢索性能與基於SWA的模型相似，如圖7所示。這為未來進一步提高Samba的檢索能力而不損害其效率和外推能力開闢了方向。此外，Samba的混合策略在所有任務中並不始終優於其他替代方案。如表2所示，MambaSWA-MLP在WinoGrande、SIQA和GSM8K等任務上表現出改進的性能。這使我們有可能投資於更複雜的方法來執行基於SWA和基於SSM模型的輸入依賴動態組合。

\

:::info 作者：

(1) Liliang Ren，Microsoft和伊利諾伊大學厄巴納-香檳分校 (liliangren@microsoft.com)；

(2) Yang Liu†，Microsoft (yaliu10@microsoft.com)；

(3) Yadong Lu†，Microsoft (yadonglu@microsoft.com)；

(4) Yelong Shen，Microsoft (yelong.shen@microsoft.com)；

(5) Chen Liang，Microsoft (chenliang1@microsoft.com)；

(6) Weizhu Chen，Microsoft (wzchen@microsoft.com)。

:::

:::info 本論文可在arxiv上獲取，採用CC BY 4.0許可證。

:::

[3] https://github.com/sustcsonglin/flash-linear-attention

\ [4] https://github.com/datamllab/LongLM/blob/master/selfextendpatch/Llama.py

\ [5] https://github.com/jzhang38/TinyLlama