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

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

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

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

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

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

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

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

Unlimiformer 可以被注入到任何現(xiàn)有的編碼器 - 解碼器 transformer 中，能夠處理長度不限的輸入。給定一個長的輸入序列，Unlimiformer 可以在所有輸入 token 的隱藏狀態(tài)上構(gòu)建一個數(shù)據(jù)存儲。然后，解碼器的標(biāo)準(zhǔn)交叉注意力機制能夠查詢數(shù)據(jù)存儲，并關(guān)注前 k 個輸入 token。數(shù)據(jù)存儲可以存儲在 GPU 或 CPU 內(nèi)存中，能夠次線性查詢。

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

Unlimiformer 技術(shù)原理

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

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

編碼

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

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

在標(biāo)準(zhǔn)的交叉注意力機制中，transformer 的解碼器關(guān)注編碼器的最終隱狀態(tài)，編碼器通常截斷輸入，并僅對輸入序列中的前 k 個 token 進行編碼。

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

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