[KongchangAI]
· 3 min read· 1,649 words

D²-OPD: A New Method for Online Distillation that Decouples Direction and Distance

D²-OPD: A New Method for Online Distillation that Decouples Direction and Distance

D²-OPD significantly improves on-policy distillation performance across multiple settings by decoupling the direction and distance of extrapolative distillation.

D²-OPD addresses the coupling of extrapolation direction and distance in on-policy distillation by proposing an explicit decoupling scheme. It introduces a per-prefix hard KL constraint to control extrapolation distance, ensuring that the target distribution does not deviate too far from the teacher model; it also adaptively combines the teacher's post-training shift with the initial student-teacher gap to determine the extrapolation direction, avoiding dominance by a single signal. In the optimization stage, forward KL divergence is used to provide learning signals for low-probability tokens. Across three settings on mathematical reasoning benchmarks, D²-OPD achieved the highest scores, surpassing the strongest baseline by up to 9.45 points. Ablation experiments verified that each component contributes positively. The article also discusses the method's advantages in engineering feasibility and notes that hyperparameter strategies such as the choice of the hard KL threshold require further research.

In the field of model distillation, On-Policy Distillation (OPD) has become an effective training paradigm by training the student model to match the teacher model's next-token distribution on sequence prefixes generated by the student itself. However, when researchers attempt to enhance the learning signal through extrapolation, a fundamental contradiction emerges: the direction and distance of extrapolation are coupled together, so that the teacher model's log-probability shift simultaneously determines both the direction and the magnitude of the target movement. The recently proposed D²-OPD method achieves significant improvements across multiple settings by explicitly decoupling these two factors.

The Coupling Dilemma of Extrapolative Distillation

In conventional on-policy distillation, the student model is trained to match the teacher model's next-token distribution on prefixes it generates itself. Extrapolative OPD further extends the teacher model's log-probability shift relative to a reference policy, thereby constructing a stronger learning objective. The intuition behind this approach is: if the teacher model is more confident than the reference policy in a certain direction, then extrapolating along this direction can help the student model learn more distinctive behavioral patterns.

The problem is that existing methods bind the direction and distance of extrapolation together. The larger the teacher model's shift, the farther the target moves. This coupling introduces two risks: on the one hand, when the teacher model's shift signal is weak, the extrapolation distance is insufficient and the learning signal is not strong enough; on the other hand, when the shift signal is strong, the extrapolation may overshoot, causing the target to deviate from a reasonable range or even introducing noise. This "one-size-fits-all" extrapolation strategy lacks adaptivity across different prefixes and different tokens.

The core insight of D²-OPD comes precisely from reflecting on this coupling relationship: direction and distance should be controlled separately as two independent design dimensions.

D²-OPD method diagram

How D²-OPD Decouples Direction and Distance

The design of D²-OPD contains two key components, responsible for distance constraint and direction determination respectively.

Per-Prefix Hard KL Constraint for Distance Control

To limit how far the target distribution deviates from the teacher model, D²-OPD introduces a per-prefix hard KL constraint. This constraint sets an explicit upper bound for each prefix, ensuring that the extrapolated target does not drift too far from the teacher model. Unlike soft constraints, hard constraints provide more direct control—no matter how strong the teacher model's shift signal is, the KL divergence between the target distribution and the teacher distribution will not exceed a preset threshold.

The benefit of this design lies in stability. In on-policy distillation, the quality of prefixes generated by the student model varies considerably; without constraining each prefix individually, extrapolation on some prefixes may spiral out of control. The per-prefix hard KL constraint acts as a safety valve for each training sample, preventing target drift while guaranteeing the strength of the learning signal.

KL divergence measures the difference between two probability distributions and is often used as a loss function in knowledge distillation. Different directions have different properties: forward KL (teacher||student) imposes a larger penalty on regions where the student's coverage is insufficient, tending to make the student cover all modes of the teacher distribution (mass-covering); reverse KL (student||teacher) tends to make the student focus on the teacher's high-probability modes (mode-seeking) and may ignore low-probability tokens. D²-OPD explicitly uses forward KL to optimize the student, which is complementary to its hard KL constraint: the hard constraint bounds the target's deviation from the teacher, while forward KL ensures that the student does not miss tokens in the teacher distribution that have low probability but may be important.

Adaptive Direction Combination

The direction is determined by adaptively combining two signals: the teacher model's post-training shift and the initial student-teacher gap. The teacher's post-training shift reflects what the teacher model has learned relative to the reference policy, while the initial student-teacher gap reveals where the student model currently differs from the teacher model.

D²-OPD adaptively weights and combines these two signals to determine the direction of extrapolation. Intuitively, when the teacher's post-training shift provides a clear directional signal, the combination weight tends toward that signal; when the shift signal is ambiguous or inconsistent with the initial gap, the initial gap information serves as a correction. This combination strategy makes the direction selection more robust, avoiding the bias introduced by a single dominant signal.

