GPT-4的32k輸入框還是不夠用？Unlimiformer把上下文長度拉到無限長

作者：機器之心 2023-05-19 10:16:27

Unlimiformer 可以被注入到任何現有的編碼器 - 解碼器 transformer 中，能夠處理長度不限的輸入。

Transformer 是時下最強大的 seq2seq 架構。預訓練 transformer 通常具有 512（例如 BERT）或 1024 個（例如 BART）token 的個上下文窗口，這對于目前許多文本摘要數據集（XSum、CNN/DM）來說是足夠長的。

但 16384 并不是生成所需上下文長度的上限：涉及長篇敘事的任務，如書籍摘要（Krys-′cinski et al.，2021）或敘事問答（Kociskyet al.，2018），通常輸入超過 10 萬個 token。維基百科文章生成的挑戰集（Liu*et al.，2018）包含超過 50 萬個 token 的輸入。生成式問答中的開放域任務可以從更大的輸入中綜合信息，例如回答關于維基百科上所有健在作者的文章的聚合屬性的問題。圖 1 根據常見的上下文窗口長度繪制了幾個流行的摘要和問答數據集的大小；最長的輸入比 Longformer 的上下文窗口長 34 倍以上。

在這些超長輸入的情況下，vanilla transformer 無法進行縮放，因為原生注意力機制具有平方級的復雜度。長輸入 transformer 雖然比標準 transformer 更高效，但仍需要大量的計算資源，這些資源隨著上下文窗口大小的增加而增加。此外，增加上下文窗口需要用新的上下文窗口大小從頭開始重新訓練模型，計算上和環境上的代價都不小。

在「Unlimiformer: Long-Range Transformers with Unlimited Length Input」一文中，來自卡內基梅隆大學的研究者引入了 Unlimiformer。這是一種基于檢索的方法，這種方法增強了預訓練的語言模型，以在測試時接受無限長度的輸入。

論文鏈接：https://arxiv.org/pdf/2305.01625v1.pdf

Unlimiformer 可以被注入到任何現有的編碼器 - 解碼器 transformer 中，能夠處理長度不限的輸入。給定一個長的輸入序列，Unlimiformer 可以在所有輸入 token 的隱藏狀態上構建一個數據存儲。然后，解碼器的標準交叉注意力機制能夠查詢數據存儲，并關注前 k 個輸入 token。數據存儲可以存儲在 GPU 或 CPU 內存中，能夠次線性查詢。

Unlimiformer 可以直接應用于經過訓練的模型，并且可以在沒有任何進一步訓練的情況下改進現有的 checkpoint。Unlimiformer 經過微調后，性能會得到進一步提高。本文證明，Unlimiformer 可以應用于多個基礎模型，如 BART（Lewis et al.，2020a）或 PRIMERA（Xiao et al.，2022），且無需添加權重和重新訓練。在各種長程 seq2seq 數據集中，Unlimiformer 不僅在這些數據集上比 Longformer（Beltagy et al.，2020b）、SLED（Ivgi et al.，2022）和 Memorizing transformers（Wu et al.，2021）等強長程 Transformer 表現更好，而且本文還發現 Unlimiform 可以應用于 Longformer 編碼器模型之上，以進行進一步改進。

Unlimiformer 技術原理

由于編碼器上下文窗口的大小是固定的，Transformer 的最大輸入長度受到限制。然而，在解碼過程中，不同的信息可能是相關的；此外，不同的注意力頭可能會關注不同類型的信息（Clark et al.，2019）。因此，固定的上下文窗口可能會在注意力不那么關注的 token 上浪費精力。

在每個解碼步驟中，Unlimiformer 中每個注意力頭都會從全部輸入中選擇一個單獨的上下文窗口。通過將 Unlimiformer 查找注入解碼器來實現：在進入交叉注意力模塊之前，該模型在外部數據存儲中執行 k 最近鄰 (kNN) 搜索，在每個解碼器層中的每個注意力頭中選一組 token 來參與。

編碼

為了將比模型的上下文窗口長度更長的輸入序列進行編碼，本文按照 Ivgi et al. (2022) 的方法對輸入的重疊塊進行編碼 (Ivgi et al. ,2022)，只保留每個 chunk 的輸出的中間一半，以確保編碼過程前后都有足夠的上下文。最后，本文使用 Faiss (Johnson et al., 2019) 等庫對數據存儲中的編碼輸入進行索引（Johnson et al.，2019）。

檢索增強的交叉注意力機制

在標準的交叉注意力機制中，transformer 的解碼器關注編碼器的最終隱狀態，編碼器通常截斷輸入，并僅對輸入序列中的前 k 個 token 進行編碼。

本文不是只關注輸入的這前 k 個 token，對于每個交叉注意頭，都檢索更長的輸入系列的前 k 個隱狀態，并只關注這前 k 個。這樣就能從整個輸入序列中檢索關鍵字，而不是截斷關鍵字。在計算和 GPU 內存方面，本文的方法也比處理所有輸入 token 更便宜，同時通常還能保留 99% 以上的注意力性能。

圖 2 顯示了本文對 seq2seq transformer 架構的更改。使用編碼器對完整輸入進行塊編碼，并將其存儲在數據存儲中；然后，解碼時查詢編碼的隱狀態數據存儲。kNN 搜索是非參數的，并且可以被注入到任何預訓練的 seq2seq transformer 中，詳情如下。