ProFit: Leveraging High-Value Signals in SFT via Probability-Guided Token Selection

TL;DR

ProFit tackles single-reference overfitting in supervised fine-tuning by exploiting token-level probability as a proxy for semantic importance. It introduces a probability-guided masking mechanism that retains high-probability tokens carrying core reasoning while masking low-probability, non-essential tokens, backed by a semantic analysis and a gradient-bound justification. The approach yields consistent improvements over standard SFT and other baselines across reasoning and mathematics benchmarks on diverse LLM families, and demonstrates stable training dynamics and favorable RL initialization characteristics. This work provides a practical, data-efficient alternative to multi-reference fine-tuning with potential for broad applicability in reasoning-intensive tasks.

Abstract

Supervised fine-tuning (SFT) is a fundamental post-training strategy to align Large Language Models (LLMs) with human intent. However, traditional SFT often ignores the one-to-many nature of language by forcing alignment with a single reference answer, leading to the model overfitting to non-core expressions. Although our empirical analysis suggests that introducing multiple reference answers can mitigate this issue, the prohibitive data and computational costs necessitate a strategic shift: prioritizing the mitigation of single-reference overfitting over the costly pursuit of answer diversity. To achieve this, we reveal the intrinsic connection between token probability and semantic importance: high-probability tokens carry the core logical framework, while low-probability tokens are mostly replaceable expressions. Based on this insight, we propose ProFit, which selectively masks low-probability tokens to prevent surface-level overfitting. Extensive experiments confirm that ProFit consistently outperforms traditional SFT baselines on general reasoning and mathematical benchmarks.

Figures (11)

• Figure 1: Breaking the trade-off between training cost and semantic diversity. While Multi-reference SFT offers semantic richness at prohibitive data and computational costs, standard SFT is efficient but semantically limited. ProFit achieves the best of both: by focusing supervision on high-value tokens, it captures core semantic integrity without sacrificing the efficiency of single-reference training.
• Figure 2: Performance comparison on diverse benchmarks. While multi-reference training (SFT w/ 3 ans) offers sporadic gains, it suffers from optimization instability and stagnation on complex tasks. In contrast, ProFit achieves superior and robust performance across all metrics by selectively extracting high-value signals from a single reference.
• Figure 3: Probability density estimation of semantic tokens. We categorize tokens into semantically Core and Trivial groups. While core tokens are heavily concentrated in high-confidence zones, trivial tokens exhibit a significant long-tail distribution, disproportionately dominating the low-probability spectrum. a hypothesis test confirms this significant distributional difference ($p = 1 \times 10^{-6}$).
• Figure 4: Ablation study on the probability threshold $\tau$. The dashed line represents the performance of the standard SFT baseline. (a) Training exclusively on low-probability tokens ($p(y_t^*) < \tau$) results in performance consistently below the baseline, indicating that non-core expressions are insufficient for constructing effective reasoning chains. (b) Conversely, the proposed strategy ($p(y_t^*) > \tau$), which masks low-probability noise, consistently outperforms the baseline across all tasks, validating the effectiveness of focusing on core logic.
• Figure 5: Average performance variation across different LoRA ranks ($r \in \{4, \dots, 1024\}$). The dashed lines represent the baseline performance of full-parameter fine-tuning for each corresponding setting. While core tokens ($p(y_t^*) > 0.1$) exhibit monotonic improvement driven by capacity, non-core tokens ($p(y_t^*) < 0.1$) and standard SFT show a U-shaped trend, revealing optimization interference at medium ranks.

Theorems & Definitions (2)

• Theorem 1
• Proof A.1