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

📄 arXiv: 2605.29303v1 📥 PDF

作者: Qi Liu, Mingdi Sun, Yongyi He, Zhi Zheng, Tong Xu, Yi Zheng, Zhefeng Wang, Enhong Chen

分类: cs.AI

发布日期: 2026-05-28

备注: 17 pages

🔗 代码/项目: GITHUB

💡 一句话要点

提出基于熵-KL散度的token掩码方法EKSFT，用于大语言模型选择性微调。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 选择性微调 熵 KL散度 token掩码 强化学习 分布偏移

📋 核心要点

  1. 现有SFT方法在数据量较少时易过拟合，导致模型偏离预训练分布，影响后续RL阶段的探索效率。
  2. EKSFT通过掩码高熵或高KL散度的token，避免模型记忆特定内容，从而更好地保留预训练分布。
  3. 实验表明，EKSFT在数学推理任务上优于标准SFT，且能提升后续RL微调的性能。

📝 摘要（中文）

监督微调（SFT）和强化学习（RL）已成为大型语言模型后训练的标准范式。该范式为RL探索提供了一个冷启动，避免了纯RL中on-policy采样产生不足正样本的低效问题。然而，在实践中，与RL阶段相比，现有方法通常使用少量数据进行SFT初始化，这可能导致模型过度拟合有限的样本，并偏离其预训练分布。这种分布偏移阻碍了模型在后续RL训练中有效探索的能力。为了解决这个挑战，我们提出在低数据情况下，SFT应该优先激活任务相关的能力，而不是记忆特定的内容。因此，我们提出了EKSFT（Entropy-KL Selective Fine-Tuning），它选择性地掩盖那些表现出高熵或与参考模型具有高KL散度的token。通过从模仿中排除这些高不确定性、分布偏移的token，EKSFT注入了任务特定的知识，同时保持了模型预训练分布的完整性。在数学推理基准上的经验评估表明，EKSFT始终优于标准SFT。从EKSFT模型进一步进行RL微调可以获得始终更好的后RL性能，表明RL阶段的探索得到了改善。我们的代码和数据集可在https://github.com/MINE-USTC/EKSFT获得。

🔬 方法详解

问题定义：论文旨在解决大型语言模型在低数据量下进行监督微调（SFT）时容易过拟合的问题。现有方法在SFT阶段使用少量数据，导致模型偏离其预训练分布，从而影响后续强化学习（RL）阶段的探索效率。这种分布偏移使得模型难以有效地学习新的任务，并可能导致性能下降。

核心思路：论文的核心思路是，在低数据量的情况下，SFT应该更加注重激活模型中与任务相关的能力，而不是简单地记忆训练数据中的特定内容。为了实现这一目标，论文提出了一种选择性微调方法，该方法通过掩码那些对模型预训练分布影响较大的token，从而避免模型过度拟合训练数据。

技术框架：EKSFT (Entropy-KL Selective Fine-Tuning) 包含以下主要步骤：1) 使用少量任务相关数据进行SFT；2) 计算每个token的熵和KL散度，熵衡量模型对该token预测的不确定性，KL散度衡量SFT后的模型与预训练模型在该token上的分布差异；3) 根据熵和KL散度选择需要掩码的token；4) 使用掩码后的数据重新进行SFT。

关键创新：EKSFT的关键创新在于其选择性token掩码策略。与传统的SFT方法不同，EKSFT不是简单地对所有token进行微调，而是有选择地掩码那些对模型预训练分布影响较大的token。这种策略可以有效地避免模型过度拟合训练数据，并保留模型的泛化能力。

关键设计：EKSFT的关键设计包括：1) 熵的计算方法，通常使用模型预测概率分布的负熵；2) KL散度的计算方法，通常使用SFT后模型和预训练模型在token预测上的概率分布；3) 掩码token的选择策略，例如设置熵和KL散度的阈值，或者选择熵和KL散度最高的top-k个token；4) 掩码策略，例如将选中的token替换为特殊token [MASK]，或者直接从训练数据中删除这些token。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，EKSFT在数学推理基准上始终优于标准SFT。进一步的RL微调实验表明，从EKSFT模型出发进行RL训练，可以获得持续更好的后RL性能，这表明EKSFT能够改善RL阶段的探索。具体性能提升数据在论文中给出，证明了EKSFT的有效性。

🎯 应用场景

EKSFT方法可应用于各种需要对大型语言模型进行微调的场景，尤其是在数据量有限的情况下。例如，可以用于特定领域的文本生成、对话系统、代码生成等任务。该方法能够提升模型在低资源场景下的性能，并提高模型在后续强化学习阶段的探索效率，具有重要的实际应用价值。

📄 摘要（原文）

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.