PaperlayerView exact revision on arXiv

Source-linked AI summary

Do We Really Need KL Divergence for On-Policy Distillation of Large Language Models?

Wenze Lin, Jiyuan Long, Jiale Zhao, Shenzhi Wang, Xitai Jiang, Ce Luo, Rui Lan, Qianli Ma, Fukang Wen, Hui Wu, Liyuan Chen, Shuoling Liu, Jiangpeng Yan, Gao Huang

arXiv:2609.33791v1cs.LGcs.AI

TL;DR

OPD inherits reverse KL as its standard loss, but this paper asks whether the loss is necessary for effective teacher guidance. It tests directional rewards and critical high-disagreement tokens, then applies the findings in C-MOPD, which consistently outperforms MOPD on math and code benchmarks.

  • Problem

    OPD uses reverse KL to align the student with the teacher on student-generated rollouts, raising the question of whether KL is necessary for effective teacher guidance.

    Source: §1 · p. 1, §7 · p. 10

  • Method

    The paper replaces reverse KL with binary directional rewards, identifies high-disagreement tokens as critical, and applies the findings in C-MOPD by guiding each update with multiple teachers.

    Source: §1 · p. 2, §6.1 · pp. 8–9

  • Results

    C-MOPD consistently outperforms MOPD on math and code benchmarks, while binary directional rewards achieve comparable or better results than reverse-KL OPD.

    Source: §1 · p. 2, §4.2.2 · p. 4, §7 · p. 10

  • Takeaways & Limitations

    The findings suggest that OPD can rely on directional updates for a small subset of high-disagreement tokens rather than requiring KL-based distribution matching.

    Source: §1 · p. 2, §7 · p. 10

Abstract

from arXiv · show

Since the advent of knowledge distillation, KL divergence has been the standard loss in distillation. Recently, on-policy distillation (OPD) has emerged as an efficient post-training paradigm for LLMs. As a distillation method, OPD naturally inherits KL divergence as its standard loss. However, in this work, we find that KL divergence may not be necessary for OPD. We show that simply preserving the update direction is sufficient for effective OPD. As long as the update direction is toward the teacher, OPD works. More precisely, it is not the direction of every token, but the direction of a small subset of tokens where the teacher and student disagree strongly. We first show that simply assigning a reward of (+1) to tokens where the teacher probability is higher than the student probability and (-1) where it is lower, which merely encourages updates toward the teacher, reproduces almost the same training mode as OPD with reverse KL. We further show that only the direction of a small subset of tokens with large teacher-student disagreement is critical, and training works as long as their update direction is toward the teacher, even if other tokens are pulled away from the teacher. And as an application of these findings, we introduce Consensus Multi-Teacher On-Policy Distillation (C-MOPD) to improve Multi-Teacher On-Policy Distillation (MOPD). Unlike MOPD, which routes each sample to a single teacher and may cause capability conflicts across domains, C-MOPD lets every sample be supervised by all teachers. Experiments show that C-MOPD consistently outperforms MOPD on both math and code benchmarks. Our code is available at https://github.com/LeapLabTHU/KL-Free-OPD.

1. Introduction

Although OPD inherits reverse KL as its standard loss, the paper argues that teacher-directed updates can suffice and that only strongly disagreeing tokens need the correct direction. It applies this finding in C-MOPD, where every update is guided by all teachers or cautiously follows the domain-specific teacher.

1. Introduction

  • A binary reward of +1 or −1 based on whether teacher probability exceeds student probability reproduces almost the same training mode as reverse-KL OPD, with comparable or better code and math results.The broader finding is that direction toward the teacher, not update magnitude, is sufficient for effective OPD.Source: §1 · pp. 1–2
  • Only a small subset of high-disagreement tokens needs updates toward the teacher; in some cases, this subset is less than 2% of all tokens.The paper reports that training fails when these critical tokens are instead updated away from the teacher, while low-disagreement tokens can be pulled away without severely affecting OPD.Source: §1 · p. 2
  • C-MOPD consistently outperforms MOPD, and it guides updates using all teachers rather than routing each sample to one teacher.When teachers conflict, it follows the domain-specific teacher only if its direction does not strongly conflict with the others; otherwise it assigns zero reward.Source: §1 · p. 2

