OISD: On-Policy Internal Self-Distillation of Language Models
On the map: LLM reasoning → Prompting and chain-of-thought reasoning
also in 1 more
Abstract (arXiv)
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
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Placed in Prompting and chain-of-thought reasoning, one of 17 directions in LLM reasoning. Builds on 24 of its 36 references that are in the map. First posted on arXiv 27 May 2026.
Builds on
24 of 36 references.
Closest ancestors · 2
Same direction, newest first
- 2026Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information17 Apr
- 2025CODI: Compressing Chain-of-Thought into Continuous Space via Self-Distillation28 Feb
Other references · 22
- 2022Training language models to follow instructions with human feedbackcited by 15,971
- 2017Proximal Policy Optimization Algorithmscited by 15,215
- 2015Distilling the Knowledge in a Neural Networkcited by 10,324
- 2022ReAct: Synergizing Reasoning and Acting in Language Modelscited by 8,695
- 2020Measuring Massive Multitask Language Understandingcited by 7,963
- 2016Layer Normalizationcited by 5,669
- 2023Let's Verify Step by Stepcited by 3,789
- 2023GPQA: A Graduate-Level Google-Proof Q&A Benchmarkcited by 2,879
- 2023On-Policy Distillation of Language Models: Learning from Self-Generated Mistakescited by 939
- 2015Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkscited by 803
- 2023Eliciting Latent Predictions from Transformers with the Tuned Lenscited by 623
- 2014Deeply-Supervised Netscited by 367
10 more
- 2026Reinforcement Learning via Self-Distillationcited by 332
- 2024LayerSkip: Enabling Early Exit Inference and Self-Speculative Decodingcited by 307
- 2024ProcessBench: Identifying Process Errors in Mathematical Reasoningcited by 292
- 2026On-Policy Context Distillation for Language Modelscited by 186
- 2025PRMBench: A Fine-grained and Challenging Benchmark for Process-Level Reward Modelscited by 99
- 2024Speculative Knowledge Distillation: Bridging the Teacher-Student Gap Through Interleaved Samplingcited by 90
- 2025Bottom-up Policy Optimization: Your Language Model Policy Secretly Contains Internal Policiescited by 23
- 2025How Do LLMs Use Their Depth?cited by 16
- 2025Reinforcement Learning Fine-Tuning Enhances Activation Intensity and Diversity in the Internal Circuitry of LLMscited by 9
- 2026Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversitycited by 9
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2026On the Geometry of On-Policy Distillation
- 2026Proxy OPD: On-Policy Distillation with Transferable Relative Proxy Update
- 2026Learning to Foresee: Unveiling the Unlocking Efficiency of On-Policy Distillation
- 2026The Many Faces of On-Policy Distillation: Pitfalls, Mechanisms, and Fixes
- 2026SAPD: Step-Aligned Privileged Distillation
- 2026A Brief Overview: On-Policy Self-Distillation In Large Language Models
- 2026SOD: Step-wise On-policy Distillation for Small Language Model Agents
- 2026CataOPD: Catalytic On-Policy Distillation for Large Language Model Reasoning