Skip to main content
José David Baena

On this page

Distilled DPO: Training a Small Model on Its Own Mistakes

By José David Baena

Published on
Updated /13 mins read

Your distillation pipeline looks finished. A smaller student, trained on a strong teacher's transcripts, clears every held-out eval you wrote. Then it ships, meets a prompt that's just slightly off from anything in its training set, and produces a malformed JSON field or a flat refusal the teacher would never have generated — and every token after that first bad one is now conditioned on a mistake the training data never contained.

That gap has a name, and a real distillation project already ran into the two halves of it at scale.

DeepSeek-R1 shows why supervision has two separate jobs

DeepSeek-R1's authors used large-scale reinforcement learning to discover reasoning behavior — verification, reflection, strategy adaptation — in a large model, then used that model's reasoning traces to fine-tune smaller Qwen and Llama students. They reported pass@1 of 55.5% on AIME 2024 for the distilled 7B student, and 72.6% on AIME 2024, 94.3% on MATH-500, and 57.2% on LiveCodeBench for the 32B student (DeepSeek-R1). Those numbers belong to named benchmarks and settings — they don't establish that reasoning transferred faithfully, or that it transferred safely. What they do establish is a sharper distinction: behavior discovery (finding a strategy) is a different job from behavior transfer (teaching another model to reproduce it), and the paper's own comparison found that direct distillation from DeepSeek-R1 beat applying RL directly to a smaller Qwen2.5-32B base. A large teacher can find a productive trajectory a smaller model struggles to discover from reward alone; once found, that trajectory becomes data the student can practice without repeating the search.

But distillation via static transcripts has a ceiling, and it's a documented one. Gudibande et al. found that models imitating a proprietary LLM's outputs on limited transcripts tend to mimic its style — confident phrasing, formatting conventions — without closing the underlying capability gap on tasks the transcripts didn't cover (The False Promise of Imitating Proprietary LLMs). Imitation teaches the vocabulary of the task. It doesn't teach recovery from a state the transcripts never modeled.

Off-policy transcripts can't teach recovery from your own mistakes

Training data sampled from the teacher and a fixed curriculum is off-policy with respect to the student that will actually get deployed — the transcripts describe someone else's trajectories, not this student's. That's not automatically bad; off-policy data is stable, reproducible, and cheap to build. The problem shows up the moment the student generates a prefix its training data never contained. If a transcript teaches "check refund status before promising a date" and the student begins, at inference time, "your refund will arrive tomorrow" — every later token is now conditioned on an unsupported promise the teacher never modeled. This is exposure mismatch, and Google DeepMind's Generalized Knowledge Distillation (GKD) work addresses it directly by training on the student's own generated sequences with teacher feedback, rather than only on fixed transcripts (GKD: On-Policy Distillation of Language Models).

The fix, in operational terms: run the current student checkpoint against approved prompts, flag the outputs that are wrong or confidently uncertain, get each one independently corrected, and turn the pair into training signal. Four things matter about how you build that loop:

  • Bound it. A global row budget and a per-slice cap stop one noisy failure category from consuming the whole release. An unbounded loop just chases production noise and memorizes one checkpoint's idiosyncratic errors.
  • Keep it current. "Current" is load-bearing — once training changes the policy, its failure distribution changes too. A rollout dataset is a versioned observation of one student checkpoint, not a timeless set of hard prompts. Retrain, and last week's harvested failures describe a model that no longer exists.
  • Isolate consent and holdout data before you read anything. A collector that can see hidden eval prompts converts evaluation into training, silently.
  • Don't invent a fixed on-policy fraction. If you mix stable off-policy data back in to protect general capability, the split is a release parameter you measure, not a constant you import — more on this below.

Here's a minimal version of the failure-detection side of that loop. It mirrors the shape of the synthetic customer-support project used to teach these ideas in the source material — a fictional refund/payment-security domain, not a real company or real customer data:

from pydantic import BaseModel, Field
import json
import time
 
class PreferencePair(BaseModel):
    pair_id: str
    prompt: str
    chosen: str = Field(..., description="Correction (preferred)")     # ①
    rejected: str = Field(..., description="Student's own failure")    # ②
    failure_category: str                                              # ③
    student_checkpoint: str
    timestamp: float = Field(default_factory=time.time)
 
 
def inspect(prompt: str, response: str) -> tuple[bool, str | None]:
    if "json" in prompt.lower():
        try:
            json.loads(response)
        except json.JSONDecodeError:
            return True, "json_syntax_error"          # ④
 
    refusals = ("i cannot fulfill", "i'm sorry, but i am unable")
    if any(p in response.lower() for p in refusals) and "escalation" not in prompt.lower():
        return True, "over_refusal_error"             # ⑤
 
    if len(response.strip()) < 10:
        return True, "stub_generation_error"
 
    return False, None

