Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models — part 2/3

Cleaned-text code points [15963, 31396). Bibliography omitted here; source.md includes it.

##### On-policy sampling from the student.

Given a problem x x , the student generates an on-policy response y ^ = ( y ^ 1 , … , y ^ | y ^ | ) ∼ p S ( ⋅ ∣ x ) . \hat{y}=(\hat{y}_{1},\ldots,\hat{y}_{|\hat{y}|})\sim p_{S}(\cdot\mid x). Both policies then evaluate this student-generated trajectory. At each position n n , they induce next-token distributions over y n ∈ 𝒱 y_{n}\in\mathcal{V} conditioned on the same student prefix: p S ​ ( y n ∣ x , y ^ < n ) , p T ​ ( y n ∣ x , y ⋆ , y ^ < n ) , p_{S}\!\left(y_{n}\mid x,\hat{y}_{<n}\right),\qquad p_{T}\!\left(y_{n}\mid x,y^{\star},\hat{y}_{<n}\right), where y ^ < n ≜ ( y ^ 1 , … , y ^ n − 1 ) \hat{y}_{<n}\triangleq(\hat{y}_{1},\ldots,\hat{y}_{n-1}) .

##### Training objective: Full-vocabulary logit distillation.

We instantiate a full-vocabulary divergence objective that matches the teacher and student next-token distributions at each position. Given a student-generated sequence y ^ \hat{y} , define the trajectory-averaged, token-wise divergence D ( p T ∥ p S ) ( y ^ ∣ x ) ≜ 1 | y ^ | ∑ n = 1 | y ^ | D ( p T ( ⋅ ∣ x , y ⋆ , y ^ < n ) ∥ p S ( ⋅ ∣ x , y ^ < n ) ) , \begin{split}D\bigl(p_{T}\,\|\,p_{S}\bigr)(\hat{y}\mid x)&\triangleq\frac{1}{|\hat{y}|}\sum_{n=1}^{|\hat{y}|}D\biggl(p_{T}\!\left(\cdot\mid x,y^{\star},\hat{y}_{<n}\right)\\ &\qquad\big\|\;p_{S}\!\left(\cdot\mid x,\hat{y}_{<n}\right)\biggr),\end{split} (6) where p S ( ⋅ ∣ x , y ^ < n ) p_{S}(\cdot\mid x,\hat{y}_{<n}) and p T ( ⋅ ∣ x , y ⋆ , y ^ < n ) p_{T}(\cdot\mid x,y^{\star},\hat{y}_{<n}) denote distributions over the next token y n ∈ 𝒱 y_{n}\in\mathcal{V} . Here, D D can be any distribution divergence measure such as the generalized Jensen-Shannon divergence JSD β \operatorname{JSD}_{\beta} , defined for a weight β ∈ [ 0 , 1 ] \beta\in[0,1] as: JSD β ( p T ∥ p S ) = β D K ​ L ( p T ∥ m ) + ( 1 − β ) D K ​ L ( p S ∥ m ) \operatorname{JSD}_{\beta}(p_{T}\|p_{S})=\beta D_{KL}(p_{T}\|m)+(1-\beta)D_{KL}(p_{S}\|m) (7) where m = β ​ p T + ( 1 − β ) ​ p S m=\beta p_{T}+(1-\beta)p_{S} is the interpolated mixture distribution. This full-vocabulary formulation provides dense, token-level feedback: the teacher, informed by y ⋆ y^{\star} , exposes the student to the entire distribution over plausible next tokens and guides it toward reasoning paths that lead to the correct answer.

We minimize the expected divergence between teacher and student over on-policy student samples: ℒ ( θ ) = 𝔼 ( x , y ⋆ ) ∼ 𝒮 [ 𝔼 y ^ ∼ p S ( ⋅ ∣ x ) [ D ( p T ∥ p S ) ( y ^ ∣ x ) ] ] . \mathcal{L}(\theta)=\mathbb{E}_{(x,y^{\star})\sim\mathcal{S}}\left[\mathbb{E}_{\hat{y}\sim p_{S}(\cdot\mid x)}\left[D\bigl(p_{T}\,\|\,p_{S}\bigr)(\hat{y}\mid x)\right]\right]. (8) Gradients are backpropagated only through the student policy p S p_{S} , while the teacher p T p_{T} acts as a fixed full-distribution target conditioned on privileged information ( x , y ⋆ ) (x,y^{\star}) .

