工作记忆与上下文管理

引言

LLM 的上下文窗口（context window）是 Agent 的工作记忆。正如人类的工作记忆容量有限（Miller's 7±2 法则），LLM 的上下文窗口也有 token 数量限制。如何高效管理这一有限资源，直接决定了 Agent 的能力上限。

上下文窗口即工作记忆

Miller 定律与 Token 限制

George Miller（1956）发现人类短期记忆容量约为 7±2 个组块（chunks）。类比到 LLM：

模型	上下文窗口	等效"组块"
GPT-3.5	4K tokens	~3000 词
GPT-4	8K/32K tokens	~6000/24000 词
GPT-4 Turbo	128K tokens	~96000 词
Claude 3.5	200K tokens	~150000 词
Gemini 1.5 Pro	1M tokens	~750000 词

上下文的构成

一个典型 Agent 的上下文包含：

[System Prompt]         ~500-2000 tokens
[工具定义]               ~200-1000 tokens/tool
[对话历史]               变长
[检索到的上下文]          变长
[当前用户输入]            变长
─────────────────────────
总计必须 < 上下文窗口限制

注意力机制基础

Transformer 的自注意力机制是上下文处理的核心：

\[\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]

其中：

• \(Q\)（Query）、\(K\)（Key）、\(V\)（Value）是输入的线性变换
• \(d_k\) 是 Key 的维度，用于缩放避免梯度消失
• softmax 产生注意力权重分布

关键问题：自注意力的计算复杂度为 \(O(n^2)\)，其中 \(n\) 是序列长度。这意味着更长的上下文需要更多的计算资源。

"Lost in the Middle" 问题

Liu et al.（2023）发现 LLM 对上下文中部的信息利用率显著低于头部和尾部：

注意力利用率
^
|███                          ███
|████                        ████
|█████                      █████
|██████                    ██████
|████████              ████████
|██████████████████████████████
+─────────────────────────────→ 位置
 开头        中间         结尾

启示：将最重要的信息放在上下文的开头和结尾。

上下文压缩策略

1. 对话历史摘要

将过长的对话历史压缩为摘要：

def compress_history(messages, max_tokens=2000):
    """将对话历史压缩为摘要"""
    if count_tokens(messages) <= max_tokens:
        return messages

    # 保留系统提示和最近的消息
    system = messages[0]
    recent = messages[-4:]  # 保留最近 2 轮对话

    # 对中间部分生成摘要
    middle = messages[1:-4]
    summary = llm.summarize(middle)

    return [system, {"role": "system", "content": f"对话摘要：{summary}"}] + recent

2. 滑动窗口

保留最近 N 轮对话，丢弃更早的内容：

def sliding_window(messages, window_size=10):
    """滑动窗口策略"""
    system = [m for m in messages if m["role"] == "system"]
    conversation = [m for m in messages if m["role"] != "system"]

    # 保留最近 window_size 轮
    recent = conversation[-window_size * 2:]
    return system + recent

3. 递归摘要（Recursive Summarization）

def recursive_summarize(messages, chunk_size=10):
    """递归摘要：将历史分块摘要后再摘要"""
    if len(messages) <= chunk_size:
        return summarize(messages)

    chunks = [messages[i:i+chunk_size] for i in range(0, len(messages), chunk_size)]
    summaries = [summarize(chunk) for chunk in chunks]
    return recursive_summarize(summaries, chunk_size)

4. 检索增强的上下文管理

不存储完整历史，而是按需检索相关片段：

def retrieval_augmented_context(query, history_store, top_k=5):
    """基于当前查询检索相关历史"""
    # 将查询向量化
    query_embedding = embed(query)

    # 从历史存储中检索最相关的片段
    relevant = history_store.similarity_search(query_embedding, top_k=top_k)

    # 组装上下文
    context = "\n".join([f"[{r.timestamp}] {r.content}" for r in relevant])
    return context

上下文管理的高级策略

Token 预算分配

pie title 上下文 Token 预算分配（128K 窗口示例）
    "系统提示" : 5
    "工具定义" : 10
    "检索上下文" : 30
    "对话历史" : 35
    "当前输入" : 10
    "预留生成空间" : 10

动态上下文管理

