跳转至

Can Compact Language Models Search Like Agents? Distillation-Guided Policy Optimization for Preserving Agentic RAG Capabilities

会议: ACL 2026
arXiv: 2508.20324
代码: https://github.com/omron-sinicx/dgpo
领域: 信息检索 / LLM Agent
关键词: Agentic RAG, 知识蒸馏, 强化学习, 紧凑模型, PPO

一句话总结

本文提出 DGPO:用教师 demonstration 做冷启动 KD 初始化,再在 PPO 阶段对"错误样本"施加 KL 蒸馏惩罚,让 0.5B 紧凑模型获得 Agentic RAG 能力,在 7 个 QA benchmark 上平均 EM 从 0.006 提升到 0.329,部分数据集甚至反超 3B 教师。

研究背景与动机

领域现状:Agentic RAG(如 Search-R1、ReAct)已成为 LLM 调用外部检索的主流范式,模型需要交错执行 <think><search><answer> 三类动作完成多跳问答。这类系统对 LLM 的推理能力要求很高,因此现有工作几乎都依赖数十亿参数以上的大模型。

现有痛点:作者尝试把 Search-R1 那一套 PPO 训练直接搬到 0.5–1B 的紧凑模型上,发现两条主流路径都失效。一是 RL 路径——紧凑模型初始 EM 几乎为 0(Qwen2.5-0.5B 在 7 个数据集上平均仅 0.006),导致 PPO/GRPO 长期拿不到正奖励,要么收敛极慢要么早期就崩溃;二是 KD 路径——纯 TGO 蒸馏存在 exposure bias,纯 SGO on-policy 蒸馏又被噪声样本带偏,DistiLLM/TAID 等动态调度方法对学生–教师 capacity gap 也很敏感。

核心矛盾:紧凑模型的 SGO 质量太差,既不能支撑 RL 的探索(无奖励信号),又不能支撑 on-policy 蒸馏(噪声目标);而离线 TGO 蒸馏又解决不了训练–推理分布不一致问题。根本矛盾是「冷启动质量」与「探索能力」无法兼得

本文目标:把 0.5–1B 紧凑模型训练成能像 agent 一样多轮搜索的检索模型,同时保证训练稳定性,并提供一套细粒度评估指标,定位 agentic 能力的具体短板。

切入角度:作者重新思考 PPO 中 reference model 的作用——传统 PPO 把它当作 KL 正则锚点(防止策略漂移),而本文把教师模型变成"主动指导者":学生答对时让它自由探索,学生答错时用 KL 蒸馏把它"拉回"教师轨迹。

核心 idea:用「KD 冷启动 + 选择性 KL 蒸馏惩罚」把蒸馏融进 PPO 内部,让 reference model 从被动的正则化器变成主动的教学者,从而在紧凑模型上获得稳定且超越教师的 agentic RAG 能力。

方法详解

整体框架

DGPO 分两个阶段:(1) 冷启动 KD 初始化——用教师生成的高质量 TGO 轨迹(仅保留答对的)对学生做离线蒸馏,让学生先学到合理的 <think>/<search>/<answer> 行为骨架;(2) 蒸馏引导的 RL——以冷启动学生为初始策略做 PPO,但把 reward 重新设计:答对时给标量 +1,答错时把 reward 替换为 \(-\beta D_{\text{KL}}[\pi_\theta(y\mid x;\mathcal{R})\|\pi_g(y\mid x;\mathcal{R})]\),让错误样本通过模仿教师获得密集学习信号。

整个过程不需要手工调度器(不像 TAID/DistiLLM 要调 \(\alpha\) 插值系数),两阶段衔接由一个性能阈值触发即可。

