简介

在大语言模型 (LLM) 时代, 上下文学习 (In-Context Learning, ICL) 已成为一种主导范式。其理念看似简单: 你无需微调模型的权重，只需在提示词 (Prompt) 中提供几个例子 (范例) ，模型就能学会其中的模式。

例如，如果你想让 LLM 将英语翻译成 SQL，你的提示词可能是这样的:

Input: Show me users over 20. Output: SELECT * FROM users WHERE age > 20;

Input: Count the products. Output: SELECT count(*) FROM products;

Input: [你的实际查询] Output: …

ICL 的成功在很大程度上取决于你选择了哪些示例。传统上，我们将此视为一个简单的搜索问题: 在训练集中找到与当前输入最相似的 \(k\) 个示例。我们抓取前 5 或 10 个，把它们堆叠起来，然后喂给 LLM。

但问题在于: 这种方法假设示例之间是相互独立的。它忽略了示例之间的“化学反应”。如果第二个示例重复了第一个示例的信息怎么办？如果示例 A 和示例 B 的组合会让模型感到困惑，即使它们单独来看都很相关呢？

在这篇文章中，我们将深入探讨微软和约翰霍普金斯大学的研究人员发表的一篇引人入胜的论文，题为 “Learning to Retrieve Iteratively for In-Context Learning” (学习为上下文学习进行迭代检索) 。他们提出了一种名为 迭代检索 (Iterative Retrieval, IterR) 的新框架。他们的模型不是一次性抓取一把示例，而是一个接一个地挑选，并在每次选择后更新其对上下文的理解。

我们将探讨他们如何将其构建为一个强化学习问题，如何创建一个“有状态”的检索器，以及为什么这种方法在语义解析等复杂任务上显著优于标准方法。

背景: “快照式”检索的缺陷

要理解解决方案，我们首先需要了解当前标准的局限性。

在典型的检索增强生成 (RAG) 或 ICL 设置中，我们有一个查询 \(x\)。我们希望从数据集 \(\mathcal{D}\) 中找到一组范例。标准的检索器 \(R\) 基于评分函数 \(S\) (通常是余弦相似度或 BM25 关键词匹配) 选择 top-\(k\) 个项目:

展示基于评分函数 S 的标准 top-k 检索的公式。

这是一种“快照式”方法。检索器查看查询，计算所有文档的分数，并返回分数最高的文档。它不关心所选项目的顺序，也不关心项目之间如何关联。

论文作者认为，寻找最佳提示词集实际上是一个 组合优化问题 。 理想的检索器应该最大化语言模型 (LM) 在给定范例序列的情况下生成正确输出 \(y\) 的概率:

展示在给定范例序列下最大化 LM 概率的公式。

由于检查每一种可能的示例组合在计算上是不可能的 (NP 难问题) ，我们需要一种比仅仅“挑选最相似的”更聪明的近似方法。

可视化差异

下图完美地展示了思维方式的转变。

图 1: 上图: 单次检索调用下的 ICL。下图: 我们提出的迭代检索器下的 ICL。

  • 上方 (标准) : 模型基于输入 \(x\) 一次性检索一批示例 \((x_1, y_1), (x_2, y_2)\)。
  • 下方 (迭代) : 模型检索一个示例，更新其内部状态 (\(s_0 \to s_1\)) ，考虑它已经学到了什么，然后再决定接下来检索什么。

核心方法: 迭代检索

研究人员建议将提示词构建过程建模为 马尔可夫决策过程 (MDP) 。 如果你熟悉强化学习 (RL) ，这听起来会很耳熟。如果不熟悉，可以把它想象成一个回合制游戏，其中“代理” (检索器) 采取“动作” (选择示例) 以最大化“得分” (LLM 的表现) 。

1. IterR 的架构

迭代检索器( IterR )由几个关键组件组成: 状态向量、策略网络和状态转移函数。

状态 (\(s\))

状态代表“当前上下文”。它跟踪检索器目前为止选择了什么。初始状态 \(s_0\) 是一个可学习的参数。

动作 (检索)