class ContextManager:
    def __init__(self, max_tokens=128000):
        self.max_tokens = max_tokens
        self.budget = {
            "system": int(max_tokens * 0.05),
            "tools": int(max_tokens * 0.10),
            "retrieval": int(max_tokens * 0.30),
            "history": int(max_tokens * 0.35),
            "input": int(max_tokens * 0.10),
            "generation": int(max_tokens * 0.10),
        }

    def build_context(self, system_prompt, tools, query, history, retrieval_results):
        context = []

        # 1. 系统提示（必须完整）
        context.append(truncate(system_prompt, self.budget["system"]))

        # 2. 工具定义（按相关性排序）
        relevant_tools = rank_tools_by_relevance(tools, query)
        context.extend(fit_within_budget(relevant_tools, self.budget["tools"]))

        # 3. 检索结果（按相关性排序）
        context.extend(fit_within_budget(retrieval_results, self.budget["retrieval"]))

        # 4. 对话历史（最近的优先）
        compressed_history = self.compress_history(history, self.budget["history"])
        context.extend(compressed_history)

        # 5. 当前输入
        context.append(query)

        return context

    def compress_history(self, history, budget):
        """智能压缩对话历史"""
        if count_tokens(history) <= budget:
            return history

        # 策略：保留最近 N 轮 + 摘要
        recent = history[-6:]  # 最近 3 轮
        remaining_budget = budget - count_tokens(recent)

        if remaining_budget > 0:
            older = history[:-6]
            summary = summarize_to_budget(older, remaining_budget)
            return [summary] + recent
        else:
            return recent[-4:]  # 退化到保留最近 2 轮

超长上下文的策略（128K+）

对于支持超长上下文的模型：

1. 分层组织：用清晰的分隔符和标题组织上下文各部分
2. 重要性标记：对关键信息添加显式标记（如 [IMPORTANT]）
3. 结构化格式：使用 XML/JSON 结构化上下文，便于模型解析
4. 冗余消除：去除重复信息，避免浪费 token

def structure_context(sections):
    """使用 XML 标签结构化上下文"""
    context = ""
    for section in sections:
        context += f"<{section.tag} priority='{section.priority}'>\n"
        context += section.content
        context += f"\n</{section.tag}>\n\n"
    return context

上下文缓存（Context Caching）

Anthropic Prompt Caching

# Anthropic 的 prompt caching：对静态前缀进行缓存
response = client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    system=[
        {
            "type": "text",
            "text": long_system_prompt,  # 这部分会被缓存
            "cache_control": {"type": "ephemeral"}
        }
    ],
    messages=[{"role": "user", "content": user_query}]
)
# 缓存命中时，延迟降低 ~85%，成本降低 ~90%

Google Gemini Context Caching

# Gemini 的上下文缓存
cached_content = genai.caching.CachedContent.create(
    model="gemini-1.5-pro",
    contents=[large_document],
    ttl=datetime.timedelta(hours=1),
)

model = genai.GenerativeModel.from_cached_content(cached_content)
response = model.generate_content("分析这份文档...")

注意力优化技术

Sparse Attention

不计算所有 token 对之间的注意力，而是只关注局部窗口和全局锚点：

• Longformer：局部滑动窗口 + 全局注意力
• BigBird：随机注意力 + 局部注意力 + 全局注意力

KV Cache 优化

• PagedAttention（vLLM）：类似操作系统的分页内存管理
• Grouped-Query Attention（GQA）：减少 KV 头数以降低缓存大小
• Multi-Query Attention（MQA）：所有头共享同一组 KV

实践建议

上下文管理清单

• [ ] 明确 token 预算分配方案
• [ ] 实现对话历史的自动压缩
• [ ] 对工具定义进行动态加载（只加载相关工具）
• [ ] 将重要信息放在上下文的开头和结尾
• [ ] 使用结构化格式组织上下文
• [ ] 考虑使用 prompt caching 降低延迟和成本
• [ ] 监控实际 token 使用量并优化

常见陷阱

1. 不监控 token 使用：导致意外的上下文溢出
2. 简单截断：粗暴截断可能丢失关键信息
3. 忽视 Lost in the Middle：将所有检索结果堆在中间
4. 过度压缩：摘要丢失了关键细节

延伸阅读

• 思维链与推理模式 - CoT 如何利用上下文窗口进行推理
• Liu, N. F., et al. (2023). "Lost in the Middle: How Language Models Use Long Contexts"
• Packer, C., et al. (2023). "MemGPT: Towards LLMs as Operating Systems"
• Munkhdalai, T., et al. (2024). "Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention"

---

评论