PowerOPD：用有界幂变换稳定 On-Policy Distillation

On-Policy Distillation（OPD）通过让学生模型在自己的生成轨迹上向教师学习，有效缓解了传统离线蒸馏的曝光偏差问题。然而，标准 OPD 使用 log-ratio 作为奖励信号，这一看似自然的构造却隐藏着严重的训练不稳定性——梯度方差爆炸。

来自 Sea AI Lab 和西湖大学的研究团队提出 PowerOPD，通过 Box-Cox 幂变换将有界性内建到奖励设计中，从根本上解决了这一病理问题。在代码生成任务上，该方法实现了相比 vanilla OPD 高达 +6.37 Avg@8 的性能提升，同时减少了 59.2% 的训练时间和 23.1% 的 GPU 内存占用。

---

问题根源：Log-Ratio 的无界性陷阱

标准 OPD 的奖励设计

On-Policy Distillation 的核心思想是用学生采样的 token 来估计 Reverse KL 散度：

\[\mathcal{L}_{\text{OPD}} = -\mathbb{E}_{y \sim \pi_\theta} \left[ \sum_t \log \frac{\pi_T(y_t | y_{<t})}{\pi_\theta(y_t | y_{<t})} \right]\]

其中 \(\pi_T\) 是教师策略，\(\pi_\theta\) 是学生策略。这个 log-ratio 奖励 \(\log \frac{\pi_T}{\pi_\theta}\) 被广泛用于 MiniLLM、GKD 等经典方法中。

训练病理的三重表现

论文揭示了这种设计的根本性缺陷：

问题	表现	根因
样本效率低下	需要大量样本才能收敛	高方差导致梯度信噪比低
生成动态不稳定	训练过程中 perplexity 剧烈波动	奖励无界导致策略更新幅度不可控
性能差距	相比全词汇表 OPD 存在明显差距	极端梯度值主导了优化方向

关键洞察：早期位置的梯度集中

作者通过梯度可视化发现了一个关键现象：高方差梯度集中在序列的早期位置。这是因为：

1. 早期 token 的生成不确定性更高，\(\pi_\theta(y_t | y_{<t})\) 可能非常小
2. 当学生概率接近零而教师概率非零时，\(\log \frac{\pi_T}{\pi_\theta} \to +\infty\)
3. 这种极端值导致梯度爆炸，破坏了训练的稳定性

核心问题：Log-ratio 在构造上就是无界的，这导致了极端高方差的梯度估计。

---

PowerOPD：内建有界性的幂变换

Box-Cox 幂变换

PowerOPD 的核心创新是用 Box-Cox 幂变换 替代 log-ratio：

\[r_\alpha(x) = \begin{cases} \frac{x^\alpha - 1}{\alpha} & \alpha > 0 \\ \log x & \alpha \to 0 \end{cases}\]

其中 \(x = \frac{\pi_T(y_t | y_{<t})}{\pi_\theta(y_t | y_{<t})}\) 是概率比值。

关键性质：有界性与符号一致性

对于 \(\alpha > 0\)，该变换具有以下优良性质：

1. 有界性（Boundedness）

\[-\frac{1}{\alpha} \leq r_\alpha(x) < \infty\]

当学生概率 \(\pi_\theta \to 0\) 时，变换后的奖励收敛到有限值 \(-\frac{1}{\alpha}\)，而非负无穷。

2. 符号一致性（Sign-Consistency）

\[\text{sign}(r_\alpha(x)) = \text{sign}(\log x) = \text{sign}(x - 1)\]

这保证了奖励的激励方向与原始 OPD 一致：当教师概率高于学生时给予正奖励，反之给予负奖励。

梯度方差分析

论文给出了不同 \(\alpha\) 值下梯度方差的理论分析。定义变换后的策略梯度：

\[\nabla_\theta J = \mathbb{E}_{y \sim \pi_\theta} \left[ \sum_t r_\alpha\left(\frac{\pi_T}{\pi_\theta}\right) \nabla_\theta \log \pi_\theta(y_t | y_{<t}) \right]\]

方差上界：

\[\text{Var}[\nabla_\theta J] \leq C \cdot \mathbb{E}\left[ r_\alpha^2\left(\frac{\pi_T}{\pi_\theta}\right) \cdot \|\nabla_\theta \log \pi_\theta\|^2 \right]\]

由于 \(r_\alpha\) 的有界性，当 \(\alpha > 0\) 时，方差上界远低于标准 OPD（\(\alpha \to 0\) 的极限情况）。

超参数 \(\alpha\) 的选择

\(\alpha\) 控制有界性的强度：

• \(\alpha \to 0\)：退化为标准 log-ratio，无界但信息丰富
• \(\alpha = 1\)：线性变换 \(r_1(x) = x - 1\)，严格有界但可能丢失细粒度信号
• 中间值（如 \(\alpha = 0.5\)）：在稳定性和信息量之间取得平衡

论文通过实验发现 \(\alpha \in [0.3, 0.7]\) 是较好的选择范围。

---

算法实现

完整训练流程

# PowerOPD 伪代码
for batch in dataloader:
    # 1. 学生生成轨迹
    y ~ π_θ(y | x)
    
    # 2. 教师评估概率
    π_T_probs = teacher(y, x)
    π_θ_probs = student(y, x)
    
    # 3. 计算功率变换奖励
    ratio = π_T_probs / (π_θ_probs + ε)
    reward = (ratio ** alpha - 1) / alpha  # Box-Cox 变换
    
    # 4. 策略梯度更新
    loss = -reward * log_π_θ
    loss.backward()
    optimizer.step()

与现有 OPD 方法的对比

