Codec-Gauge:为Transformer KV缓存学习压缩友好的变换框架

arXiv cs.LG 论文

摘要

Codec-Gauge 为Transformer KV缓存学习小型正交通道变换(度量),以在固定比特率下提高压缩保真度,显著降低多个模型和后端上的KL散度。

arXiv:2607.20538v1 公告类型:新 摘要:长上下文Transformer推理越来越依赖KV缓存压缩或量化。先前的旋转和变换编码结果表明,每个键/值向量的通道基础会影响固定后端如何保真地保留模型行为。我们引入了Codec-Gauge,这是一个训练后的缓存坐标层,它围绕现有的压缩和量化后端学习小型正交通道变换。其频率分布目标结合了令牌通道DCT谱质心损失和平滑速率代理,以将KV能量集中到面向低频率编解码器的布局中。我们使用实际测量的字节数和滚动压缩历史评分来评估实际压缩和解压缩。在六个模型的每个值3、4和6比特下,学习到的度量将zfp KL散度平均降低了44.0%,相对于原始坐标,并优于随机、Hadamard、DCT和PCA/KLT控制。同样的度量改善了块均匀和KIVI风格量化的质量保持。在27B模型和长上下文任务提示上的实验重复了质量趋势,而串行存储和时序测量验证了实现的压缩缓存路径。这些结果确立了缓存坐标几何作为实用的训练后变量,用于提高压缩保真度,而无需改变模型权重、注意力语义或后端编码规则。
查看原文
查看缓存全文

缓存时间: 2026/07/24 05:10

# Codec-Gauge: 学习面向Transformer KV缓存的压缩友好型变换
Source: https://arxiv.org/html/2607.20538
###### 摘要

长上下文Transformer推理越来越依赖KV缓存压缩或量化。旋转和变换编码结果表明,每个键/值向量的通道基会影响固定后端忠实保持模型行为的程度。我们提出*Codec-Gauge*,一个后训练的缓存坐标层,学习现有后端周围的小型正交通道变换(即度量)。其频率分布目标结合了令牌-通道DCT谱质心损失和平滑速率代理,将KV能量集中在低频编码器面向的布局中。我们使用实测字节和滚动压缩历史评分评估实际的压缩/解压缩。在六个模型上,以3、4和6比特/值的设置,学习得到的度量平均降低zfp KL散度44.0%,优于原始坐标以及随机、Hadamard、DCT和PCA/KLT对照;同样的度量也改善了块均匀和KIVI风格量化的质量保持。27B扩展和任务提示似然验证了这一质量趋势,而串行存储/时序测量验证了实现的压缩缓存路径。

## 引言

