WaveFilter: 通过小波引导的KV缓存过滤增强扩散LLMs的长上下文能力

# WaveFilter: 通过小波引导的KV缓存过滤增强扩散语言模型的长上下文能力

来源：https://arxiv.org/html/2606.00724
杨金楠1,4，王岩2，毕震3，吴克浩1，李晓洁1，楼俊刚3，李泽超1†，刘静4† 1南京理工大学，2阿里巴巴集团，3湖州师范大学，4中国科学院自动化研究所 †通讯作者

###### 摘要

扩散大语言模型（DLMs）在各种任务中展现出显著优势。然而，受限于其多步迭代推理机制，它们在长上下文任务中的计算开销和推理延迟已成为制约其大规模部署的核心瓶颈。在处理长序列时，现有的键值（KV）缓存机制经常面临生成质量急剧下降的困境，其核心挑战在于如何在超长上下文中精确且高效地筛选关键token。受人类阅读过程的启发，我们提出了**WaveFilter**，一个通用且无需训练的缓存框架。该框架创新性地引入小波变换对长序列进行分解，以实现关键token的精确识别，并据此构建稀疏KV缓存以计算最终的上下文表示。实验结果表明，WaveFilter作为一个即插即用的通用框架，显著提升了现有主流KV缓存方法在复杂长上下文任务中的性能。

WaveFilter: 通过小波引导的KV缓存过滤增强扩散语言模型的长上下文能力

杨金楠1,4，王岩2，毕震3，吴克浩1，李晓洁1，楼俊刚3，李泽超1†，刘静4† 1南京理工大学，2阿里巴巴集团，3湖州师范大学，4中国科学院自动化研究所 †通讯作者

## 1 引言

