PaperlayerView exact revision on arXiv

Source-linked AI summary

Which Tokens Should SFT Actually Learn? A Token-Trimming Perspective on Mathematical Reasoning

Yaning Jia, Chunhui Zhang, Wenxuan Xu, Xingjian Diao, Xiaoyuan Wang, Soroush Vosoughi

arXiv:2609.09707v1cs.AI

TL;DR

Uniform SFT treats all mathematical-reasoning tokens alike despite unequal learning signals. TrimSFT uses Gaussian logit-gap reweighting to trim both extremes, and it consistently outperforms standard SFT across the evaluated models and benchmarks.

  • Problem

    Standard SFT uniformly weights target tokens despite their unequal learning signals in mathematical reasoning.

  • Method

    TrimSFT applies Gaussian token-level weights based on the logit gap between each gold token and its strongest competitor, without an extra forward pass.

  • Results

    TrimSFT consistently improves over standard SFT, achieves the best average performance on five of six models, and gains up to +26.9 points on MATH500.

  • Takeaways & Limitations

    Reasoning SFT can benefit from concentrating supervision in an intermediate logit-gap region while trimming both extremes.

  • Takeaways & Limitations

    The method is evaluated primarily on mathematical reasoning, so its effectiveness in code generation, open-ended instruction following, and long-form reasoning remains to be studied.

Abstract

from arXiv · show

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.

1 Introduction

Standard SFT weights every target token equally despite unequal learning signals, motivating TrimSFT’s two-sided token-level reweighting for mathematical reasoning.

  • Standard SFT applies uniform token-level cross-entropy even though reasoning tokens provide different learning signals.
  • TrimSFT scales loss using the gap between the gold token and its strongest competitor, trimming both mastered and weakly supported tokens.It concentrates supervision in an intermediate logit-gap region.
  • TrimSFT requires neither a reference model nor an additional forward pass because weights come from the model’s own logits.
  • TrimSFT is evaluated on six base models across five mathematical reasoning benchmarks and consistently improves over standard SFT.
  • +26.9 points over SFT is the largest reported MATH500 gain, while TrimSFT achieves the best average performance on five of six models.The contribution summary also reports improvements in pass@8 and best-of-8.
  • Token-level analysis suggests TrimSFT reshapes model confidence more evenly than uniform SFT or monotonic reweighting methods.

2 Related Work

Prior token-reweighting methods modify SFT using probability-based or auxiliary signals, whereas TrimSFT uses a scale-sensitive logit gap and a bounded, non-monotonic weighting profile.

  • Recent SFT research explores token-level reweighting and selection because tokens can differ in training value.
  • Some methods use auxiliary signals or pre-computed masks, including reference-model scoring in Rho-1 and counterfactual selection in CFT.
  • DFT and Focal Loss rescale cross-entropy using gold-token probability, but they emphasize opposite confidence regimes.Focal Loss down-weights high-probability tokens, whereas DFT increases their relative contribution.
  • TrimSFT instead uses logit gaps and Gaussian band-pass weighting to target a bounded intermediate region rather than a monotonic trend.

3 Method

TrimSFT reweights token-level SFT loss by logit gaps, emphasizing an intermediate confidence region while down-weighting both weakly supported and already mastered tokens.

  • 3 Method: Standard SFT minimizes token-level cross-entropy with uniform weight across target tokens.
  • 3.1 Trimmed Logit-Gap SFT: TrimSFT reduces supervision for very large gaps and for very small or negative gaps, concentrating learning in an intermediate region.
  • 3 Method: The logit gap is the margin between the gold token’s logit and its strongest competitor, measuring local prediction separation.A positive gap means the gold token is currently top-1.
  • 3.1 Trimmed Logit-Gap SFT: A Gaussian weight peaks near margin m and decays on both sides, with m setting the center and τ controlling the region’s width.
  • 3.1 Trimmed Logit-Gap SFT: The detached weights are computed in the same forward pass and only rescale cross-entropy, requiring no auxiliary model or extra forward pass.Stop-gradient prevents the weights from contributing a gradient through the model parameters.
  • 3.3 A Unified View of Token Reweighting: TrimSFT differs from probability-based and monotonic reweighting by using scale-sensitive logit gaps and trimming both confidence extremes.
  • 3.2 Half-Trim Variants: Half-trim variants retain full weight on one side of margin m while applying Gaussian decay on the other.TrimSFT-E preserves high-gap tokens; TrimSFT-H preserves low-gap tokens.
  • 3.2 Half-Trim Variants: Comparing the half-trim variants with full TrimSFT isolates whether trimming high-gap tokens, low-gap tokens, or both contributes to performance.

