隨著模型規模的擴大，推理成本和訪存效率已成為限制大模型規模應用的關鍵瓶頸。近期，字節跳動豆包大模型團隊提出了全新的稀疏模型架構 UltraMem，該架構有效解決了 MoE 推理時高額的訪存問題，推理速度較 MoE 架構提升 2-6 倍，推理成本最高可降低 83%。該研究還揭示了新架構的 Scaling Law，證明其不僅具備優異的 Scaling 特性，更在性能上超越了 MoE。

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

在 Transformer 架構下，模型的性能與其參數數量和計算復雜度呈對數關系。隨著 LLM 規模不斷增大，推理成本會急劇增加，速度變慢。

盡管 MoE 架構已經成功將計算和參數解耦，但在推理時，較小的 batch size 就會激活全部專家，導致訪存急劇上升，進而使推理延遲大幅增加。

為此，字節跳動豆包大模型 Foundation 團隊提出 UltraMem，一種同樣將計算和參數解耦的稀疏模型架構，在保證模型效果的前提下解決了推理的訪存問題。

實驗結果表明，在參數和激活條件相同的情況下，UltraMem 在模型效果上超越了 MoE，并將推理速度提升了 2-6 倍。此外，在常見 batch size 規模下，UltraMem 的訪存成本幾乎與同計算量的 Dense 模型相當。

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

Ultra-Sparse Memory Network

論文鏈接：https://arxiv.org/abs/2411.12364

1. 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 能力差。

2. UltraMem 兼顧訪存和效果 

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

① 優化模型結構

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

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

② 優化 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 S_grid 可以做如下分解：

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

③ 隱式擴展稀疏參數

更多的稀疏參數通常會帶來更好的效果，但過多的參數又會給顯存和部署帶來麻煩。為此，研究團隊提出了 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 數量。

3. 實驗結果：推理速度較 MoE 最高提升 6 倍 

① 模型性能評估

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

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

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

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

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

② 消融實驗

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

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

4. 寫在最后 

UltraMem 能有效地應用于對延遲要求較高的推理場景（例如代碼補全），避免了類似 MoE 的訪存瓶頸。即使在通用場景下，UltraMem 相較于 MoE 都展現出了顯著的速度優勢，除非 batch size 上萬的極端場景。

當前，針對 UltraMem 的技術演進仍存在若干值得探索的方向，包括但不限于：如何高效優化稀疏參數、如何提升稀疏模型推理能力、如何更優地激活稀疏參數等。這些技術方向或將成為后續研究的重要切入點。