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

In the rapidly evolving landscape of artificial intelligence, fine-tuning large language models (LLMs) to excel in specific domains is a critical capability. Supervised Fine-Tuning (SFT) stands as a foundational technique, allowing models to adapt to custom datasets and tasks. However, a significant hurdle often arises: the impracticality and prohibitive cost of generating high-quality Chain-of-Thought (CoT) reasoning traces, which are step-by-step explanations of how a model arrives at an answer, for training data. While reasoning is a cornerstone capability of advanced models like the Amazon Nova 2 family, crucial for tackling complex problems in areas such as coding and mathematics, its absence in SFT datasets has historically forced developers to train models using only inputs and outputs, often leading to a degradation of the model’s inherent reasoning prowess and a phenomenon known as catastrophic forgetting.

Amazon Web Services (AWS) has introduced Self-Distilled Reasoning (SDR) as a groundbreaking solution to this challenge. This innovative approach allows developers to harness the full power of "thinking models" for their specific domains through SFT and Reinforcement Fine-Tuning (RFT), even when golden CoT traces are unavailable. SDR elegantly addresses the "reasoning suppression problem" by leveraging the base Amazon Nova 2 Lite model’s own chain of thought as a surrogate for missing reasoning data. Through extensive experimentation across diverse benchmarks, SDR has demonstrated its capacity not only to improve target performance but also to effectively mitigate catastrophic forgetting, aligning with a growing body of research highlighting the importance of self-distillation in model training.

The Foundational Challenge: Reasoning Suppression in SFT

Supervised Fine-Tuning involves training a pre-trained model on a smaller, task-specific dataset to adapt its behavior. While powerful, SFT on datasets lacking explicit reasoning traces presents a critical dilemma. When models are trained without intermediate reasoning steps, the training objective—which calculates loss on both reasoning and output tokens—effectively penalizes any generated tokens that do not directly contribute to the final output. This inadvertently teaches the model to bypass its internal reasoning mechanisms, even when reasoning mode is explicitly activated during inference. This phenomenon, termed the reasoning suppression problem, exemplifies "shortcut learning," where models prioritize superficial correlations over robust, generalized reasoning patterns. The consequences are stark: a model might lose its ability to generate coherent reasoning traces during inference, severely impacting its performance on complex problems.

The impact of reasoning is quantifiable and significant. Experiments have consistently shown that enabling reasoning during both training and inference yields substantial performance benefits. For instance, on the LLaVA CoT dataset, activating reasoning during both phases led to performance gains of up to 17.3% (from 47.90% to 65.17%) for a full merge weight, underscoring the critical role of these intermediate thought processes. Without proper intervention, domain-specific SFT can severely compromise a model’s general reasoning abilities, creating a significant impediment to developing truly intelligent and versatile AI systems.

Early Mitigations: The Role and Limits of Model Merging

Prior to SDR, model merging served as a primary strategy to address the reasoning suppression issue and recover lost capabilities. This technique involves creating a weighted interpolation between a fine-tuned model checkpoint and its original base model. The rationale is that by blending the fine-tuned model’s domain-specific knowledge with the base model’s general intelligence, previously acquired skills—such as mathematical or coding capabilities—could be restored. This approach proved somewhat effective in recovering reasoning ability, correlating with the restoration of general performance metrics.

However, model merging is not without its drawbacks. It often entails a trade-off: gains obtained on the SFT checkpoint can be sacrificed due to catastrophic forgetting. This means that while general skills might be partially recovered, target performance on the fine-tuned domain could suffer. Furthermore, finding the optimal merge weight to balance these competing objectives requires careful tuning and adds a post-hoc step to the training pipeline, complicating the customization workflow. For instance, while merging could recover math performance, it might not achieve the same target performance as a pure SFT model without merging, highlighting a persistent tension between specialization and generalization. Benchmarks like ToolACE (tool-calling), CoCoHD (long-context structured extraction), GovReport (abstractive summarization), Invoice-OCR (multimodal document understanding), and CaptionGen (video captioning) have been instrumental in evaluating these trade-offs, demonstrating the need for a more robust and integrated solution.

Self-Distilled Reasoning (SDR): A Paradigm Shift in Fine-Tuning

