天天看點

還在卷長文本?谷歌最新論文直接把文本幹到…無限長了

作者:硬梗娛樂

當人們還在比拼上下文視窗的時候,谷歌釋出了這樣一篇論文《Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention》。論文中寫到,團隊發明了一種新的注意力技術,叫做“無限注意力”(Infini-attention),通過這項技術,能使transformer大模型在有限的計算資源條件下,處理無限長度的輸入。

在一個transformer模型中,注意力的作用是允許模型根據目前位置的輸入元素(如詞元或token)來配置設定權重給序列中其他所有位置的元素。而上下文視窗則限制了注意力機制的實際操作範圍,即模型在計算注意力時僅考慮目前元素周圍特定範圍内(前後若幹位置)的其他元素。

無限注意力允許允許模型在處理無限長輸入序列的時候,仍能保持對上下文資訊的通路能力,再也不需要在處理新輸入時丢棄前一輸入段的注意力狀态了。那麼也就是說,它的上下文視窗可以是……無限。

無限注意力機制背後的關鍵技術叫做壓縮記憶系統,這是一種能夠以緊湊形式存儲和檢索大量資訊的結構,通過改變自身參數來捕獲新資訊,確定資訊可以在之後被恢複。單從運作的邏輯上來講,壓縮記憶系統和咱們日常生活裡壓縮檔案是一模一樣的。

壓縮記憶系統最大的作用是克服transformer标準注意力機制在處理長序列時存在的記憶體足迹和計算時間的二次複雜度問題,隻需要通過使用固定數量的參數存儲和召回資訊,确儲存儲和計算成本保持在可控範圍内。因為其參數數量不随輸入序列的增長而變化,那也就是說,無論輸入序列長度有多長,也不會對影響模型的複雜度。

接下來無限注意力機制會将輸入序列劃分為一系列嬌小的、連續的子序列,每個段具有固定的長度,使得模型在處理這些較短的段時,能夠保持較低的記憶體需求和計算複雜度。這種分段方法避免了一次性加載和處理整個無限長序列的挑戰,允許模型以流式(streaming)方式逐漸處理輸入,即每次僅處理一個或幾個段,而非一次性加載全部資料。

在每個分段内部,無限注意力模型采用局部注意力機制來處理該段内的上下文資訊。局部注意力限制了模型對目前段内token之間的注意力計算範圍,通常采用因果(causal)或自回歸(autoregressive)的形式,確定模型在處理目前token時,隻能看到該令牌之前的所有token,而不能看到未來(即目前token之後)的任何token。

在輸出結果時,無限注意力模型為了生成最終的上下文輸出,要從壓縮記憶中檢索到的長期記憶資訊與目前局部注意力計算出的上下文結合起來。這種融合確定模型既考慮了目前輸入段的局部依賴,又充分利用了曆史輸入的長期上下文。

當你了解了無限注意力機制後再回到标題,無限注意力模型能夠以流式方式處理極端長的輸入序列,無需一次性加載整個無限長的輸入,而是會根據曆史記錄進行分批次處理。那對于模型來說,就能夠在有限的記憶體和計算資源限制下,适應并處理無限長度的上下文。

論文首先在長上下文語言模組化基準上評估了無限注意力模型的表現,與包括transformer-XL在内的多種模型進行了對比。

還在卷長文本?谷歌最新論文直接把文本幹到…無限長了

采用無限注意力的模型在PG19(長文檔資料集)以Arxiv-math(數學資料集)上都取得了遠超于transformer-XL的結果,同時實作了114倍的記憶體壓縮率,在保持低困惑度的同時提高了模型效率。

為了進一步驗證無限注意力機制的性能,論文将一個10億參數的大語言模型進行改造,把這個模型的多頭注意力(MHA)子產品換成了無限注意力,并繼續對其進行預訓練。驗證過程是,團隊要求模型在長達100萬tokens的輸入中定位并檢索隐藏的密鑰資訊。

預訓練階段,模型使用的輸入序列長度僅為4K個tokens,以适應無限注意力的處理模式。經過3萬步的預訓練後,對密鑰檢索任務進行微調。在微調階段,為了模拟實際應用中可能遇到的更長上下文環境,模型在包含5K個token的長度輸入上進行微調。

在完成預訓練和微調後,團隊對模型進行評估,在不同長度(從32K到1M)和不同密鑰位置(開始、中間、結束)的長輸入文本中檢索密鑰的準确性。實驗結果表明無限注意力模型在所有測試場景中均能成功找回隐藏的密鑰,展現出其對極長上下文資訊的卓越處理能力。

随後團隊為了證明無線注意力機制在更大參數模型上的表現,又對一個用無限注意力改造的80億參數大語言模型進行了預訓練。使用8k個token長度的輸入訓練了3萬步。模型在BookSum資料集上進行微調,輸入長度設定為32K用于微調,但在評估階段增加到500K。

根據無限注意力在50萬文本長度的圖書中裡生成的摘要,模型超越了專門為摘要任務建構的編碼器-解碼器模型及其長上下文擴充版本,實作了在BookSum資料集上的新SOTA(state-of-the-art)性能。随着輸入書籍文本量的增加,模型的摘要性能名額(如Rouge分數)呈現出明顯的上升趨勢。

還在卷長文本?谷歌最新論文直接把文本幹到…無限長了

一個有效的記憶系統不僅對大型語言模型了解長文本來說是至關重要的,雖然論文并沒有大刀闊斧地修改transformer模型的注意力機制,隻是用了類似于微創手術一樣的手法,把壓縮記憶子產品緊密地內建進了模型的标準點積注意力層(vanilla dot-product attention layer),卻徹頭徹尾改善了transformer模型在處理長序列時碰到的問題。

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

其實歸根結底,無限注意力和∞-former都是對transformer的記憶系統進行改進。不過transformer有一大缺陷是不能處理非連續的資料。因為transformer最初的設計是用于處理自然語言這樣的連續文本序列,但随着圖檔生成、音樂生成、視訊生成等多個領域應用的崛起,模型為了應對多模态的資料結構就必須能夠處理非連續的資料。谷歌若想擴大自己在多模态領域的領先地位,可能會開始資料結構方面的研究工作。

繼續閱讀