推理成本比MoE直降83％！位元組最新大模型架構收錄ICLR2025

2025-05-30 17:03 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【Mamba/多模態/擴散】交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

豆包大模型團隊 投稿

轉載自：量子位（QbitAI）
位元組出了個全新架構，把推理成本給狠狠地打了下去！
有多狠？
推理速度相比MoE架構提升2-6倍，推理成本最高可降低83%。
論文：https://arxiv.org/abs/2411.12364
這個全新的稀疏模型架構叫做UltraMem，有效地解決了目前主流的MoE架構和PKM架構所存在的侷限性。
例如MoE在做推理時，較小的batch size會啟用全部專家，導致訪存急劇上升，推理延遲增加；而PKM雖然減少了訪存開銷，但效果較差且擴充套件能力有限。
實驗結果表明，訓練規模達2000 萬value的UltraMem模型，在同等計算資源下，可同時實現業界領先的推理速度和模型效能，為構建數十億規模value或expert開闢了新路徑。
這項研究目前已經被ICLR 2025接收。
那麼UltraMem架構具體是如何做到這點的呢？我們繼續往下看。

兼顧訪存和模型效果

隨著大語言模型能力的提升，所需的計算資源呈指數級增長，這在即時應用等資源受限的環境中帶來了巨大挑戰。
為了應對計算資源的瓶頸，研究者們提出了多種解決方案，其中包括MoE和Product Key Memory (PKM)。
然而，這些方法各自存在一定的侷限性。
MoE 透過稀疏啟用專家（expert）的方式，將計算與引數解耦，從而在訓練時能夠有效減少計算量；但在推理場景中，MoE 的表現卻不盡如人意。
由於推理時模型需要逐字生成輸出，batch size和sequence length通常較小，而通常情況下，數量較少的token就可以啟用幾乎所有的專家。
這種全專家啟用的模式使得訪存需求急劇增加，進而引發嚴重的訪存瓶頸，最終導致推理延遲大幅上升。
另一方面，PKM提出了 large memory layer的概念，其中包含了大量的稀疏引數 value（每個 value 實際上是一個向量）。
在推理時，每個 token 會透過“行路由”和“列路由”定位到得分最高的幾個 value，並對其進行加權求和（weighted sum pooling），作為 memory layer 的輸出。
由於每個token在推理時僅啟用極少數的value，PKM能夠有效避免訪存瓶頸。然而，儘管PKM在訪存效率上表現優異，其模型效果卻較差，且擴充套件能力（scaling）有限，難以應對更大規模的模型需求。
總的來說，MoE和PKM雖然在一定程度上緩解了計算資源的問題，但它們在推理效率、模型效果和擴充套件能力等方面仍存在明顯不足，亟需新的解決方案來進一步最佳化大模型的推理效能。
而UltraMem則是參考了PKM的設計，但針對PKM的3個缺陷予以補充，以實現更高效的訪存、更優質的value檢索，同時，降低了視訊記憶體和部署成本。

1、最佳化模型結構

PKM的設計中，memory layer只有1層，插在整個Transformer的中間層，這對大規模訓練並不友好，並且如此龐大的稀疏引數應該儘可能多的參與到每次的殘差連線中。
因此，研究團隊拆分出多個小memory layer，以固定的間隔分佈在 Transformer layer中；並且增加了skip-layer的操作，即當前層的memory layer的輸出會加到後面某層Transformer layer 的輸出。
這使得模型可以並行地執行memory layer的訪存操作和Transformer layer的計算。

2、最佳化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可以做如下分解：
其中Srow,Scol∈Rr×n，C∈Rr×r是可學習的tucker core。這個結構下，每個value的score由r個行score和r個列score的組合乘加獲得，具備更高的複雜度。

3、隱式擴充套件稀疏引數

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

較MoE最高提速6倍

1、模型效能評估

研究團隊在151M、680M、1.6B三個尺寸的啟用引數上做了廣泛的實驗，其中MoE、PKM和UltraMem的總稀疏引數保持在啟用引數的12倍。
如下表所示，可以發現 UltraMem在680M、1.6B上具有顯著的效果優勢。
隨著稀疏引數的增加，UltraMem的效果和推理速度如何變化？
  • 下圖（b）展示了UltraMem的效果變化，橫軸為稀疏引數和稠密引數的比值，每個顏色的線代表了一種稀疏度。稀疏度定義為value的數量 / 每個token啟用的value數量。觀察發現，持續增加稀疏引數和loss的下降呈對數關係；且稀疏度越小，模型效果越好；但是稀疏度持續降低帶來的收益在逐漸飽和。
  • 下圖（c）展示了UltraMem的推理時間變化，橫軸為稀疏引數和稠密引數的比值。觀察發現，UltraMem在持續增加稀疏引數時推理時間幾乎不變，反觀MoE有顯著增長的趨勢。

2、消融實驗

研究團隊在151M啟用、1.5B總引數的稀疏模型上進行了全面的消融實驗。
從最原始的PKM開始，逐漸增加一些 trick 和上文提出的結構改進，最終能拿到C4 validation loss -0.092的顯著收益，同時稀疏引數和計算量幾乎不變。
綜上所述，研究團隊提出的UltraMem具有極小的訪存，因此，相比 MoE 實現了最高達6倍的速度提升，推理成本最高可降低83%。
同時，在效能方面，隨著模型容量的增加，在相同的引數和計算量情況下，UltraMem超過了MoE，表明其具有更強的擴充套件能力。
這項工作可以說是為開發更高效和可擴充套件的語言模型提供了一個有希望的方向。

何愷明在MIT授課的課件PPT下載


在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

ECCV 2024 論文和程式碼下載


在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集

CVPR 2024 論文和程式碼下載


在CVer公眾號後臺回覆：CVPR2024，即可下載CVPR 2024論文和程式碼開源的論文合集
Mamba、多模態和擴散模型交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-Mamba、多模態學習或者擴散模型微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者擴散模型+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）來了！想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料，歡迎掃描下方二維碼，加入CVer計算機視覺（知識星球），已彙集上萬人！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請贊和在看


相關文章