大語言模型的上下文長度限制從此不存在了？！

Meta、MIT、CMU的研究者最近剛剛發(fā)表了一篇論文，提出了一種被稱為是「高效流式語言模型」（Efficient Streaming Language Models，ESLM）的方法，可以讓有限上下文能力的語言模型能夠支持幾乎無限的上下文窗口。

https://arxiv.org/pdf/2309.17453.pdf

瞬間登頂Hacker News！

甚至，這一項目在GitHub上狂攬1.8k星。

上個動圖大家先簡單感受一下效果：

一句話來形容就是——完全停不下來！

因為LLM在預訓練期間在有限的注意力窗口的限制下進行訓練，盡管可以通過各種方法來擴展模型支持的上下文長度，但是模型可接受的序列長度仍然是有限的。

核心觀點

研究人員想要解決的問題是：

能否在不犧牲效率和性能的情況下讓大預言模型支持無限長度的上下文？

研究人員在嘗試解決無限上下文長度輸入時發(fā)現(xiàn)，限制模型能力最主要的原因是這兩個：

1. 在解碼階段，基于Transformer的 LLM 會緩存之前所有token的鍵值狀態(tài)（KV），如下圖a所示，這可能會導致內(nèi)存使用過多并增加解碼延遲。

2. 現(xiàn)有模型的長度外推能力有限，當序列長度超出預訓練期間設置的注意力窗口大小時，模型性能會嚴重下降。

解決這兩個問題，研究人員嘗試了兩種方法：

第一種方法，稱為窗口注意力（window attention）（上圖b）。

這種方法只維護最新的token的KV狀態(tài)的固定大小的滑動窗口。

但它的問題是雖然在緩存最初填滿后確保了恒定的內(nèi)存使用和解碼速度，但一旦序列長度超過緩存大小，即使只是逐出第一個token的KV，模型也會崩潰，如下圖所示。

另一種可能的策略是重新計算滑動窗口（如上圖c所示）。這種發(fā)放為每個生成的token重建最新token的KV狀態(tài)。

雖然這種方法提供了強大的穩(wěn)定性，但由于需要在窗口范圍內(nèi)計算二次注意力（quadratic attention），這種方法的速度明顯過慢，使得這種方法對于沒法滿足實際應用的需求。

在探索窗口注意力方法為什么行不通的過程中，研究人員發(fā)現(xiàn)了自回歸LLM的一個有趣的特征：

大量的注意力分數(shù)被分配給初始token，而不管這些token與語言建模任務的相關性如何，如下圖所示。

研究人員將這些占用了大量注意力的初始token稱為——「注意力池（attention sinks）」

盡管這些初始token很多時候缺乏語義含義，但它們卻占用了很高的注意力分數(shù)（attention scores）。

研究人員認為主要原因是因為Softmax操作，要求所有上下文token的注意力分數(shù)總和為1。

因此，即使當前任務和許多先前的token不匹配，模型仍然需要在某個地方分配這些不需要的注意力分數(shù)，使得分數(shù)總和為1。

初始token為什么會獲得如此之高的注意力，原因也很簡單：

初始token對幾乎所有后續(xù)token都是可見的，因為自回歸語言模型的性質，使初始token更容易被訓練為「注意力池」。

基于這些原因，研究人員提出了一個新的框架：Streaming LLM。

這是一個簡單而高效的框架，使得通過有限注意力窗口訓練的LLM能夠處理無限長度的文本，而無需進行微調。

Streaming LLM利用了初始token——「注意力池」具有高注意力值的這一特點，保留初始token可以保持注意力分數(shù)分布接近正常。

因此，Streaming LLM只是將注意力集中token的KV（只需 4 個初始token就足夠了）與滑動窗口的KV保持在一起，以錨定注意力計算并穩(wěn)定模型的性能。

使用 StreamingLLM，模型包括 Llama-2-[7, 13, 70]B、MPT-[7, 30]B、Falcon-[7, 40]B 和 Pythia

最后，研究人員證實了他們提出的「注意力池」的假設，并證明語言模型可以經(jīng)過預訓練，只需要一個注意力池token即可進行流部署。

具體來說，研究人員建議在所有訓練樣本的開頭添加一個額外的可學習token可以作為指定的注意力接收器。

