財聯(lián)社2月12日訊，據(jù)豆包大模型團隊消息，字節(jié)跳動豆包大模型Foundation團隊近期提出UltraMem，一種同樣將計算和參數(shù)解耦的稀疏模型架構(gòu)，在保證模型效果的前提下解決了推理的訪存問題。據(jù)介紹，該架構(gòu)有效解決了MoE推理時高額的訪存問題，推理速度較MoE架構(gòu)提升2-6倍，推理成本最高可降低83%。

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

實驗結(jié)果表明，訓練規(guī)模達 2000 萬 value 的 UltraMem 模型，在同等計算資源下，可同時實現(xiàn)業(yè)界領(lǐng)先的推理速度和模型性能，為構(gòu)建數(shù)十億規(guī)模 value 或 expert 開辟了新路徑。

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

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

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

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

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

Ultra-Sparse Memory Network

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

1. MoE 與 PKM 的局限性

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

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

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

2. UltraMem 兼顧訪存和效果

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

① 優(yōu)化模型結(jié)構(gòu)

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

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

② 優(yōu)化 value 檢索方式

在檢索時，只有 score 最高的 m 個 value 會被激活，PKM 的 score 是通過「行 score」+「列 score」得到的。團隊進一步探索了一種更復雜的乘法方法 Tucker Decomposed Query-Key Retrieval（TDQKR）。這一方法受啟發(fā)于 Tucker Decomposition。具體來看，給定 values，shape 為（n，n，h），其中 h 為 hidden size，那么 values 的 score S_grid 可以做如下分解：

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

③ 隱式擴展稀疏參數(shù)

更多的稀疏參數(shù)通常會帶來更好的效果，但過多的參數(shù)又會給顯存和部署帶來麻煩。為此，研究團隊提出了 Implicit Value Expansion （IVE）方法隱式地擴展稀疏參數(shù)，并引入了 virtual memory 和 physical memory 的概念。

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

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

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

① 模型性能評估

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

如下表所示，可以發(fā)現(xiàn) UltraMem 在 680M、1.6B 上具有顯著的效果優(yōu)勢。

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

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

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

② 消融實驗

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

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

4. 寫在最后

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

當前，針對 UltraMem 的技術(shù)演進仍存在若干值得探索的方向，包括但不限于：如何高效優(yōu)化稀疏參數(shù)、如何提升稀疏模型推理能力、如何更優(yōu)地激活稀疏參數(shù)等。這些技術(shù)方向或?qū)⒊蔀楹罄m(xù)研究的重要切入點。