Quantize What Counts: More for Keys, Less for Values
作者: Mohsen Hariri, Alan Luo, Weicong Chen, Shaochen Zhong, Tianyi Zhang, Qifan Wang, Xia Hu, Xiaotian Han, Vipin Chaudhary
分类: cs.LG
发布日期: 2025-02-20 (更新: 2025-10-15)
🔗 代码/项目: GITHUB
💡 一句话要点
提出基于几何理论的KV量化方法以优化LLM推理性能
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 KV缓存 量化方法 推理优化 几何理论 信息密度 精度分配
📋 核心要点
- 现有的大型语言模型在推理过程中受到KV缓存的内存瓶颈影响,尤其是随着模型规模和上下文长度的增加。
- 本文提出通过理论定理指导的混合精度KV量化方法,强调在分配比特时优先考虑键的精度。
- 实验结果表明,采用键优先的比特分配策略可以在保持高准确率的同时显著节省内存资源。
📝 摘要(中文)
大型语言模型(LLMs)在推理时面临由注意力键值(KV)缓存引起的内存瓶颈。虽然KV缓存量化可以缓解这一成本,但键和值之间的比特分配通常是启发式调整,缺乏理论基础和普适性。本文提出两个定理,将混合精度KV量化锚定在Transformer模型的内在几何中。首先,键的投影在谱范数和弗罗贝尼乌斯范数上系统性地大于值矩阵,意味着键路径上的信息密度更高。其次,在给定的内存预算下,优先考虑键的精度而非值可以严格减少量化误差并更好地保持准确性。实证评估显示,键优先分配(如4位键,2位值)相比均匀分配(如均为4位)可保留高达98.3%的准确性,同时节省内存。这些结果将比特分配从随意调整转变为一个理论基础、几何驱动的高效LLM推理设计原则。
🔬 方法详解
问题定义:本文旨在解决大型语言模型推理时KV缓存引起的内存瓶颈问题。现有方法在键和值的比特分配上缺乏理论指导,导致量化效果不佳。
核心思路:论文提出的核心思路是通过理论定理来指导KV量化,强调键的精度优先于值,从而减少量化误差并提高模型准确性。
技术框架:整体架构包括两个主要模块:首先是对Transformer模型的几何特性进行分析,确定键和值的精度分配;其次是基于此分析进行KV缓存的量化设计。
关键创新:最重要的技术创新在于提出了两个定理,系统性地证明了键的投影在信息密度上优于值,从而为比特分配提供了理论基础。与现有方法相比,本文的方法更具普适性和理论指导性。
关键设计:在比特分配上,本文建议采用4位键和2位值的配置,具体参数设置经过理论推导,确保在给定内存预算下最大限度地保留模型的准确性。损失函数的设计也考虑了量化误差的最小化。
🖼️ 关键图片
📊 实验亮点
实验结果显示,采用键优先的比特分配策略(如4位键和2位值)相比于均匀分配(如均为4位),在保持高达98.3%准确率的同时,显著节省了内存资源。这一结果表明,理论指导的量化方法在实际应用中具有显著优势。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和机器翻译等大型语言模型的推理优化。通过优化KV缓存的量化策略,可以显著提高模型的推理效率,降低内存消耗,进而推动更大规模模型的实际应用。未来,随着模型规模的不断扩大,该方法可能在多种AI应用中发挥重要作用。
📄 摘要(原文)
Large Language Models (LLMs) suffer inference-time memory bottlenecks dominated by the attention Key-Value (KV) cache, which scales with model size and context length. While KV-cache quantization alleviates this cost, bit allocation between keys and values is often tuned heuristically, lacking theoretical grounding and generalizability. This paper proposes two theorems that anchor mixed-precision KV quantization in the intrinsic geometry of Transformer models. First, key projections systematically have larger spectral and Frobenius norms than value matrices, implying higher information density along the key path. Second, for any given memory budget, prioritizing precision for keys over values strictly reduces quantization error and better preserves accuracy. Empirical evaluations across various prominent LLMs and benchmarks show that key-favored allocations (e.g., 4-bit keys, 2-bit values) retain up to 98.3\% accuracy compared to uniform allocations (e.g., 4-bit for both), while conserving memory. These results transform bit allocation from ad hoc tuning into a theoretically grounded, geometry-driven design principle for efficient LLM inference. Source code is available at https://github.com/mohsenhariri/spectral-kv.