PriFT: Prior-Support Guided Supervised Fine-Tuning

2026-06-08

AI Summary

This work addresses the overfitting and limited generalization inherent in supervised fine-tuning (SFT), which often arises from fitting target tokens inconsistent with the pretraining distribution. To mitigate this, the authors propose PriFT, a method that leverages a frozen pretrained model to compute prior support for each target token—quantified via token probability (PriFT-prob) and cumulative probability mass (PriFT-mass)—and uses these signals to stably reweight the fine-tuning objective. By deriving weights from the fixed pretrained model rather than the online fine-tuned model, PriFT avoids interference from optimization dynamics and effectively preserves prior knowledge. Experiments demonstrate that PriFT significantly outperforms existing SFT approaches across mathematical reasoning, code generation, and medical question answering tasks, while also yielding better initializations for subsequent reinforcement learning stages.

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.

Problem

• supervised fine-tuning
• overfitting
• off-policy learning
• pretrained distribution
• token reweighting

Innovation

• supervised fine-tuning
• prior support
• token reweighting
• pretrained distribution
• stable optimization