谷歌TurboQuant力圖突破AI內存瓶頸

高帶寬內存的激增需求和供應鏈約束正在推高AI和HPC基礎設施成本，也推動了系統設計的創新需求。內存已成為當前AI周期的決定性約束之一。這種約束不僅限于硬件的可用性，還影響著現代AI模型的運行方式。

現在，谷歌的研究人員正在解決大語言模型推理中最消耗內存的組件之一：鍵值（KV）緩存。在最近的一篇研究論文中，該公司介紹了TurboQuant，這是一種新的壓縮方法，旨在在不犧牲模型準確性的情況下，減少推理期間KV緩存的內存占用。該方法旨在實現作者描述的壓縮和失真之間的近最優權衡，接近模型數據在不破壞其結構的情況下可以壓縮多少的理論極限。

KV緩存存儲先前token的中間向量表示，以便模型能夠生成響應而無需從頭開始重新計算先前的token。這些向量捕獲了token之間的關系，是注意力機制工作原理的核心，使緩存在長對話或文檔的速度和響應性方面變得至關重要。但這同時也使其成為內存消耗的主要來源。隨著上下文窗口增長到數萬或數十萬token，緩存相應擴展，即使是配置良好的系統也可能會被內存需求迅速壓倒。

谷歌的方法是以非常低的精度壓縮此緩存，同時保持使注意力機制起作用的數學特性。根據該公司的說法，TurboQuant可以將KV緩存內存使用減少大約六倍，在某些情況下將數據表示降低到每個值僅幾位。重要的是，這種壓縮不需要重新訓練模型或在校準數據上進行微調。作者聲稱，它可以在推理時直接應用，對準確性的影響最小。

高壓縮、最小準確性損失和無重訓練要求的結合，是使這項研究值得注意的原因。壓縮技術并不新鮮，量化已被廣泛用于縮小模型權重。但KV緩存壓縮已被證明更加困難。涉及的數據結構是高維的，對失真敏感。小誤差可以通過注意力計算傳遞并降低輸出質量。

TurboQuant通過兩部分方法解決了這個問題。第一步稱為PolarQuant，將向量表示轉換為一種可以以非常低精度更高效壓縮的形式。第二步應用基于Johnson-Lindenstrauss引理的輕量級校正機制，這是一種在保持高維空間中距離的數學技術。這一步驟補償了在壓縮期間引入的失真，有助于保持向量之間的相互關系，注意力機制使用這些關系來確定哪些token最重要。

結果是這個系統實現了比典型方法更高水平的壓縮，而沒有引入通常伴隨低精度的偏差或不穩定性。實際上，這允許更多token同時存儲在內存中，并使相同的工作負載能夠在更少的硬件上運行。

其影響可能不僅限于模型效率。TurboQuant方法可以在推理期間直接應用，不需要重新訓練、校準數據或更改模型架構，這一事實可能使其更容易集成到現有系統中，而無需重新設計模型的構建或服務方式。

該工作還達到了歷史上難以實現的壓縮水平。先前KV緩存壓縮的方法在非常低的比特表示下難以保持穩定性，而TurboQuant在每值約3至3.5比特的比特寬度下保持穩定的性能。

這里仍然存在明顯的限制。結果基于基準評估而非生產規模系統，且該方法僅針對整體內存占用的一部分。模型權重、激活和其他系統開銷保持不變，對高帶寬內存的需求不太可能消失。即便如此，該研究反映了推理時效率的不斷擴大關注。隨著模型規模的持續擴大，像TurboQuant這樣的技術表明，不僅可以新硬件來解決內存的一些壓力，還可以通過更高效地使用推理期間計算的中間表示。