Self-Distilled Reasoning (SDR) represents a significant advancement, offering an elegant and efficient way to circumvent the limitations of traditional SFT and model merging. At its core, SDR re-uses the chain of thought generated by the base Amazon Nova 2 Lite model itself, essentially allowing the model to "learn from itself." This process unfolds in three key stages:

1. Trace Generation: The base reasoning model (e.g., Amazon Nova 2 Lite) is used to generate reasoning traces for each example in the SFT dataset that lacks them. This can be done either through "basic reasoning" (model naturally produces reasoning given the question) or "guided reasoning" (model reconstructs plausible reasoning given both question and ground-truth answer). Guided reasoning is particularly useful when the base model is weak on the target domain, as it provides additional signal during trace generation.
2. Data Augmentation: The generated reasoning traces are then prepended to the original outputs in the training data, creating enriched examples that include both the thought process and the final answer.
3. Fine-Tuning with Reasoning: The augmented dataset is then used to fine-tune the model with reasoning mode explicitly enabled, ensuring that the model receives supervision on both the reasoning tokens and the output tokens.

This approach offers several compelling advantages:

• Zero Annotation Cost: No human effort is required to create the reasoning traces, eliminating a major cost and logistical barrier.
• Consistency: The generated reasoning traces are inherently consistent with the base model’s own problem-solving approach, providing a natural and effective learning signal.
• Scalability and Flexibility: The technique is applicable to existing SFT datasets across various domains and allows for the use of different teacher models, though experiments suggest that same-family teachers (e.g., Nova 2 Lite for Nova 2 Lite) yield the most balanced results.
• In-Training Regularization: SDR provides a powerful form of in-training regularization by augmenting the dataset with the model’s own reasoning. This acts as an implicit policy regularization, constraining the fine-tuned model to remain close to the base model’s policy, thereby preventing catastrophic forgetting, similar to KL-regularization in RLHF. It also provides process supervision, akin to step-by-step reward shaping, leading to more robust generalization.

Constructing Reasoning Traces with Amazon Bedrock

The implementation of SDR leverages the Amazon Bedrock Converse API to construct reasoning traces. The process involves defining prompt templates for both basic and guided reasoning. Basic reasoning mimics natural inference, while guided reasoning, by providing the ground-truth answer, prompts the model to reconstruct the logical steps. A crucial step for guided reasoning involves a secondary filtering process to ensure the model does not "leak" the answer into the reasoning trace, which could introduce noise.

The technical pipeline involves reading SFT training data, invoking the Bedrock Converse API for each prompt to generate CoT traces, and then assembling the final dataset with reasoningContent blocks. This structured approach ensures that the reasoning steps are clearly delineated and integrated into the training data format, ready for SFT.

Experimental Validation: SDR’s Superior Performance

Extensive experiments across multiple benchmarks—MedMCQA (medical question answering), CoCoHD (long-context structured extraction), and Invoice-OCR (multimodal document understanding)—have consistently validated the efficacy of SDR. The results demonstrate that SDR not only significantly improves target performance but also robustly preserves general capabilities without needing model merging.

For instance, on MedMCQA, where the base Nova 2 Lite model scored 55.60%, vanilla SFT with no reasoning led to catastrophic forgetting, with math performance dropping from 70% to 0%. Model merging partially recovered this to 8.30% while achieving 60.50% target performance. In contrast, SDR (using Nova 2 Lite Basic traces) achieved 66.60% target performance and retained 67.90% of math capability, demonstrating a clear advantage in both specialization and generalization. Similar improvements were observed on CoCoHD and Invoice-OCR, where SDR consistently delivered better target performance and dramatically higher math retention compared to SFT without reasoning or even with model merging.

Ablation Studies: SDR’s Resilience to Partial Reasoning Data

A crucial aspect of SDR’s utility lies in its effectiveness when only a portion of the SFT dataset contains reasoning traces—a common real-world scenario given the expense of curating such data. An ablation study on the LLaVA CoT benchmark, using 10k training samples, revealed critical insights into catastrophic forgetting and the power of SDR.