方法	奖励函数	有界性	计算复杂度
MiniLLM	\(\log \frac{\pi_T}{\pi_\theta}\)	❌ 无界	\(O(T \cdot V)\)
GKD	\(\log \frac{\pi_T}{\pi_\theta}\)	❌ 无界	\(O(T)\)
PowerOPD (α=0)	\(\log \frac{\pi_T}{\pi_\theta}\)	❌ 无界	\(O(T)\)
PowerOPD (α>0)	\(\frac{(\frac{\pi_T}{\pi_\theta})^\alpha - 1}{\alpha}\)	✅ 有界	\(O(T)\)

---

实验结果

主实验：代码生成任务

在 LiveCodeBench 和 HumanEval+ 等代码生成基准上，PowerOPD 展现出显著优势：

方法	Avg@8	Pass@8	相对提升
Vanilla OPD	42.15	38.24	-
PowerOPD (α=0.5)	48.52	43.95	+6.37 / +5.71

效率分析

训练速度： - PowerOPD 收敛所需迭代次数减少 35% - 总训练时间减少 59.2%

内存占用： - 峰值 GPU 内存减少 23.1% - 主要来自更稳定的梯度，可以使用更大的 batch size

消融实验

不同 α 值的影响：

α	Avg@8	训练稳定性	备注
0.0 (标准 OPD)	42.15	❌ 不稳定	baseline
0.3	46.82	✅ 稳定	较好
0.5	48.52	✅ 稳定	最佳
0.7	47.91	✅ 稳定	较好
1.0	45.23	✅ 很稳定	过于保守

关键发现： - 过小的 \(\alpha\)（接近0）保留了更多信息量但牺牲了稳定性 - 过大的 \(\alpha\)（接近1）过于稳定但丢失了教师信号的细粒度差异 - \(\alpha = 0.5\)（平方根变换）在实验中表现最佳

序列位置分析

论文验证了 PowerOPD 对早期位置梯度爆炸的缓解效果：

位置	Vanilla OPD 梯度方差	PowerOPD 梯度方差	方差降低
1-10 tokens	12.34	2.18	82.3%
11-50 tokens	3.56	1.42	60.1%
51+ tokens	1.23	0.89	27.6%

早期位置的方差降低最为显著，这与理论分析一致。

---

理论洞察

与 f-散度的联系

PowerOPD 的奖励函数可以看作是一种广散的 f-散度。定义：

\[D_\alpha(p \| q) = \mathbb{E}_{y \sim q} \left[ r_\alpha\left(\frac{p(y)}{q(y)}\right) \right]\]

当 \(\alpha \to 0\) 时，这退化为标准 KL 散度；当 \(\alpha = 1\) 时，这对应于线性散度。

重要性采样的视角

从重要性采样的角度看，标准 OPD 使用无界权重：

\[w(y) = \frac{\pi_T(y)}{\pi_\theta(y)}\]

而 PowerOPD 通过幂变换将其调整为：

\[w_\alpha(y) = \frac{w(y)^\alpha - 1}{\alpha}\]

这种”软化”的重要性权重有效降低了极端值的影响，同时保留了相对排序信息。

---

实践建议

何时使用 PowerOPD

✅ 推荐使用： - 学生模型与教师模型差距较大 - 训练初期不稳定，loss 剧烈波动 - 长序列生成任务（代码、数学推理） - 计算资源有限，需要更快收敛

⚠️ 谨慎使用： - 学生与教师已经非常接近（可能过度保守） - 需要保留完整概率分布信息的场景

超参数调优指南

1. 从 \(\alpha = 0.5\) 开始（平方根变换）
2. 如果训练仍然不稳定，尝试增大到 \(\alpha = 0.7\)
3. 如果感觉信号太弱，尝试减小到 \(\alpha = 0.3\)
4. 避免 \(\alpha > 1\)，会损失过多信息

与其他技术的兼容性

PowerOPD 可以与以下技术无缝结合：

• KL 约束：在损失中加入 \(\beta \cdot D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})\)
• 重要性采样校正：与 GKD 的离线-在线混合策略兼容
• 课程学习：逐步降低 \(\alpha\) 从保守到激进

---

开放问题与未来方向

1. 自适应 \(\alpha\)：能否根据训练动态自动调整 \(\alpha\)？
2. 位置相关的变换：早期位置用更大的 \(\alpha\)，后期用更小的？
3. 与其他方差缩减技术的结合：如基线函数、控制变量法等
4. 理论保证：能否给出收敛速率的显式上界？

---

小结

PowerOPD 的核心贡献在于揭示了标准 OPD 中 log-ratio 奖励的无界性缺陷，并通过简单的 Box-Cox 幂变换将有界性内建到奖励设计中。这一改动看似微小，却从根本上解决了梯度方差爆炸问题，带来了：

• 更稳定的训练动态
• 更高的样本效率
• 更好的最终性能
• 更低的计算开销

论文的实验表明，当 OPD 的训练不稳定问题时，PowerOPD 提供了一个即插即用的解决方案，值得在任何 OPD 实现中作为默认选项考虑。

一句话总结：用有界的幂变换替代无界的 log-ratio，让 On-Policy Distillation 真正稳定可用。

---

参考文献

• PowerOPD: Stabilizing On-Policy Distillation with Bounded Power Transformation (arXiv 2606.17199)
• MiniLLM: Knowledge Distillation of Large Language Models
• GKD: Generalized Knowledge Distillation
• On-Policy Distillation 系列综述

---

论文作者：Anhao Zhao, Junlong Tong, Yingqi Fan, Ping Nie, Wenjie Li, Xiaoyu Shen（Sea AI Lab & 西湖大学）