2. Related Work

The related work situates OPD as a post-training approach that distills teacher outputs on student-generated rollouts, with reverse KL as its de facto loss. It also describes MOPD's single-teacher routing and introduces C-MOPD's all-teacher guidance.

On-Policy Distillation with Reverse KL

  • OPD aligns the student with the teacher on the student's own rollouts, unlike classical distillation on a fixed dataset, and reverse KL is its de facto loss.The cited work describes reverse KL as mode-seeking and as preserving the on-policy property (Zhao et al., 2026; Shao et al., 2026; Song & Zheng, 2026).Source: §2 · p. 2

Multi-Teacher On-Policy Distillation

  • MOPD routes each sample to one domain-specific teacher, whereas C-MOPD lets every sample receive guidance from all teachers to avoid losing one domain's capability during updates on another.The related-work passage describes the student as gradually acquiring teachers' expertise as it sees more routed samples.Source: §2 · p. 3

3. Preliminary

OPD transfers teacher capability by aligning the student with the teacher on student-generated rollouts. Its reverse-KL objective is differentiated into a policy-gradient-like update with dense token-level reward.

On-Policy Distillation with Reverse KL

  • For each query, the student samples an output sequence from its own policy and minimizes reverse KL against the teacher at every token position.Taking the expectation over student-sampled trajectories keeps training on-policy and provides dense token-level supervision.Source: §3 · p. 3
  • Differentiating the objective yields a policy-gradient-like expression in which each token receives a dense reward equal to the log-ratio.The supplied passage identifies the log-ratio as the token reward.Source: §3 · p. 3

Multi-Teacher On-Policy Distillation

  • MOPD extends OPD to multiple domain-specific teachers by routing each query to one teacher and minimizing reverse KL on the student's own rollouts.A routing function selects the teacher corresponding to the query's domain.Source: §3 · p. 3

4. Teacher-Directional Updates Suffice for On-Policy Distillation

The paper replaces reverse KL with a binary teacher-directional reward and evaluates BinaryOPD against OPD. It reports comparable or slightly better benchmark performance and proposes repeated state visitation as an explanation for why directional feedback can suffice.

4.1. Method

  • BinaryOPD assigns +1 when teacher probability exceeds student probability and −1 when it is lower, making the reward encode update direction without magnitude.Equal-probability tokens receive reward 0 in the implementation; the student is optimized with a policy-gradient objective.Source: §4.1 · pp. 3–4

4.2.1. Experimental Setup

  • The comparison trains OPD and BinaryOPD on the same dataset with identical hyperparameters, rollout budgets, and optimization settings across math and code tasks.The setup includes five student-teacher pairs for math and two for code, and evaluates on multiple math and code benchmarks.Source: §4.2.1 · p. 4, Table 1 · p. 4

4.2.2. Experimental Results

  • 65.1 versus 64.1 on the Qwen3-4B-Non-Thinking math pair, 22.2 versus 21.2 on the Qwen3-1.7B-Base math pair, and 55.2 versus 55.4 on the code benchmark: BinaryOPD is comparable to or slightly better than OPD overall.The training dynamics also show the same trend in training accuracy, entropy, and response length.Source: §4.2.2 · pp. 4–5, Table 2 · p. 5, Table 3 · p. 6

4.3. Why can BinaryOPD work? A repeated-feedback view.

  • Repeated visits to regions already encountered let OPD use fresh directional feedback after each student change; the direction reverses once the student crosses the teacher probability.The paper argues that this closed-loop feedback can determine the needed correction iteratively without encoding its magnitude in an individual update.Source: §4.3 · p. 5

5. Only the Direction of High-Disagreement Tokens Is Critical for OPD

The section tests which token updates matter for OPD and finds that the direction of high-disagreement tokens is critical: training works when they move toward the teacher, even if low-disagreement tokens do not.

