OC-OPD: Reconciling Forward and Reverse KL Distillation with a Four-Observation Matrix

Two KL Divergences, Complementary Supervision Signals

In on-policy distillation, the student model first generates prefixes, and the teacher model then scores those prefixes. In this process, reverse KL and forward KL each play different roles: reverse KL favors refining the ranking of tokens the student is already active on, making the model more "confident" about the choices it already prefers; forward KL is responsible for allocating probability mass to candidates that the teacher supports but the student has not yet attended to sufficiently, serving to explore and fill gaps.

Using either divergence alone will bias distillation toward a certain behavioral pattern—either overly conservative or introducing too much noise. The starting point of OC-OPD is precisely to acknowledge the complementarity of these two supervision signals and attempt to reconcile them at the gradient level. The paper comes from an rss source, and the code has been made public on GitHub.

KL divergence (Kullback-Leibler Divergence) measures the difference between two probability distributions and is defined as (D_{\text{KL}}(P|Q)=\sum_x P(x)\log\frac{P(x)}{Q(x)}). It is asymmetric, hence the distinction between forward KL and reverse KL. In distillation, taking the teacher distribution (p_t) and the student distribution (p_s) as an example, forward KL is (D_{\text{KL}}(p_t|p_s)), which tends to make the student cover the probability mass of all possible teacher outputs (mass-covering) but may spread probability onto low-quality candidates; reverse KL is (D_{\text{KL}}(p_s|p_t)), which tends to make the student focus on the teacher's high-probability regions (mode-seeking) but may ignore other modes supported by the teacher. In on-policy distillation, after the student generates a prefix, the teacher scores the next-token distribution under that prefix, and the two KLs produce different gradient behaviors here: reverse KL reinforces the token ranking the student already prefers, while forward KL raises the probability of candidates that the teacher supports but the student underestimates.

Four-Observation Matrix: Treating Student and Teacher Simultaneously as Candidate Sources and Evaluators

The core construction of OC-OPD is to build a "four-observation matrix" at the prefixes generated by the student. Specifically, the student and teacher are each treated as candidate sources and evaluators, and pairwise combinations form four observations: on the diagonal are self-observations (student evaluating student, teacher evaluating teacher), used to measure their respective uncertainty; the off-diagonal cross-observations correspond respectively to the two objectives of reverse KL and forward KL.

The cleverness of this design is that it does not treat the two KLs as two independent losses to be simply weighted, but rather characterizes them within a unified observational framework. The paper further establishes an exact identity: coupling the Fisher-inverse metric of the logit gradients of the two KLs is equivalent to the sum of these contrastive observations. As a result, two different positive distributions can obtain a common local descent direction.

Fisher Information Matrix is the covariance of the log-likelihood gradient in the parameter space of a probabilistic model, characterizing local curvature or parameter sensitivity. In natural gradient methods, the inverse of the Fisher information matrix is used as a metric matrix to scale the gradient, so that the update direction descends faster in distribution space rather than Euclidean parameter space. OC-OPD's identity expresses the coupling of the two KL gradients as "Fisher-inverse metric coupling of logit gradients," meaning it uses the Fisher inverse to reconcile gradients over the student and teacher distributions, enabling different positive distributions (teacher distribution and student distribution) to find a common local descent direction. This construction connects the gradient update of KL distillation with natural gradients in information geometry, providing a theoretical basis for closed-form bargaining.

From Observation to Update: Preference, Bargaining, and Scaling

After obtaining the observations, OC-OPD completes the update in three steps. The observations are used to set preferences, determining which type of signal should be trusted more in the current state; then closed-form bargaining reconciles the gradients of forward and reverse KL, avoiding manual tuning-style weight search; finally, the distribution discrepancy is used to scale the update magnitude, allowing the supervision strength at different prefixes to adapt.

The value of this pipeline lies in interpretability: each step has a clear statistical meaning rather than being a black-box loss combination. For research scenarios that require stable training, this structured reconciliation approach is more controllable than grid-searching KL weights.

Experimental Performance under Three Student Settings

The paper evaluates OC-OPD under three heterogeneous student settings: non-thinking generation, thinking generation, and base model initialization. Results show that OC-OPD achieves clear gains on mathematical reasoning tasks and can effectively transfer to out-of-distribution scientific question answering and code generation tasks.

This consistency across settings is noteworthy—a common pain point of on-policy distillation is that hyperparameters are sensitive to model type, and changing the student initialization requires re-tuning the weights. OC-OPD replaces some manual tuning with observation-driven preferences and closed-form bargaining, which may be one reason it works across different settings. However, the paper abstract does not disclose the specific magnitude of improvements or baseline comparison details, so the actual effect still needs to be judged in conjunction with the experimental section.

Brief Commentary

The theoretical contribution of OC-OPD lies in providing an exact identity for the coupling of the two KL gradients, advancing "why the two divergences can work together" from intuition to a derivable level. Engineering-wise, it provides an on-policy distillation scheme that does not require extensive weight search. For teams working on model compression and transfer of reasoning capabilities, this is an idea worth reproducing.

Code address: https://github.com/RuoranXu/OC-OPD

One-line Summary

OC-OPD reconciles forward and reverse KL divergence through a four-observation matrix, achieving on-policy distillation without weight search.

Paragraph Summary

OC-OPD proposes an on-policy distillation method based on a four-observation matrix, treating the student and teacher simultaneously as candidate sources and evaluators to construct four observations: diagonal self-observations measure uncertainty, while off-diagonal cross-observations correspond respectively to the reverse KL and forward KL objectives. The paper establishes an exact identity that equates the Fisher-inverse metric coupling of the logit gradients of the two KLs to the sum of contrastive observations, enabling different positive distributions to share a local descent direction. The update pipeline has three steps: observations set preferences, closed-form bargaining reconciles gradients, and distribution discrepancy scales the update magnitude. Under three student settings—non-thinking generation, thinking generation, and base model initialization—clear gains are achieved on mathematical reasoning tasks, with transfer to out-of-distribution scientific question answering and code generation. The method requires no extensive weight search and provides an interpretable and stable distillation scheme.