关键设计

  1. 冷启动 KD 初始化(仅用正确 TGO):

    • 功能:把学生从「near-zero 性能」拉到能产生有意义 reward 的起点。
    • 核心思路:用混合损失 \(\mathcal{L}_{\text{distill}} = \mathcal{L}_{\text{CE}}(\pi_g,\pi_\theta) + \lambda D_{\text{KL}}[\pi_g(\cdot\mid x)\|\pi_\theta(\cdot\mid x)]\) 在教师正确轨迹上蒸馏,既学硬标签又学软分布。作者验证只用正确 TGO 比混入错误 TGO 效果更好——因为错误轨迹会传递错误的检索决策。
    • 设计动机:SFT warm-start 只能学硬标签(实验表中 SFT→PPO 仅 0.289),而本方法直接复用教师的完整 soft distribution,包含了"当存在歧义时教师是怎么权衡的"这种细粒度信息,所以单 KD 初始化就能到 0.298,已超 SFT warm-start。
  2. 选择性 KL 惩罚(Selective KL Penalty):

    • 功能:在 PPO 训练中把教师从「被动正则锚点」转为「主动错误纠正者」。
    • 核心思路:标准 PPO 的 reward 是 \(r_{\text{answer}}=\mathbb{1}[y=y^*]\),错误样本 reward 恒为 0,无学习信号;本文改成 \(r_\phi(x,y)=1\) if correct else \(-\beta D_{\text{KL}}[\pi_\theta(y\mid x;\mathcal{R})\|\pi_g(y\mid x;\mathcal{R})]\),错误样本被替换为蒸馏惩罚。这样答对的样本保留自由探索空间(不受教师约束),答错的样本被强力拉向教师。
    • 设计动机:消融实验中把这条改成 uniform KL penalty(对所有样本都施加蒸馏),平均 EM 从 0.329 掉到 0.314;去掉教师 guidance(普通 PPO)掉到 0.306。说明"对错误样本选择性纠正"既保留了 RL 的探索能力,又能避免被噪声 SGO 带偏。
  3. KD→PPO 的两阶段顺序:

    • 功能:确保 PPO 起步时已经有 reasonable 的 agentic 行为,避免在崩坏的策略上做策略梯度。
    • 核心思路:先用 5 个 epoch 做 KD 初始化,再用最多 1000 步 PPO 做 distillation-guided RL,反过来(PPO→KD)会让 PPO 在弱策略上崩溃,再做 KD 也补不回来。
    • 设计动机:消融中 invert pipeline(PPO→KD)平均 EM 仅 0.286,比 DGPO 低 4.3 个点,说明初始化顺序对紧凑模型尤其关键——一旦 PPO 阶段 reward 长期为 0 就会陷入退化策略。

损失函数 / 训练策略

KD 阶段:\(\lambda=1.0\) 平衡 CE 与 KL;PPO 阶段:\(\beta=0.001\) 控制错误样本的蒸馏强度,max 4 轮对话,每轮 top-3 文档检索(E5 retriever,2018 Wiki dump)。token masking 仅在 LLM 生成 token 上算梯度,对 <information> 段不算梯度。训练在 NVIDIA 8×H200 上跑约 1 天。

实验关键数据

主实验

Qwen 2.5(3B 教师 → 0.5B 学生)在 7 个 QA benchmark 上的 EM 对比:

方法 NQ TriviaQA HotpotQA MuSiQue Bamboogle Avg 备注
Student-0.5B 0.004 0.006 0.007 0.000 0.000 0.006 未训练
Teacher-3B 0.365 0.569 0.340 0.135 0.298 0.353 上限参考
PPO (Search-R1) 0.306 0.444 0.205 0.041 0.073 0.238 RL baseline
SFT→PPO 0.338 0.415 0.296 0.088 0.250 0.289 warm start
KD (Hinton) 0.331 0.431 0.286 0.091 0.290 0.298 离线蒸馏
DistiLLM 0.333 0.442 0.288 0.095 0.209 0.287 自适应蒸馏
TAID 0.325 0.427 0.290 0.079 0.218 0.282 调度蒸馏
DGPO 0.378 0.481 0.342 0.120 0.274 0.329 在 NQ/HotpotQA 等反超 3B 教师

跨模型族也验证了通用性:Qwen 7B→0.5B 平均 EM 从 PPO 的 0.238 提升到 0.323;Llama-3 8B→1B 从 0.250 提升到 0.389,仅差 8B 教师 4.9 个点。

消融实验

配置 Avg EM 说明
DGPO 完整 0.329 KD 冷启动 + KD→PPO + 选择性 KL
w/o 冷启动初始化 0.320 训练在 step 800 后崩溃,取崩前最高分
w/o 选择性 KL(改 uniform) 0.314 对正确样本也蒸馏,约束过强
w/o 教师 guidance(标准 PPO) 0.306 错误样本无学习信号
invert pipeline(PPO→KD) 0.286 顺序颠倒,PPO 阶段已崩坏