5.1. Method

  • Group A has high disagreement with the teacher preferred, Group C has high disagreement with the student preferred, and Group B contains the vast majority of low-disagreement tokens.The experiments use a log-ratio threshold ϵ = 0.8 and assign uniform +1 or −1 rewards to selected tokens, with unselected tokens receiving 0.Source: §5.1 · p. 6

5.2. Results

  • In the illustrated Qwen3 pair, Group B accounts for over 90% of tokens, while Groups A and C each account for only a small fraction.The section also reports experiments at ϵ = 0.2 and 0.5 and with another student-teacher pair.Source: §5.2 · pp. 6, 8
  • With positive rewards, Group A alone (<1.5% of tokens) and Groups A+B work, but adding Group C causes failure.Group C accounts for less than 5% of tokens; Group B alone improves training but underperforms GRPO.Source: §5.2 · pp. 6–7
  • With negative rewards, Group C alone works, and Groups C+B work and surpass GRPO; adding Group A causes failure.Group B alone fails, and Group A accounts for less than 5% of tokens in the failing C+B+A setting.Source: §5.2 · p. 7
  • High-disagreement tokens must be updated toward the teacher: adding a small set of tokens updated away from the teacher can make training fail.The negative-reward experiment shows the symmetric pattern; low-disagreement tokens are not critical and can be removed or included with an incorrect direction while training proceeds.Source: §5.2 · p. 8, Figure 2 · p. 8

6. Consensus Multi-Teacher On-Policy Distillation

C-MOPD supervises every sample with multiple teachers while handling teacher conflicts through reward rules. With ϵ = 0.8, it consistently outperforms MOPD across almost all math and code benchmarks.

6.1. Method

  • C-MOPD follows the finding that samples need not be routed to a single teacher: every token can be guided by all teachers while conflicts are handled without harming other teachers’ associated capabilities.The motivation is that updating on one domain’s samples may degrade capabilities in another domain.Source: §6.1 · pp. 8–9

6.2.2. Experimental Results

  • C-MOPD with ϵ = 0.8 consistently outperforms MOPD across almost all math and code benchmarks.The reported results identify ϵ = 0.8 as the best overall setting.Source: Table 5 · p. 9, §6.2.2 · pp. 9–10

6.1. Method

  • C-MOPD guides every token with all teachers, updating toward the domain-specific teacher during conflicts when that update does not hurt other teachers’ associated capabilities.The reward rules specify updates toward all teachers when their preferences agree, and toward the domain-specific teacher in specified mild-conflict cases.Source: §6.1 · p. 9

6.2.1. Experimental Setup

  • The evaluation uses a Qwen3-4B-Non-Thinking student, math and code teachers, and a mixture of math and code problems; ϵ = 0.8 is the default.The mixture includes DeepMath, CodeContests, TACO, APPS, and Codeforces problems; the study also compares the conservative ϵ = 0.0 setting.Source: §6.2.1 · p. 9

6.2.2. Experimental Results

  • The ϵ = 0.0 setting is overly conservative: its Zero Reward Ratio eventually exceeds 90%, and its performance only matches MOPD.As training proceeds, more tokens conflict across teachers, so the model refuses to update on most of them.Source: §6.2.2 · p. 9

7. Conclusion

The conclusion revisits reverse KL’s role in OPD, reporting that binary directional rewards can achieve comparable or better results and that only a small subset of high-disagreement tokens is critical. It also presents C-MOPD, which consistently outperforms MOPD on math and code benchmarks.

7. Conclusion

  • A binary directional reward achieves comparable or better results than reverse KL for OPD, while only a small subset of high-disagreement tokens is critical.The conclusion also reports that verifier signal has little effect on OPD.Source: §7 · p. 10

A. State Exploration Is Highly Redundant in OPD Training

OPD visits token-distinct contexts that are often close in a fixed teacher representation space, with little evidence of new regions emerging over training. This redundancy helps explain why BinaryOPD can work without teacher-student discrepancy magnitudes.