① chosen is never a teacher's independent answer to a fresh prompt — it's a correction of this specific failure. ② rejected is the actual student output, not a synthetic wrong answer, which is what makes it a genuine hard negative. ③ Tagging the category lets you check later whether one failure type is quietly dominating your training budget. ④ This check only proves the JSON parses — it says nothing about whether the fields inside it are correct. A response can be syntactically valid and still wrong. ⑤ Refusal detection is the most consequential and highest-risk rule here: get the keyword list wrong and you'll either miss real over-refusals or start "correcting" a model out of a refusal it should have kept.

That last caveat isn't hypothetical. XSTest was built around exactly this tension — 250 safe prompts a well-calibrated model shouldn't refuse, paired with 200 unsafe contrast prompts it generally should (XSTest). A detector tuned only to catch json_syntax_error will happily harvest a case where the correct behavior was to keep refusing, and DPO will faithfully train the student to stop.

With detection in place, harvesting one failure is mechanical:

pairs: list[PreferencePair] = []
for i, prompt in enumerate(prompts, start=1):
    response = student.generate(prompt)                    # ①
    failed, category = inspect(prompt, response)
    if not failed:
        continue                                            # ②
    correction = teacher.correct(category)
    pairs.append(PreferencePair(
        pair_id=f"pair-{i:03d}",
        prompt=prompt,
        chosen=correction,
        rejected=response,
        failure_category=category,
        student_checkpoint="weak-fixture-v1",
    ))

① student and teacher here stand in for real model calls — this is illustrative wiring, not a benchmark. ② Passing rollouts never enter the DPO set at all. DPO only ever sees things the current checkpoint actually got wrong, which is the entire point of on-policy collection: the failure distribution belongs to this checkpoint version.

Zephyr is a real, published example of exactly this shape at production scale: its authors ranked candidate outputs with a teacher model and trained a 7B chat model with DPO, reporting that it beat Llama2-Chat-70B on MT-Bench without any human annotation (Zephyr: Direct Distillation of LM Alignment). That result is evidence about their specific pipeline and benchmark — it isn't proof that AI-ranked preferences automatically encode your product's factual, legal, or safety requirements. Someone still has to be authorized to pick the winner.

The DPO objective, in the terms this pipeline produces

For prompt x, chosen response y_w, rejected response y_l, trainable policy pi_theta, and frozen reference policy pi_ref, DPO minimizes:

L_DPO = -E[(x, y_w, y_l) ~ D][
  log sigmoid(
      beta * log(pi_theta(y_w|x) / pi_ref(y_w|x))
    - beta * log(pi_theta(y_l|x) / pi_ref(y_l|x))
  )
]

This matches the objective in the original paper (Direct Preference Optimization). pi_ref is a frozen reference — normally the SFT checkpoint you're starting DPO from, so the constraint is measured against a competent baseline, not a random init. beta controls how far training is allowed to push the policy away from that reference: lower it and DPO chases the preference margin harder, at the cost of drifting further from the reference (and, as the failure modes below show, at real risk of over-optimizing); raise it and training barely moves. The paper's main experiments use beta = 0.1, and it remains the documented default in Hugging Face's TRL library — a starting point, not a law, and worth re-verifying against whatever TRL version you actually have installed (TRL DPOTrainer docs).

One thing the equation cannot do: detect a mislabeled pair. It optimizes a wrong preference exactly as faithfully as a right one, which is why the detector quality above is not a side concern — it's the whole signal.

Loading visualization…

Training with TRL turns the JSONL into an update

Once the harvested pairs exist as JSONL with prompt, chosen, and rejected fields — the schema TRL's DPOTrainer expects, and one of the more stable parts of its API across recent releases — training is mostly mechanical:

from datasets import load_dataset
from trl import DPOConfig, DPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
 
model = AutoModelForCausalLM.from_pretrained("support-distill/student-sft-v0")
ref_model = AutoModelForCausalLM.from_pretrained("support-distill/student-sft-v0")  # ①
tokenizer = AutoTokenizer.from_pretrained("support-distill/student-sft-v0")
 
dataset = load_dataset("json", data_files="distilled_dpo_pairs.jsonl")["train"]     # ②
 
config = DPOConfig(
    beta=0.1,                          # ③
    per_device_train_batch_size=4,
    learning_rate=5e-6,
    output_dir="student-dpo-v1",
)
 
trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
    processing_class=tokenizer,
)
trainer.train()

① The reference policy has to be the same checkpoint DPO starts from — the SFT model before this stage, frozen. Some setups let the trainer manage this for you instead, especially with a PEFT adapter, where "reference" just means the adapter disabled — check the docs pinned to your version for which path it takes by default. ② This is the exact schema the harvester above produced. ③ Start from the paper's own default and move only after you've measured the effect on chosen-response log-probability and your eval suite, not before.

