Fractional Decay KV-Cache：所有权感知内存管理以提升对话系统中的推理相关性

# 具备所有权感知内存管理的对话系统推理相关性改进方案
来源：https://arxiv.org/html/2608.18098

###### 摘要

键值缓存对于基于Transformer的对话系统中高效自回归推理至关重要，然而现有策略对所有缓存条目一视同仁，或采用粗粒度的驱逐启发式方法，无法适应对话主题的动态演变。本文提出分数衰减键值缓存（FD-KVC）这一新型算法，为每个键值对维护双通道评分机制：一个追踪累积重要性的*累积注意力通道*（类似于H2O），以及一个由时间衰减和基于强化学习的更新所控制的*近期加权相关性通道*。两者结合使FD-KVC既能保留历史上重要的标记，又能在对话主题转移时快速适应。由所有权损失函数驱动的自适应学习率确保收敛过程无振荡。FD-KVC完全在CPU上运行且开销可忽略。在五个包含600段对话的多样化多轮对话场景中，FD-KVC在复合后期轮次对齐度上超越了当前最优的H2O基准方法（+6.7%），具体提升包括：主题转移场景+127%，渐进演变场景+87%，混合主题对话场景+30%。FD-KVC适应新主题的速度比H2O快3.6倍，并在所有方法中实现最高的主题多样性（80.6%）。消融研究证实了各组件的贡献。

分数衰减键值缓存：具备所有权感知内存管理的对话系统推理相关性改进方案

Sukanta Ganguly NetApp Inc. [email protected]

## 1引言

基于Transformer的语言模型依赖键值缓存来避免自回归解码过程中的冗余计算（Pope et al., 2023）。在多轮对话中，缓存会累积所有先前轮次的表示，其大小随对话长度线性增长。这种增长导致两个问题：（1）*内存压力*，随着缓存消耗的RAM不断增加；（2）*相关性稀释*，因为注意力被分散到越来越多的缓存条目上，其中许多条目已失去上下文相关性。

现有方法通过滑动窗口（Xiao et al., 2024）、高频词保留（Zhang et al., 2023）或学习型压缩（Mu et al., 2023）来解决内存压力。然而，这些方法要么采用不建模时间动态的二元保留-丢弃决策，要么——如H2O等累积注意力方法——面临*缓存过时问题*：在对话早期累积高注意力分数的标记会无限期保留这些分数，导致缓存无法随主题演变而调整。

本文提出分数衰减键值缓存（FD-KVC），通过双通道评分框架解决上述问题。每个缓存的键值对维护（1）*累积注意力分数*追踪整体重要性，确保历史上有价值的标记得以保留；（2）*近期相关性分数*，该分数随时间衰减，并在条目被证明与当前查询相关时得到强化。这两个通道的混合组合驱动驱逐决策，而注意力权重通过近期性进行软调制，以聚焦当前相关上下文。

我们的主要贡献包括：

- 一种结合累积注意力（用于长期重要性）与衰减近期相关性（用于主题适应）的键值缓存条目**双通道评分模型**，实现平滑、连续的相关性跟踪。
- 一条具有动态奖励信号的**受强化学习启发的更新规则**，为上下文相关的缓存条目强化近期通道。
- 一个由收敛感知的所有权损失函数驱动的**自适应学习率**，确保快速收敛。
- 一种**CPU高效实现方案**，开销可忽略不计，无需GPU加速。
- 跨越五种对话场景的综合实验，证明在主题适应、多样性和复合对齐度方面相比H2O有显著提升。

## 2相关工作

##### 键值缓存优化。

标准键值缓存存储所有历史键值对且无限增长（Pope et al., 2023）。多查询注意力（Shazeer, 2019）和分组查询注意力（Ainslie et al., 2023）减少了每头内存消耗，但未解决时间相关性问题。PagedAttention（Kwon et al., 2023）提高了内存分配效率，但仍保留所有条目。FlashAttention（Dao et al., 2022）优化了注意力计算的计算-内存权衡，但不执行缓存驱逐。

##### 缓存驱逐策略。

