Technical Summary: OISD – On-Policy Internal Self-Distillation of Language Models

1. Problem Statement

Recent reinforcement learning (RL) post-training approaches for Large Language Models (LLMs), particularly those utilizing Group Relative Policy Optimization (GRPO) with verifiable rewards, have significantly advanced mathematical reasoning. However, these methods are predominantly final-output-centric. They optimize the final acting policy using sparse, outcome-level rewards (e.g., correctness of the final answer) while largely overlooking the predictive signals, reasoning processes, and internal representations encoded in intermediate layers.

While Knowledge Distillation (KD) offers a mechanism to provide dense token-level supervision, existing paradigms suffer from specific limitations:

• Off-Policy Distillation: Creates a distribution mismatch because the student is trained on trajectories generated by a fixed teacher, not its own generations.
• On-Policy Distillation: Reduces the mismatch but typically relies on external teacher models.
• On-Policy Self-Distillation: Recent methods remove external teachers but depend on externally injected privileged information (e.g., privileged prompts, ground-truth solutions, or auxiliary feedback) to supervise the model.

There is a gap in leveraging the model's own internal computations as a source of supervision without requiring external privileged data or multiple acting policies. Furthermore, prior research indicates that reasoning representations evolve systematically across layers (from broad candidate predictions in early layers to converged decisions in deep layers), suggesting that intermediate layers could be guided by deeper layers to improve reasoning coherence.

2. Methodology: OISD Framework

The authors propose OISD (On-Policy Internal Self-Distillation), a framework that treats the model's final layer as a "detached internal teacher" and a selected intermediate layer as an "auxiliary internal student." This occurs within a single forward pass during the rollout and GRPO optimization process.

Core Architecture

• Single Acting Policy: The final layer ($L$) remains the sole policy for generating rollouts and optimizing the GRPO objective.
• Internal Teacher-Student: During the same forward pass, the hidden state of a selected intermediate layer ($l < L$) is projected into token space using the unembedding matrix (following the logit lens framework). This creates a provisional "student" distribution derived from the same context as the "teacher" (final layer).
• No External Privilege: The supervision signals are derived entirely from the model's own internal states on the same on-policy rollout, requiring no external ground-truth solutions or privileged prompts.

Alignment Objectives

OISD employs two complementary internal alignment mechanisms to transfer predictive signals from the final layer to the intermediate layer:

1. Logit Alignment ("How to Think"):

   • Goal: Transfer high-level reasoning behaviors and predictive beliefs.
   • Mechanism: The Jensen–Shannon (JS) divergence is computed between the temperature-scaled token distributions of the intermediate layer ($p^l_t$) and the final layer ($q^L_t$).
   • Weighting: The loss is weighted by the signed GRPO advantage ($A_t$). High-advantage tokens pull the intermediate layer toward the final layer's distribution, while low-advantage tokens push it away to avoid reinforcing poor continuations.
   • Formula: $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. Attention Alignment ("Where to Look"):

   • Goal: Enforce consistent attention patterns and evidence-routing behavior.
   • Mechanism: The JS divergence is computed between the attention weights of the intermediate layer and the final layer over a sampled set of causal keys (combining local windows and global strided positions to manage memory).
   • Weighting: Similar to logit alignment, the loss is weighted by the signed advantage. Positive advantages encourage the intermediate layer to attend to the same evidence as the final layer; negative advantages reduce agreement with low-reward patterns.
   • Formula: $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

Training Objective

The total loss combines the standard GRPO loss with the two internal alignment losses:
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
where $\lambda_{\text{think}}$ and $\lambda_{\text{attn}}$ control the strength of the distillation.

3. Key Contributions

1. New Paradigm: Introduction of On-Policy Internal Self-Distillation, a reasoning RL paradigm that leverages the model's own internal computations for supervision while maintaining a single acting policy.
2. OISD Framework: A specific implementation using the final layer as a detached teacher and an intermediate layer as a student, transferring on-policy predictive signals without external privileged information.
3. Dual Alignment Mechanisms: Development of two complementary objectives:
   • Logit Alignment: Supervises the reasoning process ("how to think").
   • Attention Alignment: Supervises the evidence routing ("where to look").
4. Empirical Validation: Demonstration of substantial and consistent improvements over strong reasoning RL baselines (including GRPO and BuPO) across four mathematical reasoning benchmarks.

4. Experimental Results

The authors evaluated OISD on Qwen3-4B, Qwen3-8B, and OctoThinker models across four benchmarks: AMC23, MATH500, AIME24, and AIME25.

• Performance Gains:
  • On Qwen3-4B, OISD achieved an average score of 64.75, outperforming GRPO (55.08) by 9.67 points and BuPO (58.51) by 6.24 points.
  • Improvements were particularly significant on AIME benchmarks (e.g., +14.31 points over GRPO on AIME24).
  • On Qwen3-8B, OISD achieved the best overall average of 66.72, outperforming GRPO (64.23) and BuPO (66.36).
• Pass@K Analysis: OISD consistently achieved higher Pass@K rates across all $K$ values compared to GRPO, indicating an improvement in the overall quality of the sampled solution set, not just average accuracy.
• Ablation Studies:
  • Both Logit and Attention alignment contributed positively. Logit alignment showed the largest gain on AIME24, while Attention alignment performed slightly better on average.
  • Combining both yielded the best results, confirming their complementary nature.
• Comparison with OPSD: OISD (64.75) significantly outperformed OPSD (56.63), a state-of-the-art on-policy self-distillation method that relies on privileged ground-truth solutions. This highlights the effectiveness of OISD's purely internal supervision.
• Qualitative Analysis:
  • Training Dynamics: OISD showed steady reward improvement and an initial increase in entropy (suggesting broader exploration) before stabilizing.
  • Attention Visualization: Intermediate layers in OISD gradually aligned with the final layer's attention to critical evidence tokens (e.g., endpoints, supremum) during training.
  • Logit-Lens Visualization: OISD exhibited earlier alignment between intermediate and final predictions on reasoning-critical tokens compared to GRPO, which often produced fragmented intermediate predictions.

5. Significance and Claims

The paper claims that OISD initiates the investigation of on-policy internal self-distillation. Its primary significance lies in:

• Eliminating External Dependence: It provides dense internal supervision without requiring external teachers, privileged prompts, or ground-truth solutions during training.
• Policy Consistency: It preserves a single acting policy for both rollout generation and optimization, avoiding the distribution mismatch issues of off-policy methods.
• Enhanced Reasoning: By aligning intermediate representations with the final layer's reasoning and attention patterns, OISD encourages the model to learn stronger, more globally consistent intermediate representations.

The authors note limitations, including the focus on mathematical reasoning with verifiable answers and the lack of an in-depth study on optimal layer selection. They suggest future work should explore generalization to more complex reasoning tasks and multi-layer internal supervision.