4 Experiments

TrimSFT is evaluated across six base models and five mathematical reasoning benchmarks, with additional analyses of repeated sampling and hyperparameter choices. It consistently outperforms standard SFT, while results favor a narrow bandwidth and full two-sided trimming.

  • 4.1 Setup: TrimSFT is evaluated on six base models spanning 1.5B–8B parameters and five mathematical reasoning benchmarks, using one shared setting in the main comparison.The main setting is (m, τ) = (1.5, 0.8), without model- or benchmark-specific tuning.
  • 4.2 Main Results: TrimSFT improves over standard SFT across all six models and achieves the best average performance on five of six models.On Qwen2.5-Math-1.5B, the average score rises from 15.47 to 31.77; on Qwen2.5-Math-7B, it rises from 20.95 to 35.98.
  • 4.2 Main Results: On MATH500, TrimSFT is top-performing across all six models, with a maximum SFT gain of +26.93 points on Qwen2.5-Math-1.5B.The cited comparison is 40.02 → 66.95.
  • 4.3 Capability Ceiling and Self-Consistency: Across repeated sampling, TrimSFT improves pass@8 over standard SFT on both reported models across most benchmarks and remains stronger or competitive on best-of-8 with majority voting.These metrics characterize solution discovery and self-consistency, respectively.
  • 4.4 Ablation: Margin and Bandwidth: Performance is more sensitive to bandwidth τ than margin m: narrow bandwidth remains strong across margins, whereas broad bandwidth causes clearer degradation.With τ = 1, performance stays well above the SFT baseline across margins; increasing τ produces a clear drop, especially for m = 1.
  • 4.5 Half-Trim Variants: Full two-sided trimming provides a more robust trade-off than half-trim variants, with TrimSFT best on AIME24 while TrimSFT-H performs substantially worse across settings.TrimSFT-E can be competitive on MATH500 and AMC, especially on relatively easier or medium-difficulty benchmarks.

5 Mechanism: Logit Gap Distribution Analysis

The logit-gap analysis compares how fine-tuning objectives reshape token-level confidence. TrimSFT shifts confidence toward larger gaps while preserving a broad distribution and reducing mass near small gaps.

  • 5 Mechanism: Logit Gap Distribution Analysis: The analysis computes teacher-forced logit gaps for every response token in 100 sampled training examples and aggregates them into kernel density estimates.The distributions compare Base, SFT, DFT, FSFT, and TrimSFT.
  • 5 Mechanism: Logit Gap Distribution Analysis: Standard SFT shifts the distribution right relative to the base model, while FSFT leaves more mass near small logit gaps.DFT moves toward larger gaps and appears closer to TrimSFT than to standard SFT.
  • 5 Mechanism: Logit Gap Distribution Analysis: TrimSFT shifts the distribution toward larger gaps while preserving a broad shape and reducing mass in the small-gap region.Unlike DFT, TrimSFT achieves this through two-sided trimming rather than monotonically emphasizing high-probability tokens.
  • 5 Mechanism: Logit Gap Distribution Analysis: The distributions support the interpretation that TrimSFT reshapes token-level confidence more evenly than uniform SFT or methods that emphasize one confidence direction.Additional token-category results show more gradient mass on numeric and mathematical-symbol tokens and less on generic other tokens.

6 Conclusion

TrimSFT concentrates supervision on an intermediate logit-gap region by trimming both already well-separated and weakly supported tokens. Across six base models and five mathematical reasoning benchmarks, it consistently improves over standard SFT and yields more balanced confidence shaping.

  • TrimSFT uses logit gaps to concentrate supervision on an intermediate confidence region.It trims supervision from both already well-separated tokens and tokens with weak current support, providing a non-monotonic alternative to uniform SFT.
  • Across six base models and five mathematical reasoning benchmarks, TrimSFT consistently improves over standard SFT.The reported gains are accompanied by improved capability coverage, self-consistency, and better trade-offs than half-trim variants.
  • TrimSFT reshapes token-level confidence by moving tokens toward more confident regions while avoiding uniform over-sharpening.The conclusion attributes this behavior to selectively concentrating supervision within an intermediate logit-gap region.

Limitations

The evaluation is concentrated on mathematical reasoning benchmarks and supervised fine-tuning from base models, leaving broader-domain effectiveness for future study. The method also uses margin and bandwidth hyperparameters, although the reported ablations suggest robustness to exact margin choice.

  • TrimSFT is evaluated primarily on mathematical reasoning benchmarks with supervised fine-tuning from base models.Its effectiveness on code generation, open-ended instruction following, and long-form reasoning remains to be studied.
  • TrimSFT uses margin m and bandwidth τ to define the emphasized logit-gap region.The ablations report relative robustness to the exact margin choice, while different model scales or data distributions may benefit from adaptive weighting regions.

