上下文压缩与检索优化

KV Cache压缩：减少显存占用的核心手段

H2O：Heavy-Hitter Oracle

H2O基于一个关键观察：少量KV对贡献了绝大部分注意力权重。在典型注意力分布中，前20%的KV对占据约85%的注意力质量。

核心算法：

1. 累积注意力分数：追踪每个KV对在所有query中的累积attention score
2. 保留Heavy Hitters：只保留分数最高的$k$个KV对（如$k=20\%$）
3. 滑动窗口：始终保留最近的$w$个token（保证局部连贯性）

flowchart TD
    A[输入KV Cache<br/>长度L] --> B{计算累积注意力分数}
    B --> C[全局Top-K<br/>保留Heavy Hitters]
    B --> D[局部滑动窗口<br/>保留最近W个]
    
    C --> E[合并保留集合]
    D --> E
    
    E --> F[压缩后KV Cache<br/>长度≈K+W]
    F --> G[后续层使用<br/>压缩KV]

性能数据（在LLaMA-2-7B上）：

关键发现：保留20% Heavy Hitters + 512滑动窗口，可在减少75%显存的同时保持99%的准确率。

SnapKV：自适应压缩

SnapKV改进了H2O的静态策略，采用自适应压缩率：

• 对注意力稀疏的层使用更高压缩率
• 对注意力密集的层保留更多KV
• 根据输入动态调整保留比例

自适应公式：

$\text{retention\_ratio}_l = \alpha \cdot \text{entropy}(A_l) + \beta$

其中$A_l$为第$l$层的注意力矩阵，$\text{entropy}$衡量注意力分布的集中程度。

优势：

• 平均压缩率提升至80-85%
• 准确率损失<2%
• 支持动态长度输入

KV Cache量化

量化是另一类压缩手段，将FP16的KV Cache压缩到更低精度：

KVQuant技术：

• 分组量化：每64个token为一组共享缩放因子
• 离群值保留：对极端值使用FP16单独存储
• 动态范围调整：根据激活分布调整量化范围

实际效果：

• 4-bit量化可将128K上下文的KV Cache从64GB压缩到16GB
• 在A100 40GB上可支持64K上下文推理
• 延迟增加<5%（量化/反量化开销）

Prompt压缩：减少输入长度

LLMLingua：基于困惑度的压缩

LLMLingua采用小型语言模型（如GPT-2）对prompt进行压缩：

1. 分块计算困惑度：将prompt分块，计算每块的perplexity
2. 保留信息密度高的块：移除低困惑度（信息量少）的块
3. 语义连贯性检查：确保压缩后的prompt语义完整

压缩效果：

Selective Context：基于自信息

Selective Context利用token级别的自信息（self-information）进行压缩：

$I(x_i) = -\log P(x_i | x_{<i})$

自信息低的token（即模型容易预测的token）通常信息量小，可以移除。

优势：

• 细粒度压缩，可精确到token级别
• 保持关键实体和术语
• 支持结构化压缩（保留句子边界）

LongLLMLingua：面向长上下文的优化

针对长文档检索任务，LongLLMLingua引入：

• 问题感知压缩：根据query调整压缩策略
• 重排序增强：压缩后重新计算文档与query的相关性
• 多粒度压缩：文档级→段落级→句子级分层压缩

flowchart LR
    subgraph 原始文档["原始文档 (100K)"]
        D1[Doc1] --> D2[Doc2]
        D2 --> D3[Doc3]
        D3 --> D4[...]
    end
    
    subgraph 粗筛选["文档级筛选 (20K)"]
        S1[Doc1] --> S3[Doc3]
        S3 --> S5[Doc5]
    end
    
    subgraph 精压缩["段落级压缩 (8K)"]
        C1[Para1] --> C2[Para3]
        C2 --> C3[Para7]
    end
    
    subgraph 最终Prompt["最终Prompt (4K)"]
        F1[压缩结果]
    end
    
    原始文档 -->|BM25粗筛| 粗筛选
    粗筛选 -->|LLM重排序| 精压缩
    精压缩 -->|困惑度过滤| 最终Prompt

RAG增强：检索而非记忆

基础RAG架构

检索增强生成（RAG）将长上下文问题转化为检索问题：

