Lune

ICML2026Top-tier venue

TMS: Trajectory-Mixed Supervision for On-Policy Self Distillation

Rana Khan, Zijie Liu, Zhen Tan, Charles Fleming, Tianlong Chen

2026Year

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.

Questions to start from

Your agent calls

Luneget_paper_fulltext

Ask in Lune

Free to start. No credit card required.

lune papers fulltext 8e1c9c61-9312-4a98-aefa-dc8a9097be90

Builds on24

Related papers

Dusk over the sea between two cliffs drawn in fine vertical lines