엔트로피-KL 발산 기반 토큰 마스킹: 대규모 언어 모델의 선택적 미세 조정에 대한 새로운 접근 방식

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

지도 학습 기반 미세 조정(SFT) 후 강화 학습(RL)은 대규모 언어 모델의 표준적인 사후 훈련 패러다임이 되었습니다. 이 패러다임은 RL 탐색을 위한 초기 단계를 제공하며, 순수한 RL 방식에서 발생하는 비효율성을 방지합니다. 순수한 RL에서는 온 정책 샘플링으로 인해 충분한 양수의 샘플을 얻기 어렵습니다. 그러나 실제로 기존 접근 방식은 종종 RL 단계에 비해 SFT 초기화에 사용되는 데이터의 양이 적어 모델이 제한된 샘플에 과적합되고 사전 훈련 분포에서 벗어나게 되는 경우가 많습니다. 이러한 분포 변화는 이후 RL 학습 중에 모델의 탐색 능력을 저해합니다. 이 문제를 해결하기 위해, 우리는 데이터가 부족한 환경에서는 SFT가 특정 내용을 암기하는 것보다 작업과 관련된 기능을 활성화하는 데 우선순위를 두어야 한다고 제안합니다. 이에 따라, 우리는 엔트로피-KL 선택적 미세 조정(EKSFT)을 제안하며, 이는 기준 모델과 높은 엔트로피 또는 높은 KL 발산을 보이는 토큰을 선택적으로 마스킹합니다. EKSFT는 이러한 불확실성이 높고 분포를 변화시키는 토큰을 모방에서 제외함으로써 작업별 지식을 주입하는 동시에 모델의 사전 훈련된 분포의 무결성을 유지합니다. 수학적 추론 벤치마크에 대한 실험 결과, EKSFT가 표준 SFT보다 꾸준히 우수한 성능을 보임을 보여줍니다. EKSFT 모델로부터 추가적인 RL 미세 조정을 수행하면 일관되게 더 나은 사후 RL 성능을 얻을 수 있으며, 이는 RL 단계에서의 개선된 탐색을 나타냅니다. 우리의 코드와 데이터셋은 https://github.com/MINE-USTC/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.