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 的输入是一个 QA 问题,输出是交错 <think>/<search>/<answer> 的 agentic 轨迹,难点在于 0.5B 学生初始性能近乎为 0,既拿不到 RL 奖励也学不动 on-policy 蒸馏。它把整套训练拆成首尾相接的两段:先用教师答对的 TGO 轨迹离线蒸馏,把学生顶到能稳定产出合理行为骨架的起点;再以这个学生为初始策略跑 PPO,但重新设计奖励——答对给标量 +1,答错则把奖励替换成对教师的 KL 蒸馏惩罚 \(-\beta D_{\text{KL}}[\pi_\theta(y\mid x;\mathcal{R})\|\pi_g(y\mid x;\mathcal{R})]\),让错误样本也能从模仿教师中拿到密集信号。整个衔接靠一个性能阈值触发,不需要 TAID/DistiLLM 那种手工调 \(\alpha\) 插值系数的调度器。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["QA 问题"] --> B["冷启动 KD 初始化<br/>教师答对 TGO 轨迹离线蒸馏(CE + 软分布 KL)"]
B -->|"两阶段顺序:先 KD 再 PPO,不可颠倒"| C["蒸馏引导 PPO<br/>学生生成 think / search / answer 轨迹"]
C --> D{"答案是否正确"}
D -->|"答对"| E["奖励 +1<br/>自由探索,不受教师约束"]
D -->|"答错"| F["选择性 KL 惩罚<br/>−β·KL 拉回教师轨迹"]
E --> G["策略梯度更新"]
F --> G
G -->|"未达 1000 步则继续采样"| C
G --> H["紧凑 Agentic RAG 模型"]
关键设计¶
1. 冷启动 KD 初始化:先把学生顶过零性能门槛
紧凑模型最致命的问题是初始 EM 几乎为 0,PPO 长期拿不到正奖励,要么收敛极慢要么早期崩溃。DGPO 用教师答对的 TGO 轨迹做离线蒸馏,损失是硬标签与软分布的混合 \(\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)]\),\(\lambda\) 平衡两项。这里刻意只保留答对的轨迹,因为错误轨迹会把错误的检索决策也一并传给学生。
之所以用完整 KD 而非 SFT warm-start,是因为 SFT 只能学硬标签(实验里 SFT→PPO 仅 0.289),而软分布额外携带了"教师在歧义处如何权衡"这种细粒度信息——单靠这步 KD 初始化就能到 0.298,已经反超 SFT warm-start。
2. 选择性 KL 惩罚:让教师从正则锚点变成错误纠正者
标准 PPO 的奖励是 \(r_{\text{answer}}=\mathbb{1}[y=y^*]\),错误样本奖励恒为 0、毫无学习信号;DGPO 把它改成答对给 1、答错给 \(-\beta D_{\text{KL}}[\pi_\theta(y\mid x;\mathcal{R})\|\pi_g(y\mid x;\mathcal{R})]\),等于只对错误样本施加蒸馏惩罚。这样答对的样本保留自由探索空间不受教师约束,答错的样本被强力拉回教师轨迹,本质是用 dense 蒸馏奖励去填补 sparse 二值奖励的空白。
关键在于"选择性":消融里把它改成对所有样本统一施加蒸馏(uniform KL),平均 EM 从 0.329 掉到 0.314;完全去掉教师 guidance 退回普通 PPO 则掉到 0.306。这说明只纠正错误样本既保住了 RL 的探索能力,又避免了被噪声 SGO 带偏。
3. KD→PPO 的两阶段顺序:先有合理行为再做策略梯度
两阶段的先后不能颠倒——先跑 5 个 epoch KD 初始化,再跑最多 1000 步 distillation-guided PPO。一旦反过来让 PPO 在尚未成形的弱策略上起步,奖励会长期为 0 并迅速陷入退化策略,之后再补 KD 也救不回来。消融中 invert pipeline(PPO→KD)平均 EM 仅 0.286,比完整 DGPO 低 4.3 个点,印证了初始化顺序对紧凑模型尤其敏感。
损失函数 / 训练策略¶
KD 阶段取 \(\lambda=1.0\) 平衡 CE 与 KL;PPO 阶段取 \(\beta=0.001\) 控制错误样本的蒸馏强度,最多 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 场景。