##### Per-Token Pointwise Divergence Clipping.

In our experiments, we observe that token-level divergence is highly skewed across vocabulary entries: a small subset of stylistic tokens exhibits much higher divergence than mathematically meaningful tokens (see Table 5 ). This imbalance causes the training signal to be dominated by stylistic patterns. To address this, we apply pointwise clipping to the vocabulary-level divergence contributions. Let D f ( p T ∥ p S ) D_{f}(p_{T}\|p_{S}) denote an f f -divergence. At each token position n n and vocabulary entry v v , define: ℓ n , v ( f ) = p T ( v ∣ ⋅ ) f ( p S ( v ∣ ⋅ ) p T ( v ∣ ⋅ ) ) . \ell_{n,v}^{(f)}=p_{T}(v\mid\cdot)\;f\!\left(\frac{p_{S}(v\mid\cdot)}{p_{T}(v\mid\cdot)}\right). We compute the clipped divergence: D clip ( f ) ( p T ∥ p S ) = 1 | y ^ | ∑ n = 1 | y ^ | ∑ v ∈ 𝒱 min ( ℓ n , v ( f ) , τ ) . D_{\mathrm{clip}}^{(f)}(p_{T}\|p_{S})=\frac{1}{|\hat{y}|}\sum_{n=1}^{|\hat{y}|}\sum_{v\in\mathcal{V}}\min(\ell_{n,v}^{(f)},\tau).

##### Alternative objective: Sampled-token distillation through policy gradient.

Following recent on-policy distillation methods ( Lu and Lab, 2025 ) , we form a sampled-token reward signal (a reverse-KL signal on sampled actions) and optimize with policy gradient. For each position n n in a sampled sequence y ^ \hat{y} , define the advantage term A n ​ ( x , y ^ ) = log ⁡ p T ​ ( y ^ n ∣ x , y ⋆ , y ^ < n ) − log ⁡ p S ​ ( y ^ n ∣ x , y ^ < n ) , A_{n}(x,\hat{y})=\log p_{T}\!\left(\hat{y}_{n}\mid x,y^{\star},\hat{y}_{<n}\right)-\log p_{S}\!\left(\hat{y}_{n}\mid x,\hat{y}_{<n}\right), and optimize the policy-gradient-style objective ℒ ⁡ ( θ ) = − 𝔼 ( x , y ⋆ ) ∼ 𝒮 [ 𝔼 y ^ ∼ p S ( ⋅ ∣ x ) [ 1 | y ^ | ∑ n = 1 | y ^ | A n ( x , y ^ ) × log p S ( y ^ n ∣ x , y ^ < n ) ] ] . \begin{split}\mathcal{L}(\theta)&=-\mathbb{E}_{(x,y^{\star})\sim\mathcal{S}}\biggl[\mathbb{E}_{\hat{y}\sim p_{S}(\cdot\mid x)}\biggl[\frac{1}{|\hat{y}|}\sum_{n=1}^{|\hat{y}|}A_{n}(x,\hat{y})\\ &\qquad\times\log p_{S}\!\left(\hat{y}_{n}\mid x,\hat{y}_{<n}\right)\biggr]\biggr].\end{split} (9) A n ​ ( x , y ^ ) A_{n}(x,\hat{y}) is treated as a constant with respect to θ \theta (i.e., gradients do not flow through the advantage), so that gradients take the usual policy-gradient form A n ​ ∇ θ ​ log ⁡ p S A_{n}\nabla_{\theta}\log p_{S} . Compared to the full-vocabulary divergence objective, this on-policy shaping objective operates only on sampled tokens, using the teacher’s log-probabilities to provide dense, trajectory-level shaping signals without explicitly matching the full distribution at each step.

##### OPSD as dense-reward policy gradient and comparison to STaR.

