🕳️ LLM 中的注意力池（Attention Sinks）實現無限流暢性

社群文章釋出於 2023 年 10 月 9 日

摘要

目錄

聊天助手 LLM 的侷限性

視窗注意力

注意力池 (Attention Sinks)
注意力池 - 困惑度實驗

注意力池 - 無盡生成實驗

注意力池 - 聊天助手實驗

注意力池 - 基準測試結論

注意力池的實際應用

常見問題

瞭解更多

引用

摘要

使用帶有注意力池（attention sink）詞元的視窗注意力機制，可以讓預訓練的聊天式大語言模型（LLM），如所有 Llama、Mistral、MPT、Falcon 和 GPT-NeoX (Pythia) 模型，在數百個連續的提示中保持流暢性，這與使用 transformers 載入這些模型時的情況不同。此外，這種方法可以實現恆定的記憶體使用，而大多數用 transformers 載入的 LLM 具有線性空間複雜度，會導致記憶體問題。

使用這種形式的注意力非常簡單，只需從 attention_sinks 而不是 transformers 匯入你的模型類即可。

from attention_sinks import AutoModel

model = AutoModel.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", device_map="auto")

聊天助手 LLM 的侷限性

大型語言模型（LLM）席捲了整個行業，推動了聊天機器人和虛擬助手領域的發展。LLM 似乎特別擅長扮演（專業化的）個人助理角色，但它們也存在各種侷限性。在這篇博文中，我們將重點關注以下兩個主要限制：

視訊記憶體使用：許多 LLM（例如 Llama 2）在推理時存線上性空間複雜度的問題。在聊天助手場景中，這意味著裝置的視訊記憶體限制將制約使用者持續進行順序提示的能力。
流暢性喪失：迄今為止訓練的所有 LLM 在輸入過長時都會喪失流暢性。當這種情況發生時，模型將失去生成語言的能力，並開始生成例如無盡的換行符、任意字元（0OgOATO0OATO）、損壞的 Unicode（��）或重複的單詞（assistant: assistant: assistant: assistant:）。

大多數 LLM 在輸入長度超過其預訓練長度後會出現這種行為。例如，Llama 2 7B 在超過 4096 個詞元后會遇到此問題，而 Mistral-7B-v0.1 在大約 1 萬個詞元后會失去流暢性。

這些侷限性在實踐中很容易展現出來，例如讓 LLM 根據之前的所有詞元預測一本書的下一個詞元。這種預測的平均負對數似然損失稱為對數困惑度 (log perplexity)，它是衡量 LLM 質量的常用指標。較低的對數困惑度對應較低的平均損失，因此越低越好。視訊記憶體也很容易測量，這兩個指標都在下圖中繪製出來：

Llama-2-7b-hf	Mistral-7B-v0.1

_{有關生成這些圖表的指令碼的更多資訊，請參見困惑度。}

這些侷限性嚴重阻礙了在生產環境中使用 LLM 作為聊天助手的能力。

視窗注意力

應對問題 1（視訊記憶體使用）的一個簡單嘗試是限制輸入到 LLM 的詞元數量。在 transformers 的基礎上構建這個功能是一個相當複雜的過程，但其要點是，每當生成一個詞元時，如果當前大小超過視窗大小，past_key_values 快取就會被縮小到視窗大小。

在我的實驗中，我使用的視窗大小僅為 1024 個詞元。結果如下圖所示：

Llama-2-7b-hf	Mistral-7B-v0.1

視窗注意力確實在生成 1024 個詞元后保持了記憶體使用的恆定，但一旦超過這個視窗大小，對數困惑度立即飆升。這使得它和使用 transformers 載入模型一樣，都是不可行的方法。

注意力池

Xiao 等人（2023）注意到，當應用視窗注意力時，即使在第一個詞元被從視窗中移除後，模型也會立即失去流暢性。他們注意到自迴歸 LLM 的一個有趣現象：最初的幾個詞元佔據了驚人比例的注意力分數，即使這些詞元在語義上並不重要。

這種行為在下圖中視覺化呈現：

除了最初的兩層，幾乎所有的注意力都集中在開頭的幾個詞元上，作者稱之為**注意力池 (attention sinks)**。直觀的解釋是，如果下一個要生成的詞元與之前的任何詞元都不匹配，Softmax 操作仍然會強制注意力分數總和為 1。因此，LLM 學會將注意力分數“解除安裝”到開頭的幾個詞元上。

因此，當視窗注意力機制導致第一個詞元掉出視窗時，LLM 就無法再將注意力分數解除安裝到該詞元上。結果，注意力分數被分散到所有其他詞元上，總和仍然為 1。這導致即使與待生成詞元匹配度不高的詞元也會意外地獲得高注意力分數。其後果是：模型“崩潰”並失去流暢性。

在發現這一現象後，作者提出了一種對視窗注意力的改進，即**始終**保留序列最初的 4 個詞元，也就是注意力池詞元。這可以像這樣視覺化：

此外，在向快取詞元新增位置資訊時，該方法使用快取內的位置，而不是真實文字中的位置。因此，注意力池詞元總是靠近其餘的詞元，從而可以有效地用於解除安裝注意力。

