OISD: On-Policy Internal Self-Distillation of Language Models
cs.LG, cs.AI, cs.CV
Submitted: 2026-05-27
Updated: 2026-08-30
Comments: Findings of the Association for Computational Linguistics: EMNLP 2026
Code: https://github.com/THE-MALT-LAB/OISD
License: http://creativecommons.org/licenses/by/4.0/
The gist: Recent reinforcement learning (RL) post-training approaches primarily optimize the final output policy using sparse outcome-level rewards, while largely overlooking predictive signals encoded in
Terminology
Abstract
Recent reinforcement learning (RL) post-training approaches primarily optimize the final output policy using sparse outcome-level rewards, while largely overlooking predictive signals encoded in intermediate representations. In this paper, we introduce a new paradigm called on-policy internal self-distillation and propose the OISD framework, which improves reasoning by transferring on-policy predictive signals from the final layer to intermediate representations. During rollout and Group Relative Policy Optimization (GRPO) optimization, the final layer acts as both the policy and a detached internal teacher for selected intermediate layers, which are guided to align with it through two complementary mechanisms: logit alignment, which transfers high-level reasoning behaviors (how to think), and attention alignment, which enforces consistent attention patterns (where to look) from the final layer to the selected intermediate layer, both without requiring external privileged information. Our OISD, together with GRPO, employs signed advantage-weighted Jensen--Shannon alignment to distill informative intermediate representations while preserving policy consistency under a unified acting policy. Experimental results demonstrate the effectiveness of OISD, with substantial and consistent improvements over strong reasoning RL baselines across four mathematical reasoning tasks. The code will be released at https://github.com/THE-MALT-LAB/OISD
Sources
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- How Do LLMs Use Their Depth?
- Training Large Language Models to Reason in a Continuous Latent Space
- Distilling the Knowledge in a Neural Network
- Reinforcement Learning via Self-Distillation
- Layer Normalization
- Eliciting Latent Predictions from Transformers with the Tuned Lens
- LEAP: Layer-wise Exit-Aware Pretraining for Efficient Transformer Inference
- Reinforced Attention Learning
- Qwen3 Technical Report
- Proximal Policy Optimization Algorithms
- ReAct: Synergizing Reasoning and Acting in Language Models
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- On-Policy Context Distillation for Language Models
- Reinforcement Learning Fine-Tuning Enhances Activation Intensity and Diversity in the Internal Circuitry of LLMs
- Self-Distillation Enables Continual Learning
- Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models
- Group Sequence Policy Optimization
- Bottom-up Policy Optimization: Your Language Model Policy Secretly Contains Internal Policies
- A Survey on Latent Reasoning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks