ICLR2025|比MoE快6倍，成本暴降83％！位元組釋出超稀疏架構UltraMem

隨著模型規模的擴大，推理成本和訪存效率已成為限制大模型規模應用的關鍵瓶頸。儘管 MoE 架構已經成功將計算和引數解耦，但在推理時，較小的 batch size 就會啟用全部專家，導致訪存急劇上升，進而使推理延遲大幅增加。

早在 2019 年，Meta 提出了 PKM 架構，透過引入 large memory layer 減少訪存開銷，並於近期再次改進，提出了 Memory+。總體來看，業內已明確意識到了 MoE 的推理訪存問題，並開始研究更優的稀疏模型結構。

基於這一背景，字節跳動豆包大模型團隊近日釋出了一種全新的稀疏模型架構 UltraMem。該架構有效解決了 MoE 推理時高額的訪存問題，推理速度較 MoE 架構提升 2-6 倍，推理成本最高可降低 83%。該研究還揭示了新架構的 Scaling Law，證明其不僅具備優異的 Scaling 特性，更在效能上超越了 MoE。

實驗結果表明，訓練規模達 2000 萬 value 的 UltraMem 模型，在同等計算資源下，可同時實現業界領先的推理速度和模型效能，為構建數十億規模 value 或 expert 開闢了新路徑。

目前，該論文已被 ICLR 2025 接收。

論文標題：

Ultra-Sparse Memory Network

論文連結：

https://arxiv.org/abs/2411.12364

MoE與PKM的侷限性

LLM 的能力增長需要指數級增長的計算資源，這在即時應用等資源有限的環境中頗具挑戰。為了解決計算問題，先前的研究者提出了 MoE 和 Product Key Memory（PKM）方案，但它們都有各自的侷限性。

MoE 透過稀疏啟用 expert 解耦了計算和引數，但在推理場景中，速度其實很慢。原因在於，模型在推理時只能一個字一個字的生成，因此 batch size 和 sequence length 都很小，在這個場景下，MoE 的所有專家通常會被全部訪問到，極其容易遇到訪存瓶頸，導致推理延遲激增。

PKM 最早提出 large memory layer，其中包含了數量龐大的稀疏引數 value，這裡 value 其實就是一個向量，每個 token 會根據一個「行路由」和一個「列路由」定位到得分最高的幾個 value，啟用這些 value 後做 weighted sum pooling 作為 memory layer 的輸出。這種方法因為每個 token 在推理時僅僅只啟用極少數的 value，所以推理時不會遇到訪存瓶頸，但其效果很差，且 scaling 能力差。

UltraMem設計方案：兼顧訪存和模型效果

UltraMem 參考了 PKM 的設計，但針對 PKM 的 3 個缺陷予以補充，以實現更高效的訪存、更優質的 value 檢索，同時，降低了視訊記憶體和部署成本。

1. 最佳化模型結構

PKM 的設計中，memory layer 只有 1 層，插在整個 Transformer 的中間層，這對大規模訓練並不友好，並且如此龐大的稀疏引數應該儘可能多的參與到每次的殘差連線中。

因此，研究團隊拆分出多個小 memory layer，以固定的間隔分佈在 transformer layer 中；並且增加了 skip-layer 的操作，即當前層的 memory layer 的輸出會加到後面某層 transformer layer 的輸出。這使得模型可以並行地執行 memory layer 的訪存操作和 transformer layer 的計算。

2. 最佳化 value 檢索方式

在檢索時，只有 score 最高的 m 個 value 會被啟用，PKM 的 score 是透過「行 score」+「列 score」得到的。

團隊進一步探索了一種更復雜的乘法方法 Tucker Decomposed Query-Key Retrieval（TDQKR）。這一方法受啟發於 Tucker Decomposition。具體來看，給定 values，shape 為（n，n，h），其中 h 為 hidden size，那麼 values 的 score 可以做如下分解：

其中, 是可學習的 tucker core。這個結構下，每個 value 的 score 由 r 個行 score 和 r 個列 score 的組合乘加獲得，具備更高的複雜度。

3. 隱式擴充套件稀疏引數

更多的稀疏引數通常會帶來更好的效果，但過多的引數又會給視訊記憶體和部署帶來麻煩。為此，研究團隊提出了 Implicit Value Expansion （IVE）方法隱式地擴充套件稀疏引數，並引入了 virtual memory 和 physical memory 的概念。

以 4 倍擴充套件為例（如下圖所示），virtual memory 的數量是 physical memory 的 4 倍，給定多對（score， index）後，首先按照 virtual memory address table 做查表，4 個 virtual block 會查詢同一個 physical memory table，之後各自做 weighted sum pooling，並經過不同的線性層，最後再求和輸出。

由於最後的 Linear 和取 value 之間沒有任何非線性操作，因此每個 Linear 都可以和 physical memory table 做融合，生成一個全新的 memory table，這個例子下，實際上隱式擴充套件了 4 倍的 value 數量。

實驗結果：較MoE最高提速6倍

1. 模型效能評估

研究團隊在 151M、680M、1.6B 三個尺寸的啟用引數上做了廣泛的實驗，其中 MoE、PKM 和 UltraMem 的總稀疏引數保持在啟用引數的 12 倍。

如下表所示，可以發現 UltraMem 在 680M、1.6B 上具有顯著的效果優勢。

隨著稀疏引數的增加，UltraMem 的效果和推理速度如何變化？

下圖（b）展示了 UltraMem 的效果變化，橫軸為稀疏引數和稠密引數的比值，每個顏色的線代表了一種稀疏度。稀疏度定義為 value 的數量 / 每個 token 啟用的 value 數量。觀察發現，持續增加稀疏引數和 loss 的下降呈對數關係；且稀疏度越小，模型效果越好；但是稀疏度持續降低帶來的收益在逐漸飽和。

下圖（c）展示了 UltraMem 的推理時間變化，橫軸為稀疏引數和稠密引數的比值。觀察發現，UltraMem 在持續增加稀疏引數時推理時間幾乎不變，反觀 MoE 有顯著增長的趨勢。

2. 消融實驗

研究團隊在 151M 啟用、1.5B 總引數的稀疏模型上進行了全面的消融實驗。從最原始的 PKM 開始，逐漸增加一些 trick 和上文提出的結構改進，最終能拿到 C4 validation loss -0.092 的顯著收益，同時稀疏引數和計算量幾乎不變。

綜上所述，研究團隊提出的 UltraMem 具有極小的訪存，因此，相比 MoE 實現了最高達 6 倍的速度提升，推理成本最高可降低 83%。同時，在效能方面，隨著模型容量的增加，在相同的引數和計算量情況下，UltraMem 超過了 MoE，表明其具有更強的擴充套件能力。這項工作為開發更高效和可擴充套件的語言模型提供了一個有希望的方向。