LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers

📄 arXiv: 2611.05304v1 📥 PDF

作者: Minsoo Cheong, Sungjoo Yoo

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-11-03


💡 一句话要点

提出LoopVQ以解决循环变换器的KV缓存量化问题

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 循环变换器 KV缓存 向量量化 内存优化 语言模型 推理速度 吞吐量提升

📋 核心要点

  1. 现有循环变换器模型在KV缓存方面存在冗余,导致内存瓶颈。
  2. LoopVQ通过联合编码多个循环步骤的KV缓存,优化了量化过程。
  3. 在实验中,LoopVQ在多个任务上表现优异,显著提高了解码速度和吞吐量。

📝 摘要(中文)

循环变换器是一种参数高效的语言建模替代方案,通过在多个循环步骤中重复应用相同的层,最近的循环语言模型在性能上与标准模型相匹配。尽管这种架构有效减少了参数内存,但每个循环步骤仍然存储自己的键值(KV)缓存,导致KV缓存随着循环步骤的增加而线性增长,成为主要的内存瓶颈。为此,本文提出了LoopVQ,一种针对循环变换器的KV缓存向量量化方法,通过联合编码多个循环步骤的通道,并微调码本以匹配全精度模型的输出。在每个元素0.5位的情况下,LoopVQ在一般的推理和长上下文检索任务中超越了现有的KV缓存向量量化方法。

🔬 方法详解

问题定义:本文旨在解决循环变换器模型中KV缓存的冗余问题,现有方法在处理多个循环步骤时,KV缓存线性增长,造成内存使用效率低下。

核心思路:LoopVQ通过联合编码多个循环步骤的KV缓存,减少冗余并优化量化过程,针对不同循环步骤的量化敏感性进行微调,以提高整体性能。

技术框架:LoopVQ的整体架构包括KV缓存的联合编码模块和微调模块。联合编码模块负责将多个循环步骤的KV缓存进行编码,而微调模块则确保编码后的输出与全精度模型的输出一致。

关键创新:LoopVQ的主要创新在于其联合编码策略,能够有效减少KV缓存的冗余,与传统的单步量化方法相比,显著提高了内存效率和模型性能。

关键设计:在设计中,LoopVQ采用了0.5位每元素的量化策略,并结合TaSQ方法,确保在低位量化时仍能保持较高的BF16精度。

🖼️ 关键图片

fig_0

📊 实验亮点

在实验中,LoopVQ在0.5位每元素的情况下,超越了现有的KV缓存向量量化方法,尤其在长链推理和长上下文检索任务中表现突出。与BF16基线相比,LoopVQ在解码延迟上几乎减半,并在1.4B和2.6B Ouro模型上提高了峰值服务吞吐量,分别达到4.3倍和6.1倍的提升。

🎯 应用场景

LoopVQ的研究成果在大规模语言模型的部署中具有重要应用价值,尤其是在需要高效内存管理和快速推理的场景,如对话系统、智能助手和长文本生成等领域。未来,LoopVQ有望推动循环变换器在更多实际应用中的普及。

📄 摘要(原文)

Looped transformers are a parameter-efficient alternative to standard transformers for language modeling. By reapplying the same layers across multiple loop steps, recent looped language models match the performance of non-looped models several times their size. While this architecture effectively reduces parameter memory, each loop step still stores its own keys and values, so the key-value (KV) cache grows linearly with the number of loop steps and becomes the main memory bottleneck. We observe that the KV cache of looped models is highly redundant along the loop axis and that the loop steps differ substantially in their sensitivity to quantization. Based on these observations, we propose $\textbf{LoopVQ}$, a KV cache vector quantization method for looped transformers that codes channels of multiple loop steps jointly, and fine-tunes the codebooks to match the output of the full-precision model. At 0.5 bits per element, LoopVQ outperforms existing KV cache vector quantization methods for standard LLMs on general, long chain-of-thought reasoning, and long-context retrieval tasks with Ouro and Huginn models. Its design is orthogonal to these methods, and combined with TaSQ, it retains much of the BF16 accuracy even at 0.25 bits per element. We also implement LoopVQ in SGLang, where it nearly halves the decode latency at a 32K-token context and raises the peak serving throughput by up to $4.3\times$ and $6.1\times$ on the 1.4B and 2.6B Ouro models compared to the BF16 baseline.