大型語言模型(LLM)在文本生成過程中,需要快取所有先前計算過的鍵值(Key-Value, KV)對,這個KV快取的大小與序列長度呈線性增長,成為了服務部署時的主要記憶體瓶頸。透過將KV快取量化為更低的位元數,可以顯著降低成本。然而,當前的量化方法對所有注意力頭分配相同的位寬,忽視了不同注意力頭在模型中重要性的巨大差異。直觀地看,一個自然的想法是給重要的頭分配更多的位元,而給不那麼重要的頭分配更少的位元,即混合精度量化。但最新研究揭示了一個隱藏的陷阱:每個量化器遵循不同的失真曲線 D(b)=α·β^{-b},其中衰減率β在不同量化器設計間的變化範圍從3.6到5.3。如果錯誤地將一種量化器的失真模型應用到另一種量化器上,會導致位元分配順序完全顛倒,從而使效能甚至比均勻量化還要差。這種失敗模式被稱為“失真模型不匹配”。
為了解決這一問題,來自中國的研究團隊提出了RateQuant方法。該方法首先使用一個小的校準集來擬合每個量化器的失真模型,然後利用率失真理論中的反水填充(reverse waterfilling)方法,以閉合形式求解最優的位元分配問題。實驗結果表明,在Qwen3-8B模型上,以平均2.5位元為目標時,經過校準的RateQuant將KIVI量化方法的困惑度(perplexity)從49.3大幅降低到14.9,降幅達到70%,並在QuaRot方法的基礎上進一步提升了6.6的困惑度。整個校準過程在單個GPU上僅需1.6秒,並且在推理階段沒有任何額外開銷。
這項成果為大語言模型的高效部署提供了全新的思路,表明基於率失真理論的混合精度量化可以在不增加推理成本的前提下顯著提升模型效能。論文共18頁,包含7張圖和5張表格,為後續研究提供了豐富的參考。