大語(yǔ)言模型的上下文長(zhǎng)度限制從此不存在了？！

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

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

瞬間登頂Hacker News！

甚至，這一項(xiàng)目在GitHub上狂攬1.8k星。

上個(gè)動(dòng)圖大家先簡(jiǎn)單感受一下效果：

一句話來(lái)形容就是——完全停不下來(lái)！

因?yàn)長(zhǎng)LM在預(yù)訓(xùn)練期間在有限的注意力窗口的限制下進(jìn)行訓(xùn)練，盡管可以通過(guò)各種方法來(lái)擴(kuò)展模型支持的上下文長(zhǎng)度，但是模型可接受的序列長(zhǎng)度仍然是有限的。

核心觀點(diǎn)

研究人員想要解決的問(wèn)題是：

能否在不犧牲效率和性能的情況下讓大預(yù)言模型支持無(wú)限長(zhǎng)度的上下文？

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

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

2. 現(xiàn)有模型的長(zhǎng)度外推能力有限，當(dāng)序列長(zhǎng)度超出預(yù)訓(xùn)練期間設(shè)置的注意力窗口大小時(shí)，模型性能會(huì)嚴(yán)重下降。

解決這兩個(gè)問(wèn)題，研究人員嘗試了兩種方法：

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

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

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

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

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

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

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

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

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

研究人員認(rèn)為主要原因是因?yàn)镾oftmax操作，要求所有上下文token的注意力分?jǐn)?shù)總和為1。

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

初始token為什么會(huì)獲得如此之高的注意力，原因也很簡(jiǎn)單：

初始token對(duì)幾乎所有后續(xù)token都是可見(jiàn)的，因?yàn)樽曰貧w語(yǔ)言模型的性質(zhì)，使初始token更容易被訓(xùn)練為「注意力池」。

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

這是一個(gè)簡(jiǎn)單而高效的框架，使得通過(guò)有限注意力窗口訓(xùn)練的LLM能夠處理無(wú)限長(zhǎng)度的文本，而無(wú)需進(jìn)行微調(diào)。

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

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

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

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

具體來(lái)說(shuō)，研究人員建議在所有訓(xùn)練樣本的開(kāi)頭添加一個(gè)額外的可學(xué)習(xí)token可以作為指定的注意力接收器。

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

Streaming LLM

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

雖然窗口注意力技術(shù)在推理過(guò)程中提供了效率，但它導(dǎo)致了極高的語(yǔ)言建模復(fù)雜性。因此，該模型的性能不適合部署在流應(yīng)用程序中。

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

識(shí)別困惑度激增點(diǎn)

上圖顯示了 20K token文本的語(yǔ)言建模的復(fù)雜性。很明顯，當(dāng)文本長(zhǎng)度超過(guò)緩存大小時(shí)，由于排除初始token而導(dǎo)致困惑度度激增。這表明初始token，無(wú)論與預(yù)測(cè)token的距離如何，對(duì)于維持 LLM 的穩(wěn)定性至關(guān)重要。

為什么LLM在刪除初始token的KV時(shí)會(huì)崩潰？

研究人員將 Llama-2-7B 和模型的所有層和頭部的注意力圖可視化，如圖 2 所示。研究人員發(fā)現(xiàn)，除了底部?jī)蓪又猓Ｐ褪冀K關(guān)注所有層和頭部的初始token。

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

對(duì)于語(yǔ)言模型中初始token的重要性，有兩種可能的解釋：（1）它們的語(yǔ)義至關(guān)重要，或者（2）模型學(xué)習(xí)到對(duì)其絕對(duì)位置的偏見(jiàn)。

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

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

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

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

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

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

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

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

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

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

3.2 使用注意力接收器滾動(dòng) KV 緩存

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

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

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

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

2）滾動(dòng)KV緩存保留最新的token，這對(duì)于語(yǔ)言建模至關(guān)重要。StreamingLLM 的設(shè)計(jì)用途廣泛，可以無(wú)縫地融入任何采用相對(duì)位置編碼的自回歸語(yǔ)言模型，例如 RoPE和 ALiBi。

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

3.3 使用注意力池預(yù)訓(xùn)練 LLMS

正如上文所述，模型過(guò)度關(guān)注多個(gè)初始token的一個(gè)重要原因是缺乏指定的接收器token來(lái)卸載過(guò)多的注意力分?jǐn)?shù)。因此，該模型無(wú)意中將全局可見(jiàn)的token（主要是初始token）指定為注意力池。

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

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

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

為了進(jìn)行驗(yàn)證，研究人員在相同的設(shè)置下從頭開(kāi)始預(yù)訓(xùn)練具有 1.6 億個(gè)參數(shù)的三種語(yǔ)言模型。

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

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

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

作者介紹

Guangxuan Xiao

MIT EECS 博士一年級(jí)在讀，導(dǎo)師為韓松教授。

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

田淵棟

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

他是2021年ICML杰出論文榮譽(yù)獎(jiǎng)和2013年ICCV Marr獎(jiǎng)榮譽(yù)獎(jiǎng)的第一作者獲得者，并且是該項(xiàng)目的首席科學(xué)家和工程師。

他還獲得了 2022 年 CGO 杰出論文獎(jiǎng)。

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

BEIDI CHEN

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

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

韓松

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

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