Entropy-KL Divergence-based Token Masking: A Novel Approach for Selective Fine-tuning of Large Language Models

2026-05-27

AI Summary

In supervised fine-tuning (SFT) with limited data, models are prone to overfitting on scarce samples, causing deviation from the pretraining distribution and impairing subsequent reinforcement learning (RL) exploration. This work proposes EKSFT, a novel approach that dynamically masks tokens with high uncertainty—quantified by token-level entropy and KL divergence relative to a reference model—to enable selective imitation. By doing so, EKSFT effectively injects task-specific knowledge while preserving the pretraining distribution. The method significantly improves distributional consistency during SFT and outperforms standard SFT on mathematical reasoning benchmarks. When used to initialize RL training, EKSFT leads to faster convergence and superior final performance, demonstrating its effectiveness in enhancing exploration efficiency.

Abstract

Supervised fine-tuning (SFT) followed by reinforcement learning (RL) has become a standard post-training paradigm for large language models. This paradigm provides a cold-start for RL exploration, avoiding the inefficiency of pure RL where on-policy sampling yields insufficient positive samples. However, in practice, existing approaches often use a small amount of data for SFT initialization compared to the RL phase, which can cause the model to fit the limited samples and shift away from its pre-trained distribution. This distribution shift impedes the model's ability to effectively explore during subsequent RL training. To address this challenge, we propose that in low-data regimes, SFT should prioritize activating task-relevant capabilities rather than memorizing specific content. Along this line, we propose EKSFT (Entropy-KL Selective Fine-Tuning), which selectively masks tokens that exhibit either high entropy or high KL divergence from a reference model. By excluding these high-uncertainty, distribution-shifting tokens from imitation, EKSFT injects task-specific knowledge while preserving the integrity of the model's pre-trained distribution. Empirical evaluations on mathematical reasoning benchmarks demonstrate that EKSFT consistently outperforms standard SFT. Further RL fine-tuning from the EKSFT model yields consistently better post-RL performance, indicating improved exploration for the RL stage. Our codes and datasets are available at https://github.com/MINE-USTC/EKSFT.

Problem

• distribution shift
• supervised fine-tuning
• low-data regime
• large language models
• reinforcement learning

Innovation

• Selective Fine-tuning
• Entropy-KL Divergence
• Token Masking
• Distribution Preservation
• Low-data SFT