舉一個簡單的例子，我們考慮一個視窗大小為 10 的場景，其中包括 4 個注意力池詞元，文字是一個用空格分隔的字母表。生成時，模型看到的是：

A
A B
A B C
A B C D
A B C D E
A B C D E F
A B C D E F G
A B C D E F G H 
A B C D E F G H I
A B C D E F G H I J
A B C D F G H I J K
A B C D G H I J K L
A B C D H I J K L M
...

分配的位置如下：

0
0 1
0 1 2
0 1 2 3
0 1 2 3 4
0 1 2 3 4 5
0 1 2 3 4 5 6
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7 8
0 1 2 3 4 5 6 7 8 9
0 1 2 3 4 5 6 7 8 9
0 1 2 3 4 5 6 7 8 9
0 1 2 3 4 5 6 7 8 9
...

簡而言之，分配的位置僅取決於快取中的位置，而不是完整文字中的位置。

注意力池 - 困惑度實驗

在我使用注意力池的實驗中，我調整了我的視窗注意力實現，使其包含 4 個永不離開視窗的注意力池詞元，並將視窗大小保持在 1024。結果如下圖所示：

Llama-2-7b-hf	Mistral-7B-v0.1

^{Falcon-7B、MPT-7B 和 Pythia-6.9B 的結果請參見此處。}

結果令人矚目：使用帶有注意力池的視窗注意力的 LLM 兼具兩者的優點：恆定的空間複雜度和穩定的困惑度。Xiao 等人（2023）的研究表明，困惑度在高達 400 萬個詞元的情況下都保持穩定，之後他們的資料就用完了（！）。

請注意，在約 8000 個詞元時，注意力池方法的對數困惑度略高於（即更差）基線。這是因為注意力池僅使用 1024 個詞元的視窗大小。這個視窗大小可以增加到例如 8192 個詞元，這樣對數困惑度將在 8000 個詞元時與基線持平，*並且*將記憶體保持在約 14.85GB 的恆定水平。

注意力池 - 無盡生成實驗

批評者認為，困惑度是衡量 LLM 質量的一個不完美指標，例如因為它並不實際要求模型生成詞元。為了證明注意力池確實有效，我使用 Llama-2-7B，採用本博文中描述的三種方法生成多達 10,000 個詞元：預設（例如 transformers）、windowed 和 attention_sinks。

如果一個模型開始失去流暢性，我就會終止生成。我已將每種方法的完整日誌上傳到我的倉庫。

transformers：完整日誌：該模型在約 1900 個詞元后失去流暢性，並開始無休止地生成損壞的 Unicode 字元，如 🤖🧠👨‍�� ❌。
視窗注意力 (window attention)：完整日誌：該模型在約 1000 個詞元后失去流暢性，生成了數百個換行符，並夾雜著諸如 OOOMMO̶OANOOAMOO̶OMMO ❌ 之類的文字。
attention_sinks 完整日誌：在測試的全部 1 萬個詞元中都保持流暢 ✅。

_{有關復現這些結果的指令碼的更多資訊，請參見無盡生成中的流暢性。}

注意力池 - 聊天助手實驗

注意力池方法非常適合聊天式 LLM 應用，因為它比僅使用 transformers 載入模型時要流暢得多，並且佔用的記憶體也少得多。因此，一個自然的基準測試是在常見的聊天助手場景中對各種方法進行實驗。

在這個基準測試中，我將 MT-Bench 中的連續提示傳送給模型，並自動檢測流暢性何時喪失。這模擬了一個場景：聊天助手在同一歷史記錄中被提示數百次，在此期間模型必須處理數萬個詞元的歷史記錄。

如果一個響應滿足以下條件，我會自動將其歸類為失敗：

包含少於 26 個不同字元，並且
長度超過 1000 個詞元。

在實踐中，這種啟發式方法似乎能準確檢測流暢性的喪失。我已將結果繪製在下圖中：

Llama-2-7b-chat-hf	Mistral-7B-Instruct-v0.1

MPT-7B-chat

_{有關復現這些結果的指令碼的更多資訊，請參見聊天式 LLM 在連續提示下的流暢性。}

對於 Llama-2-7b-chat，transformers 會耗盡視訊記憶體，因此它只能處理少數幾個連續的提示。對於 MPT-7B-chat，當輸入長度超過 2048 時，transformers 會遇到一個 RuntimeError。這些圖表清楚地表明，使用 attention_sinks 載入模型對模型在連續提示下的流暢性有非常積極的影響。然而，正如 Llama-2-7B-chat-hf 的情況所示，它並不能完全避免所有流暢性問題。

注意力池 - 基準測試結論

本博文中描述的基準測試，以及我在我的 attention_sinks 倉庫中描述的針對 MPT、Pythia 和 Falcon 模型的附加基準測試，都清楚地表明，注意力池可以有效地應用於預訓練的 LLM，以應對模型的不穩定性和流暢性喪失。這種額外的穩定性不帶來任何額外成本，甚至可以實現恆定的記憶體使用，而不是大多數 LLM 的線性記憶體使用。