在每一步 \(i\)，模型需要从数据集中挑选一个范例对 \((x_i, y_i)\)。它使用 策略网络 (Policy Network) 来完成此操作。策略根据当前状态 \(s_i\) 决定哪个示例最有用。

策略由从当前状态生成的查询向量与候选示例的编码向量之间的内积 (相似度) 定义:

展示基于状态 Q(s) 和编码器 F(x) 的策略概率分布的公式。

这里:

  • \(\mathbf{Q}(s_i)\) 将当前状态转换为搜索查询。
  • \(\mathbf{F}_{\text{enc}}(x_i)\) 是一个文本编码器 (如 Contriever) ，将文本转换为向量。
  • 这看起来很像标准的密集检索，但关键区别在于 \(\mathbf{Q}(s_i)\) 在每一步都会发生变化。

为了在推理过程中实际挑选示例，模型执行贪婪选择 (挑选得分最高的项目) :

展示使用最大内积搜索贪婪选择下一个范例的公式。

转移 (更新状态)

一旦选择了一个范例，模型需要记住它。作者使用 门控循环单元 (GRU) , 一种循环神经网络 (RNN) ，来更新状态。

展示使用 GRU 更新状态的公式。

新状态 \(s_{i+1}\) 是旧状态和新选示例的向量表示的混合。这就是使检索器具备“有状态”特性的原因。

综合示例

让我们看一个来自语义解析任务 (将自然语言转换为代码) 的具体例子。

图 2: SMCALFLOW 上 BM25 与迭代检索器的比较。

在上图中:

  • 上方 (BM25): 标准检索器看到 “tailgating party” (车尾野餐会) ，就只是寻找这些关键词。它检索到的示例虽然共用了单词，但对于所需的特定逻辑可能没有帮助。
  • 下方 (迭代):
  1. 模型首先检索了一个关于 “video camp” 的示例 (\(x_1\))。这建立了一个结构模式。
  2. 状态更新 (\(s_1\))。
  3. 基于 \(s_1\)，模型意识到它需要一种不同类型的约束示例。它调整其搜索轨迹并检索了一个关于 “circus” (马戏团) 的示例 (\(x_2\))。
  4. 最终的提示词包含了一组多样化的示例，帮助 LLM 解决特定的逻辑难题。

2. 利用强化学习进行训练

我们如何训练它？我们不能使用标准的监督学习，因为我们没有完美提示词列表的“标准答案”。相反，我们必须让模型尝试构建提示词，并观察 LLM 的表现如何。

这是 强化学习 的经典设置。

模拟器

在这个 RL 设置中的“环境”是 LLM 本身 (例如 Llama-2) 。检索器挑选示例，构建提示词，将其喂给 LLM，并等待反馈。

奖励函数

定义奖励很棘手。如果我们仅仅因为答案正确给 +1 奖励，错误给 0 奖励，信号就会过于稀疏 (模型几乎学不到东西) 。

相反，作者使用了 概率变化 。 他们测量所选示例在多大程度上 增加 了 LLM 生成正确金标准输出 (\(y^*\)) 的可能性。

基于概率提升的奖励函数公式。

  • \(P_{LM}(y^* | ... (x_i, y_i))\): 添加新示例 后 正确答案的概率。
  • \(P_{LM}(y^* | ...)\): 添加示例 前 的概率。

如果新示例使正确答案的可能性更高，模型就会获得正奖励。如果它让 LLM 感到困惑，就会获得负奖励。

策略优化 (PPO)

为了更新模型权重，他们使用了 近端策略优化 (PPO) 。 PPO 是一种流行的 RL 算法，以稳定性著称。它确保策略不会在单次更新中发生过于剧烈的变化。

展示 PPO 损失函数的公式。

模型被训练为最大化预期奖励 (优势 \(\hat{A}\)) ，同时保持新策略与旧策略相近 (截断项) 。

挑战: 分层采样

这里存在一个计算障碍。数据集 \(\mathcal{D}\) 可能包含 100,000+ 个候选项。在每一步计算策略的完整概率分布太慢了。

为了解决这个问题，作者使用了 分层采样 (Stratified Sampling) 。

