DISTILLED REINFORCEMENT LEARNING FOR LLM POST-TRAINING

Research Question and Approach

Large language model (LLM) post-training typically relies on either reinforcement learning (RL), which uses coarse-grained outcome rewards, or on-policy distillation (OPD), which forces the student to match the teacher's probability distribution. The authors argue that RL suffers from poor credit assignment, while OPD suffers from a dilemma: it either provides little new knowledge when the teacher is similar to the student or produces noisy, ineffective guidance when the teacher is different. To solve this, the authors propose Distilled RL, a framework that incorporates teacher knowledge directly into the RL gradient. Instead of treating the teacher as an unconditional target, Distilled RL uses the teacher's preferences to reweight the RL learning signal at the token level, ensuring that teacher guidance is only applied to beneficial trajectories.

Key Components of Distilled RL

Distilled RL introduces three mechanisms to make teacher guidance selective and stable:

• Reverse Importance Sampling: The model measures the relative preference of the teacher for each token generated by the student, using a clipped ratio to prevent unstable gradient updates.
• Negative Sample Reset: To avoid reinforcing poor behaviors, the framework only applies teacher-based reweighting to responses with positive advantages, while resetting weights for negative samples to one.
• Sequence-level Geometric Normalization: The model normalizes token-level ratios within each response to ensure the geometric mean equals one, preventing the teacher from globally amplifying or suppressing entire sequences and focusing instead on relative token preferences.

Standard RL methods often struggle to introduce knowledge that is absent from the student's initial policy, and traditional distillation often fails when the teacher and student belong to different model families. Distilled RL bridges this gap by enabling the student to learn from a teacher's reasoning patterns without being forced into rigid imitation. Experimental results demonstrate that this method consistently outperforms both standard RL and OPD across various benchmarks, showing significant gains in both within-family and cross-family distillation scenarios.