Model token weighting method improves math and code tasks
Rethinking Token Reweighting for SFT: Suppress, Reverse, and Extrapolate Learned Features
Computation and Language
Summary
Fine-tuning language models involves teaching them to pay attention to certain words more than others, but this can sometimes reinforce mistakes or erase useful knowledge. The authors studied how existing ways of weighting tokens during training can only increase or decrease importance but cannot undo harmful behavior already learned. They introduced SCALE, a method that controls these learned features more precisely by selectively activating or suppressing parts of the model based on uncertainty signals. This approach helps models perform better at math reasoning and code generation while keeping general knowledge intact.
What this means in practice
- •For ai model developers: Improve fine-tuning of language models on math reasoning and code generation by selectively adjusting learned features without harming general knowledge.
- •For software quality teams: Enhance automated code generation tools by integrating models trained with controlled token adaptation to reduce errors and improve code correctness.
Authors
Cunchun Li, Haonan He, Yifan Gao, Minglei Li, Jingqi Ye, Qingyu Yang, Peng Ye
Abstract
Supervised fine-tuning (SFT) learns most aggressively from tokens that the model deems least likely. This helps acquire new behaviors, but also amplifies noisy or conflicting supervision and can overwrite useful pretrained knowledge. Through a unified policy-loss view, we revisit existing token-reweighting methods and show that they assign nonnegative coefficients to demonstrated tokens. Consequently, they can suppress or amplify supervised updates, but cannot reverse harmful features once learned. Moreover, larger training weights do not amount to feature extrapolation, since they change the optimization trajectory rather than scale a fixed SFT direction. We argue that reversal and extrapolation require a stable reference frame defined by a fixed SFT delta. Motivated by this, we propose SCALE (Selective Control of Adaptation via Local Entropy), an entropy-guided adaptation-strength-control method that freezes the pretrained model and the SFT delta and learns bounded token- and module-specific gates by minimizing predictive entropy alone. These gates suppress, reverse, or extrapolate frozen SFT features according to their alignment with entropy reduction. Across Qwen2.5-Math-1.5B, Qwen2.5-Math-7B, and Qwen3-4B-Base, SCALE achieves mathematical-reasoning averages of 37.84, 43.60, and 36.57, exceeding the strongest corresponding baselines while remaining competitive on general-retention benchmarks. It also attains the best average code-generation performance across HumanEval, HumanEval+, and MBPP for all three models. These results suggest that effective SFT correction can benefit from controlling how already learned residuals are used, rather than only modifying how they are learned.