能容納 50 頁(yè)文檔的輸入框不夠用,，那幾千頁(yè)呢,？

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

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

這樣的突破并不容易，因?yàn)檫@些使用 Transformer 架構(gòu)的模型都要面臨一個(gè)問(wèn)題：Transformer 處理長(zhǎng)文檔在計(jì)算上是非常昂貴的,，因?yàn)樽⒁饬Τ杀倦S輸入長(zhǎng)度呈二次增長(zhǎng),，這使得大型模型越來(lái)越難以應(yīng)用于更長(zhǎng)的輸入。

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

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

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

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

Recurrent Memory Transformer

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

冗長(zhǎng)的輸入被分成多個(gè) segment,，記憶向量（memory vector）被添加到第一個(gè) segment 嵌入之前，并與 segment token 一起處理,。對(duì)于像 BERT 這樣的純編碼器模型,，記憶只在 segment 的開(kāi)頭添加一次，這一點(diǎn)與 (Bulatov et al., 2022) 不同，純解碼器模型將記憶分為讀取和寫(xiě)入兩部分,。對(duì)于時(shí)間步長(zhǎng) τ 和 segment ,，循環(huán)按照如下步驟執(zhí)行：

其中，N 是 Transformer 的層數(shù),。前向傳播之后,，包含 segment τ 的更新記憶 token。

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

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

計(jì)算效率

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

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

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

記憶任務(wù)

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

事實(shí)記憶

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

事實(shí)檢測(cè)與記憶

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

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

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

外推能力

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

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

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

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

記憶操作的注意力模式

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

更多技術(shù)與實(shí)驗(yàn)細(xì)節(jié)請(qǐng)參閱原論文。