如果你體驗過與任何一款對話式 AI 機(jī)器人的交流，你一定能想起某些極具「挫敗感」的時刻。比如，你在前一天的對話中講述過的要點，被 AI 忘得干干凈凈……

這是因為當(dāng)前的多數(shù) LLM 只能記住有限的上下文，就像為考試而臨時抱佛腳的學(xué)生，稍加盤問就會「露出馬腳」。

想象一下，如果 AI 助手在聊天中能夠根據(jù)上下文參考幾周或幾個月前的對話，或者，你可以要求 AI 助手總結(jié)長達(dá)數(shù)千頁的報告，這樣的能力是不是令人羨慕？

為了讓 LLM 記住更多、記得更好，研究者們正在不斷努力。最近，來自 MIT、Meta AI、CMU 的研究者提出了一種名為「StreamingLLM」的方法，使語言模型能夠流暢地處理無窮無盡的文本。

• 論文地址：https://arxiv.org/pdf/2309.17453.pdf
• 項目地址：https://github.com/mit-han-lab/streaming-llm

StreamingLLM 的工作原理是識別并保存模型固有的「注意力池」（attention sinks）錨定其推理的初始 token。結(jié)合最近 token 的滾動緩存，StreamingLLM 的推理速度提高了 22 倍，而不需要犧牲任何的準(zhǔn)確性。短短幾天，該項目在 GitHub 平臺已斬獲 2.5K 星：

具體地說，StreamingLLM 使得語言模型能準(zhǔn)確無誤地記住上一場比賽的觸地得分、新生兒的名字、冗長的合同或辯論內(nèi)容，就像升級了 AI 助理的內(nèi)存，可以完美地處理更繁重的工作。

接下來讓我們看看技術(shù)細(xì)節(jié)。

方法創(chuàng)新

通常，LLM 在預(yù)訓(xùn)練時受到注意力窗口的限制。盡管為擴(kuò)大這一窗口大小、提高訓(xùn)練和推理效率，此前已有很多工作，但 LLM 可接受的序列長度仍然是有限的，這對于持久部署來說并不友好。

在這篇論文中，研究者首先介紹了 LLM 流應(yīng)用的概念，并提出了一個問題：「能否在不犧牲效率和性能的情況下以無限長輸入部署 LLM？」

將 LLM 應(yīng)用于無限長輸入流時，會面臨兩個主要挑戰(zhàn)：

1、在解碼階段，基于 transformer 的 LLM 會緩存所有先前 token 的 Key 和 Value 狀態(tài)（KV），如圖 1 (a) 所示，這可能會導(dǎo)致內(nèi)存使用過多，并增加解碼延遲；

2、現(xiàn)有模型的長度外推能力有限，即當(dāng)序列長度超過預(yù)訓(xùn)練時設(shè)定的注意力窗口大小時，其性能就會下降。

一種直觀的方法被稱為窗口注意力（Window Attention）（如圖 1 b），這種方法只在最近 token 的 KV 狀態(tài)上保持一個固定大小的滑動窗口，雖然能確保在緩存填滿后仍能保持穩(wěn)定的內(nèi)存使用率和解碼速度，但一旦序列長度超過緩存大小，甚至只是驅(qū)逐第一個 token 的 KV，模型就會崩潰。另一種方法是重新計算滑動窗口（如圖 1 c 所示），這種方法會為每個生成的 token 重建最近 token 的 KV 狀態(tài)，雖然性能強(qiáng)大，但需要在窗口內(nèi)計算二次注意力，因此速度明顯更慢，在實際的流應(yīng)用中并不理想。

在理解窗口注意力失效的過程中，研究者發(fā)現(xiàn)了自回歸 LLM 的一個有趣現(xiàn)象：如圖 2 所示，大量注意力分?jǐn)?shù)被分配給了初始 token，而不管這些 token 與語言建模任務(wù)是否相關(guān)。

研究者將這些 token 稱為「注意力池」：盡管它們?nèi)狈φZ義上的意義，但卻占據(jù)了大量的注意力分?jǐn)?shù)。研究者將這一現(xiàn)象歸因于于 Softmax（要求所有上下文 token 的注意力分?jǐn)?shù)總和為 1），即使當(dāng)前查詢在許多以前的 token 中沒有很強(qiáng)的匹配，模型仍然需要將這些不需要的注意力值分配到某處，從而使其總和為 1。初始 token 成為「池」的原因很直觀：由于自回歸語言建模的特性，初始 token 對幾乎所有后續(xù) token 都是可見的，這使得它們更容易被訓(xùn)練成注意力池。

基于上述洞察，研究者提出了 StreamingLLM，這是一個簡單而高效的框架，它可以讓使用有限注意力窗口訓(xùn)練的注意力模型在不進(jìn)行微調(diào)的情況下處理無限長的文本。

StreamingLLM 利用了注意力池具有高注意力值這一事實，保留這些注意力池可以使注意力分?jǐn)?shù)分布接近正態(tài)分布。因此，StreamingLLM 只需保留注意力池 token 的 KV 值（只需 4 個初始 token 即可）和滑動窗口的 KV 值，就能錨定注意力計算并穩(wěn)定模型的性能。

使用 StreamingLLM，包括 Llama-2-[7,13,70] B、MPT-[7,30] B、Falcon-[7,40] B 和 Pythia [2.9,6.9,12] B 在內(nèi)的模型可以可靠地模擬 400 萬個 token，甚至更多。

與唯一可行的 baseline—— 重新計算滑動窗口相比，StreamingLLM 的速度提高了 22.2 倍，而沒有損耗性能。

測評

在實驗環(huán)節(jié)，如圖 3 所示，在跨度為 20K token 的文本上，StreamingLLM 的困惑度可以與 Oracle 基線（重新計算滑動窗口）相媲美。同時，當(dāng)輸入長度超過預(yù)訓(xùn)練窗口時，密集注意力就會失效，而當(dāng)輸入長度超過緩存大小時，窗口注意力就會陷入困境，導(dǎo)致初始 token 被剔除。

圖 5 進(jìn)一步證實了 StreamingLLM 可以可靠地處理非常規(guī)規(guī)模的文本，包括 400 多萬個 token，涵蓋了各種模型系列和規(guī)模。這包括 Llama-2-[7,13,70] B、Falcon-[7,40] B、Pythia-[2.8,6.9,12] B 和 MPT-[7,30] B。

隨后，研究者證實了「注意力池」的假設(shè)，并證明語言模型可以通過預(yù)訓(xùn)練，在流式部署時只需要一個注意力池 token。具體來說，他們建議在所有訓(xùn)練樣本的開頭多加一個可學(xué)習(xí)的 token，作為指定的注意力池。通過從頭開始預(yù)訓(xùn)練 1.6 億個參數(shù)的語言模型，研究者證明了本文方法可以保持模型的性能。這與當(dāng)前的語言模型形成了鮮明對比，后者需要重新引入多個初始 token 作為注意力池才能達(dá)到相同的性能水平。

最后，研究者將 StreamingLLM 的解碼延遲和內(nèi)存使用率與重新計算滑動窗口進(jìn)行了比較，并使用 Llama-2-7B 和 Llama-2-13B 模型在單個英偉達(dá) A6000 GPU 上進(jìn)行了測試。如圖 10 所示，隨著緩存大小的增加，StreamingLLM 的解碼速度呈線性增長。后者解碼延遲則呈二次曲線上升。實驗證明，StreamingLLM 實現(xiàn)了令人印象深刻的提速，每個 token 速度的提升高達(dá) 22.2 倍。

更多研究細(xì)節(jié)，可參考原論文。