A. State Exploration Is Highly Redundant in OPD Training

  • A 29× larger state library raises cos ≥0.99 coverage only from 10.4% to 18.4%, while mean similarity to the closest earlier state rises from 0.823 to 0.891.The library expands from 2,048 to 59,392 states; the figure searches only earlier states from different prompts.Source: Figure 5 · p. 14
  • Different-prompt states have mean nearest-neighbor cosine similarity 0.893; 29.5% reach at least 0.95 and 18.5% reach at least 0.99.The different-prompt mean is nearly unchanged from the unrestricted mean, so redundancy is not just repeated prefixes of the same prompt.Source: A. State Exploration Is Highly… · p. 14, Figure 7 · p. 16
  • Cross-prompt similarity across training steps stays within [0.138, 0.149], with no clear separation of later states into new regions.The nearest-neighbor share from the same training step is about the 1/30 chance level.Source: A. State Exploration Is Highly… · pp. 14–15, Figure 6 · p. 15
  • Repeated visits to the same or nearby states mean the student is corrected many times in the same direction, helping BinaryOPD work without magnitude information.The training contexts remain token-level distinct despite their representational proximity.Source: A. State Exploration Is Highly… · pp. 13, 15, 16

B. Whether the OPD Update Direction Is Aligned with the Verifier Is Not Important

The reversed verifier signal does not hurt BinaryOPD, and aligning updates with the verifier yields no consistent improvement. The experiment reverses the usual relationship between answer correctness and token rewards.

B. Whether the OPD Update Direction Is Aligned with the Verifier Is Not Important

  • The reversed mask gives correct-answer trajectories non-positive rewards and incorrect-answer trajectories non-negative rewards, opposite to usual verifier guidance.It assigns −1 to selected tokens in correct trajectories and +1 to selected tokens in incorrect trajectories; other tokens receive 0.Source: B. Whether the OPD Update Direction Is… · p. 16
  • Verifier-Conflicted performs on par with or slightly better than OPD and BinaryOPD across student-teacher pairs.The reported training dynamics are also nearly identical across methods.Source: B. Whether the OPD Update Direction Is… · p. 17
  • 55.1 average score for Verifier Conflicted versus 54.4 for Verifier Aligned on the DeepSeek-Distill-Qwen-1.5B pair; Verifier Conflicted scores 29.1 on LiveCodeBench v6.On the Qwen3-4B-Non-Thinking math pair, both variants reach around 64.Source: B. Whether the OPD Update Direction Is… · p. 17
  • Verifier-Aligned shows no consistent improvement over OPD or BinaryOPD, while the reversed signal leaves BinaryOPD performance intact.The findings indicate that verifier alignment has little influence on the training pattern or performance in these experiments.Source: B. Whether the OPD Update Direction Is… · pp. 17–18

C. Experimental Details

The experiments use domain-specialized teacher models and a specified set of student models and datasets, with training run in Verl using the listed hyperparameters.

C. Experimental Details

  • The experiments use JustRL-1.5B, several Qwen3-4B variants, GT-Llama3.2-3B-MATH, and Nemotron-Research-Reasoning-1.5B as models.The teachers are domain-specialized models obtained by reinforcement learning on their respective domains.Source: C. Experimental Details · p. 18
  • DeepMath follows Yang et al. (2026b) and selects 57K samples with difficulty level at least 6.Qwen3-4B-Base-RL is trained from Qwen3-4B-Base with GRPO on DAPO-MATH-17K for three epochs.Source: C. Experimental Details · p. 18
  • Training uses the Verl framework and the hyperparameters listed in Table 9.The supplied passage does not state the individual hyperparameter values.Source: C. Experimental Details · p. 18, Table 9 · p. 18

D. Training Dynamics

Training-dynamics curves compare methods across experiments, and a token visualization marks large teacher-student probability disagreements. The experiments were conducted on several NVIDIA GPU types.

D. Training Dynamics

  • Across methods, training accuracy, entropy, response length, and gradient norm curves are nearly identical; incorrectly rewarding high-disagreement tokens fails to converge or clearly degrades.The comparisons include OPD, BinaryOPD, Verifier-Aligned, and Verifier-Conflicted, alongside positive- and negative-reward settings.Source: D. Training Dynamics · p. 18, Figure 8 · p. 19

E. Visualization of High-Disagreement Tokens

F. Hardware Setup

Related papers

Loading 2609.33791v1…