OISD: On-Policy Internal Self-Distillation of Language Models

• URL: http://arxiv.org/abs/2605.29089v1
• Date: Wed, 27 May 2026 20:43:10 GMT
• Title: OISD: On-Policy Internal Self-Distillation of Language Models
• Abstract summary: Recent reinforcement learning approaches primarily optimize the final output policy using sparse outcome-level rewards.<n>We propose the OISD framework, which improves reasoning by transferring on-policy predictive signals from the final layer to intermediate representations.<n> Experimental results demonstrate the effectiveness of OISD, with substantial and consistent improvements over strong reasoning RL baselines.
• Score: 20.760265566511183
• License: http://creativecommons.org/licenses/by/4.0/
• 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

Related papers

• IADD-TR: Intervention-Aware Dynamics Decoupling with Targeted Regularization for Model-Based Reinforcement Learning [57.16513298507272]
  We propose IADD-TR, a unified framework combining Intervention-Aware Dynamics Decoupling (IADD) and Targeted Regularization (TR)<n>IADD factorizes transitions into an action-intervention stage and an action-free natural evolution stage, using a zero-action anchor to resolve the non-uniqueness of this two-stage factorization for robust generalization.<n>For policy learning, we derive TR from the efficient influence function of a replay-state policy-gradient functional.
  arXiv  Detail & Related papers  (2026-08-11T08:20:02Z)
• SR-OPSD: Self-Referenced On-Policy Self-Distillation [18.590739659841926]
  On-policy self-distillation (OPSD) converts feedback into dense token-level supervision on trajectories generated by the policy to be optimized.<n>We propose the emphSelf-Referenced On-Policy Self-Distillation (SR-OPSD) to achieve state-of-the-art or competitive performance across various settings.
  arXiv  Detail & Related papers  (2026-08-10T15:40:07Z)
• TAPO: Transition-Aware Policy Optimization for LLM Agents [27.95488836373643]
  TAPO is a unified training framework that alternates between policy optimization and transition supervision.<n>It serves as a computationally lightweight, plug-and-play enhancement module for existing agent RL algorithms.<n>We conduct systematic experiments on WebShop and ALFWorld, integrating foundation models of various scales with different policy optimization algorithms.
  arXiv  Detail & Related papers  (2026-07-30T10:17:55Z)
• SEED: Self-Evolving On-Policy Distillation for Agentic Reinforcement Learning [38.54413500261524]
  Outcome-based reinforcement learning (RL) provides a practical optimization paradigm, but its sparse trajectory-level rewards offer limited guidance on intermediate decisions.<n>We propose SEED (SElf-Evolving On-Policy Distillation), a self-evolving framework that converts completed on-policy trajectories into training-time hindsight skills.
  arXiv  Detail & Related papers  (2026-07-16T09:57:18Z)
• dOPSD: On-Policy Self-Distillation for Diffusion Language Models [52.60302464420127]
  Diffusion large language models generate text by iteratively denoising a masked sequence.<n>Supervised fine-tuning is off-policy and suffers from exposure bias, while reinforcement learning gives only sparse, sequence-level rewards.<n>On-policy self-distillation (OPSD) offers a promising alternative, using one model as both student and teacher.
  arXiv  Detail & Related papers  (2026-07-05T17:47:29Z)
• Segment-Aligned Policy Optimization for Multi-Modal Reasoning [55.29606572822562]
  We propose a novel reinforcement learning paradigm that treats coherent reasoning steps, rather than tokens or full sequences as fundamental units of policy update.<n>Experiments on representative reasoning benchmarks demonstrate that SAPO consistently outperforms token-level and sequence-level policy optimization methods.<n>Our work underscores the importance of aligning reinforcement learning updates with the structure of reasoning, paving the way for more efficient and semantically grounded policy optimization in complex reasoning tasks.
  arXiv  Detail & Related papers  (2026-05-02T08:47:45Z)
• Soft Sequence Policy Optimization [0.0]
  We introduce Soft Sequence Policy Optimization (SSPO) as an off-policy reinforcement learning objective.<n>SSPO incorporates soft gating functions over token-level probability ratios within sequence-level importance weights.<n>We show that SSPO improves training stability and performance in mathematical reasoning tasks.
  arXiv  Detail & Related papers  (2026-02-22T20:21:00Z)
• f-GRPO and Beyond: Divergence-Based Reinforcement Learning Algorithms for General LLM Alignment [15.396104072574104]
  We propose f-Group Relative Policy Optimization (f-GRPO), a class of on-policy reinforcement learning, and f-Hybrid Alignment Loss (f-HAL)<n>We provide theoretical guarantees that these classes of objectives improve the average reward after alignment.
  arXiv  Detail & Related papers  (2026-02-05T18:01:52Z)
• Bottom-up Policy Optimization: Your Language Model Policy Secretly Contains Internal Policies [30.35690865689118]
  Existing reinforcement learning approaches treat large language models (LLMs) as a single unified policy, overlooking their internal mechanisms.<n>We decompose the language model policy by leveraging the intrinsic split of the Transformer residual stream and the equivalence between the composition of hidden states with the unembedding matrix.<n>This decomposition reveals Internal Layer Policies, corresponding to contributions from individual layers, and Internal Modular Policies, which align with the self-attention and feed-forward network.
  arXiv  Detail & Related papers  (2025-12-22T18:51:48Z)
• Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning [77.92320830700797]
  Reinforcement Learning has played a central role in enabling reasoning capabilities of Large Language Models.<n>We propose a tractable computational framework that tracks and leverages curvature information during policy updates.<n>The algorithm, Curvature-Aware Policy Optimization (CAPO), identifies samples that contribute to unstable updates and masks them out.
  arXiv  Detail & Related papers  (2025-10-01T12:29:32Z)
• Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient [61.440209025381016]
  Policy reinforcement learning frequently utilize deep neural networks (DNNs) to learn a shared backbone of feature representations used to compute likelihoods in an action selection layer.<n>We show that under certain constraints, a structure resembling neural collapse, which we refer to as Action Collapse (AC), emerges.<n>We propose the Action Collapse Policy Gradient (ACPG) method, which accordingly affixes a synthetic ETF as our action selection layer.
  arXiv  Detail & Related papers  (2025-09-02T18:33:11Z)
• PARL: A Unified Framework for Policy Alignment in Reinforcement Learning from Human Feedback [106.63518036538163]
  We present a novel unified bilevel optimization-based framework, textsfPARL, formulated to address the recently highlighted critical issue of policy alignment in reinforcement learning. Our framework addressed these concerns by explicitly parameterizing the distribution of the upper alignment objective (reward design) by the lower optimal variable. Our empirical results substantiate that the proposed textsfPARL can address the alignment concerns in RL by showing significant improvements.
  arXiv  Detail & Related papers  (2023-08-03T18:03:44Z)
• The Gradient Convergence Bound of Federated Multi-Agent Reinforcement Learning with Efficient Communication [20.891460617583302]
  The paper considers independent reinforcement learning (IRL) for collaborative decision-making in the paradigm of federated learning (FL) FL generates excessive communication overheads between agents and a remote central server. This paper proposes two advanced optimization schemes to improve the system's utility value.
  arXiv  Detail & Related papers  (2021-03-24T07:21:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.