The teacher's post-training shift refers to the difference in log-probabilities between the teacher model and some reference policy (typically the base model or an SFT model), reflecting the behavioral preferences the teacher has acquired through post-training (e.g., RLHF, rejection-sampling fine-tuning). The initial student-teacher gap measures the distribution difference between the student model and the teacher model on the same prefix at the start of training. The significance of combining the two is: the post-training shift provides the directional signal of "where to go," while the initial gap provides corrective information about "what the student currently lacks." When the teacher's shift is clear, it dominates; when the shift is ambiguous or conflicts with the initial gap, the initial gap helps stabilize the direction, thereby avoiding the bias introduced by a single signal.

Forward KL Optimization

After determining the target distribution, D²-OPD uses forward KL divergence to optimize the student model. A key advantage of forward KL is its ability to provide learning signals for low-probability tokens. Under reverse KL, the student model tends to ignore tokens with lower probability in the teacher distribution, whereas forward KL requires the student model to cover all modes of the teacher distribution, including those tokens that are not high-probability but may be important. This is particularly important for maintaining the diversity and generalization ability of the student model.

The core of On-Policy Distillation lies in the training data being generated by the student model itself rather than from a fixed offline dataset. At each time step, based on the sequence prefix it has generated, the student computes the teacher model's predicted distribution for the next token and uses it as a supervision signal to update its parameters. This approach mitigates the mismatch between training and inference distributions (exposure bias), because the student always learns on the state distribution it produces itself. But the cost is the need to frequently call the teacher model for forward propagation, which incurs higher computational overhead than offline distillation. D²-OPD continues this paradigm, and both its per-prefix constraint and adaptive direction mechanism are built on the prefixes generated in real time by the student.

Experimental Validation: Performance Across Three Settings

D²-OPD was evaluated under three different settings, each corresponding to a different training scenario.

The first setting is one where standard OPD performs well, and here D²-OPD needs to demonstrate that it does not lose performance due to the additional constraints. The second setting involves cases where OPD struggles under base-model initialization—a common difficult scenario in practice where the student model, initialized from the base model, has a large gap from the teacher model, and standard OPD easily falls into local optima. The third setting is evaluated when no teacher reference is available, requiring the method to possess a certain degree of bootstrapping capability.

On mathematical reasoning benchmarks, D²-OPD achieved the highest scores in all three settings, surpassing the strongest baseline by up to 9.45 points. This result covers a range of scenarios from "standard and usable" to "standard and difficult," indicating that the design of decoupling direction and distance has broad applicability. Ablation experiments further verified the contribution of each component: the hard KL constraint, adaptive direction combination, and forward KL optimization each have a positive effect on final performance.

Implications for the Design of Distillation Methods

The value of D²-OPD lies not only in its specific technical solution but also in the methodological question it reveals: in extrapolative distillation, coupling multiple factors together simplifies the design but limits the adaptability and robustness of the method. Decoupling direction and distance provides a more fine-grained control mechanism, allowing extrapolation behavior to adapt to the specific circumstances of each prefix.

From an engineering perspective, the computational overhead of the per-prefix hard KL constraint is manageable, and the adaptive direction combination does not require additional teacher forward propagation, making D²-OPD quite feasible in practical training. For scenarios that require transferring knowledge from a strong teacher model to a lightweight student model while maintaining stability, this decoupling approach is worth considering.

Of course, the current experiments are mainly focused on mathematical reasoning benchmarks, and the performance of D²-OPD on other task types (such as open-ended generation and multi-turn dialogue) remains to be verified. In addition, the strategy for selecting the threshold of the hard KL constraint is not discussed in detail in the original paper, which may be a key hyperparameter affecting the practical deployment of the method. Future research could further explore adaptive ways to determine the threshold, as well as dynamic adjustment strategies for the direction combination weights at different training stages.

One-line Summary

D²-OPD significantly improves on-policy distillation performance across multiple settings by decoupling the direction and distance of extrapolative distillation.

Paragraph Summary

D²-OPD addresses the coupling of extrapolation direction and distance in on-policy distillation by proposing an explicit decoupling scheme. It introduces a per-prefix hard KL constraint to control extrapolation distance, ensuring that the target distribution does not deviate too far from the teacher model; at the same time, it adaptively combines the teacher's post-training shift with the initial student-teacher gap to determine the extrapolation direction, avoiding dominance by a single signal. In the optimization stage, forward KL divergence is used to provide learning signals for low-probability tokens. Across three settings on mathematical reasoning benchmarks, D²-OPD achieved the highest scores, surpassing the strongest baseline by up to 9.45 points. Ablation experiments verified that each component contributes positively. The article also discusses the method's advantages in engineering feasibility and notes that hyperparameter strategies such as the choice of the hard KL threshold require further research.

#D²-OPD#online policy distillation#extrapolative distillation#hard KL constraint#forward KL divergence#knowledge distillation#mathematical reasoning
Share:

Related articles