一個多月前，OpenAI 的 GPT-4 問世。除了各種出色的直觀演示外，它還實現(xiàn)了一個重要更新：可以處理的上下文 token 長度默認(rèn)為 8k，但最長可達(dá) 32K（大約 50 頁文本）。這意味著，在向 GPT-4 提問時，我們可以輸入比之前長得多的文本。這使得 GPT-4 的應(yīng)用場景大大擴(kuò)展，能更好地處理長對話、長文本以及文件搜索和分析。

不過，這一記錄很快就被打破了：來自谷歌研究院的 CoLT5 將模型可以處理的上下文 token 長度擴(kuò)展到了 64k。

這樣的突破并不容易，因為這些使用 Transformer 架構(gòu)的模型都要面臨一個問題：Transformer 處理長文檔在計算上是非常昂貴的，因為注意力成本隨輸入長度呈二次增長，這使得大型模型越來越難以應(yīng)用于更長的輸入。

盡管如此，研究者依然在此方向上不斷突破。前幾天，一篇來自開源對話 AI 技術(shù)棧 DeepPavlov 等機構(gòu)的研究表明：通過采用一種名為 Recurrent Memory Transformer（RMT）的架構(gòu)，他們可以將 BERT 模型的有效上下文長度增加到 200 萬個 token（按照 OpenAI 的計算方式，大約相當(dāng)于 3200 頁文本），同時保持了較高的記憶檢索準(zhǔn)確性（注：Recurrent Memory Transformer 是 Aydar Bulatov 等人在 NeurIPS 2022 的一篇論文中提出的方法）。新方法允許存儲和處理局部和全局信息，并通過使用 recurrence 使信息在輸入序列的各 segment 之間流動。

作者表示，通過使用 Bulatov 等人在「Recurrent Memory Transformer」一文中介紹的簡單的基于 token 的記憶機制，他們可以將 RMT 與 BERT 這樣的預(yù)訓(xùn)練 Transformer 模型結(jié)合起來，用一個 Nvidia GTX 1080Ti GPU 就可以對超過 100 萬個 token 的序列進(jìn)行全注意和全精度操作。

論文地址：https://arxiv.org/pdf/2304.11062.pdf

不過，也有人提醒說，這并不是真正的「免費的午餐」，上述論文的提升是用「更長的推理時間 + 實質(zhì)性的質(zhì)量下降」換來的。因此，它還不能算是一次變革，但它可能成為下一個范式（token 可能無限長）的基礎(chǔ)。

Recurrent Memory Transformer

該研究采用 Bulatov 等人 2022 年提出的方法 Recurrent Memory Transformer（RMT），并將其改成即插即用的方法，主要機制如下圖所示：

冗長的輸入被分成多個 segment，記憶向量（memory vector）被添加到第一個 segment 嵌入之前，并與 segment token 一起處理。對于像 BERT 這樣的純編碼器模型，記憶只在 segment 的開頭添加一次，這一點與 (Bulatov et al., 2022) 不同，純解碼器模型將記憶分為讀取和寫入兩部分。對于時間步長 τ 和 segment 

，循環(huán)按照如下步驟執(zhí)行：

其中，N 是 Transformer 的層數(shù)。前向傳播之后，

包含 segment τ 的更新記憶 token。

輸入序列的 segment 按照順序處理。為了啟用循環(huán)連接，該研究將記憶 token 的輸出從當(dāng)前 segment 傳遞到下一個 segment 的輸入：

RMT 中的記憶和循環(huán)都僅基于全局記憶 token。這允許主干 Transformer 保持不變，從而使 RMT 的記憶增強能力與任何 Transformer 模型都兼容。

計算效率

該研究估算了不同大小和序列長度的 RMT 和 Transformer 模型所需的 FLOP。

如下圖 3 所示，如果 segment 的長度固定，RMT 可針對任何模型大小進(jìn)行線性擴(kuò)展。該研究通過將輸入序列分成 segment，并僅在 segment 邊界內(nèi)計算完整的注意力矩陣來實現(xiàn)線性擴(kuò)展。

由于 FFN 層的計算量很大，較大的 Transformer 模型往往會表現(xiàn)出較慢的隨序列長度的二次擴(kuò)展（quadratic scaling）。然而，對于大于 32000 的超長序列，它們會退回到二次擴(kuò)展。對于多于一個 segment 的序列（本研究中 > 512），RMT 比非循環(huán)模型需要更少的 FLOP，并且可以將 FLOP 的數(shù)量減少多達(dá) 295 倍。RMT 為較小的模型提供了更大的 FLOP 相對減少，但在絕對數(shù)量上，OPT-175B 模型的 FLOP 減少了 29 倍是非常顯著的。

