Google 推出 Infini-Transformer模型，可以讓大型語言模型無限長輸入，而不增加記憶體和運算需求

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 運算模式。沒過幾天，Google又放出了這項新研究。

延伸閱讀：就是要你全部的注意力！ChatGPT不小心找到一個有史以來最賺錢的商業模式
延伸閱讀：8 名Google員工發明了現代人工智慧，這是那篇論文的內幕故事

專注 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 提高了摘要性能指標。

延伸閱讀：誰發表的AI研究最具影響力？Google遙遙領先，OpenAI成果轉化率完勝DeepMind