微軟&清華最新研究,打破GPT系列開(kāi)創(chuàng)的Decoder-Only架構(gòu)—— 提出Decoder-Decoder新型架構(gòu),名為YOCO(You Only Cache Once)。 YOCO僅緩存一次鍵值對(duì),可大幅降低GPU內(nèi)存需求,且保留全局注意力能力。 一張圖來(lái)看YOCO和標(biāo)準(zhǔn)Transformer的比較。 在處理512K上下文長(zhǎng)度時(shí),標(biāo)準(zhǔn)Transformer內(nèi)存使用是YOCO的6.4倍,預(yù)填充延遲是YOCO的30.3倍,而YOCO的吞吐量提升到標(biāo)準(zhǔn)Transformer的9.6倍。 去年一張“大語(yǔ)言模型進(jìn)化樹(shù)”動(dòng)圖在學(xué)術(shù)圈瘋轉(zhuǎn),模型架構(gòu)還只有三大類:Decoder-Only、Encoder-Only、Encoder-Decoder。 那么這個(gè)新出的Decoder-Decoder架構(gòu)到底長(zhǎng)啥樣? 嗯,如網(wǎng)友所言,要讀的論文又增加了。 話不多說(shuō),一起來(lái)看。 打破Decoder-Only YOCO整體架構(gòu)設(shè)計(jì)如下,分為自解碼器(Self-Decoder)和交叉解碼器(Cross-Decoder)兩部分。 具體來(lái)說(shuō),YOCO由L個(gè)塊堆疊而成,其中前L/2層是自解碼器,其余模塊是交叉解碼器。 自解碼器利用高效自注意力(efficient self-attention)機(jī)制來(lái)獲取鍵值(KV)緩存: 接收輸入序列的嵌入表示,并使用高效自注意力來(lái)生成中間向量表示;使用因果掩碼(causal masking)保證解碼的自回歸特性;自解碼器的輸出用于生成全局KV緩存。 而交叉解碼器使用交叉注意力(cross-attention)來(lái)重用自解碼器生成的共享KV緩存: 在自解碼器生成的KV緩存基礎(chǔ)上進(jìn)行堆疊,以獲得最終的輸出向量;同樣使用因果掩碼來(lái)維持自回歸生成;允許交叉解碼器層間高效地重用KV緩存,減少了對(duì)GPU內(nèi)存的需求。 總的來(lái)說(shuō),自解碼器和交叉解碼器的模塊設(shè)計(jì)與Transformer的解碼器層類似,包含交錯(cuò)注意力和前饋網(wǎng)絡(luò)子層。不過(guò),研究人員還進(jìn)行了預(yù)RMSNorm、SwiGLU和分組查詢注意力等改進(jìn)。 兩部分之間的區(qū)別在于注意力模塊。 自解碼器使用高效自注意力,如滑動(dòng)窗口注意力(Sliding-Window Attention)或門控保留(gated retention)。 而交叉解碼器使用標(biāo)準(zhǔn)的多頭交叉注意力,Query向量通過(guò)注意力與自解碼器產(chǎn)生的全局鍵值緩存相關(guān)聯(lián)。 推理大幅度省 省 省 實(shí)驗(yàn)階段,研究人員將YOCO模型與同體量的Transformer模型進(jìn)行比較。 分析維度有四個(gè):語(yǔ)言建模評(píng)估、與Transformer比較的可擴(kuò)展性、長(zhǎng)上下文評(píng)估、推理優(yōu)勢(shì)。 語(yǔ)言建模評(píng)估 研究人員訓(xùn)練了一個(gè)3B參數(shù)的YOCO語(yǔ)言模型,并根據(jù)訓(xùn)練token數(shù)量(1T和1.6T)進(jìn)行評(píng)估。 在LM Eval Harness的多個(gè)下游任務(wù)上,YOCO與Transformer模型OpenLLaMA-3B-v2、StableLM-base-alpha-3B-v2、StableLM-3B-4E1T打得有來(lái)有回。 可擴(kuò)展性對(duì)比 接著,研究人員在160M到13B參數(shù)規(guī)模范圍內(nèi),分別訓(xùn)練了YOCO(門控保留和滑動(dòng)窗口注意力版本)和Transformer語(yǔ)言模型。 對(duì)比了它們?cè)隍?yàn)證集上的語(yǔ)言模型損失,YOCO的表現(xiàn)與Transformer基本持平: 結(jié)果證明YOCO在模型大小擴(kuò)展方面具有很強(qiáng)的可擴(kuò)展性。 長(zhǎng)上下文評(píng)估 將3B的YOCO模型擴(kuò)展到上下文為1M,在“大海撈針”等長(zhǎng)序列的needle retrieval任務(wù)上,YOCO-3B-1M的準(zhǔn)確率接近100%。 在多針檢索任務(wù)上,YOCO-3B-1M的性能優(yōu)于一些超3B的Transformer模型: 此外,YOCO模型在長(zhǎng)序列上的NLL隨著上下文長(zhǎng)度的增加而一致下降,表明YOCO能夠有效地利用長(zhǎng)距離依賴信息進(jìn)行語(yǔ)言建模: 綜上,可見(jiàn)YOCO在性能上完全不輸Transformer,關(guān)鍵來(lái)看YOCO在推理效率上取得的顯著提升。 推理優(yōu)勢(shì) 研究人員評(píng)估了YOCO在GPU內(nèi)存占用、prefilling延遲、吞吐量和服務(wù)容量等方面的優(yōu)勢(shì),評(píng)估上下文范圍為32K至1M。 如下圖所示,與Transformer相比,YOCO大幅度降低了GPU內(nèi)存占用,且YOCO的內(nèi)存消耗隨上下文長(zhǎng)度增加,增長(zhǎng)幅度很小。 例如,在1M長(zhǎng)度下,整體推理內(nèi)存使用量?jī)H為12.4GB,而傳統(tǒng)的Transformer則占用了9.38倍的GPU內(nèi)存。 下面展示了token的KV緩存對(duì)GPU內(nèi)存的占用情況。 YOCO模型只緩存一層全局的鍵值對(duì),因此與Transformer模型相比,它需要的內(nèi)存約少了L(指模型的層數(shù))倍。 例如,YOCO模型可以使用1GB的GPU內(nèi)存來(lái)處理128K token。而具有GQA的Transformer 65B大小模型,僅能支持1.6K token。 也就是說(shuō),模型越大,YOCO可以節(jié)省更多。 在預(yù)填充階段,模型并行編碼輸入token。對(duì)于512K和1M長(zhǎng)度的輸入,Transformer分別需要大約180秒和300秒。Transformer的計(jì)算復(fù)雜度為O(N^2),處理長(zhǎng)上下文需要大量的浮點(diǎn)運(yùn)算操作。 相比之下,YOCO的預(yù)填充時(shí)間為O(N),隨序列長(zhǎng)度線性增長(zhǎng)。 YOCO將Transformer的512K上下文預(yù)填充時(shí)間從180秒減少到不到6秒。 預(yù)填充階段可以在進(jìn)入交叉解碼器之前提前退出。因此,即使對(duì)于短上下文,預(yù)填充延遲的加速至少是兩倍。例如,對(duì)于32K長(zhǎng)度,YOCO比Transformer快2.87倍。 吞吐量表示模型每秒可以處理多少個(gè)token,涵蓋了預(yù)填充和生成時(shí)間。如下圖所示,與Transformer相比,YOCO在不同上下文長(zhǎng)度下實(shí)現(xiàn)了更高的吞吐量。 以512K查詢?yōu)槔,Transformer的吞吐量為4.5 token/秒,而YOCO達(dá)到了43.1token/秒,即實(shí)現(xiàn)了9.6倍的加速。 吞吐量提高的原因如前所述,YOCO減少了預(yù)填充所需的時(shí)間。其次,由于內(nèi)存消耗減少,因此可以在推理時(shí)使用更大的批量大小,這也有助于提高吞吐量。 詳細(xì)細(xì)節(jié),感興趣的家人們可以查看原論文。 論文鏈接:https://arxiv.org/abs/2405.05254 本文來(lái)源:量子位 |
原創(chuàng)欄目
IT百科
網(wǎng)友評(píng)論
聚超值•精選
在經(jīng)歷了2018-2020年“蔚小理”的上市潮后,隨著國(guó)內(nèi)新能源汽車市場(chǎng)競(jìng)爭(zhēng)日益激烈,以及全球EV需求放緩,IPO的大門正在向造車新勢(shì)力們緩緩關(guān)閉。極氪的成功上市既是對(duì)新勢(shì)力的一次激勵(lì),也是一次警示——后來(lái)者必須面對(duì)越來(lái)越嚴(yán)苛的上市條件。留給哪吒汽車、廣汽埃安們的機(jī)會(huì)可能不多了。