Reasoning Distillation: Teacher Traces as Training Data
By the end of this lesson you will be able to run a reasoning-capable teacher in thinking mode through your own endpoint, sample several attempts per problem, keep only the traces a program can check, apply the filters that matter and the ones that only look as if they matter, decontaminate the result, and describe accurately what the student you get is and is not able to do. The script that does the filtering ships with this lesson and is reused by the project.
The recipe that produced the small reasoning models
Section titled “The recipe that produced the small reasoning models”The DeepSeek-R1 paper’s central result is about reinforcement learning on a large model, but the sentence that mattered most for everybody with a laptop is the last one in its abstract: “the emergent reasoning patterns exhibited by these large-scale models can be systematically harnessed to guide and enhance the reasoning capabilities of smaller models”.
The mechanism is unglamorous. The large model solves problems and shows its working; the working is kept as text; a small model is fine-tuned on that text. The DeepSeek-R1-Distill-Qwen-7B card says the model was “finetuned with 800k samples curated with DeepSeek-R1”, on a base of Qwen2.5-Math-7B, which the card notes is “originally licensed under Apache 2.0 License”, with the code repository and weights under the MIT licence.
Read that carefully and it is exactly the sequence-level distillation of the previous lesson, with two differences. The completions contain a long chain of working before the answer, and there is a verifier: for a maths problem you can check whether the final number is right, so you can throw away the traces that reached the wrong answer instead of training on them.
That second difference is the whole lesson. Without it you are training a small model to imitate the appearance of reasoning, which it will learn very well, including on the problems the teacher got wrong.
Getting traces out of a local teacher
Section titled “Getting traces out of a local teacher”Qwen3 models carry a thinking mode. The Qwen3-30B-A3B card describes the model generating reasoning
wrapped in <think> and </think> tags before its answer, and the tokeniser makes the tags first-
class: tokenizer_config.json for Qwen3-4B, read on 2026-09-09, gives them ids 151667 and 151668,
alongside 151644 and 151645 for the chat-turn markers.
The sampling settings are not optional. The card gives Temperature 0.6, TopP 0.95, TopK 20 and MinP 0 for thinking mode, against Temperature 0.7, TopP 0.8, TopK 20 and MinP 0 for non-thinking, and warns: “DO NOT use greedy decoding, as it can lead to performance degradation and endless repetitions.” A trace generated at temperature 0 is not a cheaper version of a trace generated properly; on this family it is a different and worse thing.
How you turn thinking on depends on which engine is behind your endpoint. The transformers path
takes a chat-template argument. Through an OpenAI-compatible endpoint, the simplest route the card
documents is the in-conversation switch: the model “supports dynamic switching via /think and
/no_think tags during conversation when thinking is initially enabled”. Put the tag at the end of
the user message and you have controlled the mode without depending on which server flag your engine
spells which way.
Sample several attempts, not one
Section titled “Sample several attempts, not one”For ordinary distillation, one answer per prompt is enough. For reasoning distillation it is not, because you are about to throw most of them away.
A teacher that solves a class of problem most of the time still fails some of the time, and which
ones it fails on varies between samples. Ask four times at a proper temperature and you get four
independent attempts; keep the correct ones. That is rejection sampling, and it is the reason
the generation script in this part takes a --samples argument and refuses to accept more than one
sample at temperature zero, where every attempt would be identical.
The cost is linear: four samples is four times the generation. The benefit is that your dataset contains only correct working, which is worth far more than four times as many unfiltered traces.
The verifier is the point
Section titled “The verifier is the point”Rejection sampling with a verifier
- Problems with known answersGenerated together with their answers, or taken from a set whose labels you trust. The answer must exist before the question is asked.
- Teacher, thinking mode, several samplesFour attempts per problem at the card's recommended settings. Not greedy.
- Extract the final answerOne of three documented markers: the #### form, a boxed expression, or an "Answer:" line. Take the last match, not the first.
- VerifyA program compares the extracted number with the computed answer. No model is asked for an opinion.
- Filter what survivedLength, one closed thinking block, working actually shown, no repetition loop.
- DecontaminateDrop anything overlapping the set you will report the result on.
- Supervised fine-tuneOrdinary SFT on prompt and trace, completion-only loss. Part 13's recipe, unchanged.
Part 14’s lesson on reward functions made the case for a reward a program can compute rather than a model’s opinion, and everything it said applies here with one word changed: a verifier is a reward function used as a filter instead of as a gradient. The same answer-extraction problem, the same numeric tolerance question, the same requirement that it be exactly repeatable.
Two details from that lesson that carry over and that people get wrong:
Take the last match, not the first. A trace that says “so the total is 42, but wait, that double-counts the broken ones, Answer: 38” contains two numbers, and only the second is its answer. Extracting the first would mark a correct trace wrong and a wrong trace correct with roughly equal frequency.
Prefer the most specific marker that appears. The literature uses at least three conventions:
GSM8K’s own solutions end with a #### line, the maths papers use a boxed expression, and a prompt
that asks for an Answer: line gets one. Try them in that order and stop at the first that matches
anywhere in the text.
Filtering: the two that matter and the one that looks like it does
Section titled “Filtering: the two that matter and the one that looks like it does”A trace that reached the right answer is not automatically a trace worth training on.
One closed thinking block. A trace with an opening tag and no closing tag will teach the student to open a block and never leave it, which is a failure mode you will then have to debug in production. A trace with three blocks teaches a shape the chat template does not produce. Both are cheap to detect and worth dropping.
Working actually shown. A thinking block containing the word “obvious” followed by the correct answer is a trace that asserted rather than reasoned. Train on enough of those and the student learns to emit a short block and a guess, which looks exactly like reasoning until the problems get harder. The script in this lesson checks the length of the working separately from the length of the whole trace, and reports that rejection under its own name rather than as “too short”, because the reason is the interesting part.
Repetition loops. A model that falls into a loop produces long, fluent, useless text that often sits inside your length budget. Counting repeated eight-word sequences catches it; a length filter does not.
The filter that looks more important than it is: maximum length. It is worth having, because a
runaway trace will blow up your sequence length and your memory budget. But choosing the shortest
correct trace per problem, which is what the --prefer shortest default does, has a bigger effect
and a clearer justification: a short correct trace is more likely to be a direct route than a lucky
recovery from a wrong turn, and the student has limited capacity for what it can learn to imitate.
Decontamination is not optional here
Section titled “Decontamination is not optional here”Reasoning problems are generated from templates, scraped from the same handful of public sets, or both. The chance that a problem in your training set is also in the set you will report on is higher than for any other kind of data in this course.
The rule is the one Part 13’s dataset lesson set out and this part’s scripts implement: check exact matches, check n-gram containment, and read the pairs the containment check flags rather than trusting a threshold. An exact match is contamination and the fix is deletion. A high containment score is a judgement you have to make by reading the two texts. This part’s challenge is built around what happens when nobody made it.
The training step is ordinary
Section titled “The training step is ordinary”Once the traces are filtered, there is nothing distillation-specific left. The dataset is prompt and
completion, the trainer is the SFTTrainer from Part 13, the loss is computed on the completion
only, and the export path is the one Part 13’s merging lesson describes.
Two settings that differ from a format fine-tune:
- Sequence length. A reasoning trace is several times longer than an answer. If your maximum length truncates the trace, you are training the student on working that stops in the middle and never reaches an answer, which teaches it to do the same. Measure the length distribution of your filtered traces before choosing the setting.
- Fewer epochs. These datasets are small after filtering and the completions are long, so the number of tokens per epoch is larger than it looks. Two epochs with an evaluation after each is a better starting point than three.
RunnableAll tracks
"""Keep the teacher's reasoning traces that reach the right answer, and throw the rest away.
Purpose: reasoning distillation's filter. The teacher was asked each problem several times; this reads all of those attempts, extracts the final answer from each, checks it against the answer the problem generator computed, and keeps at most a stated number of correct traces per problem. Wrong traces are dropped, traces that reach the right answer by an unusable route are dropped, and what survives is written in the shape train-student.py reads. The verifier is the same idea as Part 14's reward functions: a program decides, not a model.Platform: all (standard library only; no model, no accelerator, no network)Minimum memory: 12 GB nominally, and far less in practice: this is text in memoryAssumes: Python 3.10 or newer. The input is the JSON Lines file written by generate-teacher-data.py over a problem file that carries an "answer" field, such as seeds/maths.jsonl from make-seed-prompts.py. distillog.py sits next to this file.
Usage: python3 rejection-sample.py --raw raw/maths-traces.jsonl --out-dir . \\ --keep-per-problem 1 --labbook labbook.md python3 rejection-sample.py --raw raw/maths-traces.jsonl --out-dir . \\ --keep-per-problem 2 --max-words 500 --prefer shortest --tasks my-tasks.json
Written under --out-dir: data-reasoning/{train,valid}.jsonl TRL conversational prompt-completion data-reasoning-mlx/{train,valid}.jsonl mlx-lm completions rejection-report.json per-problem attempts, correct, kept"""
from __future__ import annotations
import argparseimport jsonimport randomimport refrom collections import Counterfrom pathlib import Pathfrom typing import Any
import distillog
WORD_RE = re.compile(r"[a-z0-9]+")
# Three ways a model marks its final answer, in the order this course prefers them,# the same order Part 14's reward library uses. GSM8K's own solutions end with# "#### 18", the maths literature uses \boxed{18}, and the course's own prompt asks# for "Answer: 18".FINAL_PATTERNS = ( re.compile(r"####\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)"), re.compile(r"\\boxed\{\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)\s*\}"), re.compile(r"(?i)\banswer\s*[:=]\s*\$?(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),)THINK_BLOCK = re.compile(r"<think>(.*?)</think>", re.DOTALL)UNCLOSED_THINK = re.compile(r"<think>(?!.*</think>)", re.DOTALL)
def extract_final_answer(text: str) -> float | None: """The last match of the most specific pattern that appears at all.
Taking the last match matters: a trace that says "so the total is 42, but wait, Answer: 47" has two numbers and only the second is its answer. """ for pattern in FINAL_PATTERNS: matches = pattern.findall(text) if matches: try: return float(matches[-1].replace(",", "")) except ValueError: continue return None
def is_correct(predicted: float | None, expected: float, tolerance: float = 1e-6) -> bool: return predicted is not None and abs(predicted - expected) <= tolerance
def normalise(text: str) -> list[str]: return WORD_RE.findall(text.lower())
def ngrams(words: list[str], n: int) -> set[tuple[str, ...]]: if len(words) < n: return {tuple(words)} if words else set() return {tuple(words[i:i + n]) for i in range(len(words) - n + 1)}
def containment(a: set[tuple[str, ...]], b: set[tuple[str, ...]]) -> float: return len(a & b) / len(a) if a else 0.0
def trace_problem(text: str, args: argparse.Namespace) -> str | None: """Reasons to reject a trace that nonetheless reached the right answer.""" if UNCLOSED_THINK.search(text): return "unclosed-thinking-block" blocks = THINK_BLOCK.findall(text) if args.require_thinking and not blocks: return "no-thinking-block" if len(blocks) > 1: return "several-thinking-blocks" if blocks: working = normalise(blocks[0]) if len(working) < args.min_working_words: # A trace whose thinking block is two words long reached the answer by # asserting it. Training on that teaches assertion, not reasoning, and # it is checked before the length filter so the reason names the fault # rather than the symptom. return "no-working-shown" words = normalise(text) if len(words) < args.min_words: return "too-short" if len(words) > args.max_words: return "too-long" counts = Counter(tuple(words[i:i + 8]) for i in range(max(0, len(words) - 7))) if counts and counts.most_common(1)[0][1] >= 3: return "repetition-loop" return None
def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--raw", required=True, help="traces from generate-teacher-data.py") parser.add_argument("--out-dir", default=".") parser.add_argument("--keep-per-problem", type=int, default=1, help="how many correct traces to keep for each problem") parser.add_argument("--prefer", choices=["shortest", "longest", "first"], default="shortest", help="which correct traces to keep when there are more than needed") parser.add_argument("--min-words", type=int, default=20) parser.add_argument("--max-words", type=int, default=600) parser.add_argument("--min-working-words", type=int, default=15) parser.add_argument("--require-thinking", action="store_true", default=True, help="reject a trace with no <think> block (default)") parser.add_argument("--allow-no-thinking", dest="require_thinking", action="store_false") parser.add_argument("--tasks", default=None, help="evaluation set to decontaminate against") parser.add_argument("--contamination", type=float, default=0.5) parser.add_argument("--n", type=int, default=13) parser.add_argument("--valid-fraction", type=float, default=0.1) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--labbook", default=None) args = parser.parse_args()
raw_path = Path(args.raw) rows = [json.loads(line) for line in raw_path.read_text(encoding="utf-8").splitlines() if line.strip()] missing_answer = [r for r in rows if "answer" not in r] if missing_answer: raise SystemExit( f"{len(missing_answer)} trace(s) carry no 'answer' field, so nothing can verify them. " "Generate from a problem file that has one, such as seeds/maths.jsonl." ) print(f"read {len(rows)} trace(s) from {raw_path}")
by_problem: dict[str, list[dict]] = {} for row in rows: by_problem.setdefault(row["id"], []).append(row)
rejected: Counter[str] = Counter() kept: list[dict] = [] per_problem_stats = []
for pid, attempts in sorted(by_problem.items()): expected = float(attempts[0]["answer"]) correct = [] for attempt in attempts: completion = (attempt.get("completion") or "").strip() predicted = extract_final_answer(completion) if predicted is None: rejected["no-answer-found"] += 1 continue if not is_correct(predicted, expected): rejected["wrong-answer"] += 1 continue reason = trace_problem(completion, args) if reason: rejected[reason] += 1 continue correct.append({**attempt, "completion": completion, "words": len(normalise(completion))})
if args.prefer == "shortest": correct.sort(key=lambda a: a["words"]) elif args.prefer == "longest": correct.sort(key=lambda a: -a["words"]) chosen = correct[:args.keep_per_problem] kept.extend(chosen) per_problem_stats.append({"id": pid, "attempts": len(attempts), "correct": len(correct), "kept": len(chosen)})
solved = sum(1 for s in per_problem_stats if s["kept"]) print(f"problems: {len(by_problem)}, at least one usable trace for {solved}") print(f"traces kept: {len(kept)}") for reason, count in rejected.most_common(): print(f" rejected {count:>5} {reason}")
contaminated = 0 if args.tasks: spec = json.loads(Path(args.tasks).read_text(encoding="utf-8")) eval_grams = [ngrams(normalise(t.get("prompt", "")), args.n) for t in spec.get("tasks", [])] clean = [] for row in kept: grams = ngrams(normalise(f"{row['prompt']}\n{row['completion']}"), args.n) if any(containment(e, grams) >= args.contamination for e in eval_grams if e): contaminated += 1 continue clean.append(row) kept = clean print(f"decontamination: dropped {contaminated} trace(s) overlapping the evaluation set") else: print("decontamination: SKIPPED, no --tasks given.")
if not kept: raise SystemExit("no trace survived; look at rejection-report.json before loosening a threshold")
rng = random.Random(args.seed) rng.shuffle(kept) cut = max(1, int(len(kept) * args.valid_fraction)) valid, train = kept[:cut], kept[cut:]
out = Path(args.out_dir) written: dict[str, Any] = {} for name, part in (("train", train), ("valid", valid)): trl_path = out / "data-reasoning" / f"{name}.jsonl" trl_path.parent.mkdir(parents=True, exist_ok=True) with trl_path.open("w", encoding="utf-8") as handle: for row in part: handle.write(json.dumps({ "prompt": [{"role": "user", "content": row["prompt"]}], "completion": [{"role": "assistant", "content": row["completion"]}], }, ensure_ascii=False) + "\n") mlx_path = out / "data-reasoning-mlx" / f"{name}.jsonl" mlx_path.parent.mkdir(parents=True, exist_ok=True) with mlx_path.open("w", encoding="utf-8") as handle: for row in part: handle.write(json.dumps({"prompt": row["prompt"], "completion": row["completion"]}, ensure_ascii=False) + "\n") written[name] = {"count": len(part), "sha256": distillog.file_sha256(trl_path)}
report = { "raw": str(raw_path), "raw_sha256": distillog.file_sha256(raw_path), "traces": len(rows), "problems": len(by_problem), "problems_with_a_usable_trace": solved, "kept": len(kept), "contaminated_dropped": contaminated, "rejected_by_reason": dict(rejected.most_common()), "per_problem": per_problem_stats, "thresholds": { "keep_per_problem": args.keep_per_problem, "prefer": args.prefer, "min_words": args.min_words, "max_words": args.max_words, "min_working_words": args.min_working_words, "require_thinking": args.require_thinking, }, "splits": written, } (out / "rejection-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
print(f"\ntrain {len(train)} valid {len(valid)}") print(f"report written to {out / 'rejection-report.json'}") print("The problems with no usable trace are the interesting ones: they are the " "problems the teacher could not solve, and the student will not learn them here.")
if args.labbook: rec = distillog.record( labbook=args.labbook, lab="part-15/rejection-sample", stage="reject", teacher=(rows[0].get("teacher") if rows else None), student=None, dataset={ "raw": str(raw_path), "raw_sha256": report["raw_sha256"], "traces": len(rows), "problems": len(by_problem), "problems_with_a_usable_trace": solved, "kept": len(kept), "train": len(train), "valid": len(valid), "train_sha256": written["train"]["sha256"], "tasks": args.tasks, "contaminated_dropped": contaminated, }, hyperparameters=report["thresholds"], seed=args.seed, scores={"rejected_by_reason": report["rejected_by_reason"], "solve_rate": round(solved / len(by_problem), 3) if by_problem else None}, config_path=__file__, notes=None if args.tasks else "decontamination skipped: no evaluation set supplied", ) print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")
if __name__ == "__main__": main()The script takes the raw traces, verifies them, applies every filter above, decontaminates against a task file, splits, and writes both the TRL and the mlx-lm dataset shapes. Its report file lists the rejections by reason and the per-problem attempt counts, which is the thing to read: the problems where no attempt survived are the problems your teacher cannot do, and the student will not learn them from this dataset however long you train.
What the student can and cannot do afterwards
Section titled “What the student can and cannot do afterwards”This is where the honesty has to be, because reasoning distillation produces the most impressive demonstration and the most misleading impression in the whole part.
What it does. The student produces working before its answer, in the format the traces used, and on problems of the kind it was trained on it is more likely to reach the right answer than the base model was. That is a real, measurable gain and it is why the technique is used.
What it does not do. It does not give the student the teacher’s capability on problems the teacher solves and the student’s capacity cannot support. It does not generalise reliably to a different kind of reasoning: a student distilled on arithmetic word problems has learned arithmetic word problems, and the transfer to calendar puzzles or code reasoning is a measurement you have to make rather than an assumption you can carry. Part 14’s reality check makes exactly this measurement for reinforcement learning, with a held-out set of the same kind and a held-out set of a different kind, and the same design is the right one here.
And it does not make the student’s working true. The trace is text the student learned to produce; it is not a record of a computation the student performed. A student can emit five correct-looking lines and a wrong answer, or five wrong lines and a right answer, and both happen. Where the answer matters, verify it in your own code, which is what the harness in Part 10 and the verifier in this lesson both do.
Verify the answer independently of the explanation
Section titled “Verify the answer independently of the explanation”A long trace can contain a correct final answer reached through invalid intermediate steps, or a polished explanation of a wrong answer. Filter on an independent verifier where possible and inspect representative traces, especially near acceptance thresholds. Style and apparent thoroughness are weak correctness signals.
Keep the prompt, trace and extracted final answer as separate fields during data preparation. Define which fields the student must emit at deployment. If training rewards long explanations but the application has a short output budget, the model may spend its allowance before reaching the answer. Evaluate answer correctness and completion length together.
Separate source problems before generating variants. The same algebra exercise with changed names is not a fully independent held-out example. Include new problem structures and failure cases outside the teacher’s familiar prompt distribution. A successful distillation result supports a claim about the tested task family and budget. It does not establish that generated explanations reveal the model’s internal causal reasoning, or that longer traces reliably improve unrelated tasks.
Reasoning distillation is sequence-level distillation with a verifier bolted onto the filter. Put a thinking-capable teacher into thinking mode at the sampling settings its card gives, never greedy; sample several attempts per problem because you are about to discard most of them; extract the final answer with the last match of the most specific marker present; verify it against an answer that existed before the question was asked. Then filter for one closed thinking block, for working actually shown, and against repetition loops, preferring the shortest correct trace. Decontaminate against the set you will report on. The training step is Part 13’s, with a longer sequence length and fewer epochs. The resulting student produces working and does better on that family of problems, which is worth having and is not the same as having acquired the teacher’s reasoning.
Check your understanding
Sources for this lesson
6 verified · checked 2026-09-09
- 01DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning (arXiv:2501.12948)§ Abstractarxiv.org/abs/2501.129482026-09-09
- 02DeepSeek-R1-Distill-Qwen-7B model card§ Model summary; base model; licencehuggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B2026-09-09
- 03Qwen3-30B-A3B model card§ Switching between thinking and non-thinking mode; best practiceshuggingface.co/Qwen/Qwen3-30B-A3B2026-09-09
- 04Qwen3-8B model card§ Model overview; best practiceshuggingface.co/Qwen/Qwen3-8B2026-09-09
- 05Qwen3-4B tokenizer_config.jsonhuggingface.co/Qwen/Qwen3-4B/raw/main/tokenizer_config.json2026-09-09
- 06TRL documentation — SFT Trainer§ Expected dataset type and format; Train on completion onlyhuggingface.co/docs/trl/en/sft_trainer2026-09-09
Every technical claim on this page was checked against the official documentation of the tool, vendor or model publisher on the date shown, at the version pinned for the course. Where the course disagrees with folklore, the source is how you can tell which one to trust.