直接擴(kuò)展到無限長，谷歌Infini-Transformer終結(jié)上下文長度之爭 精華
谷歌又放大招了，發(fā)布下一代 Transformer 模型 Infini-Transformer。
Infini-Transformer 引入了一種有效的方法，可以將基于 Transformer 的大型語言模型 (LLM) 擴(kuò)展到無限長輸入，而不增加內(nèi)存和計算需求。使用該技術(shù)，研究者成功將一個 1B 的模型上下文長度提高到 100 萬；應(yīng)用到 8B 模型上，模型能處理 500K 的書籍摘要任務(wù)。
自 2017 年開創(chuàng)性研究論文《Attention is All You Need》問世以來，Transformer 架構(gòu)就一直主導(dǎo)著生成式人工智能領(lǐng)域。而谷歌對 Transformer 的優(yōu)化設(shè)計最近比較頻繁，幾天前，他們更新了 Transformer 架構(gòu)，發(fā)布 Mixture-of-Depths（MoD），改變了以往 Transformer 計算模式。沒過幾天，谷歌又放出了這項新研究。
專注 AI 領(lǐng)域的研究者都了解內(nèi)存的重要性，它是智能的基石，可以為 LLM 提供高效的計算。然而，Transformer 和基于 Transformer 的 LLM 由于注意力機(jī)制的固有特性，即 Transformer 中的注意力機(jī)制在內(nèi)存占用和計算時間上都表現(xiàn)出二次復(fù)雜性。例如，對于批大小為 512、上下文長度為 2048 的 500B 模型，注意力鍵 - 值 (KV) 狀態(tài)的內(nèi)存占用為 3TB。但事實上，標(biāo)準(zhǔn) Transformer 架構(gòu)有時需要將 LLM 擴(kuò)展到更長的序列（如 100 萬 token），這就帶來巨大的內(nèi)存開銷，并且隨著上下文長度的增加，部署成本也在增加。
基于此，谷歌引入了一種有效的方法，其關(guān)鍵組成部分是一種稱為 Infini-attention（無限注意力）的新注意力技術(shù)。不同于傳統(tǒng)的 Transformer 使用局部注意力丟棄舊片段，為新片段釋放內(nèi)存空間。Infini-attention 增加了壓縮內(nèi)存（compressive memory），可以將使用后的舊片段存儲到壓縮內(nèi)存中，輸出時會聚合當(dāng)前上下文信息以及壓縮內(nèi)存中的信息，因而模型可以檢索完整的上下文歷史。
該方法使 Transformer LLM 在有限內(nèi)存的情況下擴(kuò)展到無限長上下文，并以流的方式處理極長的輸入進(jìn)行計算。
實驗表明，該方法在長上下文語言建?；鶞?zhǔn)測試中的性能優(yōu)于基線，同時內(nèi)存參數(shù)減少了 100 倍以上。當(dāng)使用 100K 序列長度進(jìn)行訓(xùn)練時，該模型實現(xiàn)了更好的困惑度。此外該研究發(fā)現(xiàn)，1B 模型在 5K 序列長度的密鑰實例上進(jìn)行了微調(diào)，解決了 1M 長度的問題。最后，論文展示了具有 Infini-attention 的 8B 模型經(jīng)過持續(xù)的預(yù)訓(xùn)練和任務(wù)微調(diào)，在 500K 長度的書籍摘要任務(wù)上達(dá)到了新的 SOTA 結(jié)果。
本文貢獻(xiàn)總結(jié)如下：
- 引入了一種實用且強大的注意力機(jī)制 Infini-attention—— 具有長期壓縮內(nèi)存和局部因果注意力，可用于有效地建模長期和短期上下文依賴關(guān)系；
- Infini-attention 對標(biāo)準(zhǔn)縮放點積注意力（ standard scaled dot-product attention）進(jìn)行了最小的改變，并通過設(shè)計支持即插即用的持續(xù)預(yù)訓(xùn)練和長上下文自適應(yīng)；
- 該方法使 Transformer LLM 能夠通過流的方式處理極長的輸入，在有限的內(nèi)存和計算資源下擴(kuò)展到無限長的上下文。
- 論文鏈接：https://arxiv.org/pdf/2404.07143.pdf
- 論文標(biāo)題：Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention?
方法介紹
Infini-attention 使 Transformer LLM 能夠通過有限的內(nèi)存占用和計算有效地處理無限長的輸入。如下圖 1 所示，Infini-attention 將壓縮記憶融入到普通的注意力機(jī)制中，并在單個 Transformer 塊中構(gòu)建了掩碼局部注意力和長期線性注意力機(jī)制。
對 Transformer 注意力層進(jìn)行這種微妙但關(guān)鍵的修改可以通過持續(xù)的預(yù)訓(xùn)練和微調(diào)將現(xiàn)有 LLM 的上下文窗口擴(kuò)展到無限長。
Infini-attention 采用標(biāo)準(zhǔn)注意力計算的所有鍵、值和查詢狀態(tài)，以進(jìn)行長期記憶鞏固（memory consolidation）和檢索，并將注意力的舊 KV 狀態(tài)存儲在壓縮內(nèi)存中，而不是像標(biāo)準(zhǔn)注意力機(jī)制那樣丟棄它們。在處理后續(xù)序列時，Infini-attention 使用注意查詢狀態(tài)從內(nèi)存中檢索值。為了計算最終的上下文輸出，Infini-attention 聚合了長期記憶檢索值和局部注意力上下文。
如下圖 2 所示，研究團(tuán)隊比較了基于 Infini-attention 的 Infini-Transformer 和 Transformer-XL。與 Transformer-XL 類似，Infini-Transformer 對 segment 序列進(jìn)行操作，并計算每個 segment 中的標(biāo)準(zhǔn)因果點積注意力上下文。因此，點積注意力計算在某種意義上是局部的。
然而，局部注意力在處理下一個 segment 時會丟棄前一個 segment 的注意力狀態(tài)，但 Infini-Transformer 復(fù)用舊的 KV 注意力狀態(tài)，以通過壓縮存儲來維護(hù)整個上下文歷史。因此，Infini-Transformer 的每個注意力層都具有全局壓縮狀態(tài)和局部細(xì)粒度狀態(tài)。
與多頭注意力（MHA）類似，除了點積注意力之外，Infini-attention 還為每個注意力層維護(hù) H 個并行壓縮內(nèi)存（H 是注意力頭的數(shù)量）。
下表 1 列出了幾種模型根據(jù)模型參數(shù)和輸入 segment 長度，定義的上下文內(nèi)存占用和有效上下文長度。Infini-Transformer 支持具有有限內(nèi)存占用的無限上下文窗口。
實驗
該研究在長上下文語言建模、長度為 1M 的密鑰上下文塊檢索和 500K 長度的書籍摘要任務(wù)上評估了 Infini-Transformer 模型，這些任務(wù)具有極長的輸入序列。對于語言建模，研究者選擇從頭開始訓(xùn)練模型，而對于密鑰和書籍摘要任務(wù)，研究者采用不斷預(yù)訓(xùn)練 LLM 的方式，以證明 Infini-attention 即插即用的長上下文適應(yīng)能力。
長上下文語言建模。表 2 結(jié)果表明 Infini-Transformer 優(yōu)于 Transformer-XL 和 Memorizing Transformers 基線，并且與 Memorizing Transformer 模型相比，存儲參數(shù)減少了 114 倍。
密鑰任務(wù)。表 3 為 Infini-Transformer 在 5K 長度輸入上進(jìn)行微調(diào)后，解決了高達(dá) 1M 上下文長度的密鑰任務(wù)。實驗中輸入 token 的范圍從 32K 到 1M，對于每個測試子集，研究者控制密鑰的位置，使其位于輸入序列的開頭、中間或結(jié)尾附近。實驗報告了零樣本準(zhǔn)確率和微調(diào)準(zhǔn)確率。在對 5K 長度輸入進(jìn)行 400 個步驟的微調(diào)后，Infini-Transformer 解決了高達(dá) 1M 上下文長度的任務(wù)。
摘要任務(wù)。表 4 將 Infini-Transformer 與專門為摘要任務(wù)構(gòu)建的編碼器 - 解碼器模型進(jìn)行了比較。結(jié)果表明 Infini-Transformer 超越了之前最佳結(jié)果，并且通過處理書中的整個文本在 BookSum 上實現(xiàn)了新的 SOTA。
研究者還在圖 4 中繪制了 BookSum 數(shù)據(jù)驗證分割的總體 Rouge 分?jǐn)?shù)。根據(jù)折線趨勢表明，隨著輸入長度的增加，Infini-Transformers 提高了摘要性能指標(biāo)。
本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心
