Supervised fine-tuning improves math reasoning by ignoring easy and hard tokens
Which Tokens Should SFT Actually Learn? A Token-Trimming Perspective on Mathematical Reasoning
Artificial Intelligence
Summary
Training AI models to solve math problems usually treats all parts of an answer equally, but some parts are easier or harder for the model to learn from. The authors suggest focusing training only on the moderately difficult parts, ignoring those the model already understands well and those it gets very wrong. This method, called TrimSFT, changes the model's learning focus based on how confident it is about each token. Tests show that TrimSFT helps models perform better on various math tests.
What this means in practice
- •For machine learning engineers: Improve training of language models on math problems by focusing learning on moderately difficult tokens.
- •For natural language processing teams: Enhance mathematical reasoning capabilities of large language models by adjusting token-level loss weighting without extra model passes.
Authors
Yaning Jia, Chunhui Zhang, Wenxuan Xu, Xingjian Diao, Xiaoyuan Wang, Soroush Vosoughi
Abstract
Supervised fine-tuning (SFT) applies a uniform cross-entropy loss to all target tokens, even though different tokens provide unequal learning signals for mathematical reasoning. This uniform treatment can over-sharpen already mastered tokens while amplifying learning pressure on uncertain, low-confidence tokens, leading to suboptimal training dynamics. We propose Trimmed Logit-Gap SFT (TrimSFT), a simple token-level reweighting method that scales the SFT loss according to the logit gap between the gold token and its strongest competitor. TrimSFT trims supervision away from both extremes: tokens already mastered (large logit gap) and tokens weakly supported by the current model (small or negative logit gap), concentrating learning within an intermediate logit-gap region between them. We instantiate this principle with a Gaussian weight centered at margin m with bandwidth τ, requiring no reference model or additional forward pass. We evaluate TrimSFT on six base models from the Llama, Qwen, and DeepMath families across five mathematical reasoning benchmarks. TrimSFT consistently improves over standard SFT, achieving the best average performance on five out of six models, with gains of up to +26.9 points over SFT on MATH500. Further analyses show that the bandwidth τ matters more than the exact margin location, and that half-trim variants that remove supervision pressure from only one side yield inferior trade-offs. A token-level logit-gap distribution analysis suggests that TrimSFT reshapes model confidence in a more balanced way than uniform SFT or monotonic reweighting methods. These results suggest that reasoning SFT can benefit from trimming both extremes rather than treating all tokens uniformly.