Self‑Distilled Reasoning: Preserving Chain‑of‑Thought in Supervised Fine‑Tuning (Amazon Nova)

Exploring self-distilled reasoning for supervised fine-tuning with Amazon Nova

Reported by Amazon (Rushil Anirudh, Shiva Mahalingam, Anupam Dewan), Self‑Distilled Reasoning (SDR) is a practical technique for supervised fine‑tuning (SFT) when your training data lacks chain‑of‑thought (CoT) traces. The idea is simple: have the base model generate reasoning traces, prepend those traces to the supervised targets, and fine‑tune with reasoning mode on. Amazon’s Nova‑family experiments show SDR can prevent severe losses in multi‑step reasoning and often improve target-task accuracy, but it adds cost and needs careful validation.

Quick definitions

  • Supervised Fine‑Tuning (SFT), tuning a model on input→output pairs; here, with or without CoT traces.
  • Chain‑of‑Thought (CoT) : token sequences exposing intermediate reasoning steps before a final answer.
  • Reasoning mode, training/inference configuration that expects and scores intermediate CoT tokens as well as the final answer.
  • Reasoning suppression (catastrophic forgetting) : when SFT on answer‑only data trains a model to skip producing CoT, reducing multi‑step reasoning ability.
  • Self‑Distilled Reasoning (SDR), generate CoT traces from the base model (teacher), prepend them to targets, and fine‑tune the student with reasoning mode on.
  • Model merging, post‑hoc linear interpolation between base and fine‑tuned weights (merge weight 0.0 = base, 1.0 = fully fine‑tuned).

Why naive SFT can make models stop “thinking”

If you fine‑tune on examples that only include inputs and final answers, the loss focuses on producing those final tokens. Over time the student realizes it can hit accuracy without emitting intermediate reasoning tokens, and it learns to skip the CoT it used before fine‑tuning. Amazon saw this in practice. A Nova base model scoring about 70% on a math control benchmark fell to single digits after vanilla SFT on non‑CoT datasets (reported in their results).

This matters when multi‑step logic, robustness, or explainability are important. Losing CoT capability can make models brittle off distribution, weaken general ability, and reduce interpretability.

What SDR does (the recipe)

SDR is straightforward operationally:

  1. Use the base model to generate CoT traces for SFT examples (two styles are common: basic and guided).
  2. Filter traces to remove trivial answer leaks or artifacts.
  3. Prepend the filtered trace to the supervised target and fine‑tune with reasoning mode enabled so the model is trained to emit chain‑of‑thought tokens followed by the final answer.

Amazon provides concrete prompt patterns used for trace generation. Example system/user templates they used include:

COT_SYSTEM: “You are an expert evaluator. Provide a detailed Chain-of-Thought explaining why the given ground truth answer is correct and makes sense.”

Basic‑reasoning prompt text: “Think through this question naturally and show your reasoning: {user_text} Format your response with <thinking> tags for reasoning and <answer> tags for the final answer. Use a consistent reasoning style.”

Guided‑reasoning prompt instruction (excerpt): “Act as if you don’t know the answer yet. Question: {user_text} Answer to arrive at: {assistant_content} … Write natural reasoning that works toward this answer without assuming you know it beforehand.”

Recommended generation settings cited by Amazon for producing CoT traces:

  • temperature: 0.3
  • topP: 0.9
  • topK: 50
  • maxTokens: 4096

Headline experimental findings (as reported by Amazon)

