自蒸馏策略梯度

# 自蒸馏策略梯度

来源：https://arxiv.org/html/2606.04036

Yifeng Liu　Shiyuan Zhang¹¹脚注标记：1　Yifan Zhang¹¹脚注标记：1　Quanquan Gu²²脚注标记：2　等同贡献　美国加州大学洛杉矶分校计算机科学系；邮箱：liuyifeng@cs\.ucla\.edu　美国加州大学洛杉矶分校计算机科学系；邮箱：zsy25ucla@ucla\.edu　美国普林斯顿大学 Princeton AI Laboratory；邮箱：yifzhang@princeton\.edu　通讯作者，美国加州大学洛杉矶分校计算机科学系；邮箱：qgu@cs\.ucla\.edu

###### 摘要

在线策略自蒸馏（on-policy self-distillation）是一种有前景的密集监督来源，适用于稀疏奖励强化学习——语言模型以特权上下文为条件，对自身生成结果进行自我监督。实际上，它可以被实例化为一个辅助的全词表师生反向 Kullback-Leibler 散度损失。为此，我们提出 SDPG，一个自蒸馏策略梯度框架，将基于组相对验证器优势与归一化标准差的方法、精确的全词表在线策略自蒸馏，以及参考策略 KL 正则化相结合。实验表明，SDPG 在稳定性和性能上均优于 RLVR 和自蒸馏基线。代码已发布于 https://github.com/lauyikfung/SDPG。

提示词 x　特权上下文 c　πθ 共享模型　学生 pt=πθ(·∣x,y0]　mi=**1**[A(i)out>0]　Lout　在线策略策略梯度　β(k)L+OPD + 门控 + 调度　αLK(πθ,πref)　参考 KL 正则化（UFKL/URKL）　πref 固定　**LSDPG=Lout+β(k)L+OPD+αLK(πθ,πref)**　无 c　有 c　调制输入　策略分布　信号　损失　目标

图 1：自蒸馏策略梯度（SDPG）目标概览，融合了基于 rollout 的结果优化、来自特权上下文的门控在线策略蒸馏（OPD），以及对固定参考策略的 KL 正则化。注意，OPD 本身也是一种策略梯度形式。

## 1 引言

随着带可验证奖励的强化学习（RLVR）的发展，大语言模型（LLMs）在数学、代码生成等复杂推理任务中展现出卓越的能力。这一系列算法，如群体相对策略优化（GRPO）\(Shao et al., 2024 (https://arxiv.org/html/2606.04036#bib.bib4)\)，针对基于规则的结果奖励进行优化，已成为后训练推理模型的标准方案，消除了人类偏好标注的成本与偏差。

尽管取得了显著成效，RLVR 仍面临若干局限，包括跨 token 的序列级稀疏奖励，以及训练早期负优势带来的不稳定性。尽管近期工作如 Dr.GRPO\(Liu et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib30)\)、DAPO\(Yu et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib12)\) 和 GSPO\(Zheng et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib31)\) 通过非对称双向裁剪阈值和序列级优势解决了后者，稀疏性问题仍未得到解决。

