TMS: Trajectory-Mixed Supervision for Reward-Free, On-Policy SFT
Rana Muhammad Shahroz Khan, Zijie Liu, Zhen Tan, Charles Fleming, Tianlong Chen
TL;DR
This work tackles catastrophic forgetting in supervised fine-tuning by introducing Trajectory-Mixed Supervision (TMS), a reward-free post-training method that samples near-policy targets from a model's historical checkpoints to form a trajectory-based supervision mix. By replacing a single static reference with a multi-trajectory signal, TMS mitigates supervision-mismatch drift and preserves multiple valid solution modes, yielding retention gains close to on-policy RL while maintaining SFT-level target performance. The authors formalize PLD to diagnose mismatch and prove a bound showing forgetting increases with KL drift from the base policy. Empirically, TMS consistently improves the accuracy-retention frontier across reasoning and instruction-following tasks, scales to large models, and remains effective in multimodal and tool-use settings, offering a practical, reward-free alternative to RL for robust post-training.
Abstract
Reinforcement Learning (RL) and Supervised Fine-Tuning (SFT) are the two dominant paradigms for enhancing Large Language Model (LLM) performance on downstream tasks. While RL generally preserves broader model capabilities (retention) better than SFT, it comes with significant costs: complex reward engineering, instability, and expensive on-policy sampling. In contrast, SFT is efficient but brittle, often suffering from catastrophic forgetting due to $\textbf{Supervision Mismatch}$: the divergence between the model's evolving policy and static training labels. We address this trade-off with $\textbf{Trajectory-Mixed Supervision (TMS)}$, a reward-free framework that approximates the on-policy benefits of RL by creating a dynamic curriculum from the model's own historical checkpoints. TMS minimizes $\textit{Policy-Label Divergence (PLD)}$, preventing the mode collapse that drives forgetting in standard SFT. Experiments across reasoning (MATH, GSM8K) and instruction-following benchmarks demonstrate that TMS effectively shifts the accuracy--retention Pareto frontier. While RL remains the gold standard for retention, TMS significantly outperforms standard and iterative SFT, bridging the gap to RL without requiring reward models or verifiers. Mechanistic analysis confirms that PLD drift accurately predicts forgetting and that TMS successfully mitigates this drift.