The objective in Equation 9 can be seen as policy gradient with dense, token-level rewards. In Appendix Appendix D , we formalize this and contrast with STaR ( Zelikman et al., 2022 ) , a closely related method that also uses the same model to generate reasoning traces, then performs rejection sampling followed by SFT on correct traces. This procedure can be viewed as policy gradient with a sequence-level binary reward that assigns identical credit to all tokens and vanishes when samples are incorrect. In contrast, OPSD provides feedback at every token position regardless of final-answer correctness.

## 4 Experiments

We conduct comprehensive experiments to answer the following research questions: (1) How does OPSD compare to SFT and GRPO in reasoning performance and sample efficiency? (§ 4.2 )

(2) How does per-token pointwise KL clipping in OPSD help stabilizing training? (§ 4.3.3 )

(3) What is the effect of generation style, generation length on performance? (§ 4.3.4 )

(4) Does full-vocabulary logit distillation provide benefits over sampled-token policy gradient? (§ 4.3.5 )

### 4.1 Experimental Setup

Models and datasets. We experiment with the Qwen3 ( Team, 2025b ) model family at three scales: Qwen3-1.7B, Qwen3-4B, and Qwen3-8B, using the instruct-tuned versions. For training data, we use the mathematical reasoning subset of OpenThoughts ( Guha et al., 2025 ) , sampling up to 30K problem-solution pairs with chain-of-thought reasoning. We evaluate on competition-level mathematics benchmarks including AIME 2024, AIME 2025, HMMT 2025.

Baselines. We compare against two methods trained on the same dataset: (1) SFT , standard supervised fine-tuning on expert trajectories, which can be seen as off-policy distillation from a more powerful LLM that generated the reasoning traces; (2) GRPO ( Shao et al., 2024 ) , group relative policy optimization with binary outcome rewards verified against ground-truth answers. The max generation length is set to 16k.

Implementation details. We fix the teacher policy to be the initial policy, rather than the currently updating learning policy, as we find this helps stabilize training and implicitly acts as regularization to prevent excessive deviation from the initial policy. We use full-vocabulary logit distillation in our experiments. All experiments are conducted on A100 or H100 GPUs with LoRA ( Hu et al., 2022 ) . More experimental details are in Appendix B .

### 4.2 Main Results

Table 2 reports results on competition-level mathematical reasoning benchmarks. OPSD consistently outperforms SFT and improves over the base model across all scales, matching or exceeding GRPO in every setting. Notably, OPSD achieves these gains using only a single rollout per problem and converges within 100 steps, with each problem requiring only 1024 sampled tokens, whereas GRPO requires 8 rollouts of 16k tokens each and may exhibit performance degradation in later steps due to entropy collapse—with most of reward standard deviations within a group being zero under this OpenThoughts dataset, yielding no learning signal and wasting sampling budget. We also observe consistent performance degradation under SFT across tasks and model scales when trained on the same dataset, which we attribute to the concise reasoning style of the ground truth solutions which has reduced reasoning lengths at test time. We attribute OPSD’s token efficiency to dense token-level supervision from the teacher distribution, and we hypothesize that earlier tokens may contribute more to effective distillation as they could represent more critical branching points in the reasoning process.

As shown in Figure 3 , OPSD achieves higher token learning efficiency within 100 steps of training as compared to GRPO. Within 100 steps, GRPO’s performance stagnates with less learning signal when the outcome reward within as sampling group remains the same, leading to zero gradient. These results suggest that OPSD may extract learning signal from the same reasoning datasets more efficiently than both GRPO and SFT, while substantially reducing training time.

### 4.3 Ablation Studies & Discussions

In this section, we conduct extensive ablations to study key design choices in OPSD, including (1) the divergence objective, (2) the generation styles of the student and teacher (e.g., thinking-mode on/off), (3) the effect of per-token KL clipping, (4) the impact of student generation length, and (5) comparison between full-vocabulary logit distillation with sampled-token distillation.

#### 4.3.1 Effect of Divergence Objective

A key design choice in OPSD is the divergence used for per-token distribution matching between the privileged teacher and the student. We compare forward KL, reverse KL, and JSD on AIME25 with Qwen3-1.7B in Table 3 . All objectives are evaluated under the same pointwise clipping scheme for stability. Forward KL consistently yields the strongest gains, improving performance from 36.7 to 43.9 at step 50 and remaining above the baseline at step 100. In contrast, reverse KL and JSD provide limited or negative improvements. We therefore adopt forward KL in all remaining experiments.