长上下文Transformer推理越来越受到KV缓存内存和带宽成本的制约。在自回归解码过程中,每一层都会存储过去的键和值张量,该缓存随批量大小、上下文长度、层数、KV头数和头维度线性增长。注意力内核和内存管理器改善了访问局部性,而多查询和分组查询注意力减少了存储的KV头数((Shazeer2019 (https://arxiv.org/html/2607.20538#bib.bib1); Ainslie et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib2); Dao et al. 2022 (https://arxiv.org/html/2607.20538#bib.bib3); Kwon et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib4)));尽管如此,缓存仍然是一个需要存储、移动和反复消耗的大型运行时对象。本文研究一个互补变量:呈现给压缩或量化后端的坐标基。Codec-Gauge 从冻结模型的KV张量中学习该基,以在相同测量速率下提高保真度,同时保持模型权重、注意力语义和后端编码规则不变(图1 (https://arxiv.org/html/2607.20538#Sx1.F1))。

参照图注图 1:Codec-Gauge 概述。一次性的后训练阶段从冻结模型的KV张量中学习正交度量。推理时,度量在压缩或量化之前将KV张量映射为编码器面向的坐标,逆度量在注意力之前恢复被恢复的KV状态。大量工作通过改变缓存精度、缓存保留或模型表示来降低这一成本((Liu et al. 2024c (https://arxiv.org/html/2607.20538#bib.bib5); Hooper et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib6); Kang et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib7); Yang et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib8); Lin et al. 2025b (https://arxiv.org/html/2607.20538#bib.bib9); Zhang et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib17); Xiao et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib18); Li et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib20); DeepSeek-AI 2024 (https://arxiv.org/html/2607.20538#bib.bib29); Chang et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib30); Nawrot et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib33); Gelber et al. 2026 (https://arxiv.org/html/2607.20538#bib.bib34)))。Codec-Gauge 源于两个相关观察:变换编码表明,基变换可以将信号能量集中到更容易以固定速率保持的系数中((Ahmed et al. 1974 (https://arxiv.org/html/2607.20538#bib.bib39); Lindstrom 2014 (https://arxiv.org/html/2607.20538#bib.bib38))),而LLM量化表明,正交通道旋转可以在不改变所表示函数的情况下改变低比特误差行为((Chee et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib10); Tseng et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib11); Ashkboos et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib12); Liu et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib13); Su et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib15); Saxena and Roy 2025 (https://arxiv.org/html/2607.20538#bib.bib16)))。这些观察共同促使我们在预训练后学习呈现给固定后端的表示几何结构。尽管模型通道没有固有的空间顺序,但编码器消耗具体的存储布局,因此学习到的通道基可以改变存储KV场的平滑度和系数结构。

本文将几何结构作为一等压缩变量进行研究。*Codec-Gauge* 学习KV通道组上的正交度量,使得现有压缩和量化后端看到更有利于压缩的表示,同时模型权重、注意力语义和后端实现保持不变。所谓KV坐标,是指每个头内用来表示每个键或值向量的通道基。度量是该基的可逆变化:它可以在每个组内混合通道坐标,但从不混合令牌、头或层,也从不删除缓存条目。推理时,缓存的键和值在压缩或量化之前被映射到学习到的度量中,并在注意力之前映射回来。在无误差路径上,度量和逆度量在数值精度内保持缓存;在有损后端下,度量改变了后端所见的几何结构,从而改变了恢复误差和解压缩后保留的质量。

我们采用匹配速率的实验设计。在每个比较中,模型、令牌流、后端、测量比特预算和评估器都是固定的,只有缓存坐标映射发生变化。这使得恒等映射、随机正交、固定变换、数据驱动和学习到的度量在相同后端下具有可比性,因此改进可以归因于缓存几何结构,而不是改变的量化器、驱逐规则、打包格式、内核或调度器。

我们的主要目标是对连续KV张量进行固定速率的数值压缩。我们使用GPU zfp压缩作为主要编解码器,因为其块变换结构使得系数集中和局部平滑度与固定速率下的重建质量直接相关((Lindstrom 2014 (https://arxiv.org/html/2607.20538#bib.bib38)))。关键技术选择是在编码器面向的令牌-通道布局上使用频率分布损失:DCT谱质心项将KV能量移向较低频率,而平滑的对数幅度代理则倾向于具有较少显著变换系数的系数分布。所有质量声明都使用实际的压缩/解压缩和滚动压缩历史评分,其中新创建的KV条目经历相同的度量和后端路径。

编解码器仍在训练循环之外:拟合度量不使用重建、语言建模、logit或任务输出损失。

我们使用坐标对照来识别该训练目标的效果。恒等映射使用原始缓存基,而随机正交度量则捕获了量化中已知的通用旋转和异常值重分布效应((Chee et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib10); Tseng et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib11); Ashkboos et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib12); Liu et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib13); Sun et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib14); Su et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib15); Saxena and Roy 2025 (https://arxiv.org/html/2607.20538#bib.bib16)))。Hadamard和DCT提供了固定的结构化变换,PCA/KLT提供了数据驱动的基,而学习到的度量则在相同后端和测量速率下测试所提出的频率分布目标。这些对照将编码器面向的度量学习与原始坐标压缩和通用正交预处理分离开来。

#### 后端兼容性。

学习到的度量是每个检查点的固定变换,而不是每个令牌的有效载荷,并且既不改变后端的编码规则,也不改变其比特预算。它应用于KV张量接口;我们的测量使用显式恢复,并将该成本计入时序,将质量、存储和延迟与同一实现的路径关联起来。

在六个语言模型和多个匹配的每值比特设置下,学习到的度量在zfp上降低了KV重建误差和滚动后缀退化,优于恒等映射、随机、固定变换和PCA/KLT对照。在相同的训练坐标下,块均匀和KIVI风格量化器((Liu et al. 2024c (https://arxiv.org/html/2607.20538#bib.bib5)))也比其原始坐标版本更好地保持了输出质量。更大的27B模型检查和任务提示似然验证了坐标效应,而串行上下文长度测量验证了实际存储和恢复成本。对于zfp,学习到的度量在主要操作点上将KL散度、logit MSE、top-1翻转率和KV NRMSE分别降低了44.0%、43.3%、24.5%和18.3%。

本文做出四项贡献。首先,我们将KV缓存坐标几何结构形式化为一种后训练变量,用于在固定后端下提高压缩保真度。其次,我们提出了Codec-Gauge,一个正交的缓存坐标层,它包裹现有压缩和量化后端而不改变注意力语义。第三,我们设计了一个频率分布目标,该目标使用编码器面向的令牌-通道谱结构从冻结的KV张量中训练该层,无需语言建模损失、logit匹配、任务监督或模型权重更新。第四,我们提供配对证据,表明学习到的度量在匹配的测量速率下改善了实际的GPU zfp压缩,同时也改善了块均匀和KIVI风格量化路径的质量保留。

## 相关工作

#### KV缓存量化与旋转。

KV缓存量化通过非对称粒度、异常值处理、校正项、混合精度或服务端协同设计来降低存储键/值的精度((Liu et al. 2024c (https://arxiv.org/html/2607.20538#bib.bib5); Hooper et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib6); Kang et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib7); Yang et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib8); Lin et al. 2025b (https://arxiv.org/html/2607.20538#bib.bib9)))。坐标变换也是LLM量化的核心,包括非相干性处理、去除异常值的旋转、学习旋转、分布平坦化、自适应KV旋转以及基于Hadamard的校正((Chee et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib10); Tseng et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib11); Ashkboos et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib12); Liu et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib13); Sun et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib14); Su et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib15); Saxena and Roy 2025 (https://arxiv.org/html/2607.20538#bib.bib16)))。Codec-Gauge 使用面向固定数值编解码器的缓存几何频率目标,而非标量量化损失,然后将相同的训练坐标复用于标量低比特路径。

#### 令牌、页面和动态缓存选择。

另一类工作通过使用重击者、注意力沉点、模型内部模式、提示观察、逐层预算或查询感知稀疏性来减少保留或访问的缓存条目数量((Zhang et al. 2023 (https://arxiv.org/html/2607.20538#bib.bib17); Xiao et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib18); Ge et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib19); Li et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib20); Cai et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib21); Tang et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib22)))。InfiniGen、RocketKV和CacheGen强调动态缓存管理、多阶段压缩或缓存流式传输((Lee et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib23); Behnam et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib24); Liu et al. 2024b (https://arxiv.org/html/2607.20538#bib.bib28)))。Codec-Gauge 则作用于保留条目的密集张量坐标,因此可以与保留和分页策略叠加使用。

#### 变换与数值KV压缩。

频域和变换编码方法将KV缓存视为结构化的数值信号。FreqKV、FAEDKV和KVTC利用了频率或变换域的冗余((Kai et al. 2026 (https://arxiv.org/html/2607.20538#bib.bib25); Li et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib26); Staniszewski and Łańcucki 2026 (https://arxiv.org/html/2607.20538#bib.bib27))),而DCT和zfp展示了谱集中度和局部平滑度如何影响数组率失真行为((Ahmed et al. 1974 (https://arxiv.org/html/2607.20538#bib.bib39); Lindstrom 2014 (https://arxiv.org/html/2607.20538#bib.bib38)))。Codec-Gauge 为此设置贡献了一个学习到的坐标层:度量塑造了现有编解码器的令牌-通道布局所看到的数值场,而压缩由后端本身执行。

#### 架构、低秩和训练感知的KV表示。

若干方法通过改变模型结构或学习不同的缓存表示来减少KV状态,包括潜KV状态、低秩投影、深度冗余、动态内存压缩、针对可压缩KV的训练以及自适应正交投影((DeepSeek-AI 2024 (https://arxiv.org/html/2607.20538#bib.bib29); Chang et al. 2025 (https://arxiv.org/html/2607.20538#bib.bib30); Liu et al. 2024a (https://arxiv.org/html/2607.20538#bib.bib32); Nawrot et al. 2024 (https://arxiv.org/html/2607.20538#bib.bib33); Gelber et al. 2026 (https://arxiv.org/html/2607.20538#bib.bib34); Lin et al. 2025a (https://arxiv.org/html/2607.20538#bib.bib31)))。Codec-Gauge 保持检查点冻结:为现有模型训练一个小型可逆度量,并围绕显式压缩/解压缩路径进行评估。

## 问题陈述与概述

表 1:缓存度量训练符号。表1 (https://arxiv.org/html/2607.20538#Sx3.T1) 总结了符号。我们考虑自回归推理期间由softmax注意力Transformer产生的标准键/值缓存。对于层$l$,前$T$个令牌的缓存为:

$K_l, V_l \in \mathbb{R}^{B \times H_l \times T \times d_l}, \quad \mathcal{M}_T = \{(K_l, V_l)\}_{l=1}^L.$

FP16或BF16存储中的缓存元素数量为:

$N_{\mathrm{KV}}(T) = \sum_{l=1}^L 2 B H_l T d_l,$

它随上下文长度线性增长。我们只考虑标准softmax注意力暴露的键/值张量;循环或线性注意力状态不在本文研究的度量和压缩对象范围内。

Codec-Gauge 学习一个缓存坐标度量,而不改变令牌、注意力架构或模型权重。令 $X_{l,s}$ 表示层 $l$ 的缓存张量,其中 $s \in \{K, V\}$。对于每个层、缓存类型、头和通道组,我们学习一个可逆度量 $G_{l,s,h,r}$。组维度 $g$ 将头维度划分为 $R_l = d_l / g$ 个组;当 $g = d_l$ 时,度量作用于整个头。对于组 $r = 1, \ldots, R_l$,

$X_{l,s}[b, h, t, r] \in \mathbb{R}^g,$

相似文章

KV缓存压缩的风险

arXiv cs.LG

本文从理论上刻画了变压器中KV缓存压缩的极小极大风险,为因果掩码下的精确压缩提供了设计原则,并将其实例化到实用算法中,在LongBench上取得了有前景的结果。

KV缓存压缩比TurboQuant与逐向量香农极限高出900000倍

Hacker News Top

一篇新论文提出了一种基于概率语言Trie树和预测差分编码的顺序KV缓存压缩方法。该方法通过利用语言模型Token的序列结构而非对向量进行独立处理,实现了超越TurboQuant约91.4万倍的理论压缩比。

KV缓存压缩的消融、统计推断与验证

arXiv cs.LG

本文对KV缓存压缩方案(TurboQuant和SpectralQuant)进行了系统的比较研究,介绍了一种统计验证方法,并针对高效Transformer推理提供了特定场景下的建议。