Ignorez la pénalité de KL ! Stimuler l'exploration des jetons critiques pour améliorer le peaufinage RL.

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

February 10, 2025

Auteurs: Jean Vassoyan, Nathanaël Beau, Roman Plaud

cs.AI

Résumé

La capacité à atteindre des objectifs à long terme est un défi majeur dans le développement actuel des grands modèles de langage (GML). Pour y remédier, les GML pré-entraînés peuvent être affinés avec l'apprentissage par renforcement (AR) pour explorer des solutions qui optimisent un objectif donné. Cependant, l'exploration avec les GML est difficile, car un équilibre doit être trouvé entre la découverte de nouvelles solutions et le maintien d'une proximité suffisante avec le modèle pré-entraîné, afin de ne pas dégrader les capacités de base. Cela est généralement contrôlé avec une pénalité de Kullback-Leibler (KL). Dans cet article, nous étudions la dynamique d'exploration d'un petit modèle de langage sur une tâche arithmétique simple. Nous montrons comment différents degrés de pré-entraînement influencent l'exploration et démontrons l'importance des "tokens critiques" qui ont un impact significatif sur le résultat final. Par conséquent, nous introduisons une modification simple de la pénalité de KL qui favorise l'exploration sur les tokens critiques, augmentant ainsi l'efficacité de l'étape d'affinage par AR.

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.