InfoSFT: Learn More and Forget Less with Information-Aware Token Weighting

TL;DR

InfoSFT introduces an information-aware token weighting scheme for supervised fine-tuning that enhances generalization and preserves prior capabilities by focusing on medium-confidence tokens.

Contribution

It proposes a simple, principled token weighting method that improves fine-tuning effectiveness without complex filtering or regeneration techniques.

Findings

Outperforms vanilla SFT and likelihood-weighted baselines across multiple tasks.

Better preserves pre-existing capabilities compared to existing methods.

Requires only a one-line modification to the standard loss.

Abstract

Supervised fine-tuning (SFT) provides the standard approach for teaching LLMs new behaviors from offline expert demonstrations. However, standard SFT uniformly fits all samples -- including those with low likelihood under the base model -- which can disproportionately drive training updates toward overfitting specific samples rather than learning the target behavior. Moreover, adapting to these unlikely samples induces substantial policy shifts that degrade prior capabilities. Existing methods mitigate this by filtering, regenerating, or down-weighting low-likelihood data. In doing so, they often suppress precisely the novel behaviors the base model has yet to learn. We propose InfoSFT, a principled weighting scheme for the SFT objective that concentrates learning signals on maximally informative, medium-confidence tokens -- those neither overly familiar to the base model nor too unlikely to cause instability. Requiring only a one-line modification to the standard token-wise loss, InfoSFT demonstrably improves generalization over vanilla SFT and likelihood-weighted baselines across math, code, and chain-of-thought tasks with diverse model families, while better preserving pre-existing capabilities.