當人們還在比拼上下文窗口的時候,谷歌發布了這樣一篇論文《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 最初的設計是用于處理自然語言這樣的連續文本序列,但随着圖片生成、音樂生成、視頻生成等多個領域應用的崛起,模型爲了應對多模态的數據結構就必須能夠處理非連續的數據。谷歌若想擴大自己在多模态領域的領先地位,可能會開始數據結構方面的研究工作。