Rethinking Supervised Fine-Tuning: Emphasizing Key Answer Tokens for Improved LLM Accuracy
This addresses a specific bottleneck in LLM fine-tuning for complex reasoning tasks, offering a targeted improvement.
The paper tackles the problem of LLMs allocating too much attention to lengthy Chain-of-Thought sequences during Supervised Fine-Tuning, which reduces focus on essential answer tokens, by proposing SFTKey, a two-stage training scheme that fine-tunes only the Key portion in the second stage, resulting in an average accuracy improvement exceeding 5% over conventional SFT.
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.