On-policy distillation (OPD)

Contract: separate teacher or reference client and isolated teacher-data exchange. Caller math: Run teacher inference outside the training run (via the frozen second run) and pass teacher logprobs into your distillation loss. There is no external reward — the signal is purely how much the teacher would have preferred the student’s own choices. Whitney mapping: Create a second Whitney run for the teacher — any model, typically a different, larger one than the student — sync its sampler once, and never train or sample it: it only answers logprobs on the student’s own generated tokens. Fail closed: Do not mix teacher and student weight versions on the same sampler without explicit versioning. Reverse KL as log p - log q, i.e. advantage = teacher_logprob - student_logprob, per completion token:

def build_loss(student_sample, teacher_logprobs, cycle, request, ctx):
    prompt_tokens = request["prompt"]["token_ids"]
    sequences = student_sample["sequences"]
    data = []
    for sequence, teacher_seq_logprobs in zip(sequences, teacher_logprobs, strict=True):
        student_logprobs = sequence["logprobs"]
        per_token_advantage = [
            teacher_lp - student_lp
            for student_lp, teacher_lp in zip(student_logprobs, teacher_seq_logprobs, strict=True)
        ]
        data.append(rl_datum(prompt_tokens, sequence["tokens"], student_logprobs, per_token_advantage))
    return forward_backward_payload(data, "importance_sampling")

rl_datum and forward_backward_payload are the shared datum helpers. teacher_logprobs comes from the two-run loop — one call per sequence, no demonstration or feedback prompt prepended, since OPD’s teacher answers on the raw prompt tokens the student itself saw.

Get started

How it works

Cookbooks

Distillation

Coding agents

On-policy distillation (OPD)