微軟 & 清華最新研究，打破 GPT 系列開創的 Decoder-Only 架構——
提出 Decoder-Decoder 新型架構，名爲 YOCO（You Only Cache Once）。
YOCO僅緩存一次鍵值對，可大幅降低 GPU 内存需求，且保留全局注意力能力。
一張圖來看 YOCO 和标準 Transformer 的比較。
在處理 512K 上下文長度時，标準 Transformer 内存使用是 YOCO 的 6.4 倍，預填充延遲是 YOCO 的 30.3 倍，而 YOCO 的吞吐量提升到标準 Transformer 的 9.6 倍。
去年一張 " 大語言模型進化樹 " 動圖在學術圈瘋轉，模型架構還隻有三大類：Decoder-Only、Encoder-Only、Encoder-Decoder。
那麽這個新出的 Decoder-Decoder 架構到底長啥樣？
嗯，如網友所言，要讀的論文又增加了。
話不多說，一起來看。
打破 Decoder-Only
YOCO 整體架構設計如下，分爲自解碼器（Self-Decoder）和交叉解碼器（Cross-Decoder）兩部分。
具體來說，YOCO 由 L 個塊堆疊而成，其中前 L/2 層是自解碼器，其餘模塊是交叉解碼器。
自解碼器利用高效自注意力（efficient self-attention）機制來獲取鍵值（KV）緩存：
接收輸入序列的嵌入表示，并使用高效自注意力來生成中間向量表示；使用因果掩碼（causal masking）保證解碼的自回歸特性；自解碼器的輸出用于生成全局 KV 緩存。
而交叉解碼器使用交叉注意力（cross-attention）來重用自解碼器生成的共享 KV 緩存：
在自解碼器生成的 KV 緩存基礎上進行堆疊，以獲得最終的輸出向量；同樣使用因果掩碼來維持自回歸生成；允許交叉解碼器層間高效地重用 KV 緩存，減少了對 GPU 内存的需求。
總的來說，自解碼器和交叉解碼器的模塊設計與 Transformer 的解碼器層類似，包含交錯注意力和前饋網絡子層。不過，研究人員還進行了預 RMSNorm、SwiGLU 和分組查詢注意力等改進。
兩部分之間的區别在于注意力模塊。
自解碼器使用高效自注意力，如滑動窗口注意力（Sliding-Window Attention）或門控保留（gated retention）。
而交叉解碼器使用标準的多頭交叉注意力，Query 向量通過注意力與自解碼器産生的全局鍵值緩存相關聯。
推理大幅度省 省 省
實驗階段，研究人員将 YOCO 模型與同體量的 Transformer 模型進行比較。
分析維度有四個：語言建模評估、與 Transformer 比較的可擴展性、長上下文評估、推理優勢。
語言建模評估
研究人員訓練了一個 3B 參數的 YOCO 語言模型，并根據訓練 token 數量（1T 和 1.6T）進行評估。
在 LM Eval Harness 的多個下遊任務上，YOCO 與 Transformer 模型 OpenLLaMA-3B-v2、StableLM-base-alpha-3B-v2、StableLM-3B-4E1T 打得有來有回。
可擴展性對比
接着，研究人員在 160M 到 13B 參數規模範圍内，分别訓練了 YOCO（門控保留和滑動窗口注意力版本）和 Transformer 語言模型。
對比了它們在驗證集上的語言模型損失，YOCO 的表現與 Transformer 基本持平：
結果證明 YOCO 在模型大小擴展方面具有很強的可擴展性。
長上下文評估
将 3B 的 YOCO 模型擴展到上下文爲 1M，在 " 大海撈針 " 等長序列的 needle retrieval 任務上，YOCO-3B-1M 的準确率接近 100%。
在多針檢索任務上，YOCO-3B-1M 的性能優于一些超 3B 的 Transformer 模型：
此外，YOCO 模型在長序列上的 NLL 随着上下文長度的增加而一緻下降，表明 YOCO 能夠有效地利用長距離依賴信息進行語言建模：
綜上，可見 YOCO 在性能上完全不輸 Transformer，關鍵來看 YOCO 在推理效率上取得的顯著提升。
推理優勢
研究人員評估了 YOCO 在 GPU 内存占用、prefilling 延遲、吞吐量和服務容量等方面的優勢，評估上下文範圍爲 32K 至 1M。
如下圖所示，與 Transformer 相比，YOCO 大幅度降低了 GPU 内存占用，且 YOCO 的内存消耗随上下文長度增加，增長幅度很小。
例如，在 1M 長度下，整體推理内存使用量僅爲 12.4GB，而傳統的 Transformer 則占用了 9.38 倍的 GPU 内存。
下面展示了 token 的 KV 緩存對 GPU 内存的占用情況。
YOCO 模型隻緩存一層全局的鍵值對，因此與 Transformer 模型相比，它需要的内存約少了 L（指模型的層數）倍。
例如，YOCO 模型可以使用 1GB 的 GPU 内存來處理 128K token。而具有 GQA 的 Transformer 65B 大小模型，僅能支持 1.6K token。
也就是說，模型越大，YOCO 可以節省更多。
在預填充階段，模型并行編碼輸入 token。對于 512K 和 1M 長度的輸入，Transformer 分别需要大約 180 秒和 300 秒。Transformer 的計算複雜度爲 O ( N^2 ) ，處理長上下文需要大量的浮點運算操作。
相比之下，YOCO 的預填充時間爲 O ( N ) ，随序列長度線性增長。
YOCO 将 Transformer 的 512K 上下文預填充時間從 180 秒減少到不到 6 秒。
預填充階段可以在進入交叉解碼器之前提前退出。因此，即使對于短上下文，預填充延遲的加速至少是兩倍。例如，對于 32K 長度，YOCO 比 Transformer 快 2.87 倍。
吞吐量表示模型每秒可以處理多少個 token，涵蓋了預填充和生成時間。如下圖所示，與 Transformer 相比，YOCO 在不同上下文長度下實現了更高的吞吐量。
以 512K 查詢爲例，Transformer 的吞吐量爲 4.5 token/ 秒，而 YOCO 達到了 43.1token/ 秒，即實現了 9.6 倍的加速。
吞吐量提高的原因如前所述，YOCO 減少了預填充所需的時間。其次，由于内存消耗減少，因此可以在推理時使用更大的批量大小，這也有助于提高吞吐量。
詳細細節，感興趣的家人們可以查看原論文。
論文鏈接：https://arxiv.org/abs/2405.05254
— 完 —
點這裏關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
