谷歌大改 Transformer," 無限 " 長度上下文來了。
現在,1B 大模型上下文長度可擴展到 1M(100 萬 token,大約相當于 10 部小說),并能完成 Passkey 檢索任務。
8B 大模型在 500K 上下文長度的書籍摘要任務中,拿下最新 SOTA。
這就是谷歌最新提出的Infini-attention 機制(無限注意力)。
它能讓 Transformer 架構大模型在有限的計算資源裏處理無限長的輸入,在内存大小上實現114 倍壓縮比。
什麽概念?
就是在内存大小不變的情況下,放進去 114 倍多的信息。好比一個存放 100 本書的圖書館,通過新技術能存儲 11400 本書了。
這項最新成果立馬引發學術圈關注,大佬紛紛圍觀。
加之最近 DeepMind 也改進了 Transformer 架構,使其可以動态分配計算資源,以此提高訓練效率。
有人感慨,基于最近幾個新進展,感覺大模型越來越像一個包含高度可替換、商品化組件的軟件棧了。
引入壓縮記憶
該論文核心提出了一種新機制 Infini-attention。
它通過将壓縮記憶(compressive memory)整合到線性注意力機制中,用來處理無限長上下文。
壓縮記憶允許模型在處理新輸入時保留和重用之前的上下文信息。它通過固定數量的參數來存儲和回憶信息,而不是随着輸入序列長度的增加而增加參數量,能減少内存占用和計算成本。
線性注意力機制不同于傳統 Transformer 中的二次方複雜度注意力機制,它能通過更小的計算開銷來檢索和更新長期記憶。
在 Infini-attention 中,舊的 KV 狀态({KV}s-1)被存儲在壓縮記憶中,而不是被丢棄。
通過将查詢與壓縮記憶中存儲的鍵值進行匹配,模型就可以檢索到相關的值。
PE 表示位置嵌入,用于給模型提供序列中元素的位置信息。
對比來看 Transformer-XL,它隻緩存最後一段 KV 狀态,在處理新的序列段時就會丢棄舊的鍵值對,所以它隻能保留最近一段的上下文信息。
對比幾種不同 Transformer 模型可處理上下文的長度和内存占用情況。
Infini-attention 能在内存占用低的情況下,有效處理非常長的序列。
Infini-attention 在訓練後,分化出了兩種不同類型的注意力頭,它們協同處理長期和短期上下文信息。
專門化的頭(Specialized heads):這些頭在訓練過程中學習到了特定的功能,它們的門控得分(gating score)接近 0 或 1。這意味着它們要麽通過局部注意力機制處理當前的上下文信息,要麽從壓縮記憶中檢索信息。
混合頭(Mixer heads):這些頭的門控得分接近 0.5,它們的作用是将當前的上下文信息和長期記憶内容聚合到單一的輸出中。
研究團隊将訓練長度增加到 100K,在 Arxiv-math 數據集上進行訓練。
在長下文語言建模任務中,Infini-attention 在保持低内存占用的同時,困惑度更低。
對比來看,同樣情況下 Memorizing Transformer 存儲參數所需的内存是 Infini-attention 的 114 倍。
消融實驗比較了 " 線性 " 和 " 線性 + 增量 " 記憶兩種模式,結果顯示性能相當。
實驗結果顯示,即使在輸入隻有 5K 進行微調的情況下,Infini-Transformer 可成功搞定 1M 長度(100 萬)的 passkey 檢索任務。
在處理長達 500K 長度的書籍摘要任務時,Infini-Transformer 達到最新 SOTA。
Bard 成員參與研究
該研究由谷歌團隊帶來。
其中一位作者(Manaal Faruqui)在 Bard 團隊,領導研究 Bard 的模型質量、指令遵循等問題。
最近,DeepMind 的一項工作也關注到了高效處理長序列數據上。他們提出了兩個新的 RNN 模型,在高效處理長序列時還實現了和 Transformer 模型相當的性能和效率。
感覺到谷歌最近的研究重點之一就是長文本,論文在陸續公布。
網友覺得,很難了解哪些是真正開始推行使用的,哪些隻是一些研究員心血來潮的成果。
不過想象一下,如果有一些初創公司專門做内存數據庫,但是模型能已經能實現無限内存了,這可真是太有趣了。
論文地址:
https://arxiv.org/abs/2404.07143
參考鏈接:
[ 1 ] https://twitter.com/Joby_Fi/status/1778240236201386072
[ 2 ] https://twitter.com/omarsar0/status/1778480897198612839
[ 3 ] https://twitter.com/swyx/status/1778553757762252863