• Catastrophic Forgetting: Without SDR, reducing reasoning data below 75% caused a catastrophic collapse in general math capability. For example, at 50% reasoning data, math performance plummeted from 66.6% to 21.7%, and at 0% reasoning data, it dropped to a mere 2.5%. The model entirely lost its general mathematical reasoning.
• SDR’s Stabilizing Effect: With SDR prefill, math control performance remained remarkably stable, hovering between 64% and 72% across all data percentages, even with 0% reasoning data. This unequivocally demonstrates SDR’s ability to completely prevent catastrophic forgetting.
• Reasoning Token Generation: SDR-trained models consistently produced a higher median number of reasoning tokens during inference (400-860), indicating that the models actively engage their reasoning mechanisms. In stark contrast, models trained with less than or equal to 25% reasoning data and no prefill produced zero reasoning tokens, confirming their complete loss of step-by-step reasoning ability.
• Target Performance and Inference Reasoning: Enabling reasoning at inference time consistently delivered a roughly 2x improvement in MathVista accuracy (35-47% vs. 14-22%), irrespective of training data composition. Notably, training with SDR prefill even boosted performance when reasoning was disabled during inference, suggesting that SDR instills a more robust problem-solving approach within the model. The best configuration, 50% reasoning data plus SDR prefill, achieved 45.2% MathVista with 70.4% Math control, demonstrating maximal target performance while preventing forgetting.

Implications for AI Customization and Continual Learning

The findings from SDR have profound implications for practitioners deploying customized LLMs. Firstly, maintaining a consistent reasoning mode during both training and inference is paramount for optimal performance. Secondly, SDR provides an indispensable tool for preventing the degradation of general reasoning abilities during domain-specific SFT. Thirdly, it offers a superior alternative to model merging by removing the trade-off between target and general performance, simplifying the customization process and improving outcomes.

SDR’s integration of reasoning traces into the SFT workflow connects it to several advanced AI paradigms. It acts as a form of implicit policy regularization, akin to KL-regularization in Reinforcement Learning from Human Feedback (RLHF), by keeping the fine-tuned model’s policy close to that of the base model. This process supervision, rather than mere outcome supervision, fosters more robust generalization. Furthermore, SDR aligns with the principles of self-distillation, where a model learns from its own rich, problem-solving processes, rather than just final outputs. This is particularly valuable in continual learning settings, where models need to acquire new skills over time while preserving previously learned knowledge.

A practical nuance identified in the research concerns the choice of the teacher model for generating reasoning traces. While a stronger teacher might seem intuitively better, experiments indicated that Amazon Nova 2 Pro (Preview) traces, while improving target accuracy, caused a proportional drop in general performance. This aligns with prior findings that a significant student-teacher gap can negatively impact distillation. Consequently, Lite-generated traces for Amazon Nova 2 Lite models offer the most balanced outcome across target and general performance and are the recommended default.

Practical Recommendations for Practitioners

AWS provides clear guidance for practitioners on when and how to leverage SDR, model merging, or standard SFT based on their specific data composition and latency constraints:

• Default SFT without Reasoning, Experiencing Regression: If latency-sensitive, use model merging to recover general performance. If latency constraints allow for reasoning, SDR with Nova 2 Lite-generated traces is recommended for the best balance of target performance and general capability preservation.
• Base Model Weak on Target Domain: Consider using guided reasoning during trace generation to provide additional, ground-truth-informed signals.
• Dataset with Partial Reasoning (≥ 50%): Train with reasoning turned on; existing traces are often sufficient to preserve general skills.
• Dataset with < 50% Reasoning Traces: Pre-fill missing traces with Nova 2 Lite reasoning (a one-time, offline process). Without this, general performance is likely to collapse.
• General Capabilities Not Required: If only target performance matters, standard SFT without explicit reasoning data may suffice, as reasoning data has minimal impact on target performance in isolation.
• Continual Learning or Multi-Skill Preservation: Combine SDR with data mixing and a small (or zero) merge weight to optimize the retention of previously learned skills.
• Desire to Still Use Model Merging: Apply a smaller merge weight when SDR is utilized, as SDR already provides the regularization that merging was intended to compensate for.

In conclusion, Self-Distilled Reasoning marks a pivotal advancement in the customization of LLMs. By providing a cost-effective, robust, and integrated solution to the reasoning suppression problem and catastrophic forgetting, SDR empowers developers to unlock the full potential of sophisticated models like Amazon Nova 2. It simplifies the fine-tuning workflow, improves model reliability, and paves the way for more intelligent, adaptable, and domain-specific AI applications, particularly in complex, real-world scenarios requiring deep reasoning capabilities. The continued exploration of SDR in conjunction with other advanced fine-tuning techniques promises further enhancements in the quest for ever more capable and resilient AI systems.