通過從頭開始預訓練 1.6 億個參數(shù)語言模型，研究人員證明添加這個單一接收器token可以保持模型在流情況下的性能。這與普通模型形成鮮明對比，普通模型需要重新引入多個初始token作為注意力集中器才能達到相同的性能水平。

Streaming LLM

為什么窗口注意力方法行不通？

雖然窗口注意力技術在推理過程中提供了效率，但它導致了極高的語言建模復雜性。因此，該模型的性能不適合部署在流應用程序中。

在本節(jié)中，研究人員使用注意力接收器的概念來解釋窗口注意力的失敗，作為 StreamingLLM 背后的靈感。

識別困惑度激增點

上圖顯示了 20K token文本的語言建模的復雜性。很明顯，當文本長度超過緩存大小時，由于排除初始token而導致困惑度度激增。這表明初始token，無論與預測token的距離如何，對于維持 LLM 的穩(wěn)定性至關重要。

為什么LLM在刪除初始token的KV時會崩潰？

研究人員將 Llama-2-7B 和模型的所有層和頭部的注意力圖可視化，如圖 2 所示。研究人員發(fā)現(xiàn)，除了底部兩層之外，模型始終關注所有層和頭部的初始token。

含義很明確：刪除這些初始token的 KV 將刪除注意力計算中 SoftMax 函數(shù)（公式 1）中相當一部分分母。這種改變導致注意力分數(shù)的分布發(fā)生顯著偏離正常推理設置中的預期。

對于語言模型中初始token的重要性，有兩種可能的解釋：（1）它們的語義至關重要，或者（2）模型學習到對其絕對位置的偏見。

為了區(qū)分這些可能性，研究人員進行了實驗（上表），其中前四個token被換行token「\n」替換。

觀察結果表明，模型仍然顯著強調這些初始換行token。此外，重新引入它們可以恢復將語言建模的復雜性達到與原始初始token相當?shù)乃?。這表明起始token的絕對位置，而不是它們的語義值，具有更重要的意義。

LLM將初始token作為注意力池。為了解釋為什么模型不成比例地關注初始token——無論它們與語言建模的語義相關性如何，研究人員引入了「注意力池」的概念。

SoftMax 函數(shù)（公式 1）的本質阻止所有參與的token具有零值。即使當前嵌入有足夠的自包含信息用于其預測，也需要聚合來自所有層中所有頭的其他token的一些信息。

因此，該模型傾向于將不必要的注意力值轉儲到特定token。類似的觀察已在量化異常值領域，導致提出 SoftMax-Off-by-One作為一種潛在的補救措施。

為什么各種自回歸LLM（例如 Llama-2、MPT、Falcon 和 Pythia）在注意力集中時始終關注初始token，而不是其他token？

研究人員的解釋很簡單：由于自回歸語言模型的順序性質，初始token對所有后續(xù)token都是可見的，而后面的token僅對有限的一組后續(xù)token可見。因此，初始token更容易被訓練為注意力集中器，捕獲不必要的注意力。

研究人員注意到，LLM通常接受訓練，利用多個初始token作為注意力池，而不僅僅是一個。如上圖所示，隨著注意力的下降，引入四個初始token足以恢復 LLM 的表現(xiàn)。

相反，僅添加一兩個并不能實現(xiàn)完全恢復。研究人員認為出現(xiàn)這種模式是因為這些模型在預訓練期間沒有在所有輸入樣本中包含一致的起始token。盡管 Llama-2 確實為每個段落添加了“”token作為前綴，但它是在文本分塊之前應用的，導致大多數(shù)隨機token占據(jù)第零個位置。

缺乏統(tǒng)一的起始token導致模型使用多個初始token作為注意力池。研究人員假設，通過在所有訓練樣本的開始處合并一個穩(wěn)定的可學習token，它可以單獨充當承諾的注意力接收器，從而消除對多個初始token來確保一致流的需要。研究人員將在第 3.3 節(jié)中驗證這一假設。

3.2 使用注意力接收器滾動 KV 緩存

為了在已經(jīng)訓練好的LLM中啟用LLM Streaming，研究人員提出了一種簡單的方法，可以恢復窗口注意力的困惑度，而無需任何模型微調。

除了當前的滑動窗口token之外，研究人員在注意力計算中重新引入了一些起始token的 KV。

StreamingLLM中的KV緩存在概念上可以分為兩部分，如上圖所示：

1）注意力接收器（四個初始token）穩(wěn)定注意力計算；

