arXiv:2607.12550cs.LGcs.CL2026-07

通过联合优化张量分解与残差修复,实现近无损的注意力缓存压缩。

A JoLT for the KV cache: Near-lossless KV cache compression via joint Lagrangian allocation of Tucker ranks and a rotated residual for llms

  • 将缓存视为三阶张量,仅对令牌和特征轴做张量分解压缩。
  • 在2倍压缩下保持困惑度几乎无损,推理准确率与原始模型一致。
  • 适用于长上下文大模型推理,尤其适合资源受限场景。

键值(KV)缓存已成为Transformer推理中的主要内存开销:其大小随批量、上下文长度和层数增长,在长上下文场景下,缓存甚至成为吞吐量瓶颈。现有方法分为两类:低秩方法对缓存的二维切片(如每头矩阵或跨层特征块)进行分解;量化方法则降低每个条目的位宽。二者均未利用缓存在每一层天然为三阶张量的特性——头、令牌、特征三个轴具有显著不同的冗余程度。本文提出JoLT(联合拉格朗日张量分解),直接采用部分Tucker分解,仅压缩令牌和特征轴,保留头和层轴不变,并通过旋转后的低比特残差恢复截断能量:先进行随机正交旋转,再进行低比特量化。单个拉格朗日对偶联合分配每层组中键和值的张量秩与残差位宽,在每字节预算下实现近无损压缩。结果达到2-3倍压缩率,在分组查询注意力模型(Mistral-7B-v0.3)和多头注意力模型(LLaMA-2-13B)上,困惑度近乎无损;在2倍压缩下GSM8K准确率与针堆检索任务表现持平,3倍压缩下仍保持稳定。2倍压缩时,重建相对弗罗贝尼乌斯误差为0.009(键)和0.006(值)。采用随机化SVD的FlashJoLT版本在1024令牌上下文下实现5-13倍压缩速度提升,质量相当。

原文摘要 · Abstract (English)

The key-value (KV) cache has become the dominant memory cost of transformer inference: it grows with batch size, context length, and depth, and at long context it, rather than the model weights, sets the throughput ceiling. Existing reductions fall into two families. Low-rank methods factor two-dimensional slices of the cache, either per-head matrices or cross-layer feature blocks, and quantization methods lower the bit-width of every entry. Neither exploits the fact that the cache at a layer is naturally a third-order tensor whose three axes, the heads, the tokens, and the features, carry very different amounts of redundancy. We take this tensor view directly. Our method, JoLT (Joint Lagrangian Tucker), applies a partial Tucker decomposition that compresses only the token and feature axes while leaving the head and layer axes intact, then restores the energy that truncation discards with a rotated low-bit residual: a random orthogonal rotation followed by low-bit quantization. A single Lagrangian dual allocates the Tucker ranks and the residual bit-widths together, per layer group and separately for keys and values, under one byte budget. The result is a near-lossless 2-3x compression. Perplexity stays near-lossless on both a grouped-query-attention model (Mistral-7B-v0.3) and a multi-head-attention model (LLaMA-2-13B), and GSM8K accuracy and needle-in-a-haystack retrieval hold at the uncompressed baseline at 2x on both architectures and through 3x on the GQA model. At 2x, JoLT reconstructs the cache to relative Frobenius error 0.009 (K) and 0.006 (V) on both architectures. A randomized-SVD variant, FlashJoLT, delivers a 5-13x compression-time speedup at 1024-token context and matched quality.

KV缓存张量分解模型压缩推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。