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

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

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

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

基于此，谷歌引入了一種有效的方法，其關鍵組成部分是一種稱為 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 提高了摘要性能指標。

本文轉自 機器之心 ，作者：機器之心

原文鏈接:??https://mp.weixin.qq.com/s/4cOsPWWDLSMpnQcbUc8Wxg??