2）滾動KV緩存保留最新的token，這對于語言建模至關重要。StreamingLLM 的設計用途廣泛，可以無縫地融入任何采用相對位置編碼的自回歸語言模型，例如 RoPE和 ALiBi。

在確定相對距離并向token添加位置信息時，StreamingLLM 重點關注緩存中的位置而不是原始文本中的位置。這種區(qū)別對于 StreamingLLM 的性能至關重要。例如，如果當前緩存有token [0, 1, 2, 3, 6, 7, 8] 并且正在解碼第 9 個token，則分配的位置為 [0, 1, 2, 3, 4, 5, 6, 7]，而不是原文中的位置，即 [0, 1, 2, 3, 6, 7, 8, 9]。

3.3 使用注意力池預訓練 LLMS

正如上文所述，模型過度關注多個初始token的一個重要原因是缺乏指定的接收器token來卸載過多的注意力分數(shù)。因此，該模型無意中將全局可見的token（主要是初始token）指定為注意力池。

一種潛在的補救措施是有意包含一個全局可訓練的注意力接收器token，表示為「Sink Token」，它將作為不必要的注意力分數(shù)的存儲庫?；蛘?，用 SoftMax-off- 等變體替換傳統(tǒng)的SoftMax函數(shù)。

不需要所有上下文token的注意力分數(shù)總和為 1，也可能是有效的。

請注意，此 SoftMax 替代方案等效于在注意力計算中使用具有全零 Key 和 Value 特征的token。研究人員將此方法稱為「零池（Zero Sink）」，以使其始終適合研究人員的框架。

為了進行驗證，研究人員在相同的設置下從頭開始預訓練具有 1.6 億個參數(shù)的三種語言模型。

第一個模型利用標準 SoftMax 注意力（Vanilla），第二個模型用 SoftMax1（零池）替換常規(guī)注意力機制，并且在所有訓練樣本中預先考慮可學習的占位符token（Sink Token）。如下表所示，雖然零匯在一定程度上緩解了注意力池問題，但模型仍然依賴其他初始token作為注意力池。

引入sink token對于穩(wěn)定注意力機制非常有效。簡單地將這個接收器token與最近的token配對就足以錨定模型的性能，并且由此產(chǎn)生的評估困惑度甚至略有改善。

鑒于這些發(fā)現(xiàn)，研究人員建議在所有樣本中使用接收器token來訓練未來的法學碩士，以優(yōu)化LLM Stream的部署。

作者介紹

Guangxuan Xiao

MIT EECS 博士一年級在讀，導師為韓松教授。

2022年畢業(yè)于清華大學計算機科學與技術系，師從劉志遠教授和翟繼東教授。研究方向是機器學習應用的高效算法和系統(tǒng)，特別是基礎模型方向。

田淵棟

Meta AI Research (FAIR) 的研究科學家和高級經(jīng)理，致力于強化學習、表示學習和優(yōu)化方向的研究。

他是2021年ICML杰出論文榮譽獎和2013年ICCV Marr獎榮譽獎的第一作者獲得者，并且是該項目的首席科學家和工程師。

他還獲得了 2022 年 CGO 杰出論文獎。

此前，他于2013-2014年在谷歌自動駕駛汽車團隊工作。他于2013年在卡內(nèi)基梅隆大學機器人研究所獲得博士學位。

BEIDI CHEN

她目前是 Meta/Facebook AI Research (FAIR) 的客座研究員。之前是斯坦福大學的博士后研究員，與 Chris Ré 博士一起工作。

將于 2023 年秋季開始擔任卡內(nèi)基梅隆大學電氣與計算機工程助理教授，她的研究方向是大規(guī)模機器學習，設計和優(yōu)化隨機算法（算法-硬件協(xié)同設計）來加速大型機器學習系統(tǒng)解決現(xiàn)實問題。

韓松

韓松現(xiàn)任麻省理工學院（MIT）電子工程學院副教授。本科畢業(yè)于清華大學，在斯坦福大學取得博士學位。他提出了廣泛應用于高效人工智能計算的「深度壓縮」技術。

他還率先提出將權重稀疏性引入現(xiàn)代人工智能芯片的EIE（Efficient Inference Engine），并影響了英偉達安培GPU架構。他開創(chuàng)了TinyML研究，將機器學習引入到物聯(lián)網(wǎng)設備，實現(xiàn)了邊緣機器學習。