記憶任務(wù)

為了測試記憶能力，該研究構(gòu)建了需要記憶簡單事實和基本推理的合成數(shù)據(jù)集。任務(wù)輸入由一個或多個事實和一個只有使用所有事實才能回答的問題組成。為了增加任務(wù)難度，該研究還添加了與問題或答案無關(guān)的自然語言文本來充當(dāng)噪聲，因此模型的任務(wù)是將事實與不相關(guān)的文本分開，并使用事實來回答問題。

事實記憶

第一項任務(wù)是檢測 RMT 在記憶中長時間寫入和存儲信息的能力，如下圖 4 頂部所示。在最簡單的情況下，事實往往位于輸入開頭，而問題總是在末尾。問題和答案之間不相關(guān)文本的數(shù)量逐漸增加，以至于整個輸入不適合單個模型輸入。

事實檢測與記憶

事實檢測通過將一個事實移動到輸入中的隨機位置來增加任務(wù)難度，如上圖 4 中部所示。這需要模型首先將事實與不相關(guān)文本區(qū)分開來，把事實寫入記憶中，然后用它回答位于末尾的問題。

利用記憶事實進(jìn)行推理

記憶的另一個操作是使用記憶事實和當(dāng)前上下文進(jìn)行推理。為了評估這一功能，研究者使用了一個更復(fù)雜的任務(wù)，其中生成了兩個事實并將它們放置在了輸入序列之中，如上圖 4 底部所示。在序列末尾提出的問題是以一種「必須使用任意事實來正確回答問題」的方式來描述。

實驗結(jié)果

研究者使用 4 到 8 塊英偉達(dá) 1080ti GPU 來訓(xùn)練和評估模型。對于更長的序列，他們則使用單個 40GB 的英偉達(dá) A100 來加快評估速度。

課程學(xué)習(xí)

研究者觀察到，使用訓(xùn)練計劃能夠顯著提升解決方案的準(zhǔn)確性和穩(wěn)定性。最開始，RMT 在較短版本的任務(wù)上進(jìn)行訓(xùn)練，并在訓(xùn)練收斂時通過添加另一個 segment 來增加任務(wù)長度。課程學(xué)習(xí)過程一直持續(xù)，直到達(dá)到所需的輸入長度。

在實驗中，研究者首先從適合單個 segment 的序列開始。實際 segment 的大小為 499，但由于 BERT 的 3 個特殊 token 和 10 個記憶占位符從模型輸入中保留下來，大小為 512。他們注意到， 在較短任務(wù)上訓(xùn)練后，RMT 更容易解決更長版本任務(wù)，這得益于它使用更少訓(xùn)練步收斂到完美的解決方案。

外推能力

RMT 對不同序列長度的泛化能力如何呢？為了回答這個問題，研究者評估了在不同數(shù)量 segment 上訓(xùn)練的模型，以解決更長的任務(wù)，具體如下圖 5 所示。

他們觀察到，模型往往在較短任務(wù)上表現(xiàn)更好，唯一的例外是單 segment 推理任務(wù)，一旦模型在更長序列上訓(xùn)練，則該任務(wù)變得很難解決。一個可能的解釋是：由于任務(wù)大小超過了一個 segment，則模型不再「期待」第一個 segment 中的問題，導(dǎo)致質(zhì)量下降。

有趣的是，RMT 泛化到更長序列的能力也隨著訓(xùn)練 segment 的增加而出現(xiàn)。在 5 個或更多 segment 上訓(xùn)練后，RMT 可以近乎完美地泛化到兩倍長的任務(wù)。

為了測試泛化的局限性，研究者將驗證任務(wù)大小增至 4096 個 segment 或 2,043,904 個 token（如上圖 1 所示），RMT 在如此長的序列上表現(xiàn)得出奇的好。檢測和記憶任務(wù)最簡單，推理任務(wù)最復(fù)雜。

記憶操作的注意力模式

在下圖 6 中，通過檢查特定 segment 上的 RMT 注意力，研究者觀察到了記憶操作對應(yīng)特定的注意力模式。此外 5.2 節(jié)中極長序列上的高外推性能證明了學(xué)得記憶操作的有效性，即使使用數(shù)千次也是如此。

更多技術(shù)與實驗細(xì)節(jié)請參閱原論文。