直接擴展到無限長，GoogleInfini-Transformer終結上下文長度之爭

不知 Gemini 1.5 Pro 是否用到了這項技術。

Google又放大招了，發布下一代 Transformer 模型 Infini-Transformer。

Infini-Transformer 引入了一種有效的方法，可以將基於 Transformer 的大型語言模型 (LLM) 擴展到無限長輸入，而不增加記憶體和計算需求。使用該技術，研究者成功將一個 1B 的模型上下文長度提高到 100 萬；應用到 8B 模型上，模型能處理 500K 的書籍摘要任務。

自 2017 年開創性研究論文《Attention is All You Need》問世以來，Transformer 架構就一直主導著生成式人工智慧領域。而Google對 Transformer 的優化設計最近比較頻繁，幾天前，他們更新了 Transformer 架構，發布 Mixture-of-Depths（MoD），改變了以往 Transformer 計算模式。沒過幾天，谷歌又放了這項新研究。

專注 AI 領域的研究者都了解記憶體的重要性，它是智慧的基石，可以為 LLM 提供高效的運算。然而，Transformer 和基於 Transformer 的 LLM 由於注意力機制的固有特性，即 Transformer 中的注意力機制在記憶體佔用和計算時間上都表現出二次複雜性。例如，對於批次大小為 512、上下文長度為 2048 的 500B 模型，注意力鍵 - 值 (KV) 狀態的記憶體佔用為 3TB。但事實上，標準 Transformer 架構有時需要將 LLM 擴展到更長的序列（如 100 萬 token），這就帶來巨大的記憶體開銷，並且隨著上下文長度的增加，部署成本也在增加。

基於此，Google引入了一種有效的方法，其關鍵組成部分是一種稱為 Infini-attention（無限注意力）的新註意力技術。不同於傳統的 Transformer 使用局部注意力丟棄舊片段，為新片段釋放記憶體空間。 Infini-attention 增加了壓縮記憶體（compressive memory），可以將使用後的舊片段儲存到壓縮記憶體中，輸出時會聚合當前上下文資訊以及壓縮記憶體中的信息，因而模型可以檢索完整的上下文歷史。

該方法使 Transformer LLM 在有限記憶體的情況下擴展到無限長上下文，並以流的方式處理極長的輸入進行計算。

實驗表明，該方法在長上下文語言建模基準測試中的效能優於基線，同時記憶體參數減少了 100 倍以上。當使用 100K 序列長度進行訓練時，模型實現了更好的困惑度。此外研究發現，1B 模型在 5K 序列長度的密鑰實例上進行了微調，解決了 1M 長度的問題。最後，論文展示了具有 Infini-attention 的 8B 模型經過持續的預訓練和任務微調，在 500K 長度的書籍摘要任務上達到了新的 SOTA 結果。

本文貢獻總結如下：

• #引入了一個實用且強大的注意力機制Infini-attention——具有長期壓縮記憶體和局部因果注意力，可用於有效地建模長期和短期上下文依賴關係；
• Infini-attention 對標準縮放點積注意力（ standard scaled dot-product attention）進行了最小的改變，並透過設計支持即插即用的持續預訓練和長上下文自適應；
• ##該方法使Transformer LLM 能夠透過流的方式處理極長的輸入，在有限的記憶體和計算資源下擴展到無限長的上下文。

• 論文連結：https://arxiv.org/pdf/2404.07143.pdf
• #論文標題：Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention

##### ###############Infini-attention 使Transformer LLM 能夠透過有限的記憶體佔用和計算有效地處理無限長的輸入。如下圖 1 所示，Infini-attention 將壓縮記憶融入到普通的注意力機制中，並在單一 Transformer 區塊中建立了掩碼局部注意力和長期線性注意力機制。 ###############對 Transformer 注意力層進行這種微妙但關鍵的修改可以透過持續的預訓練和微調將現有 LLM 的上下文視窗擴展到無限長。 ###############Infini-attention 採用標準注意力計算的所有鍵、值和查詢狀態，以進行長期記憶鞏固（memory consolidation）和檢索，並將注意力的舊KV 狀態儲存在壓縮記憶體中，而不是像標準注意力機制那樣丟棄它們。在處理後續序列時，Infini-attention 使用注意查詢狀態從記憶體中檢索值。為了計算最終的上下文輸出，Infini-attention 聚合了長期記憶檢索值和局部注意力上下文。

如下圖 2 所示，研究團隊比較了基於 Infini-attention 的 Infini-Transformer 和 Transformer-XL。與 Transformer-XL 類似，Infini-Transformer 對 segment 序列進行操作，並計算每個 segment 中的標準因果點積注意力上下文。因此，點積注意力計算在某種意義上是局部的。

然而，局部注意力在處理下一個segment 時會丟棄前一個segment 的注意力狀態，但Infini-Transformer 復用舊的KV 注意力狀態，以透過壓縮儲存來維護整個上下文歷史。因此，Infini-Transformer 的每個注意力層都具有全域壓縮狀態和局部細粒度狀態。

與多頭注意力（MHA）類似，除了點積注意力之外，Infini-attention 還為每個注意力層維護H 個並行壓縮記憶體（H是注意力頭的數量）。

下表 1 列出了幾種模型根據模型參數和輸入 segment 長度，定義的上下文記憶體佔用和有效上下文長度。 Infini-Transformer 支援具有有限記憶體佔用的無限上下文視窗。

實驗

#該研究在長上下文語言建模、長度為1M 的在密鑰上下文區塊檢索和500K 長度的書籍摘要任務上評估了Infini-Transformer 模型，這些任務具有極長的輸入序列。對於語言建模，研究者選擇從頭開始訓練模型，而對於金鑰和書籍摘要任務，研究者採用不斷預訓練 LLM 的方式，以證明 Infini-attention 即插即用的長上下文適應能力。

長上下文語言建模。表 2 結果顯示 Infini-Transformer 優於 Transformer-XL 和 Memorizing Transformers 基線，且與 Memorizing Transformer 模型相比，儲存參數減少了 114 倍。

金鑰任務。表 3 為 Infini-Transformer 在 5K 長度輸入上微調後，解決了高達 1M 上下文長度的金鑰任務。實驗中輸入 token 的範圍從 32K 到 1M，對於每個測試子集，研究者控制密鑰的位置，使其位於輸入序列的開頭、中間或結尾附近。實驗報告了零樣本準確率和微調準確率。在對 5K 長度輸入進行 400 個步驟的微調後，Infini-Transformer 解決了高達 1M 上下文長度的任務。

摘要任務。表 4 將 Infini-Transformer 與專門為摘要任務構建的編碼器 - 解碼器模型進行了比較。結果表明 Infini-Transformer 超越了先前最佳結果，並且透過處理書中的整個文本在 BookSum 上實現了新的 SOTA。 

研究者在圖 4 中也繪製了 BookSum 資料驗證分割的總體 Rouge 分數。根據折線趨勢表明，隨著輸入長度的增加，Infini-Transformers 提高了摘要效能指標。

以上是直接擴展到無限長，GoogleInfini-Transformer終結上下文長度之爭的詳細內容。更多資訊請關注PHP中文網其他相關文章！