Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Rethinking Supervised Fine-Tuning: Emphasizing Key Answer Tokens for Improved LLM Accuracy

Xiaofeng Shi, Qian Kou, Yuduo Li, Hua Zhou

TL;DR

The paper tackles the issue that conventional SFT distributes learning equally across all target tokens, allowing lengthy CoT reasoning to overshadow the final answer. It introduces SFTKey, a two-stage fine-tuning framework that first trains on full Tag-formatted data with Thinking/Answer markers and then fine-tunes only the Answer tokens to boost accuracy, with a complementary SFTKey-Tag variant that emphasizes Tag-based training. A composite scoring approach balances accuracy and format adherence, demonstrating that SFTKey-Tag achieves higher overall performance (often >5% on average) than standard SFT across multiple models and datasets, with larger models benefiting most. The work provides a practical mechanism to balance CoT learning with precise answer generation, supported by thorough ablations on tagging strategies and one-stage versus two-stage training. It offers insights into token-importance-driven fine-tuning and lays groundwork for more reliable reasoning-guided fine-tuning in LLM systems, while acknowledging limitations in scale and domain coverage.

Abstract

With the rapid advancement of Large Language Models (LLMs), the Chain-of-Thought (CoT) component has become significant for complex reasoning tasks. However, in conventional Supervised Fine-Tuning (SFT), the model could allocate disproportionately more attention to CoT sequences with excessive length. This reduces focus on the much shorter but essential Key portion-the final answer, whose correctness directly determines task success and evaluation quality. To address this limitation, we propose SFTKey, a two-stage training scheme. In the first stage, conventional SFT is applied to ensure proper output format, while in the second stage, only the Key portion is fine-tuned to improve accuracy. Extensive experiments across multiple benchmarks and model families demonstrate that SFTKey achieves an average accuracy improvement exceeding 5\% over conventional SFT, while preserving the ability to generate correct formats. Overall, this study advances LLM fine-tuning by explicitly balancing CoT learning with additional optimization on answer-relevant tokens.

Rethinking Supervised Fine-Tuning: Emphasizing Key Answer Tokens for Improved LLM Accuracy

TL;DR

The paper tackles the issue that conventional SFT distributes learning equally across all target tokens, allowing lengthy CoT reasoning to overshadow the final answer. It introduces SFTKey, a two-stage fine-tuning framework that first trains on full Tag-formatted data with Thinking/Answer markers and then fine-tunes only the Answer tokens to boost accuracy, with a complementary SFTKey-Tag variant that emphasizes Tag-based training. A composite scoring approach balances accuracy and format adherence, demonstrating that SFTKey-Tag achieves higher overall performance (often >5% on average) than standard SFT across multiple models and datasets, with larger models benefiting most. The work provides a practical mechanism to balance CoT learning with precise answer generation, supported by thorough ablations on tagging strategies and one-stage versus two-stage training. It offers insights into token-importance-driven fine-tuning and lays groundwork for more reliable reasoning-guided fine-tuning in LLM systems, while acknowledging limitations in scale and domain coverage.

Abstract

With the rapid advancement of Large Language Models (LLMs), the Chain-of-Thought (CoT) component has become significant for complex reasoning tasks. However, in conventional Supervised Fine-Tuning (SFT), the model could allocate disproportionately more attention to CoT sequences with excessive length. This reduces focus on the much shorter but essential Key portion-the final answer, whose correctness directly determines task success and evaluation quality. To address this limitation, we propose SFTKey, a two-stage training scheme. In the first stage, conventional SFT is applied to ensure proper output format, while in the second stage, only the Key portion is fine-tuned to improve accuracy. Extensive experiments across multiple benchmarks and model families demonstrate that SFTKey achieves an average accuracy improvement exceeding 5\% over conventional SFT, while preserving the ability to generate correct formats. Overall, this study advances LLM fine-tuning by explicitly balancing CoT learning with additional optimization on answer-relevant tokens.
Paper Structure (25 sections, 8 equations, 3 figures, 7 tables)

This paper contains 25 sections, 8 equations, 3 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Standard Supervised Fine-Tuning (SFT)
  4. SFTKey Tag
  5. Experimental Setup
  6. Main results
  7. Ablation Study
  8. Impact of Tagging Strategy on Model Performance
  9. Evaluating SFT and Key under Identical Tagging
  10. Ablation Study on One-Stage vs. Two-Stage Training Strategies
  11. Related Works
  12. Token Importance.
  13. Token Importance based SFT.
  14. Conclusion
  15. Limitation
  16. ...and 10 more sections

Figures (3)

  • Figure 1: As illustrated in the figure, we compare four training strategies: SFT, SFT-Tag, Key-Tag, and SFTKey-Tag. In our setup, the training data are divided into two parts: the chain-of-thought (CoT) and the key (final answer). Building upon the baseline SFT, we further examine three variants: (i) SFT-Tag, which highlights the answer portion using special tags, (ii) Key-Tag, which trains exclusively on the key part, and SFTKey-Tag, a two-stage approach that combines the strategies of (i) SFT-Tag and (ii) Key-Tag.
  • Figure 2: Boxplot showing the distribution of accuracy differences between Key-Tag and SFT-Tag for individual models across multiple datasets. Each of the first five boxes represents a single model’s accuracy differences on four datasets, while the sixth box (“Avg-Acc”) represents the distribution of average accuracy differences across all models.
  • Figure 3: Comparison of answer-level loss on GSM8K for Qwen2.5-7B under different training strategies. The plot shows the loss curves for baseline SFT, SFT-Tag, and Key-Tag, highlighting the effect of structured tagging and key-focused optimization on the model’s convergence and answer accuracy.