#### 4.3.2 Effect of Generation Styles and per-token KL Clipping

Another key design choice in OPSD is the generation style of the student and teacher models, as it determines both which tokens the student learns from and the style of supervision provided by the teacher. Qwen3 models support two generation modes: Thinking Mode on (TM-on), in which the model produces self-reflective chain-of-thought tokens, and Thinking Mode off (TM-off), in which it generates responses directly. To determine which combination yields the most effective learning signal, we analyze the forward KL divergence KL ( p T ∥ p S ) \mathrm{KL}(p_{T}\|p_{S}) across all four student/teacher mode pairings, categorizing tokens into three groups: math (numerals, operators, and mathematical keywords), style (reasoning connectives), and other . Table 5 reports the mean per-token KL within each category.

Across all model sizes, the TM-off student paired with a TM-on teacher yields the largest KL on math tokens, indicating stronger supervision on mathematically relevant tokens. The reported KL values correspond to the expected divergence over the vocabulary at each position; as shown in Table 5 , this expectation is highly skewed, with stylistic tokens contributing disproportionately large values. This motivates our use of pointwise clipping to control such heavy-tailed contributions. Empirically, this configuration achieves the best downstream performance. We therefore adopt the TM-off student / TM-on teacher configuration.

#### 4.3.3 Effect of Per-Token Pointwise Clipping

As shown in Table 5 , stylistic tokens can exhibit higher KL divergence than math-related tokens, causing them to dominate the training signal. We mitigate this issue using per-token pointwise clipping. As shown in Figure 4 for Qwen3-1.7B, clipping stabilizes training and prevents performance degradation, which is particularly important given that OPSD converges rapidly within a hundred steps of training.

#### 4.3.4 Effect of Generation Length

Since our objective operates at the token level (Eq. 6 ), the number of generated tokens per sample directly determines the amount of supervision signal available to the student. Longer sequences expose the student to more teacher feedback, but they also increase computational cost and may introduce noisy or uninformative continuations. To study this trade-off, we conduct an ablation on Qwen3-1.7B by varying the generation length of on-policy sampled student responses among 1024 and 4096 tokens and use full-vocabulary logit distillation. As shown in Figure 5 , increasing the generation length does not lead to consistent improvements across either task. We attribute this to early tokens being more critical for learning: as the student generation grows longer, later tokens become increasingly predictable to the teacher when conditioned on a sufficiently long student prefix so less penalties are applied to later tokens. This phenomenon is also noted in ( Lu and Lab, 2025 ) .

#### 4.3.5 Learning Objective Comparison: Full Vocabulary Logits Distillation vs. Sampled-Token Distillation

Our objective in Eq. 6 is defined as a per-token discrepancy between the teacher and student distributions . In practice, OPSD can instantiate this objective in two ways. (1) Full-vocabulary logit distillation (as in GKD ( Agarwal et al., 2024 ) ): for each token position, we compute D ( p T ∥ p S ) D(p_{T}\,\|\,p_{S}) over the entire vocabulary via a full softmax, yielding a proper token-level f f -divergence between the two policies. (2) Sampled-token advantage policy-gradient objective (as in the on-policy distillation method of Lu and Lab (2025) ): we evaluate teacher and student log-probabilities only at the token actually sampled by the student, y ^ n \hat{y}_{n} , and use the reverse-KL term as a scalar advantage inside a policy-gradient-style loss. Thus, the first variant directly matches full token distributions, whereas the second optimizes an on-policy RL objective shaped by the teacher’s log-probabilities rather than a full-distribution divergence. We compare these variants on Qwen3-4B using a 2048-token generation budget during distillation. Table 4 summarizes the results. The full-vocabulary divergence objective provides a consistent gain over the sampled-token objective. This suggests that exposing the student to the full teacher distribution offers richer supervision than relying solely on per-token on-policy shaping. However, the full-vocabulary computation incurs higher peak memory usage due to storing vocabulary-sized logits at every position, indicating a trade-off between performance and efficiency.

## 5 Related Work

##### LLM Self-Training.