Negeer de KL-straf! Stimuleer verkenning van kritieke tokens om RL-fijnafstemming te verbeteren.

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

Samenvatting

Het vermogen om langetermijndoelen te bereiken is een belangrijke uitdaging in de huidige ontwikkeling van grote taalmodellen (LLM's). Om dit aan te pakken, kunnen vooraf getrainde LLM's worden verfijnd met reinforcement learning (RL) om oplossingen te verkennen die een gegeven doel optimaliseren. Echter, exploratie met LLM's is moeilijk, aangezien een balans moet worden gevonden tussen het ontdekken van nieuwe oplossingen en dicht genoeg bij het vooraf getrainde model blijven, zodat de basisvaardigheden niet worden aangetast. Dit wordt typisch gecontroleerd met een Kullback-Leibler (KL) straf. In dit artikel onderzoeken we de verkenningdynamiek van een klein taalmodel bij een eenvoudige rekenkundige taak. We laten zien hoe variërende gradaties van vooraf training de verkenning beïnvloeden en tonen het belang aan van "kritieke tokens" die een dramatische invloed hebben op de uiteindelijke uitkomst. Bijgevolg introduceren we een eenvoudige aanpassing aan de KL-straf die verkenning op kritieke tokens bevordert, waardoor de efficiëntie van de RL-verfijningsfase toeneemt.

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.