Search papers, labs, and topics across Lattice.
This paper introduces PriFT, a novel approach to supervised fine-tuning that enhances generalization by deriving token weights from a frozen pretrained model rather than the fine-tuned model itself. By doing so, PriFT mitigates the self-reinforcing dynamics that can occur when token weights are entangled with the optimization trajectory, leading to overfitting on poorly aligned targets. Experimental results demonstrate that PriFT consistently outperforms existing token-reweighting methods across various tasks, including mathematical reasoning and code generation, and serves as a superior initialization for subsequent reinforcement learning training.
PriFT achieves state-of-the-art performance in supervised fine-tuning by leveraging a stable token reweighting signal from a frozen pretrained model, drastically improving generalization.
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.
