Anti-Self-Distillation for Reasoning RL via Pointwise Mutual Information

Source-first digest for monthly 2026_05 rank 22, rank_id p010.

• Routing status: success, route full_markdown
• PDF extraction: not used

Motivation / Background

Reinforcement learning from verifiable rewards gives math-reasoning models a clean final-answer signal, but the reward is sparse: a rollout is marked right or wrong without saying which token or reasoning turn mattered. The paper studies whether on-policy self-distillation can provide that missing token-level credit signal without using a stronger outside teacher. In this setup, the student samples a solution, and a copy of the same model is re-run with privileged context such as a verified solution and feedback.

The surprising empirical starting point is that this self-teacher often hurts difficult math reasoning. The authors argue that the teacher becomes too much like an oracle: once it sees the verified solution, it boosts tokens already implied by that solution and suppresses the exploratory tokens that make long reasoning work. The overview and HMMT 2025 learning curve in Figure 1 frame the diagnosis and the proposed reversal.

Claims And Evidence

Support scores are support-from-paper scores, not independent reproduction scores.

Core Technical Idea

The paper's central move is to interpret the self-teacher log-ratio as conditional pointwise mutual information. Let the student distribution be \(\pi_S(\cdot \mid x, y_{\lt t})\), the privileged self-teacher be \(\pi_T(\cdot \mid x, c, y_{\lt t})\), and define \(u_t = \log \pi_T(y_t) - \log \pi_S(y_t)\). Then:

Under default self-distillation, the token reward is \(+u_t\). Tokens that become more likely once the verified solution is known get positive reward; tokens that the oracle-conditioned teacher views as unnecessary get negative reward. In the paper's diagnostic example, this means structural and answer-committing tokens are rewarded, while exploratory words such as "Wait", "Let", and "Maybe" are suppressed. Figure 2 shows both a single-rollout trace and an aggregate heatmap of this split.

The key equations in Table 1 define the full method. AntiSD does not add a learned process reward model. It changes the sign and shape of the existing self-teacher signal. Instead of descending reverse KL toward the teacher, it ascends a Jensen-Shannon divergence. The induced advantage \(-\varphi(u_t)\) is positive for deliberation tokens, negative for shortcut tokens, and bounded on the over-sampled deliberation side by the JSD softplus shape.

Method Details

AntiSD keeps the GRPO training loop and changes only the per-token auxiliary term. The base RL advantage \(A_i^{\mathrm{seq}}\) remains the group-normalized final-answer reward from GRPO. For each sampled token, the method computes the student log-probability \(s_{i,t}\), the stop-gradient privileged teacher log-probability \(t_{i,t}\), and \(u_{i,t}=t_{i,t}-s_{i,t}\). The training advantage becomes:

The entropy-triggered gate controls \(\lambda\). During five warmup steps with \(\lambda=0\), the method measures the median teacher entropy \(H_{\mathrm{warm}}\). It disables the AntiSD term when \(H < \tau_{\mathrm{down}}\), where \(\tau_{\mathrm{down}}=0.93H_{\mathrm{warm}}\), and re-enables it only when entropy recovers to \(H_{\mathrm{warm}}\). This Schmitt-trigger style rule is meant to avoid using the teacher after it has collapsed into a near-deterministic answer-template distribution.

The appendix proof adds two useful interpretations. First, the self-teacher signal telescopes over a trajectory:

So \(u_t\) can be viewed as a potential-based shaping increment. Second, the JSD derivative satisfies \(f'(\pi_S/\pi_T)=\frac{1}{2}(\log 2-\mathrm{softplus}(u))=-\varphi(u)\), giving the paper's one-step AntiSD advantage.

The shared run configuration in Table 2 matters for interpreting the claims: the reported gains are not from changing the model, dataset, rollout count, optimizer, or evaluation sampling. The privileged teacher context is a verified solution sampled from a successful rollout in the same batch when available, otherwise from the dataset, plus a correctness feedback string.

Experiments And Results

Table 3 is the main empirical result. AntiSD beats GRPO's final average on every model, with gains from +2.1 points on the strongest Olmo3-7B-TK baseline to +11.5 points on Qwen3-4B-IT-2507. Default SD is consistently worse than GRPO, and sometimes collapses dramatically, which supports the paper's argument that the privileged-context signal has the wrong sign for this setting.

The pass@k evidence in Figure 3 is important because a self-distillation method could look better by reducing diversity and improving only pass@1. The paper reports that on Qwen3-8B, AntiSD is about 13 points ahead at \(k=1\) and still roughly 7 to 10 points ahead by \(k=32\), which is more consistent with solving additional hard problems.

The code probe in Table 4 is smaller than the math study but directionally positive: AntiSD adds +1.2 on HumanEval+ and +2.3 on MBPP+. The authors interpret the smaller gap as plausible because executable-code reward is denser than final-answer math reward.

Figure 4 connects the static results to training behavior. AntiSD's train reward rises early, reaching roughly 0.95 within about 30 steps on Qwen3-8B and Qwen3-4B-IT-2507. GRPO reaches similar training reward only much later, and default SD never reaches it. The entropy traces also justify the gate: on Qwen, default SD can drive teacher and actor entropy toward collapse; on Olmo it can inflate entropy and drift.

Figure 5 and Table 5 separate the method's components. The most direct comparison is default SD versus AntiSD in the main table: reversing the sign changes Qwen3-8B average from 30.6 to 65.7. Among the remaining knobs, canonical JSD plus teacher-entropy gating is strongest on Qwen3-4B. The no-gate variant looks good early but collapses around step 90, so the reported peak is not a stable endpoint.

The continual experiment in Table 6 asks whether AntiSD has to start from the base model. It nearly matches from-base AntiSD after resuming from a saturated GRPO checkpoint, reaching 65.0 average in 30 post-resume steps versus 65.7 for from-base AntiSD. That suggests the deliberation-token signal remains useful even when GRPO's trajectory-level reward has plateaued.

The paper's limitation is scope. The main experiments are math RLVR on Qwen3 and Olmo-3 family models, plus one code-reasoning probe. The authors explicitly leave multi-turn agentic settings, richer privileged contexts, larger model scales, and multimodal conditioning as future work. This keeps the support score for broad generality lower than the support score for the math-RLVR result.

Practical Takeaways

AntiSD is most compelling when three conditions hold. First, the task has sparse outcome reward but reliable verification, as in math RLVR. Second, a privileged context can be constructed during training, such as a verified solution plus correctness feedback. Third, default self-distillation would otherwise compress or shortcut the reasoning process.

For practitioners, the method is attractive because it is operationally small: compute a self-teacher log-ratio, pass it through the JSD-derived softplus shape, subtract it from the GRPO advantage, and gate it by teacher entropy. There is no external teacher, no learned process reward model, and no extra inference-time cost. The cost is that the method depends on careful online training instrumentation and a meaningful privileged context; if that context is weak, biased, or unavailable, the PMI signal may not be useful.

The paper also offers a useful diagnostic lens: when self-distillation shortens answers or destabilizes training, inspect which tokens the privileged teacher raises and lowers. If it rewards tokens that merely follow from knowing the answer and penalizes tokens that explore alternatives, default distillation is acting like shortcut reinforcement rather than process supervision.

Reference Coverage

This digest explicitly links all anchors used above: problem setup, overview figure, PMI diagnosis, PMI trace, method equations, key equations table, hyperparameters table, main results table, pass-at-k figure, code results table, training dynamics figure, failure modes figure, Qwen3-4B ablation table, continual AntiSD table, and limitations.