Technical Summary: Temperature-Scaled On-Policy Self-Distillation (TS-OPSD)

1. Problem Statement

Reinforcement Learning from Verifiable Rewards (RLVR) has significantly enhanced the reasoning capabilities of Large Language Models (LLMs), particularly in mathematical and logical domains. However, this progress faces a critical bottleneck known as entropy collapse. As policy optimization proceeds, the model's output distribution becomes increasingly concentrated (sharp), leading to:

• Reduced Rollout Diversity: The model generates repetitive trajectories, limiting the exploration of the solution space.
• Diminished Learning Signals: Effective gradient information concentrates on a small fraction of high-entropy tokens at critical reasoning branch points. Once the policy collapses, these signals vanish, stalling further improvement.

Existing remedies attempt to mitigate this by either:

1. Modifying the RL Objective: Using entropy regularization or adaptive clipping. These methods impose external constraints that may trade off against reward maximization.
2. Adjusting Sampling Temperature: Increasing temperature during rollout collection. While this alters the behavior policy, it does not change the model's underlying parameterized distribution, creating a mismatch between the behavior policy and the optimized policy.

These approaches remain external to the model parameters. The paper posits that the exploratory effect of temperature should be internalized directly into the model weights to provide a better initialization for continued RL.

2. Methodology: TS-OPSD

The authors propose Temperature-Scaled On-Policy Self-Distillation (TS-OPSD), a lightweight "policy reheating" method designed to restore entropy without external teachers, privileged data, or additional inference costs.

Core Mechanism

TS-OPSD operates as an intervention between RL stages. Starting from an entropy-collapsed checkpoint ($\theta_{RL-collapse}$), the method constructs a self-teacher and distills a smoother distribution back into the student model ($\theta_{reheat}$).

1. Self-Teacher Construction: Given the student's current logits $z_\theta$, a self-teacher distribution $p_T$ is generated by applying a high-temperature softmax transformation ($T > 1$) to the logits. Crucially, the teacher logits are stop-gradient, and no external model is used.
   $$p_T(v | x_{<t}) = \frac{\exp(z_{\theta,v}/T)}{\sum_u \exp(z_{\theta,u}/T)}$$
2. On-Policy Distillation: The student is trained to minimize the Kullback-Leibler (KL) divergence between its current distribution $q_\theta$ and the high-temperature self-teacher $p_T$, using on-policy rollouts sampled from $q_\theta$ itself.
   $$\mathcal{L} = D_{KL}(p_T \parallel q_\theta) \quad \text{(Forward KL)}$$
   $$\mathcal{L} = D_{KL}(q_\theta \parallel p_T) \quad \text{(Reverse KL)}$$

Structural Properties

• Rank Preservation: Since $p_T$ is a monotone power transform of $q_\theta$ ($p_T(v) \propto q_\theta(v)^{1/T}$), the token ranking is preserved. The teacher and student share the same local support and semantic preference ordering; only the distributional sharpness differs.
• Internalization: Unlike standard decoding temperature adjustments, TS-OPSD updates the model weights to permanently encode the smoother, higher-entropy behavior. This eliminates the train-inference mismatch.
• Gradient Dynamics:
  • Forward KL (FKL): The gradient takes the form $q_\theta(v) - p_T(v)$. This performs "structured smoothing," re-weighting suppressed alternatives based on the model's own prior judgments.
  • Reverse KL (RKL): Under self-temperature coupling, minimizing RKL is locally equivalent to maximizing the student's entropy, with the temperature acting only as a scalar multiplier on the gradient magnitude.

3. Key Contributions

1. Policy Reheating Framework: Introduces TS-OPSD, a method to internalize the exploratory effect of temperature into model parameters, serving as a post-collapse intervention.
2. Theoretical Insight: Demonstrates that under self-temperature coupling, the teacher-student mismatch is minimized (rank-preserving), and RKL reduces to a direct entropy-maximizing update.
3. Empirical Validation: Shows that reheating an entropy-collapsed checkpoint provides a superior initialization for continued RL compared to standard continued training or rollout-level temperature adjustments.
4. Analysis of Dynamics: Characterizes the "entropy jump" phenomenon, where entropy recovery occurs in two phases (slow smoothing followed by a sharp discontinuous jump), and analyzes the reversibility of these parameter updates.

4. Experimental Results

Experiments were conducted on Qwen3-4B-Base and Qwen3-8B-Base using the DAPO-Math-17K dataset and evaluated on MATH-500, AIME24/25/26, and OlymMATH.

• Performance Gains: TS-OPSD (specifically with Forward KL, PR-FKL) consistently outperformed both the standard DAPO baseline and Rollout Reheating (DAPO w/ RR).
  • On Qwen3-8B-Base, PR-FKL achieved an Avg pass@8 of 54.74, compared to 51.60 for standard DAPO and 52.66 for Rollout Reheating.
  • Gains were most pronounced on harder benchmarks (AIME, OlymMATH).
• Entropy Recovery: After reheating, models resumed RL from a higher-entropy initialization and achieved higher final rewards than standard DAPO, which remained suppressed.
• Comparison with Rollout Reheating: While increasing sampling temperature during rollouts (RR) provided non-trivial gains, Policy Reheating (PR) yielded superior results. This suggests the benefit stems from the smoother policy distribution itself acting as a better foundation for optimization, rather than just noisier sampling.
• Preservation of Capabilities: Analyses showed that TS-OPSD performs support-preserving probability redistribution:
  • Intermediate Representations: Linear CKA scores between base and reheated layers remained near 1.0, indicating internal representations were largely unchanged.
  • Reasoning Capability: Task accuracy on MATH-500 and GPQA remained stable throughout the distillation process, confirming that entropy recovery did not degrade reasoning skills.
  • Candidate Sets: Top-10 token overlap and probability mass retention were near 1.0, indicating the model's semantic preferences were preserved while margins were softened.

5. Significance and Claims

The paper claims that entropy restoration can be decoupled from the RL objective and inserted as a modular, inter-stage operation when a policy becomes too deterministic.

• Modular Intervention: TS-OPSD offers a simple, lightweight mechanism to extend reasoning-oriented RL beyond the entropy-collapse bottleneck without altering the core RL algorithm.
• Stability in Data-Scarce Settings: In experiments with reduced data (2048 prompts), standard RL training became unstable (gradient spikes, reward degradation) due to severe entropy collapse. TS-OPSD acted as a stabilizing mechanism, allowing for longer, stable training.
• Computational Efficiency: The reheating intervention adds approximately 15–17% to the total RL training time. When normalized for batch size, the per-sample cost is comparable to or lower than standard RL steps.
• Limitations: The authors note that experiments are limited to specific model scales (4B/8B) and that the stopping criteria for reheating (entropy threshold) are currently empirical rather than principled. They suggest future work could explore cascaded schedules of alternating OPSD and RL.

In summary, TS-OPSD demonstrates that "internalizing" temperature via self-distillation effectively restores the exploratory capacity of collapsed policies, enabling continued RL to reach higher performance ceilings while preserving the reasoning structures acquired during training.