推理成本比MoE直降83%！字節(jié)最新大模型架構(gòu)入圍ICLR 2025

十三 2025-02-12 11:51:30 來源：量子位

豆包大模型團隊 投稿

量子位 | 公眾號 QbitAI

字節(jié)出了個全新架構(gòu)，把推理成本給狠狠地打了下去！

有多狠？

推理速度相比MoE架構(gòu)提升2-6倍，推理成本最高可降低83%。

這個全新的稀疏模型架構(gòu)叫做UltraMem，有效地解決了目前主流的MoE架構(gòu)和PKM架構(gòu)所存在的局限性。

例如MoE在做推理時，較小的batch size會激活全部專家，導致訪存急劇上升，推理延遲增加；而PKM雖然減少了訪存開銷，但效果較差且擴展能力有限。

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

這項研究目前已經(jīng)被ICLR 2025接收。

那么UltraMem架構(gòu)具體是如何做到這點的呢？我們繼續(xù)往下看。

兼顧訪存和模型效果

隨著大語言模型能力的提升，所需的計算資源呈指數(shù)級增長，這在實時應用等資源受限的環(huán)境中帶來了巨大挑戰(zhàn)。

為了應對計算資源的瓶頸，研究者們提出了多種解決方案，其中包括MoE和Product Key Memory (PKM)。

然而，這些方法各自存在一定的局限性。

MoE 通過稀疏激活專家（expert）的方式，將計算與參數(shù)解耦，從而在訓練時能夠有效減少計算量；但在推理場景中，MoE 的表現(xiàn)卻不盡如人意。

由于推理時模型需要逐字生成輸出，batch size和sequence length通常較小，而通常情況下，數(shù)量較少的token就可以激活幾乎所有的專家。

這種全專家激活的模式使得訪存需求急劇增加，進而引發(fā)嚴重的訪存瓶頸，最終導致推理延遲大幅上升。

另一方面，PKM提出了 large memory layer的概念，其中包含了大量的稀疏參數(shù) value（每個 value 實際上是一個向量）。

在推理時，每個 token 會通過“行路由”和“列路由”定位到得分最高的幾個 value，并對其進行加權(quán)求和（weighted sum pooling），作為 memory layer 的輸出。

由于每個token在推理時僅激活極少數(shù)的value，PKM能夠有效避免訪存瓶頸。然而，盡管PKM在訪存效率上表現(xiàn)優(yōu)異，其模型效果卻較差，且擴展能力（scaling）有限，難以應對更大規(guī)模的模型需求。

總的來說，MoE和PKM雖然在一定程度上緩解了計算資源的問題，但它們在推理效率、模型效果和擴展能力等方面仍存在明顯不足，亟需新的解決方案來進一步優(yōu)化大模型的推理性能。

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

1、優(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的計算。

2、優(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可以做如下分解：

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

3、隱式擴展稀疏參數(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ù)量。

較MoE最高提速6倍

1、模型性能評估

研究團隊在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有顯著增長的趨勢。

2、消融實驗

研究團隊在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ā)更高效和可擴展的語言模型提供了一個有希望的方向。

論文地址：
https://arxiv.org/abs/2411.12364