StreamingLLM（Xiao et al., 2024）维护固定大小的滑动窗口加注意力汇聚点。H2O（Zhang et al., 2023）保留累积注意力分数最高的重要标记。Scissorhands（Liu et al., 2023）利用重要性持续性来压缩缓存。FastGen（Ge et al., 2024）根据注意力模式自适应选择丢弃哪些键值对。SnapKV（Li et al., 2024）在生成前识别重要的键值位置。这些方法采用二元驱逐策略，未建模分数相关性或时间衰减。

##### 长上下文方法。

循环记忆Transformer（Bulatov et al., 2023）通过循环记忆机制增强Transformer。Unlimiformer（Bertsch et al., 2023）通过检索将Transformer扩展到无限长度。Gemini 1.5（Reid et al., 2024）和检索头（Wu et al., 2024）解决长上下文事实性问题。这些方法修改模型架构；相比之下，FD-KVC作为标准缓存的即插即用替代方案，无需重新训练。

##### 学习型压缩。

Gist tokens（Mu et al., 2023）学习将提示压缩为紧凑表示。虽然有效，但这需要训练专门的压缩模块。FD-KVC通过轻量级、免训练的所有权机制实现压缩。

## 3分数衰减键值缓存

### 3.1问题形式化

考虑一个包含轮次{u₁, u₂, ..., uₜ}的多轮对话。在轮次t，模型处理输入标记xₜ=(xₜ¹, ..., xₜⁿᵗ)，并通过当前和缓存表示的注意力生成响应。令Cₜ={(kᵢ, vᵢ)}ᵢ₌₁^|Cₜ|表示轮次t时的键值缓存。

标准缓存仅进行累积：Cₜ=Cₜ₋₁∪{(kⱼ, vⱼ)}ⱼ∈new，并在|Cₜ|超过最大大小时进行截断。我们寻求一种策略，在选择性保留相关条目的同时，优雅地降低过时条目的影响。

### 3.2所有权分数

每个缓存条目(kᵢ, vᵢ)被分配两个互补的分数：*累积注意力分数*cᵢ和*近期相关性分数*ρᵢ。它们共同构成混合所有权分数，用于驱逐决策和注意力调制。

#### 3.2.1累积注意力通道

累积分数追踪条目生命周期内接收的累积注意力，遵循H2O范式（Zhang et al., 2023）：

cᵢ ← cᵢ + aᵢ⁽ᵗ⁾, aᵢ⁽ᵗ⁾ = softmax(q̄ₜ⊤kᵢ / √d)  (1)

其中q̄ₜ是轮次t时的平均查询向量。新条目初始化为cᵢ=1.0。该通道提供稳定性：频繁被关注的标记累积高分数，并在各轮次间保持。

#### 3.2.2近期相关性通道

近期分数捕捉*近期*相关性，并受时间衰减影响：

ρᵢ ← ρᵢ · γ, γ ∈ (0,1)  (2)

每轮次应用一次。与累积通道不同，近期通道会“遗忘”旧的相关性，使缓存能在主题变化时适应。衰减后，基于当前查询相似性强化近期分数：

ρᵢ ← ρᵢ + αₜ · rᵢ  (3)

其中rᵢ是缓存嵌入eᵢ与平均查询嵌入之间的归一化余弦相似度，缩放到[0,1]：

rᵢ = (cos(eᵢ, ēq) - minⱼ cos(eⱼ, ēq)) / (maxⱼ cos(eⱼ, ēq) - minⱼ cos(eⱼ, ēq) + ε)  (4)

#### 3.2.3混合分数

两个通道合并为单一的混合所有权分数：

hᵢ = wₜ · ĉᵢ + wρ · ρᵢ, wₜ + wρ = 1  (5)

其中ĉᵢ = cᵢ / maxⱼ cⱼ 将累积通道归一化到[0,1]，wₜ、wρ为权重超参数。当wₜ=1时，该方法退化为H2O；当wρ=1时，仅使用衰减相关性。默认值wₜ=0.45、wρ=0.55在长期重要性和近期相关性之间取得平衡。

