DemoPSD: Disagreement-Modulated Policy Self-Distillation

Overview§

DemoPSD addresses two key issues in on-policy self-distillation (OPSD) for LLMs: privileged information leakage and suppression of exploration. The core innovation is a token-level adaptive blending of teacher and student distributions using a reverse-KL barycenter target, which balances learning from the teacher with preserving the student's own reasoning capacity.

Methodology§

Reverse-KL Barycenter Target§

Let $p_t$ and $p_s$ denote the teacher and student distributions over tokens. At each position $i$, DemoPSD targets a weighted geometric combination:

$$ q_i \propto p_{t,i}^{\lambda_i} \cdot p_{s,i}^{1-\lambda_i} $$

where $\lambda_i \in [0,1]$ controls the blending. This is equivalent to finding the point on the geodesic between $p_t$ and $p_s$ under the reverse KL divergence, i.e., minimizing $\text{KL}(q \,||\, p_t)$ and $\text{KL}(q \,||\, p_s)$ in a Pareto sense.

Disagreement-Modulated Blending§

$\lambda_i$ is set based on the disagreement between $p_t$ and $p_s$, measured via Jensen-Shannon divergence:

$$ d_i = \text{JSD}(p_{t,i} \,||\, p_{s,i}) $$

A high $d_i$ indicates the teacher is overconfident or leaking privileged info; thus $\lambda_i$ is reduced to rely more on the student. Concretely,

$$ \lambda_i = \sigma\left( \frac{d_i - \tau}{\gamma} \right) $$

where $\sigma$ is the sigmoid, $\tau$ a threshold, and $\gamma$ a temperature. This allows selective adoption: only tokens where the teacher and student agree (low disagreement) are heavily distilled.

Training Objective§

The student is trained to minimize:

$$ \mathcal{L} = -\sum_i \sum_{y} q_i(y) \log p_{s,i}(y) $$

which is the cross-entropy against the barycenter target. Since $q$ depends on $p_s$ through $\lambda_i$, the gradient flows through both the target and the student parameters.

Theoretical Guarantees§

DemoPSD provably attenuates privileged information leakage because the target distribution is anchored to the student's own distribution wherever $d_i$ is high, preventing the student from encoding answer-dependent shortcuts. Furthermore, because $q$ retains a fraction of the student's original distribution, exploration capacity is preserved, as quantified by the entropy of the student's training distribution.

Implementation Details§

• The teacher uses standard causal masking (privileged access to future tokens), while the student uses the same architecture but without future context.
• Disagreement is computed per token using activations from the final layer before softmax.
• The hyperparameters $\tau$ and $\gamma$ control the sharpness of blending; typical values: $\tau=0.5$, $\gamma=0.1$.

Code Snippet§

import torch
import torch.nn.functional as F

def compute_barycenter_target(teacher_logits, student_logits, tau=0.5, gamma=0.1):
    # teacher_logits, student_logits: (batch, seq_len, vocab)
    teacher_probs = F.softmax(teacher_logits, dim=-1)
    student_probs = F.softmax(student_logits, dim=-1)
    
    # Jensen-Shannon divergence per token
    m = 0.5 * (teacher_probs + student_probs)
    jsd = 0.5 * (F.kl_div(m.log(), teacher_probs, reduction='none').sum(-1) +
                 F.kl_div(m.log(), student_probs, reduction='none').sum(-1))
    
    # Blending coefficient
    lambda_i = torch.sigmoid((jsd - tau) / gamma).unsqueeze(-1)
    
    # Reverse-KL barycenter (geometric mean in probability space)
    barycenter = (teacher_probs ** lambda_i) * (student_probs ** (1 - lambda_i))
    barycenter = barycenter / barycenter.sum(-1, keepdim=True)
    return barycenter

Results§

On SciKnowEval (four scientific QA datasets), DemoPSD outperforms GRPO and SDPO in both in-domain accuracy and cross-domain generalization to GPQA. Training entropy remains higher than baseline OPSD, confirming preserved exploration. Ablations show that fixing $\lambda=0.5$ degrades performance, validating the advantage of adaptive blending.

Conclusion§

DemoPSD offers a principled way to mitigate overfitting and information leakage in self-distillation for LLMs, with theoretical guarantees and empirical gains.