当前位置：首页 > ds >正文

使用频域变换轻松压缩kv-cache

ds 2025/7/1 16:41:14

论文标题

FreqKV: Frequency Domain Key-Value Compression for Efficient Context Window Extension

论文地址

https://arxiv.org/pdf/2505.00570

作者背景

上海交通大学，华为诺亚方舟实验室

动机

尽管当前市面上的各大LLM都支持较长的上下文窗口，但直接扩充输入上下文的长度会导致显存开销急剧上升，如何准确、高效地处理这类数据依然是大模型的重点研究方向；

本文旨在通过压缩kv-cache，提高大模型上下文处理能力。思想上类似于之前介绍过的kv驱逐算法

大模型推理加速：通过蒸馏kv-cache降低显存开销

本文方法

一、时频域转换信息压缩

本文从时频域转换视角出发，发现kv缓存在频域空间中，能量明显集中于低频区域：

在这里插入图片描述

以下是时频域转换方法DCT/IDCT的简单介绍：

在这里插入图片描述

DCT的公式如下，其中x是时间域输入（原始kv-cache序列），y是频域输出，n表示原始cache序列中的索引，t表示变换后的频率索引，α是归一化系数。DCT的结果就是原序列和特定频率的余弦波“相乘再相加”

在这里插入图片描述

按照上述公式，当t=0时，y就是x的均方根，表示序列最基本的“趋势”；当t=1，2，…，n时，y代表越来越高的频率，即原序列中更细致、快速的变化

IDCT的公式如下，即DCT的逆操作，可将频域信息还原成时域数据

于是我们可以首先使用频率域转换方法（如DCT），将kv缓存转换到频域，然后去掉其中的高频部分，再转换回时间域，实现间接而高效的信息压缩

好比一幅风景画，80%的美感和信息都集中在大致的轮廓（低频部分），而画中细节（如每片叶子的细节，高频）可能没那么关键。
如果你想压缩保存这幅画，就可保留轮廓而忽略细节，这样既不会太影响整体画面，又能节省大量空间

二、窗口触发机制

作者采用了窗口感知的迭代压缩策略：只有当KV缓存长度达到预设的上下文窗口大小N时，才触发一次压缩。特别地，会保留序列开头的若干token不参与压缩（很多工作里都称之为“sink tokens”），因为这些token一般具有较高的权重（注意力下沉现象），可能包含关键信息（如系统提示或主题）。具体流程如下：

当缓存累计了N个token的KV后，对除去前S个“下沉”token之外的剩余N-S个KV执行频域压缩，保留Compress(N-S)条低频信息作为压缩存回缓存。
此后，新生成的token继续附加到缓存中
当缓存再次累积到N条时，就再次触发压缩——将已有的压缩token和新进来的未压缩token一起执行频域压缩，产生新的压缩KV条目
不断迭代上述过程

在这里插入图片描述

由于每次压缩都会减少旧信息的表示长度，越早的上下文将经历越多轮压缩，越近期的内容则压缩程度较低，而模型的有效KV缓存始终维持在固定的规模上限，不会无限增长；
需要注意的是，为了使模型适应使用压缩后的信息，可以在长序列微调时，每隔固定长度便插入一个压缩操作，模拟上述过程

三、位置校准与编码策略

上述操作实际上扩展了输入上下文的长度，为了避免长度扩展对原始位置编码范围的影响，作者是在应用位置编码之前进行压缩。也就是说，每次将压缩后的KV条目重新放回缓存参与后续注意力计算时，会重新施加位置编码，因而不需要使用超出模型原始上下文长度的位置信息，不需要使用位置外推

实验结果

一、困惑度

作者将LLaMA-2-7B的上下文窗口从4K分别扩展到8K、16K、32K，使用RedPajama等语料进行了极少步数的微调，然后在PG-19长文数据集和Proof-pile数学数据集上测试不同长度下的困惑度。结果表明FreqKV在长上下文下的PPL几乎没有恶化，与完全扩展位置编码并全量微调的上限水平相当，明显优于其他缓存压缩方法

在这里插入图片描述