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\) 插值系数),两阶段衔接由一个性能阈值触发即可。
关键设计¶
-
冷启动 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。
-
选择性 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 带偏。
-
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 场景。