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

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

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

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

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

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

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

Ultra-Sparse Memory Network

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

1. MoE 與 PKM 的局限性 

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

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

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

2. UltraMem 兼顧訪存和效果 

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

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

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

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

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

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

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

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

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

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

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

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

① 模型性能評(píng)估

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

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

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

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

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

② 消融實(shí)驗(yàn)

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

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

4. 寫在最后 

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

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