近期，在线策略蒸馏（OPD）方法被提出，以产生密集的 token 级信号\(Agarwal et al., 2024 (https://arxiv.org/html/2606.04036#bib.bib44); Lu and Lab, 2025 (https://arxiv.org/html/2606.04036#bib.bib47); Fu et al., 2026 (https://arxiv.org/html/2606.04036#bib.bib50)\)。此类方法维护两个模型：一个待优化的学生模型用于展开轨迹，一个教师模型通过 Kullback–Leibler 散度（KL）正则化或相关目标提供 token 级引导\(Gu et al., 2024 (https://arxiv.org/html/2606.04036#bib.bib46); Xu et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib45); Yang et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib49)\)。然而，传统蒸馏方法使用规模更大、能力更强的教师模型，在优化学生模型时会带来相当大的显存负担。此外，异构的教师信号可能影响训练过程的平滑性。

近期一系列工作通过在线策略自蒸馏解决了上述局限。在这些方法中，教师模型与学生模型完全相同，但附加了示例、直接答案和推理路径等额外知识\(Hübotter et al., 2026 (https://arxiv.org/html/2606.04036#bib.bib18); Shenfled et al., 2026 (https://arxiv.org/html/2606.04036#bib.bib5); Penaloza et al., 2026 (https://arxiv.org/html/2606.04036#bib.bib51)\)。这将稀疏且不一致的结果奖励转化为密集、逐 token 且同质的监督信号。具体而言，OPCD\(Ye et al., 2026 (https://arxiv.org/html/2606.04036#bib.bib1)\) 将上下文知识引入教师模型，并通过 KL 散度将其内化到学生模型；OPSD\(Zhao et al., 2026 (https://arxiv.org/html/2606.04036#bib.bib17)\) 应用全词表级 KL 散度以提升推理性能；TRRD\(Zhang et al., 2026b (https://arxiv.org/html/2606.04036#bib.bib52)\) 在蒸馏中引入信任域。

然而，"自蒸馏"这一表述可能遮蔽其有用的策略梯度解释。对于固定的 rollout 前缀 (x, y0<εstd>0 可避免除零。等价地，实现中通常在组内所有奖励相同时令 A(i)=0。随后，策略通过 PPO 风格\(Schulman et al., 2017 (https://arxiv.org/html/2606.04036#bib.bib19)\) 的裁剪代理目标进行优化：

LGRPO(θ)=−Ex,{y(i)}i=1G∼πold(·∣x)[1∑i|y(i)|∑i=1G∑t=1|y(i)|min(ri,tA(i),clip(ri,t,1−ε,1+ε)A(i))],

(2.2)

其中 ri,t 为重要性比率，定义如下：

ri,t=πθ(yt(i)∣x,y0<εstd>0。

**对每个训练步骤 k=1,...,T 执行以下操作：**

采样一批提示词和特权上下文 {(xj,cj)}j=1B∼D

**对每个提示词 xj 执行以下操作：**

// 从冻结的无特权行为策略中进行 rollout

采样 G 个响应 {yj(i)}i=1G∼πθ(·∣xj)

// 计算结果奖励和组相对优势

获取二值验证器奖励 Rj(i)=R(xj,yj(i))

计算 A(i)out=Rj(i)−μj/(σj+εstd)，其中 μj 和 σj 分别为 {Rj(i)}i=1G 的均值和标准差

令 mj(i)=**1**[A(i)out>0]

**对每个响应 yj(i) 和 token 位置 t 执行以下操作：**

定义前缀状态 sj,i,t=(xj,yj,0<q̄t(a)>0，只要 p̄t(a)>0。在教师分支被 detach 的情况下，反向 KL 全词表 OPD 损失

LOPD,t(θ)=DKL(pt∥q̄t)

与当前迭代处具有相同的学生侧梯度，如 detached 采样策略梯度代理

L̃PG OPD,t(θ)=−Ea∼p̄t[Adist t(a)log pt(a)]，Adist t(a)=SG[D̄t−log(p̄t(a)/q̄t(a))],

(3.3)

其中 D̄t=DKL(p̄t∥q̄t)。此外，Adist t 在 detached 学生分布下是中心化的：Ea∼p̄t[Adist t(a)]=0。

证明见附录 A.1 (https://arxiv.org/html/2606.04036#A1.SS1)。

命题 3.1 (https://arxiv.org/html/2606.04036#S3.Thmtheorem1) 是一个梯度恒等式，而非实现层面的改变：SDPG 最小化公式 (3.2 (https://arxiv.org/html/2606.04036#S3.E2)) 中的显式全词表 KL，因为这能带来更精确的梯度估计。

在采样序列上的总蒸馏损失为：

LOPD(θ)=E(x,c)∼D,{yi}i=1G∼πθ(·∣x)[1/∑i=1G|yi| · ∑i=1G∑t=1|yi| li,tOPD(θ)].

(3.4)

### 3.4 SDPG 的在线策略非归一化 KL

现在我们聚焦于公式 (3.1 (https://arxiv.org/html/2606.04036#S3.E1)) 中的 KL 正则化项 LK。一般前向和反向 KL 正则化的推导详见附录 A.2 (https://arxiv.org/html/2606.04036#A1.SS2)。由该推导可知，πref=πθ 不足以最小化前向和反向 KL 的代理损失，这是因为（归一化的）前向和反向 KL 正则化本身存在固有偏差。为解决这一不匹配问题，我们采用第 2.3 节 (https://arxiv.org/html/2606.04036#S2.SS3) 中引入的非归一化 KL 项。

为简洁起见，令 JR&D=Jout+β(k)JOPD，LR&D=Lout+β(k)LOPD 分别表示训练步骤 k 时基于奖励与蒸馏项的目标函数和损失函数。具体而言，使用非归一化前向 KL 正则化的目标如下：

JSDPG-UFKL(θ)=JR&D(θ)−αUKL(πref∥πθ),

其中 JOPD 是隐含在公式 (3.4 (https://arxiv.org/html/2606.04036#S3.E4)) 中的在线策略蒸馏目标。梯度以 πθ 上的期望表示，令 wT(x)=πθ(x)/πteacher(x)，wR=πθ(x)/πref(x)，πteacher(x)=πθ(x,c)，则：

∇θJSDPG-UFKL(θ)=∇θJR&D(θ)−αEx∼πθ[(1−wR(x)−1)∇θlog πθ(x)].

用于梯度下降最小化的对应可微代理损失项为（忽略前缀 y0]）：

mi=**1**[A(i)out>0]，L+OPD(θ)=E(x,c)∼D,{yi}i=1G∼πθ(·∣x)[1/∑i=1G|yi| · ∑i=1G∑t=1|yi| mi·li,tOPD].

(3.5)

该方法仅对组内被验证器认可的轨迹使用全词表 OPD 信号。若一个组内所有奖励相同，则中心化结果优势和 OPD 门控均消失，从而避免对无信息组进行未经验证的蒸馏。在初始阶段，门控可能经常处于非激活状态，二值结果奖励占主导。因此，具有适度难度的训练数据集或课程学习\(Wang et al., 2021 (https://arxiv.org/html/2606.04036#bib.bib13); Lee et al., 2024 (https://arxiv.org/html/2606.04036#bib.bib14); Wen et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib15); Shi et al., 2025 (https://arxiv.org/html/2606.04036#bib.bib16)\) 有助于激活蒸馏信号。若对所有响应均有 mi=1，公式 (3.5 (https://arxiv.org/html/2606.04036#S3.E5)) 退化为对所有采样前缀的标准全词表 OPD。

#### 3.5.2 β 调度器

参见图 2：β 调度示意图。

pt 与特权分布 qt 之间的早期不对齐可能使 OPD 目标变得嘈杂。为防止特权蒸馏破坏探索，我们对 β 进行预热（warmup）。在结果策略开始找到正确轨迹后，OPD 项才逐渐发挥作用。此外，在理想化的特权信息模型下，将教师以学生在部署时无法获得的信息为条件进行蒸馏，可能留下不可消除的条件互信息差距，例如 I(Yt;C∣X,Y0<C 表示特权变量时\(Yang et al., 2026a (https://arxiv.org/html/2606.04036#bib.bib8)\)。在我们的框架下，这意味着特权 OPD 目标可能因推理时不可获得的信息而持续存在偏差。因此，为释放学生模型并鼓励探索，我们在训练末期对 β 进行衰减，在学生已内化其有用信息后逐步淡出蒸馏信号。

有效蒸馏系数遵循如图 2 (https://arxiv.org/html/2606.04036#S3.F2) 所示的预热-衰减调度：

β(k)=βbase × min(1, k/Twarm)_预热 × min(1, (T−k)/Tdecay)_衰减,

其中 Twarm 和 Tdecay 分别为预热和衰减步数，T 为总训练步数。若预热和衰减窗口重叠，最大系数可能低于 βbase。

## 4 实验

本节对所提出的 SDPG 算法进行实验评估，并在基于预训练 LLMs 的挑战性数学推理任务上与基线方法进行性能对比，包括 GRPO\(Shao et al., 2024 (https://arxiv.org/html/2606.04036#bib.bib4)\) 和 RLSD\(Yang et al., 2026a (https://arxiv.org/html/2606.04036#bib.bib8)\)。

参见图　参见图

(a) AIME24　(b) A