Technical Summary: RAFT (Data Refinement and Adaptive Distillation for Domain Fine-Tuning with Alleviated Forgetting)

Problem Statement

Domain-specific supervised fine-tuning (SFT) is essential for adapting Large Language Models (LLMs) to specialized tasks but frequently results in catastrophic forgetting, where the model's general capabilities degrade significantly. The authors identify two practical gaps driving this degradation:

1. Supervision-Compatibility Gap: Domain ground-truth answers often differ in style and reasoning format from the model's natural responses. Forcing the model to fit these divergent targets causes abrupt parameter shifts that overwrite pre-trained knowledge.
2. Trajectory-Preservation Gap: Standard SFT optimizes for fixed target tokens without constraining the model's behavior on its own generated prefixes. This lack of constraint fails to preserve the model's original behavioral patterns during optimization.

Existing methods typically address only one side of this problem (e.g., data refinement via self-distillation or training constraints via regularization/KL divergence) but rarely both, leading to suboptimal trade-offs between domain adaptation and general capability retention.

Methodology

RAFT proposes a two-stage framework that couples data refinement with adaptive on-policy distillation to address both gaps simultaneously.

1. Offline Distillation (Data Refinement)

This stage constructs model-compatible supervision to reduce the distribution mismatch between ground truth and the model's output distribution.

• Self-Conditioned Rewriting: The original instruction-tuned model rewrites ground-truth answers ($y_i$) conditioned on the original answer itself ($\tilde{y}_i \sim p_\theta(\cdot | x_i, y_i)$), adapting the expression style to the model's own distribution while preserving factual content.
• Semantic Filtering: To prevent semantic drift, cosine similarity is computed between the original and rewritten answers. Rewrites are accepted only if they exceed a threshold $\tau$; otherwise, the original answer is retained.
• Multi-Expert Optimal Selection: For accepted rewrites, a stronger fusion model (e.g., Qwen3-3B) combines the original and rewritten answers. Multiple judge models evaluate the fused output for consistency with the original answer, and the highest-scoring result is selected as the final refined answer ($\hat{y}_i$).

2. Adaptive On-Policy Distillation (Training Optimization)

This stage trains the model using the refined dataset while explicitly preserving general capabilities through distillation constraints.

• Answer-Conditioned On-Policy Distillation: The student model generates trajectories ($z_i$) autoregressively. The original instruction-tuned model acts as a teacher, providing soft targets. Crucially, the teacher is conditioned on the fused answer ($\hat{y}_i$) as additional context, allowing it to provide more informative guidance on the student's generated trajectory.
• Top-K Temperature Distillation: To prevent over-smoothing and maintain output diversity, KL divergence is computed only over the top-$K$ tokens of the teacher's distribution, scaled by a temperature parameter $T$.
• EMA-Based Adaptive Balancing: Instead of fixed hyperparameters, an Exponential Moving Average (EMA) mechanism dynamically adjusts the weight ($\lambda$) between the SFT loss (cross-entropy) and the distillation loss (KL divergence). This ensures the two objectives remain balanced throughout training without manual tuning.

Key Contributions

1. Problem Identification: The paper pinpoints the coupling of supervision distribution mismatch and the lack of explicit trajectory preservation constraints as the root causes of forgetting in domain SFT.
2. Offline Distillation Strategy: Introduction of a pipeline using cosine similarity filtering and multi-model fusion to generate high-quality, distribution-consistent supervision data.
3. Adaptive On-Policy Distillation: A novel training objective combining Answer-Conditioned On-Policy Distillation, Top-K temperature scaling, and EMA-based loss balancing.
4. Empirical Validation: Demonstration that coupling data refinement with trajectory-level preservation effectively mitigates forgetting while improving domain performance.

Experimental Results

RAFT was evaluated on three instruction-tuned backbones (SmolLM3-3B, Llama-3.2-3B-Instruct, Phi-4-mini-instruct) across five domains (Business, Law, Science, Culture, Web).

• Domain Performance: RAFT improved average Domain Accuracy (D-Acc) by 23.2% compared to standard SFT.
• General Capability Retention: RAFT recovered significant portions of the degradation caused by SFT:
  • MS-Bench (subjective QA): Improved by 18.2% relative to SFT.
  • IFEval (instruction following): Improved by 10.2% relative to SFT.
• Ablation Studies:
  • Data Fusion: Using fused data alone improved D-Acc and general metrics over original data, confirming the value of style adaptation.
  • Distillation: Adding adaptive on-policy distillation to fused-data SFT provided substantial further gains in both domain accuracy and general capability.
  • Adaptive Balancing: The EMA-based strategy outperformed fixed KL coefficients, which either suppressed domain adaptation or failed to preserve general capabilities.
• Parameter Analysis: Optimal performance was found with a cosine similarity threshold $\tau=0.8$, temperature $T=1.5$, and Top-$K=512$.

Significance and Claims

The paper claims that RAFT provides an effective "recipe" for domain fine-tuning with alleviated forgetting. By treating data refinement and training constraints as coupled rather than independent problems, the framework allows models to acquire domain-specific knowledge without overwriting pre-trained general capabilities.

The authors emphasize that their approach turns the teacher model from a "conservative anchor" into a "target-aware guide." The results suggest that coupling compatible supervision with trajectory-level preservation is a practical solution to the trade-off between specialization and generalization in LLMs. The paper remains modest regarding scalability, noting that evaluations were limited to small-scale models and specific domains, and that performance on larger models or specialized tasks (e.g., code generation) requires further investigation.