微軟 & 清華最新研究，打破 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)槔琓ransformer 的吞吐量為 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

廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。