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

TL;DR

This paper introduces SFTKey, a two-stage fine-tuning method that emphasizes key answer tokens in LLMs, significantly improving accuracy by balancing reasoning and answer focus.

Contribution

The paper proposes a novel two-stage fine-tuning approach that enhances LLM accuracy by focusing on answer-critical tokens, addressing attention imbalance in conventional SFT.

Findings

SFTKey improves average accuracy by over 5% across benchmarks.

It maintains correct output formats while enhancing answer correctness.

The method effectively balances reasoning and answer token focus.

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.