Ethical Considerations

The study uses publicly available pretrained models and mathematical reasoning datasets, without new data collection or human annotation. Its experiments apply shared training and evaluation configurations, while broader capability gains may carry potential downstream misuse risks.

  • TrimSFT does not involve new data collection or human annotation.The experiments use publicly available pretrained models and mathematical reasoning datasets intended for research use.
  • Improving reasoning capabilities may also strengthen broader generation abilities that could potentially be misused downstream.
  • All methods are trained under an identical configuration and differ only in their per-token cross-entropy weights.The shared setup includes optimizer, learning rate, schedule, batch size, sequence length, precision, and number of epochs.
  • Evaluation samples eight generations per problem using a shared protocol across methods and benchmarks.Sampling uses temperature 1.0, top-p p = 1.0, and a maximum generation length of 2048 tokens.

B Detailed Results for Ablations

The ablations examine how bandwidth τ and margin m shape TrimSFT’s weighting region, then compare full and half-trim variants. Results indicate that τ controls selectivity more strongly than m, while two-sided trimming is especially advantageous on harder problems.

  • B Detailed Results for Ablations: Table 3 fixes m and varies τ, while Table 4 fixes τ and varies m in the ablation study.The appendix reports complete per-benchmark results alongside the averaged trends summarized in the main text.
  • C Understanding the Roles of m and τ: The weighting analysis computes logit gaps on 40,082 gold response tokens and derives their TrimSFT weights.The logit gap is defined as Δ_t = z_t,y_t − max_v≠y_t z_t,v; High-Weight (%) counts tokens with w_t > 0.5, and R_eff measures supervision uniformity.
  • C Understanding the Roles of m and τ: Increasing τ substantially broadens the weighted region and moves the effective token ratio toward uniform SFT.With m varied from 1 to 6 and τ = 1, the fraction of high-weight tokens remains nearly unchanged, indicating that m mainly shifts the weighting band.
  • D Half-Trim Variants: TrimSFT-E can be competitive on easier or medium-difficulty benchmarks, but TrimSFT performs better on the harder AIME24 benchmark.TrimSFT-H performs much worse than both alternatives, supporting two-sided trimming over either half-trim alone.

E Gradient-Mass Analysis

The gradient-mass analysis compares full TrimSFT with half-trim variants by measuring how optimization signal is distributed across low-, intermediate-, and high-gap tokens. It finds that low-gap trimming drives stability, while high-gap trimming adds regularization for already well-separated tokens.

  • Gradient-Mass Analysis: The analysis measures each weighting scheme’s fraction of total gradient mass in low-, middle-, and high-gap regions.The regions are defined relative to margin m and bandwidth τ: below m−τ, between m−τ and m+τ, and above m+τ.
  • Gradient-Mass Analysis: TrimSFT-H concentrates gradient mass in low-gap tokens, reaching 71.78%, 93.80%, and 99.04% as m increases.TrimSFT-H preserves full supervision on low-gap tokens.
  • Gradient-Mass Analysis: Full TrimSFT and TrimSFT-E suppress the low-gap region and retain more optimization signal in the intermediate region.The comparison uses the same margin and bandwidth across the full and half-trim variants.
  • Gradient-Mass Analysis: The results support low-gap trimming as the primary source of stability and high-gap trimming as a complementary regularization effect.This interpretation follows the asymmetric gradient-mass distributions of the half-trim variants.

F Token-Category Analysis

The token-category analysis compares standard SFT with TrimSFT by examining how much gradient mass different response-token categories receive. TrimSFT shifts more optimization signal toward numeric and mathematical-symbol tokens.

  • Token-Category Analysis: The analysis uses 100 NuminaMath-CoT training examples containing 40,082 response tokens and compares standard SFT with TrimSFT at (m, τ) = (1, 1).Tokens are grouped into four lightweight categories.
  • Token-Category Analysis: TrimSFT increases gradient-mass share for numeric tokens from 5.55% to 10.31%.The comparison is against standard SFT on Qwen2.5-Math-1.5B.
  • Token-Category Analysis: TrimSFT increases gradient-mass share for mathematical symbols from 15.76% to 22.78%.The broader Other category receives a reduced contribution under TrimSFT.
  • Token-Category Analysis: The reweighted optimization signal remains associated with math-relevant token categories rather than primarily concentrating on the broad Other category.The table defines Num. as numeric tokens and MS as mathematical symbols.
Loading 2609.09707v1…