PN-Mix: Pair-Norm Guided Mixed-Precision Key Cache Quantization for Efficient Long-Output LLM Inference
作者: Gil Jae Song, Sungjoo Yoo
分类: cs.LG, cs.CL
发布日期: 2026-10-14
备注: EMNLP 2026
💡 一句话要点
提出PN-Mix以解决长输出LLM推理中的KV缓存量化问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长输出推理 KV缓存 混合精度量化 RoPE 大型语言模型 模型优化 性能提升
📋 核心要点
- 长输出推理在大型语言模型中存在KV缓存内存和带宽瓶颈,低位量化导致模型准确性显著下降。
- PN-Mix通过在预填充阶段识别显著通道对,采用混合精度量化来优化RoPE解码器的性能。
- 实验结果表明,PN-Mix在2.26位平均精度下的表现优于其他混合精度方法,并在1.5位时表现稳定。
📝 摘要(中文)
长输出推理在大型语言模型中面临严重的KV缓存内存和带宽瓶颈。虽然KV缓存量化可以缓解这一问题,但低位量化(如2位)会显著降低模型准确性。最近的混合精度键缓存量化方法通过为重要键通道分配更高的精度来减轻这种降级,但仅使用键统计来估计通道显著性。通过对RoPE通道对的注意力对数扰动分析,我们发现具有大查询-键对范数乘积的通道对的量化误差显著影响对数变化。基于这一洞察,我们提出了PN-Mix,一种针对基于RoPE的解码器LLM的混合精度键缓存量化方法,该方法在预填充阶段识别显著通道对并在解码过程中重复使用。PN-Mix在标准推理、代码生成和AIME基准测试中以2.26位的平均精度实现了与先前混合精度方法(2.32-2.39位)相当或更好的平均准确性。
🔬 方法详解
问题定义:长输出推理中,KV缓存的内存和带宽限制导致性能瓶颈,现有的低位量化方法(如2位)会显著降低模型的准确性。
核心思路:PN-Mix通过在预填充阶段识别重要的键通道对,并为其分配更高的量化精度,从而减轻量化误差对模型性能的影响。
技术框架:PN-Mix的整体架构包括预填充阶段的显著性通道识别和后续的解码阶段,确保在解码过程中重复使用这些显著通道。
关键创新:PN-Mix的创新在于通过对RoPE通道对的注意力对数扰动分析,识别出对模型输出影响最大的通道对,从而实现更高效的混合精度量化。
关键设计:PN-Mix在量化过程中使用了2.26位的平均精度,并在1.5位时表现出优越的稳定性,相较于其他基线方法,PN-Mix在低于2.25位时性能下降更为平缓。
🖼️ 关键图片
📊 实验亮点
PN-Mix在标准推理、代码生成和AIME基准测试中以2.26位的平均精度实现了与其他混合精度方法(2.32-2.39位)相当或更好的准确性。此外,PN-Mix在1.5位时的性能下降较为平缓,显示出其优越的稳定性。
🎯 应用场景
PN-Mix的研究成果在大型语言模型的推理任务中具有广泛的应用潜力,尤其是在需要高效处理长输出的场景,如对话系统、代码生成和文本生成等领域。其优化的KV缓存量化方法能够有效提升模型的推理效率和准确性,具有重要的实际价值和未来影响。
📄 摘要(原文)
Long-output inference in large language models suffers from severe KV cache memory and bandwidth bottlenecks. While KV cache quantization alleviates this issue, low-bit quantization (e.g., 2-bit) severely degrades model accuracy. Recent mixed-precision key-cache quantization methods mitigate this degradation by allocating higher precision to important key channels, but estimate channel saliency using only key statistics. Through attention logit perturbation analysis at the RoPE channel-pair level, we show that quantization errors in channel pairs with large query--key pair-norm products contribute significantly to logit variation. Based on this insight, we propose PN-Mix, a mixed-precision key-cache quantization method for RoPE-based decoder-only LLMs, which allocates higher precision to salient channel pairs identified once during the prefill stage and reused throughout decoding. Across standard reasoning, code-generation, and AIME benchmarks, PN-Mix with 2.26-bit average precision achieves comparable or better average accuracy than prior mixed-precision methods using 2.32--2.39 bits. Furthermore, in our Pareto sweep, PN-Mix degrades gracefully down to 1.5 bits, whereas the compared baselines exhibit sharper performance drops below 2.25 bits.