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

2025-12-24

AI Summary

In supervised fine-tuning (SFT), large language models (LLMs) often overemphasize lengthy chain-of-thought (CoT) reasoning, thereby weakening modeling of critical answer tokens and degrading answer accuracy. To address this, we propose SFTKey—a two-stage token-level fine-tuning method: Stage 1 ensures CoT format compliance via standard SFT; Stage 2 applies a weighted loss mask exclusively to answer tokens, explicitly decoupling and strengthening their optimization objective. This is the first work to introduce CoT-aware, token-level weighting and a phased answer-focusing mechanism in SFT. Experiments across multiple benchmarks (e.g., GSM8K, MATH, SVAMP) and model families (e.g., LLaMA-3, Qwen, Phi-3) demonstrate an average accuracy improvement of over 5%, with substantial gains in final answer correctness—while fully preserving CoT formatting fidelity and generation capability.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

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.

Problem

• Improves LLM accuracy by focusing on key answer tokens
• Addresses attention imbalance between CoT sequences and final answers
• Proposes a two-stage fine-tuning method to enhance answer correctness

Innovation

• Two-stage fine-tuning balances CoT and key answer tokens
• Second stage focuses only on essential answer portion
• Method improves accuracy while preserving output format