Amazon ran SDR on Nova 2 family models across several benchmarks (MedMCQA, CoCoHD, Invoice‑OCR, LLaVA CoT ablations, ToolACE, GovReport, CaptionGen). Key claims from their tables and ablations:

  • Vanilla SFT can cause dramatic math collapse: reported drops from roughly 70% (base) to single digits (≈6%) on a math control benchmark after naive SFT without CoT traces (reported figures).
  • SDR preserves math control: self‑distilled models retained math performance near base levels (roughly 64-72% across ablations) even when the training data contained little or no native CoT. For example, Amazon reports SDR kept math near 67.9% in a MedMCQA run where vanilla SFT had collapsed math to near 0% (table rows as reported).
  • SDR often improves target accuracy at the same time: Amazon reports cases such as MedMCQA where target accuracy rose from 63.8% (vanilla SFT) to 66.6% (SDR with Nova 2 Lite basic traces), while math stayed preserved (reported numbers).
  • Partial‑reasoning ablation: reducing native reasoning coverage below ~75% caused severe math collapse for vanilla SFT (for example: 66.6% → 21.7% → 3.3% → 2.5% as coverage decreased). With SDR, math stayed in the 64-72% band across coverage levels (reported figures).
  • Model merging can recover reasoning but requires tuning: Amazon’s Table 1 shows a merge‑weight sweep on a LLaVA CoT task where interpolation between base and finetuned checkpoints produced different interactions with reasoning‑enabled inference. Merge weight here is linear interpolation (0.0 = base, 1.0 = fully fine‑tuned). For example, at merge weight 0.5 the target performance (reasoning on) was 60.51% versus 48.80% (reasoning off), a delta of +11.7% (table values as reported).
  • Trace lengths increase costs: SDR prepends long CoT traces. Amazon reported median reasoning token counts in the high hundreds (e.g., SDR prefill median ≈859 tokens; base Nova 2 Lite median ≈1, 254 tokens; after SFT with 100% reasoning data the median dropped to ≈367 tokens). Plan for longer sequences and greater storage and compute (reported medians).

How to interpret “merge weight” and Amazon’s “no merging required” phrasing

Amazon’s merge weight is a linear interpolation between base and fine‑tuned checkpoints (0.0 = base-only, 1.0 = fine‑tuned-only). The phrasing that SDR makes model merging unnecessary reflects the empirical point that SDR‑trained checkpoints can already preserve general reasoning and improve target metrics, so you can keep the fine‑tuned checkpoint (merge weight 1.0) rather than perform post‑hoc interpolation to recover lost capabilities. If you do combine SDR with merging, Amazon recommends smaller merge weights than you might otherwise use, because SDR already nudges the model toward the base reasoning policy.

Practical, actionable guidance (do / don’t / why)

  • Do: Use a same‑family, similar‑strength teacher by default (Amazon recommends Nova 2 Lite for Nova students). Why: It narrows the teacher‑student gap and reduces the risk of introducing artifacts during distillation.
  • Don’t: Plug in a vastly stronger teacher (e.g., Nova 2 Pro) without staged distillation. Why: A large teacher → student capacity gap can hurt distilled performance (the distillation literature documents this effect).
  • Do: Prefill generated CoT traces when your dataset’s native reasoning coverage is under ~75%. Why: Amazon’s ablations show vanilla SFT with low CoT coverage leads to catastrophic math collapse; prefill prevents that.
  • Do: Use guided traces if the base model struggles in your domain, but filter to avoid explicit answer leaks. Why: Guided traces can produce plausible reasoning that connects question and answer more reliably in specialized domains.
  • Do: Expect and budget for longer sequences, measure token counts and SFT wall time on a sample before full rollout. Why: SDR increases sequence length substantially, raising storage and compute.

Minimal validation checks to run during development

Always measure both target metrics and a general reasoning control (math or equivalent) at each checkpoint. Concrete, actionable tests:

  1. Math/control benchmark: run a math‑style reasoning benchmark every checkpoint. Stop rule: if the control metric drops >5 percentage points versus base, pause and investigate (teacher choice, trace filtering, or staged distillation).
  2. Faithfulness probes: apply Early Termination, Paraphrase, and Filler Substitution tests (from recent CoT‑faithfulness work) on a sampled set of traces to check whether the trace causally supports the answer. Stop rule: if faithfulness metrics fall substantially versus base, audit trace quality and filtering.
  3. Human spot checks: sample ~200 auto‑generated traces and inspect for hallucinated justifications, direct answer leakage, or boilerplate rationalizations. Stop rule: more than ~10% of sampled traces have obvious leakage or contradiction, refine prompts/filters.

Where SDR shines, and where it still needs proof

