TMS: Trajectory-Mixed Supervision for On-Policy Self Distillation
Rana Khan, Zijie Liu, Zhen Tan, Charles Fleming, Tianlong Chen
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 often 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 Supervision Mismatch: the divergence between the model's evolving policy and static training labels. We address this trade-off with Trajectory-Mixed Supervision (TMS), a reward-free framework that uses trajectory-aligned, near-policy supervision harvested from the model's own historical checkpoints. TMS reduces Policy-Label Divergence (PLD) within SFT-style training and preserves multiple plausible solution modes, mitigating a key source of 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 strongest retention baseline, TMS significantly outperforms standard and iterative SFT, narrowing the gap to RL without requiring reward models or verifiers. Mechanistic analysis shows that KL-to-base is the strongest cross-method predictor of forgetting, while PLD provides a complementary diagnostic of supervision mismatch within SFT-style methods.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 8e1c9c61-9312-4a98-aefa-dc8a9097be90Builds on24
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- Self-Refine: Iterative Refinement with Self-FeedbackAman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan et al.NeurIPS 2023 · 4,972 citations
Related papers
- Retaining by Doing: The Role of On-Policy Data in Mitigating ForgettingHoward Chen, Noam Razin, Karthik Narasimhan, Danqi ChenICML 2026
- Reinforcement Fine-Tuning Naturally Mitigates Forgetting in Continual Post-TrainingSong Lai, Haohan Zhao, Rong Feng, Changyi Ma et al.ICML 2026 · 46 citations
- Trust-Region Adaptive Policy OptimizationMingyu Su, Jian Guan, Yuxian Gu, Minlie Huang et al.ICLR 2026 · 2 citations
- Why Does Reinforcement Learning Generalize? A Feature-Level Mechanistic Study of Post-Training in Large Language ModelsDan Shi, Zhuowen Han, Simon Ostermann, Renren Jin et al.ACL 2026 · 1 citation
- Temporal Sampling for Forgotten Reasoning in LLMsYuetai Li, Zhangchen Xu, Fengqing Jiang, Bhaskar Ramasubramanian et al.ACL 2026 · 13 citations