### 3.3自适应学习率

学习率αₜ根据混合分数的收敛状态自适应调整。我们定义*所有权损失*：

ℒₜ = (1/|Cₜ|) ∑ᵢ₌₁^|Cₜ| ĥᵢ(1 - ĥᵢ)  (6)

其中ĥᵢ = hᵢ / maxⱼ hⱼ 是归一化的混合分数。当所有混合分数接近0或1（完全致力于驱逐或保留）时，该损失最小化。学习率调整方式为：

αₜ = α₀ / (1 + μ · ℒₜ)  (7)

其中α₀为初始学习率，μ为收敛适应因子。当所有权分数不确定（ℒₜ高）时，学习率降低以避免不稳定。当分数收敛（ℒₜ→0）时，学习率恢复至α₀以实现响应式适应。

收敛性分析：令f(h)=h(1-h)。公式(3)中的强化更新将高相关性条目的ρᵢ向上推动，通过近期通道增加hᵢ。时间衰减（公式2）降低不相关条目的ρᵢ。累积通道（公式1）确保持续被关注的标记保持较高基线。不动点为hᵢ→0（驱逐）和hᵢ→1（保留），收敛速率由γ、wₜ和αₜ控制。

### 3.4驱逐与注意力调制

混合分数低于阈值τ（相对于最大值）的条目被驱逐：

Cₜ ← {(kᵢ, vᵢ) | hᵢ ≥ τ · maxⱼ hⱼ}  (8)

当缓存在插入新条目后超过预算B时，移除混合分数最低的条目，直至|Cₜ| ≤ B。

注意力权重通过近期分数进行软调制，以聚焦近期相关上下文：

ŵᵢ = (wᵢ · (0.5 + 0.5ρ̂ᵢ)ᵝ) / ∑ⱼ wⱼ · (0.5 + 0.5ρ̂ⱼ)ᵝ  (9)

其中ρ̂ᵢ = ρᵢ / maxⱼ ρⱼ 是归一化的近期分数，β>0控制调制强度。(0.5 + 0.5ρ̂ᵢ)项确保即使低近期性条目也至少贡献其未调制权重的一半，防止因历史上重要但暂时未被引用的缓存条目而导致信息丢失。

### 3.5完整算法

算法1总结了FD-KVC流程。该算法每个对话轮次运行一次，相关性计算复杂度为O(|C|·d)，所有权更新和驱逐复杂度为O(|C|)，其中d为键维度。

算法1 FD-KVC：分数衰减键值缓存
0: 当前查询Qₜ，新键Kₜⁿᵉʷ，新值Vₜⁿᵉʷ
0: 缓存Cₜ₋₁及其分数{cᵢ, ρᵢ}
0: 超参数γ, α₀, μ, τ, β, wₜ, wρ
1:// 1. 累积注意力更新
2: 对每个条目i计算aᵢ⁽ᵗ⁾ ← softmax(q̄ₜ⊤kᵢ/√d)
3: cᵢ ← cᵢ + aᵢ⁽ᵗ⁾ // 公式(1)
4:// 2. 近期衰减
5: 对Cₜ₋₁中每个条目i执行
6: ρᵢ ← ρᵢ · γ // 公式(2)
7: 结束循环
8:// 3. 相关性强化
9: rᵢ ← normalize(cos(eᵢ, ēq)) // 公式(4)
10: ρᵢ ← ρᵢ + αₜ · rᵢ // 公式(3)
11:// 4. 自适应学习率
12: hᵢ ← wₜĉᵢ + wρρᵢ // 公式(5)
13: ℒₜ ← (1/|C|)∑ᵢ ĥᵢ(1 - ĥᵢ) // 公式(6)
14: αₜ ← α₀ / (1 + μ·ℒₜ) // 公式(7)
15:// 5. 软+硬驱逐
16: 移除满足hᵢ < τ·maxⱼhⱼ的条目 // 公式(8)
17: 插入新条目：cⱼ=1，ρⱼ∝cos(eⱼ, ēq)
18: 若|Cₜ|>B则
19: 保留