flowchart TD
    A[长文档集合] --> B[文档切分<br/>Chunking]
    B --> C[向量化<br/>Embedding]
    C --> D[向量数据库<br/>FAISS/Milvus]
    
    E[用户Query] --> F[Query向量化]
    F --> G[相似度检索<br/>Top-K]
    D --> G
    
    G --> H[检索结果<br/>K个Chunks]
    H --> I[Prompt构建<br/>Query + Chunks]
    I --> J[LLM生成<br/>短上下文]

关键参数：

• Chunk大小：通常256-512 tokens
• 重叠：20-50%以避免边界信息丢失
• Top-K：5-10个chunks

长上下文RAG优化

标准RAG在超长文档（如整本书）上的局限：

• 检索粒度太细，丢失全局结构
• 多跳推理困难（需要跨多个chunk推理）

改进方案：

1. 分层索引：

   • 第一层：章节级粗粒度索引
   • 第二层：段落级细粒度索引
   • 先粗筛再精排

2. 摘要增强检索：

   • 为每个chunk生成一句话摘要
   • 先检索摘要，再获取完整chunk
   • 减少向量噪声

3. 重排序（Reranking）：

   • 初筛：向量相似度召回Top-100
   • 精排：Cross-encoder计算query-document相关性
   • 最终保留Top-5

Self-RAG：自适应检索

Self-RAG让模型自己决定何时检索、检索什么：

1. 反思token：模型在生成过程中输出特殊token（如[Retrieve]）
2. 按需检索：只在需要事实验证时触发检索
3. 多次检索：支持多跳推理的多次迭代检索

效果对比（在Natural Questions上）：

关键优势：Self-RAG用更少的检索次数获得更高的准确率，延迟降低28%。

层次化摘要：渐进式压缩

多级摘要架构

对于超长文档（如100K+），采用金字塔式摘要：

flowchart TD
    L0[原始文档<br/>100K tokens] --> L1
    
    subgraph L1["第1层摘要 (20K)"]
        S1[段落摘要1<br/>~200字]
        S2[段落摘要2<br/>~200字]
        S3[...<br/>100个段落]
    end
    
    L1 --> L2
    
    subgraph L2["第2层摘要 (4K)"]
        SS1[章节摘要1<br/>~500字]
        SS2[章节摘要2<br/>~500字]
        SS3[...<br/>8个章节]
    end
    
    L2 --> L3[全局摘要<br/>1K tokens]

查询流程：

1. 基于全局摘要定位相关章节
2. 在章节摘要中定位相关段落
3. 读取原始段落获取细节

动态摘要生成

不预先生成所有摘要，而是按需生成：

1. 流式处理：边读取文档边生成摘要
2. 增量更新：新信息到来时更新摘要
3. 遗忘机制：移除过时或低优先级信息

Map-Reduce摘要：

• Map阶段：并行生成段落摘要
• Reduce阶段：递归合并摘要
• 支持任意长度文档

混合策略：组合多种技术

最佳实践架构

flowchart TD
    A[长文档输入<br/>>32K] --> B{长度判断}
    
    B -->|<32K| C[直接输入]
    B -->|32K-128K| D[RAG + 压缩]
    B -->|>128K| E[分层摘要 + RAG]
    
    D --> D1[LLMLingua压缩<br/>50%]
    D --> D2[H2O KV Cache<br/>保留30%]
    D1 --> D3[LLM处理]
    D2 --> D3
    
    E --> E1[三级摘要]
    E --> E2[向量检索Top-10]
    E1 --> E3[构建Prompt]
    E2 --> E3
    E3 --> E4[YaRN外推<br/>128K→256K]
    E4 --> E5[4-bit KV Cache]
    E5 --> E6[LLM生成]

性能与成本权衡

参考资料

1. H2O: Heavy-Hitter Oracle for Accurate KV Cache Compression (Zhang et al., 2023)

   • KV Cache压缩的开创性工作，Heavy-Hitter发现

2. LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models (Jiang et al., 2023)

   • 基于困惑度的Prompt压缩

3. Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks (Lewis et al., 2020)

   • RAG基础架构

4. Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection (Asai et al., 2023)

   • 自适应检索机制

5. LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios (Jiang et al., 2023)

   • 面向长上下文的RAG优化

6. SnapKV: Efficient KV Cache Compression for Long-Context Inference (Li et al., 2024)

   • 自适应KV Cache压缩