基于熵与低秩重构的高保真KV缓存摘要

# 基于低秩注意力重建的熵引导 KV 缓存摘要 来源：https://jchandra.com/posts/hae-ols/ 随着大语言模型（LLM）迈向百万级 token 上下文窗口，我们遇到了物理瓶颈：KV 缓存。为序列中的每个 token 存储 Key 和 Value 会导致显存线性增长。 在标准 Transformer 架构中，每生成一个新 token 都要访问前面所有 token 的 Key 与 Value。于是 KV 缓存随序列长度线性膨胀：$\[\mathcal{O}(N)\]$。对长上下文模型而言，这很快就会撑爆单卡显存（如 H100），只能要么分布式分片，要么激进剪枝。 现有的 Heavy-Hitter 或 Top-K 淘汰策略基于一个朴素假设：如果某个 token 现在不被关注，以后也不会被关注。然而自然语言信息本质上是上下文相关且非平稳的，一个段落里无足轻重的 token 可能在另一段成为关键锚点。 本文介绍另一种范式：**SRC（Selection-Reconstruction-Compression）流水线**。与其删除 token，不如用信息论与线性代数对其进行数学级“总结”。 为理解该假设为何失效，我们先审视注意力本身的结构。 ### 为何剪枝会失败：注意力是有结构的 注意力机制并非对 token 独立运算，而是产生 Query-Key 间的稠密交互模式，每个 token 都参与全局计算。 b4_heatmap 每行对应一个 Query token，每列对应一个 Key token，亮度表示 Query 对 Key 的关注强度。 可观察到： - 注意力高度结构化，并非稀疏 - 多个 token 联合贡献输出 - 依赖关系散布在整个序列 #### 逐 token 剪枝误差 粗看剪枝平均效果不错，但细粒度分析揭示关键失效模式。 Top-K 该图显示经 Top-K 剪枝后每个 token 的重建误差。 多数 token 误差低，似乎局部有效；但少数 token 出现**尖锐误差峰值**，即灾难性失败。 这些峰值对应： - 当下并非最显眼 - 却对下游注意力结构至关重要 且这些失败无法仅凭局部重要性预测。 > 剪枝并非均匀失效，而是**选择性且不可预测地失败**。 --- ## SRC 范式：三阶段演进 为克服上述局限，我们从启发式“决策”转向结构化“变换”流水线。 不再决定删谁，而是用一系列操作近似其贡献： `` 选择 → 重建 → 压缩 `` ### 选择：熵“回收站” 如何决定该总结哪些 token？不看幅值，看**信息不确定性**。 为每个头计算注意力权重的香农熵 H(P)。若某 token 熵高，说明模型在“分散”注意力，信息不具体，可能更冗余。 $$H(P) = -\sum_{i} p_i \log(p_i)$$ 高熵且累计重要性低的 token 移入“回收站”；低熵“锚点” token 留在高保真 Active Cache。 #### Token 熵景观 Entropy per Token 每点表示某 token 的注意力分布熵。 可见： - 部分 token**熵低**（聚焦） - 部分 token**熵高**（分散） 低熵 token 常充当**锚点**，集中注意力并承载明确语义；高熵 token 注意力分散，信息较模糊。 ### 重建：求解语义本质 回收站里的 token 要凝成一条**质心 token**。简单平均会忽略可能与之交互的特定 Query，于是我们把问题构造成**普通最小二乘（OLS）**：找权重矩阵 $W$ 使被删 token 产生的注意力输出重建误差最小。 给定参考 Query $Q_{\text{ref}}$ 与原回收 Value $V_{\text{bin}}$，求解： $$\[ W^* = \arg\min_{W} \left\| Q_{\text{ref}} W - \text{Attn}(Q_{\text{ref}}, K_{\text{bin}}, V_{\text{bin}}) \right\|_2^2 \]$$ 实践中用**Moore-Penrose 伪逆**解析求解： $$\[ W = Q_{\text{ref}}^{\dagger} \cdot \text{Attn}(Q_{\text{ref}}, K_{\text{bin}}, V_{\text{bin}}) \]$$ 如此压缩表示能精准近似原注意力输出，却只用极少 token。 #### 实现核心 OLS 逻辑 `` def summarize_bin_ols(Q_ref, bin_v): # 求解重建原 V 输出的权重 pinv_Q = torch.linalg.pinv(Q_ref) W_reconstruction = pinv_Q @ bin_v return W_reconstruction `` ### 压缩：低秩近似（SVD） 重建权重矩阵 $W$ 仍占显存。为真正省 VRAM，用**奇异值分解（SVD）**压缩 $W$。 保留前 $k$ 大奇异值及对应向量： $$ \[ W \approx U_k \Sigma_k V_k^T \]$$ 低秩分解把每个保留成分视为一条合成 key-value 对，即**质心 token**，代理整个回收站。 该过程滤掉低能分量（噪声），保留主方向（信号），在几乎不损失注意力功能的前提下获得大幅压缩。 ## 评估协议：FAIR vs REAL 展示结果前，先区分两种评估设定。 不同方法内存用法不同，直接对比易失真。摘要法引入额外 token，与剪枝法难直接比较。 因此我们采用互补的两种机制： ### FAIR：等有效容量 确保所有方法在**相同有效 KV 预算**下运行。 对 HAE，摘要引入额外 token（如质心），于是调整可用预算： `` effective_budget = k_budget − (bin_size − k_rank) `` ### REAL：实际内存占用 不调整，直接测每种方法的真实内存。 `` def kv_memory(K, V): return (K.numel() + V.numel()) * 4 / (1024 * 1024) `` 此设定反映真实部署：每存一个 token 都占内存。 ## 结果 ### 随时间演化的内存 Memory Evolution 上图展示序列推进时保留的 KV token 数。 两种行为： - **Top-K（蓝）**：严格上限，内存平坦 - **HAE（橙）**：阶梯状，内存先增后周期性压缩 #### 解读 Top-K **静态保留**： - 达预算后持续淘汰 - 被删 token 信息完全丢失 HAE **动态压缩**： - 先累积 token - 回收站满后总结压缩 - 缓存大小骤降再渐增 Step Behavior HAE 每级下降对应： `` 回收站满 → OLS 重建 → SVD 压缩 → 回插 `` --- ### KV 压缩策略对比（FAIR） KV Compression Benchmark 在 FAIR 设定下比较多种 KV 缓存压缩策略。 方法： - **Top-K**：保留注意力最高 token - **Sliding Window**：仅留最近 token - **OLS（秩k）**：无熵选择的低秩近似 - **HAE（Vanilla）**：仅熵选择无重建 - **HAE（熵+OLS）**：完整 SRC 流水线 #### 1. 完整 SRC 流水线最佳 各内存预算下，**HAE（熵+OLS）**重建误差最低。 - 低保留率（≤30%）时优势显著 - 即便激进压缩仍表现强劲 #### 2. 仅有选择不够 对比： - **HAE（Vanilla）** vs - **HAE（熵+OLS）** 可见： > 熵基选择虽优于朴素法，但重建至关重要。 无 OLS 时性能大跌，低预算下尤甚。 #### 3. 仅有压缩不足 **OLS（秩k）**基线几乎平坦： - 不区分 token 重要性 - 缺乏查询感知选择 说明： > 无选择的压缩无法捕捉有意义结构。 #### 4. Sliding Window 效果差 滑动窗口重建误差最高： - 忽略全局依赖 - 仅按时间局部性保留 #### 5. Top-K 强却脆弱 高预算下尚可，但： - 低比例时急剧劣化 - 遭遇前述选择性失效 --- ### FAIR vs REAL：精度与内存权衡 FAIR vs REAL 在两种设定下评估 SRC 流水线，揭示重建保真与内存占用权衡。 #### FAIR：重建误差 左图显示等有效容量下的重建误差（MSE）。 各保留率下 HAE 均优于 Top-K： - **30% 保留率**时误差低近**3×** - 50%、70% 时差距仍显著 - 高比例时两方法趋同 表明： > 在预算受限时，HAE 更有效地保持注意力功能。 #### REAL：内存足迹 右图显示真实内存消耗。 有趣的是，HAE 全程： > 内存占用**低于 Top-K** 原因在于： - HAE 把多 token 压成低秩表示 - 冗余信息被剔除而非原样存储 #### 综合解读 结果揭示关键区别： - **Top-K**： - 优化 token 保留 - 低预算时重建误差高 - **HAE**： - 优化功能保持 - 同时实现**更低误差**与**更低内存** #### 反思目标 发现挑战常见假设： > token 数更少未必更高效。 而应： - 按**信息密度**而非 token 计数评估内存 - 压缩在**精度与 footprint**上可双杀剪枝 ## 结论 本研究证实，通过数学引导的摘要可缓解 KV 缓存线性膨胀。分层注意力熵与低秩重建协同，将弥散信息“压实”成代表性质心，而非简单丢弃。 在合成 Transformer 上的基准显示，该方法保持了剪枝策略常丢失的注意力语义“形状”。主要代价是 OLS 与 SVD 步骤增加的计算时间。下一步工作将用定制 Triton 核融合这些操作以摊销延迟。以重建保真而非仅内存容量视角审视缓存，可构建更可持续的长上下文推理架构。 完整研究 notebook 已开放 Review：https://github.com/jayanthchandra/notebooks/blob/main/HAE_OLS.ipynb ## 引用 `` @misc{hae_kv_cache_2026, title = {Entropy-Guided KV Cache Summarization via Low-Rank Attention Reconstruction}, author = {Chandra, Jayanth}, year = {2026}, publisher = {Zenodo}, doi = {10.5281/zenodo.19657329}, url = {https://doi.org/10.5281/zenodo.19657329} } ``