PriFT: Prior-Support Guided Supervised Fine-Tuning
2026-06-08 • Computation and Language
Computation and LanguageMachine Learning
AI summaryⓘ
The authors explain that when fine-tuning language models with fixed example data, the models can sometimes overfit and not generalize well because they try to learn tokens that don’t fit their original training. To fix this, they suggest using the original pretrained model to guide which tokens are important to focus on, rather than the changing fine-tuned model. This method, called PriFT, helps keep the model from drifting too far from its original knowledge and improves fine-tuning results. They tested this idea on tasks like math reasoning and coding and found better performance compared to standard methods.
Supervised fine-tuningReinforcement learningPretrained language modelToken reweightingOverfittingPrior distributionToken probabilityOptimization trajectoryMathematical reasoningCode generation
Authors
Ke Wang, Shuangqi Li, Mathieu Salzmann, Pascal Frossard
Abstract
Supervised fine-tuning (SFT) is an efficient approach for downstream task adaptation and often serves as the initialization stage for reinforcement learning (RL), but it can show weaker generalization than RL. A key limitation is its off-policy objective: SFT fits fixed demonstrations token by token, including targets poorly aligned with the model's pretrained distribution, which can lead to overfitting. A recent line of work addresses this issue by assigning larger training weights to tokens better aligned with the current model's predictive distribution, with the intuition that fitting these tokens are less distortive to the model's pretrained knowledge and representations. However, computing the token weights from the model that is currently fine-tuned entangles token weights with the optimization trajectory, inducing a self-reinforcing dynamics as the distribution rapidly departs from the pretrained model. To address this, we propose PriFT (Prior-support guided Fine-Tuning), which derives token weights from a frozen pretrained reference to obtain a stable reweighting signal unaffected by fine-tuning. This signal estimates prior support: the extent to which each target token is supported by the pretrained distribution. Across multiple existing token-reweighting rules, replacing the reweighting signal from the online model to pretrained model consistently improves performance. We introduce two instantiations: PriFT-prob uses pretrained token probability, while PriFT-mass selects tokens by cumulative probability mass under the pretrained distribution. Extensive experiments on mathematical reasoning, code generation, and medical question answering show that PriFT achieves state-of-the-art results among SFT baselines and provides a better initialization for subsequent RL training.