任何希望使用助手式 LLM 的組織或使用者都應該考慮注意力池技術。

注意力池的實際應用

在最先進的研究和從業者可以合理使用的技術之間通常存在顯著的差距。然而，我很高興地說，注意力池可以以幾乎零額外工作量新增到任何預訓練的 LLM 中。

我釋出了 attention_sinks Python 模組，它可以作為 transformers API 的直接替代品。這個 Python 模組支援所有使用 Llama、Mistral、Falcon、MPT 和 GPT-NeoX (Pythia) 架構的模型，可以像這樣使用：

from attention_sinks import AutoModel

model = AutoModel.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", device_map="auto")

這將自動為模型新增一個注意力池 KV 快取 (Attention Sink KV Cache)，它能正確地將注意力池保留在視窗中。您可以使用以下引數配置此快取：

attention_sink_size，int，預設為 4：用作注意力池的初始詞元數量。這些詞元總是包含在注意力池 KV 快取中。
attention_sink_window_size，int，預設為 1020：滑動視窗的大小，即包含在注意力池 KV 快取中的“最近詞元”數量。較大的視窗大小會消耗更多記憶體。不建議將其設定得比 LLM 的上下文視窗大，因為 LLM 仍然只能處理最後的 context window 個詞元。

總視窗大小將是這兩個引數的總和，例如預設為 1024。

例如，載入具有更大視窗大小的 Llama-2-7B-chat 可以這樣做：

from attention_sinks import AutoModel

model = AutoModel.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    device_map="auto",
    attention_sink_size=4,
    attention_sink_window_size=4092,
)

請參閱流式演示，這是一個可以執行的指令碼，用於模擬向您選擇的 LLM 輸入數百個連續提示。（注意，您可能需要更改聊天模板）。

常見問題解答

本 FAQ 主要由 Xiao 等人（2023）撰寫。

對於 LLM 來說，“處理無限長度輸入”意味著什麼？

用 LLM 處理無限長度的文字帶來了挑戰。值得注意的是，儲存所有先前的鍵 (Key) 和值 (Value) 狀態需要大量記憶體，並且模型可能難以生成超出其訓練序列長度的文字。注意力池模型透過僅保留最近的詞元和注意力池，丟棄中間的詞元來解決這個問題。這使得模型能夠從最近的詞元生成連貫的文字，而無需重置快取——這是早期方法所不具備的能力。
LLM 的上下文視窗是否被擴充套件了？

沒有。上下文視窗保持不變。只有最近的詞元和注意力池被保留，中間的詞元被丟棄。這意味著模型只能處理最新的詞元。上下文視窗仍然受其初始預訓練的限制。例如，如果 Llama-2 是用 4096 個詞元的上下文視窗進行預訓練的，那麼基於 Llama-2 的注意力池模型的最大快取大小仍然是 4096。
我能將一篇長文，比如一本書，輸入到注意力池模型中進行摘要嗎？

雖然您可以輸入一篇長文，但模型只會識別最新的詞元。因此，如果輸入的是一本書，注意力池模型可能只會總結結尾的段落，這可能沒有太大意義。如前所述，我們既沒有擴充套件 LLM 的上下文視窗，也沒有增強它們的長期記憶。注意力池模型的優勢在於能夠從最近的詞元生成流暢的文字，而無需重新整理快取。
注意力池模型的理想用例是什麼？

注意力池模型針對流式應用進行了最佳化，例如多輪對話。它非常適合模型需要持續執行，而不需要大量記憶體或依賴過去資料的場景。一個例子是基於 LLM 的日常助手。注意力池模型可以讓模型持續執行，根據最近的對話做出回應，而無需重新整理其快取。早期的方法要麼在對話長度超過訓練長度時需要重置快取（丟失最近的上下文），要麼需要從最近的文字歷史中重新計算 KV 狀態，這可能非常耗時。
注意力池方法與最近關於上下文擴充套件的研究有什麼關係？

注意力池方法與最近的上下文擴充套件方法是正交的，並且可以與它們整合。在注意力池模型的背景下，“上下文擴充套件”指的是使用更大的快取大小來儲存更多最近詞元的可能性。有關實際演示，請參閱論文中的圖 9，其中 LongChat-7B-v1.5-32K 和 Llama-2-7B-32K-Instruct 均已適配注意力池技術。

瞭解更多

檢視以下資源以獲取有關此主題的更多資訊：

我的 attention_sinks 倉庫。
Xiao 等人（2023）的論文《Efficient Streaming Language Models with Attention Sinks》。
MIT HAN 實驗室的 StreamingLLM 研究倉庫。

引用

@article{xiao2023streamingllm,
    title={Efficient Streaming Language Models with Attention Sinks},
    author={Xiao, Guangxuan and Tian, Yuandong and Chen, Beidi and Han, Song and Lewis, Mike},
    journal={arXiv},
    year={2023}
}

社群

透過拖放到文字輸入框、貼上或點選此處上傳圖片、音訊和影片。

點選或貼上此處以上傳圖片

· 註冊或登入以發表評論

贊