最近一段時間，有關(guān)長文本能力的探討被擺到了臺面上，而就在人們比拼上下文窗口時，谷歌發(fā)布了一篇論文，談到團(tuán)隊(duì)發(fā)明了一種新的注意力技術(shù)，這項(xiàng)技術(shù)可以讓transformer大模型在有限的計(jì)算資源條件下處理無限長度的輸入。

當(dāng)人們還在比拼上下文窗口的時候，谷歌發(fā)布了這樣一篇論文《Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention》。論文中寫到，團(tuán)隊(duì)發(fā)明了一種新的注意力技術(shù)，叫做“無限注意力”（Infini-attention），通過這項(xiàng)技術(shù)，能使transformer大模型在有限的計(jì)算資源條件下，處理無限長度的輸入。

在一個transformer模型中，注意力的作用是允許模型根據(jù)當(dāng)前位置的輸入元素（如詞元或token）來分配權(quán)重給序列中其他所有位置的元素。而上下文窗口則限制了注意力機(jī)制的實(shí)際操作范圍，即模型在計(jì)算注意力時僅考慮當(dāng)前元素周圍特定范圍內(nèi)（前后若干位置）的其他元素。

無限注意力允許允許模型在處理無限長輸入序列的時候，仍能保持對上下文信息的訪問能力，再也不需要在處理新輸入時丟棄前一輸入段的注意力狀態(tài)了。那么也就是說，它的上下文窗口可以是……無限。

無限注意力機(jī)制背后的關(guān)鍵技術(shù)叫做壓縮記憶系統(tǒng)，這是一種能夠以緊湊形式存儲和檢索大量信息的結(jié)構(gòu)，通過改變自身參數(shù)來捕獲新信息，確保信息可以在之后被恢復(fù)。單從運(yùn)行的邏輯上來講，壓縮記憶系統(tǒng)和咱們?nèi)粘Ｉ罾飰嚎s文件是一模一樣的。

壓縮記憶系統(tǒng)最大的作用是克服transformer標(biāo)準(zhǔn)注意力機(jī)制在處理長序列時存在的內(nèi)存足跡和計(jì)算時間的二次復(fù)雜度問題，只需要通過使用固定數(shù)量的參數(shù)存儲和召回信息，確保存儲和計(jì)算成本保持在可控范圍內(nèi)。因?yàn)槠鋮?shù)數(shù)量不隨輸入序列的增長而變化，那也就是說，無論輸入序列長度有多長，也不會對影響模型的復(fù)雜度。

接下來無限注意力機(jī)制會將輸入序列劃分為一系列嬌小的、連續(xù)的子序列，每個段具有固定的長度，使得模型在處理這些較短的段時，能夠保持較低的內(nèi)存需求和計(jì)算復(fù)雜度。這種分段方法避免了一次性加載和處理整個無限長序列的挑戰(zhàn)，允許模型以流式（streaming）方式逐步處理輸入，即每次僅處理一個或幾個段，而非一次性加載全部數(shù)據(jù)。

在每個分段內(nèi)部，無限注意力模型采用局部注意力機(jī)制來處理該段內(nèi)的上下文信息。局部注意力限制了模型對當(dāng)前段內(nèi)token之間的注意力計(jì)算范圍，通常采用因果（causal）或自回歸（autoregressive）的形式，確保模型在處理當(dāng)前token時，只能看到該令牌之前的所有token，而不能看到未來（即當(dāng)前token之后）的任何token。

在輸出結(jié)果時，無限注意力模型為了生成最終的上下文輸出，要從壓縮記憶中檢索到的長期記憶信息與當(dāng)前局部注意力計(jì)算出的上下文結(jié)合起來。這種融合確保模型既考慮了當(dāng)前輸入段的局部依賴，又充分利用了歷史輸入的長期上下文。

當(dāng)你理解了無限注意力機(jī)制后再回到標(biāo)題，無限注意力模型能夠以流式方式處理極端長的輸入序列，無需一次性加載整個無限長的輸入，而是會根據(jù)歷史記錄進(jìn)行分批次處理。那對于模型來說，就能夠在有限的內(nèi)存和計(jì)算資源約束下，適應(yīng)并處理無限長度的上下文。

論文首先在長上下文語言建?；鶞?zhǔn)上評估了無限注意力模型的表現(xiàn)，與包括transformer-XL在內(nèi)的多種模型進(jìn)行了對比。