由于其非自回归特性和双向上下文建模能力，扩散大语言模型（DLMs）Nie等人（2025 (https://arxiv.org/html/2606.00724#bib.bib1)）在文本到图像生成、对话系统和代码生成等任务中展现出独特优势Sahoo等人（2024b (https://arxiv.org/html/2606.00724#bib.bib2)）；Gupta等人（2024 (https://arxiv.org/html/2606.00724#bib.bib3)）；Gong等人（2025 (https://arxiv.org/html/2606.00724#bib.bib4)）。然而，受限于其多步迭代推理机制，DLMs的计算复杂度和推理延迟显著高于自回归模型Li等人（2022 (https://arxiv.org/html/2606.00724#bib.bib5)，2025 (https://arxiv.org/html/2606.00724#bib.bib6)）。这种沉重的计算负担已成为限制其大规模部署的核心瓶颈。为了减轻重复计算带来的开销，研究者将键值（KV）缓存机制引入DLMsMa等人（2025 (https://arxiv.org/html/2606.00724#bib.bib7)）。通过缓存先前步骤的Key和Value向量，该方法使得后续生成能够直接复用先前的计算结果。这有效避免了已生成上下文上的冗余计算，从而降低了推理延迟并提高了生成效率。

参见说明图1：Ruler数据集的`niah_single_1`子集上不同上下文长度的性能对比。(a) 展示了LLaDA-8b-Instruct及其使用不同KV缓存方法变体的准确率(%)；(b) 展示了相应的吞吐量（Tokens/sec）。

然而，将KV缓存机制直接扩展到DLMs以处理复杂的长期上下文任务仍然面临重大挑战。一方面，如图1 (https://arxiv.org/html/2606.00724#S1.F1)a 所示，LLaDA-8B-Instruct的性能随着输入长度的增加而急剧下降，表明模型在长上下文任务中难以保持生成的鲁棒性Liu等人（2026 (https://arxiv.org/html/2606.00724#bib.bib8)）。另一方面，现有的KV缓存机制在DLMs应用于复杂长上下文任务方面的研究尚不充分。如图1 (https://arxiv.org/html/2606.00724#S1.F1)a 和图1 (https://arxiv.org/html/2606.00724#S1.F1)b所示，尽管Fast-dLLMWu等人（2025b (https://arxiv.org/html/2606.00724#bib.bib18)）和Elastic-CacheNguyen-Tri等人（2025 (https://arxiv.org/html/2606.00724#bib.bib19)）能够为短文本任务提供一定的加速优势，但其吞吐量随着上下文长度的增长而迅速恶化，并且通常伴随着准确率的进一步下降。**核心挑战在于在超长上下文序列中，精确识别和过滤对去噪过程起关键作用的token极其困难**。因此，开发一个通用、即插即用的增强框架，以赋能现有缓存方法（如Fast-dLLM和Elastic-Cache），使其能够高效且鲁棒地扩展到长上下文任务，仍然是一个亟待解决的问题。

为了在长上下文任务中准确筛选关键token，我们借鉴了人类“先浏览后扫描”的认知习惯。人类通常先快速浏览全文以构建宏观层面的上下文语义结构，随后针对特定问题进行局部扫描以定位和提取关键信息，最终实现高效准确的问答。受此认知过程启发，我们提出了**WaveFilter**，一个通用且无需训练的框架。该框架的核心在于引入**离散小波变换（DWT）** 进行KV缓存压缩。通过压缩缓存长度并滤除高频噪声，同时完整保留时域信息，WaveFilter成功促进了宏观上下文语义结构的快速构建。随后，采用多尺度递归过滤机制模拟局部扫描，精确定位与问题最相关的token，以实现高精度的问答。

具体而言，在初始时间步，首先利用DWT提取缓存的语义特征，并利用注意力机制识别查询向量所针对的初始关键token。在此基础上，对初始重要token进行多尺度递归过滤，以确定最终的关键token。这些最终token被直接用于动态构建稀疏KV缓存，随后与当前查询向量一起参与注意力计算，以精确形成最终的上下文表示。总之，本文的主要贡献如下：

- • 本文提出了WaveFilter，一个用于长上下文任务的通用且无需训练的KV缓存框架。通过模仿人类的认知阅读习惯，该框架无缝地赋能现有的KV缓存方法，有效解决了它们在长上下文任务中的性能退化问题。
- • WaveFilter创新性地引入小波变换实现多尺度token过滤：在保留关键信息的同时大幅压缩缓存，以极小的开销精确识别相关区域，从而有效解决了海量缓存中关键token识别的难题。
- • 实验结果表明，作为一个通用的即插即用框架，WaveFilter可以无缝集成到各种现有的KV缓存策略中：在保持有竞争力的生成速度的同时，提升了模型在复杂长上下文任务中的性能，始终优于单独的KV缓存方法。

## 2 预备知识

### 2.1 掩码扩散模型中的键值缓存

掩码扩散模型（MDMs）用随机掩码和迭代填空代替传统的连续噪声添加，实现了离散数据的并行生成Sohl-Dickstein等人（2015 (https://arxiv.org/html/2606.00724#bib.bib9)）；Austin等人（2021 (https://arxiv.org/html/2606.00724#bib.bib10)）；Campbell等人（2022 (https://arxiv.org/html/2606.00724#bib.bib11)）；Sahoo等人（2024a (https://arxiv.org/html/2606.00724#bib.bib38)）。为了优化逆向过程中的生成效率，KV缓存机制被集成到其Transformer主干中。如图2 (https://arxiv.org/html/2606.00724#S2.F2)a所示，在初始时间步`t`（其中`t=1`），模型对所有位置`I={1,2,...,N}`执行完整计算。在第`l`层，当前隐藏状态`h^{1,l}`通过可学习的投影矩阵`W_Q^{1,l}`、`W_K^{1,l}`和`W_V^{1,l}`被投影为查询向量`Q^{1,l}[I]`、键向量`K^{1,l}[I]`和值向量`V^{1,l}[I]`。该层的注意力输出及相应的KV缓存初始化公式为：

`A^{1,l}[I] = Softmax( (Q^{1,l}[I] (K^{1,l}[I])^T) / sqrt(d) ) V^{1,l}[I]`。 (1)
随后，初始步计算得到的KV对被保存到缓存中，初始化定义如下：

`{ \widetilde{K}^{1,l}[I] = K^{1,l}[I]; \widetilde{V}^{1,l}[I] = V^{1,l}[I]`。 (2)
在后续时间步`t > 1`，模型只对生成位置集合`\widetilde{I}`进行推理。通过复用先前时间步存储的缓存键`\widetilde{K}`和值`\widetilde{V}`，注意力计算简化为：

`A^{t,l}[\widetilde{I}] = Softmax( (Q^{t,l}[\widetilde{I}] (\widetilde{K}^{t-1,l}[I])^T) / sqrt(d) ) \widetilde{V}^{t-1,l}[I]`。 (3)
随后，使用当前步计算得到的KV对动态更新缓存：

`{ \widetilde{K}^{t,l}[\widetilde{I}] = K^{t,l}[\widetilde{I}]; \widetilde{V}^{t,l}[\widetilde{I}] = V^{t,l}[\widetilde{I}]`。 (4)
基于KV缓存的推理显著降低了逆向过程中的计算复杂度和推理延迟。通过该机制，模型在保持Transformer全局上下文建模能力的同时，大幅提高了离散序列生成的效率。

### 2.2 离散小波变换

离散小波变换（DWT）是一种用于信号分解的时频分析方法Yao等人（2022 (https://arxiv.org/html/2606.00724#bib.bib13)）；Kiruluta等人（2025 (https://arxiv.org/html/2606.00724#bib.bib14)）。DWT通过一对互补的滤波器组将信号`x[n]`分解为低频近似分量和高频细节分量。单级分解过程可公式化为：

`{ A_1[n] = sum_k x[k] * g[2n-k]; D_1[n] = sum_k x[k] * h[2n-k]`。 (5)
其中`g[n]`和`h[n]`分别表示低通和高通滤波器系数，下标`2n`表示下采样操作。DWT的核心优势在于其递归性质：在第一级分解之后，近似系数`A_1`可以作为下一级滤波器组的输入。此迭代过程构建了一个多层金字塔结构。经过`L`级分解后，原始信号由集合`{A_L, D_L, D_{L-1}, ..., D_1}`表示。在长序列建模的背景下，`A_L`捕获信号的全局语义信息，而各个细节分量`D_j`则保留不同分辨率下的局部波动。

参见说明图2：WaveFilter的流程示意图。包含四个部分：(a) **解码步骤与动机**：展示了离散扩散模型的解码机制，并引入了由粗到精的检索策略以解决长上下文任务中关键token提取的挑战。(b) **离散小波变换**：通过DWT对缓存键进行分解以提取低频分量。(c) **重要性评估与选择**：计算查询向量与缓存键低频分量之间的相关性，并利用Top-K选择来识别关键候选区域。(d) **多尺度递归过滤**：在不同尺度上递归细化候选区域，最终选择信息量最大的稀疏KV矩阵。

## 3 方法论

本节详细阐述WaveFilter框架的核心机制。为了解决长上下文任务中关键token的识别与过滤难题，本文提出了一种基于DWT的KV缓存框架。如图2 (https://arxiv.org/html/2606.00724#S2.F2)所示，整体算法流程包含两个不同的阶段（详细算法见附录A (https://arxiv.org/html/2606.00724#A1)）。第一阶段，**粗粒度全局感知（第3.1节）**，利用DWT构建低频分量，从而能够在压缩空间内快速定位包含关键信息的语义区域。第二阶段，**细粒度局部定位（第3.2节）**，通过多尺度递归过滤实现对重要token的精确提取，并据此构建稀疏KV缓存。这种由粗到精的策略不仅有效维持了速度竞争力，而且显著提升了模型的整体性能。

### 3.1 粗粒度全局感知：构建全局语义轮廓

对于长序列处理，标准注意力机制在从广泛上下文中提取与查询高度相关的关键token时面临严峻的计算挑战，这主要是由于其`O(N^2)`的时间复杂度。为了缓解这一瓶颈，本节提出了一种基于DWT的“全局语义图”构建新方法，旨在利用语义分布快速准确地定位可能包含关键token的潜在区域。

设`Q^{t,l}[\widetilde{I}] ∈ R^{\widetilde{I}×d}`表示第`l`层、时间步`t (t > 1)`的查询向量。缓存的键向量包括提示位置`I_p`和生成位置`I_g`。具体地，`\widetilde{K}^{t-1,l}[I_p] ∈ R^{I_p×d}`和`\widetilde{K}^{t-1,l}[I_g] ∈ R^{I_g×d}`分别表示第`l`层、时间步`t-1`时，针对提示和生成位置的缓存键向量。为了捕获不同感受野下的序列特征，如图2 (https://arxiv.org/html/2606.00724#S2.F2)b所示，我们首先通过采用小波基`ψ`的DWT提取`\widetilde{K}^{t-1,l}[I_p] ∈ R^{I_p×d}`的低频分量：

`\widetilde{K}_{low}^{t-1,l(B)}[I_p] = DWT(\widetilde{K}^{t-1,l}[I_p], ψ)`，