KL 페널티를 무시하세요! 강화 학습 세부 조정을 향상시키기 위해 중요한 토큰에서 탐사를 촉진하세요.

Ignore the KL Penalty! Boosting Exploration on Critical Tokens to Enhance RL Fine-Tuning

February 10, 2025

저자: Jean Vassoyan, Nathanaël Beau, Roman Plaud

cs.AI

초록

현재 대규모 언어 모델(Large Language Models, LLMs)의 개발에서 장기 목표를 달성하는 능력은 주요 과제입니다. 이를 해결하기 위해 사전 훈련된 LLMs는 강화 학습(Reinforcement Learning, RL)을 사용하여 주어진 목표를 최적화하는 해결책을 탐색하기 위해 세밀하게 조정될 수 있습니다. 그러나 LLMs와의 탐사는 어렵습니다. 새로운 해결책을 발견하고 사전 훈련된 모델에 충분히 가까이 유지하여 기본 능력을 저하시키지 않도록 균형을 유지해야 합니다. 이는 일반적으로 Kullback-Leibler (KL) 페널티로 제어됩니다. 본 논문에서는 간단한 산술 작업에서 소규모 언어 모델의 탐사 역학을 조사합니다. 사전 훈련의 정도가 탐사에 미치는 영향을 보여주고 최종 결과에 미치는 영향이 큰 "중요 토큰"의 중요성을 입증합니다. 따라서, 우리는 KL 페널티에 간단한 수정을 도입하여 중요 토큰에 대한 탐사를 선호하고 RL 세밀 조정 단계의 효율성을 높이는 방법을 제시합니다.

English

The ability to achieve long-term goals is a key challenge in the current development of large language models (LLMs). To address this, pre-trained LLMs can be fine-tuned with reinforcement learning (RL) to explore solutions that optimize a given goal. However, exploration with LLMs is difficult, as a balance has to be struck between discovering new solutions and staying close enough to the pre-trained model, so as not to degrade basic capabilities. This is typically controlled with a Kullback-Leibler (KL) penalty. In this paper, we investigate the exploration dynamics of a small language model on a simple arithmetic task. We show how varying degrees of pre-training influence exploration and demonstrate the importance of "critical tokens" which have a dramatic impact on the final outcome. Consequently, we introduce a simple modification to the KL penalty that favors exploration on critical tokens, increasing the efficiency of the RL fine-tuning stage.