From Generic Correlation to Input-Specific Credit in On-Policy Self Distillation

TL;DR

This paper analyzes on-policy self-distillation rewards in language models, revealing they measure pointwise mutual information and proposing CREDIT to isolate input-specific credit, improving performance across benchmarks.

Contribution

It provides a Bayesian filtering interpretation of self-distillation rewards and introduces CREDIT, a method to focus on input-specific information, enhancing model performance.

Findings

CREDIT isolates input-specific reward components effectively.

CREDIT achieves strong performance across multiple benchmarks.

The reward corresponds to a Bayesian filtering increment related to mutual information.

Abstract

On-policy self-distillation has emerged as a promising paradigm for post-training language models, in which the model conditions on environment feedback to serve as its own teacher, providing dense token-level rewards without external teacher models or step-level annotations. Despite its empirical success, what this reward actually measures and what kind of credit it assigns remain unclear. Under a posterior-compatibility interpretation of feedback conditioning, standard in the implicit-reward literature, we show that the self-distillation token reward is a Bayesian filtering increment whose trajectory sum is exactly the pointwise mutual information between the response and the feedback given the input. This pMI can be raised by input-specific reasoning or by input-generic shortcuts, so we further decompose the teacher log-probability along the input axis. Based on this analysis, we propose CREDIT (Contrastive REward from DIsTillation), which isolates the input-specific component with a batch-contrastive baseline. At the sequence level, CREDIT is a teacher-side surrogate for a contrastive pMI objective that also penalizes responses remaining likely under unrelated inputs. Across coding, scientific reasoning, and tool-use benchmarks on two model families, CREDIT delivers the strongest aggregate performance at negligible additional compute.