Enhancing Large Language Model Reasoning via Selective Critical Token Fine-Tuning
This addresses a bottleneck in adapting LLMs to domain-specific tasks like mathematical reasoning, offering a more efficient fine-tuning method with potential broad applicability.
The paper tackles the problem of standard supervised fine-tuning for large language models uniformly penalizing all tokens, which reduces output diversity and generalization, by proposing Critical Token Fine-tuning (CFT) that updates only critical tokens; experiments on mathematical reasoning benchmarks show CFT outperforms standard SFT while fine-tuning on less than 12% of tokens.
Large language models (LLMs) primarily rely on supervised fine-tuning (SFT) as a key method to adapt pre-trained models to domain-specific tasks such as mathematical reasoning. However, standard SFT uniformly penalizes all tokens, neglecting that only a small subset of critical tokens determines reasoning correctness. This uniform supervision often causes reduced output diversity and limited generalization. We propose Critical Token Fine-tuning (CFT), a simple yet effective approach that updates only tokens identified as functionally indispensable via counterfactual perturbations. By focusing gradient signals on these decisive reasoning steps while preserving the diversity of non-critical tokens, CFT can enhance both generation and diversity. Extensive experiments on five models across three families (Qwen, OLMo, LLaMA) and eleven mathematical reasoning benchmarks show that CFT, despite fine-tuning on less than 12% of tokens, consistently outperforms standard SFT. Moreover, CFT enables test-time scaling through improved sampling diversity and provides a stronger initialization for reinforcement learning, sustaining performance gains in later training stages while maintaining higher entropy for better exploration. These results highlight CFT as a practical and general framework for efficient and robust LLM fine-tuning.