图 4: 用于平衡探索与利用的分层采样方法。

他们不查看整个数据集或仅仅查看前几个，而是:

  1. 检索一个包含头部候选项的缓冲区。
  2. 保留最好的几个 (利用) 。
  3. 将其余的归入“层” (strata) 并从中随机采样 (探索) 。
  4. 重新归一化分数。

这确保了模型专注于好的候选项，但仍然会探索那些不太明显但可能出奇有用的选项。

实验与结果

作者在 语义解析 上测试了 IterR，这是一项要求模型将自然语言翻译成形式化表示 (如代码) 的任务。这是一个很好的测试平台，因为它要求严格遵守语法和逻辑——仅仅“接近”是不够的。

数据集

他们使用了三个数据集:

  1. SMCALFLOW: 复杂的对话数据流合成。
  2. TREEDST: 对话状态跟踪。
  3. MTOP: 面向任务的解析 (闹钟、提醒等) 。

图 3: 语义解析任务下的 (x, y) 样本对。

主要结果: 击败基线

作者将 IterR 与以下方法进行了比较:

  • BM25: 标准关键词搜索。
  • Contriever: 强大的密集检索器 (向量相似度) 。
  • EPR & CEIL: 之前最先进的范例选择方法。

结果 (表 1) 显示了迭代检索的明显胜利。

表 1: IterR 与基线在精确匹配和 SMatch 指标上的比较。

关键要点:

  • 精确匹配 (EM): IterR 显著优于所有基线。在 SMCALFLOW 上，它从 51.1% (CEIL) 跃升至 54.1% 。
  • SMatch: 这个指标衡量部分正确性 (结构是否大部分正确？) 。IterR 在这里也占据主导地位，表明它比其他方法更能帮助 LLM 理解任务的底层结构。

泛化性: 它在其他 LLM 上有效吗？

对于任何使用 RL 训练的组件，一个关键问题是: “你是否只是对模拟器过拟合了？”作者使用 Llama-2-7b 训练了 IterR。但是，如果我们使用 不同 的 LLM 进行推理，训练好的检索器还有效吗？

图 5: 使用各种推理 LLM 的性能比较。

答案是 肯定的 。

  • 家族内 (Intra-family): 它在更大的 Llama 模型 (Llama-2-70b) 上效果甚至更好。
  • 家族间 (Inter-family): 它甚至可以泛化到 Mistral-7B (一个完全不同的模型家族) ，并保持对基线的领先优势。

这表明 IterR 正在学习“什么是好的提示词组合”的普遍属性，而不仅仅是利用 Llama-2-7b 的特定偏差。

效率: 少即是多

最后，这种更智能的检索是否允许我们使用更少的示例？

图 6: 不同范例数量下的性能比较。

图 6 显示，与 EPR (蓝色) 和 CEIL (粉色) 相比，IterR (绿线) 使用更少的范例就能达到更高的性能。这是一个主要的实际好处: 更短的提示词意味着更低的延迟和更低的成本。

结论

论文 “Learning to Retrieve Iteratively for In-Context Learning” 代表了我们对提示工程思考方式的转变。它让我们从静态的、基于启发式的相似性搜索，转向 动态的、学习型的和有状态的检索 。

通过将提示词构建构建为一个顺序决策过程, IterR 框架承认:

  1. 顺序很重要: 第一个例子为第二个例子奠定了基础。
  2. 上下文很重要: 模型需要什么取决于它已经看到了什么。
  3. 反馈很重要: 检索器的最终裁判应该是下游模型的表现。

虽然该方法增加了一些训练复杂性 (强化学习以棘手著称) ，但推理成本很低——仅在检索过程中增加了一个小的 GRU (4M 参数) 。对于像语义解析或代码生成这样需要高精度的应用，这种迭代方法提供了优于标准 RAG 技术的显著优势。

随着 LLM 的不断发展，为它们 提供 信息的系统也必须随之发展。迭代检索证明，选择正确的上下文不仅仅是寻找相似的文档——它是关于为模型构建一个连贯的课程，一步一个脚印。