首頁(yè) > 科技要聞 > 科技> 正文

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

量子位 整合編輯：太平洋科技 發(fā)布于：2025-02-13 16:40

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

有多狠？

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

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

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

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

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

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

兼顧訪存和模型效果

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

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

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

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

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

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

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

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

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

總的來(lái)說(shuō)，MoE和PKM雖然在一定程度上緩解了計(jì)算資源的問(wèn)題，但它們?cè)谕评硇�、模型效果和擴(kuò)展能力等方面仍存在明顯不足，亟需新的解決方案來(lái)進(jìn)一步優(yōu)化大模型的推理性能。

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

1、優(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ì)算。

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

在檢索時(shí)，只有score最高的m個(gè)value會(huì)被激活，PKM的score是通過(guò)“行score”+“列score”得到的。

團(tuán)隊(duì)進(jìn)一步探索了一種更復(fù)雜的乘法方法Tucker Decomposed Query-Key Retrieval（TDQKR）。

這一方法受啟發(fā)于Tucker Decomposition。具體來(lái)看，給定values，shape為（n，n，h），其中h為hidden size，那么values的score S_grid可以做如下分解：

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

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

更多的稀疏參數(shù)通常會(huì)帶來(lái)更好的效果，但過(guò)多的參數(shù)又會(huì)給顯存和部署帶來(lái)麻煩。

為此，研究團(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)過(guò)不同的線性層，最后再求和輸出。

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

較MoE最高提速6倍

1、模型性能評(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ù)降低帶來(lái)的收益在逐漸飽和。

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

2、消融實(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超過(guò)了MoE，表明其具有更強(qiáng)的擴(kuò)展能力。

這項(xiàng)工作可以說(shuō)是為開發(fā)更高效和可擴(kuò)展的語(yǔ)言模型提供了一個(gè)有希望的方向。

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

本文來(lái)源：量子位

網(wǎng)友評(píng)論

聚超值•精選

推薦 手機(jī) 筆記本 影像 硬件 家居 商用 企業(yè) 出行 未來(lái)
    • 二維碼 回到頂部