采用無限注意力的模型在PG19（長文檔數(shù)據(jù)集）以Arxiv-math（數(shù)學(xué)數(shù)據(jù)集）上都取得了遠(yuǎn)超于transformer-XL的結(jié)果，同時實(shí)現(xiàn)了114倍的內(nèi)存壓縮率，在保持低困惑度的同時提高了模型效率。

為了進(jìn)一步驗(yàn)證無限注意力機(jī)制的性能，論文將一個10億參數(shù)的大語言模型進(jìn)行改造，把這個模型的多頭注意力（MHA）模塊換成了無限注意力，并繼續(xù)對其進(jìn)行預(yù)訓(xùn)練。驗(yàn)證過程是，團(tuán)隊(duì)要求模型在長達(dá)100萬tokens的輸入中定位并檢索隱藏的密鑰信息。

預(yù)訓(xùn)練階段，模型使用的輸入序列長度僅為4K個tokens，以適應(yīng)無限注意力的處理模式。經(jīng)過3萬步的預(yù)訓(xùn)練后，對密鑰檢索任務(wù)進(jìn)行微調(diào)。在微調(diào)階段，為了模擬實(shí)際應(yīng)用中可能遇到的更長上下文環(huán)境，模型在包含5K個token的長度輸入上進(jìn)行微調(diào)。

在完成預(yù)訓(xùn)練和微調(diào)后，團(tuán)隊(duì)對模型進(jìn)行評估，在不同長度（從32K到1M）和不同密鑰位置（開始、中間、結(jié)束）的長輸入文本中檢索密鑰的準(zhǔn)確性。實(shí)驗(yàn)結(jié)果表明無限注意力模型在所有測試場景中均能成功找回隱藏的密鑰，展現(xiàn)出其對極長上下文信息的卓越處理能力。

隨后團(tuán)隊(duì)為了證明無限注意力機(jī)制在更大參數(shù)模型上的表現(xiàn)，又對一個用無限注意力改造的80億參數(shù)大語言模型進(jìn)行了預(yù)訓(xùn)練。使用8k個token長度的輸入訓(xùn)練了3萬步。模型在BookSum數(shù)據(jù)集上進(jìn)行微調(diào)，輸入長度設(shè)置為32K用于微調(diào)，但在評估階段增加到500K。

根據(jù)無限注意力在50萬文本長度的圖書中里生成的摘要，模型超越了專門為摘要任務(wù)構(gòu)建的編碼器-解碼器模型及其長上下文擴(kuò)展版本，實(shí)現(xiàn)了在BookSum數(shù)據(jù)集上的新SOTA（state-of-the-art）性能。隨著輸入書籍文本量的增加，模型的摘要性能指標(biāo)（如Rouge分?jǐn)?shù)）呈現(xiàn)出明顯的上升趨勢。

一個有效的記憶系統(tǒng)不僅對大型語言模型理解長文本來說是至關(guān)重要的，雖然論文并沒有大刀闊斧地修改transformer模型的注意力機(jī)制，只是用了類似于微創(chuàng)手術(shù)一樣的手法，把壓縮記憶模塊緊密地集成進(jìn)了模型的標(biāo)準(zhǔn)點(diǎn)積注意力層（vanilla dot-product attention layer），卻徹頭徹尾改善了transformer模型在處理長序列時碰到的問題。

2022年的時候，deepmind曾發(fā)文《∞-former: Infinite Memory Transformer》，論文提出了一個叫做∞-former的模型，通過利用連續(xù)空間注意力機(jī)制對長期記憶進(jìn)行關(guān)注，讓模型的注意力復(fù)雜度變得與上下文長度無關(guān)。從方法上來看，無限注意力和∞-former是有些相似的。后者以犧牲精度為代價換取了記憶長度，可是無限注意力卻可以在極端長度的密鑰中找到關(guān)鍵信息，精準(zhǔn)度甚至比以往要高很多。

其實(shí)歸根結(jié)底，無限注意力和∞-former都是對transformer的記憶系統(tǒng)進(jìn)行改進(jìn)。不過transformer有一大缺陷是不能處理非連續(xù)的數(shù)據(jù)。因?yàn)閠ransformer最初的設(shè)計(jì)是用于處理自然語言這樣的連續(xù)文本序列，但隨著圖片生成、音樂生成、視頻生成等多個領(lǐng)域應(yīng)用的崛起，模型為了應(yīng)對多模態(tài)的數(shù)據(jù)結(jié)構(gòu)就必須能夠處理非連續(xù)的數(shù)據(jù)。谷歌若想擴(kuò)大自己在多模態(tài)領(lǐng)域的領(lǐng)先地位，可能會開始數(shù)據(jù)結(jié)構(gòu)方面的研究工作。