盡管大語言模型（LLMs）作為通用 AI 工具取得了成功，但其對計算資源的高需求使其在許多真實場景中的部署具有挑戰性。模型的大小和對話狀態受到可用高帶寬內存的限制，限制了可服務的用戶數量和最大對話長度。
目前，兩種最熱門的 LLM 架構是 Transformers 和選擇性狀態空間模型 (SSMs)，它們位于頻譜的另一端：
- Transformer：對話狀態由序列中每個元素的不同表示組成，其大小會迅速爆炸。
- SSMs：將整個序列壓縮為單個表示形式，由于容量有限，可能會忘記過去的信息。對話狀態的壓縮可釋放內存，對于在相同內存限制下運行更大的模型、一次處理更多的 tokens 或降低延遲至關重要。
為此，NVIDIA 的研究人員開發了一項稱為動態內存壓縮（Dynamic Memory Compression，DMC）的新技術，該技術可以大幅提高 LLMs 部署的效率，并在不會耗盡內存的情況下將其視野擴展到更長的序列。
DMC 開啟了第三種方式，即訓練 Transformer 模型以自適應壓縮對話狀態并實現所需的壓縮率。這樣可以在不替換熟悉的 Transformer 架構的情況下顯著減少對話狀態大小。DMC 不需要從頭開始訓練，因為可以通過少量額外訓練來改裝現有模型，這比容易出錯的無訓練方法更可靠。
哪些因素會影響 LLM 推理性能？
LLM 推理包含兩個階段：
- 預填充：用戶查詢被輸入。
- 自動回歸生成：響應一次生成一個 token。
在生成過程中，為執行自注意力操作，Transformers 會為緩存中的每個 token 附加一對表示 (鍵值對或 KVP)。系統會為每個層和注意力頭存儲不同的 KVP。
因此，KVP 緩存會隨序列長度按比例增長。由于 KVP 緩存必須與 LLM 權重一起裝入 GPU 內存，因此 KVP 緩存可能會占用顯存的大部分空間，甚至耗盡顯存。
此外，KVP 緩存越大，執行單個推理步驟所需的時間越長。這是因為計算注意力分數是受內存限制的運算。每個查詢都有自己要加載的 KVP 緩存。
注意力層或 FFN 層的線性投影的情況則不同，如果 GPU 同時并行處理許多查詢，則所有查詢都必須將 每個權重矩陣從 HBM 加載到 SRAM 中一次 。
過去的研究試圖通過 量化 KVP 緩存的表征、 共享 注意力頭或從其中 移除令牌 來減小 KVP 緩存的大小。但是，這些方法會降低原始性能，因為它們會從內存中刪除信息，而不會改變原始 LLM 的行為。
動態顯存壓縮?
動態內存壓縮 (DMC) 是一種在推理期間壓縮 KV 緩存的簡單方法，不會導致性能下降。模型使用以下公式分別決定每個 token、層和頭，以及是否應將新的鍵值對附加到 KV 緩存中 (就像在普通 Transformer 中一樣)，還是累加到最后一個：
這個方程是 DMC 的核心，可將密鑰的子序列轉換為特定的前綴和，這讓人聯想到 xLSTM 或 RWKV 等熱門 SSM。
在推理過程中，alpha 的值嚴格是二進制的。當 
在普通模型中，緩存一次擴展一個 KVP。借助 DMC，決策變量可確定是否應擴展緩存，還是應將新的緩存對與 KVP 緩存中的最后一個緩存對合并。
DMC 的一個關鍵組成部分是 改造 過程，該過程教會了 LLM 這種行為：
- 使用 2 – 8%的原始訓練數據混合訓練預先存在的 LLM (例如 Llama 系列的 LLM)。
- 初始化決策模塊，以便始終將 KVP 附加到 KVP 緩存中
- 通過對新配對與追蹤配對的平均值施加壓力，逐漸過渡到 DMC。在改裝過程中，目標壓縮率從 1x 提升到所需水平。
- 達到目標壓縮率后，進行修復，以完成改裝和整合的最后步驟。
決定是否附加或合并是獨立的。要使用梯度下降訓練 LLMs，您可以通過 Gumbel-Sigmoid 分布持續松此決策，從而在訓練期間生成部分附加和部分合并的內存元素。
為了模擬 token 無法進入中間合并階段的推理時間 DMC (圖 2)，您可以增強 Transformers 的因果關系掩碼，以越來越多地阻止對這些狀態的訪問。
總體而言，此過程可讓您在繼續預訓練期間避免出現困惑峰值。事實上，我們發現 像通用查詢注意力（general query-attention） 這樣的方法在繼續預訓練之前會顯著改變模型架構，很難恢復原始性能。
成果?
對于我們進行基準測試的所有下游應用，包括 真實性 (MMLU)、 常識問答和編碼 (HumanEval)，DMC 的性能可與普通模型相比。
至關重要的是，DMC 可以與量化和 GQA 等其他 KVP 緩存縮減方法結合使用。
| ? | DMC 壓縮率 | MMLU (真實性) | 常識問答 | HumanEval (編碼) |
| Llama-2 -7B | 1 倍 | 44.6 | 70.5 | 14.0 |
| Llama-2 -7B | 4 倍 | 44.2 | 70.2 | 16.5 |
| Llama-2 -7B | 8 倍 | 41.8 | 70.1 | 16.5 |
| Llama-2 -13B | 1 倍 | 54.5 | 73.5% | 17.5 |
| Llama-2 -13B | 4 倍 | 54.2 | 73.2 | 22.0 |
| Llama-2 -13B | 8 倍 | 52.1 | 73.3% | 21.3% |
| Llama-2 -70B+GQA | 8 倍 | 68.8 | 78.0 | 29.6 |
| Llama-2 -70B *GQA |
16 倍 | 68.8 | 77.9 | 29.9% |
| ? | Llama-2 -7B+DMC 4x | Llama-2-7B + DMC 4x (8 位量化) | Llama-2 -7B+DMC 8x | Llama-2-7B + DMC 8x (8 位量化) (8-bit quantization) |
| MMLU (真實性) | 44.2 | 44.6 | 41.8 | 41.7 |
在某些情況下，性能甚至會得到提升。我們推測這是由于緩解壓縮的現象，例如 中間丟失 的 tokens，不過還需要更多證據來證實這一點。
對話狀態的壓縮釋放了內存，而內存可用于容納更大的批量大小。由于模型受內存限制，這通常幾乎直接轉化為同等的吞吐量增加。
例如，選擇適合內存的最大批量大小 (在 NVIDIA H100 GPU 上壓縮 8 倍) 意味著每秒生成的令牌比 Vanilla 模型多 700%。
在達到一定的壓縮率后，計算注意力分數會受到計算限制，因此延遲可能會增加，同時提供更多的查詢。
在圖 4 中，x 軸顯示了適用于 Vanilla LLM 以及 DMC 2x、4x、6x 和 8x 模型的單個 GPU (7B 和 13B) 或兩個具有張量并行度 (70B) 的 GPU 上顯存的最大批量大小。
如果批量大小保持不變，延遲會降低。在實踐中，調整批量大小有助于優化延遲吞吐量權衡。
在圖 5 中，在特定批量大小下，模型受內存限制，可以相同延遲處理其他示例。除此之外，模型的計算能力也越來越受限。我們在 NVIDIA H100 GPU 上對 4K 令牌序列的最后 1K 令牌的生成進行了測量。
由于能夠以不同的速率靈活地壓縮不同的注意力頭和層，DMC 還揭示了有關 LLMs 內部功能的新見解。例如，頂層受到的壓縮最多，因為其 KVP 是最冗余的。
在圖 6 中，頭部按最高壓縮級別到最低自上而下級別排列。
結束語?
DMC 有望進一步推動 LLM 的發展。
它提供了一種自適應顯存機制，該機制介于 Transformer 中 的線性顯存和 SSM 的常量顯存之間，并能更好地在容量和空間之間進行權衡。這在不犧牲質量的情況下提高了 Transformer LLM 的吞吐量，并且能夠在相同的硬件限制條件下適應更長的上下文。這對于通過思維鏈實現推理時擴展至關重要，因為這需要生成長序列的 token。
DMC 還開辟了新的途徑，通過重新利用現有架構，以經濟實惠的方式創建原型和生產級 LLM。 使用現成的代碼 改造您自己的 Transformer LLM，或者閱讀 論文 以深入了解 Dynamic Memory Compression。
有關更多信息，請參閱以下資源：
致謝?
在此，我們要感謝 Mostofa Patwary 和 Szymon Migacz 所提供的幫助，并感謝 Przemys?aw Strzelczyk、Daniel Korzekwa 和 Bryan Catanzaro 在發表此論文時提供的有益討論和支持。這項工作得到了 UKRI Centre for Doctoral Training in Natural Language Processing 的部分支持，該中心由 UKRI (grant EP/S022481/1) 和愛丁堡大學信息學學院以及哲學、心理和語言科學學院提供資金。
?