Argument names have moved between TRL releases — tokenizer became processing_class in more recent versions, for instance — so verify against the docs matching whatever version your requirements.txt actually installs before copying this as-is.

Evaluation gates before you trust the update

  • Chosen-response log-probability under pi_ref, not just the preference margin — see the DPO-Positive failure mode below for why this matters more than it looks.
  • A held-out win rate against the pre-DPO checkpoint, judged pairwise, on prompts the harvester never saw.
  • The same task-accuracy suite you ran after SFT, so a DPO pass that fixes one failure category can't quietly regress an untouched one.
  • Unsafe compliance and over-refusal, tracked separately, never averaged. XSTest's whole design exists because pushing one of these down by brute force tends to push the other up.
  • Response length delta, given below.

What could go wrong

  • DPO can decrease both chosen and rejected log-probabilities at once. Pal et al. document this degenerate case — DPO-Positive — especially when the edit distance between chosen and rejected is small (DPO-Positive). That's a direct hit on this exact pipeline: harvested pairs are, by construction, a small correction to one specific mistake, not two wildly different responses. Watch chosen log-probability directly, not just the margin.
  • DPO can win the preference margin by getting longer, not better. Park et al. show that direct preference optimization can exploit response length as a proxy for quality (Disentangling Length from Quality in DPO). Check the length delta between chosen and rejected in your harvested set before training, and check it again in the model's outputs after.
  • Direct alignment algorithms can over-optimize a proxy signal the same way reward models do. Rafailov et al.'s later work on scaling laws for reward overoptimization applies the concept directly to DPO-style training: held-out preference accuracy can keep climbing while genuine quality falls (Scaling laws for reward model overoptimization in direct alignment algorithms). Put a stopping rule on this — don't just run a fixed number of epochs and trust the loss curve.
  • A shallow failure detector bounds your entire signal. A JSON parser can only tell you a response is syntactically valid, never whether it's semantically correct. Whatever your detector misses, DPO never sees, and whatever it flags wrong, DPO faithfully learns to avoid the right behavior instead.
  • Category imbalance skews correction toward whatever's easiest to detect. If json_syntax_error dominates the harvested set because it's trivial to check, DPO will over-correct that one failure mode and under-correct rarer, harder-to-detect ones like a factually wrong but well-formed answer.

The mixing-ratio heuristic has a boundary, not a number

Some pipelines fold a minority of stable, off-policy data back into on-policy training specifically to keep the student's general capability from narrowing around its own recent failure modes. There's no published universal ratio for this. Formally, if D_T is the fixed teacher-transcript distribution, D_S is the current student's own generated states, and lambda is the on-policy fraction, training draws from:

D_mix = (1 - lambda) * D_T + lambda * D_S,  0 <= lambda <= 1

Raising lambda gives the student more practice on states it will actually encounter in production. It also makes training more dependent on one checkpoint's particular errors — the higher you push it, the more the training signal is shaped by whatever this specific model happens to get wrong today. There's no fixed split that generalizes across projects, and a claimed "70/30" recipe from a blog post is not evidence for your model. Treat lambda as a release parameter: record it, and re-measure the effect on slices you didn't specifically train against before you call the run done.

Key takeaways

  • Off-policy transcripts teach imitation of the teacher's trajectories, not recovery from the student's own states — DeepSeek-R1's discovery/transfer split and GKD's on-policy correction both point at the same seam from different angles.
  • Bound the harvesting loop: global and per-slice budgets, holdout isolation, and checkpoint versioning. An unbounded loop chases traffic noise and memorizes one model's current quirks.
  • The DPO paper hands you a formula, not a guarantee — verify beta against the pinned TRL docs, and watch chosen-response log-probability, not just the preference margin.
  • Evaluate your failure detector as critically as the training loss. A shallow check is a ceiling on your entire signal quality.
  • There's no fixed on-policy/off-policy ratio to import. Treat it as a measured release parameter with an explicit boundary, not a constant.

What's next

The behavior release this loop produces — corrections, preference pairs, and their safety gates — is exactly the kind of artifact the compression decision matrix in this series treats as a separate, required parent: quantize or prune the resulting student, and you re-run these same quality and safety checks on the exported artifact, not just the training checkpoint.

This post adapts material from Chapter 5, "Teaching Thought and Behavior," in Distilled: The Engineering of Small, Fast, Cheap AI Models. The rest of the series is indexed at the series page.

Sources

Distillation and on-policy transfer

Preference optimization

Safety evaluation

Tooling

Share this post

HNPost to Hacker News

Follow future work

Follow public article updates through RSS. Intentionally unlisted posts stay out of the feed.

Keep reading