关键发现

  • 冷启动 KD 是稳定性关键:去掉后训练在 step 800 崩溃,说明紧凑模型不能裸跑 PPO;但崩前的最高分(0.320)跟 DGPO 差距不大,说明性能主要来自蒸馏,稳定性 来自 KD 初始化。
  • 选择性 KL > 统一 KL:把蒸馏只施加在错误样本上比对所有样本都施加好 1.5 个点,证明"自由探索 + 错误纠正"的组合优于纯模仿。
  • 细粒度 ARCap 拆解:在 NQ(单跳)上 Query Rewriting Hit ratio 反而是普通 PPO 最高(0.711 > 教师 0.682),DGPO 0.682 与教师持平;MuSiQue(多跳)上 DGPO 的 Hit ratio(0.583)和搜索步数(2.64)都最高,说明紧凑模型靠"多次搜索补偿单次推理弱"来达成多跳推理。
  • GRPO 不适合紧凑模型:GRPO 收敛快但早崩,即使加上 KD 初始化和教师 guidance 也不稳定,作者全文实验都用 PPO。

亮点与洞察

  • 重新定义 reference model 的角色:从 KL 锚点(被动)变成教学者(主动)。这一观念迁移到任何 PPO-style 微调都可能有用——只要存在一个比当前策略强的 reference,就可以把它从「正则项」升级为「错误纠正项」。
  • "分类讨论 reward"是个简单又强力的 trick:答对给标量、答错给蒸馏惩罚,本质上是用 dense distillation reward 填补 sparse binary reward 的空白。
  • ARCap 评估框架可复用:把 agentic 行为拆成 thinking / query rewriting / source referencing 三维度,并设计了"提供 ground-truth context 测纯回答能力"、"测首轮 query 的 Hit ratio"等隔离实验,这种"控制变量地评估 agent 子能力"的设计思路值得借鉴。
  • 0.5B + CPU 可跑:把 agentic RAG 从云端搬到笔记本/手机,55× 性能提升后已经接近 3B 教师。

局限与展望

  • 模型族仅在 Qwen2.5 和 Llama-3 上验证,没扩展到 Mistral/Phi 等其他紧凑模型族。
  • 教师上限是 8B,超大教师(70B+)下 capacity gap 是否仍能被 KD 弥补未知;作者承认这是算力限制。
  • 蒸馏引入约 9.5% 额外训练时间(教师 inference),虽小但在更大教师下会显著增加。
  • 仅在 QA 任务验证 agentic 行为,对 code、math、tool use 等其他 agentic 场景的迁移性留作 future work。
  • 个人补充:选择性 KL 的"正确/错误"边界依赖 EM,而 EM 在 free-form generation(如摘要、对话)里不适用,迁移这套方法到非 QA 任务时需要重新定义 reward 切分逻辑。

相关工作与启发

  • vs Search-R1 (Jin et al. 2025):Search-R1 在 7B+ 上跑 PPO 没问题,但在 0.5B 上 reward 太稀疏;本文用蒸馏给错误样本补 dense reward,专门解决紧凑模型问题。
  • vs DistiLLM (Ko et al. 2024) / TAID (Shing et al. 2025):他们用 \(\alpha\) 插值在 teacher/student 分布间做动态调度,对超参敏感;DGPO 用两阶段 + 选择性 KL,无需调度器。
  • vs GKD (Agarwal et al. 2024):纯 on-policy SGO 蒸馏,但紧凑模型 SGO 噪声大,所以 GKD 平均仅 0.240;DGPO 把 on-policy 探索和 off-policy KD 解耦到两阶段。
  • vs DeepSeek-R1 cold-start:R1 用 SFT 做 cold-start,本文用完整 KD(含 soft distribution)替换 SFT,证明在小模型上 soft target 比 hard target 更值钱。

评分

  • 新颖性: ⭐⭐⭐⭐ 「reference model 当老师」这个视角清新,但单独的 KD 初始化和选择性 KL 都不是首创。
  • 实验充分度: ⭐⭐⭐⭐⭐ 7 个 QA 数据集 × 3 个模型配置 × 5 个 ablation 维度,还配了 ARCap 子能力评估,相当完整。
  • 写作质量: ⭐⭐⭐⭐ 图 1/图 5 把 idea 讲得很清楚,Limitations 也诚实;个别公式(如 Eq.2 的 PPO objective)排版有点挤。
  • 价值: ⭐⭐⭐⭐⭐ 把 agentic RAG 拉到 0.5B 尺度可以跑端侧,实际意义大;方法可直接复用到其他 PPO-on-small-model 场景。