Amazon’s Nova experiments show SDR is a pragmatic way to keep specialized models that still “think.” It fits when your SFT dataset lacks CoT traces, when you want to avoid heavy post‑hoc merging workflows, and when explainability or multi‑step robustness matters.

Open questions and limits to keep in mind:

  • Generality: Amazon’s experiments target Nova models. SDR’s efficacy on other model families (open‑weights or different architectures) is plausible but not yet proven in those reports.
  • Trace faithfulness: model‑generated traces can be post‑hoc rationalizations. Use faithfulness diagnostics and human checks before trusting generated CoTs as supervision.
  • Compute & ops cost: generating and storing long traces at scale takes time and tokens, measure per‑sample cost and SFT wall time on a representative slice of your data.
  • Long‑run continual learning: multi‑round SFT behavior with repeated SDR remains an open empirical question (will traces accumulate artifacts?).

To reproduce quickly, a short checklist

  • Generate CoT traces with your base model using low temperature (0.3), topP 0.9, topK 50, and a high maxTokens limit (Amazon used 4096 for CoT generation).
  • Choose trace style: basic (question → CoT) or guided (question + known answer → CoT); filter traces for direct answer leakage.
  • Prepend filtered CoT traces to training targets and fine‑tune the student with reasoning mode enabled (so the training objective includes CoT tokens). If using LoRA, document rank and optimizer settings, Amazon used LoRA in some experiments but full hyperparameters were not published.
  • Run both target validation and a math/control benchmark per checkpoint; add faithfulness probes (Early Termination, Paraphrase, Filler Substitution) on a sampled trace set.
  • If you plan to merge afterwards, treat merge weight as linear interpolation (0.0 = base, 1.0 = fine‑tuned) and sweep small values if SDR was already applied.

Next experiments worth running

  • Cross‑family replication: run SDR on non‑Nova families (open models) to test generality.
  • Multi‑round continual learning: evaluate whether iterative SDR across many tasks accumulates artifacts or breaks faithfulness.
  • Trace quality sensitivity: compare SDR with human CoT, basic auto CoT, and guided auto CoT across faithfulness and target accuracy.
  • Teacher‑assistant distillation: test staged distillation when using much stronger teachers (to mitigate teacher‑student gaps).

Key takeaways, questions you should be asking (and honest answers)

  • Does supervised fine‑tuning without CoT traces break a model’s reasoning?

    Amazon reports dramatic math‑control drops (from ~70% to ~6% in some vanilla SFT runs). External literature on CoT‑faithfulness and SFT degradation supports the existence of this risk.

  • Can self‑distillation preserve reasoning while improving task performance?

    In Amazon’s Nova experiments, SDR preserved math control (roughly 64-72%) and often improved target accuracy (examples reported, e.g., MedMCQA rising from 63.8% to 66.6%). These results are promising but specific to the Nova family and the reported benchmarks.

  • Which teacher should I use to generate traces?

    Prefer a same‑family, similar‑strength teacher (Amazon recommends Nova 2 Lite for Nova students). Using a much stronger teacher (e.g., Nova 2 Pro) can increase target accuracy but risks degrading general capability unless staged distillation is used.

  • How much of my dataset needs CoT traces before I can skip prefill?

    Amazon’s ablations suggest falling below ~75% native CoT coverage risks catastrophic forgetting with vanilla SFT; they recommend pre‑filling missing traces when coverage is under that threshold.

  • Is SDR a free lunch, zero cost?

    Amazon’s “zero‑cost” phrasing refers to avoiding post‑hoc model merging in their reported runs, not to computational or storage cost. Trace generation and longer SFT sequences do add compute and operational cost.

Final note

Self‑Distilled Reasoning is a practical lever. It regularizes SFT toward producing process (reasoning) as well as outcomes (answers). Amazon’s Nova experiments show it can preserve general reasoning capability and often boost target metrics without heavy post‑hoc merging. That makes SDR worth testing when you care about multi‑step logic, explainability, or robustness, provided you validate faithfulness, measure compute costs, and choose teacher strength carefully.

If you want a one‑page implementation checklist (prompts, generation config, validation scripts, and the minimal hyperparameter sweep to try first), say the word and I’ll draft it tailored to your workload and scale constraints.

3